thodg/got

diff --git a/got/got.c b/got/got.c
index b4af862..3bdf531 100644
--- a/got/got.c
+++ b/got/got.c
@@ -336,6 +336,58 @@ check_linear_ancestry(struct got_object_id *commit_id,
 	return NULL;
 }
 
+static const struct got_error *
+check_same_branch(struct got_object_id *commit_id,
+    struct got_reference *head_ref, struct got_repository *repo)
+{
+	const struct got_error *err = NULL;
+	struct got_commit_graph *graph = NULL;
+	struct got_object_id *head_commit_id = NULL;
+	int is_same_branch = 0;
+
+	err = got_ref_resolve(&head_commit_id, repo, head_ref);
+	if (err)
+		goto done;
+
+	err = got_commit_graph_open(&graph, head_commit_id, "/", 1, repo);
+	if (err)
+		goto done;
+
+	err = got_commit_graph_iter_start(graph, head_commit_id, repo);
+	if (err)
+		goto done;
+
+	for (;;) {
+		struct got_object_id *id;
+		err = got_commit_graph_iter_next(&id, graph);
+		if (err) {
+			if (err->code == GOT_ERR_ITER_COMPLETED) {
+				err = NULL;
+				break;
+			}
+			else if (err->code != GOT_ERR_ITER_NEED_MORE)
+				break;
+			err = got_commit_graph_fetch_commits(graph, 1,
+			    repo);
+			if (err)
+				break;
+		}
+
+		if (id) {
+			if (got_object_id_cmp(id, commit_id) == 0) {
+				is_same_branch = 1;
+				break;
+			}
+		}
+	}
+done:
+	if (graph)
+		got_commit_graph_close(graph);
+	free(head_commit_id);
+	if (!err && !is_same_branch)
+		err = got_error(GOT_ERR_ANCESTRY);
+	return err;
+}
 
 static const struct got_error *
 cmd_checkout(int argc, char *argv[])
@@ -615,15 +667,25 @@ cmd_update(int argc, char *argv[])
 		if (error)
 			goto done;
 		error = check_linear_ancestry(commit_id, head_commit_id, repo);
+		free(head_commit_id);
 		if (error != NULL)
 			goto done;
+		error = check_same_branch(commit_id, head_ref, repo);
+		if (error)
+			goto done;
 		error = got_worktree_set_head_ref(worktree, head_ref);
 		if (error)
 			goto done;
 	} else {
 		error = check_linear_ancestry(commit_id,
 		    got_worktree_get_base_commit_id(worktree), repo);
-		if (error != NULL)
+		if (error != NULL) {
+			if (error->code == GOT_ERR_ANCESTRY)
+				error = got_error(GOT_ERR_BRANCH_MOVED);
+			goto done;
+		}
+		error = check_same_branch(commit_id, head_ref, repo);
+		if (error)
 			goto done;
 	}
 
diff --git a/include/got_error.h b/include/got_error.h
index ac0dd60..a5adad9 100644
--- a/include/got_error.h
+++ b/include/got_error.h
@@ -90,6 +90,7 @@
 #define GOT_ERR_COMMIT_MSG_EMPTY 74
 #define GOT_ERR_DIR_NOT_EMPTY	75
 #define GOT_ERR_COMMIT_NO_CHANGES 76
+#define GOT_ERR_BRANCH_MOVED	77
 
 static const struct got_error {
 	int code;
@@ -148,7 +149,7 @@ static const struct got_error {
 	{ GOT_ERR_FILEIDX_CSUM,	"bad file index checksum" },
 	{ GOT_ERR_PATH_PREFIX,	"worktree already contains items from a "
 				"different path prefix" },
-	{ GOT_ERR_ANCESTRY,	"new branch or rebase required" },
+	{ GOT_ERR_ANCESTRY,	"target commit is on a different branch" },
 	{ GOT_ERR_FILEIDX_BAD,	"file index is corrupt" },
 	{ GOT_ERR_BAD_REF_DATA,	"could not parse reference data" },
 	{ GOT_ERR_TREE_DUP_ENTRY,"duplicate entry in tree object" },
@@ -171,6 +172,8 @@ static const struct got_error {
 	    "changes can be committed" },
 	{ GOT_ERR_COMMIT_MSG_EMPTY, "commit message cannot be empty" },
 	{ GOT_ERR_COMMIT_NO_CHANGES, "no changes to commit" },
+	{ GOT_ERR_BRANCH_MOVED,	"work tree's branch reference has moved; "
+	    "new branch or rebase required" },
 };
 
 /*
diff --git a/regress/cmdline/update.sh b/regress/cmdline/update.sh
index 6d4a991..afaf069 100755
--- a/regress/cmdline/update.sh
+++ b/regress/cmdline/update.sh
@@ -1264,7 +1264,9 @@ function test_update_moved_branch_ref {
 	(cd $testroot/repo2 && git fetch -q --all)
 
 	echo -n > $testroot/stdout.expected
-	echo "got: new branch or rebase required" >> $testroot/stderr.expected
+	echo -n "got: work tree's branch reference has moved; " \
+		> $testroot/stderr.expected
+	echo "new branch or rebase required" >> $testroot/stderr.expected
 
 	(cd $testroot/wt && got update > $testroot/stdout 2> $testroot/stderr)
 
@@ -1354,6 +1356,47 @@ function test_update_to_another_branch {
 	test_done "$testroot" "$ret"
 }
 
+function test_update_to_commit_on_wrong_branch {
+	local testroot=`test_init update_to_commit_on_wrong_branch`
+
+	got checkout $testroot/repo $testroot/wt > /dev/null
+	ret="$?"
+	if [ "$ret" != "0" ]; then
+		test_done "$testroot" "$ret"
+		return 1
+	fi
+
+	(cd $testroot/repo && git checkout -q -b newbranch)
+	echo "modified alpha on new branch" > $testroot/repo/alpha
+	git_commit $testroot/repo -m "modified alpha on new branch"
+
+	echo -n "" > $testroot/stdout.expected
+	echo  "got: target commit is on a different branch" \
+		> $testroot/stderr.expected
+
+	local head_rev=`git_show_head $testroot/repo`
+	(cd $testroot/wt && got update -c $head_rev > $testroot/stdout \
+		2> $testroot/stderr)
+
+	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
+
+	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
+
+	test_done "$testroot" "$ret"
+}
+
 run_test test_update_basic
 run_test test_update_adds_file
 run_test test_update_deletes_file
@@ -1381,3 +1424,4 @@ run_test test_update_partial_rm
 run_test test_update_partial_dir
 run_test test_update_moved_branch_ref
 run_test test_update_to_another_branch
+run_test test_update_to_commit_on_wrong_branch