armdb 0.14.2

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

use crate::error::DbResult;
use crate::key::Key;
use crate::map_index::Page;
use crate::map_index::{Dir, OrdKey, bound_ordkey, merge_lower, merge_upper};

use super::chunked::{ChunkedPage, PageSpec, StopReason, page_chunks};
use super::key_math;
use super::scan::{OrderedScan, OwnedRows};

/// One page of mapped rows returned by [`Scope::page_by`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MappedPage<K, R> {
    /// Accepted items on this page, in scan order.
    pub items: Vec<R>,
    /// Cursor for the next page: the key of the row that filled the page
    /// (`items.len() == limit`), `None` when the scan ran out first.
    ///
    /// The scan stops on the row that fills the page, so rows rejected *after* it
    /// are re-scanned by the next call — the cursor never skips them.
    pub next: Option<K>,
}

/// Field-granular prefix / range / page builder. Opaque.
///
/// `Clone` is unconditional (`S: Copy`, `S::Key: Key: Copy`); [`Debug`](fmt::Debug)
/// is too — it is hand-written, because a derive would add a hidden
/// `S::Key: Debug` bound that [`Key`] does not promise.
///
/// `'l` carries no borrow of its own: every real lifetime lives inside
/// `S` (`TreeScan<&'l Coll>` / `MapScan<&'l View>`). It is kept as a covariance
/// marker so the collection borrow stays visible in the type — do not "fix" the
/// unused parameter away, it is load-bearing for callers that name the type.
#[derive(Clone)]
pub struct Scope<'l, S: OrderedScan + 'l, const N: usize> {
    pub(super) scan: S,
    pub(super) prefix_key: S::Key,
    pub(super) start: Bound<S::Key>,
    pub(super) end: Bound<S::Key>,
    pub(super) dir: Dir,
    pub(super) _lt: std::marker::PhantomData<&'l ()>,
}

/// Hex view of a key: [`Key`] promises only `Copy`, never `Debug`.
struct HexKey<'a, K: Key>(&'a K);

impl<K: Key> fmt::Debug for HexKey<'_, K> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("0x")?;
        for byte in Key::as_bytes(self.0) {
            write!(f, "{byte:02x}")?;
        }
        Ok(())
    }
}

struct HexBound<'a, K: Key>(&'a Bound<K>);

impl<K: Key> fmt::Debug for HexBound<'_, K> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self.0 {
            Bound::Included(k) => f.debug_tuple("Included").field(&HexKey(k)).finish(),
            Bound::Excluded(k) => f.debug_tuple("Excluded").field(&HexKey(k)).finish(),
            Bound::Unbounded => f.write_str("Unbounded"),
        }
    }
}

impl<'l, S: OrderedScan + 'l, const N: usize> fmt::Debug for Scope<'l, S, N> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // Hand-written on purpose: `#[derive(Debug)]` bounds the field projection
        // `S::Key: Debug`, and armdb keys are usually `#[repr(C)]` zerocopy structs
        // without a `Debug` derive — the impl would silently not exist for exactly
        // the keys this is meant to help debug. Keys print as hex; `scan` is an
        // opaque view and adds nothing.
        f.debug_struct("Scope")
            .field("N", &N)
            .field("prefix_key", &HexKey(&self.prefix_key))
            .field("start", &HexBound(&self.start))
            .field("end", &HexBound(&self.end))
            .field("dir", &self.dir)
            .finish_non_exhaustive()
    }
}

impl<'l, S: OrderedScan + 'l, const N: usize> Scope<'l, S, N> {
    pub(crate) fn from_parts(
        scan: S,
        prefix_key: S::Key,
        start: Bound<S::Key>,
        end: Bound<S::Key>,
        dir: Dir,
    ) -> Self {
        // The `N <= fields` half of the contract lives here, not at the call sites:
        // this is the only way to build a `Scope`, so no future collection family can
        // construct one with a bogus `N` by forgetting the assert.
        const { key_math::assert_scope_n::<S::Key>(N) };
        Self {
            scan,
            prefix_key,
            start,
            end,
            dir,
            _lt: std::marker::PhantomData,
        }
    }

    /// Narrow the scan to a sub-range on fields `[N..M)`. Last call wins.
    ///
    /// The whole contract `N < M <= KEY_SCHEME.fields()` (plus key-width) is checked
    /// in a `const { … }` block. Const-eval of such a block runs at
    /// *monomorphisation*, i.e. during codegen, so a bad `M` fails the **build** of
    /// the calling crate (`cargo build`, `cargo test`, rustdoc doctest) — and only
    /// when the call is actually instantiated. Bare `cargo check` never codegens and
    /// therefore never reports it: a CI job that only runs `cargo check` will not
    /// catch a violated `range::<M>` contract.
    ///
    /// The positive control below and the `compile_fail` test after it are a pair:
    /// `compile_fail` passes on *any* compile error, so if the shared preamble ever
    /// rots (a renamed macro, a moved path) the negative test would keep passing
    /// while checking nothing — the positive one fails loudly instead.
    ///
    /// ```
    /// # use armdb::ConstTree;
    /// # use armour_core::{KeyScheme, KeyType};
    /// # use std::ops::Bound;
    /// # #[derive(Clone, Copy, zerocopy::FromBytes, zerocopy::IntoBytes, zerocopy::Immutable)]
    /// # #[repr(C)]
    /// # struct K3 { a: u64, b: u64, c: u64 }
    /// # armdb::impl_key_zerocopy!(K3, KeyScheme::Typed(&[KeyType::U64, KeyType::U64, KeyType::U64]));
    /// fn fine(tree: &ConstTree<K3, 8>) {
    ///     let scope = tree.scope::<2>(&K3 { a: 1, b: 2, c: 0 });
    ///     // M (=3) is > N (=2) and <= fields (=3) — the const asserts pass.
    ///     let _ = scope.range::<3>(Bound::Unbounded, Bound::Unbounded);
    /// }
    /// let _f: fn(&ConstTree<K3, 8>) = fine;
    /// ```
    ///
    /// ```compile_fail
    /// # use armdb::ConstTree;
    /// # use armour_core::{KeyScheme, KeyType};
    /// # use std::ops::Bound;
    /// # #[derive(Clone, Copy, zerocopy::FromBytes, zerocopy::IntoBytes, zerocopy::Immutable)]
    /// # #[repr(C)]
    /// # struct K3 { a: u64, b: u64, c: u64 }
    /// # armdb::impl_key_zerocopy!(K3, KeyScheme::Typed(&[KeyType::U64, KeyType::U64, KeyType::U64]));
    /// fn boom(tree: &ConstTree<K3, 8>) {
    ///     let scope = tree.scope::<2>(&K3 { a: 1, b: 2, c: 0 });
    ///     // M (=1) is not > N (=2) — must fail at compile time after T2c-A.
    ///     let _ = scope.range::<1>(Bound::Unbounded, Bound::Unbounded);
    /// }
    /// // Force monomorphisation: an unreferenced `fn` body is skipped, so the
    /// // const assert never runs. A fn-pointer mention is enough.
    /// let _f: fn(&ConstTree<K3, 8>) = boom;
    /// ```
    pub fn range<const M: usize>(mut self, start: Bound<&S::Key>, end: Bound<&S::Key>) -> Self {
        const { key_math::assert_range_m::<S::Key>(N, M) };
        let (s, e) = key_math::saturate_range_bounds(&self.prefix_key, N, M, start, end);
        self.start = s;
        self.end = e;
        self
    }

    /// Scan in ascending key order.
    ///
    /// Maps and LSM trees reverse for free; on skiplist trees, the first opposing-direction seek
    /// is expected O(log n) and an m-row opposing-direction scan is expected/amortized O(log n + m).
    pub fn asc(mut self) -> Self {
        self.dir = Dir::Asc;
        self
    }

    /// Scan in descending key order.
    ///
    /// Maps and LSM trees reverse for free; on skiplist trees, the first opposing-direction seek
    /// is expected O(log n) and an m-row opposing-direction scan is expected/amortized O(log n + m).
    pub fn desc(mut self) -> Self {
        self.dir = Dir::Desc;
        self
    }

    /// Consume the builder and return a row iterator.
    ///
    /// Takes `self` so fluent temporaries work:
    /// `tree.scope::<N>(…).range::<M>(…).iter()`. `Scope` is [`Clone`], so a scope
    /// you want to both iterate and paginate can be cloned — [`page`](Self::page) /
    /// [`page_by`](Self::page_by) take `&self` and do not consume it.
    pub fn iter(self) -> S::Iter {
        self.scan.scan(self.start, self.end, self.dir)
    }

    /// Consume the builder and return a key-only iterator.
    ///
    /// For skiplist-backed trees (`VarTypedTree`, etc.), this walks keys without
    /// decoding values or reading value bytes from disk.
    ///
    /// For [`LsmTree`](crate::lsm::LsmTree), the key-only path uses
    /// [`fjall::Guard::key`]; value bytes are not loaded only when the fjall
    /// keyspace is key-value-separated — inline values may still be read.
    ///
    /// On a partially corrupt collection `keys()` and [`iter`](Self::iter) disagree
    /// by design: the codec-typed families ([`VarTypedTree`](crate::VarTypedTree),
    /// [`VarTypedMap`](crate::VarTypedMap), [`LsmTree`](crate::lsm::LsmTree)) skip a
    /// row whose value will not decode, while `keys()` never decodes and yields it.
    /// A cursor walked with `keys()` therefore covers rows that `page` skips.
    pub fn keys(self) -> S::KeyIter {
        self.scan.scan_keys(self.start, self.end, self.dir)
    }

    fn bounds_for_page(&self, after: Option<&S::Key>) -> (Bound<S::Key>, Bound<S::Key>) {
        let range_start = bound_ordkey(bound_ref(&self.start));
        let range_end = bound_ordkey(bound_ref(&self.end));
        match (after, self.dir) {
            (None, _) => (clone_bound(&self.start), clone_bound(&self.end)),
            (Some(a), Dir::Asc) => {
                let lower = merge_lower(range_start, Some(a));
                (from_ord(lower), clone_bound(&self.end))
            }
            (Some(a), Dir::Desc) => {
                let upper = merge_upper(range_end, Some(a));
                (clone_bound(&self.start), from_ord(upper))
            }
        }
    }

    /// Keyset pagination with a per-row filter.
    ///
    /// `after` is exclusive. `Ok(Some(r))` accepts a row, `Ok(None)` rejects and
    /// continues scanning, `Err` aborts. When the page fills, `next` is the key of
    /// the row that filled it; rejected rows scanned before it are simply dropped.
    ///
    /// Contract for `f`:
    /// - it runs **outside** every collection lock — Map shard mutexes are released
    ///   between batches, so the closure may touch other collections;
    /// - on skiplist-backed trees a `seize` guard is held for the whole page, so
    ///   slow work inside `f` (e.g. I/O against another collection for a join)
    ///   delays memory reclamation for that long;
    /// - writing to *this* collection while paginating is safe (no deadlock), but
    ///   whether a concurrent write shows up on the current page is unspecified.
    pub fn page_by<R>(
        &self,
        after: Option<&S::Key>,
        limit: usize,
        mut f: impl FnMut(S::Row) -> DbResult<Option<R>>,
    ) -> DbResult<MappedPage<S::Key, R>> {
        if limit == 0 {
            return Ok(MappedPage {
                items: Vec::new(),
                next: None,
            });
        }
        let (start, end) = self.bounds_for_page(after);
        // `limit` often comes straight from an RPC parameter — preallocate a page,
        // not whatever the caller asked for.
        let mut items = Vec::with_capacity(limit.min(1024));
        for row in self.scan.scan(start, end, self.dir) {
            let key = S::peek_row_key(&row);
            if let Some(r) = f(row)? {
                items.push(r);
                if items.len() == limit {
                    return Ok(MappedPage {
                        items,
                        next: Some(key),
                    });
                }
            }
        }
        Ok(MappedPage { items, next: None })
    }

    /// Keyset pagination that prepares rows in chunks before projecting them.
    ///
    /// `after` is exclusive. A [`super::StopReason::Filled`] page returns the
    /// filling row's cursor, [`super::StopReason::Exhausted`] returns `None`, and
    /// [`super::StopReason::BudgetExhausted`] returns the last examined row's
    /// cursor. A prepared suffix after a page fills is discarded and re-read on
    /// continuation, so `prepare` must be read-only or idempotent; errors return
    /// no page.
    ///
    /// As with [`page_by`](Self::page_by), Map closures run outside collection
    /// locks, while a skiplist `seize` guard spans pagination. Writes to this
    /// collection are deadlock-free, but their visibility on the current page is
    /// unspecified.
    pub fn page_by_chunks<C, R>(
        &self,
        after: Option<&S::Key>,
        spec: PageSpec,
        prepare: impl FnMut(&[S::Row]) -> DbResult<C>,
        project: impl FnMut(S::Row, &C) -> DbResult<Option<R>>,
    ) -> DbResult<ChunkedPage<S::Key, R>> {
        if spec.limit == 0 {
            // Match `page_by`: a zero limit never seeks, so no `seize` guard is
            // taken for a page that cannot hold a row.
            return Ok(ChunkedPage {
                items: Vec::new(),
                next: None,
                stop: StopReason::Filled,
            });
        }
        let (start, end) = self.bounds_for_page(after);
        page_chunks(
            self.scan.scan(start, end, self.dir),
            spec,
            S::peek_row_key,
            prepare,
            project,
        )
    }
}

impl<'l, S: OwnedRows + 'l, const N: usize> Scope<'l, S, N> {
    /// Keyset pagination without a filter — infallible, matching collection `paginate`.
    pub fn page(&self, after: Option<&S::Key>, limit: usize) -> Page<S::Key, S::Val> {
        let mapped = self
            .page_by(after, limit, |row| Ok(Some(S::own(row))))
            .expect("page closure is infallible");
        Page {
            items: mapped.items,
            next: mapped.next,
        }
    }
}

pub(super) fn clone_bound<K: Clone>(b: &Bound<K>) -> Bound<K> {
    match b {
        Bound::Included(k) => Bound::Included(k.clone()),
        Bound::Excluded(k) => Bound::Excluded(k.clone()),
        Bound::Unbounded => Bound::Unbounded,
    }
}

fn from_ord<K: Key>(b: Bound<OrdKey<K>>) -> Bound<K> {
    match b {
        Bound::Included(OrdKey(k)) => Bound::Included(k),
        Bound::Excluded(OrdKey(k)) => Bound::Excluded(k),
        Bound::Unbounded => Bound::Unbounded,
    }
}

fn bound_ref<K>(b: &Bound<K>) -> Bound<&K> {
    match b {
        Bound::Included(k) => Bound::Included(k),
        Bound::Excluded(k) => Bound::Excluded(k),
        Bound::Unbounded => Bound::Unbounded,
    }
}