armdb 0.12.0

sharded bitcask key-value storage optimized for NVMe
Documentation
use std::ops::Bound;

#[cfg(feature = "loom")]
use loom::sync::atomic::Ordering;
#[cfg(not(feature = "loom"))]
use std::sync::atomic::Ordering;

use crate::key::Key;

use super::SkipList;
use super::back_finger::BackFinger;
use super::node::SkipNode;
use super::strip_mark;

/// Key-only skip list walker with configurable bounds and direction.
///
/// `against_storage` selects the step function: forward along level 0 vs the
/// resumable `back_step` path (against physical list order). Bounds use natural byte order;
/// `reversed` matches the list's storage direction for bound checks.
///
/// # Invariant
///
/// `Item = &'g N` is tied to the list, not to the iterator, while the `seize`
/// guard that keeps those nodes alive lives *inside* the iterator. A yielded
/// reference must therefore not outlive the iterator — consumers copy out of the
/// node before stepping (see [`SkipKeyKeys`]). `Iterator` cannot express the
/// tighter `&'_ self` bound (that would be a lending iterator), so this is a
/// hand-checked invariant; the type stays `pub(crate)` because of it.
macro_rules! skip_key_iter {
    ($vis:vis) => {
        $vis struct SkipKeyIter<'g, K: Key, N: SkipNode> {
            list: &'g SkipList<N>,
            front: *mut N,
            back: Option<*mut N>,
            back_finger: BackFinger<N>,
            end: Bound<K>,
            start: Bound<K>,
            /// When true, `next` walks via `back_step` (against storage order).
            against_storage: bool,
            reversed: bool,
            done: bool,
            _guard: seize::LocalGuard<'g>,
        }
    };
}

#[cfg(any(test, feature = "loom", miri))]
skip_key_iter!(pub);
#[cfg(not(any(test, feature = "loom", miri)))]
skip_key_iter!(pub(crate));

impl<'g, K: Key, N: SkipNode> SkipKeyIter<'g, K, N> {
    pub(crate) fn bounded_dir(
        list: &'g SkipList<N>,
        start: Bound<K>,
        end: Bound<K>,
        front: *mut N,
        against_storage: bool,
        reversed: bool,
        guard: seize::LocalGuard<'g>,
    ) -> Self {
        Self {
            list,
            front,
            back: None,
            back_finger: BackFinger::empty(),
            end,
            start,
            against_storage,
            reversed,
            done: false,
            _guard: guard,
        }
    }

    fn resolve_back(&mut self) -> *mut N {
        match &self.end {
            Bound::Unbounded => self
                .list
                .back_seek_last(&mut self.back_finger, &self._guard),
            Bound::Excluded(key) => {
                self.list
                    .back_step(&mut self.back_finger, key.as_bytes(), &self._guard)
            }
            Bound::Included(key) => {
                let ge = self.list.find_first_ge(key.as_bytes(), &self._guard);
                if !ge.is_null()
                    && !unsafe { &*ge }.is_marked()
                    && unsafe { &*ge }.key_bytes() == key.as_bytes()
                {
                    ge
                } else {
                    self.list
                        .back_step(&mut self.back_finger, key.as_bytes(), &self._guard)
                }
            }
        }
    }

    #[inline(always)]
    fn check_end(&self, key: &[u8]) -> bool {
        match &self.end {
            Bound::Unbounded => true,
            Bound::Excluded(end) => {
                if self.reversed {
                    key > end.as_bytes()
                } else {
                    key < end.as_bytes()
                }
            }
            Bound::Included(end) => {
                if self.reversed {
                    key >= end.as_bytes()
                } else {
                    key <= end.as_bytes()
                }
            }
        }
    }

    #[inline(always)]
    fn check_start(&self, key: &[u8]) -> bool {
        match &self.start {
            Bound::Unbounded => true,
            Bound::Excluded(start) => {
                if self.reversed {
                    key < start.as_bytes()
                } else {
                    key > start.as_bytes()
                }
            }
            Bound::Included(start) => {
                if self.reversed {
                    key <= start.as_bytes()
                } else {
                    key >= start.as_bytes()
                }
            }
        }
    }

    fn next_with_storage(&mut self) -> Option<&'g N> {
        loop {
            if self.done || self.front.is_null() {
                return None;
            }
            let node = unsafe { &*self.front };
            self.front = strip_mark(node.tower(0).load(Ordering::Acquire));
            if node.is_marked() {
                continue;
            }
            if !self.check_end(node.key_bytes()) {
                self.done = true;
                return None;
            }
            return Some(node);
        }
    }

    /// Walks against storage order via the resumable `back_step` cursor.
    ///
    /// Unlike `ConstIter::next_back` there is no front/back convergence check:
    /// every caller today constructs the against-storage variant with a null
    /// `front`, so `check_start` alone terminates the walk. A future caller that
    /// seeds both ends must add the convergence test — the debug assert below
    /// catches that mistake in tests.
    fn next_against_storage(&mut self) -> Option<&'g N> {
        debug_assert!(
            self.front.is_null(),
            "against-storage walk assumes a null front; add a convergence check first"
        );
        if self.back.is_none() {
            self.back = Some(self.resolve_back());
            if self.front.is_null() && self.back == Some(std::ptr::null_mut()) {
                self.done = true;
            }
        }
        loop {
            let back = self.back.unwrap_or(std::ptr::null_mut());
            if self.done || back.is_null() {
                return None;
            }
            let node = unsafe { &*back };
            let key = node.key_bytes();
            self.back = Some(
                self.list
                    .back_step(&mut self.back_finger, key, &self._guard),
            );
            if node.is_marked() {
                continue;
            }
            if !self.check_start(key) {
                self.done = true;
                return None;
            }
            return Some(node);
        }
    }
}

impl<'g, K: Key, N: SkipNode> Iterator for SkipKeyIter<'g, K, N> {
    type Item = &'g N;

    fn next(&mut self) -> Option<Self::Item> {
        if self.against_storage {
            self.next_against_storage()
        } else {
            self.next_with_storage()
        }
    }
}

/// Adapts [`SkipKeyIter`] to yield decoded keys.
pub(crate) struct SkipKeyKeys<'a, K: Key, N: SkipNode> {
    inner: SkipKeyIter<'a, K, N>,
}

impl<'a, K: Key, N: SkipNode> SkipKeyKeys<'a, K, N> {
    pub(crate) fn new(inner: SkipKeyIter<'a, K, N>) -> Self {
        Self { inner }
    }
}

impl<K: Key, N: SkipNode> Iterator for SkipKeyKeys<'_, K, N> {
    type Item = K;

    fn next(&mut self) -> Option<Self::Item> {
        self.inner
            .next()
            .map(|node| K::from_bytes(node.key_bytes()))
    }
}

/// Integration-test helper (`miri_skiplist`, `loom_skiplist`) — only reachable
/// when `skiplist` itself is public, otherwise it is dead code.
#[cfg(any(feature = "loom", miri))]
#[allow(private_interfaces)]
pub fn test_all<'g, N: SkipNode>(
    list: &'g SkipList<N>,
    guard: seize::LocalGuard<'g>,
) -> SkipKeyIter<'g, [u8; 4], N> {
    SkipKeyIter::bounded_dir(
        list,
        Bound::Unbounded,
        Bound::Unbounded,
        strip_mark(unsafe { (*list.head_ptr()).tower(0).load(Ordering::Acquire) }),
        false,
        false,
        guard,
    )
}

/// Integration-test helper for the against-storage production iterator path.
#[cfg(any(feature = "loom", miri))]
#[allow(private_interfaces)]
pub fn test_all_back<'g, N: SkipNode>(
    list: &'g SkipList<N>,
    guard: seize::LocalGuard<'g>,
) -> SkipKeyIter<'g, [u8; 4], N> {
    SkipKeyIter::bounded_dir(
        list,
        Bound::Unbounded,
        Bound::Unbounded,
        std::ptr::null_mut(),
        true,
        false,
        guard,
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::disk_loc::DiskLoc;
    use crate::skiplist::SkipList;
    use crate::skiplist::node::ConstNode;

    type TestNode = ConstNode<[u8; 8], 1, DiskLoc>;
    type TestList = SkipList<TestNode>;

    fn disk() -> DiskLoc {
        DiskLoc::new(0, 0, 0)
    }

    fn key_byte(i: u8) -> [u8; 8] {
        let mut k = [0u8; 8];
        k[7] = i;
        k
    }

    fn alloc_node(key: [u8; 8]) -> *mut TestNode {
        TestNode::alloc(key, [0u8; 1], disk(), 1)
    }

    fn insert_keys<'a>(list: &'a TestList, keys: &[u8]) -> seize::LocalGuard<'a> {
        let guard = list.collector().enter();
        for &i in keys {
            let node = alloc_node(key_byte(i));
            list.insert(node, &guard);
        }
        guard
    }

    fn collect_keys(iter: SkipKeyIter<'_, [u8; 8], TestNode>) -> Vec<u8> {
        iter.map(|n| n.key_bytes()[7]).collect()
    }

    #[test]
    fn skip_key_iter_yields_keys_in_order() {
        let list: TestList = SkipList::new(false);
        let guard = insert_keys(&list, &[1, 3, 5]);
        let iter = SkipKeyIter::bounded_dir(
            &list,
            Bound::Unbounded,
            Bound::Unbounded,
            strip_mark(unsafe { (*list.head_ptr()).tower(0).load(Ordering::Acquire) }),
            false,
            false,
            guard,
        );
        assert_eq!(collect_keys(iter), vec![1, 3, 5]);
    }

    #[test]
    fn skip_key_iter_against_storage_yields_desc_on_asc_list() {
        let list: TestList = SkipList::new(false);
        let guard = insert_keys(&list, &[1, 3, 5]);
        let iter = SkipKeyIter::bounded_dir(
            &list,
            Bound::Unbounded,
            Bound::Unbounded,
            std::ptr::null_mut(),
            true,
            false,
            guard,
        );
        assert_eq!(collect_keys(iter), vec![5, 3, 1]);
    }

    #[test]
    fn skip_key_iter_range_bounds_respected() {
        let list: TestList = SkipList::new(false);
        let guard = insert_keys(&list, &[1, 2, 3, 4, 5]);
        let start_key = key_byte(2);
        let end_key = key_byte(5);
        let front = list.find_first_ge(&start_key, &guard);
        let iter = SkipKeyIter::bounded_dir(
            &list,
            Bound::Included(start_key),
            Bound::Excluded(end_key),
            front,
            false,
            false,
            guard,
        );
        assert_eq!(collect_keys(iter), vec![2, 3, 4]);
    }
}