Files
mini_lsm/mini-lsm-week-1/src/table/iterator.rs
Alex Chi Z 327f6badef separate week 1 solution
Signed-off-by: Alex Chi Z <iskyzh@gmail.com>
2024-01-16 16:00:51 +08:00

103 lines
2.9 KiB
Rust

use std::sync::Arc;
use anyhow::Result;
use super::SsTable;
use crate::block::BlockIterator;
use crate::iterators::StorageIterator;
/// An iterator over the contents of an SSTable.
pub struct SsTableIterator {
table: Arc<SsTable>,
blk_iter: BlockIterator,
blk_idx: usize,
}
impl SsTableIterator {
fn seek_to_first_inner(table: &Arc<SsTable>) -> Result<(usize, BlockIterator)> {
Ok((
0,
BlockIterator::create_and_seek_to_first(table.read_block_cached(0)?),
))
}
/// Create a new iterator and seek to the first key-value pair.
pub fn create_and_seek_to_first(table: Arc<SsTable>) -> Result<Self> {
let (blk_idx, blk_iter) = Self::seek_to_first_inner(&table)?;
let iter = Self {
blk_iter,
table,
blk_idx,
};
Ok(iter)
}
/// Seek to the first key-value pair.
pub fn seek_to_first(&mut self) -> Result<()> {
let (blk_idx, blk_iter) = Self::seek_to_first_inner(&self.table)?;
self.blk_idx = blk_idx;
self.blk_iter = blk_iter;
Ok(())
}
fn seek_to_key_inner(table: &Arc<SsTable>, key: &[u8]) -> Result<(usize, BlockIterator)> {
let mut blk_idx = table.find_block_idx(key);
let mut blk_iter =
BlockIterator::create_and_seek_to_key(table.read_block_cached(blk_idx)?, key);
if !blk_iter.is_valid() {
blk_idx += 1;
if blk_idx < table.num_of_blocks() {
blk_iter =
BlockIterator::create_and_seek_to_first(table.read_block_cached(blk_idx)?);
}
}
Ok((blk_idx, blk_iter))
}
/// Create a new iterator and seek to the first key-value pair which >= `key`.
pub fn create_and_seek_to_key(table: Arc<SsTable>, key: &[u8]) -> Result<Self> {
let (blk_idx, blk_iter) = Self::seek_to_key_inner(&table, key)?;
let iter = Self {
blk_iter,
table,
blk_idx,
};
Ok(iter)
}
/// Seek to the first key-value pair which >= `key`.
pub fn seek_to_key(&mut self, key: &[u8]) -> Result<()> {
let (blk_idx, blk_iter) = Self::seek_to_key_inner(&self.table, key)?;
self.blk_iter = blk_iter;
self.blk_idx = blk_idx;
Ok(())
}
}
impl StorageIterator for SsTableIterator {
fn value(&self) -> &[u8] {
self.blk_iter.value()
}
fn key(&self) -> &[u8] {
self.blk_iter.key()
}
fn is_valid(&self) -> bool {
self.blk_iter.is_valid()
}
fn next(&mut self) -> Result<()> {
self.blk_iter.next();
if !self.blk_iter.is_valid() {
self.blk_idx += 1;
if self.blk_idx < self.table.num_of_blocks() {
self.blk_iter = BlockIterator::create_and_seek_to_first(
self.table.read_block_cached(self.blk_idx)?,
);
}
}
Ok(())
}
}