Merge pull request #4191 from pks-t/pks/wt-ref-renames

Branch renames with worktrees

7 files changed, 257 insertions, 124 deletions

diff --git a/src/branch.c b/src/branch.c
index 7d5e9cb7f..7a83b83af 100644
--- a/src/branch.c
+++ b/src/branch.c
@@ -127,62 +127,29 @@ int git_branch_create_from_annotated(
 		repository, branch_name, commit->commit, commit->description, force);
 }
 
-int git_branch_is_checked_out(
-	const git_reference *branch)
+static int branch_equals(git_repository *repo, const char *path, void *payload)
 {
-	git_buf path = GIT_BUF_INIT, buf = GIT_BUF_INIT;
-	git_strarray worktrees;
-	git_reference *ref = NULL;
-	git_repository *repo;
-	const char *worktree;
-	int found = false;
-	size_t i;
-
-	assert(branch && git_reference_is_branch(branch));
-
-	repo = git_reference_owner(branch);
-
-	if (git_worktree_list(&worktrees, repo) < 0)
-		return -1;
-
-	for (i = 0; i < worktrees.count; i++) {
-		worktree = worktrees.strings[i];
-
-		if (git_repository_head_for_worktree(&ref, repo, worktree) < 0)
-			continue;
-
-		if (git__strcmp(ref->name, branch->name) == 0) {
-			found = true;
-			git_reference_free(ref);
-			break;
-		}
-
-		git_reference_free(ref);
-	}
-	git_strarray_free(&worktrees);
-
-	if (found)
-		return found;
+	git_reference *branch = (git_reference *) payload;
+	git_reference *head;
+	int equal;
 
-	/* Check HEAD of parent */
-	if (git_buf_joinpath(&path, repo->commondir, GIT_HEAD_FILE) < 0)
-		goto out;
-	if (git_futils_readbuffer(&buf, path.ptr) < 0)
-		goto out;
-	if (git__prefixcmp(buf.ptr, "ref: ") == 0)
-		git_buf_consume(&buf, buf.ptr + strlen("ref: "));
-	git_buf_rtrim(&buf);
+	if (git_reference__read_head(&head, repo, path) < 0 ||
+	    git_reference_type(head) != GIT_REF_SYMBOLIC)
+		return 0;
 
-	found = git__strcmp(buf.ptr, branch->name) == 0;
+	equal = !git__strcmp(head->target.symbolic, branch->name);
+	git_reference_free(head);
+	return equal;
+}
 
-out:
-	git_buf_free(&buf);
-	git_buf_free(&path);
+int git_branch_is_checked_out(const git_reference *branch)
+{
+	assert(branch && git_reference_is_branch(branch));
 
-	return found;
+	return git_repository_foreach_head(git_reference_owner(branch),
+		branch_equals, (void *) branch) == 1;
 }
 
-
 int git_branch_delete(git_reference *branch)
 {
 	int is_head;
diff --git a/src/refs.c b/src/refs.c
index 0837dc4a8..31410b75f 100644
--- a/src/refs.c
+++ b/src/refs.c
@@ -249,6 +249,40 @@ int git_reference_lookup_resolved(
 	return 0;
 }
 
+int git_reference__read_head(
+	git_reference **out,
+	git_repository *repo,
+	const char *path)
+{
+	git_buf reference = GIT_BUF_INIT;
+	char *name = NULL;
+	int error;
+
+	if ((error = git_futils_readbuffer(&reference, path)) < 0)
+		goto out;
+	git_buf_rtrim(&reference);
+
+	if (git__strncmp(reference.ptr, GIT_SYMREF, strlen(GIT_SYMREF)) == 0) {
+		git_buf_consume(&reference, reference.ptr + strlen(GIT_SYMREF));
+
+		name = git_path_basename(path);
+
+		if ((*out = git_reference__alloc_symbolic(name, reference.ptr)) == NULL) {
+			error = -1;
+			goto out;
+		}
+	} else {
+		if ((error = git_reference_lookup(out, repo, reference.ptr)) < 0)
+			goto out;
+	}
+
+out:
+	free(name);
+	git_buf_clear(&reference);
+
+	return error;
+}
+
 int git_reference_dwim(git_reference **out, git_repository *repo, const char *refname)
 {
 	int error = 0, i;
@@ -580,20 +614,53 @@ int git_reference_symbolic_set_target(
 		out, ref->db->repo, ref->name, target, 1, ref->target.symbolic, log_message);
 }
 
+typedef struct {
+    const char *old_name;
+    git_refname_t new_name;
+} rename_cb_data;
+
+static int update_wt_heads(git_repository *repo, const char *path, void *payload)
+{
+	rename_cb_data *data = (rename_cb_data *) payload;
+	git_reference *head;
+	char *gitdir = NULL;
+	int error = 0;
+
+	if (git_reference__read_head(&head, repo, path) < 0 ||
+	    git_reference_type(head) != GIT_REF_SYMBOLIC ||
+	    git__strcmp(head->target.symbolic, data->old_name) != 0 ||
+	    (gitdir = git_path_dirname(path)) == NULL)
+		goto out;
+
+	/* Update HEAD it was pointing to the reference being renamed */
+	if ((error = git_repository_create_head(gitdir, data->new_name)) < 0) {
+		giterr_set(GITERR_REFERENCE, "failed to update HEAD after renaming reference");
+		goto out;
+	}
+
+out:
+	git_reference_free(head);
+	git__free(gitdir);
+
+	return error;
+}
+
 static int reference__rename(git_reference **out, git_reference *ref, const char *new_name, int force,
 				 const git_signature *signature, const char *message)
 {
+	git_repository *repo;
 	git_refname_t normalized;
 	bool should_head_be_updated = false;
 	int error = 0;
 
 	assert(ref && new_name && signature);
 
+	repo = git_reference_owner(ref);
+
 	if ((error = reference_normalize_for_repo(
-		normalized, git_reference_owner(ref), new_name, true)) < 0)
+		normalized, repo, new_name, true)) < 0)
 		return error;
 
-
 	/* Check if we have to update HEAD. */
 	if ((error = git_branch_is_head(ref)) < 0)
 		return error;
@@ -603,14 +670,18 @@ static int reference__rename(git_reference **out, git_reference *ref, const char
 	if ((error = git_refdb_rename(out, ref->db, ref->name, normalized, force, signature, message)) < 0)
 		return error;
 
-	/* Update HEAD it was pointing to the reference being renamed */
-	if (should_head_be_updated &&
-		(error = git_repository_set_head(ref->db->repo, normalized)) < 0) {
-		giterr_set(GITERR_REFERENCE, "failed to update HEAD after renaming reference");
-		return error;
+	/* Update HEAD if it was pointing to the reference being renamed */
+	if (should_head_be_updated) {
+		error = git_repository_set_head(ref->db->repo, normalized);
+	} else {
+		rename_cb_data payload;
+		payload.old_name = ref->name;
+		memcpy(&payload.new_name, &normalized, sizeof(normalized));
+
+		error = git_repository_foreach_head(repo, update_wt_heads, &payload);
 	}
 
-	return 0;
+	return error;
 }
 
 
diff --git a/src/refs.h b/src/refs.h
index 80e655af7..0c90db3af 100644
--- a/src/refs.h
+++ b/src/refs.h
@@ -107,6 +107,20 @@ int git_reference_lookup_resolved(
 	const char *name,
 	int max_deref);
 
+/**
+ * Read reference from a file.
+ *
+ * This function will read in the file at `path`. If it is a
+ * symref, it will return a new unresolved symbolic reference
+ * with the given name pointing to the reference pointed to by
+ * the file. If it is not a symbolic reference, it will return
+ * the resolved reference.
+ */
+int git_reference__read_head(
+	git_reference **out,
+	git_repository *repo,
+	const char *path);
+
 int git_reference__log_signature(git_signature **out, git_repository *repo);
 
 /** Update a reference after a commit. */
diff --git a/src/repository.c b/src/repository.c
index 425ef796f..707b7b7bd 100644
--- a/src/repository.c
+++ b/src/repository.c
@@ -2063,47 +2063,27 @@ int git_repository_head_detached(git_repository *repo)
 	return exists;
 }
 
-static int read_worktree_head(git_buf *out, git_repository *repo, const char *name)
+static int get_worktree_file_path(git_buf *out, git_repository *repo, const char *worktree, const char *file)
 {
-	git_buf path = GIT_BUF_INIT;
-	int err;
-
-	assert(out && repo && name);
-
 	git_buf_clear(out);
-
-	if ((err = git_buf_printf(&path, "%s/worktrees/%s/HEAD", repo->commondir, name)) < 0)
-		goto out;
-	if (!git_path_exists(path.ptr))
-	{
-		err = -1;
-		goto out;
-	}
-
-	if ((err = git_futils_readbuffer(out, path.ptr)) < 0)
-		goto out;
-	git_buf_rtrim(out);
-
-out:
-	git_buf_free(&path);
-
-	return err;
+	return git_buf_printf(out, "%s/worktrees/%s/%s", repo->commondir, worktree, file);
 }
 
 int git_repository_head_detached_for_worktree(git_repository *repo, const char *name)
 {
-	git_buf buf = GIT_BUF_INIT;
-	int ret;
+	git_reference *ref = NULL;
+	int error;
 
 	assert(repo && name);
 
-	if (read_worktree_head(&buf, repo, name) < 0)
-		return -1;
+	if ((error = git_repository_head_for_worktree(&ref, repo, name)) < 0)
+		goto out;
 
-	ret = git__strncmp(buf.ptr, GIT_SYMREF, strlen(GIT_SYMREF)) != 0;
-	git_buf_free(&buf);
+	error = (git_reference_type(ref) != GIT_REF_SYMBOLIC);
+out:
+	git_reference_free(ref);
 
-	return ret;
+	return error;
 }
 
 int git_repository_head(git_reference **head_out, git_repository *repo)
@@ -2127,44 +2107,66 @@ int git_repository_head(git_reference **head_out, git_repository *repo)
 
 int git_repository_head_for_worktree(git_reference **out, git_repository *repo, const char *name)
 {
-	git_buf buf = GIT_BUF_INIT;
-	git_reference *head;
-	int err;
+	git_buf path = GIT_BUF_INIT;
+	git_reference *head = NULL;
+	int error;
 
 	assert(out && repo && name);
 
 	*out = NULL;
 
-	if (git_repository_head_detached_for_worktree(repo, name))
-		return -1;
-	if ((err = read_worktree_head(&buf, repo, name)) < 0)
+	if ((error = get_worktree_file_path(&path, repo, name, GIT_HEAD_FILE)) < 0 ||
+	    (error = git_reference__read_head(&head, repo, path.ptr)) < 0)
 		goto out;
 
-	/* We can only resolve symbolic references */
-	if (git__strncmp(buf.ptr, GIT_SYMREF, strlen(GIT_SYMREF)))
-	{
-		err = -1;
-		goto out;
+	if (git_reference_type(head) != GIT_REF_OID) {
+		git_reference *resolved;
+
+		error = git_reference_lookup_resolved(&resolved, repo, git_reference_symbolic_target(head), -1);
+		git_reference_free(head);
+		head = resolved;
 	}
-	git_buf_consume(&buf, buf.ptr + strlen(GIT_SYMREF));
 
-	if ((err = git_reference_lookup(&head, repo, buf.ptr)) < 0)
+	*out = head;
+
+out:
+	if (error)
+		git_reference_free(head);
+	git_buf_clear(&path);
+
+	return error;
+}
+
+int git_repository_foreach_head(git_repository *repo, git_repository_foreach_head_cb cb, void *payload)
+{
+	git_strarray worktrees = GIT_VECTOR_INIT;
+	git_buf path = GIT_BUF_INIT;
+	int error;
+	size_t i;
+
+	/* Execute callback for HEAD of commondir */
+	if ((error = git_buf_joinpath(&path, repo->commondir, GIT_HEAD_FILE)) < 0 ||
+	    (error = cb(repo, path.ptr, payload) != 0))
 		goto out;
-	if (git_reference_type(head) == GIT_REF_OID)
-	{
-		*out = head;
-		err = 0;
+
+	if ((error = git_worktree_list(&worktrees, repo)) < 0) {
+		error = 0;
 		goto out;
 	}
 
-	err = git_reference_lookup_resolved(
-		out, repo, git_reference_symbolic_target(head), -1);
-	git_reference_free(head);
+	/* Execute callback for all worktree HEADs */
+	for (i = 0; i < worktrees.count; i++) {
+		if (get_worktree_file_path(&path, repo, worktrees.strings[i], GIT_HEAD_FILE) < 0)
+			continue;
 
-out:
-	git_buf_free(&buf);
+		if ((error = cb(repo, path.ptr, payload)) != 0)
+			goto out;
+	}
 
-	return err;
+out:
+	git_buf_free(&path);
+	git_strarray_free(&worktrees);
+	return error;
 }
 
 int git_repository_head_unborn(git_repository *repo)
@@ -2562,6 +2564,8 @@ int git_repository_set_head(
 
 	if (ref && current->type == GIT_REF_SYMBOLIC && git__strcmp(current->target.symbolic, ref->name) &&
 	    git_reference_is_branch(ref) && git_branch_is_checked_out(ref)) {
+		giterr_set(GITERR_REPOSITORY, "cannot set HEAD to reference '%s' as it is the current HEAD "
+			"of a linked repository.", git_reference_name(ref));
 		error = -1;
 		goto cleanup;
 	}
diff --git a/src/repository.h b/src/repository.h
index 33adfa60a..f922d00f2 100644
--- a/src/repository.h
+++ b/src/repository.h
@@ -160,6 +160,26 @@ int git_repository_head_tree(git_tree **tree, git_repository *repo);
 int git_repository_create_head(const char *git_dir, const char *ref_name);
 
 /*
+ * Called for each HEAD.
+ *
+ * Can return either 0, causing the iteration over HEADs to
+ * continue, or a non-0 value causing the iteration to abort. The
+ * return value is passed back to the caller of
+ * `git_repository_foreach_head`
+ */
+typedef int (*git_repository_foreach_head_cb)(git_repository *repo, const char *path, void *payload);
+
+/*
+ * Iterate over repository and all worktree HEADs.
+ *
+ * This function will be called for the repository HEAD and for
+ * all HEADS of linked worktrees. For each HEAD, the callback is
+ * executed with the given payload. The return value equals the
+ * return value of the last executed callback function.
+ */
+int git_repository_foreach_head(git_repository *repo, git_repository_foreach_head_cb cb, void *payload);
+
+/*
  * Weak pointers to repository internals.
  *
  * The returned pointers do not need to be freed. Do not keep
diff --git a/tests/worktree/refs.c b/tests/worktree/refs.c
index b9a05606d..95b173ef5 100644
--- a/tests/worktree/refs.c
+++ b/tests/worktree/refs.c
@@ -107,28 +107,42 @@ void test_worktree_refs__set_head_fails_when_already_checked_out(void)
 
 void test_worktree_refs__delete_fails_for_checked_out_branch(void)
 {
-       git_reference *branch;
+	git_reference *branch;
 
-       cl_git_pass(git_branch_lookup(&branch, fixture.repo,
-               "testrepo-worktree", GIT_BRANCH_LOCAL));
-       cl_git_fail(git_branch_delete(branch));
+	cl_git_pass(git_branch_lookup(&branch, fixture.repo,
+		    "testrepo-worktree", GIT_BRANCH_LOCAL));
+	cl_git_fail(git_branch_delete(branch));
 
-       git_reference_free(branch);
+	git_reference_free(branch);
 }
 
 void test_worktree_refs__delete_succeeds_after_pruning_worktree(void)
 {
-       git_reference *branch;
-       git_worktree *worktree;
+	git_reference *branch;
+	git_worktree *worktree;
 
-       cl_git_pass(git_worktree_lookup(&worktree, fixture.repo, fixture.worktreename));
-       cl_git_pass(git_worktree_prune(worktree, GIT_WORKTREE_PRUNE_VALID));
-       git_worktree_free(worktree);
+	cl_git_pass(git_worktree_lookup(&worktree, fixture.repo, fixture.worktreename));
+	cl_git_pass(git_worktree_prune(worktree, GIT_WORKTREE_PRUNE_VALID));
+	git_worktree_free(worktree);
 
-       cl_git_pass(git_branch_lookup(&branch, fixture.repo,
-               "testrepo-worktree", GIT_BRANCH_LOCAL));
-       cl_git_pass(git_branch_delete(branch));
-       git_reference_free(branch);
+	cl_git_pass(git_branch_lookup(&branch, fixture.repo,
+		    "testrepo-worktree", GIT_BRANCH_LOCAL));
+	cl_git_pass(git_branch_delete(branch));
+	git_reference_free(branch);
+}
+
+void test_worktree_refs__renaming_reference_updates_worktree_heads(void)
+{
+	git_reference *head, *branch, *renamed;
+
+	cl_git_pass(git_branch_lookup(&branch, fixture.repo,
+		    "testrepo-worktree", GIT_BRANCH_LOCAL));
+	cl_git_pass(git_reference_rename(&renamed, branch, "refs/heads/renamed", 0, NULL));
+	cl_git_pass(git_repository_head(&head, fixture.worktree));
+
+	git_reference_free(head);
+	git_reference_free(branch);
+	git_reference_free(renamed);
 }
 
 void test_worktree_refs__creating_refs_uses_commondir(void)
diff --git a/tests/worktree/worktree.c b/tests/worktree/worktree.c
index 6e90e6ac0..73991bff7 100644
--- a/tests/worktree/worktree.c
+++ b/tests/worktree/worktree.c
@@ -486,3 +486,46 @@ void test_worktree_worktree__prune_both(void)
 
 	git_worktree_free(wt);
 }
+
+static int read_head_ref(git_repository *repo, const char *path, void *payload)
+{
+	git_vector *refs = (git_vector *) payload;
+	git_reference *head;
+
+	GIT_UNUSED(repo);
+
+	cl_git_pass(git_reference__read_head(&head, repo, path));
+
+	git_vector_insert(refs, head);
+
+	return 0;
+}
+
+void test_worktree_worktree__foreach_head_gives_same_results_in_wt_and_repo(void)
+{
+	git_vector repo_refs = GIT_VECTOR_INIT, worktree_refs = GIT_VECTOR_INIT;
+	git_reference *heads[2];
+	size_t i;
+
+	cl_git_pass(git_reference_lookup(&heads[0], fixture.repo, GIT_HEAD_FILE));
+	cl_git_pass(git_reference_lookup(&heads[1], fixture.worktree, GIT_HEAD_FILE));
+
+	cl_git_pass(git_repository_foreach_head(fixture.repo, read_head_ref, &repo_refs));
+	cl_git_pass(git_repository_foreach_head(fixture.worktree, read_head_ref, &worktree_refs));
+
+	cl_assert_equal_i(repo_refs.length, ARRAY_SIZE(heads));
+	cl_assert_equal_i(worktree_refs.length, ARRAY_SIZE(heads));
+
+	for (i = 0; i < ARRAY_SIZE(heads); i++) {
+		cl_assert_equal_s(heads[i]->name, ((git_reference *) repo_refs.contents[i])->name);
+		cl_assert_equal_s(heads[i]->name, ((git_reference *) repo_refs.contents[i])->name);
+		cl_assert_equal_s(heads[i]->name, ((git_reference *) worktree_refs.contents[i])->name);
+
+		git_reference_free(heads[i]);
+		git_reference_free(repo_refs.contents[i]);
+		git_reference_free(worktree_refs.contents[i]);
+	}
+
+	git_vector_free(&repo_refs);
+	git_vector_free(&worktree_refs);
+}