thodg/got

diff --git a/got/got.c b/got/got.c
index ae5dabe..88e0035 100644
--- a/got/got.c
+++ b/got/got.c
@@ -3950,7 +3950,7 @@ cmd_tree(int argc, char *argv[])
 	const struct got_error *error;
 	struct got_repository *repo = NULL;
 	struct got_worktree *worktree = NULL;
-	const char *path;
+	const char *path, *refname = NULL;
 	char *cwd = NULL, *repo_path = NULL, *in_repo_path = NULL;
 	struct got_object_id *commit_id = NULL;
 	char *commit_id_str = NULL;
@@ -4057,7 +4057,11 @@ cmd_tree(int argc, char *argv[])
 
 	if (commit_id_str == NULL) {
 		struct got_reference *head_ref;
-		error = got_ref_open(&head_ref, repo, GOT_REF_HEAD, 0);
+		if (worktree)
+			refname = got_worktree_get_head_ref_name(worktree);
+		else
+			refname = GOT_REF_HEAD;
+		error = got_ref_open(&head_ref, repo, refname, 0);
 		if (error != NULL)
 			goto done;
 		error = got_ref_resolve(&commit_id, repo, head_ref);
diff --git a/regress/cmdline/Makefile b/regress/cmdline/Makefile
index 4a9108f..0f29c0e 100644
--- a/regress/cmdline/Makefile
+++ b/regress/cmdline/Makefile
@@ -1,6 +1,6 @@
 REGRESS_TARGETS=checkout update status log add rm diff blame branch tag \
 	ref commit revert cherrypick backout rebase import histedit \
-	integrate stage unstage cat clone fetch
+	integrate stage unstage cat clone fetch tree
 NOOBJ=Yes
 
 checkout:
@@ -75,4 +75,7 @@ clone:
 fetch:
 	./fetch.sh
 
+tree:
+	./tree.sh
+
 .include <bsd.regress.mk>
diff --git a/regress/cmdline/tree.sh b/regress/cmdline/tree.sh
new file mode 100644
index 0000000..9fded8d
--- /dev/null
+++ b/regress/cmdline/tree.sh
@@ -0,0 +1,81 @@
+#!/bin/sh
+#
+# Copyright (c) 2020 Tracey Emery <tracey@openbsd.org>
+#
+# Permission to use, copy, modify, and distribute this software for any
+# purpose with or without fee is hereby granted, provided that the above
+# copyright notice and this permission notice appear in all copies.
+#
+# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+
+. ./common.sh
+
+function test_tree_basic {
+	local testroot=`test_init tree_basic`
+
+	got checkout $testroot/repo $testroot/wt > /dev/null
+
+	echo "new file" > $testroot/wt/foo
+
+	(cd $testroot/wt && got add foo > /dev/null)
+	(cd $testroot/wt && got commit -m "add foo" foo >/dev/null)
+
+	echo 'alpha' > $testroot/stdout.expected
+	echo 'beta' >> $testroot/stdout.expected
+	echo 'epsilon/' >> $testroot/stdout.expected
+	echo 'foo' >> $testroot/stdout.expected
+	echo 'gamma/' >> $testroot/stdout.expected
+
+	(cd $testroot/wt && got tree > $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"
+}
+
+function test_tree_branch {
+	local testroot=`test_init tree_branch`
+
+	got checkout $testroot/repo $testroot/wt > /dev/null
+	ret="$?"
+	if [ "$ret" != "0" ]; then
+		test_done "$testroot" "$ret"
+		return 1
+	fi
+
+	(cd $testroot/wt && got br foo > $testroot/stdout)
+
+	echo "new file" > $testroot/wt/foo
+
+	(cd $testroot/wt && got add foo > /dev/null)
+	(cd $testroot/wt && got commit -m "add foo" foo >/dev/null)
+
+	echo 'alpha' > $testroot/stdout.expected
+	echo 'beta' >> $testroot/stdout.expected
+	echo 'epsilon/' >> $testroot/stdout.expected
+	echo 'foo' >> $testroot/stdout.expected
+	echo 'gamma/' >> $testroot/stdout.expected
+
+	(cd $testroot/wt && got tree > $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_tree_basic
+run_test test_tree_branch