armdb 0.11.2

//! Upsert semantics across Db-openable collection families.
//!
//! Keep observables by backend:
//! - **Bitcask (Zero/Typed/VarTyped)**: hook count unchanged, value unchanged, `len()` unchanged.
//! - **FixedStore (Zero)**: value equality after Keep (no hook on Fixed).
//! - **LSM**: hook fires synchronously under the held stripe lock (before `atomic` returns).

#![cfg(all(feature = "typed-tree", feature = "rapira-codec"))]

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};

use armdb::codec::Codec;
use armdb::{
    Config, DbError, FixedConfig, RapiraCodec, TypedMap, TypedTree, TypedWriteHook, Upsert,
    ZeroMap, ZeroTree,
};
use rapira::Rapira;
use tempfile::tempdir;
use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout};

#[cfg(feature = "var-collections")]
use armdb::{VarTypedMap, VarTypedTree};

#[cfg(feature = "fjall")]
use armdb::lsm::{LsmConfig, LsmTree};

#[derive(
    FromBytes, IntoBytes, Immutable, KnownLayout, Clone, Copy, Debug, PartialEq, Eq, Rapira,
)]
#[repr(C)]
struct Counter {
    n: u64,
}

#[derive(Default, Clone)]
struct CountHook {
    writes: Arc<AtomicUsize>,
    /// `(old, new)` per fired hook, in the order the engine replayed them.
    log: Arc<Mutex<Vec<(Option<u64>, Option<u64>)>>>,
}

impl TypedWriteHook<[u8; 8], Counter> for CountHook {
    const NEEDS_WRITE: bool = true;
    fn on_write(&self, _k: &[u8; 8], old: Option<&Counter>, new: Option<&Counter>) {
        self.writes.fetch_add(1, Ordering::Relaxed);
        self.log
            .lock()
            .expect("hook log")
            .push((old.map(|c| c.n), new.map(|c| c.n)));
    }
}

struct U64Codec;

impl Codec<u64> for U64Codec {
    fn encode_to(&self, value: &u64, buf: &mut Vec<u8>) -> armdb::DbResult<()> {
        buf.clear();
        buf.extend_from_slice(&value.to_be_bytes());
        Ok(())
    }

    fn decode_from(&self, bytes: &[u8]) -> armdb::DbResult<u64> {
        let arr: [u8; 8] = bytes
            .try_into()
            .map_err(|_| DbError::CorruptedEntry { offset: 0 })?;
        Ok(u64::from_be_bytes(arr))
    }
}

impl CountHook {
    fn new() -> Self {
        Self {
            writes: Arc::new(AtomicUsize::new(0)),
            log: Arc::new(Mutex::new(Vec::new())),
        }
    }

    fn count(&self) -> usize {
        self.writes.load(Ordering::Relaxed)
    }

    fn log(&self) -> Vec<(Option<u64>, Option<u64>)> {
        self.log.lock().expect("hook log").clone()
    }
}

fn key(n: u64) -> [u8; 8] {
    n.to_be_bytes()
}

/// Five-cell semantics table for Bitcask collections with a counting hook.
fn assert_bitcask_counter_semantics(
    coll: &impl BitcaskCounterHarness,
    hook: &CountHook,
    k: &[u8; 8],
) {
    // absent + Put
    coll.upsert(k, |_| Upsert::Put(Counter { n: 1 })).unwrap();
    assert_eq!(coll.get(k).unwrap().n, 1);
    assert_eq!(hook.count(), 1);

    // present + Put
    coll.upsert(k, |c| {
        Upsert::Put(Counter {
            n: c.unwrap().n + 1,
        })
    })
    .unwrap();
    assert_eq!(coll.get(k).unwrap().n, 2);
    assert_eq!(hook.count(), 2);

    // Keep — hook silence + value unchanged + len unchanged
    let len_before = coll.len();
    let writes_before = hook.count();
    coll.upsert(k, |_| Upsert::Keep).unwrap();
    assert_eq!(coll.get(k).unwrap().n, 2);
    assert_eq!(hook.count(), writes_before);
    assert_eq!(coll.len(), len_before);

    // present + Delete
    coll.upsert(k, |_| Upsert::Delete).unwrap();
    assert!(coll.get(k).is_none());
    assert_eq!(hook.count(), writes_before + 1);

    // absent + Delete — no-op, no hook
    let w = hook.count();
    coll.upsert(k, |_| Upsert::Delete).unwrap();
    assert_eq!(hook.count(), w);
}

trait BitcaskCounterHarness {
    fn get(&self, k: &[u8; 8]) -> Option<Counter>;
    fn len(&self) -> usize;
    fn upsert(
        &self,
        k: &[u8; 8],
        f: impl FnOnce(Option<&Counter>) -> Upsert<Counter>,
    ) -> armdb::DbResult<()>;
}

impl<H: TypedWriteHook<[u8; 8], Counter>, D: armdb::durability::Durability> BitcaskCounterHarness
    for ZeroTree<[u8; 8], 8, Counter, H, D>
{
    fn get(&self, k: &[u8; 8]) -> Option<Counter> {
        ZeroTree::get(self, k)
    }
    fn len(&self) -> usize {
        ZeroTree::len(self)
    }
    fn upsert(
        &self,
        k: &[u8; 8],
        f: impl FnOnce(Option<&Counter>) -> Upsert<Counter>,
    ) -> armdb::DbResult<()> {
        ZeroTree::upsert(self, k, f)
    }
}

// `ZeroMap` is gated on `typed-tree` (see `lib.rs`), which the file-level
// `#![cfg]` already requires — no extra feature gate here.
impl<H: TypedWriteHook<[u8; 8], Counter>, D: armdb::durability::Durability> BitcaskCounterHarness
    for ZeroMap<[u8; 8], 8, Counter, H, D>
{
    fn get(&self, k: &[u8; 8]) -> Option<Counter> {
        ZeroMap::get(self, k)
    }
    fn len(&self) -> usize {
        ZeroMap::len(self)
    }
    fn upsert(
        &self,
        k: &[u8; 8],
        f: impl FnOnce(Option<&Counter>) -> Upsert<Counter>,
    ) -> armdb::DbResult<()> {
        ZeroMap::upsert(self, k, f)
    }
}

/// Five-cell semantics table for Typed* collections using u64 + U64Codec.
fn assert_bitcask_u64_semantics(coll: &impl BitcaskU64Harness, hook: &CountHookU64, k: &[u8; 8]) {
    coll.upsert(k, |_| Upsert::Put(1)).unwrap();
    assert_eq!(coll.get(k), Some(1));
    assert_eq!(hook.count(), 1);

    coll.upsert(k, |c| Upsert::Put(c.unwrap() + 1)).unwrap();
    assert_eq!(coll.get(k), Some(2));
    assert_eq!(hook.count(), 2);

    let len_before = coll.len();
    let writes_before = hook.count();
    coll.upsert(k, |_| Upsert::Keep).unwrap();
    assert_eq!(coll.get(k), Some(2));
    assert_eq!(hook.count(), writes_before);
    assert_eq!(coll.len(), len_before);

    coll.upsert(k, |_| Upsert::Delete).unwrap();
    assert!(coll.get(k).is_none());
    assert_eq!(hook.count(), writes_before + 1);

    let w = hook.count();
    coll.upsert(k, |_| Upsert::Delete).unwrap();
    assert_eq!(hook.count(), w);
}

#[derive(Default, Clone)]
struct CountHookU64 {
    writes: Arc<AtomicUsize>,
}

impl TypedWriteHook<[u8; 8], u64> for CountHookU64 {
    const NEEDS_WRITE: bool = true;
    fn on_write(&self, _k: &[u8; 8], _old: Option<&u64>, _new: Option<&u64>) {
        self.writes.fetch_add(1, Ordering::Relaxed);
    }
}

impl CountHookU64 {
    fn new() -> Self {
        Self {
            writes: Arc::new(AtomicUsize::new(0)),
        }
    }
    fn count(&self) -> usize {
        self.writes.load(Ordering::Relaxed)
    }
}

trait BitcaskU64Harness {
    fn get(&self, k: &[u8; 8]) -> Option<u64>;
    fn len(&self) -> usize;
    fn upsert(
        &self,
        k: &[u8; 8],
        f: impl FnOnce(Option<&u64>) -> Upsert<u64>,
    ) -> armdb::DbResult<()>;
}

impl<H: TypedWriteHook<[u8; 8], u64>> BitcaskU64Harness for TypedTree<[u8; 8], u64, U64Codec, H> {
    fn get(&self, k: &[u8; 8]) -> Option<u64> {
        TypedTree::get(self, k).map(|r| *r)
    }
    fn len(&self) -> usize {
        TypedTree::len(self)
    }
    fn upsert(
        &self,
        k: &[u8; 8],
        f: impl FnOnce(Option<&u64>) -> Upsert<u64>,
    ) -> armdb::DbResult<()> {
        TypedTree::upsert(self, k, f)
    }
}

impl<H: TypedWriteHook<[u8; 8], u64>> BitcaskU64Harness for TypedMap<[u8; 8], u64, U64Codec, H> {
    fn get(&self, k: &[u8; 8]) -> Option<u64> {
        TypedMap::get(self, k).map(|r| *r)
    }
    fn len(&self) -> usize {
        TypedMap::len(self)
    }
    fn upsert(
        &self,
        k: &[u8; 8],
        f: impl FnOnce(Option<&u64>) -> Upsert<u64>,
    ) -> armdb::DbResult<()> {
        TypedMap::upsert(self, k, f)
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Rapira)]
struct Item {
    n: u64,
}

#[cfg(feature = "var-collections")]
trait BitcaskItemHarness {
    fn get(&self, k: &[u8; 8]) -> Option<Item>;
    fn len(&self) -> usize;
    fn upsert(
        &self,
        k: &[u8; 8],
        f: impl FnOnce(Option<&Item>) -> Upsert<Item>,
    ) -> armdb::DbResult<()>;
}

#[derive(Default, Clone)]
struct CountHookItem {
    writes: Arc<AtomicUsize>,
}

#[cfg(feature = "var-collections")]
impl TypedWriteHook<[u8; 8], Item> for CountHookItem {
    const NEEDS_WRITE: bool = true;
    fn on_write(&self, _k: &[u8; 8], _old: Option<&Item>, _new: Option<&Item>) {
        self.writes.fetch_add(1, Ordering::Relaxed);
    }
}

#[cfg(feature = "var-collections")]
impl CountHookItem {
    fn new() -> Self {
        Self {
            writes: Arc::new(AtomicUsize::new(0)),
        }
    }
    fn count(&self) -> usize {
        self.writes.load(Ordering::Relaxed)
    }
}

#[cfg(feature = "var-collections")]
impl<H: TypedWriteHook<[u8; 8], Item>> BitcaskItemHarness
    for VarTypedMap<[u8; 8], Item, RapiraCodec, H>
{
    fn get(&self, k: &[u8; 8]) -> Option<Item> {
        VarTypedMap::get(self, k)
    }
    fn len(&self) -> usize {
        VarTypedMap::len(self)
    }
    fn upsert(
        &self,
        k: &[u8; 8],
        f: impl FnOnce(Option<&Item>) -> Upsert<Item>,
    ) -> armdb::DbResult<()> {
        VarTypedMap::upsert(self, k, f)
    }
}

#[cfg(feature = "var-collections")]
impl<H: TypedWriteHook<[u8; 8], Item>> BitcaskItemHarness
    for VarTypedTree<[u8; 8], Item, RapiraCodec, H>
{
    fn get(&self, k: &[u8; 8]) -> Option<Item> {
        VarTypedTree::get(self, k)
    }
    fn len(&self) -> usize {
        VarTypedTree::len(self)
    }
    fn upsert(
        &self,
        k: &[u8; 8],
        f: impl FnOnce(Option<&Item>) -> Upsert<Item>,
    ) -> armdb::DbResult<()> {
        VarTypedTree::upsert(self, k, f)
    }
}

#[cfg(feature = "var-collections")]
fn assert_bitcask_item_semantics(
    coll: &impl BitcaskItemHarness,
    hook: &CountHookItem,
    k: &[u8; 8],
) {
    coll.upsert(k, |_| Upsert::Put(Item { n: 1 })).unwrap();
    assert_eq!(coll.get(k).unwrap().n, 1);
    assert_eq!(hook.count(), 1);

    coll.upsert(k, |c| {
        Upsert::Put(Item {
            n: c.unwrap().n + 1,
        })
    })
    .unwrap();
    assert_eq!(coll.get(k).unwrap().n, 2);
    assert_eq!(hook.count(), 2);

    let len_before = coll.len();
    let writes_before = hook.count();
    coll.upsert(k, |_| Upsert::Keep).unwrap();
    assert_eq!(coll.get(k).unwrap().n, 2);
    assert_eq!(hook.count(), writes_before);
    assert_eq!(coll.len(), len_before);

    coll.upsert(k, |_| Upsert::Delete).unwrap();
    assert!(coll.get(k).is_none());
    assert_eq!(hook.count(), writes_before + 1);

    let w = hook.count();
    coll.upsert(k, |_| Upsert::Delete).unwrap();
    assert_eq!(hook.count(), w);
}

#[test]
fn zero_tree_bitcask_semantics_table() {
    let dir = tempdir().unwrap();
    let hook = CountHook::new();
    let tree =
        ZeroTree::<[u8; 8], 8, Counter, _>::open_hooked(dir.path(), Config::test(), hook.clone())
            .unwrap();
    assert_bitcask_counter_semantics(&tree, &hook, &key(1));
}

#[test]
fn zero_tree_fixed_keep_leaves_value() {
    let dir = tempdir().unwrap();
    let hook = CountHook::new();
    let tree =
        ZeroTree::<[u8; 8], 8, Counter, CountHook, armdb::durability::Fixed>::open_with_hook(
            dir.path(),
            FixedConfig::test(),
            hook.clone(),
        )
        .unwrap();
    let k = key(3);
    tree.upsert(&k, |_| Upsert::Put(Counter { n: 4 })).unwrap();
    let before = tree.get(&k).unwrap();
    assert_eq!(hook.count(), 1);

    // On FixedStore an in-place rewrite of the identical value is invisible to a
    // value comparison — it would still bump the slot version and emit a
    // replication event. The hook is the observable that catches it: `Keep` must
    // perform no write at all.
    tree.upsert(&k, |_| Upsert::Keep).unwrap();
    assert_eq!(tree.get(&k).unwrap(), before);
    assert_eq!(hook.count(), 1, "Keep must not write the slot");
    assert_eq!(hook.log(), vec![(None, Some(4))]);
}

#[test]
fn zero_map_bitcask_semantics_table() {
    let dir = tempdir().unwrap();
    let hook = CountHook::new();
    let map =
        ZeroMap::<[u8; 8], 8, Counter, _>::open_hooked(dir.path(), Config::test(), hook.clone())
            .unwrap();
    assert_bitcask_counter_semantics(&map, &hook, &key(1));
}

#[test]
fn zero_map_fixed_keep_leaves_value() {
    let dir = tempdir().unwrap();
    let hook = CountHook::new();
    let map = ZeroMap::<[u8; 8], 8, Counter, CountHook, armdb::durability::Fixed>::open_with_hook(
        dir.path(),
        FixedConfig::test(),
        hook.clone(),
    )
    .unwrap();
    let k = key(5);
    map.upsert(&k, |_| Upsert::Put(Counter { n: 7 })).unwrap();
    let before = map.get(&k).unwrap();
    assert_eq!(hook.count(), 1);

    // Same FixedStore observable as the Tree case — see the comment there.
    map.upsert(&k, |_| Upsert::Keep).unwrap();
    assert_eq!(map.get(&k).unwrap(), before);
    assert_eq!(hook.count(), 1, "Keep must not write the slot");
    assert_eq!(hook.log(), vec![(None, Some(7))]);
}

#[test]
fn typed_tree_semantics_table() {
    let dir = tempdir().unwrap();
    let hook = CountHookU64::new();
    let tree = TypedTree::<[u8; 8], u64, U64Codec, _>::open_hooked(
        dir.path(),
        Config::test(),
        U64Codec,
        hook.clone(),
    )
    .unwrap();
    assert_bitcask_u64_semantics(&tree, &hook, &key(1));
}

#[test]
fn typed_map_semantics_table() {
    let dir = tempdir().unwrap();
    let hook = CountHookU64::new();
    let map = TypedMap::<[u8; 8], u64, U64Codec, _>::open_hooked(
        dir.path(),
        Config::test(),
        U64Codec,
        hook.clone(),
    )
    .unwrap();
    assert_bitcask_u64_semantics(&map, &hook, &key(1));
}

#[cfg(feature = "var-collections")]
#[test]
fn var_typed_map_semantics_table() {
    let dir = tempdir().unwrap();
    let hook = CountHookItem::new();
    let map = VarTypedMap::<[u8; 8], Item, RapiraCodec, _>::open_hooked(
        dir.path(),
        Config::test(),
        RapiraCodec,
        hook.clone(),
    )
    .unwrap();
    assert_bitcask_item_semantics(&map, &hook, &key(1));
}

#[cfg(feature = "var-collections")]
#[test]
fn var_typed_tree_semantics_table() {
    let dir = tempdir().unwrap();
    let hook = CountHookItem::new();
    let tree = VarTypedTree::<[u8; 8], Item, RapiraCodec, _>::open_hooked(
        dir.path(),
        Config::test(),
        RapiraCodec,
        hook.clone(),
    )
    .unwrap();
    assert_bitcask_item_semantics(&tree, &hook, &key(1));
}

#[test]
fn shard_mismatch_on_foreign_key() {
    let dir = tempdir().unwrap();
    let mut cfg = Config::test();
    cfg.shard_count = 2;
    let tree = ZeroTree::<[u8; 8], 8, Counter>::open(dir.path(), cfg).unwrap();
    let k0 = key(0);
    let foreign = (0u64..10_000)
        .map(key)
        .find(|k| tree.shard_for(k) != tree.shard_for(&k0))
        .expect("foreign key");
    let err = tree
        .atomic(&k0, |s| s.upsert(&foreign, |_| Upsert::Keep))
        .unwrap_err();
    assert!(matches!(err, DbError::ShardMismatch));
}

#[test]
fn family_ctx_hook_order_inside_atomic() {
    let dir = tempdir().unwrap();
    let hook = CountHook::new();
    let tree =
        ZeroTree::<[u8; 8], 8, Counter, _>::open_hooked(dir.path(), Config::test(), hook.clone())
            .unwrap();
    let k = key(42);
    tree.atomic(&k, |s| {
        s.upsert(&k, |_| Upsert::Put(Counter { n: 1 }))?;
        assert_eq!(
            hook.count(),
            0,
            "hooks must be buffered, never fired under the held shard lock"
        );
        s.put(&k, &Counter { n: 2 })?;
        s.upsert(&k, |_| Upsert::Keep)?;
        s.upsert(&k, |_| Upsert::Delete)?;
        assert_eq!(
            hook.count(),
            0,
            "still buffered at the end of the atomic block"
        );
        Ok(())
    })
    .unwrap();
    // Replayed after unlock, interleaved with `put` in application order.
    // `Keep` wrote nothing, so it contributes no entry.
    assert_eq!(
        hook.log(),
        vec![
            (None, Some(1)),    // upsert Put — insert
            (Some(1), Some(2)), // put — overwrite
            (Some(2), None),    // upsert Delete
        ]
    );
    assert_eq!(hook.count(), 3);
    assert!(tree.get(&k).is_none());
}

#[cfg(feature = "fjall")]
mod lsm {
    use super::*;
    use armdb::CollectionMeta;
    use armour_core::{GetRefs, GetType};

    #[derive(Clone, Debug, PartialEq, Eq, Rapira, GetType, GetRefs)]
    struct LsmItem {
        n: u64,
    }

    impl CollectionMeta for LsmItem {
        type SelfId = [u8; 8];
        const NAME: &'static str = "upsert_lsm_item";
        const VERSION: u16 = 1;
    }

    #[derive(Default, Clone)]
    struct LsmCountHook {
        writes: Arc<AtomicUsize>,
    }

    impl TypedWriteHook<[u8; 8], LsmItem> for LsmCountHook {
        const NEEDS_WRITE: bool = true;
        fn on_write(&self, _k: &[u8; 8], _old: Option<&LsmItem>, _new: Option<&LsmItem>) {
            self.writes.fetch_add(1, Ordering::Relaxed);
        }
    }

    impl LsmCountHook {
        fn new() -> Self {
            Self {
                writes: Arc::new(AtomicUsize::new(0)),
            }
        }
        fn count(&self) -> usize {
            self.writes.load(Ordering::Relaxed)
        }
    }

    #[test]
    fn lsm_semantics_table() {
        let dir = tempdir().unwrap();
        let db = fjall::Database::builder(dir.path()).open().unwrap();
        let hook = LsmCountHook::new();
        let tree = LsmTree::<[u8; 8], LsmItem, RapiraCodec, _>::open(
            db,
            LsmConfig {
                lock_stripes: 1,
                ..LsmConfig::default()
            },
            RapiraCodec,
            hook.clone(),
        )
        .unwrap();
        let k = key(1);

        tree.upsert(&k, |_| Upsert::Put(LsmItem { n: 1 })).unwrap();
        assert_eq!(tree.get(&k).unwrap().n, 1);
        assert_eq!(hook.count(), 1);

        tree.upsert(&k, |c| {
            Upsert::Put(LsmItem {
                n: c.unwrap().n + 1,
            })
        })
        .unwrap();
        assert_eq!(tree.get(&k).unwrap().n, 2);
        assert_eq!(hook.count(), 2);

        let len_before = tree.len();
        let writes_before = hook.count();
        tree.upsert(&k, |_| Upsert::Keep).unwrap();
        assert_eq!(tree.get(&k).unwrap().n, 2);
        assert_eq!(hook.count(), writes_before);
        assert_eq!(tree.len(), len_before);

        tree.upsert(&k, |_| Upsert::Delete).unwrap();
        assert!(tree.get(&k).is_none());
        assert_eq!(hook.count(), writes_before + 1);

        let w = hook.count();
        tree.upsert(&k, |_| Upsert::Delete).unwrap();
        assert_eq!(hook.count(), w);
    }

    #[test]
    fn lsm_hook_fires_synchronously_under_stripe() {
        let dir = tempdir().unwrap();
        let db = fjall::Database::builder(dir.path()).open().unwrap();
        let hook = LsmCountHook::new();
        let tree = LsmTree::<[u8; 8], LsmItem, RapiraCodec, _>::open(
            db,
            LsmConfig {
                lock_stripes: 1,
                ..LsmConfig::default()
            },
            RapiraCodec,
            hook.clone(),
        )
        .unwrap();
        let k = key(9);

        tree.atomic(&k, |s| {
            assert_eq!(hook.count(), 0);
            s.upsert(&k, |_| Upsert::Put(LsmItem { n: 1 }))?;
            assert_eq!(hook.count(), 1, "hook must fire before atomic returns");
            Ok(())
        })
        .unwrap();
        assert_eq!(hook.count(), 1);
    }

    #[test]
    fn lsm_upsert_keep_rejects_foreign_stripe() {
        let dir = tempdir().unwrap();
        let db = fjall::Database::builder(dir.path()).open().unwrap();
        let tree = LsmTree::<[u8; 8], LsmItem, RapiraCodec>::open(
            db,
            LsmConfig {
                lock_stripes: 4,
                ..LsmConfig::default()
            },
            RapiraCodec,
            armdb::NoHook,
        )
        .unwrap();
        let k0 = key(0);
        let foreign = (0u64..10_000)
            .map(key)
            .find(|k| tree.shard_for(k) != tree.shard_for(&k0))
            .expect("foreign key");
        let err = tree
            .atomic(&k0, |s| s.upsert(&foreign, |_| Upsert::Keep))
            .unwrap_err();
        assert!(
            matches!(err, armdb::DbError::ShardMismatch),
            "Keep on a foreign stripe must ShardMismatch, got {err:?}"
        );
    }
}