armdb 0.12.0

use std::hash::Hash;
use std::ops::Bound;

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

use crate::access::{Access, FallibleRowFn, KeyFn, RowFn};
use crate::armour::inspection::ScanCapability;
use crate::codec::{Codec, InspectionEncoding};
use crate::error::{DbError, DbResult};
use crate::key::Key;
use crate::{TypedMap, TypedMapShard, TypedShard, TypedTree, TypedWriteHook};

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

    const STORAGE: StorageClass = StorageClass {
        model: ValueModel::Typed,
        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> {
        Ok(self.iter().count() as u64)
    }

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

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

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

    fn contains(&self, key: &K) -> DbResult<bool> {
        Ok(self.contains(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(())
    }

    /// Values live decoded in memory, so a Tree range cannot fail per row: the
    /// strict and lenient scans are the same walk. `batch` is a Map-only knob.
    fn try_scan_range(
        &self,
        start: Bound<&K>,
        end: Bound<&K>,
        _batch: usize,
        f: RowFn<'_, K, T>,
    ) -> DbResult<()> {
        self.scan_range(start, end, f)
    }

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

    fn scan_all(&self, f: FallibleRowFn<'_, K, T>) {
        for (key, value) in self.iter() {
            if f(Ok((&key, &value))).is_break() {
                break;
            }
        }
    }

    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.clone()).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> Access for TypedMap<K, T, C, H>
where
    K: Key + Hash + Eq + 'static,
    T: Clone + Send + Sync + 'static,
    C: Codec<T> + Default + Send + Sync + 'static,
    H: TypedWriteHook<K, T> + 'static,
{
    type Key = K;
    type Val = T;
    type Shard<'a>
        = TypedMapShard<'a, K, T, C, H>
    where
        Self: 'a;

    const STORAGE: StorageClass = StorageClass {
        model: ValueModel::Typed,
        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 {
        TypedMap::inspection_shard_count(self)
    }

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

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

    fn contains(&self, key: &K) -> DbResult<bool> {
        Ok(self.contains(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);
        }
        // Guard held for the whole walk: the merge iterator yields raw
        // `TypedData` pointers, same as the live `ordered_page` handler.
        let _guard = self.inspection_guard();
        for row in self.inspection_range_bounds(start, end, batch) {
            let (key, ptr) = row?;
            let data = unsafe { &*ptr };
            if f(&key, &data.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>) {
        let _ = self.for_each_while(|key, value| f(Ok((key, value))));
    }

    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.clone()).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(all(test, feature = "rapira-codec"))]
mod tests {
    use std::ops::{Bound, ControlFlow};

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

    use crate::access::{Access, CodecWire, Wire};
    use crate::codec::RapiraCodec;
    use crate::error::DbError;
    use crate::{Config, TypedMap, TypedTree};

    #[test]
    fn typed_tree_access_point_read_and_entry_len() {
        let dir = tempfile::tempdir().expect("tmp");
        let tree =
            TypedTree::<[u8; 8], u64, RapiraCodec>::open(dir.path(), Config::test(), RapiraCodec)
                .expect("open");
        let key = 7u64.to_be_bytes();
        let value = 42u64;
        tree.put(&key, value).expect("put");

        assert_eq!(
            Access::with_value(&tree, &key, |stored| *stored).expect("read"),
            Some(value)
        );

        let wire = CodecWire(RapiraCodec);
        let encoded_size = Wire::<u64>::encode(&wire, &value).expect("encode").len() as u32;
        assert_eq!(
            Access::entry_len(&tree, &key, &wire).expect("entry len"),
            Some(encoded_size)
        );
        assert_eq!(
            <TypedTree<[u8; 8], u64, RapiraCodec> as Access>::STORAGE,
            StorageClass {
                model: ValueModel::Typed,
                shape: CollectionShape::Tree,
                backend: StorageBackend::Bitcask,
            }
        );
        assert_eq!(Access::rpc_exact_count(&tree).expect("RPC count"), 1);
        assert_eq!(
            Access::inspection_exact_count(&tree).expect("inspection count"),
            1
        );
    }

    fn open_map(iterable: bool) -> (tempfile::TempDir, TypedMap<[u8; 8], u64, RapiraCodec>) {
        let dir = tempfile::tempdir().expect("tmp");
        let cfg = Config::balanced()
            .shard_count(2)
            .hints(true)
            .iterable(iterable)
            .reversed(false)
            .build();
        let map = TypedMap::open(dir.path(), cfg, RapiraCodec).expect("open");
        for id in 1..=5_u64 {
            map.put(&id.to_be_bytes(), id * 10).expect("put");
        }
        (dir, map)
    }

    /// The Inspector's ordered-page source: the strict merge iterator, driven by
    /// the caller's per-shard batch. Every batch size must produce the same
    /// ordered rows; a non-iterable Map answers `NotImplemented` (RPC `501`).
    #[test]
    fn typed_map_try_scan_range_uses_inspection_path_with_caller_batch() {
        let (_dir, map) = open_map(true);

        for batch in [1usize, 2, 64] {
            let mut rows = Vec::new();
            Access::try_scan_range(
                &map,
                Bound::Unbounded,
                Bound::Unbounded,
                batch,
                &mut |k, v| {
                    rows.push((u64::from_be_bytes(*k), *v));
                    Ok(ControlFlow::Continue(()))
                },
            )
            .expect("strict scan");
            assert_eq!(
                rows,
                vec![(1, 10), (2, 20), (3, 30), (4, 40), (5, 50)],
                "batch {batch}"
            );
        }

        // Bounds and early Break are honoured on the strict path too.
        let mut rows = Vec::new();
        Access::try_scan_range(
            &map,
            Bound::Included(&2_u64.to_be_bytes()),
            Bound::Excluded(&5_u64.to_be_bytes()),
            1,
            &mut |k, _| {
                rows.push(u64::from_be_bytes(*k));
                Ok(if rows.len() == 2 {
                    ControlFlow::Break(())
                } else {
                    ControlFlow::Continue(())
                })
            },
        )
        .expect("bounded scan");
        assert_eq!(rows, vec![2, 3]);

        let (_dir, plain) = open_map(false);
        assert!(matches!(
            Access::try_scan_range(
                &plain,
                Bound::Unbounded,
                Bound::Unbounded,
                8,
                &mut |_, _| Ok(ControlFlow::Continue(()))
            ),
            Err(DbError::NotImplemented)
        ));
    }

    /// The batch a Map hands to `inspection_range_bounds` is derived from the
    /// page limit and the shard count; the adapter reads the shard count here.
    /// A Tree ignores `batch`, so its default of 1 only has to be non-zero —
    /// `inspection_batch` debug-asserts that.
    #[test]
    fn typed_shard_count_matches_the_live_inspection_source() {
        let (_dir, map) = open_map(true);
        assert_eq!(
            Access::inspection_shard_count(&map),
            map.inspection_shard_count()
        );
        assert!(Access::inspection_shard_count(&map) >= 1);

        let dir = tempfile::tempdir().expect("tmp");
        let tree =
            TypedTree::<[u8; 8], u64, RapiraCodec>::open(dir.path(), Config::test(), RapiraCodec)
                .expect("open");
        assert_eq!(Access::inspection_shard_count(&tree), 1);
    }
}