From 2f3cc8cdd2505fc9ba9b6bc435ca822101a54542 Mon Sep 17 00:00:00 2001
From: Kerollmops <clement@meilisearch.com>
Date: Thu, 12 Dec 2024 16:15:37 +0100
Subject: [PATCH] Fix the merge_caches_sorted function

---
 crates/milli/src/update/new/extract/cache.rs | 57 +++++++++-----------
 1 file changed, 26 insertions(+), 31 deletions(-)

diff --git a/crates/milli/src/update/new/extract/cache.rs b/crates/milli/src/update/new/extract/cache.rs
index 09ca60211..62c00d2b1 100644
--- a/crates/milli/src/update/new/extract/cache.rs
+++ b/crates/milli/src/update/new/extract/cache.rs
@@ -477,21 +477,16 @@ where
     F: for<'a> FnMut(&'a [u8], DelAddRoaringBitmap) -> Result<()>,
 {
     let mut maps = Vec::new();
-    let mut readers = Vec::new();
-    let mut current_bucket = None;
-    for FrozenCache { bucket, cache, ref mut spilled } in frozen {
-        assert_eq!(*current_bucket.get_or_insert(bucket), bucket);
-        maps.push(cache);
-        readers.append(spilled);
-    }
-
-    // First manage the spilled entries by looking into the HashMaps,
-    // merge them and mark them as dummy.
     let mut heap = BinaryHeap::new();
-    for (source_index, source) in readers.into_iter().enumerate() {
-        let mut cursor = source.into_cursor()?;
-        if cursor.move_on_next()?.is_some() {
-            heap.push(Entry { cursor, source_index });
+    let mut current_bucket = None;
+    for FrozenCache { bucket, cache, spilled } in frozen {
+        assert_eq!(*current_bucket.get_or_insert(bucket), bucket);
+        maps.push((bucket, cache));
+        for reader in spilled {
+            let mut cursor = reader.into_cursor()?;
+            if cursor.move_on_next()?.is_some() {
+                heap.push(Entry { cursor, bucket });
+            }
         }
     }
 
@@ -508,25 +503,25 @@ where
 
         let mut output = DelAddRoaringBitmap::from_bytes(first_value)?;
         while let Some(mut entry) = heap.peek_mut() {
-            if let Some((key, _value)) = entry.cursor.current() {
-                if first_key == key {
-                    let new = DelAddRoaringBitmap::from_bytes(first_value)?;
-                    output = output.merge(new);
-                    // When we are done we the current value of this entry move make
-                    // it move forward and let the heap reorganize itself (on drop)
-                    if entry.cursor.move_on_next()?.is_none() {
-                        PeekMut::pop(entry);
-                    }
-                } else {
+            if let Some((key, value)) = entry.cursor.current() {
+                if first_key != key {
                     break;
                 }
+
+                let new = DelAddRoaringBitmap::from_bytes(value)?;
+                output = output.merge(new);
+                // When we are done we the current value of this entry move make
+                // it move forward and let the heap reorganize itself (on drop)
+                if entry.cursor.move_on_next()?.is_none() {
+                    PeekMut::pop(entry);
+                }
             }
         }
 
         // Once we merged all of the spilled bitmaps we must also
         // fetch the entries from the non-spilled entries (the HashMaps).
-        for (map_index, map) in maps.iter_mut().enumerate() {
-            if first_entry.source_index != map_index {
+        for (map_bucket, map) in maps.iter_mut() {
+            if first_entry.bucket != *map_bucket {
                 if let Some(new) = map.get_mut(first_key) {
                     output.union_and_clear_bbbul(new);
                 }
@@ -538,12 +533,12 @@ where
 
         // Don't forget to put the first entry back into the heap.
         if first_entry.cursor.move_on_next()?.is_some() {
-            heap.push(first_entry)
+            heap.push(first_entry);
         }
     }
 
     // Then manage the content on the HashMap entries that weren't taken (mem::take).
-    while let Some(mut map) = maps.pop() {
+    while let Some((_, mut map)) = maps.pop() {
         // Make sure we don't try to work with entries already managed by the spilled
         let mut ordered_entries: Vec<_> =
             map.iter_mut().filter(|(_, bbbul)| !bbbul.is_empty()).collect();
@@ -553,7 +548,7 @@ where
             let mut output = DelAddRoaringBitmap::empty();
             output.union_and_clear_bbbul(bbbul);
 
-            for rhs in maps.iter_mut() {
+            for (_, rhs) in maps.iter_mut() {
                 if let Some(new) = rhs.get_mut(key) {
                     output.union_and_clear_bbbul(new);
                 }
@@ -569,14 +564,14 @@ where
 
 struct Entry<R> {
     cursor: ReaderCursor<R>,
-    source_index: usize,
+    bucket: usize,
 }
 
 impl<R> Ord for Entry<R> {
     fn cmp(&self, other: &Entry<R>) -> Ordering {
         let skey = self.cursor.current().map(|(k, _)| k);
         let okey = other.cursor.current().map(|(k, _)| k);
-        skey.cmp(&okey).then(self.source_index.cmp(&other.source_index)).reverse()
+        skey.cmp(&okey).then(self.bucket.cmp(&other.bucket)).reverse()
     }
 }