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 FLOOR_KEY: &str = "stkFloor";
42const WORKTREE_KEY: &str = "stkWorktree";
45
46pub fn create_branch(branch: &str, dry_run: bool) -> Result<()> {
47 let parent = git::current_branch()?;
48 if git::local_branches()?
50 .iter()
51 .any(|existing| existing == branch)
52 {
53 bail!(
54 "branch {branch} already exists - adopt it onto {parent} \
55 with `git stk adopt {branch} --parent {parent}`"
56 );
57 }
58 if !dry_run {
59 git::create_branch(branch)?;
60 set_parent(branch, &parent)?;
61 record_base(branch, &parent);
62 }
63 anstream::println!(
64 "{} {} with parent {}",
65 if dry_run { "would create" } else { "created" },
66 style::branch(branch),
67 style::branch(&parent)
68 );
69 mark_floor_if_rooting(&parent, dry_run)?;
70 Ok(())
71}
72
73pub fn create_branch_in_worktree(branch: &str, dry_run: bool) -> Result<()> {
80 let parent = git::current_branch()?;
81 ensure_absent(branch)?;
82
83 let path = settings::worktree_path_for(branch)?;
84 if path.exists() {
85 bail!(
86 "{} already exists; remove it or pick another branch name",
87 path.display()
88 );
89 }
90
91 if !dry_run {
92 git::worktree_add_new_branch(&path, branch, &parent)?;
93 set_owned_worktree(branch, &path)?;
97 set_parent(branch, &parent)?;
98 record_base(branch, &parent);
99 }
100 anstream::println!(
101 "{} {} with parent {} in the worktree at {}",
102 if dry_run { "would create" } else { "created" },
103 style::branch(branch),
104 style::branch(&parent),
105 git::display_path(&path)
106 );
107 if !dry_run {
108 anstream::println!(
109 "{}",
110 style::dim(&format!("cd {}", git::display_path(&path)))
111 );
112 }
113 mark_floor_if_rooting(&parent, dry_run)?;
114 Ok(())
115}
116
117pub fn trunk_held_elsewhere(trunk: &str) -> Result<bool> {
126 let Some(path) = git::worktree_holding(trunk)? else {
127 return Ok(false);
128 };
129 anstream::println!(
130 "{}",
131 style::warn(&format!(
132 "skipped fetching {trunk}: it is checked out in the worktree at {}",
133 git::display_path(&path)
134 ))
135 );
136 anstream::println!(
137 "{}",
138 style::dim(&format!(
139 "using the local {trunk}; fast-forward it there to pick up the remote"
140 ))
141 );
142 Ok(true)
143}
144
145pub fn owned_worktree(branch: &str) -> Option<std::path::PathBuf> {
148 recorded_worktree(branch).filter(|path| path.exists())
149}
150
151pub fn recorded_worktree(branch: &str) -> Option<std::path::PathBuf> {
154 git::config_get(&format!("branch.{branch}.{WORKTREE_KEY}"))
155 .ok()
156 .flatten()
157 .map(std::path::PathBuf::from)
158}
159
160pub fn set_owned_worktree(branch: &str, path: &std::path::Path) -> Result<()> {
162 git::config_set(
163 &format!("branch.{branch}.{WORKTREE_KEY}"),
164 &path.to_string_lossy(),
165 )
166}
167
168pub fn unset_owned_worktree(branch: &str) -> Result<()> {
170 git::config_unset(&format!("branch.{branch}.{WORKTREE_KEY}"))
171}
172
173pub fn insert_branch(branch: &str, dry_run: bool) -> Result<()> {
178 ensure_absent(branch)?;
179 let current = git::current_branch()?;
180 let children = children_of(¤t)?;
181
182 if !dry_run {
183 snapshot::take("new --insert");
184 git::create_branch(branch)?; set_parent(branch, ¤t)?;
186 record_base(branch, ¤t);
187 for child in &children {
188 set_parent(child, branch)?;
189 record_base(child, branch);
190 }
191 }
192
193 anstream::println!(
194 "{} {} above {}",
195 if dry_run { "would insert" } else { "inserted" },
196 style::branch(branch),
197 style::branch(¤t)
198 );
199 for child in &children {
200 anstream::println!(
201 "{} {} -> {}",
202 if dry_run {
203 "would retarget"
204 } else {
205 "retargeted"
206 },
207 style::branch(child),
208 style::branch(branch)
209 );
210 }
211 mark_floor_if_rooting(¤t, dry_run)?;
212 Ok(())
213}
214
215pub fn prepend_branch(branch: &str, dry_run: bool) -> Result<()> {
219 ensure_absent(branch)?;
220 let current = git::current_branch()?;
221 let parent = stacked_parent_of(¤t)?
222 .context("current branch has no stack parent to prepend below")?;
223 if !git::worktree_is_clean()? {
224 bail!(
225 "working tree has uncommitted changes; commit or stash before `git stk new --prepend`"
226 );
227 }
228
229 if !dry_run {
230 snapshot::take("new --prepend");
231 git::checkout(&parent)?;
232 git::create_branch(branch)?; set_parent(branch, &parent)?;
234 record_base(branch, &parent);
235 set_parent(¤t, branch)?;
236 record_base(¤t, branch);
237 }
238
239 anstream::println!(
240 "{} {} between {} and {}",
241 if dry_run { "would insert" } else { "inserted" },
242 style::branch(branch),
243 style::branch(&parent),
244 style::branch(¤t)
245 );
246 anstream::println!(
247 "{} {} -> {}",
248 if dry_run {
249 "would retarget"
250 } else {
251 "retargeted"
252 },
253 style::branch(¤t),
254 style::branch(branch)
255 );
256 mark_floor_if_rooting(&parent, dry_run)?;
257 Ok(())
258}
259
260fn ensure_absent(branch: &str) -> Result<()> {
261 if git::local_branches()?
262 .iter()
263 .any(|existing| existing == branch)
264 {
265 bail!("branch {branch} already exists");
266 }
267 Ok(())
268}
269
270pub fn trunk_branch(branches: &[String]) -> Option<String> {
273 let remote = settings::remote().unwrap_or_else(|_| settings::DEFAULT_REMOTE.to_owned());
274 if let Some(default) = git::remote_default_branch(&remote) {
275 return Some(default);
276 }
277
278 ["main", "master"]
279 .iter()
280 .find(|name| branches.iter().any(|branch| branch == *name))
281 .map(|name| (*name).to_owned())
282}
283
284pub fn adopt_branch(branch: &str, parent: &str, dry_run: bool) -> Result<()> {
285 if branch == parent {
286 bail!("a branch cannot be its own stack parent");
287 }
288
289 let branches: BTreeSet<_> = git::local_branches()?.into_iter().collect();
290 if !branches.contains(branch) {
291 bail!("branch {branch} does not exist");
292 }
293 if !branches.contains(parent) {
294 bail!("parent branch {parent} does not exist");
295 }
296 if branch_and_descendants(branch)?
297 .iter()
298 .any(|descendant| descendant == parent)
299 {
300 bail!("{parent} is already below {branch} in the stack; that would form a cycle");
301 }
302
303 if !dry_run {
304 set_parent(branch, parent)?;
305 record_base(branch, parent);
306 }
307 anstream::println!(
308 "{} {} to {}",
309 if dry_run { "would attach" } else { "attached" },
310 style::branch(branch),
311 style::branch(parent)
312 );
313 if is_floor(branch)? {
318 if !dry_run {
319 clear_floor(branch)?;
320 }
321 anstream::println!(
322 "{}",
323 style::dim(&format!(
324 "{} {branch} is no longer a stack base",
325 if dry_run {
326 "would record that"
327 } else {
328 "recorded that"
329 }
330 ))
331 );
332 }
333 mark_floor_if_rooting(parent, dry_run)?;
334 Ok(())
335}
336
337pub fn detach_branch(branch: Option<&str>) -> Result<()> {
338 let branch = branch
339 .map(str::to_owned)
340 .map_or_else(git::current_branch, Ok)?;
341 unset_parent(&branch)?;
342 unset_base(&branch)?;
343 let was_floor = is_floor(&branch)?;
346 clear_floor(&branch)?;
347 anstream::println!("detached {}", style::branch(&branch));
348 if was_floor {
349 anstream::println!(
350 "{}",
351 style::dim(&format!("{branch} is no longer a stack base"))
352 );
353 }
354 Ok(())
355}
356
357pub fn rename_branch(old: &str, new: &str, dry_run: bool) -> Result<()> {
361 let children = children_of(old)?;
362
363 if !dry_run {
364 snapshot::take("rename");
365 git::rename_branch(old, new)?;
366 }
367 anstream::println!(
368 "{} {} -> {}",
369 if dry_run { "would rename" } else { "renamed" },
370 style::branch(old),
371 style::branch(new)
372 );
373
374 for child in &children {
375 if !dry_run {
376 set_parent(child, new)?;
377 }
378 anstream::println!(
379 "{} {} -> {}",
380 if dry_run {
381 "would retarget"
382 } else {
383 "retargeted"
384 },
385 style::branch(child),
386 style::branch(new)
387 );
388 }
389 Ok(())
390}
391
392pub fn set_renamed_from(branch: &str, old: &str) -> Result<()> {
395 git::config_set(&renamed_from_key(branch), old)
396}
397
398pub fn renamed_from(branch: &str) -> Result<Option<String>> {
400 git::config_get(&renamed_from_key(branch))
401}
402
403pub fn clear_renamed_from(branch: &str) -> Result<()> {
405 git::config_unset(&renamed_from_key(branch))
406}
407
408pub fn record_base(branch: &str, parent: &str) {
411 if let Ok(base) = git::merge_base(parent, branch) {
412 let _ = git::config_set(&base_key(branch), &base);
413 }
414}
415
416pub(crate) fn fork_point(branch: &str, parent: &str) -> Result<Option<String>> {
426 let recorded = base_of(branch)?.filter(|base| git::is_ancestor(base, branch).unwrap_or(false));
427 let merge_base = git::merge_base(parent, branch).ok();
428 Ok(match (recorded, merge_base) {
429 (Some(recorded), Some(merge_base)) => Some(
430 if git::is_ancestor(&merge_base, &recorded).unwrap_or(false) {
431 recorded
432 } else {
433 merge_base
434 },
435 ),
436 (recorded, merge_base) => recorded.or(merge_base),
437 })
438}
439
440pub(crate) fn base_is_current(branch: &str, parent: &str) -> Result<bool> {
444 let Some(base) = base_of(branch)? else {
445 return Ok(false);
446 };
447 if !git::is_ancestor(&base, branch).unwrap_or(false) {
448 return Ok(false);
449 }
450 Ok(match git::merge_base(parent, branch) {
451 Ok(merge_base) => git::is_ancestor(&merge_base, &base).unwrap_or(false),
452 Err(_) => true,
453 })
454}
455
456pub fn stack_root(branch: &str) -> Result<String> {
458 let parents = parent_map()?;
459 Ok(root_for(branch, &parents))
460}
461
462pub fn branch_and_descendants(branch: &str) -> Result<Vec<String>> {
463 let parents = parent_map()?;
464 let children = children_map(&parents);
465 let mut branches = vec![branch.to_owned()];
466 let mut visited = BTreeSet::from([branch.to_owned()]);
467 collect_descendants(branch, &children, &mut branches, &mut visited);
468 Ok(branches)
469}
470
471pub fn stack_line(branch: &str) -> Result<Vec<String>> {
477 let trunk = trunk_branch(&git::local_branches()?);
481 if Some(branch) == trunk.as_deref() {
482 return Ok(Vec::new());
483 }
484
485 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());
492 Ok(line)
493}
494
495pub(crate) fn line_base(branch: &str) -> Result<String> {
503 Ok(path_from_root(branch)?
504 .into_iter()
505 .next()
506 .unwrap_or_else(|| branch.to_owned()))
507}
508
509pub fn current_stack_branches(branch: &str) -> Result<Vec<String>> {
517 let base = line_base(branch)?;
518 let trunk = trunk_branch(&git::local_branches()?);
519 Ok(branch_and_descendants(&base)?
520 .into_iter()
521 .filter(|candidate| Some(candidate) != trunk.as_ref())
522 .collect())
523}
524
525pub fn listed_branches(all: bool) -> Result<BTreeSet<String>> {
530 if all {
531 Ok(parent_map()?
532 .into_iter()
533 .flat_map(|(child, parent)| [child, parent])
534 .collect())
535 } else {
536 let current = git::current_branch()?;
537 Ok(current_stack_branches(¤t)?.into_iter().collect())
538 }
539}
540
541pub fn publish_metadata(remote: &str) {
545 if let Err(error) = try_publish_metadata(remote) {
546 anstream::eprintln!(
547 "{}",
548 style::warn(&format!("could not publish stack metadata: {error:#}"))
549 );
550 }
551}
552
553fn try_publish_metadata(remote: &str) -> Result<()> {
554 let current = git::current_branch()?;
555 let trunk = trunk_branch(&git::local_branches()?);
556
557 let mut parents = serde_json::Map::new();
558 let mut floors = Vec::new();
559 for branch in current_stack_branches(¤t)? {
560 if is_floor(&branch)? {
565 floors.push(Value::String(branch));
566 } else if let Some(parent) = parent_of(&branch)? {
567 parents.insert(branch, Value::String(parent));
568 }
569 }
570 if parents.is_empty() {
571 return Ok(());
572 }
573
574 let document = json!({ "trunk": trunk, "parents": parents, "floors": floors });
575 git::write_blob_ref(METADATA_REF, METADATA_FILE, &document.to_string())?;
576 git::push_ref(remote, METADATA_REF)
577}
578
579pub fn apply_remote_metadata(remote: &str) -> Result<usize> {
582 git::fetch_ref(remote, METADATA_REF)
583 .context("no stack metadata on the remote - push it from the other machine first")?;
584 let Some(content) = git::read_ref_file(METADATA_REF, METADATA_FILE)? else {
585 bail!("the remote stack metadata is empty");
586 };
587
588 let document: Value =
589 serde_json::from_str(&content).context("failed to parse remote stack metadata")?;
590 let parents = document
591 .get("parents")
592 .and_then(Value::as_object)
593 .context("remote stack metadata is malformed")?;
594
595 let mut pairs = Vec::new();
600 for (branch, parent) in parents {
601 let Some(parent) = parent.as_str() else {
602 continue;
603 };
604 if !is_safe_ref_name(branch) || !is_safe_ref_name(parent) {
605 anstream::eprintln!(
606 "{}",
607 style::warn(&format!(
608 "skipping unsafe stack metadata entry: {branch:?} -> {parent:?}"
609 ))
610 );
611 continue;
612 }
613 pairs.push((branch.clone(), parent.to_owned()));
614 }
615
616 let publishes_floors = document.get("floors").is_some();
625 let floors: Vec<String> = document
626 .get("floors")
627 .and_then(Value::as_array)
628 .map(|floors| {
629 floors
630 .iter()
631 .filter_map(Value::as_str)
632 .filter(|floor| is_safe_ref_name(floor))
633 .map(str::to_owned)
634 .collect()
635 })
636 .unwrap_or_default();
637
638 let local: BTreeSet<String> = git::local_branches()?.into_iter().collect();
642 for branch in pairs.iter().map(|(branch, _)| branch).chain(floors.iter()) {
643 if !local.contains(branch) {
644 git::fetch_branch(remote, branch)
645 .with_context(|| format!("failed to fetch {branch} from {remote}"))?;
646 }
647 }
648
649 for floor in &floors {
650 if is_floor(floor)? {
651 continue;
652 }
653 mark_floor(floor);
654 anstream::println!("{} is now a stack base", style::branch(floor));
658 }
659 for (branch, _) in pairs
663 .iter()
664 .filter(|(branch, _)| publishes_floors && !floors.contains(branch))
665 {
666 if is_floor(branch)? {
667 clear_floor(branch)?;
668 anstream::println!("{} is no longer a stack base", style::branch(branch));
669 }
670 }
671
672 let mut attached = 0;
673 for (branch, parent) in &pairs {
674 set_parent(branch, parent)?;
675 record_base(branch, parent);
676 attached += 1;
677 anstream::println!(
678 "attached {} to {}",
679 style::branch(branch),
680 style::branch(parent)
681 );
682 }
683 Ok(attached)
684}
685
686pub(crate) fn is_safe_ref_name(name: &str) -> bool {
691 !name.is_empty()
692 && !name.starts_with('-')
693 && !name.chars().any(|c| c.is_whitespace() || c.is_control())
694}
695
696pub fn path_from_root(branch: &str) -> Result<Vec<String>> {
699 let trunk = trunk_branch(&git::local_branches()?);
700 let mut path = vec![branch.to_owned()];
701 let mut seen = BTreeSet::from([branch.to_owned()]);
702
703 let mut cursor = branch.to_owned();
704 while let Some(parent) = stacked_parent_of(&cursor)? {
705 if Some(&parent) == trunk.as_ref() || !seen.insert(parent.clone()) {
706 break;
707 }
708 path.push(parent.clone());
709 if is_floor(&parent)? {
712 break;
713 }
714 cursor = parent;
715 }
716
717 path.reverse();
718 Ok(path)
719}
720
721pub fn stacked_layers(line: &[String]) -> Result<Vec<String>> {
732 Ok(branch_parents(line)?
733 .into_iter()
734 .map(|(branch, _)| branch)
735 .collect())
736}
737
738pub fn unanchored_base(branches: &[String]) -> Result<Option<String>> {
748 let layers = stacked_layers(branches)?;
749 if layers.len() == branches.len() {
750 return Ok(None);
751 }
752 if layers.is_empty() {
753 let [lone] = branches else {
758 return Ok(None);
759 };
760 return Ok(is_floor(lone)?.then(|| lone.clone()));
761 }
762 Ok(branches
763 .iter()
764 .find(|branch| !layers.contains(branch))
765 .cloned())
766}
767
768pub fn branch_parents(branches: &[String]) -> Result<Vec<(String, String)>> {
773 let mut pairs = Vec::new();
774 for branch in branches {
775 if let Some(parent) = stacked_parent_of(branch)? {
776 pairs.push((branch.clone(), parent));
777 }
778 }
779 Ok(pairs)
780}
781
782fn parent_map() -> Result<BTreeMap<String, String>> {
783 let mut parents = BTreeMap::new();
784 for branch in git::local_branches()? {
785 if let Some(parent) = stacked_parent_of(&branch)? {
786 parents.insert(branch, parent);
787 }
788 }
789 Ok(parents)
790}
791
792fn collect_descendants(
793 branch: &str,
794 children: &BTreeMap<String, Vec<String>>,
795 branches: &mut Vec<String>,
796 visited: &mut BTreeSet<String>,
797) {
798 if let Some(branch_children) = children.get(branch) {
799 for child in branch_children {
800 if !visited.insert(child.to_owned()) {
801 continue; }
803 branches.push(child.to_owned());
804 collect_descendants(child, children, branches, visited);
805 }
806 }
807}
808
809pub(crate) fn has_stacked_branches() -> Result<bool> {
813 Ok(!parent_map()?.is_empty())
814}
815
816pub(crate) fn children_of(parent: &str) -> Result<Vec<String>> {
817 Ok(parent_map()?
818 .into_iter()
819 .filter_map(|(branch, branch_parent)| (branch_parent == parent).then_some(branch))
820 .collect())
821}
822
823fn children_map(parents: &BTreeMap<String, String>) -> BTreeMap<String, Vec<String>> {
824 let mut children: BTreeMap<String, Vec<String>> = BTreeMap::new();
825 for (branch, parent) in parents {
826 children
827 .entry(parent.to_owned())
828 .or_default()
829 .push(branch.to_owned());
830 }
831 children
832}
833
834fn root_for(branch: &str, parents: &BTreeMap<String, String>) -> String {
835 let mut root = branch.to_owned();
836 let mut seen = BTreeSet::new();
837
838 while let Some(parent) = parents.get(&root) {
839 if !seen.insert(root.clone()) {
840 break;
841 }
842 root = parent.to_owned();
843 }
844
845 root
846}
847
848fn mark_floor_if_rooting(parent: &str, dry_run: bool) -> Result<()> {
854 let trunk = trunk_branch(&git::local_branches()?);
855 if Some(parent) == trunk.as_deref() || stacked_parent_of(parent)?.is_some() || is_floor(parent)?
856 {
857 return Ok(());
858 }
859
860 if !dry_run {
868 mark_floor(parent);
869 }
870 anstream::println!(
871 "{}",
872 style::dim(&format!(
873 "{} {parent} as this stack's base; \
874 if it is a stacked branch, run `git stk detach {parent}`",
875 if dry_run { "would record" } else { "recorded" }
876 ))
877 );
878 Ok(())
879}
880
881pub fn is_floor(branch: &str) -> Result<bool> {
883 Ok(git::config_get(&floor_key(branch))?.is_some())
884}
885
886pub fn mark_floor(branch: &str) {
890 let _ = git::config_set(&floor_key(branch), "true");
891}
892
893pub fn clear_floor(branch: &str) -> Result<()> {
894 git::config_unset(&floor_key(branch))
895}
896
897pub(crate) fn parent_of(branch: &str) -> Result<Option<String>> {
898 git::config_get(&parent_key(branch))
899}
900
901pub(crate) fn stacked_parent_of(branch: &str) -> Result<Option<String>> {
907 if is_floor(branch)? {
908 return Ok(None);
909 }
910 parent_of(branch)
911}
912
913pub(crate) fn base_of(branch: &str) -> Result<Option<String>> {
914 git::config_get(&base_key(branch))
915}
916
917pub(crate) fn set_parent(branch: &str, parent: &str) -> Result<()> {
918 git::config_set(&parent_key(branch), parent)
919}
920
921pub(crate) fn unset_parent(branch: &str) -> Result<()> {
922 git::config_unset(&parent_key(branch))
923}
924
925pub(crate) fn set_base(branch: &str, base: &str) -> Result<()> {
926 git::config_set(&base_key(branch), base)
927}
928
929pub(crate) fn unset_base(branch: &str) -> Result<()> {
930 git::config_unset(&base_key(branch))
931}
932
933fn floor_key(branch: &str) -> String {
934 format!("branch.{branch}.{FLOOR_KEY}")
935}
936
937fn parent_key(branch: &str) -> String {
938 format!("branch.{branch}.{PARENT_KEY}")
939}
940
941fn base_key(branch: &str) -> String {
942 format!("branch.{branch}.{BASE_KEY}")
943}
944
945fn renamed_from_key(branch: &str) -> String {
946 format!("branch.{branch}.{RENAMED_FROM_KEY}")
947}
948
949#[cfg(test)]
950mod tests {
951 use super::is_safe_ref_name;
952
953 #[test]
954 fn safe_ref_names_pass() {
955 assert!(is_safe_ref_name("main"));
956 assert!(is_safe_ref_name("feature/a"));
957 assert!(is_safe_ref_name("user/fix-123"));
958 }
959
960 #[test]
961 fn unsafe_ref_names_are_rejected() {
962 assert!(!is_safe_ref_name("--upload-pack=touch /tmp/pwned"));
964 assert!(!is_safe_ref_name("-x"));
965 assert!(!is_safe_ref_name("a branch"));
967 assert!(!is_safe_ref_name("a\nb"));
968 assert!(!is_safe_ref_name("a\tb"));
969 assert!(!is_safe_ref_name(""));
970 }
971}