From 0f43f51975a5d6ae767dadf6db93a5808a953a7b Mon Sep 17 00:00:00 2001
From: Chris Mikkelson <cmikk@fsi.io>
Date: Thu, 18 Jul 2024 20:02:20 -0500
Subject: [PATCH] merger: remove Source paramater from Iter template.

---
 src/merger.rs | 56 ++++++++++++++++++++++++++++++++++-----------------
 1 file changed, 37 insertions(+), 19 deletions(-)

diff --git a/src/merger.rs b/src/merger.rs
index e0e34e1..cf11cf6 100644
--- a/src/merger.rs
+++ b/src/merger.rs
@@ -45,13 +45,13 @@ impl<'a> Ord for MergeEntry<'a> {
     }
 }
 
-struct MergeIter<'a, S: Source> {
-    sources: &'a [S],
+struct MergeIter<'a> {
     heap: BinaryHeap<MergeEntry<'a>>,
+    finished: Vec<Box<dyn Iter + 'a>>,
     last_key: Vec<u8>,
 }
 
-impl<'a, S: Source> Iterator for MergeIter<'a, S> {
+impl<'a> Iterator for MergeIter<'a> {
     type Item = Entry;
 
     fn next(&mut self) -> Option<Self::Item> {
@@ -71,7 +71,7 @@ impl<'a, S: Source> Iterator for MergeIter<'a, S> {
     }
 }
 
-impl<'a, S: Source> Iter for MergeIter<'a, S> {
+impl<'a> Iter for MergeIter<'a> {
     fn seek(&mut self, key: &[u8]) {
         if key > self.last_key.as_slice() {
             loop {
@@ -93,14 +93,24 @@ impl<'a, S: Source> Iter for MergeIter<'a, S> {
                 self.heap.pop();
             }
         }
-        self.heap = BinaryHeap::from_iter(self.sources.iter().filter_map(|s| {
-            let mut it = s.iter();
+
+        // backwards seek; reset heap
+        let mut finished: Vec<Box<dyn Iter + 'a>> = Vec::new();
+        let mut heap_entries: Vec<MergeEntry<'a>> = Vec::new();
+        for mut it in self
+            .heap
+            .drain()
+            .map(|me| me.it)
+            .chain(self.finished.drain(..))
+        {
             it.seek(key);
-            Some(MergeEntry {
-                e: it.next()?,
-                it: Box::new(it),
-            })
-        }));
+            match it.next() {
+                Some(e) => heap_entries.push(MergeEntry { e, it }),
+                None => finished.push(it),
+            }
+        }
+        self.heap = BinaryHeap::from_iter(heap_entries);
+        self.finished = finished;
         self.last_key.clear();
         self.last_key.extend_from_slice(key);
     }
@@ -108,16 +118,24 @@ impl<'a, S: Source> Iter for MergeIter<'a, S> {
 
 impl<'a, S: Source> Source for Merger<'a, S> {
     fn iter(&self) -> impl Iter {
+        let mut v: Vec<Box<dyn Iter>> = Vec::new();
+        let h = BinaryHeap::from_iter(self.sources.iter().filter_map(|s| {
+            let mut it = s.iter();
+            match it.next() {
+                Some(e) => Some(MergeEntry {
+                    e,
+                    it: Box::new(it),
+                }),
+                None => {
+                    v.push(Box::new(it));
+                    None
+                }
+            }
+        }));
         MergeIter {
+            finished: v,
+            heap: h,
             last_key: Vec::new(),
-            sources: self.sources.as_slice(),
-            heap: BinaryHeap::from_iter(self.sources.as_slice().iter().filter_map(|s| {
-                let mut it = s.iter();
-                Some(MergeEntry {
-                    e: it.next()?,
-                    it: Box::new(it),
-                })
-            })),
         }
     }
 }
-- 
2.50.1