1use std::collections::{BTreeMap, BTreeSet};
4
5use anyhow::{Result, bail};
6
7use super::{
8 branch_and_descendants, children_map, children_of, parent_map, parent_of, root_for,
9 trunk_branch,
10};
11use crate::git;
12use crate::prompt;
13use crate::style;
14
15fn pick_child(title: &str, children: &[String]) -> anyhow::Result<Option<String>> {
18 let painted: Vec<String> = children
19 .iter()
20 .map(|child| style::paint(style::BRANCH, child))
21 .collect();
22 Ok(prompt::pick(title, &painted)?.map(|index| children[index].clone()))
23}
24
25pub fn print_parent(branch: Option<&str>) -> Result<()> {
26 let branch = branch
27 .map(str::to_owned)
28 .map_or_else(git::current_branch, Ok)?;
29 match parent_of(&branch)? {
30 Some(parent) => println!("{parent}"),
31 None => bail!("{branch} has no stack parent"),
32 }
33 Ok(())
34}
35
36pub fn print_children(branch: Option<&str>) -> Result<()> {
37 let branch = branch
38 .map(str::to_owned)
39 .map_or_else(git::current_branch, Ok)?;
40 for child in children_of(&branch)? {
41 println!("{child}");
42 }
43 Ok(())
44}
45
46pub fn checkout_parent() -> Result<()> {
47 let current = git::current_branch()?;
48 let Some(parent) = parent_of(¤t)? else {
49 bail!("{current} has no stack parent");
50 };
51
52 git::checkout(&parent)
53}
54
55pub fn checkout_child(branch: Option<&str>) -> Result<()> {
56 let current = git::current_branch()?;
57 let children = children_of(¤t)?;
58 let child = match (branch, children.as_slice()) {
59 (Some(branch), _) => {
60 if children.iter().any(|child| child == branch) {
61 branch.to_owned()
62 } else {
63 bail!("{branch} is not a stack child of {current}");
64 }
65 }
66 (None, [child]) => child.to_owned(),
67 (None, []) => bail!("{current} has no stack children"),
68 (None, _) => {
69 match pick_child(
70 &format!("{current} has multiple stack children:"),
71 &children,
72 )? {
73 Some(child) => child,
74 None => bail!("choose one with `git stk up <branch>`"),
75 }
76 }
77 };
78
79 git::checkout(&child)
80}
81
82pub fn checkout_top() -> Result<()> {
85 let current = git::current_branch()?;
86 let mut top = current.clone();
87 loop {
88 let children = children_of(&top)?;
89 match children.as_slice() {
90 [] => break,
91 [child] => top = child.clone(),
92 _ => match pick_child(&format!("{top} has multiple stack children:"), &children)? {
94 Some(child) => top = child,
95 None => bail!("walk up from {top} with `git stk up <branch>`"),
96 },
97 }
98 }
99
100 if top == current {
101 if children_of(¤t)?.is_empty() && parent_of(¤t)?.is_none() {
102 bail!("{current} is not in a stack");
103 }
104 println!("{current} is already at the top of the stack");
105 return Ok(());
106 }
107 git::checkout(&top)
108}
109
110pub fn checkout_bottom() -> Result<()> {
113 let current = git::current_branch()?;
114 let trunk = trunk_branch(&git::local_branches()?);
115
116 let bottom = if Some(¤t) == trunk.as_ref() {
117 let children = children_of(¤t)?;
118 match children.as_slice() {
119 [child] => child.clone(),
120 [] => bail!("{current} has no stacked branches"),
121 _ => {
122 match pick_child(
123 &format!("{current} has multiple stack children:"),
124 &children,
125 )? {
126 Some(child) => child,
127 None => bail!("choose one with `git stk up <branch>`"),
128 }
129 }
130 }
131 } else {
132 let mut bottom = current.clone();
133 while let Some(parent) = parent_of(&bottom)? {
134 if Some(&parent) == trunk.as_ref() {
135 break;
136 }
137 bottom = parent;
138 }
139 bottom
140 };
141
142 if bottom == current {
143 if parent_of(¤t)?.is_none() && children_of(¤t)?.is_empty() {
144 bail!("{current} is not in a stack");
145 }
146 println!("{current} is already at the bottom of the stack");
147 return Ok(());
148 }
149 git::checkout(&bottom)
150}
151
152pub fn print_stack() -> Result<()> {
153 let current = git::current_branch()?;
154 let parents = parent_map()?;
155 let root = root_for(¤t, &parents);
156 let children = children_map(&parents);
157 let trunk = trunk_branch(&git::local_branches()?);
158
159 let mut lines = Vec::new();
160 collect_tree_lines(
161 &root,
162 ¤t,
163 trunk.as_deref(),
164 &children,
165 0,
166 &mut BTreeSet::new(),
167 &mut lines,
168 );
169
170 for line in lines.iter().rev() {
173 anstream::println!("{line}");
174 }
175
176 for branch in branch_and_descendants(&root)? {
177 if let Some(parent) = parents.get(&branch)
178 && let Some(hint) = behind_parent_hint(&branch, parent)
179 {
180 anstream::println!("{} {hint}", style::paint(style::HINT, "hint:"));
181 }
182 }
183 Ok(())
184}
185
186pub fn behind_parent_hint(branch: &str, parent: &str) -> Option<String> {
189 let behind = git::commits_behind(branch, parent)
190 .ok()
191 .filter(|count| *count > 0)?;
192 Some(format!(
193 "{branch} is {behind} commit{} behind {parent} - run `git stk restack`",
194 if behind == 1 { "" } else { "s" }
195 ))
196}
197
198#[allow(clippy::too_many_arguments)]
199fn collect_tree_lines(
200 branch: &str,
201 current: &str,
202 trunk: Option<&str>,
203 children: &BTreeMap<String, Vec<String>>,
204 depth: usize,
205 seen: &mut BTreeSet<String>,
206 lines: &mut Vec<String>,
207) {
208 let mut line = " ".repeat(depth);
210 if branch == current {
211 line.push_str(&style::paint(style::CURRENT, &format!("\u{25c9} {branch}")));
212 } else {
213 line.push_str("\u{25cb} ");
214 line.push_str(&style::paint(style::BRANCH, branch));
215 }
216 if Some(branch) == trunk {
217 line.push_str(&style::paint(style::DIM, " (trunk)"));
218 }
219 lines.push(line);
220
221 if !seen.insert(branch.to_owned()) {
222 lines.push(format!("{}<cycle detected>", " ".repeat(depth + 1)));
223 return;
224 }
225
226 if let Some(branch_children) = children.get(branch) {
227 for child in branch_children {
228 collect_tree_lines(child, current, trunk, children, depth + 1, seen, lines);
229 }
230 }
231}