Cover more cases when applying branches to prevent failures.

Otherwise, we think we applied the branch, but it's not enough
due to ambiguity. Then the stack or segment wouldn't actually
show up unless we make the right choice of adding it as dependent
branch instead of stack.

Author: Byron

crates/but-core/src/ref_metadata.rs

@@ -36,8 +36,27 @@ pub struct Workspace {
     pub push_remote: Option<String>,
 }
 
+impl Workspace {}
+
 /// Mutations
 impl Workspace {
+    /// Remove the named segment `branch`, which removes the whole stack if it's empty after removing a segment
+    /// of that name.
+    /// Returns `true` if it was removed or `false` if it wasn't found.
+    pub fn remove_segment(&mut self, branch: &FullNameRef) -> bool {
+        let Some((stack_idx, segment_idx)) = self.find_owner_indexes_by_name(branch) else {
+            return false;
+        };
+
+        let stack = &mut self.stacks[stack_idx];
+        stack.branches.remove(segment_idx);
+
+        if stack.branches.is_empty() {
+            self.stacks.remove(stack_idx);
+        }
+        true
+    }
+
     /// Insert `branch` as new stack if it's not yet contained in the workspace and if `order` is not `None` or push
     /// it to the end of the stack list.
     /// Note that `order` is only relevant at insertion time.
@@ -68,6 +87,28 @@ impl Workspace {
         }
         true
     }
+
+    /// Insert `branch` as new segment if it's not yet contained in the workspace,
+    /// and insert it above the given `anchor` segment name, which maybe the tip of a stack or any segment within one
+    /// Returns `true` if the ref was newly added, or `false` if it already existed, or `None` if `anchor` didn't exist.
+    pub fn insert_new_segment_above_anchor_if_not_present(
+        &mut self,
+        branch: &FullNameRef,
+        anchor: &FullNameRef,
+    ) -> Option<bool> {
+        if self.contains_ref(branch) {
+            return Some(false);
+        };
+        let (stack_idx, segment_idx) = self.find_owner_indexes_by_name(anchor)?;
+        self.stacks[stack_idx].branches.insert(
+            segment_idx,
+            WorkspaceStackBranch {
+                ref_name: branch.to_owned(),
+                archived: false,
+            },
+        );
+        Some(true)
+    }
 }
 
 /// Access

crates/but-core/tests/core/ref_metadata.rs

@@ -2,7 +2,7 @@ mod workspace {
     use but_core::ref_metadata::Workspace;
 
     #[test]
-    fn add_new_stack_if_not_present() {
+    fn add_new_stack_if_not_present_journey() {
         let mut ws = Workspace::default();
         assert_eq!(ws.stacks.len(), 0);
 
@@ -18,6 +18,100 @@ mod workspace {
         let c_ref = r("refs/heads/C");
         assert!(ws.add_or_insert_new_stack_if_not_present(c_ref, None));
         assert_eq!(ws.stack_names().collect::<Vec<_>>(), [b_ref, a_ref, c_ref]);
+
+        assert!(ws.remove_segment(a_ref));
+        assert!(ws.remove_segment(b_ref));
+        assert!(!ws.remove_segment(b_ref));
+        assert!(ws.remove_segment(c_ref));
+        assert!(!ws.remove_segment(c_ref));
+
+        // Everything should be removed.
+        insta::assert_debug_snapshot!(ws, @r"
+        Workspace {
+            ref_info: RefInfo { created_at: None, updated_at: None },
+            stacks: [],
+            target_ref: None,
+            push_remote: None,
+        }
+        ");
+    }
+
+    #[test]
+    fn insert_new_segment_above_anchor_if_not_present_journey() {
+        let mut ws = Workspace::default();
+        assert_eq!(ws.stacks.len(), 0);
+
+        let a_ref = r("refs/heads/A");
+        let b_ref = r("refs/heads/B");
+        assert_eq!(
+            ws.insert_new_segment_above_anchor_if_not_present(b_ref, a_ref),
+            None,
+            "anchor doesn't exist"
+        );
+        assert!(ws.add_or_insert_new_stack_if_not_present(a_ref, None));
+        assert_eq!(
+            ws.insert_new_segment_above_anchor_if_not_present(b_ref, a_ref),
+            Some(true),
+            "anchor existed and it was added"
+        );
+        assert_eq!(
+            ws.insert_new_segment_above_anchor_if_not_present(b_ref, a_ref),
+            Some(false),
+            "anchor existed and it was NOT added as it already existed"
+        );
+
+        let c_ref = r("refs/heads/C");
+        assert_eq!(
+            ws.insert_new_segment_above_anchor_if_not_present(c_ref, a_ref),
+            Some(true)
+        );
+
+        insta::assert_snapshot!(but_testsupport::sanitize_uuids_and_timestamps(format!("{ws:#?}")), @r#"
+        Workspace {
+            ref_info: RefInfo { created_at: None, updated_at: None },
+            stacks: [
+                WorkspaceStack {
+                    id: 1,
+                    branches: [
+                        WorkspaceStackBranch {
+                            ref_name: FullName(
+                                "refs/heads/B",
+                            ),
+                            archived: false,
+                        },
+                        WorkspaceStackBranch {
+                            ref_name: FullName(
+                                "refs/heads/C",
+                            ),
+                            archived: false,
+                        },
+                        WorkspaceStackBranch {
+                            ref_name: FullName(
+                                "refs/heads/A",
+                            ),
+                            archived: false,
+                        },
+                    ],
+                },
+            ],
+            target_ref: None,
+            push_remote: None,
+        }
+        "#);
+
+        assert!(ws.remove_segment(b_ref));
+        assert!(ws.remove_segment(a_ref));
+        assert!(ws.remove_segment(c_ref));
+
+        // Everything should be removed.
+        insta::assert_debug_snapshot!(ws, @r"
+        Workspace {
+            ref_info: RefInfo { created_at: None, updated_at: None },
+            stacks: [],
+            target_ref: None,
+            push_remote: None,
+        }
+        ");
     }
 
     fn r(name: &str) -> &gix::refs::FullNameRef {

crates/but-workspace/src/branch/apply.rs

@@ -157,6 +157,7 @@ pub(crate) mod function {
                 repo,
                 checkout::Options {
                     uncommitted_changes,
+                    skip_head_update: false,
                 },
             )?;
             let graph = workspace.graph.redo_traversal_with_overlay(
@@ -335,8 +336,9 @@ pub(crate) mod function {
             );
         }
         // At this point, the workspace-metadata already knows the new branch(es), but the workspace itself
-        // doesn't see one or more of to-be-applied branches. These are, however, part of the graph by now,
-        // and we want to try to create a workspace merge.
+        // doesn't see one or more of to-be-applied branches (to become stacks).
+        // These are, however, part of the graph by now, and we want to try to create a workspace
+        // merge.
         let mut in_memory_repo = repo.clone().for_tree_diffing()?.with_object_memory();
         let merge_result = WorkspaceCommit::from_new_merge_with_metadata(
             &ws_md.stacks,
@@ -353,6 +355,110 @@ pub(crate) mod function {
             });
         }
 
+        let prev_head_id = graph
+            .entrypoint_commit()
+            .context("BUG: how is it possible that there is no head commit?")?
+            .id;
+        let mut new_head_id = merge_result.workspace_commit_id;
+        let overlay =
+            overlay.with_entrypoint(new_head_id, Some(workspace_ref_name_to_update.clone()));
+        let mut graph =
+            graph.redo_traversal_with_overlay(&in_memory_repo, meta, overlay.clone())?;
+
+        let workspace = graph.to_workspace()?;
+        let collect_unapplied_branches = |workspace: &but_graph::projection::Workspace| {
+            branches_to_apply
+                .iter()
+                .filter(|rn| !workspace.refname_is_segment(rn))
+                .collect::<Vec<_>>()
+        };
+        let unapplied_branches = collect_unapplied_branches(&workspace);
+        if !unapplied_branches.is_empty() {
+            // Now that the merge is done, try to redo the operation one last time with dependent branches instead.
+            // Only do that for the still unapplied branches, which should always find some sort of anchor.
+            let ws_mut: &mut Workspace = &mut ws_md;
+            for branch_to_remove in &unapplied_branches {
+                ws_mut.remove_segment(branch_to_remove);
+            }
+            for rn in &unapplied_branches {
+                // Here we have to check if the new ref would be able to become its own stack,
+                // or if it has to be a dependent branch. Stacks only work if the ref rests on a base
+                // outside the workspace, so if we find it in the workspace (in an ambiguous spot) it must be
+                // a dependent branch
+                if let Some(segment_to_insert_above) = workspace
+                    .stacks
+                    .iter()
+                    .flat_map(|stack| stack.segments.iter())
+                    .find_map(|segment| {
+                        segment.commits.iter().flat_map(|c| c.refs.iter()).find_map(
+                            |ambiguous_rn| {
+                                (ambiguous_rn.as_ref() == **rn)
+                                    .then_some(segment.ref_name.as_ref())
+                                    .flatten()
+                            },
+                        )
+                    })
+                {
+                    if ws_mut
+                        .insert_new_segment_above_anchor_if_not_present(
+                            rn,
+                            segment_to_insert_above.as_ref(),
+                        )
+                        .is_none()
+                    {
+                        // For now bail, until we know it's worth fixing this case automatically.
+                        bail!(
+                            "Missing reference {segment_to_insert_above} which should be known to workspace metadata to serve as insertion position for {rn}"
+                        );
+                    }
+                } else {
+                    bail!(
+                        "Unexpectedly failed to find anchor for {rn} to make it a dependent branch"
+                    )
+                }
+            }
+
+            // Redo the merge, with the different stack configuration.
+            // Note that this is the exception, typically using stacks will be fine.
+            let merge_result = WorkspaceCommit::from_new_merge_with_metadata(
+                &ws_md.stacks,
+                workspace.graph,
+                &in_memory_repo,
+                Some(branch),
+            )?;
+
+            if merge_result.has_conflicts() && on_workspace_conflict.should_abort() {
+                return Ok(Outcome {
+                    graph: Cow::Owned(graph),
+                    workspace_ref_created: needs_ws_ref_creation,
+                    workspace_merge: Some(merge_result),
+                });
+            }
+            new_head_id = merge_result.workspace_commit_id;
+            let ws_md_override = Some((workspace_ref_name_to_update.clone(), (*ws_md).clone()));
+
+            graph = graph.redo_traversal_with_overlay(
+                &in_memory_repo,
+                meta,
+                overlay
+                    .with_entrypoint(new_head_id, Some(workspace_ref_name_to_update.clone()))
+                    .with_workspace_metadata_override(ws_md_override),
+            )?;
+            let workspace = graph.to_workspace()?;
+            let unapplied_branches = collect_unapplied_branches(&workspace);
+
+            if !unapplied_branches.is_empty() {
+                bail!(
+                    "Unexpectedly failed to apply {branches} which is/are still not in the workspace",
+                    branches = unapplied_branches
+                        .iter()
+                        .map(|rn| rn.shorten().to_string())
+                        .collect::<Vec<_>>()
+                        .join(", ")
+                )
+            }
+        }
+
         // All work is done, persist and exit.
         // Note that it could be that some stacks aren't merged in,
         // while being present in the workspace metadata.
@@ -361,23 +467,14 @@ pub(crate) mod function {
             storage.persist(repo)?;
             drop(in_memory_repo);
         }
-        let prev_head_id = graph
-            .entrypoint_commit()
-            .context("BUG: how is it possible that there is no head commit?")?
-            .id;
-        let new_head_id = merge_result.workspace_commit_id;
-        let graph = graph.redo_traversal_with_overlay(
-            repo,
-            meta,
-            overlay.with_entrypoint(new_head_id, Some(workspace_ref_name_to_update.clone())),
-        )?;
         persist_metadata(meta, &branches_to_apply, &ws_md)?;
         crate::branch::safe_checkout(
             prev_head_id,
             new_head_id,
             repo,
             checkout::Options {
                 uncommitted_changes,
+                skip_head_update: true,
             },
         )?;
 
@@ -416,10 +513,9 @@ pub(crate) mod function {
         new_ref_target: gix::ObjectId,
         new_ref: Option<&gix::refs::FullNameRef>,
     ) -> anyhow::Result<()> {
-        // This also means we want HEAD to point to it.
-        let head_message = "GitButler switch to workspace during apply-branch".into();
         let edits = match new_ref {
             None => {
+                let head_message = "GitButler checkout workspace during apply-branch".into();
                 vec![RefEdit {
                     change: Change::Update {
                         log: LogChange {
@@ -435,6 +531,8 @@ pub(crate) mod function {
                 }]
             }
             Some(new_ref) => {
+                // This also means we want HEAD to point to it.
+                let head_message = "GitButler switch to workspace during apply-branch".into();
                 vec![
                     RefEdit {
                         change: Change::Update {

crates/but-workspace/src/branch/checkout/function.rs

@@ -50,6 +50,7 @@ pub fn safe_checkout(
     repo: &gix::Repository,
     Options {
         uncommitted_changes: conflicting_worktree_changes_opts,
+        skip_head_update,
     }: Options,
 ) -> anyhow::Result<Outcome> {
     let source_tree = current_head_id.attach(repo).object()?.peel_to_tree()?;
@@ -138,7 +139,7 @@ pub fn safe_checkout(
     }
 
     let mut head_update = None;
-    if new_object.kind.is_commit() {
+    if new_object.kind.is_commit() && !skip_head_update {
         let needs_update = repo
             .head()?
             .id()

crates/but-workspace/src/branch/checkout/mod.rs

@@ -17,6 +17,10 @@ pub enum UncommitedWorktreeChanges {
 pub struct Options {
     /// How to deal with uncommitted changes.
     pub uncommitted_changes: UncommitedWorktreeChanges,
+    /// If `true`, do not change `HEAD` to the new commit.
+    ///
+    /// This is typically to be avoided, but may be used if you want to change the HEAD location yourself.
+    pub skip_head_update: bool,
 }
 
 /// The successful outcome of [super::safe_checkout()] operation.

crates/but-workspace/tests/fixtures/scenario/no-ws-ref-stack-and-dependent-branch.sh

@@ -0,0 +1,21 @@
+#!/usr/bin/env bash
+
+set -eu -o pipefail
+
+source "${BASH_SOURCE[0]%/*}/shared.sh"
+
+### General Description
+
+# A target branch along with a stack and a depdendent branch at the same tip, without workspace commit.
+git init
+commit M
+setup_target_to_match_main
+git checkout -b A
+commit A1
+git branch D
+git branch E
+commit A2
+git branch B
+git branch C
+
+git checkout main