Incorporate worktree information into graph and workspace.

Author: Byron

crates/but-graph/src/api.rs

@@ -139,7 +139,7 @@ impl Graph {
                 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))
                 })
             }

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

@@ -1,13 +1,7 @@
 //! Utilities for graph-walking specifically.
-use std::{collections::BTreeSet, ops::Deref};
-
-use anyhow::{Context, bail};
-use but_core::{RefMetadata, ref_metadata};
-use gix::{hashtable::hash_map::Entry, reference::Category, traverse::commit::Either};
-use petgraph::Direction;
-
 use crate::{
     Commit, CommitFlags, CommitIndex, Edge, Graph, Segment, SegmentIndex, SegmentMetadata,
+    Worktree,
     init::{
         Goals, PetGraph,
         overlay::{OverlayMetadata, OverlayRepo},
@@ -16,6 +10,14 @@ use crate::{
     },
     is_workspace_ref_name,
 };
+use anyhow::{Context, bail};
+use bstr::{BStr, BString};
+use but_core::{RefMetadata, ref_metadata};
+use gix::{hashtable::hash_map::Entry, reference::Category, traverse::commit::Either};
+use petgraph::Direction;
+use std::collections::BTreeMap;
+use std::path::PathBuf;
+use std::{collections::BTreeSet, ops::Deref};
 
 pub(crate) type RefsById = gix::hashtable::HashMap<gix::ObjectId, Vec<gix::refs::FullName>>;
 
@@ -941,3 +943,49 @@ pub fn prune_integrated_tips(graph: &mut Graph, next: &mut Queue) -> anyhow::Res
         });
     Ok(())
 }
+
+type WorktreeByBranch = BTreeMap<gix::refs::FullName, Vec<Worktree>>;
+
+pub fn worktree_branches(repo: &gix::Repository) -> anyhow::Result<WorktreeByBranch> {
+    fn maybe_insert_head(
+        head: Option<gix::Head<'_>>,
+        out: &mut WorktreeByBranch,
+    ) -> anyhow::Result<()> {
+        let Some((head, wd)) = head.and_then(|head| {
+            head.repo.worktree().map(|wt| {
+                (
+                    head,
+                    match wt.id() {
+                        None => Worktree::Main,
+                        Some(id) => Worktree::LinkedId(id.to_owned()),
+                    },
+                )
+            })
+        }) else {
+            return Ok(());
+        };
+
+        out.entry("HEAD".try_into().expect("valid"))
+            .or_default()
+            .push(wd.to_owned());
+        let mut ref_chain = Vec::new();
+        let mut cursor = head.try_into_referent();
+        while let Some(ref_) = cursor {
+            ref_chain.push(ref_.name().to_owned());
+            cursor = ref_.follow().transpose()?;
+        }
+        for name in ref_chain {
+            out.entry(name).or_default().push(wd.to_owned());
+        }
+
+        Ok(())
+    }
+
+    let mut map = BTreeMap::new();
+    maybe_insert_head(repo.head().ok(), &mut map)?;
+    for proxy in repo.worktrees()? {
+        let repo = proxy.into_repo_with_possibly_inaccessible_worktree()?;
+        maybe_insert_head(repo.head().ok(), &mut map)?;
+    }
+    Ok(map)
+}

crates/but-graph/src/lib.rs

@@ -201,7 +201,7 @@
 mod segment;
 /// Use this for basic types like [`petgraph::Direction`], and graph algorithms.
 pub use petgraph;
-pub use segment::{Commit, CommitFlags, Segment, SegmentMetadata};
+pub use segment::{Commit, CommitFlags, RefInfo, Segment, SegmentMetadata, Worktree};
 
 mod api;
 /// Produce a graph from a Git repository.

crates/but-graph/src/segment.rs

@@ -1,6 +1,7 @@
-use bitflags::bitflags;
-
 use crate::{CommitIndex, SegmentIndex};
+use bitflags::bitflags;
+use bstr::BString;
+use std::borrow::Cow;
 
 /// A commit with must useful information extracted from the Git commit itself.
 #[derive(Clone, Eq, PartialEq)]
@@ -13,15 +14,57 @@ pub struct Commit {
     pub flags: CommitFlags,
     /// The references pointing to this commit, even after dereferencing tag objects.
     /// These can be names of tags and branches.
-    pub refs: Vec<gix::refs::FullName>,
+    pub refs: Vec<RefInfo>,
+}
+
+/// A structure to inform about a reference which was present at a commit.
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub struct RefInfo {
+    /// The name of the reference.
+    pub ref_name: gix::refs::FullName,
+    /// If `Some`, provide information about the worktree that checks out the reference at `ref_name`,
+    /// i.e. its `HEAD` points to `ref_name` directly or indirectly due to chains of .
+    ///
+    /// It is `None` if no worktree needs to be updated if this reference is changed.
+    pub worktree: Option<Worktree>,
+}
+
+/// Describes which worktree is checked out.
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub enum Worktree {
+    /// The main worktree, i.e. the primary workspace associated with this repository, is checked out.
+    ///
+    /// It cannot be removed.
+    Main,
+    /// The identifier of the worktree, which is always `.git/worktrees/<id>`,
+    /// indicating that this is a linked worktree that can be removed.
+    LinkedId(BString),
+}
+
+impl RefInfo {
+    /// Produce a string that identifies this instance concisely, and visually distinguishable.
+    pub fn debug_string(&self) -> String {
+        let ws = match self.worktree {
+            None => String::new(),
+            Some(ws) => match ws {
+                Worktree::Main => {
+                    format!("[🌳main]")
+                }
+                Worktree::LinkedId(id) => {
+                    format!("[📁{id}]")
+                }
+            },
+        };
+        format!("►{}{ws}", self.ref_name.shorten())
+    }
 }
 
 impl std::fmt::Debug for Commit {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let refs = self
             .refs
             .iter()
-            .map(|rn| format!("►{}", rn.shorten()))
+            .map(|ri| ri.debug_string())
             .collect::<Vec<_>>()
             .join(", ");
         write!(

crates/but-graph/tests/fixtures/scenarios.sh

@@ -120,6 +120,13 @@ git init multi-root
   git checkout main && git merge --allow-unrelated-histories C
 )
 
+git init ambiguous-worktrees
+(cd ambiguous-worktrees
+  commit M
+  git worktree add ../wt-outside-ambiguous-worktree
+  git worktree add wt-inside-ambiguous-worktree
+)
+
 # A single root that splits up into 4 branches and merges again
 git init four-diamond
 (cd four-diamond
@@ -205,6 +212,32 @@ git init special-branches
 
 mkdir ws
 (cd ws
+  git init ambiguous-worktrees
+  (cd ambiguous-worktrees
+    set -x
+    commit M1
+    commit M-base
+
+    git branch A
+    git worktree add -b A-inside wt-A-inside
+    git worktree add -b A-outside ../wt-A-outside
+
+    git checkout -b soon-origin-A main
+      commit A-remote
+    git checkout main
+      commit M-advanced
+      setup_target_to_match_main
+
+    git checkout -b B A
+      commit B
+    git checkout A
+    git worktree add wt-B-inside B
+
+    create_workspace_commit_once A B
+    setup_remote_tracking soon-origin-A A "move"
+    git worktree add wt-origin-A-inside origin/A
+  )
+
   git init remote-and-integrated-tracking-linear
   (cd remote-and-integrated-tracking-linear
      commit M1

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

@@ -723,6 +723,27 @@ fn special_branch_names_do_not_end_up_in_segment() -> anyhow::Result<()> {
     Ok(())
 }
 
+#[test]
+fn ambiguous_worktrees() -> anyhow::Result<()> {
+    let (repo, meta) = read_only_in_memory_scenario("ambiguous-worktrees")?;
+    insta::assert_snapshot!(visualize_commit_graph_all(&repo)?, @"* 85efbe4 (HEAD -> main, wt-outside-ambiguous-worktree, wt-inside-ambiguous-worktree) M");
+
+    let graph = Graph::from_head(&repo, &*meta, standard_options())?.validated()?;
+    // TODO: show worktree ids at least, make main workspace clear
+    insta::assert_snapshot!(graph_tree(&graph), @r"
+    └── 👉►:0[0]:main
+        └── ·85efbe4 (⌂|1) ►wt-inside-ambiguous-worktree, ►wt-outside-ambiguous-worktree
+    ");
+
+    insta::assert_snapshot!(graph_workspace(&graph.to_workspace()?), @r"
+    ⌂:0:main <> ✓!
+    └── ≡:0:main
+        └── :0:main
+            └── ·85efbe4 ►wt-inside-ambiguous-worktree, ►wt-outside-ambiguous-worktree
+    ");
+    Ok(())
+}
+
 mod with_workspace;
 
 mod utils;

crates/but-graph/tests/graph/init/with_workspace.rs

@@ -5462,6 +5462,56 @@ fn no_ws_commit_two_branches_no_target() -> anyhow::Result<()> {
     Ok(())
 }
 
+#[test]
+fn ambiguous_worktrees() -> anyhow::Result<()> {
+    let (repo, mut meta) = read_only_in_memory_scenario("ws/ambiguous-worktrees")?;
+    insta::assert_snapshot!(visualize_commit_graph_all(&repo)?, @r"
+    *   a5f94a2 (HEAD -> gitbutler/workspace) GitButler Workspace Commit
+    |\  
+    | * 3e01e28 (B) B
+    |/  
+    | * 8dc508f (origin/main, main) M-advanced
+    |/  
+    | * 197ddce (origin/A) A-remote
+    |/  
+    * 081bae9 (A-outside, A-inside, A) M-base
+    * 3183e43 M1
+    ");
+
+    add_stack_with_segments(&mut meta, 0, "A", StackState::InWorkspace, &[]);
+    let graph = Graph::from_head(&repo, &*meta, standard_options())?.validated()?;
+    // TODO: Show worktrees
+    insta::assert_snapshot!(graph_tree(&graph), @r"
+    ├── 👉📕►►►:0[0]:gitbutler/workspace
+    │   └── ·a5f94a2 (⌂|🏘|1)
+    │       ├── ►:5[1]:B
+    │       │   └── ·3e01e28 (⌂|🏘|1)
+    │       │       └── ►:3[2]:anon: →:6:
+    │       │           ├── ·081bae9 (⌂|🏘|✓|1111) ►A-inside, ►A-outside
+    │       │           └── ·3183e43 (⌂|🏘|✓|1111)
+    │       └── 📙►:6[1]:A <> origin/A →:4:
+    │           └── →:3:
+    ├── ►:1[0]:origin/main →:2:
+    │   └── ►:2[1]:main <> origin/main →:1:
+    │       └── ·8dc508f (⌂|✓|10)
+    │           └── →:3:
+    └── ►:4[0]:origin/A →:6:
+        └── 🟣197ddce (0x0|1000)
+            └── →:3:
+    ");
+
+    insta::assert_snapshot!(graph_workspace(&graph.to_workspace()?), @r"
+    📕🏘️:0:gitbutler/workspace <> ✓refs/remotes/origin/main⇣1 on 081bae9
+    ├── ≡📙:6:A <> origin/A →:4:⇣1 on 081bae9 {0}
+    │   └── 📙:6:A <> origin/A →:4:⇣1
+    │       └── 🟣197ddce
+    └── ≡:5:B on 081bae9
+        └── :5:B
+            └── ·3e01e28 (🏘️)
+    ");
+    Ok(())
+}
+
 mod edit_commit {
     use but_graph::Graph;
     use but_testsupport::{graph_tree, graph_workspace, visualize_commit_graph_all};