armdb 0.11.0

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

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

use crate::access::{Access, FallibleRowFn, KeyFn, RowFn, Wire};
use crate::armour::inspection::ScanCapability;
use crate::armour::typ_width::zero_inspection_encoding;
use crate::codec::InspectionEncoding;
use crate::durability::{Bitcask, Durability, Fixed};
use crate::error::{DbError, DbResult};
use crate::hook::{TypedWriteHook, ZeroHookAdapter};
use crate::key::Key;
use crate::{ZeroMap, ZeroMapShard, ZeroShard, ZeroTree};

pub(crate) trait DurBackend: Durability {
    const B: StorageBackend;
}

impl DurBackend for Bitcask {
    const B: StorageBackend = StorageBackend::Bitcask;
}

impl DurBackend for Fixed {
    const B: StorageBackend = StorageBackend::Fixed;
}

impl<K, const V: usize, T, H, D> Access for ZeroTree<K, V, T, H, D>
where
    K: Key + Ord + 'static,
    T: Copy
        + zerocopy::IntoBytes
        + zerocopy::FromBytes
        + zerocopy::Immutable
        + Send
        + Sync
        + 'static,
    H: TypedWriteHook<K, T> + 'static,
    D: DurBackend + 'static,
{
    type Key = K;
    type Val = T;
    type Shard<'a>
        = ZeroShard<'a, K, V, T, D, ZeroHookAdapter<K, T, H>>
    where
        Self: 'a;

    const STORAGE: StorageClass = StorageClass {
        model: ValueModel::Zero,
        shape: CollectionShape::Tree,
        backend: D::B,
    };

    fn encoding(&self, ty: &Typ) -> InspectionEncoding {
        zero_inspection_encoding(ty, V)
    }

    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 entry_len<W: Wire<T>>(&self, key: &K, _wire: &W) -> DbResult<Option<u32>> {
        Ok(if self.contains(key) {
            Some(V as u32)
        } else {
            None
        })
    }

    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(())
    }

    /// Const-width 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).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, const V: usize, T, H, D> Access for ZeroMap<K, V, T, H, D>
where
    K: Key + Send + Sync + Hash + Eq + 'static,
    T: Copy
        + zerocopy::IntoBytes
        + zerocopy::FromBytes
        + zerocopy::Immutable
        + Send
        + Sync
        + 'static,
    H: TypedWriteHook<K, T> + 'static,
    D: DurBackend + 'static,
{
    type Key = K;
    type Val = T;
    type Shard<'a>
        = ZeroMapShard<'a, K, V, T, D, ZeroHookAdapter<K, T, H>>
    where
        Self: 'a;

    const STORAGE: StorageClass = StorageClass {
        model: ValueModel::Zero,
        shape: CollectionShape::Map,
        backend: D::B,
    };

    fn encoding(&self, ty: &Typ) -> InspectionEncoding {
        zero_inspection_encoding(ty, V)
    }

    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 {
        ZeroMap::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 entry_len<W: Wire<T>>(&self, key: &K, _wire: &W) -> DbResult<Option<u32>> {
        Ok(if self.contains(key) {
            Some(V as u32)
        } else {
            None
        })
    }

    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>) {
        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).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)]
mod tests {
    use std::ops::{Bound, ControlFlow};

    use armour_core::{ScalarTyp, Typ};

    use crate::access::{Access, Wire};
    use crate::armour::typ_width::zero_inspection_encoding;
    use crate::error::{DbError, DbResult};
    use crate::{Config, ZeroMap, ZeroTree};

    struct PanicSizeWire;

    impl Wire<u32> for PanicSizeWire {
        fn encode_to(&self, _: &u32, _: &mut Vec<u8>) -> DbResult<()> {
            unreachable!()
        }

        fn decode(&self, _: &[u8]) -> DbResult<u32> {
            unreachable!()
        }

        fn size(&self, _: &u32) -> DbResult<u32> {
            panic!("entry_len must not call wire.size")
        }
    }

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

        assert_eq!(
            Access::entry_len(&tree, &key, &PanicSizeWire).expect("entry len"),
            Some(4)
        );
        let ty = Typ::Scalar(ScalarTyp::U32);
        assert_eq!(
            Access::encoding(&tree, &ty),
            zero_inspection_encoding(&ty, 4)
        );
    }

    fn open_map(iterable: bool) -> (tempfile::TempDir, ZeroMap<[u8; 8], 4, u32>) {
        let dir = tempfile::tempdir().expect("tmp");
        let cfg = Config::balanced()
            .shard_count(2)
            .hints(true)
            .iterable(iterable)
            .reversed(false)
            .build();
        let map = <ZeroMap<[u8; 8], 4, u32>>::open(dir.path(), cfg).expect("open");
        for id in 1..=5_u32 {
            map.put(&u64::from(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 zero_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}"
            );
        }

        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)
        ));
    }

    #[test]
    fn zero_shard_count_matches_the_live_inspection_source() {
        let (_dir, map) = open_map(true);
        assert_eq!(
            Access::inspection_shard_count(&map),
            map.inspection_shard_count()
        );

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