level_reader.rs

// Copyright (c) 2024-present, fjall-rs
// This source code is licensed under both the Apache 2.0 and MIT License
// (found in the LICENSE-* files in the repository)

use crate::{
    level_manifest::level::Level,
    segment::{range::Range, value_block::CachePolicy, Segment},
    InternalValue, UserKey,
};
use std::{ops::Bound, sync::Arc};

/// Reads through a disjoint level
pub struct LevelReader {
    segments: Arc<Level>,
    lo: usize,
    hi: usize,
    lo_reader: Option<Range>,
    hi_reader: Option<Range>,
    cache_policy: CachePolicy,
    prefix_hash: Option<(u64, u64)>,
}

impl LevelReader {
    #[must_use]
    pub fn new(
        level: Arc<Level>,
        range: &(Bound<UserKey>, Bound<UserKey>),
        cache_policy: CachePolicy,
        prefix_hash: Option<(u64, u64)>,
    ) -> Self {
        assert!(!level.is_empty(), "level reader cannot read empty level");

        let disjoint_level = level.as_disjoint().expect("level should be disjoint");

        let Some((lo, hi)) = disjoint_level.range_indexes(range) else {
            // NOTE: We will never emit any item
            return Self {
                segments: level,
                lo: 0,
                hi: 0,
                lo_reader: None,
                hi_reader: None,
                cache_policy,
                prefix_hash,
            };
        };

        Self::from_indexes(
            level,
            range,
            (Some(lo), Some(hi)),
            cache_policy,
            prefix_hash,
        )
    }

    #[must_use]
    pub fn from_indexes(
        level: Arc<Level>,
        range: &(Bound<UserKey>, Bound<UserKey>),
        (lo, hi): (Option<usize>, Option<usize>),
        cache_policy: CachePolicy,
        prefix_hash: Option<(u64, u64)>,
    ) -> Self {
        let lo = lo.unwrap_or_default();
        let hi = hi.unwrap_or(level.len() - 1);

        let lo_segment = level.segments.get(lo).expect("should exist");

        let lo_reader = lo_segment.range(range.clone()).cache_policy(cache_policy);

        let hi_reader = if hi > lo {
            let hi_segment = level.segments.get(hi).expect("should exist");

            Some(hi_segment.range(range.clone()).cache_policy(cache_policy))
        } else {
            None
        };

        Self {
            segments: level,
            lo,
            hi,
            lo_reader: Some(lo_reader),
            hi_reader,
            cache_policy,
            prefix_hash,
        }
    }

    fn may_segment_contain_hash(&self, segment: &Segment) -> bool {
        self.prefix_hash
            .map_or(true, |hash| segment.may_contain_hash(hash))
    }
}

impl Iterator for LevelReader {
    type Item = crate::Result<InternalValue>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if let Some(lo_reader) = &mut self.lo_reader {
                if let Some(item) = lo_reader.next() {
                    return Some(item);
                }

                // NOTE: Lo reader is empty, get next one
                self.lo_reader = None;
            } else if self.lo == self.hi {
                // NOTE: We reached the hi marker, so consume from it instead
                //
                // If it returns nothing, it is empty, so we are done
                return self.hi_reader.as_mut().and_then(|r| r.next());
            } else {
                self.lo += 1;

                if self.lo < self.hi {
                    let segment = self.segments.get(self.lo).expect("should exist");
                    if self.may_segment_contain_hash(segment) {
                        self.lo_reader = Some(segment.iter().cache_policy(self.cache_policy));
                    }
                }
            }
        }
    }
}

impl DoubleEndedIterator for LevelReader {
    fn next_back(&mut self) -> Option<Self::Item> {
        loop {
            if let Some(hi_reader) = &mut self.hi_reader {
                if let Some(item) = hi_reader.next_back() {
                    return Some(item);
                }

                // NOTE: Hi reader is empty, get the previous one
                self.hi_reader = None;
            } else if self.lo == self.hi {
                // NOTE: We reached the lo marker, so consume from it instead
                //
                // If it returns nothing, it is empty, so we are done
                return self.lo_reader.as_mut().and_then(|r| r.next_back());
            } else {
                self.hi -= 1;

                if self.lo < self.hi {
                    let segment = self.segments.get(self.hi).expect("should exist");
                    if self.may_segment_contain_hash(segment) {
                        self.hi_reader = Some(segment.iter().cache_policy(self.cache_policy));
                    }
                }
            }
        }
    }
}

#[cfg(test)]
#[allow(clippy::expect_used)]
mod tests {
    use super::*;
    use crate::{AbstractTree, Slice};
    use std::ops::Bound::Unbounded;
    use test_log::test;

    // TODO: same test for prefix & ranges

    #[test]
    fn level_reader_basic() -> crate::Result<()> {
        let tempdir = tempfile::tempdir()?;
        let tree = crate::Config::new(&tempdir).open()?;

        let ids = [
            ["a", "b", "c"],
            ["d", "e", "f"],
            ["g", "h", "i"],
            ["j", "k", "l"],
        ];

        for batch in ids {
            for id in batch {
                tree.insert(id, vec![], 0);
            }
            tree.flush_active_memtable(0)?;
        }

        let segments = tree
            .levels
            .read()
            .expect("lock is poisoned")
            .iter()
            .cloned()
            .collect::<Vec<_>>();

        let level = Arc::new(Level {
            segments,
            is_disjoint: true,
        });

        #[allow(clippy::unwrap_used)]
        {
            let multi_reader = LevelReader::new(
                level.clone(),
                &(Unbounded, Unbounded),
                CachePolicy::Read,
                None,
            );

            let mut iter = multi_reader.flatten();

            assert_eq!(Slice::from(*b"a"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"b"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"c"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"d"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"e"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"f"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"g"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"h"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"i"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"j"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"k"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"l"), iter.next().unwrap().key.user_key);
        }

        #[allow(clippy::unwrap_used)]
        {
            let multi_reader = LevelReader::new(
                level.clone(),
                &(Unbounded, Unbounded),
                CachePolicy::Read,
                None,
            );

            let mut iter = multi_reader.rev().flatten();

            assert_eq!(Slice::from(*b"l"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"k"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"j"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"i"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"h"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"g"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"f"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"e"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"d"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"c"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"b"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"a"), iter.next().unwrap().key.user_key);
        }

        #[allow(clippy::unwrap_used)]
        {
            let multi_reader =
                LevelReader::new(level, &(Unbounded, Unbounded), CachePolicy::Read, None);

            let mut iter = multi_reader.flatten();

            assert_eq!(Slice::from(*b"a"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"l"), iter.next_back().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"b"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"k"), iter.next_back().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"c"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"j"), iter.next_back().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"d"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"i"), iter.next_back().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"e"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"h"), iter.next_back().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"f"), iter.next().unwrap().key.user_key);
            assert_eq!(Slice::from(*b"g"), iter.next_back().unwrap().key.user_key);
        }

        Ok(())
    }
}