armdb 0.1.13

use std::collections::HashMap;
use std::hash::Hash;
use std::mem::size_of;
use std::sync::Arc;

use crate::Key;
use crate::byte_view::ByteView;
use crate::cache::{BlockCache, BlockKey};
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, WriteHook};
use crate::io::aligned_buf::AlignedBuf;
use crate::io::direct;
use crate::recovery::recover_var_map;
use crate::shard::ShardInner;
use crate::sync::{self, Mutex, MutexGuard};

/// A map with fixed-size keys and variable-length values.
/// Uses per-shard HashMap for O(1) lookup. Values are stored on disk,
/// cached at 4096-byte block granularity via `BlockCache`.
/// No ordered iteration — use `VarTree` if you need prefix/range scans.
///
/// Each `VarMap` owns its storage engine — one map = one database directory.
pub struct VarMap<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K> = NoHook> {
    indexes: Vec<Mutex<HashMap<K, DiskLoc>>>,
    engine: Engine,
    cache: BlockCache,
    compaction_threshold: f64,
    shard_prefix_bits: usize,
    hook: H,
}

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

impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> VarMap<K, H> {
    /// Open or create a `VarMap` with a write hook for secondary index maintenance.
    pub fn open_hooked(
        path: impl AsRef<std::path::Path>,
        config: Config,
        hook: H,
    ) -> DbResult<Self> {
        Self::open_inner(path, config, hook)
    }

    fn open_inner(path: impl AsRef<std::path::Path>, config: Config, hook: H) -> DbResult<Self> {
        let compaction_threshold = config.compaction_threshold;
        let shard_prefix_bits = config.shard_prefix_bits;
        let cache = BlockCache::new(&config.cache);
        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 {
            indexes.push(Mutex::new(HashMap::new()));
        }

        let map = Self {
            indexes,
            engine,
            cache,
            compaction_threshold,
            shard_prefix_bits,
            hook,
        };

        // 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 max_gsn = recover_var_map::<K>(
            &shard_dir_refs,
            &shard_ids,
            map.indexes(),
            hints,
            #[cfg(feature = "encryption")]
            map.engine.cipher(),
        )?;

        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>());
            }
        }
        tracing::info!(
            key_size = size_of::<K>(),
            entries = map.len(),
            "var_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, H: WriteHook<K>> CompactionIndex<K> for VarMap<K, 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(disk) = index.get_mut(key)
            && *disk == old_loc
        {
            *disk = new_loc;
            return true;
        }
        false
    }

    fn invalidate_blocks(&self, shard_id: u8, file_id: u32, total_bytes: u64) {
        self.cache.invalidate_file(shard_id, file_id, total_bytes);
    }

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

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

    /// Get a value by key. Checks block cache first, then reads from disk.
    pub fn get(&self, key: &K) -> Option<ByteView> {
        metrics::counter!("armdb.ops", "op" => "get", "tree" => "var_map").increment(1);
        #[cfg(feature = "hot-path-tracing")]
        tracing::trace!("var_map.get");
        let disk = {
            let index = sync::lock(&self.indexes[self.shard_for(key)]);
            *index.get(key)?
        };
        self.read_value_cached(&disk)
    }

    /// Get a value by key, returning `Err(KeyNotFound)` if absent.
    pub fn get_or_err(&self, key: &K) -> DbResult<ByteView> {
        self.get(key).ok_or(DbError::KeyNotFound)
    }

    /// Insert or update a key-value pair.
    pub fn put(&self, key: &K, value: &[u8]) -> DbResult<()> {
        metrics::counter!("armdb.ops", "op" => "put", "tree" => "var_map").increment(1);
        #[cfg(feature = "hot-path-tracing")]
        tracing::trace!("var_map.put");
        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 old_value = if H::NEEDS_OLD_VALUE {
            if let Some(disk) = index.get(key) {
                self.read_value_locked(disk, &inner)
            } else {
                None
            }
        } else {
            None
        };
        self.put_locked(shard_id, &mut inner, &mut index, key, value)?;
        self.hook.on_write(key, old_value.as_deref(), Some(value));
        Ok(())
    }

    /// 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: &[u8]) -> DbResult<()> {
        metrics::counter!("armdb.ops", "op" => "insert", "tree" => "var_map").increment(1);
        #[cfg(feature = "hot-path-tracing")]
        tracing::trace!("var_map.insert");
        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]);
        self.insert_locked(shard_id, &mut inner, &mut index, key, value)?;
        self.hook.on_write(key, None, Some(value));
        Ok(())
    }

    /// Delete a key. Returns `true` if the key existed.
    pub fn delete(&self, key: &K) -> DbResult<bool> {
        metrics::counter!("armdb.ops", "op" => "delete", "tree" => "var_map").increment(1);
        #[cfg(feature = "hot-path-tracing")]
        tracing::trace!("var_map.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 old_value = if H::NEEDS_OLD_VALUE {
            if let Some(disk) = index.get(key) {
                self.read_value_locked(disk, &inner)
            } else {
                None
            }
        } else {
            None
        };
        let existed = self.delete_locked(shard_id, &mut inner, &mut index, key)?;
        if existed {
            self.hook.on_write(key, old_value.as_deref(), None);
        }
        Ok(existed)
    }

    /// 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.
    ///
    /// **Caveat:** Holds the shard lock while reading the current value. On a
    /// block-cache miss this performs disk I/O under the lock, blocking all
    /// writes to the same shard. Pre-warm the cache to avoid latency spikes.
    pub fn cas(&self, key: &K, expected: &[u8], new_value: &[u8]) -> DbResult<()> {
        metrics::counter!("armdb.ops", "op" => "cas", "tree" => "var_map").increment(1);
        #[cfg(feature = "hot-path-tracing")]
        tracing::trace!("var_map.cas");
        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 disk = *index.get(key).ok_or(DbError::KeyNotFound)?;
        let current = self
            .read_value_locked(&disk, &inner)
            .ok_or(DbError::KeyNotFound)?;
        if current.as_ref() != expected {
            return Err(DbError::CasMismatch);
        }

        let (new_disk_loc, _gsn) =
            inner.append_entry(shard_id as u8, key.as_bytes(), new_value, false)?;
        inner.add_dead_bytes(
            disk.file_id as u32,
            crate::entry::entry_size(size_of::<K>(), disk.len),
        );
        index.insert(*key, new_disk_loc);

        self.hook.on_write(key, Some(&*current), Some(new_value));
        Ok(())
    }

    /// Atomically read-modify-write. Returns `Some(new_value)` if key existed, `None` otherwise.
    /// The closure receives the current value and returns a new ByteView.
    /// The closure must not be heavy (shard lock is held).
    ///
    /// **Caveat:** Holds the shard lock while reading the current value. On a
    /// block-cache miss this performs disk I/O under the lock, blocking all
    /// writes to the same shard. Pre-warm the cache to avoid latency spikes.
    pub fn update(&self, key: &K, f: impl FnOnce(&[u8]) -> ByteView) -> DbResult<Option<ByteView>> {
        self.update_inner(key, f, false)
    }

    /// Like [`update()`](Self::update), but returns `Some(old_value)` instead of the new one.
    pub fn fetch_update(
        &self,
        key: &K,
        f: impl FnOnce(&[u8]) -> ByteView,
    ) -> DbResult<Option<ByteView>> {
        self.update_inner(key, f, true)
    }

    fn update_inner(
        &self,
        key: &K,
        f: impl FnOnce(&[u8]) -> ByteView,
        return_old: bool,
    ) -> DbResult<Option<ByteView>> {
        metrics::counter!("armdb.ops", "op" => "update", "tree" => "var_map").increment(1);
        #[cfg(feature = "hot-path-tracing")]
        tracing::trace!("var_map.update");
        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 disk = match index.get(key) {
            Some(d) => *d,
            None => return Ok(None),
        };

        let current = match self.read_value_locked(&disk, &inner) {
            Some(v) => v,
            None => return Ok(None),
        };

        let new_value = f(&current);

        let (new_disk_loc, _gsn) =
            inner.append_entry(shard_id as u8, key.as_bytes(), &new_value, false)?;
        inner.add_dead_bytes(
            disk.file_id as u32,
            crate::entry::entry_size(size_of::<K>(), disk.len),
        );
        index.insert(*key, new_disk_loc);

        self.hook.on_write(key, Some(&*current), Some(&*new_value));
        Ok(Some(if return_old { current } else { new_value }))
    }

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

    /// Total number of entries across all shards.
    pub fn len(&self) -> usize {
        self.indexes.iter().map(|m| sync::lock(m).len()).sum()
    }

    pub fn is_empty(&self) -> bool {
        self.indexes.iter().all(|m| sync::lock(m).is_empty())
    }

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

    /// Pre-populate the block cache with blocks containing live values.
    ///
    /// Walks the index to collect unique block offsets, sorts them for sequential I/O,
    /// then reads each block into the cache. Only blocks with live data are loaded.
    pub fn warmup(&self) -> DbResult<()> {
        use std::collections::BTreeSet;

        let mut blocks: BTreeSet<(u8, u32, u64)> = BTreeSet::new();
        for index in self.indexes.iter() {
            let map = sync::lock(index);
            for disk in map.values() {
                let block_offset = disk.offset as u64 & !4095;
                blocks.insert((disk.shard_id, disk.file_id as u32, block_offset));
            }
        }

        for (shard_id, file_id, block_offset) in &blocks {
            let key = BlockKey {
                shard_id: *shard_id,
                file_id: *file_id,
                block_offset: *block_offset,
            };
            if self.cache.get(&key).is_some() {
                continue;
            }
            let shard = &self.engine.shards()[*shard_id as usize];
            let (buf, _) = shard.read_block(*file_id, *block_offset)?;
            self.cache.insert(key, Arc::new(buf));
        }

        Ok(())
    }

    /// Iterate all entries and optionally mutate them. Call once at startup.
    ///
    /// Block cache is populated naturally — no need to call `warmup()` separately.
    /// Hook is called for every entry, including `Keep`.
    ///
    /// Returns the number of mutated entries.
    pub fn migrate(
        &self,
        f: impl Fn(&K, &[u8]) -> crate::MigrateAction<ByteView>,
    ) -> 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 = match self.get(&key) {
                    Some(v) => v,
                    None => continue,
                };
                match f(&key, &value) {
                    MigrateAction::Keep => {
                        self.hook.on_write(&key, None, Some(&value));
                    }
                    MigrateAction::Update(new_value) => {
                        self.put(&key, &new_value)?;
                        count += 1;
                    }
                    MigrateAction::Delete => {
                        self.delete(&key)?;
                        count += 1;
                    }
                }
            }
        }

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

    /// 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 VarMapShard<'_, K, 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 mut shard = VarMapShard {
            tree: self,
            inner,
            index,
            shard_id,
        };
        f(&mut shard)
    }

    pub(crate) fn indexes(&self) -> &[Mutex<HashMap<K, DiskLoc>>] {
        &self.indexes
    }

    fn put_locked(
        &self,
        shard_id: usize,
        inner: &mut ShardInner,
        index: &mut HashMap<K, DiskLoc>,
        key: &K,
        value: &[u8],
    ) -> DbResult<()> {
        let (disk_loc, _gsn) = inner.append_entry(shard_id as u8, key.as_bytes(), value, false)?;

        if let Some(old_disk) = index.insert(*key, disk_loc) {
            inner.add_dead_bytes(
                old_disk.file_id as u32,
                crate::entry::entry_size(size_of::<K>(), old_disk.len),
            );
        }

        Ok(())
    }

    fn insert_locked(
        &self,
        shard_id: usize,
        inner: &mut ShardInner,
        index: &mut HashMap<K, DiskLoc>,
        key: &K,
        value: &[u8],
    ) -> DbResult<()> {
        if index.contains_key(key) {
            return Err(DbError::KeyExists);
        }

        let (disk_loc, _gsn) = inner.append_entry(shard_id as u8, key.as_bytes(), value, false)?;
        index.insert(*key, disk_loc);
        Ok(())
    }

    fn delete_locked(
        &self,
        shard_id: usize,
        inner: &mut ShardInner,
        index: &mut HashMap<K, DiskLoc>,
        key: &K,
    ) -> DbResult<bool> {
        if !index.contains_key(key) {
            return Ok(false);
        }

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

        if let Some(old_disk) = index.remove(key) {
            inner.add_dead_bytes(
                old_disk.file_id as u32,
                crate::entry::entry_size(size_of::<K>(), old_disk.len),
            );
            Ok(true)
        } else {
            Ok(false)
        }
    }

    /// Read a value via block cache. Checks write buffer for unflushed data.
    fn read_value_cached(&self, disk: &DiskLoc) -> Option<ByteView> {
        let block_offset = disk.offset as u64 & !4095;
        let start = (disk.offset & 4095) as usize;
        let len = disk.len as usize;
        let cache_key = BlockKey {
            shard_id: disk.shard_id,
            file_id: disk.file_id as u32,
            block_offset,
        };

        // 1. Block Cache (lock-free)
        if let Some(block) = self.cache.get(&cache_key) {
            metrics::counter!("armdb.cache.hit").increment(1);
            return Self::extract_from_block(&block, start, len, || {
                self.get_or_read_block(disk.shard_id, disk.file_id as u32, block_offset + 4096)
            });
        }

        // 2. Write buffer — for unflushed data in active file (brief shard lock)
        {
            let shard = &self.engine.shards()[disk.shard_id as usize];
            let inner = shard.lock();
            if inner.active.file_id == disk.file_id as u32
                && let Some(bytes) = inner.write_buf.read(disk.offset as u64, len)
            {
                return Some(ByteView::new(bytes));
            }
        }

        // 3. Disk read + cache
        metrics::counter!("armdb.cache.miss").increment(1);
        let block = self.get_or_read_block(disk.shard_id, disk.file_id as u32, block_offset)?;
        Self::extract_from_block(&block, start, len, || {
            self.get_or_read_block(disk.shard_id, disk.file_id as u32, block_offset + 4096)
        })
    }

    fn extract_from_block(
        block: &AlignedBuf,
        start: usize,
        len: usize,
        next_block: impl FnOnce() -> Option<Arc<AlignedBuf>>,
    ) -> Option<ByteView> {
        if start + len <= 4096 {
            Some(ByteView::new(&block[start..start + len]))
        } else {
            let next = next_block()?;
            let first_part = &block[start..];
            let second_len = len - first_part.len();
            let mut combined = Vec::with_capacity(len);
            combined.extend_from_slice(first_part);
            combined.extend_from_slice(&next[..second_len]);
            Some(ByteView::from_vec(combined))
        }
    }

    fn get_or_read_block(
        &self,
        shard_id: u8,
        file_id: u32,
        block_offset: u64,
    ) -> Option<Arc<AlignedBuf>> {
        let key = BlockKey {
            shard_id,
            file_id,
            block_offset,
        };
        if let Some(cached) = self.cache.get(&key) {
            return Some(cached);
        }
        let shard = &self.engine.shards()[shard_id as usize];
        let (buf, is_full_block) = match shard.read_block(file_id, block_offset) {
            Ok(b) => b,
            Err(e) => {
                #[cfg(feature = "hot-path-tracing")]
                tracing::error!("VarMap block read error: {:?}", e);
                let _ = e;
                return None;
            }
        };
        let arc = Arc::new(buf);
        if is_full_block {
            self.cache.insert(key, arc.clone());
        }
        Some(arc)
    }

    /// Read value when shard lock is already held. Used by CAS/update.
    fn read_value_locked(&self, disk: &DiskLoc, inner: &ShardInner) -> Option<ByteView> {
        let len = disk.len as usize;

        // 1. Write buffer (for unflushed data)
        if inner.active.file_id == disk.file_id as u32
            && let Some(bytes) = inner.write_buf.read(disk.offset as u64, len)
        {
            return Some(ByteView::new(bytes));
        }

        // 2. Block cache (lock-free)
        let block_offset = disk.offset as u64 & !4095;
        let start = (disk.offset & 4095) as usize;
        if start + len <= 4096 {
            let cache_key = BlockKey {
                shard_id: disk.shard_id,
                file_id: disk.file_id as u32,
                block_offset,
            };
            if let Some(block) = self.cache.get(&cache_key) {
                return Some(ByteView::new(&block[start..start + len]));
            }
        }

        // 3. Disk read (file handle from inner, no re-lock)
        let file: &std::fs::File = if inner.active.file_id == disk.file_id as u32 {
            &inner.active.read_file
        } else {
            let imm = inner
                .immutable
                .iter()
                .find(|f| f.file_id == disk.file_id as u32)?;
            &imm.file
        };
        let bytes = direct::pread_value(file, disk.offset as u64, len).ok()?;
        Some(ByteView::new(&bytes))
    }

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

#[cfg(feature = "replication")]
impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> crate::replication::ReplicationTarget
    for VarMap<K, H>
{
    fn apply_entry(
        &self,
        shard_id: u8,
        file_id: u32,
        entry_offset: u64,
        header: &crate::entry::EntryHeader,
        key: &[u8],
        value: &[u8],
    ) -> DbResult<()> {
        let key = K::from_bytes(key);

        let disk = DiskLoc::new(
            shard_id,
            file_id as u16,
            (entry_offset + 16 + size_of::<K>() as u64) as u32,
            header.value_len,
        );

        if header.is_tombstone() {
            sync::lock(&self.indexes[self.shard_for(&key)]).remove(&key);
        } else {
            let _ = value;
            sync::lock(&self.indexes[self.shard_for(&key)]).insert(key, disk);
        }

        Ok(())
    }

    fn try_apply_entry(
        &self,
        shard_id: u8,
        file_id: u32,
        entry_offset: u64,
        header: &crate::entry::EntryHeader,
        raw_after_header: &[u8],
    ) -> DbResult<bool> {
        if raw_after_header.len() < size_of::<K>() + header.value_len as usize {
            return Ok(false);
        }
        let key = &raw_after_header[..size_of::<K>()];
        let value = &raw_after_header[size_of::<K>()..size_of::<K>() + header.value_len as usize];
        let crc = crate::entry::compute_crc32(header.gsn, header.value_len, key, value);
        if crc != header.crc32 {
            return Ok(false);
        }
        self.apply_entry(shard_id, file_id, entry_offset, header, key, value)?;
        Ok(true)
    }

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

/// Handle for atomic multi-key operations within a single shard.
/// Obtained via [`VarMap::atomic`]. The shard + index locks are held for the
/// lifetime of this struct — keep the closure short.
pub struct VarMapShard<'a, K: Key + Send + Sync + Hash + Eq, H: WriteHook<K> = NoHook> {
    tree: &'a VarMap<K, H>,
    inner: MutexGuard<'a, ShardInner>,
    index: MutexGuard<'a, HashMap<K, DiskLoc>>,
    shard_id: usize,
}

impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> VarMapShard<'_, K, H> {
    pub fn put(&mut self, key: &K, value: &[u8]) -> DbResult<()> {
        self.check_shard(key)?;
        self.tree
            .put_locked(self.shard_id, &mut self.inner, &mut self.index, key, value)
    }

    pub fn insert(&mut self, key: &K, value: &[u8]) -> DbResult<()> {
        self.check_shard(key)?;
        self.tree
            .insert_locked(self.shard_id, &mut self.inner, &mut self.index, key, value)
    }

    pub fn delete(&mut self, key: &K) -> DbResult<bool> {
        self.check_shard(key)?;
        self.tree
            .delete_locked(self.shard_id, &mut self.inner, &mut self.index, key)
    }

    pub fn get(&self, key: &K) -> Option<ByteView> {
        let disk = *self.index.get(key)?;
        self.tree.read_value_locked(&disk, &self.inner)
    }

    pub fn get_or_err(&self, key: &K) -> DbResult<ByteView> {
        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.tree.shard_for(key) != self.shard_id {
            return Err(DbError::ShardMismatch);
        }
        Ok(())
    }
}