Use safe-checkout for BranchManager::apply_branch().

Author: Byron

crates/but-workspace/src/head.rs

@@ -2,7 +2,7 @@
 use anyhow::{Context, Result};
 use gitbutler_cherry_pick::RepositoryExt as _;
 use gitbutler_command_context::CommandContext;
-use gitbutler_oxidize::{GixRepositoryExt, ObjectIdExt, OidExt, RepoExt};
+use gitbutler_oxidize::{GixRepositoryExt, ObjectIdExt, OidExt};
 use gitbutler_repo::{RepositoryExt, SignaturePurpose};
 use gitbutler_stack::{Stack, VirtualBranchesHandle};
 use gix::merge::tree::TreatAsUnresolved;
@@ -22,7 +22,7 @@ pub fn merge_worktree_with_workspace<'a>(
 
     // The tree of where the gitbutler workspace is at
     let workspace_tree = gix_repo
-        .find_commit(super::remerged_head_commit(ctx)?.to_gix())?
+        .find_commit(super::remerged_workspace_commit_v2(ctx)?.to_gix())?
         .tree_id()?
         .detach();
 
@@ -40,16 +40,12 @@ pub fn merge_worktree_with_workspace<'a>(
     Ok((outcome, conflict_kind))
 }
 
-/// Creates and returns a merge commit of all active branch heads.
-///
-/// This is the base against which we diff the working directory to understand
-/// what files have been modified.
-///
-/// This should be used to update the `gitbutler/workspace` ref with, which is usually
-/// done from [`update_workspace_commit()`], after any of its input changes.
-/// This is namely the conflicting state, or any head of the virtual branches.
-#[instrument(level = tracing::Level::DEBUG, skip(ctx))]
-pub fn remerged_head_commit(ctx: &CommandContext) -> Result<git2::Oid> {
+/// Merge all currently stored stacks together into a new tree and return `(merged_tree, stacks, target_commit)` id accordingly.
+/// `gix_repo` should be optimised for merging.
+pub fn remerged_workspace_tree_v2<'git2_repo>(
+    ctx: &'git2_repo CommandContext,
+    gix_repo: &gix::Repository,
+) -> Result<(git2::Oid, Vec<Stack>, git2::Commit<'git2_repo>)> {
     let vb_state = VirtualBranchesHandle::new(ctx.project().gb_dir());
     let target = vb_state
         .get_default_target()
@@ -58,15 +54,14 @@ pub fn remerged_head_commit(ctx: &CommandContext) -> Result<git2::Oid> {
     let mut stacks: Vec<Stack> = vb_state.list_stacks_in_workspace()?;
 
     let target_commit = repo.find_commit(target.sha)?;
-    let mut workspace_tree = repo.find_real_tree(&target_commit, Default::default())?;
+    let workspace_tree = repo.find_real_tree(&target_commit, Default::default())?;
     let mut workspace_tree_id = workspace_tree.id().to_gix();
 
-    let gix_repo = ctx.gix_repo_for_merging()?;
     let (merge_options_fail_fast, conflict_kind) = gix_repo.merge_options_fail_fast()?;
     let merge_tree_id = repo.find_commit(target.sha)?.tree_id().to_gix();
     for stack in stacks.iter_mut() {
         stack.migrate_change_ids(ctx).ok(); // If it fails that's ok - best effort migration
-        let branch_head = repo.find_commit(stack.head_oid(&gix_repo)?.to_git2())?;
+        let branch_head = repo.find_commit(stack.head_oid(gix_repo)?.to_git2())?;
         let branch_tree_id = repo
             .find_real_tree(&branch_head, Default::default())?
             .id()
@@ -89,11 +84,26 @@ pub fn remerged_head_commit(ctx: &CommandContext) -> Result<git2::Oid> {
             vb_state.set_stack(stack.clone())?;
         }
     }
-    workspace_tree = repo.find_tree(workspace_tree_id.to_git2())?;
+    Ok((workspace_tree_id.to_git2(), stacks, target_commit))
+}
+
+/// Creates and returns a merge commit of all active branch heads.
+///
+/// This is the base against which we diff the working directory to understand
+/// what files have been modified.
+///
+/// This should be used to update the `gitbutler/workspace` ref with, which is usually
+/// done from [`update_workspace_commit()`], after any of its input changes.
+/// This is namely the conflicting state, or any head of the virtual branches.
+#[instrument(level = tracing::Level::DEBUG, skip(ctx))]
+pub fn remerged_workspace_commit_v2(ctx: &CommandContext) -> Result<git2::Oid> {
+    let repo = ctx.repo();
+    let gix_repo = ctx.gix_repo_for_merging()?;
+    let (workspace_tree_id, stacks, target_commit) = remerged_workspace_tree_v2(ctx, &gix_repo)?;
+    let workspace_tree = repo.find_tree(workspace_tree_id)?;
 
     let committer = gitbutler_repo::signature(SignaturePurpose::Committer)?;
     let author = gitbutler_repo::signature(SignaturePurpose::Author)?;
-    let gix_repo = repo.to_gix()?;
     let mut heads: Vec<git2::Commit<'_>> = stacks
         .iter()
         .filter_map(|stack| stack.head_oid(&gix_repo).ok())

crates/but-workspace/src/lib.rs

@@ -49,7 +49,9 @@ pub use tree_manipulation::{
     split_commit::{CommitFiles, CommmitSplitOutcome, split_commit},
 };
 pub mod head;
-pub use head::{merge_worktree_with_workspace, remerged_head_commit};
+pub use head::{
+    merge_worktree_with_workspace, remerged_workspace_commit_v2, remerged_workspace_tree_v2,
+};
 
 /// 🚧utilities for applying and unapplying branches 🚧.
 /// Ignore the name of this module; it's just a place to put code by now.

crates/gitbutler-branch-actions/src/base.rs

@@ -58,7 +58,7 @@ pub fn get_base_branch_data(ctx: &CommandContext) -> Result<BaseBranch> {
 fn go_back_to_integration(ctx: &CommandContext, default_target: &Target) -> Result<BaseBranch> {
     let gix_repo = ctx.gix_repo_for_merging()?;
     if ctx.app_settings().feature_flags.cv3 {
-        let workspace_commit_to_checkout = but_workspace::remerged_head_commit(ctx)?;
+        let workspace_commit_to_checkout = but_workspace::remerged_workspace_commit_v2(ctx)?;
         let tree_to_checkout_to_avoid_ref_update = gix_repo
             .find_commit(workspace_commit_to_checkout.to_gix())?
             .tree_id()?;

crates/gitbutler-branch-actions/src/branch_manager/branch_creation.rs

@@ -1,4 +1,4 @@
-use super::BranchManager;
+use super::{checkout_remerged_head, BranchManager};
 use crate::r#virtual as vbranch;
 use crate::{hunk::VirtualBranchHunk, integration::update_workspace_commit, VirtualBranchesExt};
 use anyhow::{anyhow, bail, Context, Result};
@@ -345,13 +345,19 @@ impl BranchManager<'_> {
                 .context("failed to merge trees")?;
 
             if !merges_cleanly {
-                for stack in vb_state
+                for stack_to_unapply in vb_state
                     .list_stacks_in_workspace()?
                     .iter()
                     .filter(|branch| branch.id != stack_id)
                 {
-                    unapplied_stacks.push(stack.id);
-                    self.unapply(stack.id, perm, false, Vec::new(), safe_checkout)?;
+                    unapplied_stacks.push(stack_to_unapply.id);
+                    let res =
+                        self.unapply(stack_to_unapply.id, perm, false, Vec::new(), safe_checkout);
+                    if res.is_err() {
+                        stack.in_workspace = false;
+                        vb_state.set_stack(stack.clone())?;
+                    }
+                    res?;
                 }
             }
         }
@@ -429,16 +435,25 @@ impl BranchManager<'_> {
             }
         }
 
+        // Permissions here might be wonky, just go with it though.
         let new_workspace = WorkspaceState::create(self.ctx, perm.read_permission())?;
-
-        update_uncommited_changes_with_tree(
-            self.ctx,
-            workspace_state,
-            new_workspace,
-            old_cwd,
-            Some(true),
-            perm,
-        )?;
+        if safe_checkout {
+            let res = checkout_remerged_head(self.ctx, &gix_repo);
+            if res.is_err() {
+                stack.in_workspace = false;
+                vb_state.set_stack(stack.clone())?;
+            }
+            res?;
+        } else {
+            update_uncommited_changes_with_tree(
+                self.ctx,
+                workspace_state,
+                new_workspace,
+                old_cwd,
+                Some(true),
+                perm,
+            )?;
+        }
 
         update_workspace_commit(&vb_state, self.ctx)?;
 

crates/gitbutler-branch-actions/src/branch_manager/branch_removal.rs

@@ -10,7 +10,7 @@ use gitbutler_stack::StackId;
 use gitbutler_workspace::workspace_base;
 use tracing::instrument;
 
-use super::BranchManager;
+use super::{checkout_remerged_head, BranchManager};
 use crate::r#virtual as vbranch;
 use crate::VirtualBranchesExt;
 
@@ -63,16 +63,7 @@ impl BranchManager<'_> {
         let gix_repo = self.ctx.gix_repo_for_merging()?;
         if safe_checkout {
             // This reads the just stored data and re-merges it all stacks, excluding the unapplied one.
-            let res = (|| -> Result<()> {
-                let workspace_without_stack = but_workspace::remerged_head_commit(self.ctx)?;
-                but_workspace::branch::safe_checkout_from_head(
-                    workspace_without_stack.to_gix(),
-                    &gix_repo,
-                    but_workspace::branch::checkout::Options::default(),
-                )?;
-                Ok(())
-            })();
-
+            let res = checkout_remerged_head(self.ctx, &gix_repo);
             // Undo the removal, stack is still there officially now.
             if res.is_err() {
                 stack.in_workspace = true;

crates/gitbutler-branch-actions/src/branch_manager/mod.rs

@@ -1,9 +1,24 @@
 pub use branch_creation::CreateBranchFromBranchOutcome;
 use gitbutler_command_context::CommandContext;
+use gitbutler_oxidize::OidExt;
 
 mod branch_creation;
 mod branch_removal;
 
+/// Note that this checks out the commit and sets the HEAD ref to point to it.
+pub(crate) fn checkout_remerged_head(
+    ctx: &CommandContext,
+    repo: &gix::Repository,
+) -> anyhow::Result<()> {
+    let (workspace_tree_id, _, _) = but_workspace::remerged_workspace_tree_v2(ctx, repo)?;
+    but_workspace::branch::safe_checkout_from_head(
+        workspace_tree_id.to_gix(),
+        repo,
+        but_workspace::branch::checkout::Options::default(),
+    )?;
+    Ok(())
+}
+
 pub struct BranchManager<'l> {
     ctx: &'l CommandContext,
 }

crates/gitbutler-branch-actions/src/integration.rs

@@ -84,7 +84,7 @@ pub fn update_workspace_commit(
         .list_stacks_in_workspace()
         .context("failed to list virtual branches")?;
 
-    let workspace_head = repo.find_commit(but_workspace::remerged_head_commit(ctx)?)?;
+    let workspace_head = repo.find_commit(but_workspace::remerged_workspace_commit_v2(ctx)?)?;
 
     // message that says how to get back to where they were
     let mut message = GITBUTLER_WORKSPACE_COMMIT_TITLE.to_string();

crates/gitbutler-branch-actions/src/status.rs

@@ -33,7 +33,10 @@ pub fn get_applied_status(
     ctx: &CommandContext,
     perm: Option<&mut WorktreeWritePermission>,
 ) -> Result<VirtualBranchesStatus> {
-    let diffs = gitbutler_diff::workdir(ctx.repo(), but_workspace::remerged_head_commit(ctx)?)?;
+    let diffs = gitbutler_diff::workdir(
+        ctx.repo(),
+        but_workspace::remerged_workspace_commit_v2(ctx)?,
+    )?;
     get_applied_status_cached(ctx, perm, &diffs)
 }
 

crates/gitbutler-branch-actions/src/virtual.rs

@@ -226,7 +226,10 @@ pub fn commit(
     ownership: Option<&BranchOwnershipClaims>,
 ) -> Result<git2::Oid> {
     // get the files to commit
-    let diffs = gitbutler_diff::workdir(ctx.repo(), but_workspace::remerged_head_commit(ctx)?)?;
+    let diffs = gitbutler_diff::workdir(
+        ctx.repo(),
+        but_workspace::remerged_workspace_commit_v2(ctx)?,
+    )?;
     let statuses = get_applied_status_cached(ctx, None, &diffs)
         .context("failed to get status by branch")?
         .branches;