rewrite: update references after rewriting all commits

lib/src/rewrite.rs

@@ -538,12 +538,10 @@ impl<'settings, 'repo> DescendantRebaser<'settings, 'repo> {
 
     fn rebase_one(&mut self, old_commit: Commit) -> Result<(), TreeMergeError> {
         let old_commit_id = old_commit.id().clone();
-        if let Some(new_parent_ids) = self.mut_repo.parent_mapping.get(&old_commit_id).cloned() {
+        if self.mut_repo.parent_mapping.contains_key(&old_commit_id) {
             // This is a commit that had already been rebased before `self` was created
             // (i.e. it's part of the input for this rebase). We don't need
-            // to rebase it, but we still want to update branches pointing
-            // to the old commit.
-            self.update_references(old_commit_id, new_parent_ids)?;
+            // to rebase it.
             return Ok(());
         }
         let old_parent_ids = old_commit.parent_ids();
@@ -553,9 +551,9 @@ impl<'settings, 'repo> DescendantRebaser<'settings, 'repo> {
             self.mut_repo
                 .parent_mapping
                 .insert(old_commit_id.clone(), new_parent_ids.clone());
-            self.update_references(old_commit_id, new_parent_ids)?;
             return Ok(());
-        } else if new_parent_ids == old_parent_ids {
+        }
+        if new_parent_ids == old_parent_ids {
             // The commit is already in place.
             return Ok(());
         }
@@ -591,14 +589,26 @@ impl<'settings, 'repo> DescendantRebaser<'settings, 'repo> {
         };
         self.rebased
             .insert(old_commit_id.clone(), new_commit.id().clone());
-        self.update_references(old_commit_id, vec![new_commit.id().clone()])?;
+        Ok(())
+    }
+
+    fn update_all_references(&mut self) -> Result<(), BackendError> {
+        for (old_parent_id, new_parent_ids) in self.mut_repo.parent_mapping.clone() {
+            // Call `new_parents()` here since `parent_mapping` only contains direct
+            // mappings, not transitive ones.
+            // TODO: keep parent_mapping updated with transitive mappings so we don't need
+            // to call `new_parents()` here.
+            let new_parent_ids = self.new_parents(&new_parent_ids);
+            self.update_references(old_parent_id, new_parent_ids)?;
+        }
         Ok(())
     }
 
     pub fn rebase_all(&mut self) -> Result<(), TreeMergeError> {
         while let Some(old_commit) = self.to_visit.pop() {
             self.rebase_one(old_commit)?;
         }
+        self.update_all_references()?;
         let mut view = self.mut_repo.view().store_view().clone();
         for commit_id in &self.heads_to_remove {
             view.head_ids.remove(commit_id);

lib/tests/test_rewrite.rs

@@ -707,12 +707,12 @@ fn test_rebase_descendants_multiple_swap() {
         .set_rewritten_commit(commit_b.id().clone(), commit_d.id().clone());
     tx.mut_repo()
         .set_rewritten_commit(commit_d.id().clone(), commit_b.id().clone());
-    let _ = tx.mut_repo().rebase_descendants_return_map(&settings); // Panics because of the cycle
+    let _ = tx.mut_repo().rebase_descendants(&settings); // Panics because of
+                                                         // the cycle
 }
 
-// Unlike `test_rebase_descendants_multiple_swap`, this does not currently
-// panic, but it would probably be OK if it did.
 #[test]
+#[should_panic(expected = "cycle detected")]
 fn test_rebase_descendants_multiple_no_descendants() {
     let settings = testutils::user_settings();
     let test_repo = TestRepo::init();
@@ -733,19 +733,8 @@ fn test_rebase_descendants_multiple_no_descendants() {
         .set_rewritten_commit(commit_b.id().clone(), commit_c.id().clone());
     tx.mut_repo()
         .set_rewritten_commit(commit_c.id().clone(), commit_b.id().clone());
-    let rebase_map = tx
-        .mut_repo()
-        .rebase_descendants_return_map(&settings)
-        .unwrap();
-    assert!(rebase_map.is_empty());
-
-    assert_eq!(
-        *tx.mut_repo().view().heads(),
-        hashset! {
-            commit_b.id().clone(),
-            commit_c.id().clone()
-        }
-    );
+    let _ = tx.mut_repo().rebase_descendants(&settings); // Panics because of
+                                                         // the cycle
 }
 
 #[test]
@@ -1028,11 +1017,13 @@ fn test_rebase_descendants_branch_move_two_steps() {
     let commit_b2 = tx
         .mut_repo()
         .rewrite_commit(&settings, &commit_b)
+        .set_description("different")
         .write()
         .unwrap();
     let commit_c2 = tx
         .mut_repo()
         .rewrite_commit(&settings, &commit_c)
+        .set_description("more different")
         .write()
         .unwrap();
     tx.mut_repo().rebase_descendants(&settings).unwrap();