feat(graph): Align stacks on repair

docs/reference.md

@@ -128,3 +128,4 @@ Configuration is read from the following (in precedence order):
 | stack.show-format      | --format | "silent", "branches", "branch-commits", "commits", "debug"  | How to show the stacked diffs at the end |
 | stack.show-stacked     | \-       | bool                       | Show branches as stacked on top of each other, where possible |
 | stack.auto-fixup       | --fixup  | "ignore", "move", "squash" | Default fixup operation with `--rebase` |
+| stack.auto-repair      | \-       | bool                       | Perform branch repair with `--rebase` |

src/bin/git-stack/args.rs

@@ -44,6 +44,12 @@ pub struct Args {
     )]
     pub fixup: Option<git_stack::config::Fixup>,
 
+    /// Repair diverging branches.
+    #[structopt(long, overrides_with("no-repair"))]
+    repair: bool,
+    #[structopt(long, overrides_with("repair"), hidden(true))]
+    no_repair: bool,
+
     #[structopt(short = "n", long)]
     pub dry_run: bool,
 
@@ -85,8 +91,22 @@ impl Args {
             show_format: self.format,
             show_stacked: None,
             auto_fixup: None,
+            auto_repair: None,
 
             capacity: None,
         }
     }
+
+    pub fn repair(&self) -> Option<bool> {
+        resolve_bool_arg(self.repair, self.no_repair)
+    }
+}
+
+fn resolve_bool_arg(yes: bool, no: bool) -> Option<bool> {
+    match (yes, no) {
+        (true, false) => Some(true),
+        (false, true) => Some(false),
+        (false, false) => None,
+        (_, _) => unreachable!("StructOpt should make this impossible"),
+    }
 }

src/bin/git-stack/stack.rs

@@ -17,6 +17,7 @@ struct State {
     pull: bool,
     push: bool,
     fixup: git_stack::config::Fixup,
+    repair: bool,
     dry_run: bool,
     snapshot_capacity: Option<usize>,
     protect_commit_count: Option<usize>,
@@ -59,6 +60,20 @@ impl State {
                 no_op
             }
         };
+        let repair = match (args.repair(), args.rebase) {
+            (Some(repair), _) => repair,
+            (_, true) => repo_config.auto_repair(),
+            _ => {
+                // Assume the user is only wanting to show the tree and not modify it.
+                if repo_config.auto_repair() {
+                    log::trace!(
+                        "Ignoring `auto-repair={}` without an explicit `--rebase`",
+                        repo_config.auto_repair()
+                    );
+                }
+                false
+            }
+        };
         let push = args.push;
         let protected = git_stack::git::ProtectedBranches::new(
             repo_config.protected_branches().iter().map(|s| s.as_str()),
@@ -180,6 +195,7 @@ impl State {
             pull,
             push,
             fixup,
+            repair,
             dry_run,
             snapshot_capacity,
             protect_commit_count,
@@ -288,7 +304,7 @@ pub fn stack(
     let mut success = true;
     let mut backed_up = false;
     let mut stash_id = None;
-    if state.rebase || state.fixup != git_stack::config::Fixup::Ignore {
+    if state.rebase || state.fixup != git_stack::config::Fixup::Ignore || state.repair {
         if stash_id.is_none() && !state.dry_run {
             stash_id = git_stack::git::stash_push(&mut state.repo, "branch-stash");
         }
@@ -425,6 +441,9 @@ fn plan_changes(state: &State, stack: &StackState) -> eyre::Result<git_stack::gi
         ));
     }
     git_stack::graph::fixup(&mut graph, state.fixup);
+    if state.repair {
+        git_stack::graph::realign_stacks(&mut graph);
+    }
 
     let mut script = git_stack::graph::to_script(&graph);
     script.commands.extend(
@@ -560,6 +579,9 @@ fn show(state: &State, colored_stdout: bool, colored_stderr: bool) -> eyre::Resu
                 );
             }
             git_stack::graph::fixup(&mut graph, state.fixup);
+            if state.repair {
+                git_stack::graph::realign_stacks(&mut graph);
+            }
         }
 
         git_stack::graph::pushable(&mut graph);

src/config.rs

@@ -11,6 +11,7 @@ pub struct RepoConfig {
     pub show_format: Option<Format>,
     pub show_stacked: Option<bool>,
     pub auto_fixup: Option<Fixup>,
+    pub auto_repair: Option<bool>,
 
     pub capacity: Option<usize>,
 }
@@ -24,6 +25,7 @@ static PULL_REMOTE_FIELD: &str = "stack.pull-remote";
 static FORMAT_FIELD: &str = "stack.show-format";
 static STACKED_FIELD: &str = "stack.show-stacked";
 static AUTO_FIXUP_FIELD: &str = "stack.auto-fixup";
+static AUTO_REPAIR_FIELD: &str = "stack.auto-repair";
 static BACKUP_CAPACITY_FIELD: &str = "branch-stash.capacity";
 
 static DEFAULT_PROTECTED_BRANCHES: [&str; 4] = ["main", "master", "dev", "stable"];
@@ -150,6 +152,8 @@ impl RepoConfig {
                 if let Some(value) = value.as_ref().and_then(|v| FromStr::from_str(v).ok()) {
                     config.auto_fixup = Some(value);
                 }
+            } else if key == AUTO_REPAIR_FIELD {
+                config.auto_repair = Some(value.as_ref().map(|v| v == "true").unwrap_or(true));
             } else if key == BACKUP_CAPACITY_FIELD {
                 config.capacity = value.as_deref().and_then(|s| s.parse::<usize>().ok());
             } else {
@@ -249,6 +253,8 @@ impl RepoConfig {
             .ok()
             .and_then(|s| FromStr::from_str(&s).ok());
 
+        let auto_repair = config.get_bool(AUTO_REPAIR_FIELD).ok();
+
         let capacity = config
             .get_i64(BACKUP_CAPACITY_FIELD)
             .map(|i| i as usize)
@@ -264,6 +270,7 @@ impl RepoConfig {
             show_format,
             show_stacked,
             auto_fixup,
+            auto_repair,
 
             capacity,
         }
@@ -303,6 +310,7 @@ impl RepoConfig {
         self.show_format = other.show_format.or(self.show_format);
         self.show_stacked = other.show_stacked.or(self.show_stacked);
         self.auto_fixup = other.auto_fixup.or(self.auto_fixup);
+        self.auto_repair = other.auto_repair.or(self.auto_repair);
         self.capacity = other.capacity.or(self.capacity);
 
         self
@@ -350,6 +358,10 @@ impl RepoConfig {
         self.auto_fixup.unwrap_or_else(Default::default)
     }
 
+    pub fn auto_repair(&self) -> bool {
+        self.auto_repair.unwrap_or(true)
+    }
+
     pub fn capacity(&self) -> Option<usize> {
         let capacity = self.capacity.unwrap_or(DEFAULT_CAPACITY);
         (capacity != 0).then(|| capacity)
@@ -415,6 +427,12 @@ impl std::fmt::Display for RepoConfig {
             AUTO_FIXUP_FIELD.split_once(".").unwrap().1,
             self.auto_fixup()
         )?;
+        writeln!(
+            f,
+            "\t{}={}",
+            AUTO_REPAIR_FIELD.split_once(".").unwrap().1,
+            self.auto_repair()
+        )?;
         writeln!(f, "[{}]", BACKUP_CAPACITY_FIELD.split_once(".").unwrap().0)?;
         writeln!(
             f,

src/graph/ops.rs

@@ -93,7 +93,7 @@ fn protect_large_branches_recursive(
                 protect_large_branches_recursive(graph, child_id, count + 1, max, large_branches);
         }
         if needs_protection {
-            let mut node = graph.get_mut(node_id).expect("all children exist");
+            let node = graph.get_mut(node_id).expect("all children exist");
             node.action = crate::graph::Action::Protected;
         }
     } else {
@@ -108,7 +108,7 @@ fn mark_branch_protected(graph: &mut Graph, node_id: git2::Oid, branches: &mut V
     let mut protected_queue = VecDeque::new();
     protected_queue.push_back(node_id);
     while let Some(current_id) = protected_queue.pop_front() {
-        let mut current = graph.get_mut(current_id).expect("all children exist");
+        let current = graph.get_mut(current_id).expect("all children exist");
         current.action = crate::graph::Action::Protected;
 
         if current.branches.is_empty() {
@@ -431,7 +431,7 @@ pub fn drop_squashed_by_tree_id(
     }
     while let Some(current_id) = protected_queue.pop_front() {
         let current_children = graph
-            .get_mut(current_id)
+            .get(current_id)
             .expect("all children exist")
             .children
             .clone();
@@ -544,7 +544,7 @@ pub fn fixup(graph: &mut Graph, effect: crate::config::Fixup) {
     }
     while let Some(current_id) = protected_queue.pop_front() {
         let current_children = graph
-            .get_mut(current_id)
+            .get(current_id)
             .expect("all children exist")
             .children
             .clone();
@@ -570,7 +570,7 @@ fn fixup_branch(
 
     let mut outstanding = std::collections::BTreeMap::new();
     let node_children = graph
-        .get_mut(node_id)
+        .get(node_id)
         .expect("all children exist")
         .children
         .clone();
@@ -611,7 +611,7 @@ fn fixup_node(
     debug_assert_ne!(effect, crate::config::Fixup::Ignore);
 
     let node_children = graph
-        .get_mut(node_id)
+        .get(node_id)
         .expect("all children exist")
         .children
         .clone();
@@ -712,6 +712,67 @@ fn splice_after(graph: &mut Graph, node_id: git2::Oid, fixup_ids: Vec<git2::Oid>
     }
 }
 
+/// When a branch has extra commits, update dependent branches to the latest
+pub fn realign_stacks(graph: &mut Graph) {
+    let mut protected_queue = VecDeque::new();
+    let root_action = graph.root().action;
+    if root_action.is_protected() {
+        protected_queue.push_back(graph.root_id());
+    }
+    while let Some(current_id) = protected_queue.pop_front() {
+        let current_children = graph
+            .get(current_id)
+            .expect("all children exist")
+            .children
+            .clone();
+
+        for child_id in current_children {
+            let child_action = graph.get(child_id).expect("all children exist").action;
+            if child_action.is_protected() || child_action.is_delete() {
+                protected_queue.push_back(child_id);
+            } else {
+                realign_stack(graph, child_id);
+            }
+        }
+    }
+}
+
+fn realign_stack(graph: &mut Graph, node_id: git2::Oid) {
+    let mut children = std::collections::BTreeSet::new();
+
+    let mut current_id = node_id;
+    loop {
+        let current = graph.get_mut(current_id).expect("all children exist");
+        if current.branches.is_empty() {
+            let mut current_children: Vec<_> = current.children.iter().copied().collect();
+            match current_children.len() {
+                0 => {
+                    current.children.extend(children);
+                    return;
+                }
+                1 => {
+                    current_id = current_children.into_iter().next().unwrap();
+                }
+                _ => {
+                    // Assuming the more recent work is a continuation of the existing stack and
+                    // aligning the other stacks to be on top of it
+                    //
+                    // This should be safe in light of our rebases since we don't preserve the time
+                    current_children.sort_unstable_by_key(|id| {
+                        graph.get(*id).expect("all children exist").commit.time
+                    });
+                    let newest = current_children.pop().unwrap();
+                    children.extend(current_children);
+                    current_id = newest;
+                }
+            }
+        } else {
+            current.children.extend(children);
+            return;
+        }
+    }
+}
+
 pub fn to_script(graph: &Graph) -> crate::git::Script {
     let mut script = crate::git::Script::new();