thodg/got

diff --git a/got/got.c b/got/got.c
index 8dddba2..16ce520 100644
--- a/got/got.c
+++ b/got/got.c
@@ -1681,7 +1681,7 @@ done:
 __dead static void
 usage_diff(void)
 {
-	fprintf(stderr, "usage: %s diff [-C number] [-r repository-path] "
+	fprintf(stderr, "usage: %s diff [-C number] [-r repository-path] [-s] "
 	    "[object1 object2 | path]\n", getprogname());
 	exit(1);
 }
@@ -1692,6 +1692,7 @@ struct print_diff_arg {
 	int diff_context;
 	const char *id_str;
 	int header_shown;
+	int diff_staged;
 };
 
 static const struct got_error *
@@ -1706,19 +1707,48 @@ print_diff(void *arg, unsigned char status, unsigned char staged_status,
 	char *abspath = NULL;
 	struct stat sb;
 
-	if (staged_status == GOT_STATUS_DELETE)
-		return NULL;
-
-	if (status != GOT_STATUS_MODIFY && status != GOT_STATUS_ADD &&
-	    status != GOT_STATUS_DELETE && status != GOT_STATUS_CONFLICT)
-		return NULL;
+	if (a->diff_staged) {
+		if (staged_status != GOT_STATUS_MODIFY &&
+		    staged_status != GOT_STATUS_ADD &&
+		    staged_status != GOT_STATUS_DELETE)
+			return NULL;
+	} else {
+		if (staged_status == GOT_STATUS_DELETE)
+			return NULL;
+		if (status != GOT_STATUS_MODIFY &&
+		    status != GOT_STATUS_ADD &&
+		    status != GOT_STATUS_DELETE &&
+		    status != GOT_STATUS_CONFLICT)
+			return NULL;
+	}
 
 	if (!a->header_shown) {
-		printf("diff %s %s\n", a->id_str,
-		    got_worktree_get_root_path(a->worktree));
+		printf("diff %s %s%s\n", a->id_str,
+		    got_worktree_get_root_path(a->worktree),
+		    a->diff_staged ? " (staged changes)" : "");
 		a->header_shown = 1;
 	}
 
+	if (a->diff_staged) {
+		const char *label1 = NULL, *label2 = NULL;
+		switch (staged_status) {
+		case GOT_STATUS_MODIFY:
+			label1 = path;
+			label2 = path;
+			break;
+		case GOT_STATUS_ADD:
+			label2 = path;
+			break;
+		case GOT_STATUS_DELETE:
+			label1 = path;
+			break;
+		default:
+			return got_error(GOT_ERR_FILE_STATUS);
+		}
+		return got_diff_objects_as_blobs(blob_id, staged_blob_id,
+		    label1, label2, a->diff_context, a->repo, stdout);
+	}
+
 	if (staged_status == GOT_STATUS_ADD ||
 	    staged_status == GOT_STATUS_MODIFY)
 		err = got_object_open_as_blob(&blob1, a->repo, staged_blob_id,
@@ -1769,7 +1799,7 @@ cmd_diff(int argc, char *argv[])
 	const char *id_str1 = NULL, *id_str2 = NULL;
 	char *label1 = NULL, *label2 = NULL;
 	int type1, type2;
-	int diff_context = 3, ch;
+	int diff_context = 3, diff_staged = 0, ch;
 	const char *errstr;
 	char *path = NULL;
 
@@ -1779,7 +1809,7 @@ cmd_diff(int argc, char *argv[])
 		err(1, "pledge");
 #endif
 
-	while ((ch = getopt(argc, argv, "C:r:")) != -1) {
+	while ((ch = getopt(argc, argv, "C:r:s")) != -1) {
 		switch (ch) {
 		case 'C':
 			diff_context = strtonum(optarg, 1, INT_MAX, &errstr);
@@ -1792,6 +1822,9 @@ cmd_diff(int argc, char *argv[])
 				err(1, "-r option");
 			got_path_strip_trailing_slashes(repo_path);
 			break;
+		case 's':
+			diff_staged = 1;
+			break;
 		default:
 			usage_diff();
 			/* NOTREACHED */
@@ -1835,6 +1868,9 @@ cmd_diff(int argc, char *argv[])
 			}
 		}
 	} else if (argc == 2) {
+		if (diff_staged)
+			errx(1, "-s option can't be used when diffing "
+			    "objects in repository");
 		id_str1 = argv[0];
 		id_str2 = argv[1];
 		if (worktree && repo_path == NULL) {
@@ -1880,6 +1916,7 @@ cmd_diff(int argc, char *argv[])
 		arg.diff_context = diff_context;
 		arg.id_str = id_str;
 		arg.header_shown = 0;
+		arg.diff_staged = diff_staged;
 
 		error = got_pathlist_append(&paths, path, NULL);
 		if (error)
diff --git a/lib/worktree.c b/lib/worktree.c
index 67da6f9..e5ed395 100644
--- a/lib/worktree.c
+++ b/lib/worktree.c
@@ -2211,23 +2211,34 @@ report_file_status(struct got_fileindex_entry *ie, const char *abspath,
 	unsigned char staged_status = get_staged_status(ie);
 	struct stat sb;
 	struct got_object_id blob_id, commit_id, staged_blob_id;
+	struct got_object_id *blob_idp = NULL, *commit_idp = NULL;
+	struct got_object_id *staged_blob_idp = NULL;
 
 	err = get_file_status(&status, &sb, ie, abspath, repo);
-	if (err == NULL && (status != GOT_STATUS_NO_CHANGE ||
-	    staged_status != GOT_STATUS_NO_CHANGE)) {
+	if (err)
+		return err;
+
+	if (status == GOT_STATUS_NO_CHANGE &&
+	    staged_status == GOT_STATUS_NO_CHANGE)
+		return NULL;
+
+	if (got_fileindex_entry_has_blob(ie)) {
 		memcpy(blob_id.sha1, ie->blob_sha1, SHA1_DIGEST_LENGTH);
+		blob_idp = &blob_id;
+	}
+	if (got_fileindex_entry_has_commit(ie)) {
 		memcpy(commit_id.sha1, ie->commit_sha1, SHA1_DIGEST_LENGTH);
-		if (staged_status == GOT_STATUS_ADD ||
-		    staged_status == GOT_STATUS_MODIFY) {
-			memcpy(staged_blob_id.sha1, ie->staged_blob_sha1,
-			    SHA1_DIGEST_LENGTH);
-			err = (*status_cb)(status_arg, status, staged_status,
-			    ie->path, &blob_id, &staged_blob_id, &commit_id);
-		} else
-			err = (*status_cb)(status_arg, status, staged_status,
-			    ie->path, &blob_id, NULL, &commit_id);
+		commit_idp = &commit_id;
 	}
-	return err;
+	if (staged_status == GOT_STATUS_ADD ||
+	    staged_status == GOT_STATUS_MODIFY) {
+		memcpy(staged_blob_id.sha1, ie->staged_blob_sha1,
+		    SHA1_DIGEST_LENGTH);
+		staged_blob_idp = &staged_blob_id;
+	}
+
+	return (*status_cb)(status_arg, status, staged_status,
+	    ie->path, blob_idp, staged_blob_idp, commit_idp);
 }
 
 static const struct got_error *
diff --git a/regress/cmdline/stage.sh b/regress/cmdline/stage.sh
index 7003e97..2b77027 100755
--- a/regress/cmdline/stage.sh
+++ b/regress/cmdline/stage.sh
@@ -532,6 +532,16 @@ function test_stage_diff {
 	echo "new file" > $testroot/wt/foo
 	(cd $testroot/wt && got add foo > /dev/null)
 
+	(cd $testroot/wt && got diff -s > $testroot/stdout)
+	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 ' M alpha' > $testroot/stdout.expected
 	echo ' D beta' >> $testroot/stdout.expected
 	echo ' A foo' >> $testroot/stdout.expected
@@ -576,6 +586,46 @@ function test_stage_diff {
 	ret="$?"
 	if [ "$ret" != "0" ]; then
 		diff -u $testroot/stdout.expected $testroot/stdout
+		test_done "$testroot" "$ret"
+		return 1
+	fi
+
+	(cd $testroot/wt && got diff -s > $testroot/stdout)
+
+	echo "diff $head_commit $testroot/wt (staged changes)" \
+		> $testroot/stdout.expected
+	echo -n 'blob - ' >> $testroot/stdout.expected
+	got tree -r $testroot/repo -i | grep 'alpha$' | cut -d' ' -f 1 \
+		>> $testroot/stdout.expected
+	echo -n 'blob + ' >> $testroot/stdout.expected
+	(cd $testroot/wt && got stage -l alpha) | cut -d' ' -f 1 \
+		>> $testroot/stdout.expected
+	echo '--- alpha' >> $testroot/stdout.expected
+	echo '+++ alpha' >> $testroot/stdout.expected
+	echo '@@ -1 +1 @@' >> $testroot/stdout.expected
+	echo '-alpha' >> $testroot/stdout.expected
+	echo '+modified file' >> $testroot/stdout.expected
+	echo -n 'blob - ' >> $testroot/stdout.expected
+	got tree -r $testroot/repo -i | grep 'beta$' | cut -d' ' -f 1 \
+		>> $testroot/stdout.expected
+	echo 'blob + /dev/null' >> $testroot/stdout.expected
+	echo '--- beta' >> $testroot/stdout.expected
+	echo '+++ /dev/null' >> $testroot/stdout.expected
+	echo '@@ -1 +0,0 @@' >> $testroot/stdout.expected
+	echo '-beta' >> $testroot/stdout.expected
+	echo 'blob - /dev/null' >> $testroot/stdout.expected
+	echo -n 'blob + ' >> $testroot/stdout.expected
+	(cd $testroot/wt && got stage -l foo) | cut -d' ' -f 1 \
+		>> $testroot/stdout.expected
+	echo '--- /dev/null' >> $testroot/stdout.expected
+	echo '+++ foo' >> $testroot/stdout.expected
+	echo '@@ -0,0 +1 @@' >> $testroot/stdout.expected
+	echo '+new file' >> $testroot/stdout.expected
+
+	cmp -s $testroot/stdout.expected $testroot/stdout
+	ret="$?"
+	if [ "$ret" != "0" ]; then
+		diff -u $testroot/stdout.expected $testroot/stdout
 	fi
 	test_done "$testroot" "$ret"