From 1d80bbb70a3fafa2d6559a1966b07f9bccc6eccb Mon Sep 17 00:00:00 2001
From: Yuya Nishihara <yuya@tcha.org>
Date: Tue, 19 Dec 2023 20:09:23 +0900
Subject: [PATCH] index: leverage ancestor iterator to collect segments to be
 squashed

I think "for" loop is easier to follow. Maybe it could be rewritten further to
.find_map() loop, but that would be too clever.

I also made ancestor_index_segments() pub(super) since it doesn't make sense
to only provide ancestor_files_without_local().
---
 lib/src/default_index/composite.rs |  8 +++++--
 lib/src/default_index/mutable.rs   | 34 ++++++++++++------------------
 2 files changed, 20 insertions(+), 22 deletions(-)

diff --git a/lib/src/default_index/composite.rs b/lib/src/default_index/composite.rs
index 943a084c4..67903921c 100644
--- a/lib/src/default_index/composite.rs
+++ b/lib/src/default_index/composite.rs
@@ -91,12 +91,16 @@ impl<'a> CompositeIndex<'a> {
         CompositeIndex(segment)
     }
 
-    fn ancestor_files_without_local(&self) -> impl Iterator<Item = &'a Arc<ReadonlyIndexSegment>> {
+    /// Iterates parent and its ancestor readonly index segments.
+    pub(super) fn ancestor_files_without_local(
+        &self,
+    ) -> impl Iterator<Item = &'a Arc<ReadonlyIndexSegment>> {
         let parent_file = self.0.segment_parent_file();
         iter::successors(parent_file, |file| file.segment_parent_file())
     }
 
-    fn ancestor_index_segments(&self) -> impl Iterator<Item = &'a dyn IndexSegment> {
+    /// Iterates self and its ancestor index segments.
+    pub(super) fn ancestor_index_segments(&self) -> impl Iterator<Item = &'a dyn IndexSegment> {
         iter::once(self.0).chain(
             self.ancestor_files_without_local()
                 .map(|file| file.as_ref() as &dyn IndexSegment),
diff --git a/lib/src/default_index/mutable.rs b/lib/src/default_index/mutable.rs
index 3b45036f0..55ae68c87 100644
--- a/lib/src/default_index/mutable.rs
+++ b/lib/src/default_index/mutable.rs
@@ -240,33 +240,27 @@ impl MutableIndexSegment {
     fn maybe_squash_with_ancestors(self) -> MutableIndexSegment {
         let mut num_new_commits = self.segment_num_commits();
         let mut files_to_squash = vec![];
-        let mut maybe_parent_file = self.parent_file.clone();
-        let mut squashed;
-        loop {
-            match maybe_parent_file {
-                Some(parent_file) => {
-                    // TODO: We should probably also squash if the parent file has less than N
-                    // commits, regardless of how many (few) are in `self`.
-                    if 2 * num_new_commits < parent_file.segment_num_commits() {
-                        squashed = MutableIndexSegment::incremental(parent_file);
-                        break;
-                    }
-                    num_new_commits += parent_file.segment_num_commits();
-                    files_to_squash.push(parent_file.clone());
-                    maybe_parent_file = parent_file.segment_parent_file().cloned();
-                }
-                None => {
-                    squashed =
-                        MutableIndexSegment::full(self.commit_id_length, self.change_id_length);
-                    break;
-                }
+        let mut base_parent_file = None;
+        for parent_file in self.as_composite().ancestor_files_without_local() {
+            // TODO: We should probably also squash if the parent file has less than N
+            // commits, regardless of how many (few) are in `self`.
+            if 2 * num_new_commits < parent_file.segment_num_commits() {
+                base_parent_file = Some(parent_file.clone());
+                break;
             }
+            num_new_commits += parent_file.segment_num_commits();
+            files_to_squash.push(parent_file.clone());
         }
 
         if files_to_squash.is_empty() {
             return self;
         }
 
+        let mut squashed = if let Some(parent_file) = base_parent_file {
+            MutableIndexSegment::incremental(parent_file)
+        } else {
+            MutableIndexSegment::full(self.commit_id_length, self.change_id_length)
+        };
         for parent_file in files_to_squash.iter().rev() {
             squashed.add_commits_from(parent_file.as_ref());
         }