1use std::collections::{BTreeMap, BTreeSet};
6
7use anyhow::{Context, Result, bail};
8use serde_json::{Value, json};
9
10use crate::git;
11use crate::settings;
12use crate::style;
13
14const METADATA_REF: &str = "refs/stk/metadata";
17const METADATA_FILE: &str = "stack.json";
18
19mod nav;
20mod restack;
21mod snapshot;
22
23pub use nav::{
24 NavOutput, behind_parent_hint, checkout_bottom, checkout_child, checkout_parent, checkout_top,
25 print_all_stacks, print_children, print_parent, print_stack,
26};
27pub use restack::{abort_restack, continue_restack, restack};
28pub use snapshot::{take as snapshot, undo};
29
30const PARENT_KEY: &str = "stkParent";
31const BASE_KEY: &str = "stkBase";
32const RENAMED_FROM_KEY: &str = "stkRenamedFrom";
35const WORKTREE_KEY: &str = "stkWorktree";
38
39pub fn create_branch(branch: &str, dry_run: bool) -> Result<()> {
40 let parent = git::current_branch()?;
41 if git::local_branches()?
43 .iter()
44 .any(|existing| existing == branch)
45 {
46 bail!(
47 "branch {branch} already exists - adopt it onto {parent} \
48 with `git stk adopt {branch} --parent {parent}`"
49 );
50 }
51 if !dry_run {
52 git::create_branch(branch)?;
53 set_parent(branch, &parent)?;
54 record_base(branch, &parent);
55 }
56 anstream::println!(
57 "{} {} with parent {}",
58 if dry_run { "would create" } else { "created" },
59 style::branch(branch),
60 style::branch(&parent)
61 );
62 Ok(())
63}
64
65pub fn create_branch_in_worktree(branch: &str, dry_run: bool) -> Result<()> {
72 let parent = git::current_branch()?;
73 ensure_absent(branch)?;
74
75 let path = settings::worktree_path_for(branch)?;
76 if path.exists() {
77 bail!(
78 "{} already exists; remove it or pick another branch name",
79 path.display()
80 );
81 }
82
83 if !dry_run {
84 git::worktree_add_new_branch(&path, branch, &parent)?;
85 set_owned_worktree(branch, &path)?;
89 set_parent(branch, &parent)?;
90 record_base(branch, &parent);
91 }
92
93 anstream::println!(
94 "{} {} with parent {} in the worktree at {}",
95 if dry_run { "would create" } else { "created" },
96 style::branch(branch),
97 style::branch(&parent),
98 git::display_path(&path)
99 );
100 if !dry_run {
101 anstream::println!(
102 "{}",
103 style::dim(&format!("cd {}", git::display_path(&path)))
104 );
105 }
106 Ok(())
107}
108
109pub fn trunk_held_elsewhere(trunk: &str) -> Result<bool> {
118 let Some(path) = git::worktree_holding(trunk)? else {
119 return Ok(false);
120 };
121 anstream::println!(
122 "{}",
123 style::warn(&format!(
124 "skipped fetching {trunk}: it is checked out in the worktree at {}",
125 git::display_path(&path)
126 ))
127 );
128 anstream::println!(
129 "{}",
130 style::dim(&format!(
131 "using the local {trunk}; fast-forward it there to pick up the remote"
132 ))
133 );
134 Ok(true)
135}
136
137pub fn owned_worktree(branch: &str) -> Option<std::path::PathBuf> {
140 recorded_worktree(branch).filter(|path| path.exists())
141}
142
143pub fn recorded_worktree(branch: &str) -> Option<std::path::PathBuf> {
146 git::config_get(&format!("branch.{branch}.{WORKTREE_KEY}"))
147 .ok()
148 .flatten()
149 .map(std::path::PathBuf::from)
150}
151
152pub fn set_owned_worktree(branch: &str, path: &std::path::Path) -> Result<()> {
154 git::config_set(
155 &format!("branch.{branch}.{WORKTREE_KEY}"),
156 &path.to_string_lossy(),
157 )
158}
159
160pub fn unset_owned_worktree(branch: &str) -> Result<()> {
162 git::config_unset(&format!("branch.{branch}.{WORKTREE_KEY}"))
163}
164
165pub fn insert_branch(branch: &str, dry_run: bool) -> Result<()> {
170 ensure_absent(branch)?;
171 let current = git::current_branch()?;
172 let children = children_of(¤t)?;
173
174 if !dry_run {
175 snapshot::take("new --insert");
176 git::create_branch(branch)?; set_parent(branch, ¤t)?;
178 record_base(branch, ¤t);
179 for child in &children {
180 set_parent(child, branch)?;
181 record_base(child, branch);
182 }
183 }
184
185 anstream::println!(
186 "{} {} above {}",
187 if dry_run { "would insert" } else { "inserted" },
188 style::branch(branch),
189 style::branch(¤t)
190 );
191 for child in &children {
192 anstream::println!(
193 "{} {} -> {}",
194 if dry_run {
195 "would retarget"
196 } else {
197 "retargeted"
198 },
199 style::branch(child),
200 style::branch(branch)
201 );
202 }
203 Ok(())
204}
205
206pub fn prepend_branch(branch: &str, dry_run: bool) -> Result<()> {
210 ensure_absent(branch)?;
211 let current = git::current_branch()?;
212 let parent =
213 parent_of(¤t)?.context("current branch has no stack parent to prepend below")?;
214 if !git::worktree_is_clean()? {
215 bail!(
216 "working tree has uncommitted changes; commit or stash before `git stk new --prepend`"
217 );
218 }
219
220 if !dry_run {
221 snapshot::take("new --prepend");
222 git::checkout(&parent)?;
223 git::create_branch(branch)?; set_parent(branch, &parent)?;
225 record_base(branch, &parent);
226 set_parent(¤t, branch)?;
227 record_base(¤t, branch);
228 }
229
230 anstream::println!(
231 "{} {} between {} and {}",
232 if dry_run { "would insert" } else { "inserted" },
233 style::branch(branch),
234 style::branch(&parent),
235 style::branch(¤t)
236 );
237 anstream::println!(
238 "{} {} -> {}",
239 if dry_run {
240 "would retarget"
241 } else {
242 "retargeted"
243 },
244 style::branch(¤t),
245 style::branch(branch)
246 );
247 Ok(())
248}
249
250fn ensure_absent(branch: &str) -> Result<()> {
251 if git::local_branches()?
252 .iter()
253 .any(|existing| existing == branch)
254 {
255 bail!("branch {branch} already exists");
256 }
257 Ok(())
258}
259
260pub fn trunk_branch(branches: &[String]) -> Option<String> {
263 let remote = settings::remote().unwrap_or_else(|_| settings::DEFAULT_REMOTE.to_owned());
264 if let Some(default) = git::remote_default_branch(&remote) {
265 return Some(default);
266 }
267
268 ["main", "master"]
269 .iter()
270 .find(|name| branches.iter().any(|branch| branch == *name))
271 .map(|name| (*name).to_owned())
272}
273
274pub fn adopt_branch(branch: &str, parent: &str, dry_run: bool) -> Result<()> {
275 if branch == parent {
276 bail!("a branch cannot be its own stack parent");
277 }
278
279 let branches: BTreeSet<_> = git::local_branches()?.into_iter().collect();
280 if !branches.contains(branch) {
281 bail!("branch {branch} does not exist");
282 }
283 if !branches.contains(parent) {
284 bail!("parent branch {parent} does not exist");
285 }
286 if branch_and_descendants(branch)?
287 .iter()
288 .any(|descendant| descendant == parent)
289 {
290 bail!("{parent} is already below {branch} in the stack; that would form a cycle");
291 }
292
293 if !dry_run {
294 set_parent(branch, parent)?;
295 record_base(branch, parent);
296 }
297 anstream::println!(
298 "{} {} to {}",
299 if dry_run { "would attach" } else { "attached" },
300 style::branch(branch),
301 style::branch(parent)
302 );
303 Ok(())
304}
305
306pub fn detach_branch(branch: Option<&str>) -> Result<()> {
307 let branch = branch
308 .map(str::to_owned)
309 .map_or_else(git::current_branch, Ok)?;
310 unset_parent(&branch)?;
311 unset_base(&branch)?;
312 anstream::println!("detached {}", style::branch(&branch));
313 Ok(())
314}
315
316pub fn rename_branch(old: &str, new: &str, dry_run: bool) -> Result<()> {
320 let children = children_of(old)?;
321
322 if !dry_run {
323 snapshot::take("rename");
324 git::rename_branch(old, new)?;
325 }
326 anstream::println!(
327 "{} {} -> {}",
328 if dry_run { "would rename" } else { "renamed" },
329 style::branch(old),
330 style::branch(new)
331 );
332
333 for child in &children {
334 if !dry_run {
335 set_parent(child, new)?;
336 }
337 anstream::println!(
338 "{} {} -> {}",
339 if dry_run {
340 "would retarget"
341 } else {
342 "retargeted"
343 },
344 style::branch(child),
345 style::branch(new)
346 );
347 }
348 Ok(())
349}
350
351pub fn set_renamed_from(branch: &str, old: &str) -> Result<()> {
354 git::config_set(&renamed_from_key(branch), old)
355}
356
357pub fn renamed_from(branch: &str) -> Result<Option<String>> {
359 git::config_get(&renamed_from_key(branch))
360}
361
362pub fn clear_renamed_from(branch: &str) -> Result<()> {
364 git::config_unset(&renamed_from_key(branch))
365}
366
367pub fn record_base(branch: &str, parent: &str) {
370 if let Ok(base) = git::merge_base(parent, branch) {
371 let _ = git::config_set(&base_key(branch), &base);
372 }
373}
374
375pub(crate) fn fork_point(branch: &str, parent: &str) -> Result<Option<String>> {
385 let recorded = base_of(branch)?.filter(|base| git::is_ancestor(base, branch).unwrap_or(false));
386 let merge_base = git::merge_base(parent, branch).ok();
387 Ok(match (recorded, merge_base) {
388 (Some(recorded), Some(merge_base)) => Some(
389 if git::is_ancestor(&merge_base, &recorded).unwrap_or(false) {
390 recorded
391 } else {
392 merge_base
393 },
394 ),
395 (recorded, merge_base) => recorded.or(merge_base),
396 })
397}
398
399pub(crate) fn base_is_current(branch: &str, parent: &str) -> Result<bool> {
403 let Some(base) = base_of(branch)? else {
404 return Ok(false);
405 };
406 if !git::is_ancestor(&base, branch).unwrap_or(false) {
407 return Ok(false);
408 }
409 Ok(match git::merge_base(parent, branch) {
410 Ok(merge_base) => git::is_ancestor(&merge_base, &base).unwrap_or(false),
411 Err(_) => true,
412 })
413}
414
415pub fn stack_root(branch: &str) -> Result<String> {
417 let parents = parent_map()?;
418 Ok(root_for(branch, &parents))
419}
420
421pub fn branch_and_descendants(branch: &str) -> Result<Vec<String>> {
422 let parents = parent_map()?;
423 let children = children_map(&parents);
424 let mut branches = vec![branch.to_owned()];
425 let mut visited = BTreeSet::from([branch.to_owned()]);
426 collect_descendants(branch, &children, &mut branches, &mut visited);
427 Ok(branches)
428}
429
430pub fn stack_line(branch: &str) -> Result<Vec<String>> {
436 let trunk = trunk_branch(&git::local_branches()?);
440 if Some(branch) == trunk.as_deref() {
441 return Ok(Vec::new());
442 }
443
444 let mut line = path_from_root(branch)?; let above = branch_and_descendants(branch)?; line.extend(above.into_iter().skip(1)); line.retain(|candidate| Some(candidate) != trunk.as_ref());
451 Ok(line)
452}
453
454pub(crate) fn line_base(branch: &str) -> Result<String> {
462 Ok(path_from_root(branch)?
463 .into_iter()
464 .next()
465 .unwrap_or_else(|| branch.to_owned()))
466}
467
468pub fn current_stack_branches(branch: &str) -> Result<Vec<String>> {
476 let base = line_base(branch)?;
477 let trunk = trunk_branch(&git::local_branches()?);
478 Ok(branch_and_descendants(&base)?
479 .into_iter()
480 .filter(|candidate| Some(candidate) != trunk.as_ref())
481 .collect())
482}
483
484pub fn listed_branches(all: bool) -> Result<BTreeSet<String>> {
489 if all {
490 Ok(parent_map()?
491 .into_iter()
492 .flat_map(|(child, parent)| [child, parent])
493 .collect())
494 } else {
495 let current = git::current_branch()?;
496 Ok(current_stack_branches(¤t)?.into_iter().collect())
497 }
498}
499
500pub fn publish_metadata(remote: &str) {
504 if let Err(error) = try_publish_metadata(remote) {
505 anstream::eprintln!(
506 "{}",
507 style::warn(&format!("could not publish stack metadata: {error:#}"))
508 );
509 }
510}
511
512fn try_publish_metadata(remote: &str) -> Result<()> {
513 let current = git::current_branch()?;
514 let trunk = trunk_branch(&git::local_branches()?);
515
516 let mut parents = serde_json::Map::new();
517 for branch in current_stack_branches(¤t)? {
518 if let Some(parent) = parent_of(&branch)? {
519 parents.insert(branch, Value::String(parent));
520 }
521 }
522 if parents.is_empty() {
523 return Ok(());
524 }
525
526 let document = json!({ "trunk": trunk, "parents": parents });
527 git::write_blob_ref(METADATA_REF, METADATA_FILE, &document.to_string())?;
528 git::push_ref(remote, METADATA_REF)
529}
530
531pub fn apply_remote_metadata(remote: &str) -> Result<usize> {
534 git::fetch_ref(remote, METADATA_REF)
535 .context("no stack metadata on the remote - push it from the other machine first")?;
536 let Some(content) = git::read_ref_file(METADATA_REF, METADATA_FILE)? else {
537 bail!("the remote stack metadata is empty");
538 };
539
540 let document: Value =
541 serde_json::from_str(&content).context("failed to parse remote stack metadata")?;
542 let parents = document
543 .get("parents")
544 .and_then(Value::as_object)
545 .context("remote stack metadata is malformed")?;
546
547 let mut pairs = Vec::new();
552 for (branch, parent) in parents {
553 let Some(parent) = parent.as_str() else {
554 continue;
555 };
556 if !is_safe_ref_name(branch) || !is_safe_ref_name(parent) {
557 anstream::eprintln!(
558 "{}",
559 style::warn(&format!(
560 "skipping unsafe stack metadata entry: {branch:?} -> {parent:?}"
561 ))
562 );
563 continue;
564 }
565 pairs.push((branch.clone(), parent.to_owned()));
566 }
567
568 let local: BTreeSet<String> = git::local_branches()?.into_iter().collect();
571 for (branch, _) in &pairs {
572 if !local.contains(branch) {
573 git::fetch_branch(remote, branch)
574 .with_context(|| format!("failed to fetch {branch} from {remote}"))?;
575 }
576 }
577
578 let mut attached = 0;
579 for (branch, parent) in &pairs {
580 set_parent(branch, parent)?;
581 record_base(branch, parent);
582 attached += 1;
583 anstream::println!(
584 "attached {} to {}",
585 style::branch(branch),
586 style::branch(parent)
587 );
588 }
589 Ok(attached)
590}
591
592pub(crate) fn is_safe_ref_name(name: &str) -> bool {
597 !name.is_empty()
598 && !name.starts_with('-')
599 && !name.chars().any(|c| c.is_whitespace() || c.is_control())
600}
601
602pub fn path_from_root(branch: &str) -> Result<Vec<String>> {
605 let trunk = trunk_branch(&git::local_branches()?);
606 let mut path = vec![branch.to_owned()];
607 let mut seen = BTreeSet::from([branch.to_owned()]);
608
609 let mut cursor = branch.to_owned();
610 while let Some(parent) = parent_of(&cursor)? {
611 if Some(&parent) == trunk.as_ref() || !seen.insert(parent.clone()) {
612 break;
613 }
614 path.push(parent.clone());
615 cursor = parent;
616 }
617
618 path.reverse();
619 Ok(path)
620}
621
622pub fn branch_parents(branches: &[String]) -> Result<Vec<(String, String)>> {
625 let mut pairs = Vec::new();
626 for branch in branches {
627 if let Some(parent) = parent_of(branch)? {
628 pairs.push((branch.clone(), parent));
629 }
630 }
631 Ok(pairs)
632}
633
634fn parent_map() -> Result<BTreeMap<String, String>> {
635 let mut parents = BTreeMap::new();
636 for branch in git::local_branches()? {
637 if let Some(parent) = parent_of(&branch)? {
638 parents.insert(branch, parent);
639 }
640 }
641 Ok(parents)
642}
643
644fn collect_descendants(
645 branch: &str,
646 children: &BTreeMap<String, Vec<String>>,
647 branches: &mut Vec<String>,
648 visited: &mut BTreeSet<String>,
649) {
650 if let Some(branch_children) = children.get(branch) {
651 for child in branch_children {
652 if !visited.insert(child.to_owned()) {
653 continue; }
655 branches.push(child.to_owned());
656 collect_descendants(child, children, branches, visited);
657 }
658 }
659}
660
661pub(crate) fn children_of(parent: &str) -> Result<Vec<String>> {
662 Ok(parent_map()?
663 .into_iter()
664 .filter_map(|(branch, branch_parent)| (branch_parent == parent).then_some(branch))
665 .collect())
666}
667
668fn children_map(parents: &BTreeMap<String, String>) -> BTreeMap<String, Vec<String>> {
669 let mut children: BTreeMap<String, Vec<String>> = BTreeMap::new();
670 for (branch, parent) in parents {
671 children
672 .entry(parent.to_owned())
673 .or_default()
674 .push(branch.to_owned());
675 }
676 children
677}
678
679fn root_for(branch: &str, parents: &BTreeMap<String, String>) -> String {
680 let mut root = branch.to_owned();
681 let mut seen = BTreeSet::new();
682
683 while let Some(parent) = parents.get(&root) {
684 if !seen.insert(root.clone()) {
685 break;
686 }
687 root = parent.to_owned();
688 }
689
690 root
691}
692
693pub(crate) fn parent_of(branch: &str) -> Result<Option<String>> {
694 git::config_get(&parent_key(branch))
695}
696
697pub(crate) fn base_of(branch: &str) -> Result<Option<String>> {
698 git::config_get(&base_key(branch))
699}
700
701pub(crate) fn set_parent(branch: &str, parent: &str) -> Result<()> {
702 git::config_set(&parent_key(branch), parent)
703}
704
705pub(crate) fn unset_parent(branch: &str) -> Result<()> {
706 git::config_unset(&parent_key(branch))
707}
708
709pub(crate) fn set_base(branch: &str, base: &str) -> Result<()> {
710 git::config_set(&base_key(branch), base)
711}
712
713pub(crate) fn unset_base(branch: &str) -> Result<()> {
714 git::config_unset(&base_key(branch))
715}
716
717fn parent_key(branch: &str) -> String {
718 format!("branch.{branch}.{PARENT_KEY}")
719}
720
721fn base_key(branch: &str) -> String {
722 format!("branch.{branch}.{BASE_KEY}")
723}
724
725fn renamed_from_key(branch: &str) -> String {
726 format!("branch.{branch}.{RENAMED_FROM_KEY}")
727}
728
729#[cfg(test)]
730mod tests {
731 use super::is_safe_ref_name;
732
733 #[test]
734 fn safe_ref_names_pass() {
735 assert!(is_safe_ref_name("main"));
736 assert!(is_safe_ref_name("feature/a"));
737 assert!(is_safe_ref_name("user/fix-123"));
738 }
739
740 #[test]
741 fn unsafe_ref_names_are_rejected() {
742 assert!(!is_safe_ref_name("--upload-pack=touch /tmp/pwned"));
744 assert!(!is_safe_ref_name("-x"));
745 assert!(!is_safe_ref_name("a branch"));
747 assert!(!is_safe_ref_name("a\nb"));
748 assert!(!is_safe_ref_name("a\tb"));
749 assert!(!is_safe_ref_name(""));
750 }
751}