thodg/got

diff --git a/got/got.1 b/got/got.1
index f7751d1..32d8fc7 100644
--- a/got/got.1
+++ b/got/got.1
@@ -532,8 +532,8 @@ which must be an existing reference.
 Care should be taken not to create loops between references when
 this option is used.
 .El
-.It Cm branch Oo Fl c Ar commit Oc Oo Fl r Ar repository-path Oc Oo Fl l Oc Oo Fl d Ar name Oc Op Ar name
-Manage branches in a repository.
+.It Cm branch Oo Fl c Ar commit Oc Oo Fl r Ar repository-path Oc Oo Fl l Oc Oo Fl d Ar name Oc Oo Fl n Oc Op Ar name
+Create, list, or delete branches.
 .Pp
 Branches are managed via references which live in the
 .Dq refs/heads/
@@ -544,6 +544,7 @@ command operates on references in this namespace only.
 .Pp
 If invoked in a work tree without any arguments, print the name of the
 work tree's current branch.
+.Pp
 If a
 .Ar name
 argument is passed, attempt to create a branch reference with the given name.
@@ -551,6 +552,22 @@ By default the new branch reference will point at the latest commit on the
 work tree's current branch if invoked in a work tree, and otherwise to a commit
 resolved via the repository's HEAD reference.
 .Pp
+If invoked in a work tree, once the branch was created successfully
+switch the work tree's head reference to the newly created branch and
+update files across the entire work tree, just like
+.Cm got update -b Ar name
+would do.
+Show the status of each affected file, using the following status codes:
+.Bl -column YXZ description
+.It U Ta file was updated and contained no local changes
+.It G Ta file was updated and local changes were merged cleanly
+.It C Ta file was updated and conflicts occurred during merge
+.It D Ta file was deleted
+.It A Ta new file was added
+.It \(a~ Ta versioned file is obstructed by a non-regular file
+.It ! Ta a missing versioned file was restored
+.El
+.Pp
 The options for
 .Cm got branch
 are as follows:
@@ -583,6 +600,8 @@ Only the branch reference is deleted.
 Any commit, tree, and blob objects belonging to the branch
 remain in the repository and may be removed separately with
 Git's garbage collector.
+.It Fl n
+Do not switch and update the work tree after creating a new branch.
 .El
 .It Cm br
 Short alias for
diff --git a/got/got.c b/got/got.c
index 7a50d3a..6511394 100644
--- a/got/got.c
+++ b/got/got.c
@@ -3284,8 +3284,8 @@ __dead static void
 usage_branch(void)
 {
 	fprintf(stderr,
-	    "usage: %s branch [-c commit] [-r repository] [-l] | -d name | "
-	    "[name]\n", getprogname());
+	    "usage: %s branch [-c commit] [-d] [-r repository] [-l] [-n] "
+	        "[name]\n", getprogname());
 	exit(1);
 }
 
@@ -3470,10 +3470,17 @@ cmd_branch(int argc, char *argv[])
 	struct got_repository *repo = NULL;
 	struct got_worktree *worktree = NULL;
 	char *cwd = NULL, *repo_path = NULL;
-	int ch, do_list = 0, do_show = 0;
+	int ch, do_list = 0, do_show = 0, do_update = 1;
 	const char *delref = NULL, *commit_id_arg = NULL;
+	struct got_reference *ref = NULL;
+	struct got_pathlist_head paths;
+	struct got_pathlist_entry *pe;
+	struct got_object_id *commit_id = NULL;
+	char *commit_id_str = NULL;
 
-	while ((ch = getopt(argc, argv, "c:d:r:l")) != -1) {
+	TAILQ_INIT(&paths);
+
+	while ((ch = getopt(argc, argv, "c:d:r:ln")) != -1) {
 		switch (ch) {
 		case 'c':
 			commit_id_arg = optarg;
@@ -3491,6 +3498,9 @@ cmd_branch(int argc, char *argv[])
 		case 'l':
 			do_list = 1;
 			break;
+		case 'n':
+			do_update = 0;
+			break;
 		default:
 			usage_branch();
 			/* NOTREACHED */
@@ -3570,7 +3580,6 @@ cmd_branch(int argc, char *argv[])
 	else if (delref)
 		error = delete_branch(repo, worktree, delref);
 	else {
-		struct got_object_id *commit_id;
 		if (commit_id_arg == NULL)
 			commit_id_arg = worktree ?
 			    got_worktree_get_head_ref_name(worktree) :
@@ -3580,15 +3589,59 @@ cmd_branch(int argc, char *argv[])
 		if (error)
 			goto done;
 		error = add_branch(repo, argv[0], commit_id);
-		free(commit_id);
+		if (error)
+			goto done;
+		if (worktree && do_update) {
+			int did_something = 0;
+			char *branch_refname = NULL;
+
+			error = got_object_id_str(&commit_id_str, commit_id);
+			if (error)
+				goto done;
+			error = get_worktree_paths_from_argv(&paths, 0, NULL,
+			    worktree);
+			if (error)
+				goto done;
+			if (asprintf(&branch_refname, "refs/heads/%s", argv[0])
+			    == -1) {
+				error = got_error_from_errno("asprintf");
+				goto done;
+			}
+			error = got_ref_open(&ref, repo, branch_refname, 0);
+			free(branch_refname);
+			if (error)
+				goto done;
+			error = switch_head_ref(ref, commit_id, worktree,
+			    repo);
+			if (error)
+				goto done;
+			error = got_worktree_set_base_commit_id(worktree, repo,
+			    commit_id);
+			if (error)
+				goto done;
+			error = got_worktree_checkout_files(worktree, &paths,
+			    repo, update_progress, &did_something,
+			    check_cancelled, NULL);
+			if (error)
+				goto done;
+			if (did_something)
+				printf("Updated to commit %s\n", commit_id_str);
+		}
 	}
 done:
+	if (ref)
+		got_ref_close(ref);
 	if (repo)
 		got_repo_close(repo);
 	if (worktree)
 		got_worktree_close(worktree);
 	free(cwd);
 	free(repo_path);
+	free(commit_id);
+	free(commit_id_str);
+	TAILQ_FOREACH(pe, &paths, entry)
+		free((char *)pe->path);
+	got_pathlist_free(&paths);
 	return error;
 }
 
diff --git a/regress/cmdline/branch.sh b/regress/cmdline/branch.sh
index 2a7d11e..c079325 100755
--- a/regress/cmdline/branch.sh
+++ b/regress/cmdline/branch.sh
@@ -18,6 +18,7 @@
 
 function test_branch_create {
 	local testroot=`test_init branch_create`
+	local commit_id0=`git_show_head $testroot/repo`
 
 	# Create a branch based on repository's HEAD reference
 	got branch -r $testroot/repo newbranch
@@ -58,7 +59,7 @@ function test_branch_create {
 	fi
 
 	# Create a branch based on the work tree's branch
-	(cd $testroot/wt && got branch anotherbranch)
+	(cd $testroot/wt && got branch -n anotherbranch)
 	ret="$?"
 	if [ "$ret" != "0" ]; then
 		test_done "$testroot" "$ret"
@@ -74,7 +75,7 @@ function test_branch_create {
 	fi
 
 	# Create a branch based on another specific branch
-	(cd $testroot/wt && got branch -c master yetanotherbranch)
+	(cd $testroot/wt && got branch -n -c master yetanotherbranch)
 	ret="$?"
 	if [ "$ret" != "0" ]; then
 		test_done "$testroot" "$ret"
@@ -103,6 +104,24 @@ function test_branch_create {
 	ret="$?"
 	if [ "$ret" != "0" ]; then
 		echo "git checkout command failed unexpectedly"
+		test_done "$testroot" "$ret"
+		return 1
+	fi
+
+	# Create a branch and let the work tree be updated to it
+	(cd $testroot/wt && got branch -c $commit_id0 updatebranch \
+		> $testroot/stdout)
+
+	echo -n "Switching work tree from refs/heads/newbranch to " \
+		> $testroot/stdout.expected
+	echo "refs/heads/updatebranch" >> $testroot/stdout.expected
+	echo "U  gamma/delta" >> $testroot/stdout.expected
+	echo "Updated to commit $commit_id0" >> $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"
 }
diff --git a/regress/cmdline/rebase.sh b/regress/cmdline/rebase.sh
index 3c441f1..0537513 100755
--- a/regress/cmdline/rebase.sh
+++ b/regress/cmdline/rebase.sh
@@ -741,7 +741,7 @@ function test_rebase_no_commits_to_rebase {
 		return 1
 	fi
 
-	(cd $testroot/wt && got branch newbranch)
+	(cd $testroot/wt && got branch -n newbranch)
 
 	echo "modified alpha on master" > $testroot/wt/alpha
 	(cd $testroot/wt && got commit -m 'test rebase_no_commits_to_rebase' \
@@ -848,7 +848,7 @@ function test_rebase_forward {
 		return 1
 	fi
 
-	(cd $testroot/wt && got branch > $testroot/stdout)
+	(cd $testroot/wt && got branch -n > $testroot/stdout)
 	echo "master" > $testroot/stdout.expected
 	cmp -s $testroot/stdout.expected $testroot/stdout
 	ret="$?"