feat(code): part 3 iterators

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

mini-lsm/src/table/builder.rs

@@ -1,6 +1,7 @@
+use std::path::Path;
+
 use anyhow::Result;
 use bytes::BufMut;
-use std::path::Path;
 
 use super::{BlockMeta, FileObject, SsTable};
 use crate::block::BlockBuilder;
@@ -11,36 +12,29 @@ pub struct SsTableBuilder {
     first_key: Vec<u8>,
     data: Vec<u8>,
     pub(super) meta: Vec<BlockMeta>,
-    target_size: usize,
     block_size: usize,
 }
 
 impl SsTableBuilder {
     /// Create a builder based on target SST size and target block size.
-    pub fn new(target_size: usize, block_size: usize) -> Self {
+    pub fn new(block_size: usize) -> Self {
         Self {
             data: Vec::new(),
             meta: Vec::new(),
             first_key: Vec::new(),
-            target_size,
             block_size,
             builder: BlockBuilder::new(block_size),
         }
     }
 
-    /// Adds a key-value pair to SSTable, return false when SST full.
-    #[must_use]
-    pub fn add(&mut self, key: &[u8], value: &[u8]) -> bool {
-        if self.data.len() > self.target_size {
-            return false;
-        }
-
+    /// Adds a key-value pair to SSTable
+    pub fn add(&mut self, key: &[u8], value: &[u8]) {
         if self.first_key.is_empty() {
             self.first_key = key.to_vec();
         }
 
         if self.builder.add(key, value) {
-            return true;
+            return;
         }
         // create a new block builder and append block data
         self.finish_block();
@@ -48,8 +42,11 @@ impl SsTableBuilder {
         // add the key-value pair to the next block
         assert!(self.builder.add(key, value));
         self.first_key = key.to_vec();
+    }
 
-        true
+    /// Get the estimated size of the SSTable.
+    pub fn estimated_size(&self) -> usize {
+        self.data.len()
     }
 
     fn finish_block(&mut self) {
@@ -62,7 +59,8 @@ impl SsTableBuilder {
         self.data.extend(encoded_block);
     }
 
-    /// Builds the SSTable and writes it to the given path. No need to actually write to disk until chapter 4 block cache.
+    /// Builds the SSTable and writes it to the given path. No need to actually write to disk until
+    /// chapter 4 block cache.
     pub fn build(mut self, path: impl AsRef<Path>) -> Result<SsTable> {
         self.finish_block();
         let mut buf = self.data;

mini-lsm/src/table/iterator.rs

@@ -1,8 +1,10 @@
-use anyhow::Result;
 use std::sync::Arc;
 
+use anyhow::Result;
+
 use super::SsTable;
 use crate::block::BlockIterator;
+use crate::iterators::impls::StorageIterator;
 
 /// An iterator over the contents of an SSTable.
 pub struct SsTableIterator {
@@ -68,25 +70,22 @@ impl SsTableIterator {
         self.blk_idx = blk_idx;
         Ok(())
     }
+}
 
-    /// Get the current key.
-    pub fn key(&self) -> &[u8] {
-        self.blk_iter.key()
-    }
-
-    /// Get the current value.
-    pub fn value(&self) -> &[u8] {
+impl StorageIterator for SsTableIterator {
+    fn value(&self) -> &[u8] {
         self.blk_iter.value()
     }
 
-    /// Check if the iterator is valid.
-    pub fn is_valid(&self) -> bool {
+    fn key(&self) -> &[u8] {
+        self.blk_iter.key()
+    }
+
+    fn is_valid(&self) -> bool {
         self.blk_iter.is_valid()
     }
 
-    /// Move to the next key-value pair.
-    #[allow(clippy::should_implement_trait)]
-    pub fn next(&mut self) -> Result<()> {
+    fn next(&mut self) -> Result<()> {
         self.blk_iter.next();
         if !self.blk_iter.is_valid() {
             self.blk_idx += 1;

mini-lsm/src/table/tests.rs

@@ -2,41 +2,30 @@ use std::sync::Arc;
 
 use bytes::Bytes;
 
+use super::*;
+use crate::iterators::impls::StorageIterator;
 use crate::table::SsTableBuilder;
 
-use super::{SsTable, SsTableIterator};
-
 #[test]
 fn test_sst_build_single_key() {
-    let mut builder = SsTableBuilder::new(16, 16);
-    assert!(builder.add(b"233", b"233333"));
+    let mut builder = SsTableBuilder::new(16);
+    builder.add(b"233", b"233333");
     builder.build("").unwrap();
 }
 
 #[test]
 fn test_sst_build_two_blocks() {
-    let mut builder = SsTableBuilder::new(1024, 16);
-    assert!(builder.add(b"11", b"11"));
-    assert!(builder.add(b"22", b"22"));
-    assert!(builder.add(b"33", b"11"));
-    assert!(builder.add(b"44", b"22"));
-    assert!(builder.add(b"55", b"11"));
-    assert!(builder.add(b"66", b"22"));
+    let mut builder = SsTableBuilder::new(16);
+    builder.add(b"11", b"11");
+    builder.add(b"22", b"22");
+    builder.add(b"33", b"11");
+    builder.add(b"44", b"22");
+    builder.add(b"55", b"11");
+    builder.add(b"66", b"22");
     assert!(builder.meta.len() >= 2);
     builder.build("").unwrap();
 }
 
-#[test]
-fn test_sst_build_full() {
-    let mut builder = SsTableBuilder::new(32, 16);
-    assert!(builder.add(b"11", b"11"));
-    assert!(builder.add(b"22", b"22"));
-    assert!(builder.add(b"33", b"11"));
-    assert!(builder.add(b"44", b"22"));
-    assert!(!builder.add(b"55", b"11"));
-    builder.build("").unwrap();
-}
-
 fn key_of(idx: usize) -> Vec<u8> {
     format!("key_{:03}", idx * 5).into_bytes()
 }
@@ -50,11 +39,11 @@ fn num_of_keys() -> usize {
 }
 
 fn generate_sst() -> SsTable {
-    let mut builder = SsTableBuilder::new(65536, 128);
+    let mut builder = SsTableBuilder::new(128);
     for idx in 0..num_of_keys() {
         let key = key_of(idx);
         let value = value_of(idx);
-        assert!(builder.add(&key[..], &value[..]));
+        builder.add(&key[..], &value[..]);
     }
     builder.build("").unwrap()
 }