armdb 0.1.13

//! Follower apply trait.  Concrete impls for ConstTree<Fixed> and ConstMap<Fixed>
//! live in Tasks 22-23.

use std::mem::size_of;
use std::path::PathBuf;

use std::hash::Hash;

use crate::const_map::ConstMap;
use crate::const_tree::ConstTree;
use crate::durability::{Durability, Fixed};
use crate::error::DbResult;
use crate::fixed::slot::{self, STATUS_OCCUPIED, is_newer, status_of, version_of};
use crate::hook::WriteHook;
use crate::key::Key;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ApplyOutcome {
    /// Leader's version is newer — event applied.
    Applied,
    /// Own version ≥ incoming (modular compare) — event skipped.
    Skipped,
}

pub trait FixedReplicationTarget: Send + Sync {
    fn shard_count(&self) -> usize;
    fn key_len(&self) -> usize;
    fn value_len(&self) -> usize;
    fn shard_prefix_bits(&self) -> u8;

    /// Current number of OCCUPIED slots for `shard_id`. Used by the client
    /// to decide whether to set `FLAG_EMPTY_STATE` in SyncRequest.
    fn shard_occupied_count(&self, shard_id: u8) -> u32;

    /// Directory for per-shard cursor files.
    fn shard_dir(&self, shard_id: u8) -> PathBuf;

    /// Grow the shard to at least `min_slot_count` slots.  No-op when follower
    /// already has more slots.
    fn grow_shard_to(&self, shard_id: u8, min_slot_count: u32) -> DbResult<()>;

    /// Apply an OCCUPIED slot event. Handles modular version compare,
    /// stale-cleanup of old key (rare path), slot pwrite, and index upsert.
    fn apply_occupied(
        &self,
        shard_id: u8,
        slot_id: u32,
        meta: u32,
        key: &[u8],
        value: &[u8],
    ) -> DbResult<ApplyOutcome>;

    /// Apply a DELETED slot event.  Conditional `index.remove` — only if
    /// the index still points to this slot_id.
    fn apply_deleted(
        &self,
        shard_id: u8,
        slot_id: u32,
        meta: u32,
        key: &[u8],
    ) -> DbResult<ApplyOutcome>;
}

// ============================================================================
// Impl for ConstTree<K, V, H, Fixed>
// ============================================================================

impl<K, const V: usize, H> FixedReplicationTarget for ConstTree<K, V, H, Fixed>
where
    K: Key + 'static,
    H: WriteHook<K> + 'static,
{
    fn shard_count(&self) -> usize {
        self.fixed_durability().shard_count()
    }

    fn key_len(&self) -> usize {
        size_of::<K>()
    }

    fn value_len(&self) -> usize {
        V
    }

    fn shard_prefix_bits(&self) -> u8 {
        self.fixed_durability().shard_prefix_bits() as u8
    }

    fn shard_occupied_count(&self, shard_id: u8) -> u32 {
        let inner = self.fixed_durability().lock_shard(shard_id as usize);
        inner.bitmap.occupied()
    }

    fn shard_dir(&self, shard_id: u8) -> PathBuf {
        self.fixed_durability()
            .engine
            .path()
            .join(format!("shard_{:03}", shard_id))
    }

    fn grow_shard_to(&self, shard_id: u8, min_slot_count: u32) -> DbResult<()> {
        let mut inner = self.fixed_durability().lock_shard(shard_id as usize);
        while inner.slot_count() < min_slot_count {
            inner.grow()?;
        }
        if inner.slot_count() > min_slot_count {
            tracing::warn!(
                shard_id,
                follower_count = inner.slot_count(),
                leader_count = min_slot_count,
                "follower has more slots than leader — ignoring (no truncate)"
            );
        }
        Ok(())
    }

    fn apply_occupied(
        &self,
        shard_id: u8,
        slot_id: u32,
        meta: u32,
        key: &[u8],
        value: &[u8],
    ) -> DbResult<ApplyOutcome> {
        debug_assert_eq!(key.len(), size_of::<K>());
        debug_assert_eq!(value.len(), V);

        // Decode key + value.
        let key_typed: K = K::from_bytes(key);
        let mut value_arr = [0u8; V];
        value_arr.copy_from_slice(value);

        let mut inner = self.fixed_durability().lock_shard(shard_id as usize);

        // Grow if needed (leader advanced slot_count after handshake).
        while inner.slot_count() <= slot_id {
            inner.grow()?;
        }

        // Modular version compare.
        let old_meta = inner.versions[slot_id as usize];
        if !is_newer(version_of(meta), version_of(old_meta)) {
            metrics::counter!(
                "armdb.fixed.skipped_events",
                "shard" => shard_id.to_string()
            )
            .increment(1);
            return Ok(ApplyOutcome::Skipped);
        }

        // ── Step 1: stale cleanup of old key (Bug 2 fix) ──
        // If prior slot state was OCCUPIED, a different key_X may be sitting
        // in the index pointing at this slot_id.  We must remove that stale
        // mapping. Fast-path: if new_key already points here, old_key == new_key.
        if status_of(old_meta) == STATUS_OCCUPIED && self.get_slot_id(&key_typed) != Some(slot_id) {
            let buf = inner.read_slot_header_and_key(slot_id, size_of::<K>())?;
            let old_key_bytes =
                &buf[slot::SLOT_HEADER_SIZE..slot::SLOT_HEADER_SIZE + size_of::<K>()];
            let old_key = K::from_bytes(old_key_bytes);
            if old_key.as_bytes() != key_typed.as_bytes()
                && self.remove_key_if_slot_matches(&old_key, slot_id)
            {
                metrics::counter!(
                    "armdb.fixed.stale_cleanup_count",
                    "shard" => shard_id.to_string()
                )
                .increment(1);
            }
        }

        // ── Step 2: persist slot ──
        inner.apply_foreign_slot(slot_id, meta, key, value)?;
        inner.bitmap.set(slot_id);

        // Drop shard lock before touching the lock-free index — avoids holding
        // the mutex across SkipList work.
        drop(inner);

        // ── Step 3: upsert index ──
        self.upsert_replicated(&key_typed, value_arr, slot_id);

        metrics::counter!(
            "armdb.fixed.applied_events",
            "shard" => shard_id.to_string(),
            "kind" => "occupied"
        )
        .increment(1);
        Ok(ApplyOutcome::Applied)
    }

    fn apply_deleted(
        &self,
        shard_id: u8,
        slot_id: u32,
        meta: u32,
        key: &[u8],
    ) -> DbResult<ApplyOutcome> {
        debug_assert_eq!(key.len(), size_of::<K>());
        let key_typed: K = K::from_bytes(key);

        let mut inner = self.fixed_durability().lock_shard(shard_id as usize);
        while inner.slot_count() <= slot_id {
            inner.grow()?;
        }

        let old_meta = inner.versions[slot_id as usize];
        if !is_newer(version_of(meta), version_of(old_meta)) {
            metrics::counter!(
                "armdb.fixed.skipped_events",
                "shard" => shard_id.to_string()
            )
            .increment(1);
            return Ok(ApplyOutcome::Skipped);
        }

        inner.apply_foreign_delete(slot_id, meta)?;
        inner.bitmap.clear(slot_id);

        drop(inner);

        // ── Bug 1 fix: only remove if index still points to this slot. ──
        // The key may have "moved" to a different slot (DELETE + INSERT
        // with reuse); we must not destroy a valid mapping.
        let _removed = self.remove_key_if_slot_matches(&key_typed, slot_id);

        metrics::counter!(
            "armdb.fixed.applied_events",
            "shard" => shard_id.to_string(),
            "kind" => "deleted"
        )
        .increment(1);
        Ok(ApplyOutcome::Applied)
    }
}

// ============================================================================
// Impl for ConstMap<K, V, H, Fixed>
// ============================================================================

impl<K, const V: usize, H> FixedReplicationTarget for ConstMap<K, V, H, Fixed>
where
    K: Key + Send + Sync + Hash + Eq + 'static,
    H: WriteHook<K> + 'static,
{
    fn shard_count(&self) -> usize {
        self.fixed_durability().shard_count()
    }

    fn key_len(&self) -> usize {
        size_of::<K>()
    }

    fn value_len(&self) -> usize {
        V
    }

    fn shard_prefix_bits(&self) -> u8 {
        self.fixed_durability().shard_prefix_bits() as u8
    }

    fn shard_occupied_count(&self, shard_id: u8) -> u32 {
        let inner = self.fixed_durability().lock_shard(shard_id as usize);
        inner.bitmap.occupied()
    }

    fn shard_dir(&self, shard_id: u8) -> PathBuf {
        self.fixed_durability()
            .engine
            .path()
            .join(format!("shard_{:03}", shard_id))
    }

    fn grow_shard_to(&self, shard_id: u8, min_slot_count: u32) -> DbResult<()> {
        let mut inner = self.fixed_durability().lock_shard(shard_id as usize);
        while inner.slot_count() < min_slot_count {
            inner.grow()?;
        }
        if inner.slot_count() > min_slot_count {
            tracing::warn!(
                shard_id,
                follower_count = inner.slot_count(),
                leader_count = min_slot_count,
                "follower has more slots than leader — ignoring (no truncate)"
            );
        }
        Ok(())
    }

    fn apply_occupied(
        &self,
        shard_id: u8,
        slot_id: u32,
        meta: u32,
        key: &[u8],
        value: &[u8],
    ) -> DbResult<ApplyOutcome> {
        debug_assert_eq!(key.len(), size_of::<K>());
        debug_assert_eq!(value.len(), V);

        // Decode key + value.
        let key_typed: K = K::from_bytes(key);
        let mut value_arr = [0u8; V];
        value_arr.copy_from_slice(value);

        let mut inner = self.fixed_durability().lock_shard(shard_id as usize);

        // Grow if needed (leader advanced slot_count after handshake).
        while inner.slot_count() <= slot_id {
            inner.grow()?;
        }

        // Modular version compare.
        let old_meta = inner.versions[slot_id as usize];
        if !is_newer(version_of(meta), version_of(old_meta)) {
            metrics::counter!(
                "armdb.fixed.skipped_events",
                "shard" => shard_id.to_string()
            )
            .increment(1);
            return Ok(ApplyOutcome::Skipped);
        }

        // ── Step 1: stale cleanup of old key (Bug 2 fix) ──
        // If prior slot state was OCCUPIED, a different key_X may be sitting
        // in the index pointing at this slot_id.  We must remove that stale
        // mapping. Fast-path: if new_key already points here, old_key == new_key.
        if status_of(old_meta) == STATUS_OCCUPIED && self.get_slot_id(&key_typed) != Some(slot_id) {
            let buf = inner.read_slot_header_and_key(slot_id, size_of::<K>())?;
            let old_key_bytes =
                &buf[slot::SLOT_HEADER_SIZE..slot::SLOT_HEADER_SIZE + size_of::<K>()];
            let old_key = K::from_bytes(old_key_bytes);
            if old_key.as_bytes() != key_typed.as_bytes()
                && self.remove_key_if_slot_matches(&old_key, slot_id)
            {
                metrics::counter!(
                    "armdb.fixed.stale_cleanup_count",
                    "shard" => shard_id.to_string()
                )
                .increment(1);
            }
        }

        // ── Step 2: persist slot ──
        inner.apply_foreign_slot(slot_id, meta, key, value)?;
        inner.bitmap.set(slot_id);

        // Drop shard lock before touching the lock-free index — avoids holding
        // the mutex across HashMap work.
        drop(inner);

        // ── Step 3: upsert index ──
        self.upsert_replicated(&key_typed, value_arr, slot_id);

        metrics::counter!(
            "armdb.fixed.applied_events",
            "shard" => shard_id.to_string(),
            "kind" => "occupied"
        )
        .increment(1);
        Ok(ApplyOutcome::Applied)
    }

    fn apply_deleted(
        &self,
        shard_id: u8,
        slot_id: u32,
        meta: u32,
        key: &[u8],
    ) -> DbResult<ApplyOutcome> {
        debug_assert_eq!(key.len(), size_of::<K>());
        let key_typed: K = K::from_bytes(key);

        let mut inner = self.fixed_durability().lock_shard(shard_id as usize);
        while inner.slot_count() <= slot_id {
            inner.grow()?;
        }

        let old_meta = inner.versions[slot_id as usize];
        if !is_newer(version_of(meta), version_of(old_meta)) {
            metrics::counter!(
                "armdb.fixed.skipped_events",
                "shard" => shard_id.to_string()
            )
            .increment(1);
            return Ok(ApplyOutcome::Skipped);
        }

        inner.apply_foreign_delete(slot_id, meta)?;
        inner.bitmap.clear(slot_id);

        drop(inner);

        // ── Bug 1 fix: only remove if index still points to this slot. ──
        // The key may have "moved" to a different slot (DELETE + INSERT
        // with reuse); we must not destroy a valid mapping.
        let _removed = self.remove_key_if_slot_matches(&key_typed, slot_id);

        metrics::counter!(
            "armdb.fixed.applied_events",
            "shard" => shard_id.to_string(),
            "kind" => "deleted"
        )
        .increment(1);
        Ok(ApplyOutcome::Applied)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::FixedConfig;
    use crate::fixed::FixedTree;
    use crate::fixed::slot::{STATUS_OCCUPIED, pack_meta};
    use tempfile::tempdir;

    fn cfg() -> FixedConfig {
        FixedConfig {
            shard_count: 1,
            grow_step: 64,
            ..FixedConfig::test()
        }
    }

    /// Key/value types are `[u8; 8]` so the key matches `size_of::<K>() == 8`
    /// and `V = 8` bytes — smallest tree that still exercises the full path.
    type T = FixedTree<[u8; 8], 8>;

    fn shard_of(tree: &T, key: &[u8; 8]) -> usize {
        tree.shard_for(key)
    }

    #[test]
    fn test_apply_occupied_insert() {
        let dir = tempdir().unwrap();
        let tree: T = FixedTree::open(dir.path(), cfg()).unwrap();

        let key = 1u64.to_be_bytes();
        let value = 42u64.to_be_bytes();
        let meta = pack_meta(STATUS_OCCUPIED, 1);
        let shard = shard_of(&tree, &key) as u8;

        let out = tree.apply_occupied(shard, 5, meta, &key, &value).unwrap();
        assert_eq!(out, ApplyOutcome::Applied);

        // Index must reflect the new mapping.
        assert_eq!(tree.get(&key), Some(value));
        assert_eq!(tree.get_slot_id(&key), Some(5));
    }

    #[test]
    fn test_apply_deleted_bug1_moved_key() {
        // Bug 1: an old DELETE event for `key` at slot_A must NOT wipe a newer
        // mapping of `key` at slot_B that was installed after a reuse cycle.
        let dir = tempdir().unwrap();
        let tree: T = FixedTree::open(dir.path(), cfg()).unwrap();

        let key = 7u64.to_be_bytes();
        let value = 7u64.to_be_bytes();
        let shard = shard_of(&tree, &key) as u8;

        // Install key at slot_B first by replicating an OCCUPIED event.
        let occ = pack_meta(STATUS_OCCUPIED, 10);
        tree.apply_occupied(shard, 13, occ, &key, &value).unwrap();
        assert_eq!(tree.get_slot_id(&key), Some(13));

        // Now replay an older DELETE event for the SAME key on a DIFFERENT slot.
        // The index points at slot 13, so the DELETE targeting slot 5 must be a
        // no-op on the index (mapping preserved).
        //
        // We need versions[5] to have a lower version than `meta` so the modular
        // compare lets the event through. versions[5] starts at 0 → `is_newer(v, 0)`
        // is true when 0 < v < 2^29.
        let del = pack_meta(crate::fixed::slot::STATUS_DELETED, 1);
        let out = tree.apply_deleted(shard, 5, del, &key).unwrap();
        assert_eq!(out, ApplyOutcome::Applied);

        // Mapping untouched: still at slot 13.
        assert_eq!(
            tree.get_slot_id(&key),
            Some(13),
            "Bug 1 regression: stale DELETE removed a live mapping"
        );
        assert_eq!(tree.get(&key), Some(value));
    }

    #[test]
    fn test_apply_occupied_bug2_stale_cleanup() {
        // Bug 2: an OCCUPIED event overwrites slot_A with key_Y. The follower's
        // index currently has key_X → slot_A (stale) and key_Y → slot_C (live
        // from an older put). After apply_occupied, key_X must be removed and
        // key_Y must point at slot_A.
        let dir = tempdir().unwrap();
        let tree: T = FixedTree::open(dir.path(), cfg()).unwrap();

        // Pick keys that route to the same shard (shard_count=1 so always shard 0).
        let key_x = 100u64.to_be_bytes();
        let key_y = 200u64.to_be_bytes();
        let val_x = 11u64.to_be_bytes();
        let val_y = 22u64.to_be_bytes();
        let val_y_new = 33u64.to_be_bytes();
        let shard = 0u8;

        // Seed: install key_X → slot 3, key_Y → slot 7 via replicated events
        // (mirrors what a previous leader state would have produced).
        let occ_x = pack_meta(STATUS_OCCUPIED, 2);
        tree.apply_occupied(shard, 3, occ_x, &key_x, &val_x)
            .unwrap();
        let occ_y = pack_meta(STATUS_OCCUPIED, 4);
        tree.apply_occupied(shard, 7, occ_y, &key_y, &val_y)
            .unwrap();
        assert_eq!(tree.get_slot_id(&key_x), Some(3));
        assert_eq!(tree.get_slot_id(&key_y), Some(7));

        // Now replay an OCCUPIED event that reassigns slot 3 to key_Y.
        // Needs a version strictly greater (modular) than versions[3] (=2).
        let occ_reassign = pack_meta(STATUS_OCCUPIED, 5);
        let out = tree
            .apply_occupied(shard, 3, occ_reassign, &key_y, &val_y_new)
            .unwrap();
        assert_eq!(out, ApplyOutcome::Applied);

        // Bug 2 regression check: key_X must be gone from the index.
        assert_eq!(
            tree.get_slot_id(&key_x),
            None,
            "Bug 2 regression: stale key_X mapping to slot_A not cleaned up"
        );
        assert!(!tree.contains(&key_x));

        // key_Y now points at slot 3 with the new value.
        assert_eq!(tree.get_slot_id(&key_y), Some(3));
        assert_eq!(tree.get(&key_y), Some(val_y_new));
    }

    #[test]
    fn test_apply_modular_skipped() {
        let dir = tempdir().unwrap();
        let tree: T = FixedTree::open(dir.path(), cfg()).unwrap();

        let key = 9u64.to_be_bytes();
        let value = 9u64.to_be_bytes();
        let shard = shard_of(&tree, &key) as u8;

        // Establish versions[slot] at a high version.
        let high = pack_meta(STATUS_OCCUPIED, 100);
        tree.apply_occupied(shard, 0, high, &key, &value).unwrap();

        // Replay an OLDER event (lower version) — must be Skipped.
        let low = pack_meta(STATUS_OCCUPIED, 50);
        let out = tree.apply_occupied(shard, 0, low, &key, &value).unwrap();
        assert_eq!(out, ApplyOutcome::Skipped);

        // Same for DELETE.
        let low_del = pack_meta(crate::fixed::slot::STATUS_DELETED, 50);
        let out = tree.apply_deleted(shard, 0, low_del, &key).unwrap();
        assert_eq!(out, ApplyOutcome::Skipped);

        // State unchanged.
        assert_eq!(tree.get(&key), Some(value));
    }

    // -------------------------------------------------------------------------
    // ConstMap tests
    // -------------------------------------------------------------------------

    type M = crate::fixed::FixedMap<[u8; 8], 8>;

    #[test]
    fn test_map_apply_occupied_and_deleted() {
        // Confirm the ConstMap impl compiles and the basic apply_occupied + get
        // + apply_deleted round-trip works correctly.
        let dir = tempdir().unwrap();
        let map: M = M::open(dir.path(), cfg()).unwrap();

        let key = 42u64.to_be_bytes();
        let value = 99u64.to_be_bytes();
        let shard = map.shard_for(&key) as u8;
        let meta_occ = pack_meta(STATUS_OCCUPIED, 1);

        // apply_occupied → value readable via get().
        let out = map
            .apply_occupied(shard, 0, meta_occ, &key, &value)
            .unwrap();
        assert_eq!(out, ApplyOutcome::Applied);
        assert_eq!(map.get(&key), Some(value));
        assert_eq!(map.get_slot_id(&key), Some(0));

        // apply_deleted → key removed from index.
        let meta_del = pack_meta(crate::fixed::slot::STATUS_DELETED, 2);
        let out = map.apply_deleted(shard, 0, meta_del, &key).unwrap();
        assert_eq!(out, ApplyOutcome::Applied);
        assert_eq!(map.get(&key), None);
    }
}