git: fix import_head not getting the HEAD for the current workspace

When you import HEAD, we want to do that using a worktree if the current
workspace is one.

cli/src/cli_util.rs

@@ -941,8 +941,9 @@ impl WorkspaceCommandHelper {
         assert!(self.may_update_working_copy);
         let command = self.env.command.clone();
         let workspace_id = self.workspace_id().clone();
+        let git_repo = self.git_either_colocated_or_backend()?;
         let mut tx = self.start_transaction();
-        git::import_head(tx.repo_mut(), &workspace_id)?;
+        git::import_head(tx.repo_mut(), &git_repo, &workspace_id)?;
         if !tx.repo().has_changes() {
             return Ok(());
         }
@@ -1088,13 +1089,36 @@ impl WorkspaceCommandHelper {
         self.colocated_git_worktree.is_some()
     }
 
-    pub fn open_colocated_git_worktree(&self) -> Result<Option<git2::Repository>, git2::Error> {
+    pub fn open_colocated_git_worktree_git2(
+        &self,
+    ) -> Result<Option<git2::Repository>, git2::Error> {
         self.colocated_git_worktree
             .as_deref()
             .map(|wt| git2::Repository::open(wt))
             .transpose()
     }
 
+    pub fn open_colocated_git_worktree_gix(
+        &self,
+    ) -> Result<Option<gix::Repository>, gix::open::Error> {
+        self.colocated_git_worktree
+            .as_deref()
+            .map(|wt| gix::open(wt))
+            .transpose()
+    }
+
+    pub fn git_either_colocated_or_backend(&self) -> Result<gix::Repository, CommandError> {
+        Ok(self
+            .open_colocated_git_worktree_gix()
+            .map_err(user_error)
+            .transpose()
+            .unwrap_or_else(|| {
+                self.git_backend()
+                    .ok_or_else(|| user_error("JJ repo is not backed by git"))
+                    .map(|g| g.git_repo())
+            })?)
+    }
+
     pub fn format_file_path(&self, file: &RepoPath) -> String {
         self.path_converter().format_file_path(file)
     }
@@ -1770,7 +1794,7 @@ See https://martinvonz.github.io/jj/latest/working-copy/#stale-working-copy \
             .map(|commit_id| tx.repo().store().get_commit(commit_id))
             .transpose()?;
 
-        if let Some(git_worktree) = self.open_colocated_git_worktree()? {
+        if let Some(git_worktree) = self.open_colocated_git_worktree_git2()? {
             if let Some(wc_commit) = &maybe_new_wc_commit {
                 git::reset_head(tx.repo_mut(), &git_worktree, self.workspace_id(), wc_commit)?;
             }

cli/src/commands/git/import.rs

@@ -33,10 +33,11 @@ pub fn cmd_git_import(
 ) -> Result<(), CommandError> {
     let mut workspace_command = command.workspace_helper(ui)?;
     let workspace_id = workspace_command.workspace_id().clone();
+    let git_repo = workspace_command.git_either_colocated_or_backend()?;
     let mut tx = workspace_command.start_transaction();
     // In non-colocated repo, HEAD@git will never be moved internally by jj.
     // That's why cmd_git_export() doesn't export the HEAD ref.
-    git::import_head(tx.repo_mut(), &workspace_id)?;
+    git::import_head(tx.repo_mut(), &git_repo, &workspace_id)?;
     let stats = git::import_refs(tx.repo_mut(), &command.settings().git_settings())?;
     print_git_import_stats(ui, tx.repo(), &stats, true)?;
     tx.finish(ui, "import git refs")?;

cli/src/commands/git/init.rs

@@ -175,8 +175,9 @@ pub fn do_init(
             maybe_set_repository_level_trunk_alias(ui, &workspace_command)?;
             if !workspace_command.working_copy_shared_with_git() {
                 let workspace_id = workspace_command.workspace_id().clone();
+                let git_repo = workspace_command.git_either_colocated_or_backend()?;
                 let mut tx = workspace_command.start_transaction();
-                jj_lib::git::import_head(tx.repo_mut(), &workspace_id)?;
+                jj_lib::git::import_head(tx.repo_mut(), &git_repo, &workspace_id)?;
                 if let Some(git_head_id) = tx
                     .repo()
                     .view()

cli/tests/test_git_colocated.rs

@@ -1273,8 +1273,7 @@ fn test_colocated_view_proto_deprecated() {
          working_copies: ["renamed", "second"]
     "#);
 
-    // A bit strange: apparently we are finding HEAD for the
-    // default workspace (aka renamed) even when it is not present in the view
+    // One tiny command, and the git HEAD is restored for this workspace
     insta::assert_snapshot!(log_heads(&test_env, &repo_path), @r#"
     ○  renamed
     │
@@ -1287,7 +1286,7 @@ fn test_colocated_view_proto_deprecated() {
          working_copies: ["renamed", "second"]
     "#);
 
-    // Even weirder: even in the other workspace.
+    // Touch it ever so slightly and it's back here too
     insta::assert_snapshot!(log_heads(&test_env, &second_path), @r#"
     ○  second
     ○  renamed

lib/src/git.rs

@@ -868,14 +868,14 @@ fn remotely_pinned_commit_ids(view: &View) -> Vec<CommitId> {
 /// the child of the new `HEAD@git` revision.
 pub fn import_head(
     mut_repo: &mut MutableRepo,
+    git_worktree: &gix::Repository,
     workspace_id: &WorkspaceId,
 ) -> Result<(), GitImportError> {
     let store = mut_repo.store();
     let git_backend = get_git_backend(store).ok_or(GitImportError::UnexpectedBackend)?;
-    let git_repo = git_backend.git_repo();
 
     let old_git_head = mut_repo.view().git_head(workspace_id);
-    let new_git_head_id = if let Ok(oid) = git_repo.head_id() {
+    let new_git_head_id = if let Ok(oid) = git_worktree.head_id() {
         Some(CommitId::from_bytes(oid.as_bytes()))
     } else {
         None

lib/tests/test_git.rs

@@ -111,6 +111,10 @@ fn get_git_repo(repo: &Arc<ReadonlyRepo>) -> git2::Repository {
     get_git_backend(repo).open_git_repo().unwrap()
 }
 
+fn get_git_repo_gix(repo: &Arc<ReadonlyRepo>) -> gix::Repository {
+    get_git_backend(repo).git_repo()
+}
+
 #[test]
 fn test_import_refs() {
     let settings = testutils::user_settings();
@@ -121,6 +125,7 @@ fn test_import_refs() {
     let test_repo = TestRepo::init_with_backend(TestRepoBackend::Git);
     let repo = &test_repo.repo;
     let git_repo = get_git_repo(repo);
+    let gix = get_git_repo_gix(repo);
     let workspace_id = &WorkspaceId::default();
 
     let commit1 = empty_git_commit(&git_repo, "refs/heads/main", &[]);
@@ -137,7 +142,7 @@ fn test_import_refs() {
     git_repo.set_head("refs/heads/main").unwrap();
 
     let mut tx = repo.start_transaction(&settings);
-    git::import_head(tx.repo_mut(), workspace_id).unwrap();
+    git::import_head(tx.repo_mut(), &gix, workspace_id).unwrap();
     let stats = git::import_refs(tx.repo_mut(), &git_settings).unwrap();
     tx.repo_mut().rebase_descendants(&settings).unwrap();
     let repo = tx.commit("test");
@@ -389,13 +394,14 @@ fn test_import_refs_reimport_git_head_does_not_count() {
     let test_repo = TestRepo::init_with_backend(TestRepoBackend::Git);
     let repo = &test_repo.repo;
     let git_repo = get_git_repo(repo);
+    let gix = get_git_repo_gix(repo);
     let workspace_id = &WorkspaceId::default();
 
     let commit = empty_git_commit(&git_repo, "refs/heads/main", &[]);
     git_repo.set_head_detached(commit.id()).unwrap();
 
     let mut tx = repo.start_transaction(&settings);
-    git::import_head(tx.repo_mut(), workspace_id).unwrap();
+    git::import_head(tx.repo_mut(), &gix, workspace_id).unwrap();
     git::import_refs(tx.repo_mut(), &git_settings).unwrap();
     tx.repo_mut().rebase_descendants(&settings).unwrap();
 
@@ -406,7 +412,7 @@ fn test_import_refs_reimport_git_head_does_not_count() {
         .unwrap()
         .delete()
         .unwrap();
-    git::import_head(tx.repo_mut(), workspace_id).unwrap();
+    git::import_head(tx.repo_mut(), &gix, workspace_id).unwrap();
     git::import_refs(tx.repo_mut(), &git_settings).unwrap();
     tx.repo_mut().rebase_descendants(&settings).unwrap();
     assert!(!tx.repo_mut().view().heads().contains(&jj_id(&commit)));
@@ -420,6 +426,7 @@ fn test_import_refs_reimport_git_head_without_ref() {
     let test_repo = TestRepo::init_with_backend(TestRepoBackend::Git);
     let repo = &test_repo.repo;
     let git_repo = get_git_repo(repo);
+    let gix = get_git_repo_gix(repo);
     let workspace_id = &WorkspaceId::default();
 
     // First, HEAD points to commit1.
@@ -429,7 +436,7 @@ fn test_import_refs_reimport_git_head_without_ref() {
     git_repo.set_head_detached(git_id(&commit1)).unwrap();
 
     // Import HEAD.
-    git::import_head(tx.repo_mut(), workspace_id).unwrap();
+    git::import_head(tx.repo_mut(), &gix, workspace_id).unwrap();
     git::import_refs(tx.repo_mut(), &git_settings).unwrap();
     tx.repo_mut().rebase_descendants(&settings).unwrap();
     assert!(tx.repo_mut().view().heads().contains(commit1.id()));
@@ -442,7 +449,7 @@ fn test_import_refs_reimport_git_head_without_ref() {
     // would be moved by `git checkout` command. This isn't always true because the
     // detached HEAD commit could be rewritten by e.g. `git commit --amend` command,
     // but it should be safer than abandoning old checkout branch.
-    git::import_head(tx.repo_mut(), workspace_id).unwrap();
+    git::import_head(tx.repo_mut(), &gix, workspace_id).unwrap();
     git::import_refs(tx.repo_mut(), &git_settings).unwrap();
     tx.repo_mut().rebase_descendants(&settings).unwrap();
     assert!(tx.repo_mut().view().heads().contains(commit1.id()));
@@ -457,6 +464,7 @@ fn test_import_refs_reimport_git_head_with_moved_ref() {
     let test_repo = TestRepo::init_with_backend(TestRepoBackend::Git);
     let repo = &test_repo.repo;
     let git_repo = get_git_repo(repo);
+    let gix = get_git_repo_gix(repo);
     let workspace_id = &WorkspaceId::default();
 
     // First, both HEAD and main point to commit1.
@@ -469,7 +477,7 @@ fn test_import_refs_reimport_git_head_with_moved_ref() {
     git_repo.set_head_detached(git_id(&commit1)).unwrap();
 
     // Import HEAD and main.
-    git::import_head(tx.repo_mut(), workspace_id).unwrap();
+    git::import_head(tx.repo_mut(), &gix, workspace_id).unwrap();
     git::import_refs(tx.repo_mut(), &git_settings).unwrap();
     tx.repo_mut().rebase_descendants(&settings).unwrap();
     assert!(tx.repo_mut().view().heads().contains(commit1.id()));
@@ -482,13 +490,13 @@ fn test_import_refs_reimport_git_head_with_moved_ref() {
     git_repo.set_head_detached(git_id(&commit2)).unwrap();
 
     // Reimport HEAD and main, which abandons the old main branch.
-    git::import_head(tx.repo_mut(), workspace_id).unwrap();
+    git::import_head(tx.repo_mut(), &gix, workspace_id).unwrap();
     git::import_refs(tx.repo_mut(), &git_settings).unwrap();
     tx.repo_mut().rebase_descendants(&settings).unwrap();
     assert!(!tx.repo_mut().view().heads().contains(commit1.id()));
     assert!(tx.repo_mut().view().heads().contains(commit2.id()));
     // Reimport HEAD and main, which abandons the old main bookmark.
-    git::import_head(tx.repo_mut(), workspace_id).unwrap();
+    git::import_head(tx.repo_mut(), &gix, workspace_id).unwrap();
     git::import_refs(tx.repo_mut(), &git_settings).unwrap();
     tx.repo_mut().rebase_descendants(&settings).unwrap();
     assert!(!tx.repo_mut().view().heads().contains(commit1.id()));
@@ -972,6 +980,7 @@ fn test_import_refs_reimport_git_head_with_fixed_ref() {
     let test_repo = TestRepo::init_with_backend(TestRepoBackend::Git);
     let repo = &test_repo.repo;
     let git_repo = get_git_repo(repo);
+    let gix = get_git_repo_gix(repo);
     let workspace_id = &WorkspaceId::default();
 
     // First, both HEAD and main point to commit1.
@@ -984,7 +993,7 @@ fn test_import_refs_reimport_git_head_with_fixed_ref() {
     git_repo.set_head_detached(git_id(&commit1)).unwrap();
 
     // Import HEAD and main.
-    git::import_head(tx.repo_mut(), workspace_id).unwrap();
+    git::import_head(tx.repo_mut(), &gix, workspace_id).unwrap();
     git::import_refs(tx.repo_mut(), &git_settings).unwrap();
     tx.repo_mut().rebase_descendants(&settings).unwrap();
     assert!(tx.repo_mut().view().heads().contains(commit1.id()));
@@ -994,7 +1003,7 @@ fn test_import_refs_reimport_git_head_with_fixed_ref() {
     git_repo.set_head_detached(git_id(&commit2)).unwrap();
 
     // Reimport HEAD, which shouldn't abandon the old HEAD branch.
-    git::import_head(tx.repo_mut(), workspace_id).unwrap();
+    git::import_head(tx.repo_mut(), &gix, workspace_id).unwrap();
     git::import_refs(tx.repo_mut(), &git_settings).unwrap();
     tx.repo_mut().rebase_descendants(&settings).unwrap();
     assert!(tx.repo_mut().view().heads().contains(commit1.id()));
@@ -1391,6 +1400,7 @@ fn test_import_refs_missing_git_commit() {
     let repo = &test_workspace.repo;
     let workspace_id = &WorkspaceId::default();
     let git_repo = get_git_repo(repo);
+    let gix = get_git_repo_gix(repo);
 
     let commit1 = empty_git_commit(&git_repo, "refs/heads/main", &[]);
     let commit2 = empty_git_commit(&git_repo, "refs/heads/main", &[&commit1]);
@@ -1422,7 +1432,7 @@ fn test_import_refs_missing_git_commit() {
         .unwrap();
     git_repo.set_head_detached(commit2.id()).unwrap();
     let mut tx = repo.start_transaction(&settings);
-    let result = git::import_head(tx.repo_mut(), workspace_id);
+    let result = git::import_head(tx.repo_mut(), &gix, workspace_id);
     assert_matches!(
         result,
         Err(GitImportError::MissingHeadTarget {
@@ -1454,7 +1464,7 @@ fn test_import_refs_missing_git_commit() {
     git_repo.set_head_detached(commit1.id()).unwrap();
     fs::rename(&object_file, &backup_object_file).unwrap();
     let mut tx = repo.start_transaction(&settings);
-    let result = git::import_head(tx.repo_mut(), workspace_id);
+    let result = git::import_head(tx.repo_mut(), &gix, workspace_id);
     assert!(result.is_ok());
 }
 
@@ -1473,9 +1483,10 @@ fn test_import_refs_detached_head() {
         .unwrap();
     test_data.git_repo.set_head_detached(commit1.id()).unwrap();
     let workspace_id = &WorkspaceId::default();
+    let gix = get_git_repo_gix(&test_data.repo);
 
     let mut tx = test_data.repo.start_transaction(&test_data.settings);
-    git::import_head(tx.repo_mut(), workspace_id).unwrap();
+    git::import_head(tx.repo_mut(), &gix, workspace_id).unwrap();
     git::import_refs(tx.repo_mut(), &git_settings).unwrap();
     tx.repo_mut()
         .rebase_descendants(&test_data.settings)
@@ -1499,11 +1510,12 @@ fn test_export_refs_no_detach() {
     let git_settings = GitSettings::default();
     let git_repo = test_data.git_repo;
     let workspace_id = &WorkspaceId::default();
+    let gix = get_git_repo_gix(&test_data.repo);
     let commit1 = empty_git_commit(&git_repo, "refs/heads/main", &[]);
     git_repo.set_head("refs/heads/main").unwrap();
     let mut tx = test_data.repo.start_transaction(&test_data.settings);
     let mut_repo = tx.repo_mut();
-    git::import_head(mut_repo, workspace_id).unwrap();
+    git::import_head(mut_repo, &gix, workspace_id).unwrap();
     git::import_refs(mut_repo, &git_settings).unwrap();
     mut_repo.rebase_descendants(&test_data.settings).unwrap();
 
@@ -1526,6 +1538,7 @@ fn test_export_refs_bookmark_changed() {
     let test_data = GitRepoData::create();
     let git_settings = GitSettings::default();
     let git_repo = test_data.git_repo;
+    let gix = get_git_repo_gix(&test_data.repo);
     let workspace_id = &WorkspaceId::default();
     let commit = empty_git_commit(&git_repo, "refs/heads/main", &[]);
     git_repo
@@ -1535,7 +1548,8 @@ fn test_export_refs_bookmark_changed() {
 
     let mut tx = test_data.repo.start_transaction(&test_data.settings);
     let mut_repo = tx.repo_mut();
-    git::import_head(mut_repo, workspace_id).unwrap();
+
+    git::import_head(mut_repo, &gix, workspace_id).unwrap();
     git::import_refs(mut_repo, &git_settings).unwrap();
     mut_repo.rebase_descendants(&test_data.settings).unwrap();
     assert!(git::export_refs(mut_repo).unwrap().is_empty());
@@ -1570,12 +1584,13 @@ fn test_export_refs_current_bookmark_changed() {
     let test_data = GitRepoData::create();
     let git_settings = GitSettings::default();
     let git_repo = test_data.git_repo;
+    let gix = get_git_repo_gix(&test_data.repo);
     let workspace_id = &WorkspaceId::default();
     let commit1 = empty_git_commit(&git_repo, "refs/heads/main", &[]);
     git_repo.set_head("refs/heads/main").unwrap();
     let mut tx = test_data.repo.start_transaction(&test_data.settings);
     let mut_repo = tx.repo_mut();
-    git::import_head(mut_repo, workspace_id).unwrap();
+    git::import_head(mut_repo, &gix, workspace_id).unwrap();
     git::import_refs(mut_repo, &git_settings).unwrap();
     mut_repo.rebase_descendants(&test_data.settings).unwrap();
     assert!(git::export_refs(mut_repo).unwrap().is_empty());
@@ -1609,11 +1624,12 @@ fn test_export_refs_unborn_git_bookmark(move_placeholder_ref: bool) {
     let test_data = GitRepoData::create();
     let git_settings = GitSettings::default();
     let git_repo = test_data.git_repo;
+    let gix = get_git_repo_gix(&test_data.repo);
     let workspace_id = &WorkspaceId::default();
     git_repo.set_head("refs/heads/main").unwrap();
     let mut tx = test_data.repo.start_transaction(&test_data.settings);
     let mut_repo = tx.repo_mut();
-    git::import_head(mut_repo, workspace_id).unwrap();
+    git::import_head(mut_repo, &gix, workspace_id).unwrap();
     git::import_refs(mut_repo, &git_settings).unwrap();
     mut_repo.rebase_descendants(&test_data.settings).unwrap();
     assert!(git::export_refs(mut_repo).unwrap().is_empty());