Refactor MutRepo to make DescendantRebaser private (pt 2)

lib/src/repo.rs

@@ -859,10 +859,6 @@ impl MutableRepo {
             .insert(new_id);
     }
 
-    pub fn clear_rewritten_commits(&mut self) {
-        self.rewritten_commits.clear();
-    }
-
     /// Record a commit as having been abandoned in this transaction.
     ///
     /// This record is used by `rebase_descendants` to know which commits have
@@ -877,29 +873,17 @@ impl MutableRepo {
         self.abandoned_commits.insert(old_id);
     }
 
-    pub fn clear_abandoned_commits(&mut self) {
+    fn clear_descendant_rebaser_plans(&mut self) {
+        self.rewritten_commits.clear();
         self.abandoned_commits.clear();
     }
 
     pub fn has_rewrites(&self) -> bool {
         !(self.rewritten_commits.is_empty() && self.abandoned_commits.is_empty())
     }
 
-    /// Creates a `DescendantRebaser` to rebase descendants of the recorded
-    /// rewritten and abandoned commits.
-    // TODO(ilyagr): Inline this. It's only used in tests.
-    pub fn create_descendant_rebaser<'settings, 'repo>(
-        &'repo mut self,
-        settings: &'settings UserSettings,
-    ) -> DescendantRebaser<'settings, 'repo> {
-        DescendantRebaser::new(
-            settings,
-            self,
-            self.rewritten_commits.clone(),
-            self.abandoned_commits.clone(),
-        )
-    }
-
+    /// After the rebaser returned by this function is dropped,
+    /// self.clear_descendant_rebaser_plans() needs to be called.
     fn rebase_descendants_return_rebaser<'settings, 'repo>(
         &'repo mut self,
         settings: &'settings UserSettings,
@@ -909,7 +893,12 @@ impl MutableRepo {
             // Optimization
             return Ok(None);
         }
-        let mut rebaser = self.create_descendant_rebaser(settings);
+        let mut rebaser = DescendantRebaser::new(
+            settings,
+            self,
+            self.rewritten_commits.clone(),
+            self.abandoned_commits.clone(),
+        );
         *rebaser.mut_options() = options;
         rebaser.rebase_all()?;
         Ok(Some(rebaser))
@@ -923,9 +912,39 @@ impl MutableRepo {
         settings: &UserSettings,
         options: RebaseOptions,
     ) -> Result<usize, TreeMergeError> {
-        Ok(self
+        let result = self
+            .rebase_descendants_return_rebaser(settings, options)?
+            .map_or(0, |rebaser| rebaser.rebased().len());
+        self.clear_descendant_rebaser_plans();
+        Ok(result)
+    }
+
+    /// This is similar to `rebase_descendants_return_map`, but the return value
+    /// needs more explaining.
+    ///
+    /// If the `options.empty` is the default, this function will only
+    /// rebase commits, and the return value is what you'd expect it to be.
+    ///
+    /// Otherwise, this function may rebase some commits and abandon others. The
+    /// behavior is such that only commits with a single parent will ever be
+    /// abandoned. In the returned map, an abandoned commit will look as a
+    /// key-value pair where the key is the abandoned commit and the value is
+    /// **its parent**. One can tell this case apart since the change ids of the
+    /// key and the value will not match. The parent will inherit the
+    /// descendants and the branches of the abandoned commit.
+    pub fn rebase_descendants_with_options_return_map(
+        &mut self,
+        settings: &UserSettings,
+        options: RebaseOptions,
+    ) -> Result<HashMap<CommitId, CommitId>, TreeMergeError> {
+        let result = Ok(self
+            // We do not set RebaseOptions here, since this function does not currently return
+            // enough information to describe the results of a rebase if some commits got
+            // abandoned
             .rebase_descendants_return_rebaser(settings, options)?
-            .map_or(0, |rebaser| rebaser.rebased().len()))
+            .map_or(HashMap::new(), |rebaser| rebaser.rebased().clone()));
+        self.clear_descendant_rebaser_plans();
+        result
     }
 
     pub fn rebase_descendants(&mut self, settings: &UserSettings) -> Result<usize, TreeMergeError> {
@@ -936,12 +955,7 @@ impl MutableRepo {
         &mut self,
         settings: &UserSettings,
     ) -> Result<HashMap<CommitId, CommitId>, TreeMergeError> {
-        Ok(self
-            // We do not set RebaseOptions here, since this function does not currently return
-            // enough information to describe the results of a rebase if some commits got
-            // abandoned
-            .rebase_descendants_return_rebaser(settings, Default::default())?
-            .map_or(HashMap::new(), |rebaser| rebaser.rebased().clone()))
+        self.rebase_descendants_with_options_return_map(settings, Default::default())
     }
 
     pub fn set_wc_commit(

lib/src/rewrite.rs

@@ -242,7 +242,7 @@ pub struct RebaseOptions {
 }
 
 /// Rebases descendants of a commit onto a new commit (or several).
-pub struct DescendantRebaser<'settings, 'repo> {
+pub(crate) struct DescendantRebaser<'settings, 'repo> {
     settings: &'settings UserSettings,
     mut_repo: &'repo mut MutableRepo,
     // The commit identified by the key has been replaced by all the ones in the value, typically
@@ -521,7 +521,7 @@ impl<'settings, 'repo> DescendantRebaser<'settings, 'repo> {
 
     // TODO: Perhaps change the interface since it's not just about rebasing
     // commits.
-    pub fn rebase_next(&mut self) -> Result<Option<RebasedDescendant>, TreeMergeError> {
+    fn rebase_next(&mut self) -> Result<Option<RebasedDescendant>, TreeMergeError> {
         while let Some(old_commit) = self.to_visit.pop() {
             let old_commit_id = old_commit.id().clone();
             if let Some(new_parent_ids) = self.parent_mapping.get(&old_commit_id).cloned() {
@@ -587,9 +587,12 @@ impl<'settings, 'repo> DescendantRebaser<'settings, 'repo> {
                 new_commit,
             }));
         }
-        // TODO: As the TODO above says, we should probably change the API. Even if we
-        // don't, we should at least make this code not do any work if you call
-        // rebase_next() after we've returned None.
+        Ok(None)
+    }
+
+    pub fn rebase_all(&mut self) -> Result<(), TreeMergeError> {
+        while self.rebase_next()?.is_some() {}
+        // TODO: As the TODO above says, we should probably change the API.
         let mut view = self.mut_repo.view().store_view().clone();
         for commit_id in &self.heads_to_remove {
             view.head_ids.remove(commit_id);
@@ -600,13 +603,6 @@ impl<'settings, 'repo> DescendantRebaser<'settings, 'repo> {
         self.heads_to_remove.clear();
         self.heads_to_add.clear();
         self.mut_repo.set_view(view);
-        self.mut_repo.clear_rewritten_commits();
-        self.mut_repo.clear_abandoned_commits();
-        Ok(None)
-    }
-
-    pub fn rebase_all(&mut self) -> Result<(), TreeMergeError> {
-        while self.rebase_next()?.is_some() {}
         Ok(())
     }
 }

lib/tests/test_commit_builder.rs

@@ -19,7 +19,7 @@ use jj_lib::repo::Repo;
 use jj_lib::repo_path::{RepoPath, RepoPathBuf};
 use jj_lib::settings::UserSettings;
 use test_case::test_case;
-use testutils::{assert_rebased, create_tree, CommitGraphBuilder, TestRepo, TestRepoBackend};
+use testutils::{assert_rebased_onto, create_tree, CommitGraphBuilder, TestRepo, TestRepoBackend};
 
 fn to_owned_path_vec(paths: &[&RepoPath]) -> Vec<RepoPathBuf> {
     paths.iter().map(|&path| path.to_owned()).collect()
@@ -269,8 +269,11 @@ fn test_commit_builder_descendants(backend: TestRepoBackend) {
         )
         .write()
         .unwrap();
-    let mut rebaser = tx.mut_repo().create_descendant_rebaser(&settings);
-    assert!(rebaser.rebase_next().unwrap().is_none());
+    let rebase_map = tx
+        .mut_repo()
+        .rebase_descendants_return_map(&settings)
+        .unwrap();
+    assert_eq!(rebase_map.len(), 0);
 
     // Test with for_rewrite_from()
     let mut tx = repo.start_transaction(&settings);
@@ -279,9 +282,12 @@ fn test_commit_builder_descendants(backend: TestRepoBackend) {
         .rewrite_commit(&settings, &commit2)
         .write()
         .unwrap();
-    let mut rebaser = tx.mut_repo().create_descendant_rebaser(&settings);
-    assert_rebased(rebaser.rebase_next().unwrap(), &commit3, &[&commit4]);
-    assert!(rebaser.rebase_next().unwrap().is_none());
+    let rebase_map = tx
+        .mut_repo()
+        .rebase_descendants_return_map(&settings)
+        .unwrap();
+    assert_rebased_onto(tx.mut_repo(), &rebase_map, &commit3, &[commit4.id()]);
+    assert_eq!(rebase_map.len(), 1);
 
     // Test with for_rewrite_from() but new change id
     let mut tx = repo.start_transaction(&settings);
@@ -290,6 +296,9 @@ fn test_commit_builder_descendants(backend: TestRepoBackend) {
         .generate_new_change_id()
         .write()
         .unwrap();
-    let mut rebaser = tx.mut_repo().create_descendant_rebaser(&settings);
-    assert!(rebaser.rebase_next().unwrap().is_none());
+    let rebase_map = tx
+        .mut_repo()
+        .rebase_descendants_return_map(&settings)
+        .unwrap();
+    assert!(rebase_map.is_empty());
 }

lib/tests/test_mut_repo.rs

@@ -17,7 +17,7 @@ use jj_lib::op_store::{RefTarget, RemoteRef, RemoteRefState, WorkspaceId};
 use jj_lib::repo::Repo;
 use maplit::hashset;
 use testutils::{
-    assert_rebased, create_random_commit, write_random_commit, CommitGraphBuilder, TestRepo,
+    assert_rebased_onto, create_random_commit, write_random_commit, CommitGraphBuilder, TestRepo,
 };
 
 #[test]
@@ -517,9 +517,7 @@ fn test_has_changed() {
 
 #[test]
 fn test_rebase_descendants_simple() {
-    // Tests that MutableRepo::create_descendant_rebaser() creates a
-    // DescendantRebaser that rebases descendants of rewritten and abandoned
-    // commits.
+    // There are many additional tests of this functionality in `test_rewrite.rs`.
     let settings = testutils::user_settings();
     let test_repo = TestRepo::init();
     let repo = &test_repo.repo;
@@ -539,24 +537,29 @@ fn test_rebase_descendants_simple() {
     let commit6 = graph_builder.commit_with_parents(&[&commit1]);
     mut_repo.record_rewritten_commit(commit2.id().clone(), commit6.id().clone());
     mut_repo.record_abandoned_commit(commit4.id().clone());
-    let mut rebaser = mut_repo.create_descendant_rebaser(&settings);
+    let rebase_map = tx
+        .mut_repo()
+        .rebase_descendants_return_map(&settings)
+        .unwrap();
     // Commit 3 got rebased onto commit 2's replacement, i.e. commit 6
-    assert_rebased(rebaser.rebase_next().unwrap(), &commit3, &[&commit6]);
+    assert_rebased_onto(tx.mut_repo(), &rebase_map, &commit3, &[commit6.id()]);
     // Commit 5 got rebased onto commit 4's parent, i.e. commit 1
-    assert_rebased(rebaser.rebase_next().unwrap(), &commit5, &[&commit1]);
-    assert!(rebaser.rebase_next().unwrap().is_none());
+    assert_rebased_onto(tx.mut_repo(), &rebase_map, &commit5, &[commit1.id()]);
+    assert_eq!(rebase_map.len(), 2);
+
     // No more descendants to rebase if we try again.
-    assert!(mut_repo
-        .create_descendant_rebaser(&settings)
-        .rebase_next()
-        .unwrap()
-        .is_none());
+    let rebase_map = tx
+        .mut_repo()
+        .rebase_descendants_return_map(&settings)
+        .unwrap();
+    assert_eq!(rebase_map.len(), 0);
 }
 
 #[test]
 fn test_rebase_descendants_conflicting_rewrite() {
-    // Tests MutableRepo::create_descendant_rebaser() when a commit has been marked
-    // as rewritten to several other commits.
+    // Test rebasing descendants when one commit was rewritten to several other
+    // commits. There are many additional tests of this functionality in
+    // `test_rewrite.rs`.
     let settings = testutils::user_settings();
     let test_repo = TestRepo::init();
     let repo = &test_repo.repo;
@@ -575,16 +578,13 @@ fn test_rebase_descendants_conflicting_rewrite() {
     let commit5 = graph_builder.commit_with_parents(&[&commit1]);
     mut_repo.record_rewritten_commit(commit2.id().clone(), commit4.id().clone());
     mut_repo.record_rewritten_commit(commit2.id().clone(), commit5.id().clone());
-    let mut rebaser = mut_repo.create_descendant_rebaser(&settings);
     // Commit 3 does *not* get rebased because it's unclear if it should go onto
     // commit 4 or commit 5
-    assert!(rebaser.rebase_next().unwrap().is_none());
-    // No more descendants to rebase if we try again.
-    assert!(mut_repo
-        .create_descendant_rebaser(&settings)
-        .rebase_next()
-        .unwrap()
-        .is_none());
+    let rebase_map = tx
+        .mut_repo()
+        .rebase_descendants_return_map(&settings)
+        .unwrap();
+    assert!(rebase_map.is_empty());
 }
 
 #[test]