kip_db 0.1.2-alpha.26.fix1

轻量级、异步 基于LSM Leveled Compaction K-V数据库
Documentation
use crate::kernel::lsm::iterator::{Iter, Seek, SeekIter};
use crate::kernel::lsm::mem_table::KeyValue;
use crate::kernel::lsm::table::btree_table::BTreeTable;
use bytes::Bytes;
use std::collections::btree_map::Range;
use std::collections::Bound;

pub(crate) struct BTreeTableIter<'a> {
    inner: Option<Range<'a, Bytes, KeyValue>>,
    table: &'a BTreeTable,
}

impl<'a> BTreeTableIter<'a> {
    pub(crate) fn new(table: &'a BTreeTable) -> BTreeTableIter<'a> {
        let mut iter = BTreeTableIter { inner: None, table };
        iter._seek(Seek::First);
        iter
    }

    fn _seek(&mut self, seek: Seek) {
        self.inner = match seek {
            Seek::First => Some(
                self.table
                    .inner
                    .range::<Bytes, (Bound<Bytes>, Bound<Bytes>)>((
                        Bound::Unbounded,
                        Bound::Unbounded,
                    )),
            ),
            Seek::Last => None,
            Seek::Backward(key) => Some(
                self.table
                    .inner
                    .range::<Bytes, (Bound<Bytes>, Bound<Bytes>)>((
                        Bound::Included(Bytes::copy_from_slice(key)),
                        Bound::Unbounded,
                    )),
            ),
        };
    }
}

impl<'a> Iter<'a> for BTreeTableIter<'a> {
    type Item = KeyValue;

    fn try_next(&mut self) -> crate::kernel::KernelResult<Option<Self::Item>> {
        Ok(self
            .inner
            .as_mut()
            .and_then(|iter| iter.next())
            .map(item_clone))
    }

    fn is_valid(&self) -> bool {
        true
    }
}

impl<'a> SeekIter<'a> for BTreeTableIter<'a> {
    fn seek(&mut self, seek: Seek<'_>) -> crate::kernel::KernelResult<()> {
        self._seek(seek);

        Ok(())
    }
}

fn item_clone((_, value): (&Bytes, &KeyValue)) -> KeyValue {
    value.clone()
}

#[cfg(test)]
mod tests {
    use crate::kernel::lsm::iterator::Seek;
    use crate::kernel::lsm::table::btree_table::BTreeTable;
    use crate::kernel::lsm::table::Table;
    use crate::kernel::KernelResult;
    use bytes::Bytes;

    #[test]
    fn test_iterator() -> KernelResult<()> {
        let vec = vec![
            (Bytes::from(vec![b'1']), None),
            (Bytes::from(vec![b'2']), Some(Bytes::from(vec![b'1']))),
            (Bytes::from(vec![b'3']), None),
            (Bytes::from(vec![b'4']), None),
            (Bytes::from(vec![b'5']), Some(Bytes::from(vec![b'2']))),
            (Bytes::from(vec![b'6']), None),
        ];
        let table = BTreeTable::new(0, 0, vec.clone());
        let mut iter = table.iter()?;

        for test_data in vec.clone() {
            assert_eq!(iter.try_next()?, Some(test_data))
        }

        assert_eq!(iter.try_next()?, None);

        iter.seek(Seek::First)?;
        assert_eq!(iter.try_next()?, Some(vec[0].clone()));

        iter.seek(Seek::Backward(&[b'3']))?;
        assert_eq!(iter.try_next()?, Some(vec[2].clone()));

        iter.seek(Seek::Last)?;
        assert_eq!(iter.try_next()?, None);

        Ok(())
    }
}