From 7cef879ef605d740d27e025fde490e40e4c483fc Mon Sep 17 00:00:00 2001
From: Ilya Grigoriev <ilyagr@users.noreply.github.com>
Date: Sun, 17 Dec 2023 21:47:01 -0800
Subject: [PATCH] lib `repo.rs` & `rewrite.rs`: Move clearing of
 rewritten/abandoned commits

This commit is a little out of place in this sequence, but
it seems to make more sense for MutRepo to own these maps.

@yuja [pointed out] that any tests written using `create_descendant_rebaser` now
need to do this cleanup, but there are no longer any such tests after the
previous commits and a follow-up commit removes `create_descendant_rebaser`
entirely.

[pointed out]: https://github.com/martinvonz/jj/pull/2737#discussion_r1435754370
---
 lib/src/repo.rs    | 21 ++++++++++++---------
 lib/src/rewrite.rs |  2 --
 2 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/lib/src/repo.rs b/lib/src/repo.rs
index e7bf6e41b4..05050843b5 100644
--- a/lib/src/repo.rs
+++ b/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,7 +873,8 @@ 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();
     }
 
@@ -900,6 +897,8 @@ impl MutableRepo {
         )
     }
 
+    /// 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,
@@ -923,9 +922,11 @@ 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()))
+            .map_or(0, |rebaser| rebaser.rebased().len());
+        self.clear_descendant_rebaser_plans();
+        Ok(result)
     }
 
     pub fn rebase_descendants(&mut self, settings: &UserSettings) -> Result<usize, TreeMergeError> {
@@ -936,12 +937,14 @@ impl MutableRepo {
         &mut self,
         settings: &UserSettings,
     ) -> Result<HashMap<CommitId, CommitId>, TreeMergeError> {
-        Ok(self
+        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, Default::default())?
-            .map_or(HashMap::new(), |rebaser| rebaser.rebased().clone()))
+            .map_or(HashMap::new(), |rebaser| rebaser.rebased().clone()));
+        self.clear_descendant_rebaser_plans();
+        result
     }
 
     pub fn set_wc_commit(
diff --git a/lib/src/rewrite.rs b/lib/src/rewrite.rs
index f6d449b4ee..e179fb29db 100644
--- a/lib/src/rewrite.rs
+++ b/lib/src/rewrite.rs
@@ -600,8 +600,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)
     }