armdb 0.1.13

#[cfg(feature = "var-collections")]
use armdb::VarTree;
use armdb::{Config, ConstTree, MigrateAction, WriteHook};
use parking_lot::Mutex;
use std::collections::HashMap;
use std::sync::Arc;
use tempfile::tempdir;

#[derive(Clone, Debug, PartialEq)]
struct HookEvent {
    key: [u8; 8],
    old: Option<Vec<u8>>,
    new: Option<Vec<u8>>,
}

type Events = Arc<Mutex<Vec<HookEvent>>>;

fn take_events(events: &Events) -> Vec<HookEvent> {
    std::mem::take(&mut *events.lock())
}

/// Hook with NEEDS_OLD_VALUE = true.
struct TrackingHook {
    events: Events,
}

impl TrackingHook {
    fn new(events: Events) -> Self {
        Self { events }
    }
}

impl WriteHook<[u8; 8]> for TrackingHook {
    const NEEDS_OLD_VALUE: bool = true;

    fn on_write(&self, key: &[u8; 8], old: Option<&[u8]>, new: Option<&[u8]>) {
        self.events.lock().push(HookEvent {
            key: *key,
            old: old.map(|v| v.to_vec()),
            new: new.map(|v| v.to_vec()),
        });
    }
}

#[cfg(feature = "var-collections")]
/// Hook with NEEDS_OLD_VALUE = false — old is always None.
struct NoOldHook {
    events: Events,
}

#[cfg(feature = "var-collections")]
impl NoOldHook {
    fn new(events: Events) -> Self {
        Self { events }
    }
}

#[cfg(feature = "var-collections")]
impl WriteHook<[u8; 8]> for NoOldHook {
    const NEEDS_OLD_VALUE: bool = false;

    fn on_write(&self, key: &[u8; 8], old: Option<&[u8]>, new: Option<&[u8]>) {
        self.events.lock().push(HookEvent {
            key: *key,
            old: old.map(|v| v.to_vec()),
            new: new.map(|v| v.to_vec()),
        });
    }
}

#[test]
fn test_hook_on_put() {
    let dir = tempdir().unwrap();
    let events: Events = Arc::new(Mutex::new(Vec::new()));
    let tree = ConstTree::<[u8; 8], 8, TrackingHook>::open_hooked(
        dir.path(),
        Config::test(),
        TrackingHook::new(events.clone()),
    )
    .unwrap();

    let key = 1u64.to_be_bytes();
    let val1 = 100u64.to_be_bytes();
    let val2 = 200u64.to_be_bytes();

    // First put — no old value
    tree.put(&key, &val1).unwrap();
    let evs = take_events(&events);
    assert_eq!(evs.len(), 1);
    assert_eq!(evs[0].key, key);
    assert_eq!(evs[0].old, None);
    assert_eq!(evs[0].new, Some(val1.to_vec()));

    // Overwrite — old value present
    tree.put(&key, &val2).unwrap();
    let evs = take_events(&events);
    assert_eq!(evs.len(), 1);
    assert_eq!(evs[0].key, key);
    assert_eq!(evs[0].old, Some(val1.to_vec()));
    assert_eq!(evs[0].new, Some(val2.to_vec()));
}

#[test]
fn test_hook_on_delete() {
    let dir = tempdir().unwrap();
    let events: Events = Arc::new(Mutex::new(Vec::new()));
    let tree = ConstTree::<[u8; 8], 8, TrackingHook>::open_hooked(
        dir.path(),
        Config::test(),
        TrackingHook::new(events.clone()),
    )
    .unwrap();

    let key = 1u64.to_be_bytes();
    let val = 100u64.to_be_bytes();

    tree.put(&key, &val).unwrap();
    events.lock().clear();

    // Delete existing key
    tree.delete(&key).unwrap();
    let evs = take_events(&events);
    assert_eq!(evs.len(), 1);
    assert_eq!(evs[0].key, key);
    assert_eq!(evs[0].old, Some(val.to_vec()));
    assert_eq!(evs[0].new, None);
}

#[test]
fn test_hook_on_insert() {
    let dir = tempdir().unwrap();
    let events: Events = Arc::new(Mutex::new(Vec::new()));
    let tree = ConstTree::<[u8; 8], 8, TrackingHook>::open_hooked(
        dir.path(),
        Config::test(),
        TrackingHook::new(events.clone()),
    )
    .unwrap();

    let key = 1u64.to_be_bytes();
    let val = 100u64.to_be_bytes();

    tree.insert(&key, &val).unwrap();
    let evs = take_events(&events);
    assert_eq!(evs.len(), 1);
    assert_eq!(evs[0].key, key);
    assert_eq!(evs[0].old, None);
    assert_eq!(evs[0].new, Some(val.to_vec()));
}

#[test]
fn test_hook_on_cas() {
    let dir = tempdir().unwrap();
    let events: Events = Arc::new(Mutex::new(Vec::new()));
    let tree = ConstTree::<[u8; 8], 8, TrackingHook>::open_hooked(
        dir.path(),
        Config::test(),
        TrackingHook::new(events.clone()),
    )
    .unwrap();

    let key = 1u64.to_be_bytes();
    let val1 = 100u64.to_be_bytes();
    let val2 = 200u64.to_be_bytes();

    tree.put(&key, &val1).unwrap();
    events.lock().clear();

    tree.cas(&key, &val1, &val2).unwrap();
    let evs = take_events(&events);
    assert_eq!(evs.len(), 1);
    assert_eq!(evs[0].key, key);
    assert_eq!(evs[0].old, Some(val1.to_vec()));
    assert_eq!(evs[0].new, Some(val2.to_vec()));
}

#[test]
fn test_hook_not_called_on_compaction() {
    let dir = tempdir().unwrap();
    let events: Events = Arc::new(Mutex::new(Vec::new()));
    let mut config = Config::test();
    config.max_file_size = 4096;
    config.compaction_threshold = 0.1;

    let tree = ConstTree::<[u8; 8], 8, TrackingHook>::open_hooked(
        dir.path(),
        config,
        TrackingHook::new(events.clone()),
    )
    .unwrap();

    for i in 0..100u64 {
        let key = i.to_be_bytes();
        let val = i.to_be_bytes();
        tree.put(&key, &val).unwrap();
    }
    for i in 0..50u64 {
        let key = i.to_be_bytes();
        let val = (i + 1000).to_be_bytes();
        tree.put(&key, &val).unwrap();
    }

    // Clear events before compaction
    events.lock().clear();

    tree.compact().unwrap();

    let evs = take_events(&events);
    assert!(
        evs.is_empty(),
        "hook should not be called during compaction, but got {} events",
        evs.len()
    );

    // Verify data is still correct
    for i in 0..100u64 {
        let key = i.to_be_bytes();
        let expected = if i < 50 {
            (i + 1000).to_be_bytes()
        } else {
            i.to_be_bytes()
        };
        assert_eq!(tree.get(&key).unwrap(), expected);
    }
}

#[cfg(feature = "var-collections")]
#[test]
fn test_hook_needs_old_value_false_var_tree() {
    // NEEDS_OLD_VALUE = false only matters for VarTree (skips disk I/O).
    // ConstTree always provides old value (inline in memory, zero cost).
    let dir = tempdir().unwrap();
    let events: Events = Arc::new(Mutex::new(Vec::new()));
    let tree = VarTree::<[u8; 8], NoOldHook>::open_hooked(
        dir.path(),
        Config::test(),
        NoOldHook::new(events.clone()),
    )
    .unwrap();

    let key = 1u64.to_be_bytes();

    // First put
    tree.put(&key, b"hello").unwrap();
    let evs = take_events(&events);
    assert_eq!(evs.len(), 1);
    assert_eq!(evs[0].old, None);

    // Overwrite — old still None because NEEDS_OLD_VALUE = false
    tree.put(&key, b"world").unwrap();
    let evs = take_events(&events);
    assert_eq!(evs.len(), 1);
    assert_eq!(evs[0].old, None);
    assert_eq!(evs[0].new, Some(b"world".to_vec()));
}

#[cfg(feature = "var-collections")]
#[test]
fn test_hook_var_tree() {
    let dir = tempdir().unwrap();
    let events: Events = Arc::new(Mutex::new(Vec::new()));
    let tree = VarTree::<[u8; 8], TrackingHook>::open_hooked(
        dir.path(),
        Config::test(),
        TrackingHook::new(events.clone()),
    )
    .unwrap();

    let key = 1u64.to_be_bytes();

    // Put
    tree.put(&key, b"hello").unwrap();
    let evs = take_events(&events);
    assert_eq!(evs.len(), 1);
    assert_eq!(evs[0].old, None);
    assert_eq!(evs[0].new, Some(b"hello".to_vec()));

    // Overwrite
    tree.put(&key, b"world").unwrap();
    let evs = take_events(&events);
    assert_eq!(evs.len(), 1);
    assert_eq!(evs[0].old, Some(b"hello".to_vec()));
    assert_eq!(evs[0].new, Some(b"world".to_vec()));

    // Delete
    tree.delete(&key).unwrap();
    let evs = take_events(&events);
    assert_eq!(evs.len(), 1);
    assert_eq!(evs[0].old, Some(b"world".to_vec()));
    assert_eq!(evs[0].new, None);
}

#[test]
fn test_hook_delete_nonexistent() {
    let dir = tempdir().unwrap();
    let events: Events = Arc::new(Mutex::new(Vec::new()));
    let tree = ConstTree::<[u8; 8], 8, TrackingHook>::open_hooked(
        dir.path(),
        Config::test(),
        TrackingHook::new(events.clone()),
    )
    .unwrap();

    let key = 99u64.to_be_bytes();

    // Delete non-existent key — no hook call
    let result = tree.delete(&key).unwrap();
    assert!(result.is_none());
    let evs = take_events(&events);
    assert!(
        evs.is_empty(),
        "hook should not fire on delete of missing key"
    );
}

// ---------------------------------------------------------------------------
// on_init hook types
// ---------------------------------------------------------------------------

type InitEvents = Arc<Mutex<HashMap<[u8; 8], Vec<u8>>>>;

/// Init-only hook: on_init tracks entries, on_write is a no-op.
struct InitOnlyHook {
    inits: InitEvents,
}

impl InitOnlyHook {
    fn new(inits: InitEvents) -> Self {
        Self { inits }
    }
}

impl WriteHook<[u8; 8]> for InitOnlyHook {
    const NEEDS_OLD_VALUE: bool = false;
    const NEEDS_INIT: bool = true;

    fn on_write(&self, _key: &[u8; 8], _old: Option<&[u8]>, _new: Option<&[u8]>) {
        // no-op — init-only hook
    }

    fn on_init(&self, key: &[u8; 8], value: &[u8]) {
        self.inits.lock().insert(*key, value.to_vec());
    }
}

/// Init+write hook: tracks both on_init and on_write calls.
struct InitWriteHook {
    inits: InitEvents,
    writes: Events,
}

impl InitWriteHook {
    fn new(inits: InitEvents, writes: Events) -> Self {
        Self { inits, writes }
    }
}

impl WriteHook<[u8; 8]> for InitWriteHook {
    const NEEDS_OLD_VALUE: bool = true;
    const NEEDS_INIT: bool = true;

    fn on_write(&self, key: &[u8; 8], old: Option<&[u8]>, new: Option<&[u8]>) {
        self.writes.lock().push(HookEvent {
            key: *key,
            old: old.map(|v| v.to_vec()),
            new: new.map(|v| v.to_vec()),
        });
    }

    fn on_init(&self, key: &[u8; 8], value: &[u8]) {
        self.inits.lock().insert(*key, value.to_vec());
    }
}

// ---------------------------------------------------------------------------
// on_init tests — init-only
// ---------------------------------------------------------------------------

/// on_init fires for all live entries via migrate(Keep).
#[test]
fn test_on_init_fires_via_migrate() {
    let dir = tempdir().unwrap();
    let inits: InitEvents = Arc::new(Mutex::new(HashMap::new()));

    let tree = ConstTree::<[u8; 8], 8, InitOnlyHook>::open_hooked(
        dir.path(),
        Config::test(),
        InitOnlyHook::new(inits.clone()),
    )
    .unwrap();

    // Insert entries
    for i in 0..10u64 {
        tree.put(&i.to_be_bytes(), &(i * 100).to_be_bytes())
            .unwrap();
    }

    // migrate(Keep) triggers on_init for every entry
    let mutated = tree.migrate(|_, _| MigrateAction::Keep).unwrap();
    assert_eq!(mutated, 0); // Keep doesn't count as mutation

    let inits = inits.lock();
    assert_eq!(inits.len(), 10);
    for i in 0..10u64 {
        let key = i.to_be_bytes();
        assert_eq!(inits[&key], (i * 100).to_be_bytes().to_vec());
    }
}

/// on_init after reopen: write data without hooks, close, reopen with init hook.
#[test]
fn test_on_init_after_reopen() {
    let dir = tempdir().unwrap();

    // Phase 1: write data without any hook
    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), Config::test()).unwrap();
        for i in 0..20u64 {
            tree.put(&i.to_be_bytes(), &(i * 10).to_be_bytes()).unwrap();
        }
        tree.close().unwrap();
    }

    // Phase 2: reopen with init hook, call migrate to trigger on_init
    let inits: InitEvents = Arc::new(Mutex::new(HashMap::new()));
    let tree = ConstTree::<[u8; 8], 8, InitOnlyHook>::open_hooked(
        dir.path(),
        Config::test(),
        InitOnlyHook::new(inits.clone()),
    )
    .unwrap();

    tree.migrate(|_, _| MigrateAction::Keep).unwrap();

    let inits = inits.lock();
    assert_eq!(inits.len(), 20);
    for i in 0..20u64 {
        let key = i.to_be_bytes();
        assert_eq!(inits[&key], (i * 10).to_be_bytes().to_vec());
    }
}

/// on_init skips deleted entries — only live entries are visited.
#[test]
fn test_on_init_skips_deleted() {
    let dir = tempdir().unwrap();

    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), Config::test()).unwrap();
        for i in 0..10u64 {
            tree.put(&i.to_be_bytes(), &i.to_be_bytes()).unwrap();
        }
        // Delete even keys
        for i in (0..10u64).step_by(2) {
            tree.delete(&i.to_be_bytes()).unwrap();
        }
        tree.close().unwrap();
    }

    let inits: InitEvents = Arc::new(Mutex::new(HashMap::new()));
    let tree = ConstTree::<[u8; 8], 8, InitOnlyHook>::open_hooked(
        dir.path(),
        Config::test(),
        InitOnlyHook::new(inits.clone()),
    )
    .unwrap();

    tree.migrate(|_, _| MigrateAction::Keep).unwrap();

    let inits = inits.lock();
    assert_eq!(inits.len(), 5, "only odd keys should survive");
    for i in (1..10u64).step_by(2) {
        assert!(inits.contains_key(&i.to_be_bytes()));
    }
}

/// on_write does NOT fire during migrate (hook-free put/delete).
#[test]
fn test_on_write_not_fired_during_migrate() {
    let dir = tempdir().unwrap();
    let inits: InitEvents = Arc::new(Mutex::new(HashMap::new()));
    let writes: Events = Arc::new(Mutex::new(Vec::new()));

    let tree = ConstTree::<[u8; 8], 8, InitWriteHook>::open_hooked(
        dir.path(),
        Config::test(),
        InitWriteHook::new(inits.clone(), writes.clone()),
    )
    .unwrap();

    for i in 0..5u64 {
        tree.put(&i.to_be_bytes(), &i.to_be_bytes()).unwrap();
    }
    writes.lock().clear(); // clear on_write events from puts

    // migrate with Update — should fire on_init but NOT on_write
    tree.migrate(|key, _val| {
        let new_val = [key[7] + 1, 0, 0, 0, 0, 0, 0, 0];
        MigrateAction::Update(new_val)
    })
    .unwrap();

    let write_evs = take_events(&writes);
    assert!(
        write_evs.is_empty(),
        "on_write should not fire during migrate, got {} events",
        write_evs.len()
    );

    let inits = inits.lock();
    assert_eq!(
        inits.len(),
        5,
        "on_init should fire for all updated entries"
    );
}

/// on_init fires with the NEW value for MigrateAction::Update.
#[test]
fn test_on_init_receives_updated_value() {
    let dir = tempdir().unwrap();
    let inits: InitEvents = Arc::new(Mutex::new(HashMap::new()));

    let tree = ConstTree::<[u8; 8], 8, InitOnlyHook>::open_hooked(
        dir.path(),
        Config::test(),
        InitOnlyHook::new(inits.clone()),
    )
    .unwrap();

    for i in 0..5u64 {
        tree.put(&i.to_be_bytes(), &i.to_be_bytes()).unwrap();
    }

    // Migrate: double every value
    tree.migrate(|_key, val| {
        let v = u64::from_be_bytes(*val);
        MigrateAction::Update((v * 2).to_be_bytes())
    })
    .unwrap();

    let inits = inits.lock();
    for i in 0..5u64 {
        let key = i.to_be_bytes();
        assert_eq!(
            inits[&key],
            (i * 2).to_be_bytes().to_vec(),
            "on_init should receive the updated (new) value"
        );
    }
}

/// on_init does NOT fire for MigrateAction::Delete.
#[test]
fn test_on_init_not_fired_on_migrate_delete() {
    let dir = tempdir().unwrap();
    let inits: InitEvents = Arc::new(Mutex::new(HashMap::new()));

    let tree = ConstTree::<[u8; 8], 8, InitOnlyHook>::open_hooked(
        dir.path(),
        Config::test(),
        InitOnlyHook::new(inits.clone()),
    )
    .unwrap();

    for i in 0..10u64 {
        tree.put(&i.to_be_bytes(), &i.to_be_bytes()).unwrap();
    }

    // Migrate: delete even keys, keep odd
    tree.migrate(|key, _| {
        let v = u64::from_be_bytes(*key);
        if v % 2 == 0 {
            MigrateAction::Delete
        } else {
            MigrateAction::Keep
        }
    })
    .unwrap();

    let inits = inits.lock();
    assert_eq!(inits.len(), 5, "on_init should only fire for kept entries");
    for i in (1..10u64).step_by(2) {
        assert!(inits.contains_key(&i.to_be_bytes()));
    }
}

// ---------------------------------------------------------------------------
// on_init + on_write tests
// ---------------------------------------------------------------------------

/// Both on_init and on_write work: on_init fires during migrate, on_write on subsequent ops.
#[test]
fn test_on_init_plus_on_write() {
    let dir = tempdir().unwrap();

    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), Config::test()).unwrap();
        for i in 0..5u64 {
            tree.put(&i.to_be_bytes(), &i.to_be_bytes()).unwrap();
        }
        tree.close().unwrap();
    }

    let inits: InitEvents = Arc::new(Mutex::new(HashMap::new()));
    let writes: Events = Arc::new(Mutex::new(Vec::new()));
    let tree = ConstTree::<[u8; 8], 8, InitWriteHook>::open_hooked(
        dir.path(),
        Config::test(),
        InitWriteHook::new(inits.clone(), writes.clone()),
    )
    .unwrap();

    // Phase 1: migrate → on_init fires, on_write does NOT
    tree.migrate(|_, _| MigrateAction::Keep).unwrap();

    assert_eq!(inits.lock().len(), 5, "on_init should fire for 5 entries");
    assert!(
        writes.lock().is_empty(),
        "on_write should not fire during migrate"
    );

    // Phase 2: normal operations → on_write fires
    let key = 0u64.to_be_bytes();
    let new_val = 999u64.to_be_bytes();
    tree.put(&key, &new_val).unwrap();

    let write_evs = take_events(&writes);
    assert_eq!(write_evs.len(), 1);
    assert_eq!(write_evs[0].key, key);
    assert_eq!(write_evs[0].old, Some(0u64.to_be_bytes().to_vec()));
    assert_eq!(write_evs[0].new, Some(new_val.to_vec()));
}

/// on_init does not fire when NEEDS_INIT = false (TrackingHook).
#[test]
fn test_on_init_not_fired_when_needs_init_false() {
    let dir = tempdir().unwrap();

    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), Config::test()).unwrap();
        for i in 0..5u64 {
            tree.put(&i.to_be_bytes(), &i.to_be_bytes()).unwrap();
        }
        tree.close().unwrap();
    }

    // TrackingHook has NEEDS_INIT = false (default)
    let events: Events = Arc::new(Mutex::new(Vec::new()));
    let tree = ConstTree::<[u8; 8], 8, TrackingHook>::open_hooked(
        dir.path(),
        Config::test(),
        TrackingHook::new(events.clone()),
    )
    .unwrap();

    // migrate(Keep) — on_init should NOT fire because NEEDS_INIT = false
    tree.migrate(|_, _| MigrateAction::Keep).unwrap();

    let evs = take_events(&events);
    assert!(
        evs.is_empty(),
        "on_write should not fire during migrate, and on_init is disabled"
    );
}