git: extract Git HEAD handling bits from import_some_refs()

I'm going to make WorkspaceCommandHelper::maybe_snapshot() snapshot the working
copy before importing refs. git::import_some_refs() can rebase the working copy
branch and therefore @ can be moved. git::import_head() doesn't, and it should
be invoked before snapshotting.

git::import_head() is inserted to some of the git:import_refs() callers where
HEAD seems to matter. I feel it's a bit odd that the HEAD ref is imported to
non-colocated repo, but "jj init --git-repo" relies on that, and I think the
existence of HEAD@git is harmless. It's merely a ref to the revision checked
out somewhere else.

cli/src/cli_util.rs

@@ -824,6 +824,7 @@ impl WorkspaceCommandHelper {
         let git_settings = self.settings.git_settings();
         let mut tx = self.start_transaction();
         // Automated import shouldn't fail because of reserved remote name.
+        git::import_head(tx.mut_repo())?;
         let stats = git::import_some_refs(tx.mut_repo(), &git_settings, |ref_name| {
             !git::is_reserved_git_remote_ref(ref_name)
         })?;

cli/src/commands/git.rs

@@ -980,6 +980,9 @@ fn cmd_git_import(
 ) -> Result<(), CommandError> {
     let mut workspace_command = command.workspace_helper(ui)?;
     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.mut_repo())?;
     let stats = git::import_refs(tx.mut_repo(), &command.settings().git_settings())?;
     print_git_import_stats(ui, &stats)?;
     tx.finish(ui, "import git refs")?;

cli/src/commands/init.rs

@@ -73,6 +73,7 @@ pub(crate) fn cmd_init(
         workspace_command.maybe_snapshot(ui)?;
         if !workspace_command.working_copy_shared_with_git() {
             let mut tx = workspace_command.start_transaction();
+            jj_lib::git::import_head(tx.mut_repo())?;
             let stats = jj_lib::git::import_some_refs(
                 tx.mut_repo(),
                 &command.settings().git_settings(),

lib/src/git.rs

@@ -240,35 +240,22 @@ pub fn import_some_refs(
     let git_backend = get_git_backend(store).ok_or(GitImportError::UnexpectedBackend)?;
     let git_repo = git_backend.git_repo();
 
-    // TODO: Should this be a separate function? We may not always want to import
-    // the Git HEAD (and add it to our set of heads).
-    let old_git_head = mut_repo.view().git_head();
-    let changed_git_head = if let Ok(head_git_commit) = git_repo.head_commit() {
-        // The current HEAD is not added to `hidable_git_heads` because HEAD move
-        // doesn't automatically mean the old HEAD branch has been rewritten.
-        let head_commit_id = CommitId::from_bytes(head_git_commit.id().as_bytes());
-        let new_head_target = RefTarget::normal(head_commit_id);
-        (*old_git_head != new_head_target).then_some(new_head_target)
-    } else {
-        old_git_head.is_present().then(RefTarget::absent)
-    };
-
     let RefsToImport {
         changed_git_refs,
         changed_remote_refs,
     } = diff_refs_to_import(mut_repo.view(), &git_repo, git_ref_filter)?;
 
-    // Bulk-import all reachable Git commits to the backend to reduce overhead of
-    // table merging and ref updates.
+    // Bulk-import all reachable Git commits to the backend to reduce overhead
+    // of table merging and ref updates.
+    //
+    // changed_remote_refs might contain new_targets that are not in
+    // changed_git_refs, but such targets should have already been imported to
+    // the backend.
     let index = mut_repo.index();
-    let missing_head_ids = itertools::chain(
-        &changed_git_head,
-        changed_git_refs.iter().map(|(_, new_target)| new_target),
-        // changed_remote_refs might contain new_targets that are not in changed_git_refs,
-        // but such targets should have already been imported to the backend.
-    )
-    .flat_map(|target| target.added_ids())
-    .filter(|&id| !index.has_id(id));
+    let missing_head_ids = changed_git_refs
+        .iter()
+        .flat_map(|(_, new_target)| new_target.added_ids())
+        .filter(|&id| !index.has_id(id));
     let heads_imported = git_backend.import_head_commits(missing_head_ids).is_ok();
 
     // Import new remote heads
@@ -280,15 +267,6 @@ pub fn import_some_refs(
         }
         store.get_commit(id)
     };
-    if let Some(new_head_target) = &changed_git_head {
-        for id in new_head_target.added_ids() {
-            let commit = get_commit(id).map_err(|err| GitImportError::MissingHeadTarget {
-                id: id.clone(),
-                err,
-            })?;
-            head_commits.push(commit);
-        }
-    }
     for (ref_name, (_, new_target)) in &changed_remote_refs {
         for id in new_target.added_ids() {
             let commit = get_commit(id).map_err(|err| GitImportError::MissingRefAncestor {
@@ -305,9 +283,6 @@ pub fn import_some_refs(
         .map_err(GitImportError::InternalBackend)?;
 
     // Apply the change that happened in git since last time we imported refs.
-    if let Some(new_head_target) = changed_git_head {
-        mut_repo.set_git_head_target(new_head_target);
-    }
     for (full_name, new_target) in changed_git_refs {
         mut_repo.set_git_ref_target(&full_name, new_target);
     }
@@ -528,6 +503,67 @@ fn pinned_commit_ids(view: &View) -> impl Iterator<Item = &CommitId> {
     .flat_map(|target| target.added_ids())
 }
 
+/// Imports `HEAD@git` from the underlying Git repo.
+///
+/// Unlike `import_refs()`, the old HEAD branch is not abandoned because HEAD
+/// move doesn't always mean the old HEAD branch has been rewritten.
+///
+/// Unlike `reset_head()`, this function doesn't move the working-copy commit to
+/// the child of the new `HEAD@git` revision.
+pub fn import_head(mut_repo: &mut MutableRepo) -> 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();
+    let changed_git_head = if let Ok(head_git_commit) = git_repo.head_commit() {
+        let head_commit_id = CommitId::from_bytes(head_git_commit.id().as_bytes());
+        let new_head_target = RefTarget::normal(head_commit_id);
+        (*old_git_head != new_head_target).then_some(new_head_target)
+    } else {
+        old_git_head.is_present().then(RefTarget::absent)
+    };
+
+    // Bulk-import all reachable Git commits to the backend to reduce overhead of
+    // table merging and ref updates.
+    let index = mut_repo.index();
+    let missing_head_ids = changed_git_head
+        .iter()
+        .flat_map(|target| target.added_ids())
+        .filter(|&id| !index.has_id(id));
+    let heads_imported = git_backend.import_head_commits(missing_head_ids).is_ok();
+
+    // Import new remote heads
+    let mut head_commits = Vec::new();
+    let get_commit = |id| {
+        // If bulk-import failed, try again to find bad head or ref.
+        if !heads_imported && !index.has_id(id) {
+            git_backend.import_head_commits([id])?;
+        }
+        store.get_commit(id)
+    };
+    if let Some(new_head_target) = &changed_git_head {
+        for id in new_head_target.added_ids() {
+            let commit = get_commit(id).map_err(|err| GitImportError::MissingHeadTarget {
+                id: id.clone(),
+                err,
+            })?;
+            head_commits.push(commit);
+        }
+    }
+    // It's unlikely the imported commits were missing, but I/O-related error
+    // can still occur.
+    mut_repo
+        .add_heads(&head_commits)
+        .map_err(GitImportError::InternalBackend)?;
+
+    // Apply the change that happened in git since last time we imported refs.
+    if let Some(new_head_target) = changed_git_head {
+        mut_repo.set_git_head_target(new_head_target);
+    }
+    Ok(())
+}
+
 #[derive(Error, Debug)]
 pub enum GitExportError {
     #[error("Git error: {0}")]

lib/tests/test_git.rs

@@ -109,6 +109,7 @@ fn test_import_refs() {
     git_repo.set_head("refs/heads/main").unwrap();
 
     let mut tx = repo.start_transaction(&settings);
+    git::import_head(tx.mut_repo()).unwrap();
     let stats = git::import_refs(tx.mut_repo(), &git_settings).unwrap();
     tx.mut_repo().rebase_descendants(&settings).unwrap();
     let repo = tx.commit("test");
@@ -362,6 +363,7 @@ fn test_import_refs_reimport_git_head_counts() {
     git_repo.set_head_detached(commit.id()).unwrap();
 
     let mut tx = repo.start_transaction(&settings);
+    git::import_head(tx.mut_repo()).unwrap();
     git::import_refs(tx.mut_repo(), &git_settings).unwrap();
     tx.mut_repo().rebase_descendants(&settings).unwrap();
 
@@ -372,6 +374,7 @@ fn test_import_refs_reimport_git_head_counts() {
         .unwrap()
         .delete()
         .unwrap();
+    git::import_head(tx.mut_repo()).unwrap();
     git::import_refs(tx.mut_repo(), &git_settings).unwrap();
     tx.mut_repo().rebase_descendants(&settings).unwrap();
     assert!(tx.mut_repo().view().heads().contains(&jj_id(&commit)));
@@ -393,6 +396,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.mut_repo()).unwrap();
     git::import_refs(tx.mut_repo(), &git_settings).unwrap();
     tx.mut_repo().rebase_descendants(&settings).unwrap();
     assert!(tx.mut_repo().view().heads().contains(commit1.id()));
@@ -405,6 +409,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.mut_repo()).unwrap();
     git::import_refs(tx.mut_repo(), &git_settings).unwrap();
     tx.mut_repo().rebase_descendants(&settings).unwrap();
     assert!(tx.mut_repo().view().heads().contains(commit1.id()));
@@ -430,6 +435,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.mut_repo()).unwrap();
     git::import_refs(tx.mut_repo(), &git_settings).unwrap();
     tx.mut_repo().rebase_descendants(&settings).unwrap();
     assert!(tx.mut_repo().view().heads().contains(commit1.id()));
@@ -442,6 +448,7 @@ 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.mut_repo()).unwrap();
     git::import_refs(tx.mut_repo(), &git_settings).unwrap();
     tx.mut_repo().rebase_descendants(&settings).unwrap();
     assert!(!tx.mut_repo().view().heads().contains(commit1.id()));
@@ -788,6 +795,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.mut_repo()).unwrap();
     git::import_refs(tx.mut_repo(), &git_settings).unwrap();
     tx.mut_repo().rebase_descendants(&settings).unwrap();
     assert!(tx.mut_repo().view().heads().contains(commit1.id()));
@@ -797,6 +805,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.mut_repo()).unwrap();
     git::import_refs(tx.mut_repo(), &git_settings).unwrap();
     tx.mut_repo().rebase_descendants(&settings).unwrap();
     assert!(tx.mut_repo().view().heads().contains(commit1.id()));
@@ -1216,7 +1225,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_refs(tx.mut_repo(), &git_settings);
+    let result = git::import_head(tx.mut_repo());
     assert_matches!(
         result,
         Err(GitImportError::MissingHeadTarget {
@@ -1248,7 +1257,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_refs(tx.mut_repo(), &git_settings);
+    let result = git::import_head(tx.mut_repo());
     assert!(result.is_ok());
 }
 
@@ -1268,6 +1277,7 @@ fn test_import_refs_detached_head() {
     test_data.git_repo.set_head_detached(commit1.id()).unwrap();
 
     let mut tx = test_data.repo.start_transaction(&test_data.settings);
+    git::import_head(tx.mut_repo()).unwrap();
     git::import_refs(tx.mut_repo(), &git_settings).unwrap();
     tx.mut_repo()
         .rebase_descendants(&test_data.settings)
@@ -1291,6 +1301,7 @@ fn test_export_refs_no_detach() {
     git_repo.set_head("refs/heads/main").unwrap();
     let mut tx = test_data.repo.start_transaction(&test_data.settings);
     let mut_repo = tx.mut_repo();
+    git::import_head(mut_repo).unwrap();
     git::import_refs(mut_repo, &git_settings).unwrap();
     mut_repo.rebase_descendants(&test_data.settings).unwrap();
 
@@ -1321,6 +1332,7 @@ fn test_export_refs_branch_changed() {
 
     let mut tx = test_data.repo.start_transaction(&test_data.settings);
     let mut_repo = tx.mut_repo();
+    git::import_head(mut_repo).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());
@@ -1358,6 +1370,7 @@ fn test_export_refs_current_branch_changed() {
     git_repo.set_head("refs/heads/main").unwrap();
     let mut tx = test_data.repo.start_transaction(&test_data.settings);
     let mut_repo = tx.mut_repo();
+    git::import_head(mut_repo).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());
@@ -1393,6 +1406,7 @@ fn test_export_refs_unborn_git_branch() {
     git_repo.set_head("refs/heads/main").unwrap();
     let mut tx = test_data.repo.start_transaction(&test_data.settings);
     let mut_repo = tx.mut_repo();
+    git::import_head(mut_repo).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());