Use common primatives, and make sure conflicted trees are handled correctly

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

@@ -12,7 +12,7 @@ use crate::VirtualBranchesExt as _;
 /// Checks out the combined trees of all branches in the workspace.
 ///
 /// This function will fail if the applied branches conflict with each other.
-pub(crate) fn checkout_branch_trees<'a>(
+pub fn checkout_branch_trees<'a>(
     ctx: &'a CommandContext,
     _perm: &mut WorktreeWritePermission,
 ) -> Result<git2::Tree<'a>> {

crates/gitbutler-edit-mode/src/lib.rs

@@ -4,6 +4,9 @@ use anyhow::{bail, Context, Result};
 use bstr::ByteSlice;
 use git2::build::CheckoutBuilder;
 use gitbutler_branch::{signature, Branch, SignaturePurpose, VirtualBranchesHandle};
+use gitbutler_branch_actions::branch_trees::{
+    checkout_branch_trees, compute_updated_branch_head, BranchHeadAndTree,
+};
 use gitbutler_branch_actions::internal::list_virtual_branches;
 use gitbutler_branch_actions::{update_workspace_commit, RemoteBranchFile};
 use gitbutler_cherry_pick::{ConflictedTreeKey, RepositoryExt as _};
@@ -19,46 +22,10 @@ use gitbutler_operating_modes::{
 };
 use gitbutler_project::access::{WorktreeReadPermission, WorktreeWritePermission};
 use gitbutler_reference::{ReferenceName, Refname};
-use gitbutler_repo::{
+use gitbutler_repo::{rebase::cherry_rebase, RepositoryExt};
-    rebase::{cherry_rebase, cherry_rebase_group},
-    RepositoryExt,
-};
 
 pub mod commands;
 
-pub const EDIT_UNCOMMITED_FILES_REF: &str = "refs/gitbutler/edit_uncommited_files";
-
-fn save_uncommited_files(ctx: &CommandContext) -> Result<()> {
-    let repository = ctx.repository();
-
-    // Create a tree of all uncommited files
-    let tree = repository.create_wd_tree()?;
-
-    // Commit tree and reference it
-    let author_signature =
-        signature(SignaturePurpose::Author).context("Failed to get gitbutler signature")?;
-    let committer_signature =
-        signature(SignaturePurpose::Committer).context("Failed to get gitbutler signature")?;
-    let head = repository.head().context("Failed to get head")?;
-    let head_commit = head.peel_to_commit().context("Failed to get head commit")?;
-    let commit = repository
-        .commit(
-            None,
-            &author_signature,
-            &committer_signature,
-            "Edit mode saved changes",
-            &tree,
-            &[&head_commit],
-        )
-        .context("Failed to write stash commit")?;
-
-    repository
-        .reference(EDIT_UNCOMMITED_FILES_REF, commit, true, "")
-        .context("Failed to reference uncommited files")?;
-
-    Ok(())
-}
-
 fn get_commit_index(repository: &git2::Repository, commit: &git2::Commit) -> Result<git2::Index> {
     let commit_tree = commit.tree().context("Failed to get commit's tree")?;
     // Checkout the commit as unstaged changes
@@ -182,7 +149,6 @@ pub(crate) fn enter_edit_mode(
         bail!("Can not enter edit mode for a reference which does not have a cooresponding virtual branch")
     }
 
-    save_uncommited_files(ctx).context("Failed to save uncommited files")?;
     checkout_edit_branch(ctx, commit).context("Failed to checkout edit branch")?;
     write_edit_mode_metadata(ctx, &edit_mode_metadata).context("Failed to persist metadata")?;
 
@@ -191,36 +157,16 @@ pub(crate) fn enter_edit_mode(
 
 pub(crate) fn abort_and_return_to_workspace(
     ctx: &CommandContext,
-    _perm: &mut WorktreeWritePermission,
+    perm: &mut WorktreeWritePermission,
 ) -> Result<()> {
     let repository = ctx.repository();
 
     // Checkout gitbutler workspace branch
-    {
+    repository
-        repository
+        .set_head(WORKSPACE_BRANCH_REF)
-            .set_head(WORKSPACE_BRANCH_REF)
+        .context("Failed to set head reference")?;
-            .context("Failed to set head reference")?;
-        repository
-            .checkout_head(Some(CheckoutBuilder::new().force().remove_untracked(true)))
-            .context("Failed to checkout gitbutler/workspace")?;
-    }
 
-    // Checkout any stashed changes.
+    checkout_branch_trees(ctx, perm)?;
-    {
-        let stashed_workspace_changes_reference = repository
-            .find_reference(EDIT_UNCOMMITED_FILES_REF)
-            .context("Failed to find stashed workspace changes")?;
-        let stashed_workspace_changes_commit = stashed_workspace_changes_reference
-            .peel_to_commit()
-            .context("Failed to get stashed changes commit")?;
-
-        repository
-            .checkout_tree(
-                stashed_workspace_changes_commit.tree()?.as_object(),
-                Some(CheckoutBuilder::new().force().remove_untracked(true)),
-            )
-            .context("Failed to checkout workspace changes tree")?;
-    }
 
     Ok(())
 }
@@ -237,12 +183,6 @@ pub(crate) fn save_and_return_to_workspace(
     let commit = repository
         .find_commit(edit_mode_metadata.commit_oid)
         .context("Failed to find commit")?;
-    let stashed_workspace_changes_reference = repository
-        .find_reference(EDIT_UNCOMMITED_FILES_REF)
-        .context("Failed to find stashed workspace changes")?;
-    let stashed_workspace_changes_commit = stashed_workspace_changes_reference
-        .peel_to_commit()
-        .context("Failed to get stashed changes commit")?;
 
     let Some(mut virtual_branch) =
         find_virtual_branch_by_reference(ctx, &edit_mode_metadata.branch_reference)?
@@ -255,6 +195,7 @@ pub(crate) fn save_and_return_to_workspace(
     // Recommit commit
     let tree = repository.create_wd_tree()?;
 
+    let (_, committer) = repository.signatures()?;
     let commit_headers = commit
         .gitbutler_headers()
         .map(|commit_headers| CommitHeadersV2 {
@@ -266,7 +207,7 @@ pub(crate) fn save_and_return_to_workspace(
         .commit_with_signature(
             None,
             &commit.author(),
-            &commit.committer(),
+            &committer, // Use a new committer
             &commit.message_bstr().to_str_lossy(),
             &tree,
             &parents.iter().collect::<Vec<_>>(),
@@ -278,58 +219,35 @@ pub(crate) fn save_and_return_to_workspace(
     let new_branch_head = cherry_rebase(ctx, new_commit_oid, commit.id(), virtual_branch.head)
         .context("Failed to rebase commits onto new commit")?
         .unwrap_or(new_commit_oid);
+
+    // Update virtual_branch
+    let fearless_rebasing = ctx.project().succeeding_rebases;
+    let BranchHeadAndTree {
+        head: new_branch_head,
+        tree: new_branch_tree,
+    } = compute_updated_branch_head(
+        repository,
+        &virtual_branch,
+        new_branch_head,
+        fearless_rebasing,
+    )?;
+
+    virtual_branch.head = new_branch_head;
+    virtual_branch.tree = new_branch_tree;
+    virtual_branch.updated_timestamp_ms = gitbutler_time::time::now_ms();
+    vb_state
+        .set_branch(virtual_branch)
+        .context("Failed to update vbstate")?;
+
+    // Switch branch to gitbutler/workspace
     repository
-        .reference(
+        .set_head(WORKSPACE_BRANCH_REF)
-            &edit_mode_metadata.branch_reference,
+        .context("Failed to set head reference")?;
-            new_branch_head,
-            true,
-            "",
-        )
-        .context("Failed to reference new commit branch head")?;
 
-    // Move back to gitbutler/workspace and restore stashed changes
+    // Checkout the applied branches
-    {
+    checkout_branch_trees(ctx, perm)?;
-        repository
+    update_workspace_commit(&vb_state, ctx)?;
-            .set_head(WORKSPACE_BRANCH_REF)
+    list_virtual_branches(ctx, perm)?;
-            .context("Failed to set head reference")?;
-        repository
-            .checkout_head(Some(CheckoutBuilder::new().force().remove_untracked(true)))
-            .context("Failed to checkout gitbutler/workspace")?;
-
-        virtual_branch.head = new_branch_head;
-        virtual_branch.updated_timestamp_ms = gitbutler_time::time::now_ms();
-        vb_state
-            .set_branch(virtual_branch)
-            .context("Failed to update vbstate")?;
-
-        let workspace_commit_oid = update_workspace_commit(&vb_state, ctx)
-            .context("Failed to update gitbutler workspace")?;
-
-        let rebased_stashed_workspace_changes_commit = cherry_rebase_group(
-            repository,
-            workspace_commit_oid,
-            &[stashed_workspace_changes_commit.id()],
-            true,
-        )
-        .context("Failed to rebase stashed workspace commit changes")?;
-
-        let commit_thing = repository
-            .find_commit(rebased_stashed_workspace_changes_commit)
-            .context("Failed to find commit of rebased stashed workspace changes commit oid")?;
-
-        let tree_thing = repository
-            .find_real_tree(&commit_thing, Default::default())
-            .context("Failed to get tree of commit of rebased stashed workspace changes")?;
-
-        repository
-            .checkout_tree(
-                tree_thing.as_object(),
-                Some(CheckoutBuilder::new().force().remove_untracked(true)),
-            )
-            .context("Failed to checkout stashed changes tree")?;
-
-        list_virtual_branches(ctx, perm).context("Failed to list virtual branches")?;
-    }
 
     Ok(())
 }