thodg/got

diff --git a/include/got_error.h b/include/got_error.h
index 5e0b55d..8450cbb 100644
--- a/include/got_error.h
+++ b/include/got_error.h
@@ -231,7 +231,7 @@ static const struct got_error {
 	{ GOT_ERR_NO_MERGED_PATHS, "empty list of merged paths" },
 	{ GOT_ERR_COMMIT_BRANCH, "will not commit to a branch outside the "
 	    "\"refs/heads/\" reference namespace" },
-	{ GOT_ERR_FILE_STAGED, "file has staged changes" },
+	{ GOT_ERR_FILE_STAGED, "file is staged" },
 };
 
 /*
diff --git a/lib/worktree.c b/lib/worktree.c
index 9b859b2..08c378a 100644
--- a/lib/worktree.c
+++ b/lib/worktree.c
@@ -2639,10 +2639,10 @@ revert_file(struct got_worktree *worktree, struct got_fileindex *fileindex,
 	struct got_fileindex_entry *ie;
 	struct got_tree_object *tree = NULL;
 	struct got_object_id *tree_id = NULL;
-	const struct got_tree_entry *te;
+	const struct got_tree_entry *te = NULL;
 	char *tree_path = NULL, *te_name;
 	struct got_blob_object *blob = NULL;
-	unsigned char status;
+	unsigned char status, staged_status;
 	struct stat sb;
 
 	err = got_path_skip_common_ancestor(&relpath,
@@ -2695,11 +2695,19 @@ revert_file(struct got_worktree *worktree, struct got_fileindex *fileindex,
 	if (status == GOT_STATUS_MISSING || status == GOT_STATUS_DELETE)
 		sb.st_mode = got_fileindex_perms_to_st(ie);
 
+	staged_status = get_staged_status(ie);
+	if (status == GOT_STATUS_DELETE &&
+	    staged_status != GOT_STATUS_NO_CHANGE) {
+		err = got_error_path(ie->path, GOT_ERR_FILE_STAGED);
+		goto done;
+	}
+
 	err = got_object_id_by_path(&tree_id, repo, worktree->base_commit_id,
 	    tree_path);
 	if (err) {
 		if (!(err->code == GOT_ERR_NO_TREE_ENTRY &&
-		    status == GOT_STATUS_ADD))
+		    (status == GOT_STATUS_ADD ||
+		    staged_status == GOT_STATUS_ADD)))
 			goto done;
 	} else {
 		err = got_object_open_as_tree(&tree, repo, tree_id);
@@ -2713,7 +2721,8 @@ revert_file(struct got_worktree *worktree, struct got_fileindex *fileindex,
 		}
 
 		te = got_object_tree_find_entry(tree, te_name);
-		if (te == NULL && status != GOT_STATUS_ADD) {
+		if (te == NULL && status != GOT_STATUS_ADD &&
+		    staged_status != GOT_STATUS_ADD) {
 			err = got_error(GOT_ERR_NO_TREE_ENTRY);
 			goto done;
 		}
@@ -2731,13 +2740,19 @@ revert_file(struct got_worktree *worktree, struct got_fileindex *fileindex,
 	case GOT_STATUS_CONFLICT:
 	case GOT_STATUS_MISSING: {
 		struct got_object_id id;
-		memcpy(id.sha1, ie->blob_sha1, SHA1_DIGEST_LENGTH);
+		if (staged_status == GOT_STATUS_ADD ||
+		    staged_status == GOT_STATUS_MODIFY) {
+			memcpy(id.sha1, ie->staged_blob_sha1,
+			    SHA1_DIGEST_LENGTH);
+		} else
+			memcpy(id.sha1, ie->blob_sha1,
+			    SHA1_DIGEST_LENGTH);
 		err = got_object_open_as_blob(&blob, repo, &id, 8192);
 		if (err)
 			goto done;
 		err = install_blob(worktree, ondisk_path, ie->path,
-		    te->mode, sb.st_mode, blob, 0, 1, repo, progress_cb,
-		    progress_arg);
+		    te ? te->mode : GOT_DEFAULT_FILE_MODE, sb.st_mode,
+		    blob, 0, 1, repo, progress_cb, progress_arg);
 		if (err)
 			goto done;
 		if (status == GOT_STATUS_DELETE) {
diff --git a/regress/cmdline/stage.sh b/regress/cmdline/stage.sh
index 0631a5b..8cd9094 100755
--- a/regress/cmdline/stage.sh
+++ b/regress/cmdline/stage.sh
@@ -227,8 +227,7 @@ function test_stage_rm_already_staged_file {
 	fi
 
 	for f in alpha foo; do
-		echo "got: $f: file has staged changes" \
-			> $testroot/stderr.expected
+		echo "got: $f: file is staged" > $testroot/stderr.expected
 		(cd $testroot/wt && got rm $f \
 			> $testroot/stdout 2> $testroot/stderr)
 		ret="$?"
@@ -259,7 +258,189 @@ function test_stage_rm_already_staged_file {
 	test_done "$testroot" "$ret"
 }
 
+function test_stage_revert {
+	local testroot=`test_init stage_revert`
+
+	got checkout $testroot/repo $testroot/wt > /dev/null
+	ret="$?"
+	if [ "$ret" != "0" ]; then
+		test_done "$testroot" "$ret"
+		return 1
+	fi
+
+	echo "modified alpha" > $testroot/wt/alpha
+	(cd $testroot/wt && got rm beta > /dev/null)
+	echo "new file" > $testroot/wt/foo
+	(cd $testroot/wt && got add foo > /dev/null)
+	(cd $testroot/wt && got stage alpha beta foo > /dev/null)
+
+	echo "modified file again" >> $testroot/wt/alpha
+	echo "modified added file again" >> $testroot/wt/foo
+
+	(cd $testroot/wt && got revert alpha > $testroot/stdout)
+	ret="$?"
+	if [ "$ret" != "0" ]; then
+		echo "revert command failed unexpectedly" >&2
+		test_done "$testroot" "$ret"
+		return 1
+	fi
+
+	echo "R  alpha" > $testroot/stdout.expected
+	cmp -s $testroot/stdout.expected $testroot/stdout
+	ret="$?"
+	if [ "$ret" != "0" ]; then
+		diff -u $testroot/stdout.expected $testroot/stdout
+		test_done "$testroot" "$ret"
+		return 1
+	fi
+
+	echo "modified alpha" > $testroot/content.expected
+	cat $testroot/wt/alpha > $testroot/content
+	cmp -s $testroot/content.expected $testroot/content
+	ret="$?"
+	if [ "$ret" != "0" ]; then
+		diff -u $testroot/content.expected $testroot/content
+		test_done "$testroot" "$ret"
+		return 1
+	fi
+
+	echo ' M alpha' > $testroot/stdout.expected
+	echo ' D beta' >> $testroot/stdout.expected
+	echo 'MA foo' >> $testroot/stdout.expected
+	(cd $testroot/wt && got status > $testroot/stdout)
+	cmp -s $testroot/stdout.expected $testroot/stdout
+	ret="$?"
+	if [ "$ret" != "0" ]; then
+		diff -u $testroot/stdout.expected $testroot/stdout
+		test_done "$testroot" "$ret"
+		return 1
+	fi
+
+	(cd $testroot/wt && got revert alpha > $testroot/stdout)
+	ret="$?"
+	if [ "$ret" != "0" ]; then
+		echo "revert command failed unexpectedly" >&2
+		test_done "$testroot" "$ret"
+		return 1
+	fi
+
+	echo -n > $testroot/stdout.expected
+	cmp -s $testroot/stdout.expected $testroot/stdout
+	ret="$?"
+	if [ "$ret" != "0" ]; then
+		diff -u $testroot/stdout.expected $testroot/stdout
+		test_done "$testroot" "$ret"
+		return 1
+	fi
+
+	echo "modified alpha" > $testroot/content.expected
+	cat $testroot/wt/alpha > $testroot/content
+	cmp -s $testroot/content.expected $testroot/content
+	ret="$?"
+	if [ "$ret" != "0" ]; then
+		diff -u $testroot/content.expected $testroot/content
+		test_done "$testroot" "$ret"
+		return 1
+	fi
+
+	(cd $testroot/wt && got revert beta > $testroot/stdout \
+		2> $testroot/stderr)
+	ret="$?"
+	if [ "$ret" == "0" ]; then
+		echo "revert command succeeded unexpectedly" >&2
+		test_done "$testroot" "1"
+		return 1
+	fi
+
+	echo "got: beta: file is staged" > $testroot/stderr.expected
+	cmp -s $testroot/stderr.expected $testroot/stderr
+	ret="$?"
+	if [ "$ret" != "0" ]; then
+		diff -u $testroot/stderr.expected $testroot/stderr
+		test_done "$testroot" "$ret"
+		return 1
+	fi
+
+	(cd $testroot/wt && got revert foo > $testroot/stdout)
+	ret="$?"
+	if [ "$ret" != "0" ]; then
+		echo "revert command failed unexpectedly" >&2
+		test_done "$testroot" "$ret"
+		return 1
+	fi
+
+	echo "R  foo" > $testroot/stdout.expected
+	cmp -s $testroot/stdout.expected $testroot/stdout
+	ret="$?"
+	if [ "$ret" != "0" ]; then
+		diff -u $testroot/stdout.expected $testroot/stdout
+		test_done "$testroot" "$ret"
+		return 1
+	fi
+
+	echo "new file" > $testroot/content.expected
+	cat $testroot/wt/foo > $testroot/content
+	cmp -s $testroot/content.expected $testroot/content
+	ret="$?"
+	if [ "$ret" != "0" ]; then
+		diff -u $testroot/content.expected $testroot/content
+		test_done "$testroot" "$ret"
+		return 1
+	fi
+
+	echo ' M alpha' > $testroot/stdout.expected
+	echo ' D beta' >> $testroot/stdout.expected
+	echo ' A foo' >> $testroot/stdout.expected
+	(cd $testroot/wt && got status > $testroot/stdout)
+	cmp -s $testroot/stdout.expected $testroot/stdout
+	ret="$?"
+	if [ "$ret" != "0" ]; then
+		diff -u $testroot/stdout.expected $testroot/stdout
+		test_done "$testroot" "$ret"
+		return 1
+	fi
+
+	(cd $testroot/wt && got revert foo > $testroot/stdout)
+	ret="$?"
+	if [ "$ret" != "0" ]; then
+		echo "revert command failed unexpectedly" >&2
+		test_done "$testroot" "$ret"
+		return 1
+	fi
+
+	echo -n > $testroot/stdout.expected
+	cmp -s $testroot/stdout.expected $testroot/stdout
+	ret="$?"
+	if [ "$ret" != "0" ]; then
+		diff -u $testroot/stdout.expected $testroot/stdout
+		test_done "$testroot" "$ret"
+		return 1
+	fi
+
+	echo "new file" > $testroot/content.expected
+	cat $testroot/wt/foo > $testroot/content
+	cmp -s $testroot/content.expected $testroot/content
+	ret="$?"
+	if [ "$ret" != "0" ]; then
+		diff -u $testroot/content.expected $testroot/content
+		test_done "$testroot" "$ret"
+		return 1
+	fi
+
+	echo ' M alpha' > $testroot/stdout.expected
+	echo ' D beta' >> $testroot/stdout.expected
+	echo ' A foo' >> $testroot/stdout.expected
+	(cd $testroot/wt && got status > $testroot/stdout)
+	cmp -s $testroot/stdout.expected $testroot/stdout
+	ret="$?"
+	if [ "$ret" != "0" ]; then
+		diff -u $testroot/stdout.expected $testroot/stdout
+	fi
+	test_done "$testroot" "$ret"
+}
+
 run_test test_stage_basic
 run_test test_stage_status
 run_test test_stage_add_already_staged_file
 run_test test_stage_rm_already_staged_file
+run_test test_stage_revert