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(char::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/// Find all relevant commits: branch tips, HEAD, and merge-base fork points.
67async fn find_commits<'a>(
68    all_branches: &[&'a str],
69    trunk: &str,
70    tips: &HashMap<&str, String>,
71    head_commit: Option<&String>,
72) -> Result<HashSet<String>, git::Error> {
73    let mut commits: HashSet<String> = tips.values().cloned().collect();
74
75    if let Some(head) = head_commit {
76        commits.insert(head.clone());
77    }
78
79    // Compute merge-base of each non-trunk branch with trunk.
80    let mut trunk_merge_bases: HashMap<&'a str, String> = HashMap::new();
81    for name in all_branches {
82        if *name == trunk {
83            continue;
84        }
85        let mb = git(["merge-base", trunk, name]).await?;
86        let mb_short = git(["rev-parse", "--short=7", mb.trim()])
87            .await?
88            .trim()
89            .to_owned();
90        trunk_merge_bases.insert(name, mb_short.clone());
91        commits.insert(mb_short);
92    }
93
94    // Group branches by their trunk merge-base to find those that need pairwise comparison.
95    let mut by_merge_base: HashMap<&str, Vec<&str>> = HashMap::new();
96    for (branch, mb) in &trunk_merge_bases {
97        by_merge_base.entry(mb.as_str()).or_default().push(branch);
98    }
99
100    // For branches sharing a trunk merge-base, compute pairwise merge-bases.
101    for branches in by_merge_base.values() {
102        if branches.len() > 1 {
103            for i in 0..branches.len() {
104                for j in (i + 1)..branches.len() {
105                    let mb = git(["merge-base", branches[i], branches[j]]).await?;
106                    let mb_short = git(["rev-parse", "--short=7", mb.trim()])
107                        .await?
108                        .trim()
109                        .to_owned();
110                    commits.insert(mb_short);
111                }
112            }
113        }
114    }
115
116    Ok(commits)
117}
118
119/// For each commit, find its parent (closest ancestor in the commit set).
120async fn find_parents(
121    commits: &[String],
122) -> Result<HashMap<String, Option<String>>, git::Error> {
123    let mut parents: HashMap<String, Option<String>> = HashMap::new();
124
125    for commit in commits {
126        let mut best_parent: Option<String> = None;
127        let mut best_distance = usize::MAX;
128
129        for other in commits {
130            if other == commit {
131                continue;
132            }
133
134            let is_ancestor = git(["merge-base", "--is-ancestor", other, commit]).await;
135            if is_ancestor.is_ok() {
136                let dist = count_commits(other, commit).await?;
137                if dist < best_distance && dist > 0 {
138                    best_distance = dist;
139                    best_parent = Some(other.clone());
140                }
141            }
142        }
143
144        parents.insert(commit.clone(), best_parent);
145    }
146
147    Ok(parents)
148}
149
150/// Collect all local branches and build the topology tree with merge-base commits
151/// as first-class nodes.
152pub async fn collect(trunk: &str) -> Result<Topology, git::Error> {
153    let prs = fetch_prs().await;
154
155    let head_commit = git(["rev-parse", "--short=7", "HEAD"])
156        .await
157        .map(|s| s.trim().to_owned())
158        .ok();
159
160    let output = git(["for-each-ref", "--format=%(refname:short)", "refs/heads/"]).await?;
161    let all_branches: Vec<&str> = output.lines().collect();
162
163    let mut tips: HashMap<&str, String> = HashMap::new();
164    for name in &all_branches {
165        let tip = git(["rev-parse", "--short=7", name])
166            .await?
167            .trim()
168            .to_owned();
169        tips.insert(name, tip);
170    }
171
172    if all_branches.len() <= 1 {
173        let trunk_tip = tips.get(trunk).cloned().unwrap_or_default();
174        let mut branches = vec![BranchInfo {
175            name: trunk.to_owned(),
176            ahead: 0,
177            pr: prs.get(trunk).copied(),
178        }];
179        if head_commit.as_deref() == Some(trunk_tip.as_str()) {
180            branches.push(BranchInfo {
181                name: "[HEAD]".to_owned(),
182                ahead: 0,
183                pr: None,
184            });
185        }
186        return Ok(Topology {
187            trunk: trunk.to_owned(),
188            root: Node {
189                commit: trunk_tip.clone(),
190                branches,
191                children: vec![],
192            },
193        });
194    }
195
196    let commits: HashSet<String> =
197        find_commits(&all_branches, trunk, &tips, head_commit.as_ref()).await?;
198    let commits_vec: Vec<String> = commits.into_iter().collect();
199    let parents = find_parents(&commits_vec).await?;
200
201    let root_commit = commits_vec
202        .iter()
203        .find(|c| parents.get(*c) == Some(&None))
204        .cloned()
205        .unwrap_or_else(|| tips.get(trunk).cloned().unwrap_or_default());
206
207    let mut commit_to_branches: HashMap<String, Vec<String>> = HashMap::new();
208    for (name, tip) in &tips {
209        commit_to_branches
210            .entry(tip.clone())
211            .or_default()
212            .push((*name).to_owned());
213    }
214    if let Some(ref head) = head_commit {
215        commit_to_branches
216            .entry(head.clone())
217            .or_default()
218            .push("[HEAD]".to_owned());
219    }
220
221    let root = build_node(
222        &root_commit,
223        &commits_vec,
224        &parents,
225        &commit_to_branches,
226        &prs,
227    )
228    .await?;
229
230    Ok(Topology {
231        trunk: trunk.to_owned(),
232        root,
233    })
234}
235
236/// Recursively build a node and its children.
237async fn build_node(
238    commit: &str,
239    all_commits: &[String],
240    parents: &HashMap<String, Option<String>>,
241    commit_to_branches: &HashMap<String, Vec<String>>,
242    prs: &HashMap<String, u32>,
243) -> Result<Node, git::Error> {
244    // Find children (commits whose parent is this commit).
245    let mut child_commits: Vec<&String> = all_commits
246        .iter()
247        .filter(|c| parents.get(*c) == Some(&Some(commit.to_owned())))
248        .collect();
249
250    // Sort children for deterministic output.
251    child_commits.sort();
252
253    // Build children recursively.
254    let mut children = Vec::new();
255    for child in child_commits {
256        let child_node = Box::pin(build_node(
257            child,
258            all_commits,
259            parents,
260            commit_to_branches,
261            prs,
262        ))
263        .await?;
264        children.push(child_node);
265    }
266
267    // Calculate ahead count (distance from parent).
268    let ahead = if let Some(Some(parent)) = parents.get(commit) {
269        count_commits(parent, commit).await?
270    } else {
271        0
272    };
273
274    // Get all branches pointing to this commit. Branches sorted alphabetically,
275    // with [HEAD] always last.
276    let branches = commit_to_branches
277        .get(commit)
278        .map(|names| {
279            let mut infos: Vec<BranchInfo> = names
280                .iter()
281                .filter(|name| *name != "[HEAD]")
282                .map(|name| BranchInfo {
283                    name: name.clone(),
284                    ahead,
285                    pr: prs.get(name).copied(),
286                })
287                .collect();
288            infos.sort_by(|a, b| a.name.cmp(&b.name));
289            // Append [HEAD] at the end if present.
290            if names.iter().any(|n| n == "[HEAD]") {
291                infos.push(BranchInfo {
292                    name: "[HEAD]".to_owned(),
293                    ahead,
294                    pr: None,
295                });
296            }
297            infos
298        })
299        .unwrap_or_default();
300
301    Ok(Node {
302        commit: commit.to_owned(),
303        branches,
304        children,
305    })
306}
307
308/// Count commits in range `from..to`.
309async fn count_commits(from: &str, to: &str) -> Result<usize, git::Error> {
310    let output = git(["rev-list", "--count", &format!("{from}..{to}")]).await?;
311    Ok(output.trim().parse().unwrap_or(0))
312}