armdb 0.6.0

//! Integration tests for Db::atomic2/atomic3/atomic4.
#![cfg(all(
    feature = "armour",
    feature = "typed-tree",
    feature = "rapira-codec",
    feature = "var-collections"
))]

use std::path::Path;

use armdb::armour::Db;
use armdb::{Config, ConstTree};

// Independent Const trees, [u8;8] -> [u8;4]. Opened STANDALONE (no Db opener for
// Const); 4 shards so two keys can land in different shards.
fn open_tree(dir: &Path, name: &str) -> ConstTree<[u8; 8], 4> {
    let mut cfg = Config::test();
    cfg.shard_count = 4;
    ConstTree::<[u8; 8], 4>::open(dir.join(name), cfg).expect("open const tree")
}

#[test]
fn atomic2_writes_both_collections() {
    let dir = tempfile::tempdir().unwrap();
    let db = Db::open_test(dir.path()).unwrap();
    let a = open_tree(dir.path(), "a");
    let b = open_tree(dir.path(), "b");

    let ka = [1u8; 8];
    let kb = [2u8; 8];
    db.atomic2(&a, &[ka], &b, &[kb], |ta, tb| {
        ta.put(&ka, &[11u8; 4])?;
        tb.put(&kb, &[22u8; 4])?;
        Ok(())
    })
    .unwrap();

    assert_eq!(a.get(&ka), Some([11u8; 4]));
    assert_eq!(b.get(&kb), Some([22u8; 4]));
}

#[test]
fn atomic2_rejects_same_collection_twice() {
    let dir = tempfile::tempdir().unwrap();
    let db = Db::open_test(dir.path()).unwrap();
    let a = open_tree(dir.path(), "a");
    let ka = [1u8; 8];
    let err = db.atomic2(&a, &[ka], &a, &[ka], |_, _| Ok(()));
    assert!(matches!(err, Err(armdb::DbError::DuplicateCollectionInTx)));
}

#[test]
fn atomic3_and_atomic4_write_all() {
    let dir = tempfile::tempdir().unwrap();
    let db = Db::open_test(dir.path()).unwrap();
    let a = open_tree(dir.path(), "a");
    let b = open_tree(dir.path(), "b");
    let c = open_tree(dir.path(), "c");
    let d = open_tree(dir.path(), "d");
    let (ka, kb, kc, kd) = ([1u8; 8], [2u8; 8], [3u8; 8], [4u8; 8]);

    db.atomic3(&a, &[ka], &b, &[kb], &c, &[kc], |ta, tb, tc| {
        ta.put(&ka, &[1u8; 4])?;
        tb.put(&kb, &[2u8; 4])?;
        tc.put(&kc, &[3u8; 4])?;
        Ok(())
    })
    .unwrap();
    db.atomic4(
        &a,
        &[ka],
        &b,
        &[kb],
        &c,
        &[kc],
        &d,
        &[kd],
        |ta, tb, tc, td| {
            ta.delete(&ka)?;
            tb.put(&kb, &[9u8; 4])?;
            tc.put(&kc, &[9u8; 4])?;
            td.put(&kd, &[4u8; 4])?;
            Ok(())
        },
    )
    .unwrap();

    assert_eq!(a.get(&ka), None);
    assert_eq!(b.get(&kb), Some([9u8; 4]));
    assert_eq!(d.get(&kd), Some([4u8; 4]));
}

use armdb::ConstMap;

fn open_map(dir: &Path, name: &str) -> ConstMap<[u8; 8], 4> {
    let mut cfg = Config::test();
    cfg.shard_count = 4;
    ConstMap::<[u8; 8], 4>::open(dir.join(name), cfg).expect("open const map")
}

#[test]
fn atomic2_tree_plus_map_heterogeneous() {
    let dir = tempfile::tempdir().unwrap();
    let db = Db::open_test(dir.path()).unwrap();
    let t = open_tree(dir.path(), "t");
    let m = open_map(dir.path(), "m");
    let kt = [1u8; 8];
    let km = [2u8; 8];

    db.atomic2(&t, &[kt], &m, &[km], |tt, tm| {
        tt.put(&kt, &[7u8; 4])?;
        tm.put(&km, &[8u8; 4])?;
        Ok(())
    })
    .unwrap();

    assert_eq!(t.get(&kt), Some([7u8; 4]));
    assert_eq!(m.get(&km), Some([8u8; 4]));
}

use armdb::ZeroTree;

fn open_zero(dir: &Path, name: &str) -> ZeroTree<[u8; 8], 4, u32> {
    let mut cfg = Config::test();
    cfg.shard_count = 4;
    ZeroTree::<[u8; 8], 4, u32>::open(dir.join(name), cfg).expect("open zero tree")
}

#[test]
fn atomic2_zero_trees_typed_values() {
    let dir = tempfile::tempdir().unwrap();
    let db = Db::open_test(dir.path()).unwrap();
    let z1 = open_zero(dir.path(), "z1");
    let z2 = open_zero(dir.path(), "z2");
    let k1 = [1u8; 8];
    let k2 = [2u8; 8];

    db.atomic2(&z1, &[k1], &z2, &[k2], |t1, t2| {
        t1.put(&k1, &111u32)?;
        t2.put(&k2, &222u32)?;
        assert_eq!(t1.try_get(&k1)?, Some(111u32));
        Ok(())
    })
    .unwrap();

    assert_eq!(z1.get(&k1), Some(111u32));
    assert_eq!(z2.get(&k2), Some(222u32));
}

// ---- Typed collections (Task 9/10) --------------------------------------
use armdb::{CollectionMeta, NoHook, RapiraCodec};
use armour_core::GetType;
use rapira::Rapira;

#[derive(Clone, Debug, PartialEq, Rapira, GetType)]
struct TItemA {
    value: u64,
}
impl CollectionMeta for TItemA {
    type SelfId = [u8; 8];
    const NAME: &'static str = "typed_a";
}

#[derive(Clone, Debug, PartialEq, Rapira, GetType)]
struct TItemB {
    value: u64,
}
impl CollectionMeta for TItemB {
    type SelfId = [u8; 8];
    const NAME: &'static str = "typed_b";
}

#[test]
fn atomic2_typed_trees() {
    let dir = tempfile::tempdir().unwrap();
    let db = Db::open_test(dir.path()).unwrap();
    let a = db
        .open_typed_tree::<TItemA, RapiraCodec, _>(Config::test(), NoHook, &[])
        .unwrap();
    let b = db
        .open_typed_tree::<TItemB, RapiraCodec, _>(Config::test(), NoHook, &[])
        .unwrap();
    let ka = [1u8; 8];
    let kb = [2u8; 8];

    db.atomic2(a.as_ref(), &[ka], b.as_ref(), &[kb], |ta, tb| {
        ta.put(&ka, TItemA { value: 10 })?;
        tb.put(&kb, TItemB { value: 20 })?;
        assert_eq!(ta.try_get(&ka)?, Some(&TItemA { value: 10 }));
        Ok(())
    })
    .unwrap();

    assert_eq!(a.get(&ka).unwrap().value, 10);
    assert_eq!(b.get(&kb).unwrap().value, 20);
}

#[derive(Clone, Debug, PartialEq, Rapira, GetType)]
struct TMapItem {
    value: u64,
}
impl CollectionMeta for TMapItem {
    type SelfId = [u8; 8];
    const NAME: &'static str = "typed_map_item";
}

#[test]
fn atomic2_typed_map_plus_tree() {
    let dir = tempfile::tempdir().unwrap();
    let db = Db::open_test(dir.path()).unwrap();
    let m = db
        .open_typed_map::<TMapItem, RapiraCodec, _>(Config::test(), NoHook, &[])
        .unwrap();
    let t = db
        .open_typed_tree::<TItemA, RapiraCodec, _>(Config::test(), NoHook, &[])
        .unwrap();
    let km = [3u8; 8];
    let kt = [4u8; 8];

    db.atomic2(m.as_ref(), &[km], t.as_ref(), &[kt], |tm, tt| {
        tm.put(&km, TMapItem { value: 30 })?;
        tt.put(&kt, TItemA { value: 40 })?;
        assert_eq!(tm.try_get(&km)?, Some(&TMapItem { value: 30 }));
        Ok(())
    })
    .unwrap();

    assert_eq!(m.get(&km).unwrap().value, 30);
    assert_eq!(t.get(&kt).unwrap().value, 40);
}

// ---- Var collections (Task 11/12) ---------------------------------------
use armdb::VarTree;

fn open_var(dir: &Path, name: &str) -> VarTree<[u8; 8]> {
    let mut cfg = Config::test();
    cfg.shard_count = 4;
    VarTree::<[u8; 8]>::open(dir.join(name), cfg).expect("open var tree")
}

#[test]
fn atomic2_var_trees() {
    let dir = tempfile::tempdir().unwrap();
    let db = Db::open_test(dir.path()).unwrap();
    let a = open_var(dir.path(), "va");
    let b = open_var(dir.path(), "vb");
    let ka = [1u8; 8];
    let kb = [2u8; 8];

    db.atomic2(&a, &[ka], &b, &[kb], |ta, tb| {
        ta.put(&ka, b"hello")?;
        tb.put(&kb, b"world")?;
        assert_eq!(ta.try_get(&ka)?.as_deref(), Some(&b"hello"[..]));
        Ok(())
    })
    .unwrap();

    assert_eq!(a.get(&ka).as_deref(), Some(&b"hello"[..]));
    assert_eq!(b.get(&kb).as_deref(), Some(&b"world"[..]));
}

use armdb::VarMap;

fn open_var_map(dir: &Path, name: &str) -> VarMap<[u8; 8]> {
    let mut cfg = Config::test();
    cfg.shard_count = 4;
    VarMap::<[u8; 8]>::open(dir.join(name), cfg).expect("open var map")
}

#[test]
fn atomic2_var_map_plus_var_tree() {
    let dir = tempfile::tempdir().unwrap();
    let db = Db::open_test(dir.path()).unwrap();
    let m = open_var_map(dir.path(), "vmap");
    let t = open_var(dir.path(), "vtree2");
    let km = [5u8; 8];
    let kt = [6u8; 8];

    db.atomic2(&m, &[km], &t, &[kt], |tm, tt| {
        tm.put(&km, b"mapval")?;
        tt.put(&kt, b"treeval")?;
        assert_eq!(tm.try_get(&km)?.as_deref(), Some(&b"mapval"[..]));
        Ok(())
    })
    .unwrap();

    assert_eq!(m.get(&km).as_deref(), Some(&b"mapval"[..]));
    assert_eq!(t.get(&kt).as_deref(), Some(&b"treeval"[..]));
}

// ---- VarTyped collections (Task 13) -------------------------------------
use armdb::VarTypedTree;

#[derive(Clone, Debug, PartialEq, Rapira)]
struct VTItem {
    label: String,
    n: u64,
}

fn open_var_typed(dir: &Path, name: &str) -> VarTypedTree<[u8; 8], VTItem, RapiraCodec> {
    let mut cfg = Config::test();
    cfg.shard_count = 4;
    VarTypedTree::<[u8; 8], VTItem, RapiraCodec>::open(dir.join(name), cfg, RapiraCodec)
        .expect("open var typed tree")
}

#[test]
fn atomic2_var_typed_trees() {
    let dir = tempfile::tempdir().unwrap();
    let db = Db::open_test(dir.path()).unwrap();
    let a = open_var_typed(dir.path(), "vt_a");
    let b = open_var_typed(dir.path(), "vt_b");
    let ka = [1u8; 8];
    let kb = [2u8; 8];

    db.atomic2(&a, &[ka], &b, &[kb], |ta, tb| {
        ta.put(
            &ka,
            &VTItem {
                label: "x".into(),
                n: 1,
            },
        )?;
        tb.put(
            &kb,
            &VTItem {
                label: "y".into(),
                n: 2,
            },
        )?;
        assert_eq!(
            ta.try_get(&ka)?,
            Some(VTItem {
                label: "x".into(),
                n: 1
            })
        );
        Ok(())
    })
    .unwrap();

    assert_eq!(
        a.get(&ka),
        Some(VTItem {
            label: "x".into(),
            n: 1
        })
    );
    assert_eq!(
        b.get(&kb),
        Some(VTItem {
            label: "y".into(),
            n: 2
        })
    );
}

// ---- Task 14: contract tests --------------------------------------------

// R1: every accessor enforces the locked-shard scope, including try_contains.
#[test]
fn out_of_scope_key_errors_on_all_accessors() {
    let dir = tempfile::tempdir().unwrap();
    let db = Db::open_test(dir.path()).unwrap();
    let a = open_tree(dir.path(), "a");
    let b = open_tree(dir.path(), "b");
    let ka = [1u8; 8];
    let kb = [2u8; 8];
    // A key NOT passed to atomic2 (different shard than ka, best-effort).
    let mut other = [9u8; 8];
    for i in 0..=255u8 {
        other[0] = i;
        if a.shard_for(&other) != a.shard_for(&ka) {
            break;
        }
    }
    db.atomic2(&a, &[ka], &b, &[kb], |ta, _tb| {
        if a.shard_for(&other) != a.shard_for(&ka) {
            assert!(matches!(
                ta.try_get(&other),
                Err(armdb::DbError::ShardMismatch)
            ));
            assert!(matches!(
                ta.try_contains(&other),
                Err(armdb::DbError::ShardMismatch)
            ));
            assert!(matches!(
                ta.put(&other, &[0u8; 4]),
                Err(armdb::DbError::ShardMismatch)
            ));
            // genuinely-absent in-scope key: Ok(None)/Ok(false), not Err
            assert_eq!(ta.try_get(&ka).unwrap(), None);
            assert!(!ta.try_contains(&ka).unwrap());
        }
        Ok(())
    })
    .unwrap();
}

// Err from closure: applied writes remain, hooks replay, Err propagates (3a).
#[test]
fn err_from_closure_keeps_writes_and_replays_hooks() {
    let dir = tempfile::tempdir().unwrap();
    let db = Db::open_test(dir.path()).unwrap();
    let a = open_tree(dir.path(), "a");
    let b = open_tree(dir.path(), "b");
    let ka = [1u8; 8];
    let kb = [2u8; 8];
    let r: Result<(), _> = db.atomic2(&a, &[ka], &b, &[kb], |ta, _tb| {
        ta.put(&ka, &[5u8; 4])?;
        Err(armdb::DbError::TxConflict)
    });
    assert!(matches!(r, Err(armdb::DbError::TxConflict)));
    assert_eq!(a.get(&ka), Some([5u8; 4])); // no rollback
}

// Mixed backends: a Bitcask ConstTree + a Fixed ConstMap in one transaction.
#[test]
fn atomic2_mixed_backend_bitcask_tree_plus_fixed_map() {
    use armdb::{FixedConfig, FixedMap};
    let dir = tempfile::tempdir().unwrap();
    let db = Db::open_test(dir.path()).unwrap();
    let t = open_tree(dir.path(), "bt"); // Bitcask ConstTree
    let mut fcfg = FixedConfig::test();
    fcfg.shard_count = 4;
    let m: FixedMap<[u8; 8], 4> =
        FixedMap::<[u8; 8], 4>::open(dir.path().join("fm"), fcfg).expect("open fixed map");
    let kt = [1u8; 8];
    let km = [2u8; 8];

    db.atomic2(&t, &[kt], &m, &[km], |tt, tm| {
        tt.put(&kt, &[3u8; 4])?;
        tm.put(&km, &[4u8; 4])?;
        Ok(())
    })
    .unwrap();

    assert_eq!(t.get(&kt), Some([3u8; 4]));
    assert_eq!(m.get(&km), Some([4u8; 4]));
}

// ---- Task 15: stress + reader visibility --------------------------------

// Stress: N threads run atomic2 transfers across two collections; the total
// balance is invariant. Run with --retries 2.
#[test]
fn stress_atomic2_preserves_balance_sum() {
    use std::sync::Arc;
    use std::thread;
    let dir = tempfile::tempdir().unwrap();
    let db = Arc::new(Db::open_test(dir.path()).unwrap());
    let acc_a = Arc::new(open_tree(dir.path(), "acc_a")); // [u8;8] -> [u8;4] LE balance
    let acc_b = Arc::new(open_tree(dir.path(), "acc_b"));

    let k = [1u8; 8];
    acc_a.put(&k, &100u32.to_le_bytes()).unwrap();
    acc_b.put(&k, &100u32.to_le_bytes()).unwrap();

    let mut handles = Vec::new();
    for t in 0..8 {
        let (db, acc_a, acc_b) = (db.clone(), acc_a.clone(), acc_b.clone());
        handles.push(thread::spawn(move || {
            for _ in 0..200 {
                // Even threads move a->b, odd threads b->a. Skip when source is 0.
                let _ = db.atomic2(acc_a.as_ref(), &[k], acc_b.as_ref(), &[k], |ta, tb| {
                    let va = u32::from_le_bytes(ta.get_or_err(&k)?);
                    let vb = u32::from_le_bytes(tb.get_or_err(&k)?);
                    if t % 2 == 0 {
                        if va > 0 {
                            ta.put(&k, &(va - 1).to_le_bytes())?;
                            tb.put(&k, &(vb + 1).to_le_bytes())?;
                        }
                    } else if vb > 0 {
                        tb.put(&k, &(vb - 1).to_le_bytes())?;
                        ta.put(&k, &(va + 1).to_le_bytes())?;
                    }
                    Ok(())
                });
            }
        }));
    }
    for h in handles {
        h.join().unwrap();
    }
    let s1 = u32::from_le_bytes(acc_a.get(&k).unwrap());
    let s2 = u32::from_le_bytes(acc_b.get(&k).unwrap());
    assert_eq!(s1 + s2, 200); // write isolation preserves the total
}

// R5: reader visibility — a concurrent reader of a tree participant may observe
// one collection updated and the other not. Assert no panic / no torn value, and
// document the non-guarantee (we do NOT assert atomic visibility).
#[test]
fn reader_sees_no_torn_values_during_atomic2() {
    use std::sync::Arc;
    use std::sync::atomic::{AtomicBool, Ordering};
    use std::thread;
    let dir = tempfile::tempdir().unwrap();
    let db = Arc::new(Db::open_test(dir.path()).unwrap());
    let a = Arc::new(open_tree(dir.path(), "ra"));
    let b = Arc::new(open_tree(dir.path(), "rb"));
    let k = [1u8; 8];
    a.put(&k, &[0u8; 4]).unwrap();
    b.put(&k, &[0u8; 4]).unwrap();

    let stop = Arc::new(AtomicBool::new(false));

    let writer = {
        let (db, a, b, stop) = (db.clone(), a.clone(), b.clone(), stop.clone());
        thread::spawn(move || {
            for i in 0u32..2000 {
                let v = (i % 250) as u8;
                db.atomic2(a.as_ref(), &[k], b.as_ref(), &[k], |ta, tb| {
                    ta.put(&k, &[v; 4])?;
                    tb.put(&k, &[v; 4])?;
                    Ok(())
                })
                .unwrap();
            }
            stop.store(true, Ordering::Relaxed);
        })
    };

    let reader = {
        let (a, b, stop) = (a.clone(), b.clone(), stop.clone());
        thread::spawn(move || {
            while !stop.load(Ordering::Relaxed) {
                // Every observed value must be a fully-written [b;4] (no torn read).
                if let Some(va) = a.get(&k) {
                    assert!(va.iter().all(|&x| x == va[0]));
                }
                if let Some(vb) = b.get(&k) {
                    assert!(vb.iter().all(|&x| x == vb[0]));
                }
            }
        })
    };

    writer.join().unwrap();
    reader.join().unwrap();
}