mini_lsm/table/
iterator.rs

1use std::sync::Arc;
2
3use anyhow::Result;
4
5use super::SsTable;
6use crate::block::BlockIterator;
7use crate::iterators::StorageIterator;
8use crate::key::KeySlice;
9
10/// An iterator over the contents of an SSTable.
11pub struct SsTableIterator {
12    table: Arc<SsTable>,
13    blk_iter: BlockIterator,
14    blk_idx: usize,
15}
16
17impl SsTableIterator {
18    fn seek_to_first_inner(table: &Arc<SsTable>) -> Result<(usize, BlockIterator)> {
19        Ok((
20            0,
21            BlockIterator::create_and_seek_to_first(table.read_block_cached(0)?),
22        ))
23    }
24
25    /// Create a new iterator and seek to the first key-value pair.
26    pub fn create_and_seek_to_first(table: Arc<SsTable>) -> Result<Self> {
27        let (blk_idx, blk_iter) = Self::seek_to_first_inner(&table)?;
28        let iter = Self {
29            blk_iter,
30            table,
31            blk_idx,
32        };
33        Ok(iter)
34    }
35
36    /// Seek to the first key-value pair.
37    pub fn seek_to_first(&mut self) -> Result<()> {
38        let (blk_idx, blk_iter) = Self::seek_to_first_inner(&self.table)?;
39        self.blk_idx = blk_idx;
40        self.blk_iter = blk_iter;
41        Ok(())
42    }
43
44    fn seek_to_key_inner(table: &Arc<SsTable>, key: KeySlice) -> Result<(usize, BlockIterator)> {
45        let mut blk_idx = table.find_block_idx(key);
46        let mut blk_iter =
47            BlockIterator::create_and_seek_to_key(table.read_block_cached(blk_idx)?, key);
48        if !blk_iter.is_valid() {
49            blk_idx += 1;
50            if blk_idx < table.num_of_blocks() {
51                blk_iter =
52                    BlockIterator::create_and_seek_to_first(table.read_block_cached(blk_idx)?);
53            }
54        }
55        Ok((blk_idx, blk_iter))
56    }
57
58    /// Create a new iterator and seek to the first key-value pair which >= `key`.
59    pub fn create_and_seek_to_key(table: Arc<SsTable>, key: KeySlice) -> Result<Self> {
60        let (blk_idx, blk_iter) = Self::seek_to_key_inner(&table, key)?;
61        let iter = Self {
62            blk_iter,
63            table,
64            blk_idx,
65        };
66        Ok(iter)
67    }
68
69    /// Seek to the first key-value pair which >= `key`.
70    pub fn seek_to_key(&mut self, key: KeySlice) -> Result<()> {
71        let (blk_idx, blk_iter) = Self::seek_to_key_inner(&self.table, key)?;
72        self.blk_iter = blk_iter;
73        self.blk_idx = blk_idx;
74        Ok(())
75    }
76}
77
78impl StorageIterator for SsTableIterator {
79    type KeyType<'a> = KeySlice<'a>;
80
81    fn value(&self) -> &[u8] {
82        self.blk_iter.value()
83    }
84
85    fn key(&self) -> KeySlice {
86        self.blk_iter.key()
87    }
88
89    fn is_valid(&self) -> bool {
90        self.blk_iter.is_valid()
91    }
92
93    fn next(&mut self) -> Result<()> {
94        self.blk_iter.next();
95        if !self.blk_iter.is_valid() {
96            self.blk_idx += 1;
97            if self.blk_idx < self.table.num_of_blocks() {
98                self.blk_iter = BlockIterator::create_and_seek_to_first(
99                    self.table.read_block_cached(self.blk_idx)?,
100                );
101            }
102        }
103        Ok(())
104    }
105}