git_branches/
branch.rs

1use grit::git::{self, git};
2use std::collections::{HashMap, HashSet};
3
4/// A node in the topology tree. Every node is a commit; branches add metadata.
5#[derive(Debug)]
6pub struct Node {
7    pub commit: String,           // commit SHA (short, 7 chars)
8    pub branches: Vec<BranchInfo>, // branches pointing to this commit (may be empty)
9    pub children: Vec<Node>,
10}
11
12/// Branch metadata attached to a node.
13#[derive(Debug)]
14pub struct BranchInfo {
15    pub name: String,
16    pub ahead: usize,    // commits from parent node to this tip
17    pub pr: Option<u32>, // PR number if branch has an open PR
18}
19
20/// The topology of all local branches.
21#[derive(Debug)]
22pub struct Topology {
23    pub trunk: String, // trunk branch name
24    pub root: Node,    // root node (common ancestor of all branches)
25}
26
27/// Fetch open PRs and return a map from branch name to PR number.
28async fn fetch_prs() -> HashMap<String, u32> {
29    let output = tokio::process::Command::new("gh")
30        .args(["pr", "list", "--json", "headRefName,number"])
31        .output()
32        .await;
33    let Ok(output) = output else {
34        return HashMap::new();
35    };
36    if !output.status.success() {
37        return HashMap::new();
38    }
39    let json = String::from_utf8_lossy(&output.stdout);
40
41    // Parse JSON: [{"headRefName": "branch", "number": 123}, ...]
42    let mut prs = HashMap::new();
43    // Simple JSON parsing without adding a dependency.
44    for line in json.split('{') {
45        let Some(name_start) = line.find("\"headRefName\":\"") else {
46            continue;
47        };
48        let rest = &line[name_start + 15..];
49        let Some(name_end) = rest.find('"') else {
50            continue;
51        };
52        let name = &rest[..name_end];
53
54        let Some(num_start) = line.find("\"number\":") else {
55            continue;
56        };
57        let rest = &line[num_start + 9..];
58        let num_str: String = rest.chars().take_while(|c| c.is_ascii_digit()).collect();
59        if let Ok(num) = num_str.parse() {
60            prs.insert(name.to_owned(), num);
61        }
62    }
63    prs
64}
65
66/// Collect all local branches and build the topology tree with merge-base commits
67/// as first-class nodes.
68pub async fn collect(trunk: &str) -> Result<Topology, git::Error> {
69    // Fetch open PRs.
70    let prs = fetch_prs().await;
71
72    // Get HEAD commit.
73    let head_commit = git(["rev-parse", "--short=7", "HEAD"])
74        .await
75        .map(|s| s.trim().to_owned())
76        .ok();
77
78    // Phase 1: Collect all branch tips (including trunk).
79    let output = git(["for-each-ref", "--format=%(refname:short)", "refs/heads/"]).await?;
80    let all_branches: Vec<&str> = output.lines().collect();
81
82    let mut tips: HashMap<&str, String> = HashMap::new();
83    for name in &all_branches {
84        let tip = git(["rev-parse", "--short=7", name]).await?.trim().to_owned();
85        tips.insert(name, tip);
86    }
87
88    // Handle empty repo or trunk-only case.
89    if all_branches.len() <= 1 {
90        let trunk_tip = tips.get(trunk).cloned().unwrap_or_default();
91        let mut branches = vec![BranchInfo {
92            name: trunk.to_owned(),
93            ahead: 0,
94            pr: prs.get(trunk).copied(),
95        }];
96        if head_commit.as_deref() == Some(trunk_tip.as_str()) {
97            branches.push(BranchInfo {
98                name: "[HEAD]".to_owned(),
99                ahead: 0,
100                pr: None,
101            });
102        }
103        return Ok(Topology {
104            trunk: trunk.to_owned(),
105            root: Node {
106                commit: trunk_tip.clone(),
107                branches,
108                children: vec![],
109            },
110        });
111    }
112
113    // Phase 2: Find all merge-bases (fork points).
114    let mut commits: HashSet<String> = HashSet::new();
115
116    // Add all branch tips.
117    for tip in tips.values() {
118        commits.insert(tip.clone());
119    }
120
121    // Add HEAD commit so it appears even when detached.
122    if let Some(ref head) = head_commit {
123        commits.insert(head.clone());
124    }
125
126    // Compute merge-base of each non-trunk branch with trunk.
127    let mut trunk_merge_bases: HashMap<&str, String> = HashMap::new();
128    for name in &all_branches {
129        if *name == trunk {
130            continue;
131        }
132        let mb = git(["merge-base", trunk, name]).await?;
133        let mb_short = git(["rev-parse", "--short=7", mb.trim()]).await?.trim().to_owned();
134        trunk_merge_bases.insert(name, mb_short.clone());
135        commits.insert(mb_short);
136    }
137
138    // Group branches by their trunk merge-base to find those that need pairwise comparison.
139    let mut by_merge_base: HashMap<&str, Vec<&str>> = HashMap::new();
140    for (branch, mb) in &trunk_merge_bases {
141        by_merge_base.entry(mb.as_str()).or_default().push(branch);
142    }
143
144    // For branches sharing a trunk merge-base, compute pairwise merge-bases.
145    for (_mb, branches) in &by_merge_base {
146        if branches.len() > 1 {
147            for i in 0..branches.len() {
148                for j in (i + 1)..branches.len() {
149                    let mb = git(["merge-base", branches[i], branches[j]]).await?;
150                    let mb_short = git(["rev-parse", "--short=7", mb.trim()]).await?.trim().to_owned();
151                    commits.insert(mb_short);
152                }
153            }
154        }
155    }
156
157    // Phase 3: Build parent relationships.
158    // For each commit, find its parent (closest ancestor in our commit set).
159    let commits_vec: Vec<String> = commits.into_iter().collect();
160    let mut parents: HashMap<String, Option<String>> = HashMap::new();
161
162    for commit in &commits_vec {
163        let mut best_parent: Option<String> = None;
164        let mut best_distance = usize::MAX;
165
166        for other in &commits_vec {
167            if other == commit {
168                continue;
169            }
170
171            // Is other an ancestor of commit?
172            let is_ancestor = git(["merge-base", "--is-ancestor", other, commit]).await;
173            if is_ancestor.is_ok() {
174                let dist = count_commits(other, commit).await?;
175                if dist < best_distance && dist > 0 {
176                    best_distance = dist;
177                    best_parent = Some(other.clone());
178                }
179            }
180        }
181
182        parents.insert(commit.clone(), best_parent);
183    }
184
185    // Phase 4: Build tree from root.
186    // Find root (commit with no parent in our set).
187    let root_commit = commits_vec
188        .iter()
189        .find(|c| parents.get(*c) == Some(&None))
190        .cloned()
191        .unwrap_or_else(|| tips.get(trunk).cloned().unwrap_or_default());
192
193    // Build a map from commit to branch names (multiple branches may point to same commit).
194    let mut commit_to_branches: HashMap<String, Vec<String>> = HashMap::new();
195    for (name, tip) in &tips {
196        commit_to_branches
197            .entry(tip.clone())
198            .or_default()
199            .push((*name).to_owned());
200    }
201
202    // Add [HEAD] as a pseudo-branch at its commit.
203    if let Some(ref head) = head_commit {
204        commit_to_branches
205            .entry(head.clone())
206            .or_default()
207            .push("[HEAD]".to_owned());
208    }
209
210    let root = build_node(&root_commit, &commits_vec, &parents, &commit_to_branches, &prs).await?;
211
212    Ok(Topology {
213        trunk: trunk.to_owned(),
214        root,
215    })
216}
217
218/// Recursively build a node and its children.
219async fn build_node(
220    commit: &str,
221    all_commits: &[String],
222    parents: &HashMap<String, Option<String>>,
223    commit_to_branches: &HashMap<String, Vec<String>>,
224    prs: &HashMap<String, u32>,
225) -> Result<Node, git::Error> {
226    // Find children (commits whose parent is this commit).
227    let mut child_commits: Vec<&String> = all_commits
228        .iter()
229        .filter(|c| parents.get(*c) == Some(&Some(commit.to_owned())))
230        .collect();
231
232    // Sort children for deterministic output.
233    child_commits.sort();
234
235    // Build children recursively.
236    let mut children = Vec::new();
237    for child in child_commits {
238        let child_node =
239            Box::pin(build_node(child, all_commits, parents, commit_to_branches, prs)).await?;
240        children.push(child_node);
241    }
242
243    // Calculate ahead count (distance from parent).
244    let ahead = if let Some(Some(parent)) = parents.get(commit) {
245        count_commits(parent, commit).await?
246    } else {
247        0
248    };
249
250    // Get all branches pointing to this commit. Branches sorted alphabetically,
251    // with [HEAD] always last.
252    let branches = commit_to_branches
253        .get(commit)
254        .map(|names| {
255            let mut infos: Vec<BranchInfo> = names
256                .iter()
257                .filter(|name| *name != "[HEAD]")
258                .map(|name| BranchInfo {
259                    name: name.clone(),
260                    ahead,
261                    pr: prs.get(name).copied(),
262                })
263                .collect();
264            infos.sort_by(|a, b| a.name.cmp(&b.name));
265            // Append [HEAD] at the end if present.
266            if names.iter().any(|n| n == "[HEAD]") {
267                infos.push(BranchInfo {
268                    name: "[HEAD]".to_owned(),
269                    ahead,
270                    pr: None,
271                });
272            }
273            infos
274        })
275        .unwrap_or_default();
276
277    Ok(Node {
278        commit: commit.to_owned(),
279        branches,
280        children,
281    })
282}
283
284/// Count commits in range `from..to`.
285async fn count_commits(from: &str, to: &str) -> Result<usize, git::Error> {
286    let output = git(["rev-list", "--count", &format!("{from}..{to}")]).await?;
287    Ok(output.trim().parse().unwrap_or(0))
288}