checkin part 2 solution

Signed-off-by: Alex Chi <iskyzh@gmail.com>

mini-lsm/src/iterators/concat_iterator.rs

@@ -0,0 +1,122 @@
+use std::sync::Arc;
+
+use anyhow::Result;
+
+use crate::table::{SsTable, SsTableIterator};
+
+use super::StorageIterator;
+
+/// Concat multiple iterators ordered in key order and their key ranges do not overlap. We do not want to create the
+/// iterators when initializing this iterator to reduce the overhead of seeking.
+pub struct SstConcatIterator {
+    current: Option<SsTableIterator>,
+    next_sst_idx: usize,
+    sstables: Vec<Arc<SsTable>>,
+}
+
+impl SstConcatIterator {
+    fn check_sst_valid(sstables: &[Arc<SsTable>]) {
+        for sst in sstables {
+            assert!(sst.first_key() <= sst.last_key());
+        }
+        if !sstables.is_empty() {
+            for i in 0..(sstables.len() - 1) {
+                assert!(sstables[i].last_key() < sstables[i + 1].first_key());
+            }
+        }
+    }
+
+    pub fn create_and_seek_to_first(sstables: Vec<Arc<SsTable>>) -> Result<Self> {
+        Self::check_sst_valid(&sstables);
+        if sstables.is_empty() {
+            return Ok(Self {
+                current: None,
+                next_sst_idx: 0,
+                sstables,
+            });
+        }
+        let mut iter = Self {
+            current: Some(SsTableIterator::create_and_seek_to_first(
+                sstables[0].clone(),
+            )?),
+            next_sst_idx: 1,
+            sstables,
+        };
+        iter.move_until_valid()?;
+        Ok(iter)
+    }
+
+    pub fn create_and_seek_to_key(sstables: Vec<Arc<SsTable>>, key: &[u8]) -> Result<Self> {
+        Self::check_sst_valid(&sstables);
+        let idx: usize = sstables
+            .partition_point(|table| table.first_key() <= key)
+            .saturating_sub(1);
+        if idx >= sstables.len() {
+            return Ok(Self {
+                current: None,
+                next_sst_idx: sstables.len(),
+                sstables,
+            });
+        }
+        let mut iter = Self {
+            current: Some(SsTableIterator::create_and_seek_to_key(
+                sstables[idx].clone(),
+                key,
+            )?),
+            next_sst_idx: idx + 1,
+            sstables,
+        };
+        iter.move_until_valid()?;
+        Ok(iter)
+    }
+
+    fn move_until_valid(&mut self) -> Result<()> {
+        loop {
+            if let Some(iter) = self.current.as_mut() {
+                if iter.is_valid() {
+                    break;
+                }
+                if self.next_sst_idx >= self.sstables.len() {
+                    self.current = None;
+                } else {
+                    self.current = Some(SsTableIterator::create_and_seek_to_first(
+                        self.sstables[self.next_sst_idx].clone(),
+                    )?);
+                    self.next_sst_idx += 1;
+                }
+            } else {
+                break;
+            }
+        }
+        Ok(())
+    }
+}
+
+impl StorageIterator for SstConcatIterator {
+    fn key(&self) -> &[u8] {
+        self.current.as_ref().unwrap().key()
+    }
+
+    fn value(&self) -> &[u8] {
+        self.current.as_ref().unwrap().value()
+    }
+
+    fn is_valid(&self) -> bool {
+        if let Some(current) = &self.current {
+            assert!(current.is_valid());
+            true
+        } else {
+            false
+        }
+    }
+
+    fn next(&mut self) -> Result<()> {
+        self.current.as_mut().unwrap().next()?;
+        self.move_until_valid()?;
+        Ok(())
+    }
+
+    fn num_active_iterators(&self) -> usize {
+        1
+    }
+}

mini-lsm/src/iterators/merge_iterator.rs

@@ -77,13 +77,11 @@ impl<I: StorageIterator> MergeIterator<I> {
 
 impl<I: StorageIterator> StorageIterator for MergeIterator<I> {
     fn key(&self) -> &[u8] {
-        unsafe { self.current.as_ref().unwrap_unchecked() }.1.key()
+        self.current.as_ref().unwrap().1.key()
     }
 
     fn value(&self) -> &[u8] {
-        unsafe { self.current.as_ref().unwrap_unchecked() }
-            .1
-            .value()
+        self.current.as_ref().unwrap().1.value()
     }
 
     fn is_valid(&self) -> bool {
@@ -94,7 +92,7 @@ impl<I: StorageIterator> StorageIterator for MergeIterator<I> {
     }
 
     fn next(&mut self) -> Result<()> {
-        let current = unsafe { self.current.as_mut().unwrap_unchecked() };
+        let current = self.current.as_mut().unwrap();
         // Pop the item out of the heap if they have the same value.
         while let Some(mut inner_iter) = self.iters.peek_mut() {
             debug_assert!(
@@ -136,4 +134,16 @@ impl<I: StorageIterator> StorageIterator for MergeIterator<I> {
 
         Ok(())
     }
+
+    fn num_active_iterators(&self) -> usize {
+        self.iters
+            .iter()
+            .map(|x| x.1.num_active_iterators())
+            .sum::<usize>()
+            + self
+                .current
+                .as_ref()
+                .map(|x| x.1.num_active_iterators())
+                .unwrap_or(0)
+    }
 }

mini-lsm/src/iterators/two_merge_iterator.rs

@@ -77,4 +77,8 @@ impl<A: StorageIterator, B: StorageIterator> StorageIterator for TwoMergeIterato
         self.choose_a = Self::choose_a(&self.a, &self.b);
         Ok(())
     }
+
+    fn num_active_iterators(&self) -> usize {
+        self.a.num_active_iterators() + self.b.num_active_iterators()
+    }
 }