armdb 0.5.3

use std::hash::Hash;
use std::mem::size_of;
use std::ops::Bound;

use crate::Key;

use crate::codec::Codec;
use crate::compaction::{CompactionIndex, compact_shard};
use crate::config::Config;
use crate::disk_loc::DiskLoc;
use crate::engine::Engine;
use crate::error::{DbError, DbResult};
use crate::hook::{NoHook, TypedWriteHook};
use crate::map_index::{
    Dir, KeyPage, MapIter, MapIterView, OrdKey, Page, ShardIndex, ShardSource, dir_from,
    key_index_batch,
};
use crate::recovery::recover_typed_map;
use crate::shard::ShardInner;
use crate::skiplist::node::TypedData;
use crate::sync::{self, Mutex, MutexGuard};
use crate::typed_tree::TypedRef;

// ---------------------------------------------------------------------------
// TypedMapEntry
// ---------------------------------------------------------------------------

pub(crate) struct TypedMapEntry<T> {
    pub(crate) ptr: *mut TypedData<T>,
}

// SAFETY: TypedData<T> is Send+Sync when T is, and we manage the pointer lifetime via seize.
unsafe impl<T: Send> Send for TypedMapEntry<T> {}
unsafe impl<T: Sync> Sync for TypedMapEntry<T> {}

/// Deferred hook/retire work returned by _locked methods.
/// Caller drops shard + index locks, fires hook, then retires pointers.
struct MutationResult<T> {
    /// Raw pointer for TypedRef return. None if key was new / not found.
    ref_ptr: Option<*const TypedData<T>>,
    /// Hook args: (old_value_ptr, new_value_ptr). None = don't fire hook.
    hook: Option<(Option<*const T>, Option<*const T>)>,
    /// Pointer to retire via seize after hook. None = nothing to retire.
    retire: Option<*mut TypedData<T>>,
}

// ---------------------------------------------------------------------------
// TypedMap
// ---------------------------------------------------------------------------

/// A map with fixed-size keys and typed values `T`. Values are encoded via
/// a [`Codec`] for disk persistence but stored as `T` in memory — reads never
/// touch disk and return [`TypedRef<T>`](TypedRef) (guard-protected reference).
///
/// Each `TypedMap` owns its storage engine — one map = one database directory.
///
/// # Clone-free reads
///
/// Returns [`TypedRef<T>`](TypedRef) — a guard-protected reference. The `seize` guard
/// prevents reclamation of the old data while the reference exists.
/// `put()`, `delete()`, and `update()` return the old value **without requiring `T: Clone`**.
/// Ordered iteration ([`iter_view`](Self::iter_view) → `iter`/`range`/`keys`) is likewise
/// clone-free; keys-only iteration never even dereferences `T`.
///
/// The methods that require `T: Clone` are [`compact()`](Self::compact) and the
/// value-collecting [`paginate`](TypedMapIterView::paginate) (keys-only
/// `keys_paginate` does not).
///
/// # Write hooks
///
/// Uses [`TypedWriteHook<K, T>`] — the hook receives `&T` directly (not encoded bytes).
/// `on_write` fires on `put`/`insert`/`delete`/`cas`/`compare_delete`/`update`
/// and on each mutation inside `atomic()` (replayed after the lock is released, in application order).
/// Old value is always provided (it lives in
/// memory) — `NEEDS_OLD_VALUE` is ignored.
///
/// `on_init` fires once per live entry during `migrate()` or `replay_init()`
/// (enable via `NEEDS_INIT = true`).
///
/// # Ordered iteration
///
/// HashMap-based, O(1) lookup. Ordered scans are available when opened with
/// `iterable(true)` via [`iter_view`](Self::iter_view). Use [`TypedTree`](crate::TypedTree)
/// if you need prefix/range scans without the iterable flag.
///
/// # Usage
///
/// ```ignore
/// let map = TypedMap::<[u8; 16], MyValue, RapiraCodec>::open(
///     "data/sessions",
///     Config::balanced().build(),
///     RapiraCodec,
/// )?;
/// map.put(&key, value)?;
/// if let Some(r) = map.get(&key) {
///     println!("{:?}", &*r);  // TypedRef<MyValue> derefs to &MyValue
/// }
/// map.close()?;
/// ```
pub struct TypedMap<
    K: Key + Send + Sync + Hash + Eq,
    T: Send + Sync,
    C: Codec<T>,
    H: TypedWriteHook<K, T> = NoHook,
> {
    indexes: Vec<Mutex<ShardIndex<K, TypedMapEntry<T>>>>,
    collector: seize::Collector,
    engine: Engine,
    codec: C,
    compaction_threshold: f64,
    shard_prefix_bits: usize,
    hook: H,
    iterable: bool,
    reversed: bool,
    len: std::sync::atomic::AtomicUsize,
    encode_bufs: Vec<Mutex<Vec<u8>>>,
}

impl<K: Key + Send + Sync + Hash + Eq, T: Send + Sync, C: Codec<T> + Sync> TypedMap<K, T, C> {
    /// Open or create a `TypedMap` at the given path.
    /// Recovers the index from existing data files on disk.
    pub fn open(path: impl AsRef<std::path::Path>, config: Config, codec: C) -> DbResult<Self> {
        Self::open_hooked(path, config, codec, NoHook)
    }
}

impl<K: Key + Send + Sync + Hash + Eq, T: Send + Sync, C: Codec<T> + Sync, H: TypedWriteHook<K, T>>
    TypedMap<K, T, C, H>
{
    /// Open or create a `TypedMap` with a write hook for secondary index maintenance.
    pub fn open_hooked(
        path: impl AsRef<std::path::Path>,
        config: Config,
        codec: C,
        hook: H,
    ) -> DbResult<Self> {
        let config = config.with_resolved_hints(false);
        Self::open_inner(path, config, codec, hook)
    }

    fn open_inner(
        path: impl AsRef<std::path::Path>,
        config: Config,
        codec: C,
        hook: H,
    ) -> DbResult<Self> {
        let compaction_threshold = config.compaction_threshold;
        let shard_prefix_bits = config.shard_prefix_bits;
        let iterable = config.iterable;
        let reversed = config.reversed;
        let engine = Engine::open(path, config)?;

        let shard_count = engine.shards().len();
        let mut indexes = Vec::with_capacity(shard_count);
        for _ in 0..shard_count {
            // Built once after recovery via `set_iterable_and_rebuild` — see the
            // ConstMap `open_inner` for the rationale.
            indexes.push(Mutex::new(ShardIndex::new(false)));
        }

        let encode_bufs: Vec<Mutex<Vec<u8>>> =
            (0..shard_count).map(|_| Mutex::new(Vec::new())).collect();

        let map = Self {
            indexes,
            collector: seize::Collector::new(),
            engine,
            codec,
            compaction_threshold,
            shard_prefix_bits,
            hook,
            iterable,
            reversed,
            len: std::sync::atomic::AtomicUsize::new(0),
            encode_bufs,
        };

        // Recover index from disk
        let shard_dirs = map.engine.shard_dirs();
        let shard_dir_refs = Engine::shard_dir_refs(&shard_dirs);
        let shard_ids = map.engine.shard_ids();

        let hints = map.engine.hints();
        let outcome = recover_typed_map::<K, T, C>(
            &shard_dir_refs,
            &shard_ids,
            map.indexes(),
            &map.codec,
            hints,
            #[cfg(feature = "encryption")]
            map.engine.cipher(),
        )?;
        for tail in &outcome.active_tails {
            map.engine.shards()[tail.shard_idx].apply_recovery_tail(tail)?;
        }
        for (shard_idx, dead) in outcome.shard_dead_bytes {
            map.engine.shards()[shard_idx].install_dead_bytes(dead);
        }
        let max_gsn = outcome.max_gsn;

        map.engine
            .gsn()
            .fetch_max(max_gsn + 1, std::sync::atomic::Ordering::Relaxed);
        if hints {
            for shard in map.engine.shards().iter() {
                shard.set_key_len(size_of::<K>());
            }
        }
        if iterable {
            for idx in &map.indexes {
                sync::lock(idx).set_iterable_and_rebuild();
            }
        }
        let initial_len: usize = map.indexes.iter().map(|m| sync::lock(m).len()).sum();
        map.len
            .store(initial_len, std::sync::atomic::Ordering::Relaxed);

        tracing::info!(
            key_size = size_of::<K>(),
            entries = initial_len,
            "typed_map recovered"
        );

        Ok(map)
    }

    /// Graceful shutdown: write hint files (if enabled), flush write buffers + fsync.
    pub fn close(self) -> DbResult<()> {
        if self.engine.hints() {
            self.sync_hints()?;
        }
        self.engine.flush()
    }

    /// Flush all shard write buffers to disk (without fsync).
    pub fn flush_buffers(&self) -> DbResult<()> {
        self.engine.flush_buffers()
    }

    /// Get the database configuration.
    pub fn config(&self) -> &Config {
        self.engine.config()
    }
}

impl<
    K: Key + Send + Sync + Hash + Eq,
    T: Clone + Send + Sync,
    C: Codec<T> + Sync,
    H: TypedWriteHook<K, T>,
> CompactionIndex<K> for TypedMap<K, T, C, H>
{
    fn update_if_match(&self, key: &K, old_loc: DiskLoc, new_loc: DiskLoc) -> bool {
        let mut index = sync::lock(&self.indexes[self.shard_for(key)]);
        if let Some(entry) = index.get_mut(key) {
            let data = unsafe { &*entry.ptr };
            if data.disk == old_loc {
                let new_data = Box::into_raw(Box::new(TypedData {
                    disk: new_loc,
                    value: data.value.clone(),
                }));
                let old_ptr = entry.ptr;
                entry.ptr = new_data;
                unsafe {
                    self.collector
                        .retire(old_ptr, seize::reclaim::boxed::<TypedData<T>>);
                }
                return true;
            }
        }
        false
    }

    fn contains_key(&self, key: &K) -> bool {
        self.contains(key)
    }

    fn is_live(&self, shard_id: u8, key: &K, loc: DiskLoc) -> bool {
        let index = sync::lock(&self.indexes[shard_id as usize]);
        match index.get(key) {
            Some(entry) => unsafe { &*entry.ptr }.disk == loc,
            None => false,
        }
    }
}

impl<K: Key + Send + Sync + Hash + Eq, T: Send + Sync, C: Codec<T> + Sync, H: TypedWriteHook<K, T>>
    TypedMap<K, T, C, H>
{
    /// Trigger a compaction pass across all shards.
    pub fn compact(&self) -> DbResult<usize>
    where
        T: Clone,
    {
        let mut total_compacted = 0;
        for shard in self.engine.shards().iter() {
            total_compacted += compact_shard(shard, self, self.compaction_threshold)?;
        }
        Ok(total_compacted)
    }

    // -- Reads ----------------------------------------------------------------

    /// Get a guard-protected reference to a value by key. O(1) lookup.
    /// Brief shard lock to read pointer, then lock-free access via seize guard.
    pub fn get(&self, key: &K) -> Option<TypedRef<'_, T>> {
        metrics::counter!("armdb.ops", "op" => "get", "tree" => "typed_map").increment(1);
        let guard = self.collector.enter();
        let data_ptr = {
            let index = sync::lock(&self.indexes[self.shard_for(key)]);
            let entry = index.get(key)?;
            entry.ptr as *const TypedData<T>
        };
        Some(TypedRef::new(guard, data_ptr))
    }

    /// Get a guard-protected reference to a value by key, returning `Err(KeyNotFound)` if absent.
    pub fn get_or_err(&self, key: &K) -> DbResult<TypedRef<'_, T>> {
        self.get(key).ok_or(DbError::KeyNotFound)
    }

    /// Check if a key exists.
    pub fn contains(&self, key: &K) -> bool {
        let index = sync::lock(&self.indexes[self.shard_for(key)]);
        index.contains_key(key)
    }

    // -- Writes ---------------------------------------------------------------

    /// Insert or update a key-value pair. Returns a `TypedRef` to the old value
    /// if the key existed (valid while the guard lives).
    pub fn put(&self, key: &K, value: T) -> DbResult<Option<TypedRef<'_, T>>> {
        metrics::counter!("armdb.ops", "op" => "put", "tree" => "typed_map").increment(1);
        let shard_id = self.shard_for(key);
        let guard = self.collector.enter();
        let mr = {
            let mut inner = self.engine.shards()[shard_id].lock();
            let mut index = sync::lock(&self.indexes[shard_id]);
            self.put_locked_inner::<true>(shard_id, &mut inner, &mut index, &guard, key, value)?
        };
        self.fire_hook_and_retire(key, &mr);
        Ok(mr.ref_ptr.map(|p| TypedRef::new(guard, p)))
    }

    /// Insert a key-value pair only if the key does not exist.
    /// Returns `Err(KeyExists)` if the key is already present.
    pub fn insert(&self, key: &K, value: T) -> DbResult<()> {
        metrics::counter!("armdb.ops", "op" => "insert", "tree" => "typed_map").increment(1);
        let shard_id = self.shard_for(key);
        let guard = self.collector.enter();
        let mr = {
            let mut inner = self.engine.shards()[shard_id].lock();
            let mut index = sync::lock(&self.indexes[shard_id]);
            self.insert_locked::<true>(shard_id, &mut inner, &mut index, &guard, key, value)?
        };
        self.fire_hook_and_retire(key, &mr);
        Ok(())
    }

    /// Delete a key. Returns a `TypedRef` to the old value if the key existed.
    pub fn delete(&self, key: &K) -> DbResult<Option<TypedRef<'_, T>>> {
        metrics::counter!("armdb.ops", "op" => "delete", "tree" => "typed_map").increment(1);
        let shard_id = self.shard_for(key);
        let guard = self.collector.enter();
        let mr = {
            let mut inner = self.engine.shards()[shard_id].lock();
            let mut index = sync::lock(&self.indexes[shard_id]);
            self.delete_locked_inner::<true>(shard_id, &mut inner, &mut index, &guard, key)?
        };
        self.fire_hook_and_retire(key, &mr);
        Ok(mr.ref_ptr.map(|p| TypedRef::new(guard, p)))
    }

    /// Compare-and-swap: if current value == expected, replace with new_value.
    /// Returns `Ok(())` on success, `Err(CasMismatch)` if current != expected,
    /// `Err(KeyNotFound)` if key doesn't exist.
    pub fn cas(&self, key: &K, expected: &T, new_value: T) -> DbResult<()>
    where
        T: PartialEq,
    {
        metrics::counter!("armdb.ops", "op" => "cas", "tree" => "typed_map").increment(1);
        let shard_id = self.shard_for(key);
        let _guard = self.collector.enter();

        let (hook_old, hook_new, retire_ptr) = {
            let mut inner = self.engine.shards()[shard_id].lock();
            let mut index = sync::lock(&self.indexes[shard_id]);

            let entry = index.get(key).ok_or(DbError::KeyNotFound)?;
            let current_data = unsafe { &*entry.ptr };
            if current_data.value != *expected {
                return Err(DbError::CasMismatch);
            }

            let mut buf = sync::lock(&self.encode_bufs[shard_id]);
            self.codec.encode_to(&new_value, &mut buf)?;
            let (disk_loc, _gsn) =
                inner.append_entry(shard_id as u8, key.as_bytes(), &buf, false)?;

            let new_data = Box::new(TypedData {
                disk: disk_loc,
                value: new_value,
            });

            let new_data_ptr = Box::into_raw(new_data);
            let old_ptr = entry.ptr;
            index.get_mut(key).expect("key exists").ptr = new_data_ptr;

            let old_disk = unsafe { (*old_ptr).disk };
            inner.add_dead_bytes(
                old_disk.file_id,
                crate::entry::entry_size(size_of::<K>(), old_disk.len),
            );

            (
                &current_data.value as *const T,
                unsafe { &(*new_data_ptr).value as *const T },
                old_ptr,
            )
        };

        unsafe {
            self.hook.on_write(key, Some(&*hook_old), Some(&*hook_new));
            self.collector
                .retire(retire_ptr, seize::reclaim::boxed::<TypedData<T>>);
        }

        Ok(())
    }

    /// Compare-and-delete: if the current value == `expected`, delete the key.
    /// Returns `Ok(())` on success, `Err(CasMismatch)` if current != expected,
    /// `Err(KeyNotFound)` if the key doesn't exist.
    pub fn compare_delete(&self, key: &K, expected: &T) -> DbResult<()>
    where
        T: PartialEq,
    {
        metrics::counter!("armdb.ops", "op" => "compare_delete", "tree" => "typed_map")
            .increment(1);
        let shard_id = self.shard_for(key);
        let guard = self.collector.enter();

        let mr = {
            let mut inner = self.engine.shards()[shard_id].lock();
            let mut index = sync::lock(&self.indexes[shard_id]);

            let entry = index.get(key).ok_or(DbError::KeyNotFound)?;
            let current_data = unsafe { &*entry.ptr };
            if current_data.value != *expected {
                return Err(DbError::CasMismatch);
            }

            self.delete_locked_inner::<true>(shard_id, &mut inner, &mut index, &guard, key)?
        };

        self.fire_hook_and_retire(key, &mr);
        Ok(())
    }

    /// Compare-and-delete by pointer identity: delete only if the live entry still
    /// points at `expected`. Used by [`TypedMapIterView::retain`] without requiring `T: PartialEq`.
    pub(crate) fn compare_delete_if_ptr(
        &self,
        key: &K,
        expected: *const TypedData<T>,
    ) -> DbResult<()> {
        metrics::counter!("armdb.ops", "op" => "compare_delete", "tree" => "typed_map")
            .increment(1);
        let shard_id = self.shard_for(key);
        let guard = self.collector.enter();

        let mr = {
            let mut inner = self.engine.shards()[shard_id].lock();
            let mut index = sync::lock(&self.indexes[shard_id]);

            let entry = index.get(key).ok_or(DbError::KeyNotFound)?;
            if !std::ptr::eq(entry.ptr, expected as *mut TypedData<T>) {
                return Err(DbError::CasMismatch);
            }

            self.delete_locked_inner::<true>(shard_id, &mut inner, &mut index, &guard, key)?
        };

        self.fire_hook_and_retire(key, &mr);
        Ok(())
    }

    /// Atomically read-modify-write. Returns `Some(TypedRef)` to the **new** value
    /// if key existed, `None` otherwise.
    /// The closure must not be heavy (shard lock is held).
    pub fn update(&self, key: &K, f: impl FnOnce(&T) -> T) -> DbResult<Option<TypedRef<'_, T>>> {
        self.update_inner(key, f, false)
    }

    /// Like [`update()`](Self::update), but returns `Some(TypedRef)` to the **old** value.
    pub fn fetch_update(
        &self,
        key: &K,
        f: impl FnOnce(&T) -> T,
    ) -> DbResult<Option<TypedRef<'_, T>>> {
        self.update_inner(key, f, true)
    }

    fn update_inner(
        &self,
        key: &K,
        f: impl FnOnce(&T) -> T,
        return_old: bool,
    ) -> DbResult<Option<TypedRef<'_, T>>> {
        metrics::counter!("armdb.ops", "op" => "update", "tree" => "typed_map").increment(1);
        let shard_id = self.shard_for(key);
        let guard = self.collector.enter();

        let result = {
            let mut inner = self.engine.shards()[shard_id].lock();
            let mut index = sync::lock(&self.indexes[shard_id]);

            let entry = match index.get(key) {
                Some(e) => e,
                None => return Ok(None),
            };

            let old_data = unsafe { &*entry.ptr };
            let new_value = f(&old_data.value);

            let mut buf = sync::lock(&self.encode_bufs[shard_id]);
            self.codec.encode_to(&new_value, &mut buf)?;
            let (disk_loc, _gsn) =
                inner.append_entry(shard_id as u8, key.as_bytes(), &buf, false)?;

            let new_data = Box::new(TypedData {
                disk: disk_loc,
                value: new_value,
            });

            let new_data_ptr = Box::into_raw(new_data);
            let old_ptr = entry.ptr;
            index.get_mut(key).expect("key exists").ptr = new_data_ptr;

            let old_disk = unsafe { (*old_ptr).disk };
            inner.add_dead_bytes(
                old_disk.file_id,
                crate::entry::entry_size(size_of::<K>(), old_disk.len),
            );

            (old_ptr as *const TypedData<T>, new_data_ptr, old_ptr)
        };

        let (old_data, new_data_ptr, retire_ptr) = result;
        unsafe {
            self.hook
                .on_write(key, Some(&(*old_data).value), Some(&(*new_data_ptr).value));
            self.collector
                .retire(retire_ptr, seize::reclaim::boxed::<TypedData<T>>);
        }

        let ref_ptr = if return_old {
            old_data
        } else {
            new_data_ptr as *const TypedData<T>
        };
        Ok(Some(TypedRef::new(guard, ref_ptr)))
    }

    // -- Atomic ---------------------------------------------------------------

    /// Atomically execute multiple operations on a single shard.
    /// All keys must route to the same shard as `shard_key`.
    /// The closure must be short — shard lock is held for its duration.
    pub fn atomic<R>(
        &self,
        shard_key: &K,
        f: impl FnOnce(&mut TypedMapShard<'_, K, T, C, H>) -> DbResult<R>,
    ) -> DbResult<R> {
        let shard_id = self.shard_for(shard_key);
        let inner = self.engine.shards()[shard_id].lock();
        let index = sync::lock(&self.indexes[shard_id]);
        let guard = self.collector.enter();
        let mut shard = TypedMapShard {
            map: self,
            inner,
            index,
            shard_id,
            guard,
            events: Vec::new(),
        };
        let result = f(&mut shard);
        let TypedMapShard {
            inner,
            index,
            guard,
            events,
            ..
        } = shard;
        drop(index);
        drop(inner);
        if H::NEEDS_WRITE {
            for (key, mr) in &events {
                self.fire_hook_and_retire(key, mr);
            }
        }
        drop(guard);
        result
    }

    // -- Info -----------------------------------------------------------------

    pub fn len(&self) -> usize {
        self.len.load(std::sync::atomic::Ordering::Relaxed)
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Write hint files for all active shard files. Call during graceful shutdown.
    pub fn sync_hints(&self) -> DbResult<()> {
        for shard in self.engine.shards().iter() {
            shard.write_active_hint(size_of::<K>())?;
        }
        Ok(())
    }

    pub fn shard_for(&self, key: &K) -> usize {
        if self.shard_prefix_bits == 0 || self.shard_prefix_bits >= size_of::<K>() * 8 {
            let hash = xxhash_rust::xxh3::xxh3_64(key.as_bytes());
            return (hash as usize) % self.engine.shards().len();
        }

        let full_bytes = self.shard_prefix_bits / 8;
        let extra_bits = self.shard_prefix_bits % 8;

        let hash = if extra_bits == 0 {
            xxhash_rust::xxh3::xxh3_64(&key.as_bytes()[..full_bytes])
        } else {
            let mut buf = K::zeroed();
            buf.as_bytes_mut()[..full_bytes].copy_from_slice(&key.as_bytes()[..full_bytes]);
            let mask = !((1u8 << (8 - extra_bits)) - 1);
            buf.as_bytes_mut()[full_bytes] = key.as_bytes()[full_bytes] & mask;
            xxhash_rust::xxh3::xxh3_64(&buf.as_bytes()[..full_bytes + 1])
        };

        (hash as usize) % self.engine.shards().len()
    }

    // -- Ordered iteration ----------------------------------------------------

    /// Ordered scan view. `None` unless the collection was opened with
    /// `iterable(true)`. Direction follows `Config::reversed` (default DESC).
    ///
    /// **Occasional admin/maintenance API, not a hot path.** Use it to browse a
    /// collection or bulk-clean stale rows (`retain`); for scan-heavy access reach
    /// for the `*Tree` sibling — a Map's strength is O(1) point lookup, and the
    /// companion key index costs extra memory plus O(log N) per write. Unlike a
    /// `Tree` scan it is **not a point-in-time snapshot**: the k-way merge drops
    /// each shard lock between batches, so concurrent writes may be only partially
    /// reflected (run on a quiet DB if you need a coherent view).
    pub fn iter_view(&self) -> Option<TypedMapIterView<'_, K, T, C, H>> {
        self.iterable.then(|| TypedMapIterView {
            map: self,
            inner: MapIterView::new(self, dir_from(self.reversed)),
        })
    }

    // -- Read-only scan -------------------------------------------------------

    /// Read-only pass over all live entries. Order is unspecified; entries
    /// inserted/removed concurrently may be missed or visited once — intended
    /// for schema validation on a quiet database.
    pub fn for_each(&self, mut f: impl FnMut(&K, &T)) {
        let _guard = self.collector.enter();
        for i in 0..self.engine.shards().len() {
            let pairs: Vec<(K, *const TypedData<T>)> = {
                let index = sync::lock(&self.indexes[i]);
                index
                    .iter()
                    .map(|(k, entry)| (*k, entry.ptr as *const TypedData<T>))
                    .collect()
            };
            for (key, ptr) in pairs {
                let data = unsafe { &*ptr };
                f(&key, &data.value);
            }
        }
    }

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

    /// Iterate all entries and optionally mutate them. Call once at startup.
    ///
    /// **Must run before any other writer obtains a handle to this collection.**
    /// Concurrent writes during migration produce undefined results.
    ///
    /// The callback receives each (key, &T) and returns `MigrateAction`:
    /// - `Keep` — no change (fires `on_init` if `NEEDS_INIT`)
    /// - `Update(new_value)` — replace value (hook-free write, fires `on_init`)
    /// - `Delete` — remove entry (hook-free tombstone)
    ///
    /// Returns the number of mutated entries.
    pub fn migrate(&self, f: impl Fn(&K, &T) -> crate::MigrateAction<T>) -> DbResult<usize> {
        use crate::MigrateAction;

        let mut count = 0;
        for i in 0..self.engine.shards().len() {
            let keys: Vec<K> = {
                let index = sync::lock(&self.indexes[i]);
                index.keys().copied().collect()
            };
            for key in keys {
                let value_ref = match self.get(&key) {
                    Some(v) => v,
                    None => continue,
                };
                let action = f(&key, &*value_ref);
                drop(value_ref);

                match action {
                    MigrateAction::Keep => {
                        if H::NEEDS_INIT
                            && let Some(v) = self.get(&key)
                        {
                            self.hook.on_init(&key, &*v);
                        }
                    }
                    MigrateAction::Update(value) => {
                        let shard_id = self.shard_for(&key);
                        {
                            let mut inner = self.engine.shards()[shard_id].lock();
                            let mut index = sync::lock(&self.indexes[shard_id]);
                            let guard = self.collector.enter();
                            self.put_locked_inner::<false>(
                                shard_id, &mut inner, &mut index, &guard, &key, value,
                            )?;
                        }
                        if H::NEEDS_INIT
                            && let Some(v) = self.get(&key)
                        {
                            self.hook.on_init(&key, &*v);
                        }
                        count += 1;
                    }
                    MigrateAction::Delete => {
                        let shard_id = self.shard_for(&key);
                        let mut inner = self.engine.shards()[shard_id].lock();
                        let mut index = sync::lock(&self.indexes[shard_id]);
                        let guard = self.collector.enter();
                        self.delete_locked_inner::<false>(
                            shard_id, &mut inner, &mut index, &guard, &key,
                        )?;
                        count += 1;
                    }
                }
            }
        }

        tracing::info!(mutations = count, "typed_map migration complete");
        Ok(count)
    }

    pub(crate) fn replay_init(&self) {
        if !H::NEEDS_INIT {
            return;
        }
        let _guard = self.collector.enter();
        for shard in &self.indexes {
            let pairs: Vec<(K, *const TypedData<T>)> = {
                let index = sync::lock(shard);
                index
                    .iter()
                    .map(|(k, entry)| (*k, entry.ptr as *const TypedData<T>))
                    .collect()
            };
            for (key, ptr) in pairs {
                let data = unsafe { &*ptr };
                self.hook.on_init(&key, &data.value);
            }
        }
    }

    // -- Private --------------------------------------------------------------

    pub(crate) fn indexes(&self) -> &[Mutex<ShardIndex<K, TypedMapEntry<T>>>] {
        &self.indexes
    }

    /// Fire deferred on_write and retire old data. Called after lock release.
    ///
    /// SAFETY: caller must hold a live `seize::LocalGuard` that was entered
    /// before the mutation. This guarantees old/new value pointers are valid
    /// (not yet reclaimed) even if a reentrant hook retires them.
    fn fire_hook_and_retire(&self, key: &K, mr: &MutationResult<T>) {
        if let Some((old, new)) = mr.hook {
            unsafe {
                self.hook.on_write(key, old.map(|p| &*p), new.map(|p| &*p));
            }
        }
        if let Some(ptr) = mr.retire {
            unsafe {
                self.collector
                    .retire(ptr, seize::reclaim::boxed::<TypedData<T>>);
            }
        }
    }

    fn put_locked_inner<const HOOKS: bool>(
        &self,
        shard_id: usize,
        inner: &mut ShardInner,
        index: &mut ShardIndex<K, TypedMapEntry<T>>,
        _guard: &seize::LocalGuard<'_>,
        key: &K,
        value: T,
    ) -> DbResult<MutationResult<T>> {
        let mut buf = sync::lock(&self.encode_bufs[shard_id]);
        self.codec.encode_to(&value, &mut buf)?;
        let (disk_loc, _gsn) = inner.append_entry(shard_id as u8, key.as_bytes(), &buf, false)?;

        let new_data_ptr = Box::into_raw(Box::new(TypedData {
            disk: disk_loc,
            value,
        }));

        if let Some(old_entry) = index.upsert(*key, TypedMapEntry { ptr: new_data_ptr }) {
            let old_data = old_entry.ptr as *const TypedData<T>;

            let old_disk = unsafe { (*old_entry.ptr).disk };
            inner.add_dead_bytes(
                old_disk.file_id,
                crate::entry::entry_size(size_of::<K>(), old_disk.len),
            );

            if HOOKS {
                return Ok(MutationResult {
                    ref_ptr: Some(old_data),
                    hook: Some((
                        Some(unsafe { &(*old_data).value as *const T }),
                        Some(unsafe { &(*new_data_ptr).value as *const T }),
                    )),
                    retire: Some(old_entry.ptr),
                });
            }
            unsafe {
                self.collector
                    .retire(old_entry.ptr, seize::reclaim::boxed::<TypedData<T>>);
            }
            Ok(MutationResult {
                ref_ptr: Some(old_data),
                hook: None,
                retire: None,
            })
        } else {
            self.len.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
            let hook = if HOOKS {
                Some((None, Some(unsafe { &(*new_data_ptr).value as *const T })))
            } else {
                None
            };
            Ok(MutationResult {
                ref_ptr: None,
                hook,
                retire: None,
            })
        }
    }

    fn insert_locked<const HOOKS: bool>(
        &self,
        shard_id: usize,
        inner: &mut ShardInner,
        index: &mut ShardIndex<K, TypedMapEntry<T>>,
        _guard: &seize::LocalGuard<'_>,
        key: &K,
        value: T,
    ) -> DbResult<MutationResult<T>> {
        if index.contains_key(key) {
            return Err(DbError::KeyExists);
        }

        let mut buf = sync::lock(&self.encode_bufs[shard_id]);
        self.codec.encode_to(&value, &mut buf)?;
        let (disk_loc, _gsn) = inner.append_entry(shard_id as u8, key.as_bytes(), &buf, false)?;

        let new_data_ptr = Box::into_raw(Box::new(TypedData {
            disk: disk_loc,
            value,
        }));

        index.upsert(*key, TypedMapEntry { ptr: new_data_ptr });
        self.len.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        let hook = if HOOKS {
            Some((None, Some(unsafe { &(*new_data_ptr).value as *const T })))
        } else {
            None
        };
        Ok(MutationResult {
            ref_ptr: None,
            hook,
            retire: None,
        })
    }

    fn delete_locked_inner<const HOOKS: bool>(
        &self,
        shard_id: usize,
        inner: &mut ShardInner,
        index: &mut ShardIndex<K, TypedMapEntry<T>>,
        _guard: &seize::LocalGuard<'_>,
        key: &K,
    ) -> DbResult<MutationResult<T>> {
        let entry = match index.get(key) {
            Some(e) => e,
            None => {
                return Ok(MutationResult {
                    ref_ptr: None,
                    hook: None,
                    retire: None,
                });
            }
        };

        let old_data = entry.ptr as *const TypedData<T>;

        inner.append_entry(shard_id as u8, key.as_bytes(), &[], true)?;

        let removed = index.remove(key).expect("key exists after probe");
        self.len.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);

        let old_disk = unsafe { (*removed.ptr).disk };
        inner.add_dead_bytes(
            old_disk.file_id,
            crate::entry::entry_size(size_of::<K>(), old_disk.len),
        );

        if HOOKS {
            Ok(MutationResult {
                ref_ptr: Some(old_data),
                hook: Some((Some(unsafe { &(*old_data).value as *const T }), None)),
                retire: Some(removed.ptr),
            })
        } else {
            unsafe {
                self.collector
                    .retire(removed.ptr, seize::reclaim::boxed::<TypedData<T>>);
            }
            Ok(MutationResult {
                ref_ptr: Some(old_data),
                hook: None,
                retire: None,
            })
        }
    }
}

impl<K: Key + Send + Sync + Hash + Eq, T: Send + Sync, C: Codec<T>, H: TypedWriteHook<K, T>> Drop
    for TypedMap<K, T, C, H>
{
    fn drop(&mut self) {
        for index_mutex in &self.indexes {
            let map = sync::lock(index_mutex);
            for (_, entry) in map.iter() {
                unsafe {
                    drop(Box::from_raw(entry.ptr));
                }
            }
        }
    }
}

impl<K: Key + Send + Sync + Hash + Eq, T: Send + Sync, C: Codec<T> + Sync, H: TypedWriteHook<K, T>>
    ShardSource for TypedMap<K, T, C, H>
{
    type Key = K;
    type Val = *const TypedData<T>;

    fn shard_count(&self) -> usize {
        self.engine.shards().len()
    }

    fn next_batch(
        &self,
        shard: usize,
        after: Option<&K>,
        dir: Dir,
        n: usize,
        range_start: Bound<OrdKey<K>>,
        range_end: Bound<OrdKey<K>>,
    ) -> Vec<(K, *const TypedData<T>)> {
        let index = sync::lock(&self.indexes[shard]);
        let Some(ki) = index.key_index() else {
            return Vec::new();
        };
        let keys = key_index_batch(ki, after, dir, n, range_start, range_end);
        let mut out = Vec::with_capacity(keys.len());
        for k in keys {
            if let Some(e) = index.get(&k) {
                out.push((k, e.ptr as *const TypedData<T>));
            }
        }
        out
    }

    fn next_keys_batch(
        &self,
        shard: usize,
        after: Option<&K>,
        dir: Dir,
        n: usize,
        range_start: Bound<OrdKey<K>>,
        range_end: Bound<OrdKey<K>>,
    ) -> Vec<K> {
        let index = sync::lock(&self.indexes[shard]);
        let Some(ki) = index.key_index() else {
            return Vec::new();
        };
        key_index_batch(ki, after, dir, n, range_start, range_end)
    }
}

/// View over an iterable [`TypedMap`]; exposes ordered iteration yielding
/// guard-protected [`TypedRef`] values.
pub struct TypedMapIterView<
    'a,
    K: Key + Send + Sync + Hash + Eq,
    T: Send + Sync,
    C: Codec<T> + Sync,
    H: TypedWriteHook<K, T>,
> {
    map: &'a TypedMap<K, T, C, H>,
    inner: MapIterView<'a, TypedMap<K, T, C, H>>,
}

impl<'a, K, T, C, H> TypedMapIterView<'a, K, T, C, H>
where
    K: Key + Send + Sync + Hash + Eq,
    T: Send + Sync,
    C: Codec<T> + Sync,
    H: TypedWriteHook<K, T>,
{
    /// Full ordered scan in the collection's configured direction.
    pub fn iter(&self) -> TypedMapIter<'a, K, T, C, H> {
        let guard = self.map.collector.enter();
        TypedMapIter {
            map: self.map,
            inner: self.inner.iter(),
            _guard: guard,
        }
    }

    /// Iterate entries in `[start, end)` — start inclusive, end exclusive.
    pub fn range(&self, start: &K, end: &K) -> TypedMapIter<'a, K, T, C, H> {
        let guard = self.map.collector.enter();
        TypedMapIter {
            map: self.map,
            inner: self.inner.range(start, end),
            _guard: guard,
        }
    }

    /// Iterate entries in range defined by `start` and `end` bounds.
    pub fn range_bounds(&self, start: Bound<&K>, end: Bound<&K>) -> TypedMapIter<'a, K, T, C, H> {
        let guard = self.map.collector.enter();
        TypedMapIter {
            map: self.map,
            inner: self.inner.range_bounds(start, end),
            _guard: guard,
        }
    }

    /// Iterate entries whose keys start with `prefix`.
    pub fn prefix_iter(&self, prefix: &[u8]) -> TypedMapIter<'a, K, T, C, H> {
        let guard = self.map.collector.enter();
        TypedMapIter {
            map: self.map,
            inner: self.inner.prefix_iter(prefix),
            _guard: guard,
        }
    }

    /// Keyset pagination in the collection's configured direction.
    ///
    /// Requires `T: Clone` because each item is collected as an owned value.
    pub fn paginate(&self, after: Option<&K>, limit: usize) -> Page<K, T>
    where
        T: Clone,
    {
        let iter = match (after, dir_from(self.map.reversed)) {
            (None, _) => self.iter(),
            (Some(k), Dir::Asc) => self.range_bounds(Bound::Excluded(k), Bound::Unbounded),
            (Some(k), Dir::Desc) => self.range_bounds(Bound::Unbounded, Bound::Excluded(k)),
        };
        let items: Vec<_> = iter
            .take(limit)
            .map(|(k, r)| (k, r.clone_value()))
            .collect();
        let next = if items.len() == limit {
            items.last().map(|(k, _)| *k)
        } else {
            None
        };
        Page { items, next }
    }

    // -- Keys-only iteration -------------------------------------------------
    //
    // Keys are owned `K` values, so — unlike value iteration — these need no
    // `seize` guard (no pointer is dereferenced) and never touch the stored `T`.

    /// Full ordered **keys-only** scan — no value dereferenced, no guard needed.
    pub fn keys(&self) -> impl Iterator<Item = K> + '_ {
        self.inner.keys()
    }

    /// Keys-only iteration over `[start, end)`.
    pub fn keys_range(&self, start: &K, end: &K) -> impl Iterator<Item = K> + '_ {
        self.inner.keys_range(start, end)
    }

    /// Keys-only iteration over a bounded range.
    pub fn keys_range_bounds(
        &self,
        start: Bound<&K>,
        end: Bound<&K>,
    ) -> impl Iterator<Item = K> + '_ {
        self.inner.keys_range_bounds(start, end)
    }

    /// Keys-only iteration over keys starting with `prefix`.
    pub fn keys_prefix(&self, prefix: &[u8]) -> impl Iterator<Item = K> + '_ {
        self.inner.keys_prefix(prefix)
    }

    /// Keys-only pagination — cheap, materializes no values.
    pub fn keys_paginate(&self, after: Option<&K>, limit: usize) -> KeyPage<K> {
        self.inner.keys_paginate(after, limit)
    }

    /// Keep entries where `f` returns `true`; delete the rest via pointer-identity
    /// compare-delete (no `T: PartialEq` required).
    ///
    /// If the entry was replaced or removed concurrently (`CasMismatch` /
    /// `KeyNotFound`), the delete is skipped. Not atomic across the collection; a
    /// genuine storage error aborts the scan and is returned.
    pub fn retain(&self, mut f: impl FnMut(&K, &T) -> bool) -> DbResult<()> {
        let _guard = self.map.collector.enter();
        for (key, ptr) in self.inner.iter() {
            let val = unsafe { &(*ptr).value };
            if f(&key, val) {
                continue;
            }
            match self.map.compare_delete_if_ptr(&key, ptr) {
                Ok(()) => {}
                Err(DbError::CasMismatch) | Err(DbError::KeyNotFound) => {
                    tracing::debug!("retain: compare_delete skipped (concurrent change)");
                }
                Err(e) => return Err(e),
            }
        }
        Ok(())
    }
}

/// Iterator over entries in an iterable [`TypedMap`].
///
/// Weakly-consistent: concurrent inserts/updates may be visible during iteration.
/// The iterator holds a `seize` guard for its lifetime so batched raw pointers
/// remain valid; each returned [`TypedRef`] carries an additional guard so values
/// outlive the iterator.
pub struct TypedMapIter<
    'a,
    K: Key + Send + Sync + Hash + Eq,
    T: Send + Sync,
    C: Codec<T> + Sync,
    H: TypedWriteHook<K, T>,
> {
    map: &'a TypedMap<K, T, C, H>,
    inner: MapIter<'a, TypedMap<K, T, C, H>>,
    _guard: seize::LocalGuard<'a>,
}

impl<'a, K, T, C, H> Iterator for TypedMapIter<'a, K, T, C, H>
where
    K: Key + Send + Sync + Hash + Eq,
    T: Send + Sync,
    C: Codec<T> + Sync,
    H: TypedWriteHook<K, T>,
{
    type Item = (K, TypedRef<'a, T>);

    fn next(&mut self) -> Option<Self::Item> {
        let (k, ptr) = self.inner.next()?;
        let guard = self.map.collector.enter();
        Some((k, TypedRef::new(guard, ptr)))
    }
}

// ---------------------------------------------------------------------------
// TypedMapShard
// ---------------------------------------------------------------------------

/// Handle for atomic multi-key operations within a single shard.
/// Obtained via [`TypedMap::atomic`]. The shard + index locks are held for the
/// lifetime of this struct — keep the closure short.
pub struct TypedMapShard<
    'a,
    K: Key + Send + Sync + Hash + Eq,
    T: Send + Sync,
    C: Codec<T>,
    H: TypedWriteHook<K, T> = NoHook,
> {
    map: &'a TypedMap<K, T, C, H>,
    inner: MutexGuard<'a, ShardInner>,
    index: MutexGuard<'a, ShardIndex<K, TypedMapEntry<T>>>,
    shard_id: usize,
    guard: seize::LocalGuard<'a>,
    events: Vec<(K, MutationResult<T>)>,
}

impl<K: Key + Send + Sync + Hash + Eq, T: Send + Sync, C: Codec<T> + Sync, H: TypedWriteHook<K, T>>
    TypedMapShard<'_, K, T, C, H>
{
    pub fn put(&mut self, key: &K, value: T) -> DbResult<Option<TypedRef<'_, T>>> {
        self.check_shard(key)?;
        let mr = self.map.put_locked_inner::<true>(
            self.shard_id,
            &mut self.inner,
            &mut self.index,
            &self.guard,
            key,
            value,
        )?;
        let ref_ptr = mr.ref_ptr;
        if H::NEEDS_WRITE {
            self.events.push((*key, mr));
        } else if let Some(ptr) = mr.retire {
            unsafe {
                self.map
                    .collector
                    .retire(ptr, seize::reclaim::boxed::<TypedData<T>>);
            }
        }
        let guard = self.map.collector.enter();
        Ok(ref_ptr.map(|p| TypedRef::new(guard, p)))
    }

    pub fn insert(&mut self, key: &K, value: T) -> DbResult<()> {
        self.check_shard(key)?;
        let mr = self.map.insert_locked::<true>(
            self.shard_id,
            &mut self.inner,
            &mut self.index,
            &self.guard,
            key,
            value,
        )?;
        if H::NEEDS_WRITE {
            self.events.push((*key, mr));
        } else if let Some(ptr) = mr.retire {
            unsafe {
                self.map
                    .collector
                    .retire(ptr, seize::reclaim::boxed::<TypedData<T>>);
            }
        }
        Ok(())
    }

    pub fn delete(&mut self, key: &K) -> DbResult<Option<TypedRef<'_, T>>> {
        self.check_shard(key)?;
        let mr = self.map.delete_locked_inner::<true>(
            self.shard_id,
            &mut self.inner,
            &mut self.index,
            &self.guard,
            key,
        )?;
        let ref_ptr = mr.ref_ptr;
        if H::NEEDS_WRITE {
            self.events.push((*key, mr));
        } else if let Some(ptr) = mr.retire {
            unsafe {
                self.map
                    .collector
                    .retire(ptr, seize::reclaim::boxed::<TypedData<T>>);
            }
        }
        let guard = self.map.collector.enter();
        Ok(ref_ptr.map(|p| TypedRef::new(guard, p)))
    }

    pub fn get(&self, key: &K) -> Option<&T> {
        let entry = self.index.get(key)?;
        Some(unsafe { &(*entry.ptr).value })
    }

    pub fn get_or_err(&self, key: &K) -> DbResult<&T> {
        self.get(key).ok_or(DbError::KeyNotFound)
    }

    pub fn contains(&self, key: &K) -> bool {
        self.index.contains_key(key)
    }

    fn check_shard(&self, key: &K) -> DbResult<()> {
        if self.map.shard_for(key) != self.shard_id {
            return Err(DbError::ShardMismatch);
        }
        Ok(())
    }
}

#[cfg(feature = "armour")]
impl<T, C, H> crate::armour::collection::Collection for TypedMap<T::SelfId, T, C, H>
where
    T: crate::CollectionMeta + Clone + Send + Sync,
    C: crate::Codec<T> + Sync,
    H: crate::hook::TypedWriteHook<T::SelfId, T>,
    T::SelfId: crate::Key + Send + Sync + std::hash::Hash + Eq,
{
    fn name(&self) -> &str {
        T::NAME
    }
    fn len(&self) -> usize {
        self.len()
    }
    fn compact(&self) -> crate::DbResult<usize> {
        self.compact()
    }
    fn flush(&self) -> crate::DbResult<()> {
        self.flush_buffers()?;
        self.sync_hints()?;
        Ok(())
    }
    fn periodic_flush(&self) -> crate::DbResult<()> {
        self.flush_buffers()
    }
}

// ==========================================================================
// MultiTx — cross-collection transaction support (feature `armour`)
//
// Engine-direct map family: per shard holds (engine ShardInner, index HashMap);
// one seize guard; no durability sync seam; replays via fire_hook_and_retire.
// ==========================================================================

#[cfg(feature = "armour")]
type TypedMapGuard<'a, K, T> = (
    usize,
    MutexGuard<'a, ShardInner>,
    MutexGuard<'a, ShardIndex<K, TypedMapEntry<T>>>,
);

/// Multi-shard transaction handle for [`TypedMap`] inside `Db::atomicN`.
#[cfg(feature = "armour")]
pub struct TypedMapTx<'a, K, T, C, H = NoHook>
where
    K: Key + Send + Sync + Hash + Eq,
    T: Send + Sync,
    C: Codec<T>,
    H: TypedWriteHook<K, T>,
{
    map: &'a TypedMap<K, T, C, H>,
    shards: Vec<TypedMapGuard<'a, K, T>>,
    seize: seize::LocalGuard<'a>,
    log: Vec<(K, MutationResult<T>)>,
}

#[cfg(feature = "armour")]
impl<K, T, C, H> TypedMapTx<'_, K, T, C, H>
where
    K: Key + Send + Sync + Hash + Eq,
    T: Send + Sync,
    C: Codec<T> + Sync,
    H: TypedWriteHook<K, T>,
{
    fn position(&self, key: &K) -> DbResult<usize> {
        let sid = self.map.shard_for(key);
        self.shards
            .iter()
            .position(|(s, _, _)| *s == sid)
            .ok_or(DbError::ShardMismatch)
    }

    pub fn try_get(&self, key: &K) -> DbResult<Option<&T>> {
        let i = self.position(key)?;
        Ok(self.shards[i]
            .2
            .get(key)
            .map(|e| unsafe { &(*e.ptr).value }))
    }

    pub fn try_contains(&self, key: &K) -> DbResult<bool> {
        let i = self.position(key)?;
        Ok(self.shards[i].2.contains_key(key))
    }

    pub fn get_or_err(&self, key: &K) -> DbResult<&T> {
        self.try_get(key)?.ok_or(DbError::KeyNotFound)
    }

    fn retire_if_no_hook(&self, mr: &MutationResult<T>) {
        if let Some(ptr) = mr.retire {
            unsafe {
                self.map
                    .collector
                    .retire(ptr, seize::reclaim::boxed::<TypedData<T>>);
            }
        }
    }

    pub fn put(&mut self, key: &K, value: T) -> DbResult<()> {
        let i = self.position(key)?;
        let (sid, inner, index) = &mut self.shards[i];
        let mr = self
            .map
            .put_locked_inner::<true>(*sid, inner, index, &self.seize, key, value)?;
        if H::NEEDS_WRITE {
            self.log.push((*key, mr));
        } else {
            self.retire_if_no_hook(&mr);
        }
        Ok(())
    }

    pub fn insert(&mut self, key: &K, value: T) -> DbResult<()> {
        let i = self.position(key)?;
        let (sid, inner, index) = &mut self.shards[i];
        let mr = self
            .map
            .insert_locked::<true>(*sid, inner, index, &self.seize, key, value)?;
        if H::NEEDS_WRITE {
            self.log.push((*key, mr));
        } else {
            self.retire_if_no_hook(&mr);
        }
        Ok(())
    }

    pub fn delete(&mut self, key: &K) -> DbResult<()> {
        let i = self.position(key)?;
        let (sid, inner, index) = &mut self.shards[i];
        let mr = self
            .map
            .delete_locked_inner::<true>(*sid, inner, index, &self.seize, key)?;
        if H::NEEDS_WRITE {
            self.log.push((*key, mr));
        } else {
            self.retire_if_no_hook(&mr);
        }
        Ok(())
    }
}

#[cfg(feature = "armour")]
impl<K, T, C, H> crate::armour::multi_tx::MultiTx for TypedMap<K, T, C, H>
where
    K: Key + Send + Sync + Hash + Eq,
    T: Send + Sync,
    C: Codec<T> + Sync,
    H: TypedWriteHook<K, T>,
{
    type Key = K;
    type Tx<'a>
        = TypedMapTx<'a, K, T, C, H>
    where
        Self: 'a;

    fn shard_for_key(&self, key: &K) -> usize {
        self.shard_for(key)
    }

    fn begin_tx(&self) -> TypedMapTx<'_, K, T, C, H> {
        TypedMapTx {
            map: self,
            shards: Vec::new(),
            seize: self.collector.enter(),
            log: Vec::new(),
        }
    }

    fn lock_shard_into<'a>(&'a self, shard_id: usize, tx: &mut TypedMapTx<'a, K, T, C, H>) {
        // engine shard inner first, then the per-shard index — matching atomic().
        let inner = self.engine.shards()[shard_id].lock();
        let index = sync::lock(&self.indexes[shard_id]);
        tx.shards.push((shard_id, inner, index));
    }

    fn release_locks(
        &self,
        tx: &mut TypedMapTx<'_, K, T, C, H>,
    ) -> crate::armour::multi_tx::SyncNeeds {
        tx.shards.clear(); // drops index + engine guards; no per-write sync seam
        crate::armour::multi_tx::SyncNeeds::none()
    }

    fn run_sync(&self, _needs: crate::armour::multi_tx::SyncNeeds) -> DbResult<()> {
        Ok(())
    }

    fn replay_hooks(&self, tx: TypedMapTx<'_, K, T, C, H>) {
        if H::NEEDS_WRITE {
            for (k, mr) in &tx.log {
                self.fire_hook_and_retire(k, mr);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::codec::Codec;
    use crate::config::Config;
    use crate::hook::TypedWriteHook;
    use std::collections::HashMap as StdHashMap;
    use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
    use tempfile::tempdir;

    struct U64Codec;

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

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

    type Map = TypedMap<[u8; 8], u64, U64Codec>;

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

    fn key_on_other_shard(map: &Map, anchor: &[u8; 8]) -> [u8; 8] {
        let anchor_shard = map.shard_for(anchor);
        (2u64..1000)
            .map(|i| i.to_be_bytes())
            .find(|k| map.shard_for(k) != anchor_shard)
            .expect("key on different shard")
    }

    #[derive(Default)]
    struct TRecHook {
        writes: AtomicUsize,
        #[allow(clippy::type_complexity)]
        seq: crate::sync::Mutex<Vec<(u64, Option<u64>, Option<u64>)>>,
    }

    impl TypedWriteHook<[u8; 8], u64> for TRecHook {
        fn on_write(&self, key: &[u8; 8], old: Option<&u64>, new: Option<&u64>) {
            self.writes.fetch_add(1, AtomicOrdering::Relaxed);
            crate::sync::lock(&self.seq).push((
                u64::from_be_bytes(*key),
                old.copied(),
                new.copied(),
            ));
        }
    }

    fn open_typed_map_hooked<H: TypedWriteHook<[u8; 8], u64>>(
        dir: &std::path::Path,
        cfg: Config,
        hook: H,
    ) -> TypedMap<[u8; 8], u64, U64Codec, H> {
        TypedMap::open_hooked(dir, cfg, U64Codec, hook).expect("open hooked")
    }

    #[test]
    fn typed_map_atomic_fires_hooks_in_order() {
        let dir = tempdir().unwrap();
        let mut cfg = Config::test();
        cfg.shard_count = 1;
        let map = open_typed_map_hooked(dir.path(), cfg, TRecHook::default());
        let k = 5u64.to_be_bytes();
        map.atomic(&k, |s| {
            s.put(&k, 10)?; // new key
            s.put(&k, 20)?; // update, old=10
            s.delete(&k)?; // delete, old=20
            Ok(())
        })
        .expect("atomic");
        assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 3);
        let seq = crate::sync::lock(&map.hook.seq).clone();
        assert_eq!(seq[0], (5, None, Some(10)));
        assert_eq!(seq[1], (5, Some(10), Some(20)));
        assert_eq!(seq[2], (5, Some(20), None));
    }

    #[test]
    fn typed_map_atomic_fires_for_applied_on_err() {
        let dir = tempdir().unwrap();
        let mut cfg = Config::test();
        cfg.shard_count = 1;
        let map = open_typed_map_hooked(dir.path(), cfg, TRecHook::default());
        let k = 1u64.to_be_bytes();
        let r: DbResult<()> = map.atomic(&k, |s| {
            s.put(&k, 99)?;
            Err(DbError::KeyNotFound)
        });
        assert!(r.is_err());
        assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 1);
    }

    #[test]
    fn typed_map_iter_no_clone_bound() {
        #[derive(rapira::Rapira, PartialEq, Debug)]
        struct NoClone {
            n: u32,
            s: String,
        }

        let dir = tempdir().unwrap();
        let cfg = Config::balanced()
            .shard_count(2)
            .hints(true)
            .iterable(true)
            .build();
        let m = TypedMap::<[u8; 2], NoClone, crate::RapiraCodec>::open(
            dir.path(),
            cfg,
            crate::RapiraCodec,
        )
        .unwrap();
        m.put(
            &[1, 0],
            NoClone {
                n: 1,
                s: "a".into(),
            },
        )
        .unwrap();
        m.put(
            &[2, 0],
            NoClone {
                n: 2,
                s: "b".into(),
            },
        )
        .unwrap();
        let v = m.iter_view().unwrap();
        let ns: Vec<u32> = v.iter().map(|(_, r)| r.n).collect();
        assert_eq!(ns, [2, 1]);
        // keys-only iteration is clone-free and guard-free (never derefs `T`).
        let keys: Vec<[u8; 2]> = v.keys().collect();
        assert_eq!(keys, [[2, 0], [1, 0]]);
        v.retain(|_k, val| val.n % 2 == 0).unwrap();
        assert_eq!(m.iter_view().unwrap().iter().count(), 1);
    }

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

        let key_a = 1u64.to_be_bytes();
        let key_b = key_on_other_shard(&map, &key_a);
        let key_c = key_on_other_shard(&map, &key_b);

        let expected = StdHashMap::from([(key_a, 10u64), (key_b, 20u64), (key_c, 30u64)]);
        for (key, value) in &expected {
            map.put(key, *value).unwrap();
        }

        let mut seen = StdHashMap::new();
        map.for_each(|key, value| {
            seen.insert(*key, *value);
        });

        assert_eq!(seen, expected);
        assert_ne!(map.shard_for(&key_a), map.shard_for(&key_b));
    }
}