rewrite: when rebasing forward, also rebase "side branches"

As the updates test case shows, when rebasing forward, we missed
commits that fork off from the section between the source and the
destination.

As part of the fix, I also restructured the code a bit to prepare for
support for rebasing descendants of multiple rewritten commits.

lib/src/rewrite.rs

@@ -15,6 +15,7 @@
 use std::collections::{HashMap, HashSet};
 
 use itertools::Itertools;
+use maplit::hashmap;
 
 use crate::backend::CommitId;
 use crate::commit::Commit;
@@ -105,10 +106,13 @@ pub fn back_out_commit(
 pub struct DescendantRebaser<'settings, 'repo> {
     settings: &'settings UserSettings,
     mut_repo: &'repo mut MutableRepo,
-    old_parent_id: CommitId,
-    new_parent_ids: Vec<CommitId>,
+    replacements: HashMap<CommitId, Vec<CommitId>>,
     // In reverse order, so we can remove the last one to rebase first.
     to_rebase: Vec<CommitId>,
+    // Ancestors of the destinations. These were also in `to_rebase` to start with, but we don't
+    // actually rebase them. Instead, we record them in `replacements` when we visit them. That
+    // way, their descendants will be rebased correctly.
+    ancestors: HashSet<CommitId>,
     rebased: HashMap<CommitId, CommitId>,
 }
 
@@ -119,26 +123,41 @@ impl<'settings, 'repo> DescendantRebaser<'settings, 'repo> {
         old_parent_id: CommitId,
         new_parent_ids: Vec<CommitId>,
     ) -> DescendantRebaser<'settings, 'repo> {
-        let expression = RevsetExpression::commit(old_parent_id.clone())
+        let replacements = hashmap! { old_parent_id => new_parent_ids};
+        let old_commits_expression =
+            RevsetExpression::commits(replacements.keys().cloned().collect());
+        let new_commits_expression =
+            RevsetExpression::commits(replacements.values().flatten().cloned().collect());
+
+        let to_rebase_expression = old_commits_expression
             .descendants(&RevsetExpression::all_non_obsolete_heads())
-            .minus(
-                &RevsetExpression::commit(old_parent_id.clone())
-                    .union(&RevsetExpression::commits(new_parent_ids.clone()))
-                    .ancestors(),
-            );
-        let revset = expression.evaluate(mut_repo.as_repo_ref()).unwrap();
+            .minus(&old_commits_expression);
+        let to_rebase_revset = to_rebase_expression
+            .evaluate(mut_repo.as_repo_ref())
+            .unwrap();
         let mut to_rebase = vec![];
-        for index_entry in revset.iter() {
+        for index_entry in to_rebase_revset.iter() {
             to_rebase.push(index_entry.commit_id());
         }
-        drop(revset);
+        drop(to_rebase_revset);
+
+        let ancestors_expression =
+            to_rebase_expression.intersection(&new_commits_expression.ancestors());
+        let ancestors_revset = ancestors_expression
+            .evaluate(mut_repo.as_repo_ref())
+            .unwrap();
+        let mut ancestors = HashSet::new();
+        for index_entry in ancestors_revset.iter() {
+            ancestors.insert(index_entry.commit_id());
+        }
+        drop(ancestors_revset);
 
         DescendantRebaser {
             settings,
             mut_repo,
-            old_parent_id,
-            new_parent_ids,
+            replacements,
             to_rebase,
+            ancestors,
             rebased: Default::default(),
         }
     }
@@ -156,15 +175,19 @@ impl<'settings, 'repo> DescendantRebaser<'settings, 'repo> {
             let mut new_parent_ids = vec![];
             let old_parent_ids = old_commit.parent_ids();
             for old_parent_id in &old_parent_ids {
-                if old_parent_id == &self.old_parent_id {
-                    new_parent_ids.extend(self.new_parent_ids.clone());
+                if let Some(replacements) = self.replacements.get(old_parent_id) {
+                    new_parent_ids.extend(replacements.clone());
                 } else if let Some(new_parent_id) = self.rebased.get(old_parent_id) {
                     new_parent_ids.push(new_parent_id.clone());
                 } else {
                     new_parent_ids.push(old_parent_id.clone());
                 };
             }
-            if new_parent_ids == old_parent_ids {
+            if self.ancestors.contains(&old_commit_id) {
+                // Update the `replacements` map so descendants are rebased correctly.
+                self.replacements.insert(old_commit_id, new_parent_ids);
+                RebasedDescendant::AncestorOfDestination(old_commit)
+            } else if new_parent_ids == old_parent_ids {
                 RebasedDescendant::AlreadyInPlace(old_commit)
             } else {
                 // Don't create commit where one parent is an ancestor of another.
@@ -202,6 +225,7 @@ impl<'settings, 'repo> DescendantRebaser<'settings, 'repo> {
 #[derive(Debug, PartialEq, Eq, Clone)]
 pub enum RebasedDescendant {
     AlreadyInPlace(Commit),
+    AncestorOfDestination(Commit),
     Rebased {
         old_commit: Commit,
         new_commit: Commit,

lib/tests/test_rewrite.rs

@@ -29,6 +29,14 @@ fn assert_in_place(rebased: Option<RebasedDescendant>, expected_old_commit: &Com
     }
 }
 
+fn assert_ancestor(rebased: Option<RebasedDescendant>, expected_old_commit: &Commit) {
+    if let Some(RebasedDescendant::AncestorOfDestination(old_commit)) = rebased {
+        assert_eq!(old_commit, *expected_old_commit);
+    } else {
+        panic!("expected ancestor commit: {:?}", rebased);
+    }
+}
+
 fn assert_rebased(
     rebased: Option<RebasedDescendant>,
     expected_old_commit: &Commit,
@@ -42,7 +50,7 @@ fn assert_rebased(
         assert_eq!(old_commit, *expected_old_commit);
         assert_eq!(new_commit.change_id(), expected_old_commit.change_id());
         assert_eq!(&new_commit.parent_ids(), expected_new_parents);
-        return new_commit;
+        new_commit
     } else {
         panic!("expected rebased commit: {:?}", rebased);
     }
@@ -93,11 +101,15 @@ fn test_rebase_descendants_forward(use_git: bool) {
     let settings = testutils::user_settings();
     let (_temp_dir, repo) = testutils::init_repo(&settings, use_git);
 
-    // Commit 2 was replaced by commit 3. Commit 5 should be rebased (commit 4 is
-    // already in place).
+    // Commit 2 was replaced by commit 6. Commits 3 and 5 should be rebased onto 6.
+    // Commit 4 does not get rebased because it's an ancestor of the
+    // destination. Commit 7 does not get replaced because it's already in
+    // place.
     //
-    // 4
-    // 3 5
+    // 7
+    // 6 5
+    // |/
+    // 4 3
     // |/
     // 2
     // 1
@@ -106,19 +118,24 @@ fn test_rebase_descendants_forward(use_git: bool) {
     let commit1 = graph_builder.initial_commit();
     let commit2 = graph_builder.commit_with_parents(&[&commit1]);
     let commit3 = graph_builder.commit_with_parents(&[&commit2]);
-    let commit4 = graph_builder.commit_with_parents(&[&commit3]);
-    let commit5 = graph_builder.commit_with_parents(&[&commit2]);
+    let commit4 = graph_builder.commit_with_parents(&[&commit2]);
+    let commit5 = graph_builder.commit_with_parents(&[&commit4]);
+    let commit6 = graph_builder.commit_with_parents(&[&commit4]);
+    let commit7 = graph_builder.commit_with_parents(&[&commit6]);
 
     let mut rebaser = DescendantRebaser::new(
         &settings,
         tx.mut_repo(),
         commit2.id().clone(),
-        vec![commit3.id().clone()],
+        vec![commit6.id().clone()],
     );
-    assert_in_place(rebaser.rebase_next(), &commit4);
-    assert_rebased(rebaser.rebase_next(), &commit5, &[commit3.id().clone()]);
+    assert_rebased(rebaser.rebase_next(), &commit3, &[commit6.id().clone()]);
+    assert_ancestor(rebaser.rebase_next(), &commit4);
+    assert_rebased(rebaser.rebase_next(), &commit5, &[commit6.id().clone()]);
+    assert_ancestor(rebaser.rebase_next(), &commit6);
+    assert_in_place(rebaser.rebase_next(), &commit7);
     assert!(rebaser.rebase_next().is_none());
-    assert_eq!(rebaser.rebased().len(), 1);
+    assert_eq!(rebaser.rebased().len(), 2);
 
     tx.discard();
 }