kglite 0.16.8

//! Mmap-resident `id_indices.bin` store with overlay for mutations.
//!
//! Replaces the eager `zstd::decode_all` + 124M-entry `HashMap::insert`
//! load path. Reads come from a memory-mapped flat binary on the disk;
//! mutations land in an in-memory overlay that takes precedence over
//! the base. On save, overlay + base are merged into a fresh `.bin`.
//!
//! ## File format `id_indices.bin`
//!
//! ```text
//! Header (32 bytes):
//!   [ 0.. 8]  magic           = b"KGLIIDXR"  (R = raw, mmap-friendly)
//!   [ 8..12]  version         = u32 LE (= 2)
//!   [12..16]  num_types       = u32 LE
//!   [16..24]  dir_offset      = u64 LE   (always 32)
//!   [24..32]  data_offset     = u64 LE   (32 + 48 * num_types)
//!
//! Directory at [dir_offset]: 48 bytes per entry, sorted by type_key:
//!   [ 0.. 8]  type_key:    u64 LE  (InternedKey)
//!   [ 8.. 9]  variant:     u8      (0 = Integer, 1 = General)
//!   [ 9..16]  padding:     [u8; 7]
//!   [16..24]  num_entries: u64 LE
//!   [24..32]  payload_off: u64 LE   (file-relative)
//!   [32..40]  payload_len: u64 LE
//!   [40..48]  padding:     u64
//!
//! Data section at [data_offset]:
//!   Integer (variant=0):
//!     [payload_off..payload_off + 4*num_entries]               keys: [u32 sorted asc]
//!     [payload_off + 4*num_entries..payload_off + payload_len] idxs: [u32]
//!   General (variant=1):
//!     Postcard of HashMap<Value, NodeIndex>, length = payload_len
//! ```
//!
//! Lookup is `O(log n)` binary search on `keys` for the Integer variant
//! (cache-friendly, ~24 comparisons even at 13M entries) and a single
//! `HashMap` probe for the General variant (decoded at load, where each
//! General payload is decoded to validate it).

use crate::datatypes::Value;
use crate::graph::schema::{InternedKey, StringInterner, TypeIdIndex};
use crate::graph::storage::disk::id_index_layer::TypeEntry;
use crate::serde_codec;
use memmap2::Mmap;
use petgraph::graph::NodeIndex;
use rustc_hash::FxHashMap;
use std::collections::HashMap;
use std::path::Path;
use std::sync::{Arc, RwLock};

const MAGIC: &[u8; 8] = b"KGLIIDXR";
const VERSION: u32 = 2;
const HEADER_BYTES: usize = 32;
const DIR_ENTRY_BYTES: usize = 48;
const MAX_GENERAL_INDEX_DECODE_BYTES: u64 = 2 * 1024 * 1024 * 1024;

fn invalid_index(message: &str) -> std::io::Error {
    std::io::Error::new(
        std::io::ErrorKind::InvalidData,
        format!("invalid id_indices.bin: {message}"),
    )
}

fn read_le_u32(bytes: &[u8], index: usize) -> Option<u32> {
    let start = index.checked_mul(4)?;
    Some(u32::from_le_bytes(
        bytes.get(start..start.checked_add(4)?)?.try_into().ok()?,
    ))
}

fn le_u32_binary_search(bytes: &[u8], wanted: u32) -> Option<usize> {
    let mut low = 0usize;
    let mut high = bytes.len() / 4;
    while low < high {
        let mid = low + (high - low) / 2;
        match read_le_u32(bytes, mid)?.cmp(&wanted) {
            std::cmp::Ordering::Less => low = mid + 1,
            std::cmp::Ordering::Greater => high = mid,
            std::cmp::Ordering::Equal => return Some(mid),
        }
    }
    None
}

/// Mmap-backed read-only view of `id_indices.bin`.
pub struct IdIndexBase {
    mmap: Arc<Mmap>,
    /// type_name -> directory entry. Built once at load (88k entries × ~50 bytes ≈ 4 MB).
    /// Strings owned to keep the API HashMap-compatible without lifetime gymnastics.
    dir: HashMap<String, BaseEntry>,
    /// Decoded General payloads. Filled at load, which decodes every General
    /// entry to validate it; `general_map` decodes only on a miss.
    /// Integer variant never enters here — it's read directly from mmap.
    general_cache: RwLock<HashMap<String, Arc<FxHashMap<Value, NodeIndex>>>>,
}

#[derive(Clone, Copy)]
struct BaseEntry {
    variant: u8,
    num_entries: u32,
    payload_off: u64,
    payload_len: u64,
}

impl IdIndexBase {
    /// Load `id_indices.bin` from `dir`. Returns `Ok(None)` if absent, shorter
    /// than the header, or magic mismatch.
    pub fn load_from(dir: &Path, interner: &StringInterner) -> std::io::Result<Option<Self>> {
        let path = dir.join("id_indices.bin");
        if !path.exists() {
            return Ok(None);
        }
        let file = std::fs::File::open(&path)?;
        let len = file.metadata()?.len() as usize;
        if len < HEADER_BYTES {
            return Ok(None);
        }
        // SAFETY: GraphDirectoryLock serializes disk-graph writers, which
        // publish a new immutable generation instead of truncating the
        // generation selected by this reader. This inode therefore remains
        // stable for the mapping's lifetime.
        let mmap = unsafe { Mmap::map(&file)? };
        if &mmap[..8] != MAGIC {
            return Ok(None);
        }
        let version = u32::from_le_bytes(mmap[8..12].try_into().unwrap());
        match version {
            1 => {
                return Err(crate::graph::io::file::pre_014_bincode_error(
                    "id_indices.bin v1",
                ));
            }
            VERSION => {}
            _ => return Err(invalid_index("unsupported raw index version")),
        }
        let num_types = u32::from_le_bytes(mmap[12..16].try_into().unwrap()) as usize;
        let dir_offset = usize::try_from(u64::from_le_bytes(mmap[16..24].try_into().unwrap()))
            .map_err(|_| invalid_index("directory offset exceeds usize"))?;
        let data_offset = usize::try_from(u64::from_le_bytes(mmap[24..32].try_into().unwrap()))
            .map_err(|_| invalid_index("data offset exceeds usize"))?;
        let dir_bytes = DIR_ENTRY_BYTES
            .checked_mul(num_types)
            .ok_or_else(|| invalid_index("directory size overflow"))?;
        let need = dir_offset
            .checked_add(dir_bytes)
            .ok_or_else(|| invalid_index("directory range overflow"))?;
        if dir_offset != HEADER_BYTES || data_offset != need || need > len {
            return Err(invalid_index("invalid directory/data boundary"));
        }

        let mut dir_map: HashMap<String, BaseEntry> = HashMap::with_capacity(num_types);
        let mut general_cache_map = HashMap::new();
        let mut previous_key = None;
        let mut expected_payload = data_offset;
        for i in 0..num_types {
            let off = dir_offset + i * DIR_ENTRY_BYTES;
            let type_key = u64::from_le_bytes(mmap[off..off + 8].try_into().unwrap());
            let variant = mmap[off + 8];
            let num_entries_u64 = u64::from_le_bytes(mmap[off + 16..off + 24].try_into().unwrap());
            let payload_off = u64::from_le_bytes(mmap[off + 24..off + 32].try_into().unwrap());
            let payload_len = u64::from_le_bytes(mmap[off + 32..off + 40].try_into().unwrap());
            if previous_key.is_some_and(|previous| type_key <= previous) {
                return Err(invalid_index("directory keys are not strictly increasing"));
            }
            previous_key = Some(type_key);
            if !matches!(variant, 0 | 1) {
                return Err(invalid_index("directory contains an unknown variant"));
            }
            let num_entries = u32::try_from(num_entries_u64)
                .map_err(|_| invalid_index("entry count exceeds u32"))?;
            let payload_off_usize = usize::try_from(payload_off)
                .map_err(|_| invalid_index("payload offset exceeds usize"))?;
            let payload_len_usize = usize::try_from(payload_len)
                .map_err(|_| invalid_index("payload length exceeds usize"))?;
            let payload_end = payload_off_usize
                .checked_add(payload_len_usize)
                .ok_or_else(|| invalid_index("payload range overflow"))?;
            if payload_off_usize != expected_payload || payload_end > len {
                return Err(invalid_index(
                    "payloads overlap, contain gaps, or exceed the file",
                ));
            }
            if variant == 0 {
                let expected_len = num_entries_u64
                    .checked_mul(8)
                    .ok_or_else(|| invalid_index("integer payload size overflow"))?;
                if payload_len != expected_len {
                    return Err(invalid_index("integer payload has invalid size"));
                }
                let keys_end = payload_off_usize + num_entries as usize * 4;
                let mut previous = None;
                for index in 0..num_entries as usize {
                    let key = read_le_u32(&mmap[payload_off_usize..keys_end], index).unwrap();
                    if previous.is_some_and(|prior| key <= prior) {
                        return Err(invalid_index("integer keys are not strictly increasing"));
                    }
                    previous = Some(key);
                }
            } else if payload_len > MAX_GENERAL_INDEX_DECODE_BYTES {
                return Err(invalid_index("general payload exceeds decode limit"));
            }
            expected_payload = payload_end;
            let Some(name) = interner.try_resolve(InternedKey::from_u64(type_key)) else {
                // Directories written before the writer resolved its keys
                // (through 0.15.0) carry entries for type names the interner
                // sidecar never received — a type declared with no rows, or a
                // label a query merely mentioned. Those entries are always
                // empty, and an id index is a cache the read path rebuilds on
                // demand, so dropping one recovers the graph at no cost rather
                // than making the whole directory unreadable. A *populated*
                // entry under an unresolvable key is not that: it is a
                // mismatched or damaged sidecar, and still fails the load.
                if num_entries == 0 {
                    continue;
                }
                return Err(invalid_index("directory contains an unresolved type key"));
            };
            if variant == 1 {
                let blob = &mmap[payload_off_usize..payload_end];
                let map: FxHashMap<Value, NodeIndex> = serde_codec::decode_exact_with(
                    serde_codec::CURRENT_CODEC,
                    blob,
                    blob.len() as u64,
                    serde_codec::DecodeLimits::new(
                        MAX_GENERAL_INDEX_DECODE_BYTES,
                        MAX_GENERAL_INDEX_DECODE_BYTES,
                    ),
                )
                .map_err(|_| invalid_index("general payload Postcard is malformed"))?;
                if map.len() != num_entries as usize {
                    return Err(invalid_index(
                        "general payload has duplicate or missing keys",
                    ));
                }
                general_cache_map.insert(name.to_string(), Arc::new(map));
            }
            if dir_map
                .insert(
                    name.to_string(),
                    BaseEntry {
                        variant,
                        num_entries,
                        payload_off,
                        payload_len,
                    },
                )
                .is_some()
            {
                return Err(invalid_index("duplicate resolved type name"));
            }
        }
        if expected_payload != len {
            return Err(invalid_index(
                "payload directory does not cover the file exactly",
            ));
        }

        Ok(Some(Self {
            mmap: Arc::new(mmap),
            dir: dir_map,
            general_cache: RwLock::new(general_cache_map),
        }))
    }

    pub fn contains(&self, name: &str) -> bool {
        self.dir.contains_key(name)
    }

    pub fn lookup(&self, name: &str, id: &Value) -> Option<NodeIndex> {
        let entry = self.dir.get(name)?;
        match entry.variant {
            0 => self.lookup_integer(entry, id),
            1 => self.lookup_general(name, entry, id),
            _ => None,
        }
    }

    /// Materialize a base entry into an owned `TypeIdIndex` (used on save and
    /// on first mutation when the entry must be promoted into the overlay).
    pub fn materialize(&self, name: &str) -> Option<TypeIdIndex> {
        let entry = self.dir.get(name)?;
        match entry.variant {
            0 => {
                let (keys, idxs) = self.integer_bytes(entry)?;
                let mut map: FxHashMap<u32, NodeIndex> = FxHashMap::with_capacity_and_hasher(
                    entry.num_entries as usize,
                    Default::default(),
                );
                for index in 0..entry.num_entries as usize {
                    map.insert(
                        read_le_u32(keys, index)?,
                        NodeIndex::new(read_le_u32(idxs, index)? as usize),
                    );
                }
                Some(TypeIdIndex::Integer(map))
            }
            1 => {
                let map = self.general_map(name, entry)?;
                Some(TypeIdIndex::General((*map).clone()))
            }
            _ => None,
        }
    }

    fn integer_bytes(&self, entry: &BaseEntry) -> Option<(&[u8], &[u8])> {
        let n = entry.num_entries as usize;
        let off = entry.payload_off as usize;
        let half = n * 4;
        if entry.payload_len != (half * 2) as u64 {
            return None;
        }
        let bytes = self.mmap.get(off..off + half * 2)?;
        Some(bytes.split_at(half))
    }

    fn lookup_integer(&self, entry: &BaseEntry, id: &Value) -> Option<NodeIndex> {
        let key_u32 = coerce_to_u32(id)?;
        let (keys, idxs) = self.integer_bytes(entry)?;
        let index = le_u32_binary_search(keys, key_u32)?;
        Some(NodeIndex::new(read_le_u32(idxs, index)? as usize))
    }

    fn lookup_general(&self, name: &str, entry: &BaseEntry, id: &Value) -> Option<NodeIndex> {
        let map = self.general_map(name, entry)?;
        if let Some(&idx) = map.get(id) {
            return Some(idx);
        }
        // Mirror TypeIdIndex::General numeric coercion fallbacks (Int64 ↔
        // UniqueId, Float64 → Int/UniqueId). No string→u32 coercion — a
        // String id matches only by exact value (handled above).
        match id {
            Value::Int64(i) => {
                if *i >= 0 && *i <= u32::MAX as i64 {
                    return map.get(&Value::UniqueId(*i as u32)).copied();
                }
                None
            }
            Value::UniqueId(u) => map.get(&Value::Int64(*u as i64)).copied(),
            Value::Float64(f) => {
                if f.fract() == 0.0 {
                    let i = *f as i64;
                    if let Some(&idx) = map.get(&Value::Int64(i)) {
                        return Some(idx);
                    }
                    if i >= 0 && i <= u32::MAX as i64 {
                        return map.get(&Value::UniqueId(i as u32)).copied();
                    }
                }
                None
            }
            _ => None,
        }
    }

    fn general_map(
        &self,
        name: &str,
        entry: &BaseEntry,
    ) -> Option<Arc<FxHashMap<Value, NodeIndex>>> {
        if let Some(arc) = self.general_cache.read().unwrap().get(name).cloned() {
            return Some(arc);
        }
        let off = entry.payload_off as usize;
        let len = entry.payload_len as usize;
        let blob = self.mmap.get(off..off + len)?;
        let map: FxHashMap<Value, NodeIndex> = serde_codec::decode_exact_with(
            serde_codec::CURRENT_CODEC,
            blob,
            blob.len() as u64,
            serde_codec::DecodeLimits::new(
                MAX_GENERAL_INDEX_DECODE_BYTES,
                MAX_GENERAL_INDEX_DECODE_BYTES,
            ),
        )
        .ok()?;
        if map.len() != entry.num_entries as usize {
            return None;
        }
        let arc = Arc::new(map);
        self.general_cache
            .write()
            .unwrap()
            .insert(name.to_string(), Arc::clone(&arc));
        Some(arc)
    }
}

/// HashMap-shaped wrapper around an optional mmap base + in-memory overlay.
///
/// Reads consult overlay first (covers post-load mutations), then base.
/// Mutations only ever land in overlay; `removed` tracks types that the
/// caller explicitly cleared so that base entries are masked.
#[derive(Default)]
pub struct IdIndexStore {
    /// In-memory layer: indices built/mutated post-load, plus lazily-cached
    /// indices the read path builds on a miss. Behind a `RwLock` so the
    /// read path can build + cache through `&self` — `DirGraph` is shared
    /// as `Arc<DirGraph>` and reads run on multiple threads (GIL-release),
    /// so this must be thread-safe.
    overlay: RwLock<HashMap<String, TypeEntry>>,
    /// Types that exist in `base` but were removed/invalidated post-load.
    removed: std::collections::HashSet<String>,
    base: Option<Arc<IdIndexBase>>,
}

impl Clone for IdIndexStore {
    /// **The fork seam for `id_indices`.**
    ///
    /// Instead of deep-copying a map with one entry per node of every
    /// materialised type — 3.7 ms at 1M, and 90% of what a plain graph's fork
    /// still cost at that point — this converts each of *our own* entries into
    /// a shared base in place and hands the child an empty delta over the same
    /// allocation. Both graphs then read identical content; only the
    /// representation changed.
    ///
    /// It has to happen here, taking the write lock through `&self`, because
    /// every fork reaches this field as a `&self` clone: by the time write
    /// entry holds a `&mut DirGraph` the copy has already been made. That is
    /// why `overlay`'s `RwLock` is load-bearing beyond thread safety.
    fn clone(&self) -> Self {
        let mut overlay = self.overlay.write().unwrap();
        let shared: HashMap<String, TypeEntry> = overlay
            .iter_mut()
            .map(|(name, entry)| (name.clone(), TypeEntry::layered_over(entry.share())))
            .collect();
        Self {
            overlay: RwLock::new(shared),
            removed: self.removed.clone(),
            base: self.base.clone(),
        }
    }
}

impl IdIndexStore {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn from_base(base: IdIndexBase) -> Self {
        Self {
            overlay: RwLock::new(HashMap::new()),
            removed: std::collections::HashSet::new(),
            base: Some(Arc::new(base)),
        }
    }

    pub fn contains_key(&self, name: &str) -> bool {
        if self.overlay.read().unwrap().contains_key(name) {
            return true;
        }
        if self.removed.contains(name) {
            return false;
        }
        self.base.as_ref().is_some_and(|b| b.contains(name))
    }

    /// Look up `id` for `name`. If the type isn't indexed anywhere (neither
    /// overlay nor base, or it was invalidated), build the index via `build`
    /// — which scans the graph — and cache it in the overlay, so the read
    /// path is O(1) on every subsequent lookup. The build runs at most once
    /// per type until the next invalidation. Returns None when the id simply
    /// isn't present (no scan).
    ///
    /// Without this, the read path (`MATCH (n {id:X})`, `MERGE` match) falls
    /// back to a full scan whenever the index is absent — after `add_nodes` /
    /// `CREATE` / `DELETE`. (issue #20)
    pub fn lookup_or_build(
        &self,
        name: &str,
        id: &Value,
        build: impl FnOnce() -> TypeIdIndex,
    ) -> Option<NodeIndex> {
        {
            let ov = self.overlay.read().unwrap();
            if let Some(idx) = ov.get(name) {
                return idx.get(id);
            }
        }
        if !self.removed.contains(name) {
            if let Some(base) = self.base.as_deref() {
                if base.contains(name) {
                    return base.lookup(name, id);
                }
            }
        }
        // Not indexed anywhere — build once and cache (idempotent under a
        // concurrent race: the first writer wins, both indices are equal).
        let built = build();
        let mut ov = self.overlay.write().unwrap();
        ov.entry(name.to_string())
            .or_insert_with(|| TypeEntry::from(built))
            .get(id)
    }

    /// Ensure `name` is indexed (overlay or base) — the `&self` pre-warm
    /// counterpart of the self-healing read path, so callers can pre-build an
    /// id index without `&mut` and its `Arc::make_mut` deep copy. Racing
    /// builders are harmless: the first writer wins and both indices are equal.
    pub fn ensure(&self, name: &str, build: impl FnOnce() -> TypeIdIndex) {
        if self.contains_key(name) {
            return;
        }
        let built = build();
        let mut ov = self.overlay.write().unwrap();
        ov.entry(name.to_string())
            .or_insert_with(|| TypeEntry::from(built));
    }

    /// Look up without building — None when the type isn't indexed.
    pub fn lookup(&self, name: &str, id: &Value) -> Option<NodeIndex> {
        {
            let ov = self.overlay.read().unwrap();
            if let Some(idx) = ov.get(name) {
                return idx.get(id);
            }
        }
        if self.removed.contains(name) {
            return None;
        }
        self.base.as_deref().and_then(|b| {
            if b.contains(name) {
                b.lookup(name, id)
            } else {
                None
            }
        })
    }

    /// Borrow `source`'s and `target`'s **overlay-resident** id indices in
    /// place for the length of one bulk pass, and run `f` against them.
    ///
    /// This is the probing counterpart to [`Self::materialize_type`]: one
    /// lock acquisition and two borrowed entries, where materializing pays a
    /// map insert per node *of the whole type* before the first row is looked
    /// at. On a property-free `add_connections` at 100k nodes / 24k edges,
    /// materializing was 53% of the call (samply, 2026-08-15) and grew with
    /// the graph while the row count stayed fixed.
    ///
    /// Returns `None` — without calling `f` — unless **both** types are in the
    /// overlay, which is every heap-resident graph and every type a loaded
    /// graph has since mutated. A base (mmap) entry is deliberately excluded:
    /// its Integer variant answers a probe with a binary search over the
    /// mapped file, so trading one materialization for R such probes is a
    /// regime question rather than a win, and the caller keeps the
    /// materializing path there.
    ///
    /// `f` must not re-enter the store — the read lock is held for its whole
    /// execution.
    pub fn with_overlay_type_pair<R>(
        &self,
        source: &str,
        target: &str,
        f: impl FnOnce(&TypeEntry, &TypeEntry) -> R,
    ) -> Option<R> {
        let overlay = self.overlay.read().unwrap();
        let source_entry = overlay.get(source)?;
        let target_entry = if source == target {
            source_entry
        } else {
            overlay.get(target)?
        };
        Some(f(source_entry, target_entry))
    }

    /// Materialize the full `id → NodeIndex` map for a type, or None when the
    /// type isn't indexed. Used by `CombinedTypeLookup::from_id_indices`, which
    /// resolves connection rows' endpoints against the endpoint types.
    pub fn materialize_type(&self, name: &str) -> Option<FxHashMap<Value, NodeIndex>> {
        {
            let ov = self.overlay.read().unwrap();
            if let Some(entry) = ov.get(name) {
                return Some(entry.materialize().iter().collect());
            }
        }
        if self.removed.contains(name) {
            return None;
        }
        let base = self.base.as_deref()?;
        if base.contains(name) {
            base.materialize(name).map(|ti| ti.iter().collect())
        } else {
            None
        }
    }

    pub fn insert(&mut self, name: String, idx: TypeIdIndex) {
        self.removed.remove(&name);
        self.overlay
            .get_mut()
            .unwrap()
            .insert(name, TypeEntry::from(idx));
    }

    /// Number of ids indexed for `name` in the mutable overlay, or `None`
    /// when the type is not overlay-resident. Deliberately does not consult
    /// the mmap'd base: the only caller uses this to decide whether an
    /// in-place edit is safe, and base entries are never edited in place.
    pub fn overlay_len(&self, name: &str) -> Option<usize> {
        self.overlay
            .read()
            .unwrap()
            .get(name)
            .map(|entry| entry.len())
    }

    /// Drop `entries` (`id → node`) from `name`'s index in place, instead of
    /// invalidating the whole type.
    ///
    /// Deleting one node used to `remove()` the entire type index, so the next
    /// id lookup rebuilt it by scanning every node of the type — an O(N_type)
    /// cost charged to a single-node delete. The create path maintains the
    /// index incrementally the same way (the `pk_id` match in the create
    /// executor).
    ///
    /// Falls back to whole-type invalidation, and returns `false`, whenever the
    /// index is not overlay-resident — an unbuilt type has nothing to edit, and
    /// a base-resident type lives in an immutable mmap. Each entry is removed
    /// only if it still resolves to the given node, so a re-pointed id is left
    /// intact.
    ///
    /// The caller is responsible for the duplicate-id precondition: this edits
    /// exactly the ids it is given, whereas a rebuild re-derives the whole map
    /// and would surface a shadowed duplicate. See `detach_delete_nodes`.
    pub fn evict_entries(&mut self, name: &str, entries: &[(Value, NodeIndex)]) -> bool {
        let overlay = self.overlay.get_mut().unwrap();
        let Some(entry) = overlay.get_mut(name) else {
            self.remove(name);
            return false;
        };
        for (id, idx) in entries {
            entry.remove_matching(id, *idx);
        }
        true
    }

    pub fn remove(&mut self, name: &str) -> Option<TypeIdIndex> {
        let prev = self
            .overlay
            .get_mut()
            .unwrap()
            .remove(name)
            .map(|entry| entry.materialize());
        if self.base.as_ref().is_some_and(|b| b.contains(name)) {
            self.removed.insert(name.to_string());
        }
        prev
    }

    pub fn clear(&mut self) {
        self.overlay.get_mut().unwrap().clear();
        if let Some(base) = &self.base {
            self.removed.extend(base.dir.keys().cloned());
        }
    }

    pub fn len(&self) -> usize {
        let overlay = self.overlay.read().unwrap();
        let base_count = self
            .base
            .as_ref()
            .map(|b| b.dir.keys().filter(|k| !self.removed.contains(*k)).count())
            .unwrap_or(0);
        let overlay_only = overlay
            .keys()
            .filter(|k| self.base.as_ref().map(|b| !b.contains(k)).unwrap_or(true))
            .count();
        base_count + overlay_only
    }

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

    /// Owned snapshot of every live `TypeIdIndex` (overlay first, then base
    /// entries that aren't shadowed/removed). Cold path — used by N-Triples
    /// export. Returns owned indices because the read lock can't be held
    /// across the caller's iteration.
    pub fn values(&self) -> Vec<TypeIdIndex> {
        self.snapshot().into_iter().map(|(_, v)| v).collect()
    }

    /// Owned `(name, TypeIdIndex)` snapshot of every live entry. Cold path —
    /// used by save.
    pub fn iter(&self) -> Vec<(String, TypeIdIndex)> {
        self.snapshot()
    }

    fn snapshot(&self) -> Vec<(String, TypeIdIndex)> {
        let overlay = self.overlay.read().unwrap();
        let mut out: Vec<(String, TypeIdIndex)> = overlay
            .iter()
            .map(|(k, v)| (k.clone(), v.materialize()))
            .collect();
        if let Some(base) = self.base.as_deref() {
            for k in base.dir.keys() {
                if !overlay.contains_key(k.as_str()) && !self.removed.contains(k.as_str()) {
                    if let Some(materialized) = base.materialize(k) {
                        out.push((k.clone(), materialized));
                    }
                }
            }
        }
        out
    }

    /// HashMap-`entry`-shaped accessor: materialize any base entry into the
    /// overlay (or default-construct), then hand back a `&mut` to it. Used
    /// wherever the index is maintained incrementally: the N-Triples and RDF
    /// loaders' per-entity build, and the create paths. `&mut self` gives
    /// exclusive access, so `get_mut()` is uncontended (no lock cost).
    pub fn entry_or_default(&mut self, name: String) -> &mut TypeEntry {
        let needs_materialize = {
            let overlay = self.overlay.get_mut().unwrap();
            !overlay.contains_key(&name) && !self.removed.contains(&name)
        };
        if needs_materialize {
            if let Some(base) = self.base.as_deref() {
                if let Some(materialized) = base.materialize(&name) {
                    self.overlay
                        .get_mut()
                        .unwrap()
                        .insert(name.clone(), TypeEntry::from(materialized));
                }
            }
        }
        self.removed.remove(&name);
        self.overlay.get_mut().unwrap().entry(name).or_default()
    }

    /// Fold every shared base back in where this graph is its last holder.
    ///
    /// Called at write entry beside `GraphBackend::try_compact`, so the
    /// "hold a view, write, drop the view, write again" sequence returns to the
    /// flat representation on the very next write. Per entry the fold is a plain
    /// map overwrite: unlike the topology overlay there is no slot to predict,
    /// because the delta already recorded the real `NodeIndex` values the graph
    /// handed out. What it must not do is edit a base another graph is reading,
    /// which is what `Arc::get_mut` inside `TypeEntry::try_compact` gates.
    pub fn try_compact(&mut self) {
        for entry in self.overlay.get_mut().unwrap().values_mut() {
            entry.try_compact();
        }
    }

    /// Replace the entire store with a fresh HashMap (used by load fallback
    /// for legacy `.bin.zst`-only graphs and by `reindex()`).
    pub fn replace_with(&mut self, map: HashMap<String, TypeIdIndex>) {
        *self.overlay.get_mut().unwrap() = map
            .into_iter()
            .map(|(name, index)| (name, TypeEntry::from(index)))
            .collect();
        self.removed.clear();
        self.base = None;
    }
}

/// Coerce a `Value` to `u32` for binary search on the Integer variant.
/// Mirrors the matching branches in `TypeIdIndex::get`.
fn coerce_to_u32(id: &Value) -> Option<u32> {
    match id {
        Value::UniqueId(u) => Some(*u),
        Value::Int64(i) => {
            if *i >= 0 && *i <= u32::MAX as i64 {
                Some(*i as u32)
            } else {
                None
            }
        }
        Value::Float64(f) => {
            if f.fract() == 0.0 {
                let i = *f as i64;
                if i >= 0 && i <= u32::MAX as i64 {
                    Some(i as u32)
                } else {
                    None
                }
            } else {
                None
            }
        }
        // No string→u32 coercion: a String id matches only by exact value.
        _ => None,
    }
}

/// Write `id_indices.bin` (raw mmap layout). Iterates the store's union view
/// (overlay + base) so saves capture both fresh mutations and unchanged
/// base entries.
///
/// Directory keys are `InternedKey` hashes and the loader turns each one back
/// into a type name through the interner sidecar that ships with the same
/// snapshot. Names are therefore *resolved*, never interned here: deriving a
/// key from the name alone — the old `interner.clone().try_get_or_intern`,
/// which registered the name in a throwaway clone — manufactured keys for
/// names the persisted interner never carried, and a single such entry made
/// the whole directory unloadable ("directory contains an unresolved type
/// key").
///
/// Only an empty index can carry an unregistered name: creating a node interns
/// its type (`mutation/batch.rs`), so any index holding ids names an interned
/// type. Empty entries do reach the store — the read path caches a
/// build-on-miss index for a label a query merely mentioned
/// ([`IdIndexStore::lookup_or_build`]) — and they are pure cache, so dropping
/// them loses nothing. A non-empty unregistered entry would be a broken
/// invariant, so it fails the save rather than shipping a directory that
/// cannot be read back.
pub fn write_id_indices_bin(
    dir: &Path,
    store: &IdIndexStore,
    interner: &StringInterner,
) -> Result<(), String> {
    let mut entries: Vec<(u64, TypeIdIndex)> = Vec::new();
    for (name, materialized) in store.iter() {
        let Some(key) = interner.try_resolve_to_key(&name) else {
            if materialized.is_empty() {
                continue;
            }
            return Err(format!(
                "id index for type '{name}' holds {} ids but the type name is \
                 not in the graph's interner; refusing to write an \
                 id_indices.bin that cannot be read back",
                materialized.len()
            ));
        };
        entries.push((key.as_u64(), materialized));
    }
    entries.sort_by_key(|(k, _)| *k);

    let num_types = entries.len();
    let header_size = HEADER_BYTES;
    let dir_size = DIR_ENTRY_BYTES * num_types;
    let data_offset = header_size + dir_size;

    // Pre-compute payload offsets/lengths so we can emit the directory first.
    struct Plan {
        type_key: u64,
        variant: u8,
        num_entries: u64,
        payload_off: u64,
        payload_len: u64,
        data: Vec<u8>,
    }

    let mut plans: Vec<Plan> = Vec::with_capacity(num_types);
    let mut cursor = data_offset as u64;

    for (type_key, idx) in &entries {
        match idx {
            TypeIdIndex::Integer(map) => {
                let mut pairs: Vec<(u32, u32)> =
                    map.iter().map(|(k, v)| (*k, v.index() as u32)).collect();
                pairs.sort_by_key(|(k, _)| *k);
                let n = pairs.len();
                let mut data = Vec::with_capacity(n * 8);
                for (k, _) in &pairs {
                    data.extend_from_slice(&k.to_le_bytes());
                }
                for (_, v) in &pairs {
                    data.extend_from_slice(&v.to_le_bytes());
                }
                let len = data.len() as u64;
                plans.push(Plan {
                    type_key: *type_key,
                    variant: 0,
                    num_entries: n as u64,
                    payload_off: cursor,
                    payload_len: len,
                    data,
                });
                cursor += len;
            }
            TypeIdIndex::General(map) => {
                let blob = serde_codec::encode_versioned(
                    serde_codec::CURRENT_CODEC,
                    map,
                    MAX_GENERAL_INDEX_DECODE_BYTES,
                )
                .map_err(|e| format!("id_indices General-variant codec failed: {e}"))?;
                let len = blob.len() as u64;
                plans.push(Plan {
                    type_key: *type_key,
                    variant: 1,
                    num_entries: map.len() as u64,
                    payload_off: cursor,
                    payload_len: len,
                    data: blob,
                });
                cursor += len;
            }
        }
    }

    let total = cursor as usize;
    let mut out = Vec::with_capacity(total);
    // Header
    out.extend_from_slice(MAGIC);
    out.extend_from_slice(&VERSION.to_le_bytes());
    out.extend_from_slice(&(num_types as u32).to_le_bytes());
    out.extend_from_slice(&(HEADER_BYTES as u64).to_le_bytes());
    out.extend_from_slice(&(data_offset as u64).to_le_bytes());

    // Directory
    for plan in &plans {
        out.extend_from_slice(&plan.type_key.to_le_bytes());
        out.push(plan.variant);
        out.extend_from_slice(&[0u8; 7]);
        out.extend_from_slice(&plan.num_entries.to_le_bytes());
        out.extend_from_slice(&plan.payload_off.to_le_bytes());
        out.extend_from_slice(&plan.payload_len.to_le_bytes());
        out.extend_from_slice(&[0u8; 8]);
    }

    // Data
    for plan in plans {
        out.extend_from_slice(&plan.data);
    }

    debug_assert_eq!(out.len(), total);

    std::fs::write(dir.join("id_indices.bin"), out)
        .map_err(|e| format!("Failed to write id_indices.bin: {}", e))?;
    Ok(())
}

#[cfg(test)]
mod validation_tests {
    use super::*;
    use crate::graph::storage::disk::temp_owner::{TempGraphDir, TrackedOwner};

    fn integer_fixture(type_key: u64, pairs: &[(u32, u32)]) -> Vec<u8> {
        let data_offset = HEADER_BYTES + DIR_ENTRY_BYTES;
        let mut bytes = Vec::new();
        bytes.extend_from_slice(MAGIC);
        bytes.extend_from_slice(&VERSION.to_le_bytes());
        bytes.extend_from_slice(&1u32.to_le_bytes());
        bytes.extend_from_slice(&(HEADER_BYTES as u64).to_le_bytes());
        bytes.extend_from_slice(&(data_offset as u64).to_le_bytes());
        bytes.extend_from_slice(&type_key.to_le_bytes());
        bytes.push(0);
        bytes.extend_from_slice(&[0; 7]);
        bytes.extend_from_slice(&(pairs.len() as u64).to_le_bytes());
        bytes.extend_from_slice(&(data_offset as u64).to_le_bytes());
        bytes.extend_from_slice(&((pairs.len() * 8) as u64).to_le_bytes());
        bytes.extend_from_slice(&[0; 8]);
        for (key, _) in pairs {
            bytes.extend_from_slice(&key.to_le_bytes());
        }
        for (_, node) in pairs {
            bytes.extend_from_slice(&node.to_le_bytes());
        }
        bytes
    }

    /// A loaded [`IdIndexBase`] together with the temp directory its `mmap`
    /// points into. Field order is the contract: `base` drops before `temp`,
    /// and `temp`'s guard asserts it.
    ///
    /// The previous helper returned the base alone, so the `TempDir` local
    /// was dropped the moment `load` returned and every assertion below ran
    /// against an unlinked inode — valid on Unix, and therefore silent.
    struct LoadedIndex {
        base: TrackedOwner<IdIndexBase>,
        /// Held only for its `Drop`: it asserts `base` above is gone.
        _temp: TempGraphDir,
    }

    impl LoadedIndex {
        fn base(&self) -> &IdIndexBase {
            &self.base
        }
    }

    fn load(bytes: &[u8], interner: &StringInterner) -> std::io::Result<Option<LoadedIndex>> {
        let temp = TempGraphDir::new();
        std::fs::write(temp.path().join("id_indices.bin"), bytes).unwrap();
        let Some(base) = IdIndexBase::load_from(temp.path(), interner)? else {
            return Ok(None);
        };
        let base = temp.own("IdIndexBase", base);
        Ok(Some(LoadedIndex { base, _temp: temp }))
    }

    fn assert_invalid(bytes: &[u8], interner: &StringInterner) {
        let outcome = std::panic::catch_unwind(|| load(bytes, interner));
        match outcome.expect("invalid index must not panic") {
            Err(error) => assert_eq!(error.kind(), std::io::ErrorKind::InvalidData),
            Ok(_) => panic!("invalid index loaded successfully"),
        }
    }

    #[test]
    fn integer_fixture_reads_canonical_little_endian_bytes() {
        let mut interner = StringInterner::new();
        let key = interner.get_or_intern("Person").as_u64();
        let loaded = load(&integer_fixture(key, &[(7, 70), (42, 420)]), &interner)
            .unwrap()
            .unwrap();
        let base = loaded.base();
        assert_eq!(
            base.lookup("Person", &Value::UniqueId(7)),
            Some(NodeIndex::new(70))
        );
        assert_eq!(
            base.lookup("Person", &Value::UniqueId(42)),
            Some(NodeIndex::new(420))
        );
    }

    #[test]
    fn rejects_invalid_header_directory_and_variant() {
        let mut interner = StringInterner::new();
        let key = interner.get_or_intern("Person").as_u64();
        let valid = integer_fixture(key, &[(7, 70)]);

        let mut huge_count = valid.clone();
        huge_count[12..16].copy_from_slice(&u32::MAX.to_le_bytes());
        assert_invalid(&huge_count, &interner);
        let mut bad_dir = valid.clone();
        bad_dir[16..24].copy_from_slice(&u64::MAX.to_le_bytes());
        assert_invalid(&bad_dir, &interner);
        let mut bad_data = valid.clone();
        bad_data[24..32].copy_from_slice(&33u64.to_le_bytes());
        assert_invalid(&bad_data, &interner);
        let mut bad_variant = valid.clone();
        bad_variant[40] = 2;
        assert_invalid(&bad_variant, &interner);
    }

    #[test]
    fn rejects_bad_counts_ranges_and_integer_ordering() {
        let mut interner = StringInterner::new();
        let key = interner.get_or_intern("Person").as_u64();
        let valid = integer_fixture(key, &[(7, 70), (42, 420)]);

        let mut too_many = valid.clone();
        too_many[48..56].copy_from_slice(&(u32::MAX as u64 + 1).to_le_bytes());
        assert_invalid(&too_many, &interner);
        let mut past_eof = valid.clone();
        past_eof[56..64].copy_from_slice(&u64::MAX.to_le_bytes());
        assert_invalid(&past_eof, &interner);
        let mut wrong_len = valid.clone();
        wrong_len[64..72].copy_from_slice(&15u64.to_le_bytes());
        assert_invalid(&wrong_len, &interner);
        assert_invalid(&integer_fixture(key, &[(42, 1), (7, 2)]), &interner);
        assert_invalid(&integer_fixture(key, &[(7, 1), (7, 2)]), &interner);
    }

    #[test]
    fn rejects_malformed_general_postcard_during_load() {
        let mut interner = StringInterner::new();
        let key = interner.get_or_intern("StringIds").as_u64();
        let mut bytes = integer_fixture(key, &[(1, 1)]);
        bytes[40] = 1;
        bytes[64..72].copy_from_slice(&16u64.to_le_bytes());
        bytes.truncate(HEADER_BYTES + DIR_ENTRY_BYTES);
        bytes.extend_from_slice(&1u64.to_le_bytes());
        bytes.extend_from_slice(&[0xff; 8]);
        assert_invalid(&bytes, &interner);
    }

    #[test]
    fn writer_round_trip_accepts_unaligned_integer_payload_after_general() {
        let temp = tempfile::tempdir().unwrap();
        let mut interner = StringInterner::new();
        let candidates = ["Alpha", "Beta"];
        for name in candidates {
            interner.get_or_intern(name);
        }
        let mut ordered = candidates;
        ordered.sort_by_key(|name| InternedKey::from_str(name).as_u64());
        let general_name = ordered[0];
        let integer_name = ordered[1];
        let general = TypeIdIndex::General(FxHashMap::from_iter([(
            Value::String("x".into()),
            NodeIndex::new(3),
        )]));
        let integer = TypeIdIndex::Integer(FxHashMap::from_iter([(7, NodeIndex::new(4))]));
        let mut store = IdIndexStore::default();
        store.replace_with(HashMap::from([
            (general_name.to_string(), general),
            (integer_name.to_string(), integer),
        ]));
        write_id_indices_bin(temp.path(), &store, &interner).unwrap();

        let raw = std::fs::read(temp.path().join("id_indices.bin")).unwrap();
        let second_payload_off = u64::from_le_bytes(
            raw[HEADER_BYTES + DIR_ENTRY_BYTES + 24..HEADER_BYTES + DIR_ENTRY_BYTES + 32]
                .try_into()
                .unwrap(),
        );
        assert_ne!(
            second_payload_off % 4,
            0,
            "fixture must exercise an unaligned integer payload"
        );

        let base = IdIndexBase::load_from(temp.path(), &interner)
            .unwrap()
            .unwrap();
        assert_eq!(
            base.lookup(general_name, &Value::String("x".into())),
            Some(NodeIndex::new(3))
        );
        assert_eq!(
            base.lookup(integer_name, &Value::UniqueId(7)),
            Some(NodeIndex::new(4))
        );
    }

    /// A directory saved before the writer resolved its keys carries an empty
    /// index under a type key the interner sidecar never received. Both
    /// variants of that entry are recovered rather than failing the load; a
    /// populated entry under an unresolvable key still fails (asserted in
    /// `rejects_unsupported_unresolved_and_trailing_data`).
    #[test]
    fn an_empty_entry_with_an_unresolved_type_key_is_recovered() {
        let mut interner = StringInterner::new();
        interner.get_or_intern("Person");
        let stale = InternedKey::from_str("NeverInterned").as_u64();

        let loaded = load(&integer_fixture(stale, &[]), &interner)
            .expect("a stale empty entry must not fail the load")
            .unwrap();
        assert!(!loaded.base().contains("NeverInterned"));

        // The General variant is what a type with no rows actually produced:
        // `num_entries` is 0, but the Postcard payload of an empty map is not.
        let blob = serde_codec::encode_versioned(
            serde_codec::CURRENT_CODEC,
            &HashMap::<Value, NodeIndex>::new(),
            MAX_GENERAL_INDEX_DECODE_BYTES,
        )
        .unwrap();
        let mut bytes = integer_fixture(stale, &[]);
        bytes[40] = 1;
        bytes[64..72].copy_from_slice(&(blob.len() as u64).to_le_bytes());
        bytes.extend_from_slice(&blob);
        let loaded = load(&bytes, &interner)
            .expect("a stale empty General entry must not fail the load")
            .unwrap();
        assert!(!loaded.base().contains("NeverInterned"));
    }

    /// The writer resolves names against the interner it is handed, so it can
    /// never emit a directory key that the matching `interner.bin.zst` fails
    /// to resolve. An unregistered name is dropped when its index is empty
    /// (pure cache, rebuilt on demand) and fails the save when it is not.
    #[test]
    fn writer_never_emits_a_key_the_interner_cannot_resolve() {
        let temp = tempfile::tempdir().unwrap();
        let mut interner = StringInterner::new();
        interner.get_or_intern("Known");

        let mut store = IdIndexStore::default();
        store.replace_with(HashMap::from([
            (
                "Known".to_string(),
                TypeIdIndex::Integer(FxHashMap::from_iter([(7, NodeIndex::new(1))])),
            ),
            // Never interned: an id index cached for a type with no rows.
            ("Unregistered".to_string(), TypeIdIndex::default()),
        ]));
        write_id_indices_bin(temp.path(), &store, &interner).unwrap();

        // Asserted on the bytes, not through the loader: the loader also
        // recovers such an entry (directories written before this fix carry
        // them), so a round trip alone would not pin the writer.
        let raw = std::fs::read(temp.path().join("id_indices.bin")).unwrap();
        assert_eq!(
            u32::from_le_bytes(raw[12..16].try_into().unwrap()),
            1,
            "the unregistered name must not reach the directory at all"
        );

        let base = IdIndexBase::load_from(temp.path(), &interner)
            .expect("the written directory must load")
            .unwrap();
        assert_eq!(
            base.lookup("Known", &Value::UniqueId(7)),
            Some(NodeIndex::new(1))
        );
        assert!(!base.contains("Unregistered"));

        store.insert(
            "Unregistered".to_string(),
            TypeIdIndex::Integer(FxHashMap::from_iter([(1, NodeIndex::new(0))])),
        );
        let error = write_id_indices_bin(temp.path(), &store, &interner).unwrap_err();
        assert!(error.contains("Unregistered"), "{error}");
        assert!(error.contains("cannot be read back"), "{error}");
    }

    #[test]
    fn rejects_unsupported_unresolved_and_trailing_data() {
        let mut interner = StringInterner::new();
        let key = interner.get_or_intern("Person").as_u64();
        let valid = integer_fixture(key, &[(7, 70)]);
        let mut version = valid.clone();
        version[8..12].copy_from_slice(&(VERSION + 1).to_le_bytes());
        assert_invalid(&version, &interner);
        let mut trailing = valid.clone();
        trailing.push(0);
        assert_invalid(&trailing, &interner);
        assert_invalid(&integer_fixture(key.wrapping_add(1), &[(7, 70)]), &interner);
    }
}