Incorporate worktree information into graph and workspace.

Author: Byron

Cargo.lock

@@ -135,10 +135,10 @@ pub fn create_branch(
                 {
                     let segment = stack.segments.first().context("BUG: no empty stacks")?;
                     segment
-                    .ref_name
+                    .ref_info
                     .as_ref()
-                    .map(|rn| but_workspace::branch::create_reference::Anchor::AtSegment {
-                        ref_name: Cow::Borrowed(rn.as_ref()),
+                    .map(|ri| but_workspace::branch::create_reference::Anchor::AtSegment {
+                        ref_name: Cow::Borrowed(ri.ref_name.as_ref()),
                         position: Above,
                     })
                     .or_else(|| {

crates/but-api/src/commands/stack.rs

@@ -133,13 +133,13 @@ impl Graph {
         name: &gix::refs::FullNameRef,
     ) -> Option<(&Segment, &Commit)> {
         self.inner.node_weights().find_map(|s| {
-            if s.ref_name.as_ref().is_some_and(|rn| rn.as_ref() == name) {
+            if s.ref_name().is_some_and(|rn| rn == name) {
                 self.tip_skip_empty(s.id).map(|c| (s, c))
             } else {
                 s.commits.iter().find_map(|c| {
                     c.refs
                         .iter()
-                        .any(|rn| rn.as_ref() == name)
+                        .any(|ri| ri.ref_name.as_ref() == name)
                         .then_some((s, c))
                 })
             }
@@ -156,7 +156,7 @@ impl Graph {
     pub fn named_segment_by_ref_name(&self, name: &gix::refs::FullNameRef) -> Option<&Segment> {
         self.inner
             .node_weights()
-            .find(|s| s.ref_name.as_ref().is_some_and(|rn| rn.as_ref() == name))
+            .find(|s| s.ref_name().is_some_and(|rn| rn == name))
     }
 
     /// Starting a `segment`, ignore all segments that have no commit and return the first commit

crates/but-graph/src/api.rs

@@ -96,14 +96,14 @@ impl Graph {
             Ok(())
         };
         for node in self.inner.node_weights_mut() {
-            if let Some(rn) = node.ref_name.as_mut() {
-                anon(rn)?;
+            if let Some(ri) = node.ref_info.as_mut() {
+                anon(&mut ri.ref_name)?;
             }
             if let Some(rn) = node.remote_tracking_ref_name.as_mut() {
                 anon(rn)?;
             }
-            for rn in node.commits.iter_mut().flat_map(|c| c.refs.iter_mut()) {
-                anon(rn)?;
+            for ri in node.commits.iter_mut().flat_map(|c| c.refs.iter_mut()) {
+                anon(&mut ri.ref_name)?;
             }
             if let Some(SegmentMetadata::Workspace(md)) = node.metadata.as_mut() {
                 if let Some(rn) = md.target_ref.as_mut() {
@@ -154,7 +154,9 @@ impl Graph {
                     commit
                         .refs
                         .iter()
-                        .map(|rn| format!("►{}", { Self::ref_debug_string(rn) }))
+                        .map(|ri| format!("►{}", {
+                            Self::ref_debug_string(ri.ref_name.as_ref(), ri.worktree.as_ref())
+                        }))
                         .collect::<Vec<_>>()
                         .join(", ")
                 )
@@ -163,34 +165,43 @@ impl Graph {
     }
 
     /// Shorten the given `name` so it's still clear if it is a special ref (like tag) or not.
-    pub fn ref_debug_string(name: &gix::refs::FullName) -> String {
-        let (cat, sn) = name.category_and_short_name().expect("valid refs");
+    pub fn ref_debug_string(
+        ref_name: &gix::refs::FullNameRef,
+        worktree: Option<&crate::Worktree>,
+    ) -> String {
+        let (cat, sn) = ref_name.category_and_short_name().expect("valid refs");
         // Only shorten those that look good and are unambiguous enough.
-        if matches!(cat, Category::LocalBranch | Category::RemoteBranch) {
-            sn
-        } else {
-            name.as_bstr()
-                .strip_prefix(b"refs/")
-                .map(|n| n.as_bstr())
-                .unwrap_or(name.as_bstr())
-        }
-        .to_string()
+        format!(
+            "{}{ws}",
+            if matches!(cat, Category::LocalBranch | Category::RemoteBranch) {
+                sn
+            } else {
+                ref_name
+                    .as_bstr()
+                    .strip_prefix(b"refs/")
+                    .map(|n| n.as_bstr())
+                    .unwrap_or(ref_name.as_bstr())
+            },
+            ws = worktree
+                .map(|ws| ws.debug_string(ref_name))
+                .unwrap_or_default()
+        )
     }
 
     /// Return a useful one-line string showing the relationship between `ref_name`, `remote_ref_name` and how
     /// they are linked with `sibling_id`.
     pub fn ref_and_remote_debug_string(
-        ref_name: Option<&gix::refs::FullName>,
+        ref_info: Option<&crate::RefInfo>,
         remote_ref_name: Option<&gix::refs::FullName>,
         sibling_id: Option<SegmentIndex>,
     ) -> String {
         format!(
             "{ref_name}{remote}",
-            ref_name = ref_name
+            ref_name = ref_info
                 .as_ref()
-                .map(|rn| format!(
+                .map(|ri| format!(
                     "{}{maybe_id}",
-                    Graph::ref_debug_string(rn),
+                    Graph::ref_debug_string(ri.ref_name.as_ref(), ri.worktree.as_ref()),
                     maybe_id = sibling_id
                         .filter(|_| remote_ref_name.is_none())
                         .map(|id| format!(" →:{}:", id.index()))
@@ -206,7 +217,7 @@ impl Graph {
                 .as_ref()
                 .map(|remote_ref_name| format!(
                     " <> {remote_name}{maybe_id}",
-                    remote_name = Graph::ref_debug_string(remote_ref_name),
+                    remote_name = Graph::ref_debug_string(remote_ref_name.as_ref(), None),
                     maybe_id = sibling_id
                         .map(|id| format!(" →:{}:", id.index()))
                         .unwrap_or_default()
@@ -289,14 +300,14 @@ impl Graph {
             let name = format!(
                 "{ref_name_and_remote}{maybe_centering_newline}",
                 ref_name_and_remote = Self::ref_and_remote_debug_string(
-                    s.ref_name.as_ref(),
+                    s.ref_info.as_ref(),
                     s.remote_tracking_ref_name.as_ref(),
                     s.sibling_segment_id
                 ),
                 maybe_centering_newline = if s.commits.is_empty() { "" } else { "\n" },
             );
             // Reduce noise by preferring ref-based entry-points.
-            let show_segment_entrypoint = s.ref_name.is_some()
+            let show_segment_entrypoint = s.ref_info.is_some()
                 && entrypoint.is_some_and(|(s, cidx)| s == sidx && matches!(cidx, None | Some(0)));
             let mut commits = s
                 .commits

crates/but-graph/src/debug.rs

@@ -6,6 +6,7 @@ use gix::{
     prelude::{ObjectIdExt, ReferenceExt},
     refs::Category,
 };
+use std::collections::BTreeMap;
 use tracing::instrument;
 
 use crate::{CommitFlags, CommitIndex, Edge, Graph, Segment, SegmentIndex, SegmentMetadata};
@@ -154,10 +155,17 @@ impl Graph {
                 // It's OK to default-initialise this here as overlays are only used when redoing
                 // the traversal.
                 let (_repo, meta, _entrypoint) = Overlay::default().into_parts(repo, meta);
+                let wt_by_branch = {
+                    // Assume linked worktrees are never unborn!
+                    let mut m = BTreeMap::new();
+                    m.insert(ref_name.clone(), vec![crate::Worktree::Main]);
+                    m
+                };
                 graph.insert_segment_set_entrypoint(branch_segment_from_name_and_meta(
                     Some((ref_name, None)),
                     &meta,
                     None,
+                    &wt_by_branch,
                 )?);
                 return Ok(graph);
             }
@@ -182,7 +190,7 @@ impl Graph {
             if let Some((rn, first_commit)) = s
                 .commits
                 .first_mut()
-                .and_then(|first_commit| s.ref_name.take().map(|rn| (rn, first_commit)))
+                .and_then(|first_commit| s.ref_info.take().map(|rn| (rn, first_commit)))
             {
                 first_commit.refs.push(rn);
             }
@@ -350,6 +358,8 @@ impl Graph {
         let tip_is_not_workspace_commit = !workspaces
             .iter()
             .any(|(_, wsrn, _)| Some(wsrn) == ref_name.as_ref());
+        let worktree_by_branch = worktree_branches(repo.for_worktree_only())?;
+
         let mut ctx = post::Context {
             repo,
             symbolic_remote_names: &symbolic_remote_names,
@@ -358,12 +368,14 @@ impl Graph {
             refs_by_id,
             hard_limit: false,
             dangerously_skip_postprocessing_for_debugging,
+            worktree_by_branch,
         };
         if tip_is_not_workspace_commit {
             let current = graph.insert_segment_set_entrypoint(branch_segment_from_name_and_meta(
                 None,
                 meta,
                 Some((&ctx.refs_by_id, tip)),
+                &ctx.worktree_by_branch,
             )?);
             _ = next.push_back_exhausted((
                 tip,
@@ -406,8 +418,12 @@ impl Graph {
                     max_limit.with_indirect_goal(tip, &mut goals),
                 )
             };
-            let mut ws_segment =
-                branch_segment_from_name_and_meta(Some((ws_ref, None)), meta, None)?;
+            let mut ws_segment = branch_segment_from_name_and_meta(
+                Some((ws_ref, None)),
+                meta,
+                None,
+                &ctx.worktree_by_branch,
+            )?;
             // The limits for the target ref and the worktree ref are synced so they can always find each other,
             // while being able to stop when the entrypoint is included.
             ws_segment.metadata = Some(SegmentMetadata::Workspace(ws_meta));
@@ -439,6 +455,7 @@ impl Graph {
                     Some((target_ref, None)),
                     meta,
                     None,
+                    &ctx.worktree_by_branch,
                 )?);
                 let (local_sidx, local_goal) =
                     if let Some((local_ref_name, target_local_tip)) = local_tip_info {
@@ -448,12 +465,13 @@ impl Graph {
                                 Some(target_segment),
                                 meta,
                                 Some((&ctx.refs_by_id, target_local_tip)),
+                                &ctx.worktree_by_branch,
                             )?);
                         // We use auto-naming based on ambiguity - if the name ends up something else,
                         // remove the nodes sibling link.
                         let has_sibling_link = {
                             let s = &mut graph[local_sidx];
-                            if s.ref_name.as_ref().is_none_or(|rn| rn != &local_ref_name) {
+                            if s.ref_name().is_none_or(|rn| rn != local_ref_name.as_ref()) {
                                 s.sibling_segment_id = None;
                                 false
                             } else {
@@ -504,6 +522,7 @@ impl Graph {
                     None,
                     meta,
                     Some((&ctx.refs_by_id, extra_target)),
+                    &ctx.worktree_by_branch,
                 )?);
                 _ = next.push_front_exhausted((
                     extra_target,
@@ -550,15 +569,19 @@ impl Graph {
                     None,
                     meta,
                     Some((&ctx.refs_by_id, segment_tip.detach())),
+                    &ctx.worktree_by_branch,
                 )?;
 
                 // However, if this is a remote segment that is explicitly mentioned, and we couldn't name
                 // it, then just fix it up here as we really want that name.
                 let is_remote = segment_name
                     .category()
                     .is_some_and(|c| c == Category::RemoteBranch);
-                if segment.ref_name.is_none() && is_remote {
-                    segment.ref_name = Some(segment_name.clone());
+                if segment.ref_info.is_none() && is_remote {
+                    segment.ref_info = Some(crate::RefInfo::from_ref(
+                        segment_name.clone(),
+                        &ctx.worktree_by_branch,
+                    ));
                     segment.metadata = meta
                         .branch_opt(segment_name.as_ref())?
                         .map(SegmentMetadata::Branch);
@@ -577,6 +600,7 @@ impl Graph {
             &mut graph,
             &mut next,
             (ws_tips, repo, meta),
+            &ctx.worktree_by_branch,
         )?;
         max_commits_recharge_location.sort();
         while let Some((id, mut propagated_flags, instruction, mut limit)) = next.pop_front() {
@@ -614,6 +638,7 @@ impl Graph {
                             &info,
                             &ctx.refs_by_id,
                             meta,
+                            &ctx.worktree_by_branch,
                         )?
                         .unwrap_or(src_sidx);
                         e.insert(src_sidx);
@@ -641,6 +666,7 @@ impl Graph {
                             None,
                             meta,
                             Some((&ctx.refs_by_id, id)),
+                            &ctx.worktree_by_branch,
                         )?;
                         let segment_below = graph.connect_new_segment(
                             parent_above,
@@ -672,6 +698,7 @@ impl Graph {
                 limit,
                 &mut goals,
                 &next,
+                &ctx.worktree_by_branch,
             )?;
 
             let segment = &mut graph[segment_idx_for_id];
@@ -700,8 +727,9 @@ impl Graph {
                     refs_at_commit_before_removal
                         .clone()
                         .into_iter()
-                        .filter(|rn| segment.ref_name.as_ref() != Some(rn))
+                        .filter(|rn| segment.ref_name() != Some(rn.as_ref()))
                         .collect(),
+                    &ctx.worktree_by_branch,
                 )?,
             );
 
@@ -738,7 +766,7 @@ impl Graph {
                     .tip_skip_empty(tip_sidx)
                     .context("BUG: entrypoint must eventually point to a commit")?
                     .id;
-                let ref_name = self[tip_sidx].ref_name.clone();
+                let ref_name = self[tip_sidx].ref_info.clone().map(|ri| ri.ref_name);
                 (tip, ref_name)
             }
         };

crates/but-graph/src/init/mod.rs

@@ -182,6 +182,10 @@ impl<'repo> OverlayRepo<'repo> {
         self.inner
     }
 
+    pub fn for_worktree_only(&self) -> &'repo gix::Repository {
+        self.inner
+    }
+
     pub fn remote_names(&self) -> gix::remote::Names<'repo> {
         self.inner.remote_names()
     }