Handle commiting to workspace correctly

Author: Caleb-T-Owens

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

@@ -1,22 +1,25 @@
+use std::cmp::Ordering;
+use std::collections::HashMap;
 use std::{path::PathBuf, vec};
 
 use anyhow::{anyhow, Context, Result};
-use bstr::ByteSlice;
 use gitbutler_branch::BranchCreateRequest;
 use gitbutler_branch::{self, GITBUTLER_WORKSPACE_REFERENCE};
 use gitbutler_cherry_pick::RepositoryExt as _;
 use gitbutler_command_context::CommandContext;
 use gitbutler_commit::commit_ext::CommitExt;
+use gitbutler_diff::diff_files_into_hunks;
 use gitbutler_error::error::Marker;
-use gitbutler_operating_modes::OPEN_WORKSPACE_REFS;
+use gitbutler_operating_modes::{OPEN_WORKSPACE_REFS, WORKSPACE_BRANCH_REF};
 use gitbutler_oxidize::{git2_to_gix_object_id, gix_to_git2_oid, GixRepositoryExt};
-use gitbutler_project::access::WorktreeWritePermission;
+use gitbutler_project::access::{WorktreeReadPermission, WorktreeWritePermission};
 use gitbutler_repo::logging::{LogUntil, RepositoryExt as _};
-use gitbutler_repo::RepositoryExt;
+use gitbutler_repo::rebase::cherry_rebase_group;
 use gitbutler_repo::SignaturePurpose;
-use gitbutler_stack::{Stack, VirtualBranchesHandle};
+use gitbutler_stack::{Stack, StackId, VirtualBranchesHandle};
 use tracing::instrument;
 
+use crate::compute_workspace_dependencies;
 use crate::{branch_manager::BranchManagerExt, conflicts, VirtualBranchesExt};
 
 const WORKSPACE_HEAD: &str = "Workspace Head";
@@ -322,28 +325,29 @@ fn verify_current_branch_name(ctx: &CommandContext) -> Result<&CommandContext> {
     }
 }
 
-// TODO(ST): Probably there should not be an implicit vbranch creation here.
-fn verify_head_is_clean(ctx: &CommandContext, perm: &mut WorktreeWritePermission) -> Result<()> {
-    let head_commit = ctx
-        .repo()
-        .head()
-        .context("failed to get head")?
-        .peel_to_commit()
-        .context("failed to peel to commit")?;
+#[derive(Debug)]
+pub enum WorkspaceState {
+    OffWorkspaceCommit {
+        workspace_commit: git2::Oid,
+        extra_commits: Vec<git2::Oid>,
+    },
+    OnWorkspaceCommit,
+}
 
+pub fn workspace_state(
+    ctx: &CommandContext,
+    _perm: &WorktreeReadPermission,
+) -> Result<WorkspaceState> {
+    let repository = ctx.repo();
     let vb_handle = VirtualBranchesHandle::new(ctx.project().gb_dir());
-    let default_target = vb_handle
-        .get_default_target()
-        .context("failed to get default target")?;
+    let default_target = vb_handle.get_default_target()?;
 
-    let commits = ctx
-        .repo()
-        .log(
-            head_commit.id(),
-            LogUntil::Commit(default_target.sha),
-            false,
-        )
-        .context("failed to get log")?;
+    let head_commit = repository.head()?.peel_to_commit()?;
+    let commits = repository.log(
+        head_commit.id(),
+        LogUntil::Commit(default_target.sha),
+        false,
+    )?;
 
     let workspace_index = commits
         .iter()
@@ -353,69 +357,125 @@ fn verify_head_is_clean(ctx: &CommandContext, perm: &mut WorktreeWritePermission
                     || message.starts_with(GITBUTLER_INTEGRATION_COMMIT_TITLE)
             })
         })
-        .context("GitButler workspace commit not found")?;
+        .context("")?;
     let workspace_commit = &commits[workspace_index];
-    let mut extra_commits = commits[..workspace_index].to_vec();
-    extra_commits.reverse();
+    let extra_commits = commits[..workspace_index].to_vec();
 
     if extra_commits.is_empty() {
         // no extra commits found, so we're good
-        return Ok(());
+        return Ok(WorkspaceState::OnWorkspaceCommit);
     }
 
-    ctx.repo()
-        .reset(workspace_commit.as_object(), git2::ResetType::Soft, None)
-        .context("failed to reset to workspace commit")?;
-
-    let branch_manager = ctx.branch_manager();
-    let mut new_branch = branch_manager
-        .create_virtual_branch(
-            &BranchCreateRequest {
-                name: extra_commits
-                    .last()
-                    .map(|commit| commit.message_bstr().to_string()),
-                ..Default::default()
-            },
-            perm,
-        )
-        .context("failed to create virtual branch")?;
-
-    // rebasing the extra commits onto the new branch
-    let mut head = new_branch.head();
-    for commit in extra_commits {
-        let new_branch_head = ctx
-            .repo()
-            .find_commit(head)
-            .context("failed to find new branch head")?;
-
-        let rebased_commit_oid = ctx
-            .repo()
-            .commit_with_signature(
-                None,
-                &commit.author(),
-                &commit.committer(),
-                &commit.message_bstr().to_str_lossy(),
-                &commit.tree().unwrap(),
-                &[&new_branch_head],
-                None,
-            )
-            .context(format!(
-                "failed to rebase commit {} onto new branch",
-                commit.id()
-            ))?;
+    Ok(WorkspaceState::OffWorkspaceCommit {
+        workspace_commit: workspace_commit.id(),
+        extra_commits: extra_commits
+            .iter()
+            .map(git2::Commit::id)
+            .collect::<Vec<_>>(),
+    })
+}
+
+// TODO(ST): Probably there should not be an implicit vbranch creation here.
+fn verify_head_is_clean(ctx: &CommandContext, perm: &mut WorktreeWritePermission) -> Result<()> {
+    let repository = ctx.repo();
+    let head_commit = repository.head()?.peel_to_commit()?;
+
+    let WorkspaceState::OffWorkspaceCommit {
+        workspace_commit,
+        extra_commits,
+    } = dbg!(workspace_state(ctx, perm.read_permission())?)
+    else {
+        return Ok(());
+    };
 
-        let rebased_commit = ctx.repo().find_commit(rebased_commit_oid).context(format!(
-            "failed to find rebased commit {}",
-            rebased_commit_oid
-        ))?;
+    let best_stack_id = find_best_stack_for_changes(ctx, perm, head_commit.id(), workspace_commit)?;
 
-        new_branch.set_stack_head(ctx, rebased_commit.id(), Some(rebased_commit.tree_id()))?;
+    if let Some(best_stack_id) = best_stack_id {
+        let vb_handle = VirtualBranchesHandle::new(ctx.project().gb_dir());
+        let mut stack = vb_handle.get_stack_in_workspace(best_stack_id)?;
 
-        head = rebased_commit.id();
+        let new_head = cherry_rebase_group(repository, stack.head(), &extra_commits, false)?;
+
+        stack.set_stack_head(
+            ctx,
+            new_head,
+            Some(repository.find_commit(new_head)?.tree_id()),
+        )?;
+
+        update_workspace_commit(&vb_handle, ctx)?;
+    } else {
+        // There is no stack which can hold the commits so we should just unroll those changes
+        repository.reference(WORKSPACE_BRANCH_REF, workspace_commit, true, "")?;
+        repository.set_head(WORKSPACE_BRANCH_REF)?;
     }
+
     Ok(())
 }
 
+fn find_best_stack_for_changes(
+    ctx: &CommandContext,
+    perm: &mut WorktreeWritePermission,
+    head_commit: git2::Oid,
+    workspace_commit: git2::Oid,
+) -> Result<Option<StackId>> {
+    let vb_state = VirtualBranchesHandle::new(ctx.project().gb_dir());
+    let default_target = vb_state.get_default_target()?;
+    let repository = ctx.repo();
+    let stacks = vb_state.list_stacks_in_workspace()?;
+
+    let head_commit = repository.find_commit(head_commit)?;
+
+    let diffs = gitbutler_diff::trees(
+        ctx.repo(),
+        &repository.find_commit(workspace_commit)?.tree()?,
+        &head_commit.tree()?,
+        true,
+    )?;
+    let base_diffs: HashMap<_, _> = diff_files_into_hunks(&diffs).collect();
+    let workspace_dependencies =
+        compute_workspace_dependencies(ctx, &default_target.sha, &base_diffs, &stacks)?;
+
+    match workspace_dependencies.commit_dependent_diffs.len().cmp(&1) {
+        Ordering::Greater => {
+            // The commits are locked to multiple stacks. We can't correctly assign it
+            // to any one stack, so the commits should be undone.
+            Ok(None)
+        }
+        Ordering::Equal => {
+            // There is one stack which the commits are locked to, so the commits
+            // should be added to that particular stack.
+            let stack_id = workspace_dependencies
+                .commit_dependent_diffs
+                .keys()
+                .next()
+                .expect("Values was asserted length 1 above");
+            Ok(Some(*stack_id))
+        }
+        Ordering::Less => {
+            // We should return the branch selected for changes, or create a new default branch.
+            let mut stacks = vb_state.list_stacks_in_workspace()?;
+            stacks.sort_by_key(|stack| stack.selected_for_changes.unwrap_or(0));
+
+            if let Some(stack) = stacks.last() {
+                return Ok(Some(stack.id));
+            }
+
+            let branch_manager = ctx.branch_manager();
+            let new_stack = branch_manager
+                .create_virtual_branch(
+                    &BranchCreateRequest {
+                        name: Some(head_commit.message_bstr().to_string()),
+                        ..Default::default()
+                    },
+                    perm,
+                )
+                .context("failed to create virtual branch")?;
+
+            Ok(Some(new_stack.id))
+        }
+    }
+}
+
 fn invalid_head_err(head_name: &str) -> anyhow::Error {
     anyhow!(
         "project is on {head_name}. Please checkout {} to continue",

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

@@ -38,7 +38,7 @@ pub use dependencies::compute_workspace_dependencies;
 pub mod upstream_integration;
 
 mod integration;
-pub use integration::{update_workspace_commit, verify_branch};
+pub use integration::{update_workspace_commit, verify_branch, workspace_state, WorkspaceState};
 
 mod file;
 pub use file::{Get, RemoteBranchFile};

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

@@ -8,7 +8,8 @@ use crate::{
     remote::branch_to_remote_branch,
     stack::stack_series,
     status::{get_applied_status, get_applied_status_cached},
-    Get, RemoteBranchData, VirtualBranchHunkRange, VirtualBranchHunkRangeMap, VirtualBranchesExt,
+    verify_branch, Get, RemoteBranchData, VirtualBranchHunkRange, VirtualBranchHunkRangeMap,
+    VirtualBranchesExt,
 };
 use anyhow::{anyhow, bail, Context, Result};
 use bstr::{BString, ByteSlice};
@@ -326,6 +327,9 @@ pub fn list_virtual_branches_cached(
 ) -> Result<StackListResult> {
     assure_open_workspace_mode(ctx)
         .context("Listing virtual branches requires open workspace mode")?;
+    // Make sure that the workspace commit is the head of the branch before listing.
+    verify_branch(ctx, perm)?;
+
     let mut branches: Vec<VirtualBranch> = Vec::new();
 
     let vb_state = ctx.project().virtual_branches();

crates/gitbutler-branch-actions/tests/virtual_branches/create_virtual_branch_from_branch.rs

@@ -382,6 +382,8 @@ mod conflict_cases {
         let branch_refname =
             gitbutler_branch_actions::save_and_unapply_virutal_branch(ctx, branch.id).unwrap();
 
+        gitbutler_branch_actions::list_virtual_branches(ctx).unwrap();
+
         // Make X and set base branch to X
         let mut tree_builder = git_repository
             .treebuilder(Some(

crates/gitbutler-stack/src/stack.rs

@@ -511,6 +511,9 @@ impl Stack {
         self.updated_timestamp_ms = gitbutler_time::time::now_ms();
         #[allow(deprecated)] // this is the only place where this is allowed
         self.set_head(commit_id);
+        // TODO(CTO): Determine whether this does anything. If we're not
+        // calling `checkout_branch_trees` right after this, then it will
+        // likly get overridden by the next `list_virtual_branches`.
         if let Some(tree) = tree {
             self.tree = tree;
         }

crates/gitbutler-watcher/src/handler.rs

@@ -1,7 +1,9 @@
 use std::{path::PathBuf, sync::Arc};
 
 use anyhow::{Context, Result};
-use gitbutler_branch_actions::{internal::StackListResult, VirtualBranches};
+use gitbutler_branch_actions::{
+    internal::StackListResult, verify_branch, workspace_state, VirtualBranches, WorkspaceState,
+};
 use gitbutler_command_context::CommandContext;
 use gitbutler_diff::DiffByPathMap;
 use gitbutler_error::error::Marker;
@@ -196,6 +198,29 @@ impl Handler {
         Ok(())
     }
 
+    fn handle_commited_workspace(&self, ctx: &CommandContext) -> Result<()> {
+        // Don't worry about non-workspace settings (yet)
+        if !in_open_workspace_mode(ctx) {
+            return Ok(());
+        }
+
+        let workspace_state = {
+            let guard = ctx.project().exclusive_worktree_access();
+            let permission = guard.read_permission();
+            workspace_state(ctx, permission)?
+        };
+        if matches!(workspace_state, WorkspaceState::OffWorkspaceCommit { .. }) {
+            {
+                let mut guard = ctx.project().exclusive_worktree_access();
+                let permission = guard.write_permission();
+                verify_branch(ctx, permission)?;
+            }
+            self.calculate_virtual_branches(ctx, None)?;
+        }
+
+        Ok(())
+    }
+
     pub fn git_files_change(&self, paths: Vec<PathBuf>, ctx: &CommandContext) -> Result<()> {
         for path in paths {
             let Some(file_name) = path.to_str() else {
@@ -206,6 +231,7 @@ impl Handler {
                     self.emit_app_event(Change::GitFetch(ctx.project().id))?;
                 }
                 "logs/HEAD" => {
+                    self.handle_commited_workspace(ctx)?;
                     self.emit_app_event(Change::GitActivity(ctx.project().id))?;
                 }
                 "index" => {