armdb 0.6.0

#![cfg(feature = "typed-tree")]

use armdb::codec::Codec;
use armdb::{Config, DbError, MigrateAction, TypedTree};
use tempfile::tempdir;

// ---------------------------------------------------------------------------
// Test codec: u64 as big-endian 8 bytes
// ---------------------------------------------------------------------------

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

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

type Tree = TypedTree<[u8; 8], u64, U64Codec>;

fn open(dir: &std::path::Path) -> Tree {
    TypedTree::<[u8; 8], u64, U64Codec>::open(dir, Config::test(), U64Codec).unwrap()
}

fn populate(tree: &Tree, n: u64) {
    for i in 0..n {
        tree.put(&i.to_be_bytes(), i * 10).unwrap();
    }
}

// ---------------------------------------------------------------------------
// Basic CRUD
// ---------------------------------------------------------------------------

#[test]
fn test_put_get() {
    let dir = tempdir().unwrap();
    let tree = open(dir.path());

    assert!(tree.get(&1u64.to_be_bytes()).is_none());

    tree.put(&1u64.to_be_bytes(), 100).unwrap();
    let r = tree.get(&1u64.to_be_bytes()).unwrap();
    assert_eq!(*r, 100);
}

#[test]
fn test_put_returns_old() {
    let dir = tempdir().unwrap();
    let tree = open(dir.path());

    let old = tree.put(&1u64.to_be_bytes(), 100).unwrap();
    assert!(old.is_none());

    let old = tree.put(&1u64.to_be_bytes(), 200).unwrap();
    assert_eq!(*old.unwrap(), 100);

    assert_eq!(*tree.get(&1u64.to_be_bytes()).unwrap(), 200);
}

#[test]
fn test_insert() {
    let dir = tempdir().unwrap();
    let tree = open(dir.path());

    tree.insert(&1u64.to_be_bytes(), 100).unwrap();
    assert_eq!(*tree.get(&1u64.to_be_bytes()).unwrap(), 100);

    let err = tree.insert(&1u64.to_be_bytes(), 200).unwrap_err();
    assert!(matches!(err, DbError::KeyExists));

    // Value unchanged
    assert_eq!(*tree.get(&1u64.to_be_bytes()).unwrap(), 100);
}

#[test]
fn test_delete() {
    let dir = tempdir().unwrap();
    let tree = open(dir.path());

    let old = tree.delete(&1u64.to_be_bytes()).unwrap();
    assert!(old.is_none());

    tree.put(&1u64.to_be_bytes(), 100).unwrap();
    let old = tree.delete(&1u64.to_be_bytes()).unwrap();
    assert_eq!(*old.unwrap(), 100);

    assert!(tree.get(&1u64.to_be_bytes()).is_none());
}

#[test]
fn test_contains() {
    let dir = tempdir().unwrap();
    let tree = open(dir.path());

    assert!(!tree.contains(&1u64.to_be_bytes()));
    tree.put(&1u64.to_be_bytes(), 42).unwrap();
    assert!(tree.contains(&1u64.to_be_bytes()));
}

#[test]
fn test_len() {
    let dir = tempdir().unwrap();
    let tree = open(dir.path());

    assert_eq!(tree.len(), 0);
    assert!(tree.is_empty());

    populate(&tree, 10);
    assert_eq!(tree.len(), 10);
    assert!(!tree.is_empty());
}

// ---------------------------------------------------------------------------
// CAS and Update
// ---------------------------------------------------------------------------

#[test]
fn test_cas() {
    let dir = tempdir().unwrap();
    let tree = open(dir.path());

    tree.put(&1u64.to_be_bytes(), 100).unwrap();

    // Mismatch
    let err = tree.cas(&1u64.to_be_bytes(), &999, 200).unwrap_err();
    assert!(matches!(err, DbError::CasMismatch));
    assert_eq!(*tree.get(&1u64.to_be_bytes()).unwrap(), 100);

    // Match
    tree.cas(&1u64.to_be_bytes(), &100, 200).unwrap();
    assert_eq!(*tree.get(&1u64.to_be_bytes()).unwrap(), 200);

    // Not found
    let err = tree.cas(&99u64.to_be_bytes(), &0, 1).unwrap_err();
    assert!(matches!(err, DbError::KeyNotFound));
}

#[test]
fn test_update() {
    let dir = tempdir().unwrap();
    let tree = open(dir.path());

    // Key not found
    let r = tree.update(&1u64.to_be_bytes(), |v| v + 1).unwrap();
    assert!(r.is_none());

    tree.put(&1u64.to_be_bytes(), 100).unwrap();

    // Update returns ref to NEW value
    let r = tree.update(&1u64.to_be_bytes(), |v| v * 2).unwrap();
    assert_eq!(*r.unwrap(), 200);
    assert_eq!(*tree.get(&1u64.to_be_bytes()).unwrap(), 200);
}

// ---------------------------------------------------------------------------
// Iteration
// ---------------------------------------------------------------------------

#[test]
fn test_iter_all() {
    let dir = tempdir().unwrap();
    let tree = open(dir.path());
    populate(&tree, 20);

    let entries: Vec<_> = tree.iter().collect();
    assert_eq!(entries.len(), 20);

    // reversed=true → descending
    for i in 0..20u64 {
        let expected = 19 - i;
        assert_eq!(entries[i as usize].0, expected.to_be_bytes());
        assert_eq!(*entries[i as usize].1, expected * 10);
    }
}

#[test]
fn test_iter_empty() {
    let dir = tempdir().unwrap();
    let tree = open(dir.path());

    assert_eq!(tree.iter().count(), 0);
    assert!(tree.iter().next().is_none());
    assert!(tree.iter().next_back().is_none());
}

#[test]
fn test_range() {
    let dir = tempdir().unwrap();
    let tree = open(dir.path());
    populate(&tree, 20);

    let start = 5u64.to_be_bytes();
    let end = 10u64.to_be_bytes();
    let entries: Vec<_> = tree.range(&start, &end).collect();

    assert_eq!(entries.len(), 5); // reversed: [9, 8, 7, 6, 5]
    for (i, (key, val)) in entries.iter().enumerate() {
        let expected = 9 - i as u64;
        assert_eq!(*key, expected.to_be_bytes());
        assert_eq!(**val, expected * 10);
    }
}

#[test]
fn test_prefix_iter() {
    let dir = tempdir().unwrap();
    let tree: TypedTree<[u8; 16], u64, U64Codec> =
        TypedTree::open(dir.path(), Config::test(), U64Codec).unwrap();

    // Insert entries with prefix [0,0,0,1,...] and [0,0,0,2,...]
    for i in 0u64..5 {
        let mut key = [0u8; 16];
        key[3] = 1;
        key[8..].copy_from_slice(&i.to_be_bytes());
        tree.put(&key, i).unwrap();
    }
    for i in 0u64..3 {
        let mut key = [0u8; 16];
        key[3] = 2;
        key[8..].copy_from_slice(&i.to_be_bytes());
        tree.put(&key, 100 + i).unwrap();
    }

    let entries: Vec<_> = tree.prefix_iter(&[0, 0, 0, 1]).collect();
    assert_eq!(entries.len(), 5);

    let entries: Vec<_> = tree.prefix_iter(&[0, 0, 0, 2]).collect();
    assert_eq!(entries.len(), 3);
}

// ---------------------------------------------------------------------------
// DoubleEndedIterator
// ---------------------------------------------------------------------------

#[test]
fn test_iter_rev() {
    let dir = tempdir().unwrap();
    let tree = open(dir.path());
    populate(&tree, 10);

    // .rev() on reversed=true → ascending
    let entries: Vec<_> = tree.iter().rev().collect();
    assert_eq!(entries.len(), 10);
    for i in 0..10u64 {
        assert_eq!(entries[i as usize].0, i.to_be_bytes());
    }
}

#[test]
fn test_next_back_mixed() {
    let dir = tempdir().unwrap();
    let tree = open(dir.path());
    populate(&tree, 5); // keys: 0,1,2,3,4

    // reversed=true: forward yields [4,3,2,1,0], backward yields [0,1,2,3,4]
    let mut iter = tree.iter();

    let (k, _) = iter.next().unwrap();
    assert_eq!(k, 4u64.to_be_bytes());

    let (k, _) = iter.next_back().unwrap();
    assert_eq!(k, 0u64.to_be_bytes());

    let (k, _) = iter.next().unwrap();
    assert_eq!(k, 3u64.to_be_bytes());

    let (k, _) = iter.next_back().unwrap();
    assert_eq!(k, 1u64.to_be_bytes());

    let (k, _) = iter.next().unwrap();
    assert_eq!(k, 2u64.to_be_bytes());

    assert!(iter.next().is_none());
    assert!(iter.next_back().is_none());
}

#[test]
fn test_range_rev() {
    let dir = tempdir().unwrap();
    let tree = open(dir.path());
    populate(&tree, 20);

    // range [5,10).rev() on reversed → ascending: 5,6,7,8,9
    let start = 5u64.to_be_bytes();
    let end = 10u64.to_be_bytes();
    let entries: Vec<_> = tree.range(&start, &end).rev().collect();

    assert_eq!(entries.len(), 5);
    for (i, (key, _)) in entries.iter().enumerate() {
        let expected = i as u64 + 5;
        assert_eq!(*key, expected.to_be_bytes());
    }
}

// ---------------------------------------------------------------------------
// Reversed prefix iteration (explicit reversed=true)
// ---------------------------------------------------------------------------

#[test]
fn test_reversed_prefix_iter() {
    let dir = tempdir().unwrap();
    let mut config = Config::test();
    config.reversed = true;
    let tree: TypedTree<[u8; 16], u64, U64Codec> =
        TypedTree::open(dir.path(), config, U64Codec).unwrap();

    // Insert 5 entries with prefix [0,0,0,1,...]
    for i in 0u64..5 {
        let mut key = [0u8; 16];
        key[3] = 1;
        key[8..].copy_from_slice(&i.to_be_bytes());
        tree.put(&key, i).unwrap();
    }

    let entries: Vec<_> = tree.prefix_iter(&[0, 0, 0, 1]).collect();
    assert_eq!(entries.len(), 5);

    // Reversed: descending within prefix
    assert_eq!(*entries[0].1, 4);
    assert_eq!(*entries[4].1, 0);
}

#[test]
fn test_reversed_take() {
    let dir = tempdir().unwrap();
    let mut config = Config::test();
    config.reversed = true;
    let tree: TypedTree<[u8; 16], u64, U64Codec> =
        TypedTree::open(dir.path(), config, U64Codec).unwrap();

    for i in 0u64..100 {
        let mut key = [0u8; 16];
        key[3] = 1;
        key[8..].copy_from_slice(&i.to_be_bytes());
        tree.put(&key, i).unwrap();
    }

    // Take first 3 from reversed prefix = 3 newest
    let top3: Vec<_> = tree.prefix_iter(&[0, 0, 0, 1]).take(3).collect();
    assert_eq!(top3.len(), 3);
    assert_eq!(*top3[0].1, 99);
    assert_eq!(*top3[1].1, 98);
    assert_eq!(*top3[2].1, 97);
}

// ---------------------------------------------------------------------------
// Single element edge case
// ---------------------------------------------------------------------------

#[test]
fn test_single_element() {
    let dir = tempdir().unwrap();
    let tree = open(dir.path());
    tree.put(&1u64.to_be_bytes(), 42).unwrap();

    let entries: Vec<_> = tree.iter().collect();
    assert_eq!(entries.len(), 1);
    assert_eq!(*entries[0].1, 42);

    let rev: Vec<_> = tree.iter().rev().collect();
    assert_eq!(rev.len(), 1);
    assert_eq!(*rev[0].1, 42);
}

// ---------------------------------------------------------------------------
// Recovery (close + reopen)
// ---------------------------------------------------------------------------

#[test]
fn test_recovery() {
    let dir = tempdir().unwrap();

    {
        let tree = open(dir.path());
        populate(&tree, 50);
        tree.close().unwrap();
    }

    // Reopen
    let tree = open(dir.path());
    assert_eq!(tree.len(), 50);

    for i in 0..50u64 {
        let r = tree.get(&i.to_be_bytes()).unwrap();
        assert_eq!(*r, i * 10);
    }
}

#[test]
fn test_recovery_with_deletes() {
    let dir = tempdir().unwrap();

    {
        let tree = open(dir.path());
        populate(&tree, 20);
        for i in (0..20u64).step_by(2) {
            tree.delete(&i.to_be_bytes()).unwrap();
        }
        tree.close().unwrap();
    }

    let tree = open(dir.path());
    assert_eq!(tree.len(), 10);

    // Only odd keys remain
    for i in 0..20u64 {
        if i % 2 == 1 {
            assert!(tree.contains(&i.to_be_bytes()));
        } else {
            assert!(!tree.contains(&i.to_be_bytes()));
        }
    }
}

// ---------------------------------------------------------------------------
// Migrate
// ---------------------------------------------------------------------------

#[test]
fn test_migrate() {
    let dir = tempdir().unwrap();
    let tree = open(dir.path());
    populate(&tree, 10);

    let mutations = tree
        .migrate(|_key, value| {
            if *value < 50 {
                MigrateAction::Update(*value + 1000)
            } else {
                MigrateAction::Keep
            }
        })
        .unwrap();

    // Values 0,10,20,30,40 are < 50 → updated
    assert_eq!(mutations, 5);

    assert_eq!(*tree.get(&0u64.to_be_bytes()).unwrap(), 1000);
    assert_eq!(*tree.get(&3u64.to_be_bytes()).unwrap(), 1030);
    assert_eq!(*tree.get(&5u64.to_be_bytes()).unwrap(), 50); // unchanged
}

// ---------------------------------------------------------------------------
// Compact
// ---------------------------------------------------------------------------

#[test]
fn test_compact() {
    let dir = tempdir().unwrap();
    let tree = open(dir.path());

    // Insert then overwrite to create dead bytes
    populate(&tree, 100);
    for i in 0..100u64 {
        tree.put(&i.to_be_bytes(), i * 100).unwrap();
    }

    tree.flush_buffers().unwrap();
    let compacted = tree.compact().unwrap();
    // Compaction should process some entries
    let _ = compacted;

    // Data integrity preserved
    for i in 0..100u64 {
        assert_eq!(*tree.get(&i.to_be_bytes()).unwrap(), i * 100);
    }
}

// ---------------------------------------------------------------------------
// TypedRef lifetime
// ---------------------------------------------------------------------------

#[test]
fn test_typed_ref_deref() {
    let dir = tempdir().unwrap();
    let tree = open(dir.path());
    tree.put(&1u64.to_be_bytes(), 42).unwrap();

    let r = tree.get(&1u64.to_be_bytes()).unwrap();

    // TypedRef derefs to &u64
    let val: &u64 = &r;
    assert_eq!(*val, 42);

    // Can use in expressions
    assert_eq!(*r + 1, 43);
}

#[test]
fn test_put_old_ref_valid() {
    let dir = tempdir().unwrap();
    let tree = open(dir.path());

    tree.put(&1u64.to_be_bytes(), 100).unwrap();
    let old = tree.put(&1u64.to_be_bytes(), 200).unwrap().unwrap();

    // Old ref is still valid (seize guard alive)
    assert_eq!(*old, 100);

    // New value is accessible
    assert_eq!(*tree.get(&1u64.to_be_bytes()).unwrap(), 200);

    // Old ref still valid after new get
    assert_eq!(*old, 100);
}

#[test]
fn test_delete_old_ref_valid() {
    let dir = tempdir().unwrap();
    let tree = open(dir.path());

    tree.put(&1u64.to_be_bytes(), 42).unwrap();
    let old = tree.delete(&1u64.to_be_bytes()).unwrap().unwrap();

    // Deleted value still accessible through TypedRef
    assert_eq!(*old, 42);

    // Key gone
    assert!(tree.get(&1u64.to_be_bytes()).is_none());
}

// ---------------------------------------------------------------------------
// compare_delete
// ---------------------------------------------------------------------------

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

struct DeleteHook {
    writes: std::sync::Arc<AtomicUsize>,
    last_old: std::sync::Arc<std::sync::Mutex<Option<u64>>>,
    last_new: std::sync::Arc<std::sync::Mutex<Option<u64>>>,
}

impl Default for DeleteHook {
    fn default() -> Self {
        Self {
            writes: std::sync::Arc::new(AtomicUsize::new(0)),
            last_old: std::sync::Arc::new(std::sync::Mutex::new(None)),
            last_new: std::sync::Arc::new(std::sync::Mutex::new(None)),
        }
    }
}

impl armdb::TypedWriteHook<[u8; 8], u64> for DeleteHook {
    fn on_write(&self, _key: &[u8; 8], old: Option<&u64>, new: Option<&u64>) {
        self.writes.fetch_add(1, Ordering::Relaxed);
        *self.last_old.lock().unwrap() = old.copied();
        *self.last_new.lock().unwrap() = new.copied();
    }
}

#[test]
fn compare_delete_match_mismatch_absent() {
    let dir = tempdir().unwrap();
    let tree = open(dir.path());

    let k = 1u64.to_be_bytes();
    tree.put(&k, 42).unwrap();

    let err = tree.compare_delete(&k, &99).unwrap_err();
    assert!(matches!(err, DbError::CasMismatch));
    assert_eq!(tree.get(&k).map(|r| *r), Some(42));

    tree.compare_delete(&k, &42).unwrap();
    assert!(tree.get(&k).is_none());

    let err = tree.compare_delete(&k, &42).unwrap_err();
    assert!(matches!(err, DbError::KeyNotFound));
}

#[test]
fn compare_delete_fires_hook_and_reuses_delete_path() {
    let dir = tempdir().unwrap();
    let hook = DeleteHook::default();
    let writes = std::sync::Arc::clone(&hook.writes);
    let last_old = std::sync::Arc::clone(&hook.last_old);
    let last_new = std::sync::Arc::clone(&hook.last_new);
    let tree = TypedTree::<[u8; 8], u64, U64Codec, DeleteHook>::open_hooked(
        dir.path(),
        Config::test(),
        U64Codec,
        hook,
    )
    .unwrap();

    let k = 5u64.to_be_bytes();
    tree.put(&k, 100).unwrap();

    let writes_before = writes.load(Ordering::Relaxed);
    tree.compare_delete(&k, &100).unwrap();

    // exactly one on_write(Some(old), None) for the delete
    assert_eq!(writes.load(Ordering::Relaxed) - writes_before, 1);
    assert_eq!(*last_old.lock().unwrap(), Some(100));
    assert_eq!(*last_new.lock().unwrap(), None);

    // Reusing delete_locked_inner (retire: None) must not double-free:
    // subsequent operations on the same key/shard stay sound.
    tree.put(&k, 7).unwrap();
    assert_eq!(tree.get(&k).map(|r| *r), Some(7));
}