armdb 0.12.0

use std::hash::Hash;
use std::ops::{Bound, ControlFlow};

use armour_core::{CollectionShape, StorageBackend, StorageClass, Typ, ValueModel};

use crate::access::{Access, FallibleRowFn, KeyFn, RowFn, Wire};
use crate::armour::inspection::ScanCapability;
use crate::codec::{Codec, InspectionEncoding};
use crate::error::{DbError, DbResult};
use crate::key::Key;
use crate::{
    CollectionMeta, TypedWriteHook, VarTypedMap, VarTypedMapShard, VarTypedShard, VarTypedTree,
};

impl<K, T, C, H> VarTypedTree<K, T, C, H>
where
    K: Key + Ord + 'static,
    T: CollectionMeta + Send + Sync + 'static,
    C: Codec<T> + Clone + Send + Sync + 'static,
    H: TypedWriteHook<K, T> + 'static,
{
    /// Shared traversal behind `Access::scan_all` and `Access::undecodable_rows`.
    ///
    /// `warn` is off for the counting pass: it reports the same rows as a
    /// `UndecodableRows` finding, and warning about each of them twice per
    /// validation run is noise.
    fn walk_rows(&self, warn: bool, f: FallibleRowFn<'_, K, T>) {
        let codec = self.codec();
        for (key, bytes) in self.as_inner().iter() {
            let step = match codec.decode_from(&bytes) {
                Ok(value) => f(Ok((&key, &value))),
                Err(error) => {
                    if warn {
                        tracing::warn!(
                            collection = T::NAME,
                            error = ?error,
                            "schema scan: decode error in VarTypedTree, skipping entry — relations not checked"
                        );
                    }
                    f(Err(error))
                }
            };
            if step.is_break() {
                break;
            }
        }
    }
}

impl<K, T, C, H> Access for VarTypedTree<K, T, C, H>
where
    K: Key + Ord + 'static,
    T: CollectionMeta + Send + Sync + 'static,
    C: Codec<T> + Clone + Send + Sync + 'static,
    H: TypedWriteHook<K, T> + 'static,
{
    type Key = K;
    type Val = T;
    type Shard<'a>
        = VarTypedShard<'a, K, T, C, H>
    where
        Self: 'a;

    const STORAGE: StorageClass = StorageClass {
        model: ValueModel::VarTyped,
        shape: CollectionShape::Tree,
        backend: StorageBackend::Bitcask,
    };

    fn encoding(&self, _ty: &Typ) -> InspectionEncoding {
        C::INSPECTION_ENCODING
    }

    fn approx_len(&self) -> u64 {
        self.len() as u64
    }

    fn rpc_exact_count(&self) -> DbResult<u64> {
        Ok(self.iter().count() as u64)
    }

    fn inspection_exact_count(&self) -> DbResult<u64> {
        self.try_count_live_keys()
    }

    fn scan_capability(&self) -> ScanCapability {
        ScanCapability::Ordered
    }

    fn reversed(&self) -> bool {
        self.config().reversed
    }

    fn with_value<R>(&self, key: &K, f: impl FnOnce(&T) -> R) -> DbResult<Option<R>> {
        Ok(self.try_get(key)?.map(|value| f(&value)))
    }

    fn contains(&self, key: &K) -> DbResult<bool> {
        Ok(self.contains(key))
    }

    fn entry_len<W: Wire<T>>(&self, key: &K, _wire: &W) -> DbResult<Option<u32>> {
        Ok(self.entry_len(key))
    }

    fn with_first<R>(&self, f: impl FnOnce(&K, &T) -> R) -> DbResult<Option<R>> {
        Ok(self.first().map(|(key, value)| f(&key, &value)))
    }

    fn with_last<R>(&self, f: impl FnOnce(&K, &T) -> R) -> DbResult<Option<R>> {
        Ok(self.last().map(|(key, value)| f(&key, &value)))
    }

    fn scan_range(&self, start: Bound<&K>, end: Bound<&K>, f: RowFn<'_, K, T>) -> DbResult<()> {
        for (key, value) in self.range_bounds(start, end) {
            if f(&key, &value)?.is_break() {
                break;
            }
        }
        Ok(())
    }

    /// Ordered single-structure scan — `batch` is a Map-only knob.
    fn try_scan_range(
        &self,
        start: Bound<&K>,
        end: Bound<&K>,
        _batch: usize,
        f: RowFn<'_, K, T>,
    ) -> DbResult<()> {
        for row in self.try_range_bounds(start, end) {
            let (key, value) = row?;
            if f(&key, &value)?.is_break() {
                break;
            }
        }
        Ok(())
    }

    fn scan_keys(&self, start: Bound<&K>, end: Bound<&K>, f: KeyFn<'_, K>) -> DbResult<()> {
        // Decoding walk, matching the live lenient `range_keys`: a row whose
        // value will not decode is skipped, so its key stays off the wire.
        // The raw inner walk would newly expose corrupt-row keys — a second
        // wire-visible change, which this PR series does not carry.
        for (key, _) in self.range_bounds(start, end) {
            if f(&key)?.is_break() {
                break;
            }
        }
        Ok(())
    }

    fn scan_all(&self, f: FallibleRowFn<'_, K, T>) {
        self.walk_rows(true, f);
    }

    fn undecodable_rows(&self) -> u64 {
        let mut n = 0u64;
        self.walk_rows(false, &mut |row| {
            if row.is_err() {
                n += 1;
            }
            ControlFlow::Continue(())
        });
        n
    }

    fn effective_projection_limits(&self) -> Option<(usize, usize)> {
        Some(self.effective_projection_limits())
    }

    fn put(&self, key: &K, value: T) -> DbResult<()> {
        self.put(key, &value).map(drop)
    }

    fn put_ref(&self, key: &K, value: &T) -> DbResult<()> {
        self.put(key, value).map(drop)
    }

    fn delete(&self, key: &K) -> DbResult<()> {
        self.delete(key).map(drop)
    }

    fn atomic_at<R>(
        &self,
        key: &K,
        f: impl FnOnce(&mut Self::Shard<'_>) -> DbResult<R>,
    ) -> DbResult<R> {
        self.atomic(key, f)
    }
}

impl<K, T, C, H> VarTypedMap<K, T, C, H>
where
    K: Key + Send + Sync + Hash + Eq + 'static,
    T: CollectionMeta + Send + Sync + 'static,
    C: Codec<T> + Clone + Send + Sync + 'static,
    H: TypedWriteHook<K, T> + 'static,
{
    /// Shared traversal behind `Access::scan_all` and `Access::undecodable_rows`.
    ///
    /// Runs on `VarMap::for_each_while`, so a `Break` ends the walk instead of
    /// silently reading every remaining value off disk. `warn` is off for the
    /// counting pass — same rationale as the `VarTypedTree` walk above.
    fn walk_rows(&self, warn: bool, f: FallibleRowFn<'_, K, T>) {
        let codec = self.codec();
        let report = |error: DbError, f: FallibleRowFn<'_, K, T>| {
            if warn {
                // Caller-agnostic wording: `scan_all` backs both the schema
                // validator (skips this row, reports `UndecodableRows`) and the
                // Inspection `unordered_snapshot` (aborts on the first `Err`
                // row) — this message must be true for both outcomes.
                tracing::warn!(
                    collection = T::NAME,
                    error = ?error,
                    "unreadable row in VarTypedMap"
                );
            }
            f(Err(error))
        };
        self.as_inner().for_each_while(|row| match row {
            Ok((key, bytes)) => match codec.decode_from(&bytes) {
                Ok(value) => f(Ok((&key, &value))),
                Err(error) => report(error, f),
            },
            // A value the storage layer could not read at all: same contract as
            // a decode failure — the row's relations are not checked.
            Err(error) => report(error, f),
        });
    }
}

impl<K, T, C, H> Access for VarTypedMap<K, T, C, H>
where
    K: Key + Send + Sync + Hash + Eq + 'static,
    T: CollectionMeta + Send + Sync + 'static,
    C: Codec<T> + Clone + Send + Sync + 'static,
    H: TypedWriteHook<K, T> + 'static,
{
    type Key = K;
    type Val = T;
    type Shard<'a>
        = VarTypedMapShard<'a, K, T, C, H>
    where
        Self: 'a;

    const STORAGE: StorageClass = StorageClass {
        model: ValueModel::VarTyped,
        shape: CollectionShape::Map,
        backend: StorageBackend::Bitcask,
    };

    fn encoding(&self, _ty: &Typ) -> InspectionEncoding {
        C::INSPECTION_ENCODING
    }

    fn approx_len(&self) -> u64 {
        self.len() as u64
    }

    fn rpc_exact_count(&self) -> DbResult<u64> {
        Ok(self.len() as u64)
    }

    fn inspection_exact_count(&self) -> DbResult<u64> {
        Ok(self.len() as u64)
    }

    fn scan_capability(&self) -> ScanCapability {
        if self.iter_view().is_some() {
            ScanCapability::Ordered
        } else {
            ScanCapability::UnorderedSnapshot
        }
    }

    fn inspection_shard_count(&self) -> usize {
        VarTypedMap::inspection_shard_count(self)
    }

    fn reversed(&self) -> bool {
        self.config().reversed
    }

    fn with_value<R>(&self, key: &K, f: impl FnOnce(&T) -> R) -> DbResult<Option<R>> {
        Ok(self.try_get(key)?.map(|value| f(&value)))
    }

    fn contains(&self, key: &K) -> DbResult<bool> {
        Ok(self.contains(key))
    }

    fn entry_len<W: Wire<T>>(&self, key: &K, _wire: &W) -> DbResult<Option<u32>> {
        Ok(self.entry_len(key))
    }

    fn with_first<R>(&self, f: impl FnOnce(&K, &T) -> R) -> DbResult<Option<R>> {
        Ok(self.first()?.map(|(key, value)| f(&key, &value)))
    }

    fn with_last<R>(&self, f: impl FnOnce(&K, &T) -> R) -> DbResult<Option<R>> {
        Ok(self.last()?.map(|(key, value)| f(&key, &value)))
    }

    fn scan_range(&self, start: Bound<&K>, end: Bound<&K>, f: RowFn<'_, K, T>) -> DbResult<()> {
        let view = self.iter_view().ok_or(DbError::NotImplemented)?;
        for (key, value) in view.range_bounds(start, end) {
            if f(&key, &value)?.is_break() {
                break;
            }
        }
        Ok(())
    }

    /// The Inspector's ordered-page path: the strict `inspection_range_bounds`
    /// merge iterator, fed the caller's per-shard `batch`.
    fn try_scan_range(
        &self,
        start: Bound<&K>,
        end: Bound<&K>,
        batch: usize,
        f: RowFn<'_, K, T>,
    ) -> DbResult<()> {
        if self.iter_view().is_none() {
            return Err(DbError::NotImplemented);
        }
        for row in self.inspection_range_bounds(start, end, batch) {
            let (key, value) = row?;
            if f(&key, &value)?.is_break() {
                break;
            }
        }
        Ok(())
    }

    fn scan_keys(&self, start: Bound<&K>, end: Bound<&K>, f: KeyFn<'_, K>) -> DbResult<()> {
        let view = self.iter_view().ok_or(DbError::NotImplemented)?;
        for key in view.keys_range_bounds(start, end) {
            if f(&key)?.is_break() {
                break;
            }
        }
        Ok(())
    }

    fn scan_all(&self, f: FallibleRowFn<'_, K, T>) {
        self.walk_rows(true, f);
    }

    fn undecodable_rows(&self) -> u64 {
        let mut n = 0u64;
        self.walk_rows(false, &mut |row| {
            if row.is_err() {
                n += 1;
            }
            ControlFlow::Continue(())
        });
        n
    }

    fn effective_projection_limits(&self) -> Option<(usize, usize)> {
        Some(self.effective_projection_limits())
    }

    fn put(&self, key: &K, value: T) -> DbResult<()> {
        self.put(key, &value).map(drop)
    }

    fn put_ref(&self, key: &K, value: &T) -> DbResult<()> {
        self.put(key, value).map(drop)
    }

    fn delete(&self, key: &K) -> DbResult<()> {
        self.delete(key).map(drop)
    }

    fn atomic_at<R>(
        &self,
        key: &K,
        f: impl FnOnce(&mut Self::Shard<'_>) -> DbResult<R>,
    ) -> DbResult<R> {
        self.atomic(key, f)
    }
}

#[cfg(test)]
#[cfg(feature = "rapira-codec")]
mod tests {
    use std::ops::{Bound, ControlFlow};

    use armour_core::{GetRefs, GetType};
    use rapira::Rapira;

    use crate::access::Access;
    use crate::codec::RapiraCodec;
    use crate::{CollectionMeta, Config, VarTypedMap, VarTypedTree};

    #[derive(Clone, Copy, Debug, PartialEq, Eq, GetType, GetRefs, Rapira)]
    struct Row {
        value: u64,
    }

    impl CollectionMeta for Row {
        type SelfId = [u8; 8];
        const NAME: &'static str = "access_var_typed";
    }

    fn corrupt_tree() -> (tempfile::TempDir, VarTypedTree<[u8; 8], Row, RapiraCodec>) {
        let dir = tempfile::tempdir().expect("tmp");
        let tree = VarTypedTree::open(dir.path(), Config::test(), RapiraCodec).expect("open");
        tree.put(&1_u64.to_be_bytes(), &Row { value: 1 })
            .expect("put good row");
        // Raw garbage — readable as a live key, undecodable as Row.
        tree.as_inner()
            .put(&2_u64.to_be_bytes(), &[0xff, 0xfe, 0xfd])
            .expect("put poison bytes");
        tree.put(&3_u64.to_be_bytes(), &Row { value: 3 })
            .expect("put good row");
        (dir, tree)
    }

    #[test]
    fn var_typed_tree_exact_counts_disagree_on_corrupt_row() {
        let (_dir, tree) = corrupt_tree();

        assert_eq!(Access::rpc_exact_count(&tree).expect("RPC count"), 2);
        assert_eq!(
            Access::inspection_exact_count(&tree).expect("inspection count"),
            3
        );
    }

    /// `scan_keys` (the RPC `range_keys` path) must skip an undecodable row's
    /// key exactly like `scan_range` skips its whole row — walking the raw
    /// inner key range instead would newly expose corrupt-row keys on the
    /// wire, a second wire-visible change PR 3c does not carry (review round 1).
    #[test]
    fn var_typed_tree_scan_keys_skips_undecodable_rows() {
        let (_dir, tree) = corrupt_tree();
        let mut keys = Vec::new();

        Access::scan_keys(&tree, Bound::Unbounded, Bound::Unbounded, &mut |k| {
            keys.push(u64::from_be_bytes(*k));
            Ok(ControlFlow::Continue(()))
        })
        .expect("scan_keys");

        // `corrupt_tree` opens with `Config::test()`, whose default `reversed:
        // true` walks descending — same order `scan_range`'s lenient walk uses.
        assert_eq!(
            keys,
            vec![3, 1],
            "the undecodable row's key must stay off the wire, matching the lenient scan_range"
        );
    }

    #[test]
    fn var_typed_tree_try_scan_range_fails_where_scan_range_skips() {
        let (_dir, tree) = corrupt_tree();
        let mut lenient_rows = 0;

        Access::scan_range(&tree, Bound::Unbounded, Bound::Unbounded, &mut |_, _| {
            lenient_rows += 1;
            Ok(ControlFlow::Continue(()))
        })
        .expect("lenient scan");

        assert_eq!(lenient_rows, 2);
        assert!(
            Access::try_scan_range(
                &tree,
                Bound::Unbounded,
                Bound::Unbounded,
                8,
                &mut |_, _| Ok(ControlFlow::Continue(())),
            )
            .is_err()
        );
    }

    fn corrupt_map() -> (tempfile::TempDir, VarTypedMap<[u8; 8], Row, RapiraCodec>) {
        let dir = tempfile::tempdir().expect("tmp");
        let cfg = Config::balanced()
            .shard_count(2)
            .hints(true)
            .iterable(true)
            .reversed(false)
            .build();
        let map = VarTypedMap::open(dir.path(), cfg, RapiraCodec).expect("open");
        for id in 1..=5_u64 {
            map.put(&id.to_be_bytes(), &Row { value: id })
                .expect("put good row");
        }
        // Raw garbage — readable as a live key, undecodable as Row.
        map.as_inner()
            .put(&6_u64.to_be_bytes(), &[0xff, 0xfe, 0xfd])
            .expect("put poison bytes");
        (dir, map)
    }

    /// The strict Map scan is the Inspector's `ordered_page` source: it must
    /// surface an undecodable row instead of shortening the page, and it must
    /// honour the per-shard batch the caller derived from the page limit.
    #[test]
    fn var_typed_map_try_scan_range_is_strict_and_batch_driven() {
        let (_dir, map) = corrupt_map();

        // Lenient RPC policy: the corrupt row is skipped, the scan succeeds.
        let mut lenient_rows = 0;
        Access::scan_range(&map, Bound::Unbounded, Bound::Unbounded, &mut |_, _| {
            lenient_rows += 1;
            Ok(ControlFlow::Continue(()))
        })
        .expect("lenient scan");
        assert_eq!(lenient_rows, 5);

        // Strict Inspector policy: the same row fails the scan.
        for batch in [1usize, 64] {
            assert!(
                Access::try_scan_range(
                    &map,
                    Bound::Unbounded,
                    Bound::Unbounded,
                    batch,
                    &mut |_, _| Ok(ControlFlow::Continue(())),
                )
                .is_err(),
                "batch {batch} must still propagate the undecodable row"
            );
        }

        // Every batch size yields the same ordered prefix before the bad row.
        for batch in [1usize, 2, 64] {
            let mut keys = Vec::new();
            let bad = Access::try_scan_range(
                &map,
                Bound::Unbounded,
                Bound::Included(&5_u64.to_be_bytes()),
                batch,
                &mut |k, _| {
                    keys.push(u64::from_be_bytes(*k));
                    Ok(ControlFlow::Continue(()))
                },
            );
            assert!(bad.is_ok(), "batch {batch}");
            assert_eq!(keys, vec![1, 2, 3, 4, 5], "batch {batch}");
        }
    }

    /// `Break` must end the walk, not just mute the visitor: a bounded consumer
    /// stays O(rows consumed), not O(collection size). Counted at the storage
    /// seam, since a suppressed visitor looks identical from the outside.
    #[test]
    fn var_typed_map_scan_all_break_stops_storage_reads() {
        let (_dir, map) = corrupt_map();

        let before = crate::var_map::test_value_reads();
        let mut visited = 0;
        Access::scan_all(&map, &mut |_row| {
            visited += 1;
            ControlFlow::Break(())
        });
        let reads = crate::var_map::test_value_reads() - before;

        assert_eq!(visited, 1);
        assert_eq!(reads, 1, "no value may be read after the visitor broke");

        // A full pass reads every live row and reports the corrupt one.
        let before = crate::var_map::test_value_reads();
        let mut rows = 0;
        let mut bad = 0;
        Access::scan_all(&map, &mut |row| {
            rows += 1;
            if row.is_err() {
                bad += 1;
            }
            ControlFlow::Continue(())
        });
        assert_eq!((rows, bad), (6, 1));
        assert_eq!(crate::var_map::test_value_reads() - before, 6);
        assert_eq!(Access::undecodable_rows(&map), 1);
    }
}