kglite 0.16.3

// src/graph/storage/disk/edge_properties.rs
//
// Disk-backed edge property storage. Replaces the heap-only
// `HashMap<u32, Vec<(InternedKey, Value)>>` that blew RAM at Wikidata
// scale (30–60 GB). PR2 of the disk-graph-improvement-plan.
//
// Layout: per-edge columnar slots indexed by edge_idx.
//
//   edge_prop_offsets.bin  MmapOrVec<u64>, (max_edge_idx + 1) entries.
//                          offsets[i]..offsets[i+1] = byte range in heap
//                          for edge i's Postcard-serialized props blob.
//                          offsets[i] == offsets[i+1] means "no props".
//   edge_prop_heap.bin     MmapBytes, variable-length. Each populated
//                          slot holds a Postcard `Vec<(u64, Value)>` —
//                          raw InternedKey hashes, no interner needed
//                          on the read path.
//
// Runtime state: columnar `base` (read-only mmap) + HashMap `overlay`
// for edges mutated since last save. The overlay grows with mutation
// count, not graph size — bounded by workload, per the bounded-memory
// constraint in `feedback_bounded_memory.md`.
//
// Sequential access pattern: iterating edges in edge_idx order reads
// offsets and heap linearly. Random single-edge lookups incur one
// page fault per file, which is unavoidable and still cheaper than
// the current zstd-decode-whole-HashMap-at-load path.
//
// Pre-0.14 format 0/1 graphs are rejected before payload decoding.
// Current columnar reads store raw hashes and never touch the interner.

use crate::datatypes::values::Value;
use crate::graph::schema::{InternedKey, StringInterner};
use crate::graph::storage::mapped::mmap_vec::{MmapBytes, MmapOrVec};
use std::borrow::Cow;
use std::collections::HashMap;
use std::io;
use std::path::Path;
use std::sync::Arc;

/// Filename constants for the columnar base. Both live under the graph's
/// data directory alongside the existing CSR/column files.
pub const OFFSETS_FILE: &str = "edge_prop_offsets.bin";
pub const HEAP_FILE: &str = "edge_prop_heap.bin";

/// Retired pre-0.14 combined edge-property file, removed on current saves.
pub const LEGACY_FILE: &str = "edge_properties.bin.zst";

/// Lengths needed to mmap the columnar files at load time. Persisted
/// in `DiskGraphMeta` so the loader can `load_mapped` without scanning
/// file sizes at runtime.
#[derive(Debug, serde::Serialize, serde::Deserialize, Default, Clone, Copy)]
pub struct EdgePropertyStoreMeta {
    /// Number of u64 offsets written. Equal to `(upper_bound + 1)` passed
    /// to `save_to`. `offsets_len * 8` is the byte size of the offsets file.
    pub offsets_len: usize,
    /// Heap byte length. The mmap'd file may be padded but only
    /// `heap_len` bytes are valid.
    pub heap_len: usize,
}

/// Disk-backed columnar snapshot of edge properties. Read-only after load.
#[derive(Debug)]
struct ColumnarBase {
    /// offsets[edge_idx]..offsets[edge_idx + 1] = byte range in heap.
    offsets: MmapOrVec<u64>,
    /// Concatenated Postcard blobs. Each populated slot is a
    /// `Vec<(u64, Value)>` — u64 is the raw `InternedKey` hash.
    heap: MmapBytes,
    codec: crate::serde_codec::CodecVersion,
}

impl ColumnarBase {
    /// Byte slice for a single edge's props blob, or empty slice if
    /// the edge has no properties in this base.
    fn slot(&self, edge_idx: u32) -> Option<&[u8]> {
        let i = edge_idx as usize;
        // offsets has length (upper_bound + 1). Edges past that are
        // "not in base" — caller must consult overlay.
        if i + 1 >= self.offsets.len() {
            return None;
        }
        let start = self.offsets.get(i) as usize;
        let end = self.offsets.get(i + 1) as usize;
        if start == end {
            // Empty slot (sparsity encoding).
            return Some(&[]);
        }
        Some(self.heap.slice(start, end))
    }

    /// Upper bound on edge_idx values covered by this base (exclusive).
    fn len(&self) -> u32 {
        // (upper_bound + 1) offsets written; final trailing offset is
        // the total heap length. So the number of edges covered is
        // offsets.len().saturating_sub(1).
        self.offsets.len().saturating_sub(1) as u32
    }
}

/// Decode a single slot's bytes into `(InternedKey, Value)` pairs.
/// Assumes bytes were produced by `encode_props_into` at save time.
fn decode_props(
    codec: crate::serde_codec::CodecVersion,
    bytes: &[u8],
) -> Option<Vec<(InternedKey, Value)>> {
    let raw: Vec<(u64, Value)> = crate::serde_codec::decode_exact_with(
        codec,
        bytes,
        bytes.len() as u64,
        crate::serde_codec::DecodeLimits::new(bytes.len() as u64, bytes.len() as u64),
    )
    .ok()?;
    Some(
        raw.into_iter()
            .map(|(k, v)| (InternedKey::from_u64(k), v))
            .collect(),
    )
}

/// Encode `(InternedKey, Value)` pairs directly into the provided heap
/// buffer. Avoids the per-edge `Vec<u8>` allocation that dominated
/// save-path cost during PR2 phase-2 benchmarking. The interner is not
/// consulted — we store the raw u64 hash.
fn encode_props_into(props: &[(InternedKey, Value)], heap: &mut Vec<u8>) -> io::Result<()> {
    let raw: Vec<(u64, &Value)> = props.iter().map(|(k, v)| (k.as_u64(), v)).collect();
    let encoded =
        crate::serde_codec::encode_versioned(crate::serde_codec::CURRENT_CODEC, &raw, u64::MAX)
            .map_err(io::Error::other)?;
    heap.extend_from_slice(&encoded);
    Ok(())
}

/// Edge-property store: columnar disk base + in-memory mutation overlay.
///
/// `None` overlay entries are tombstones (edge was deleted or its props
/// were explicitly emptied). `Some(vec)` entries replace whatever the
/// base had for that edge. Entries absent from the overlay fall through
/// to the base.
#[derive(Debug, Default)]
pub struct EdgePropertyStore {
    base: Option<Arc<ColumnarBase>>,
    overlay: HashMap<u32, Option<Vec<(InternedKey, Value)>>>,
}

impl EdgePropertyStore {
    /// Empty store — no base, empty overlay. Used by fresh in-memory builds.
    pub fn new() -> Self {
        Self::default()
    }

    /// Construct from a pre-populated HashMap.
    pub fn from_overlay(map: HashMap<u32, Vec<(InternedKey, Value)>>) -> Self {
        Self {
            base: None,
            overlay: map.into_iter().map(|(k, v)| (k, Some(v))).collect(),
        }
    }

    /// Lookup an edge's current properties.
    /// - Overlay hit: returns `Cow::Borrowed` (zero copy).
    /// - Overlay tombstone: returns `None`.
    /// - Base hit: deserializes the columnar slot into an owned `Vec`
    ///   and wraps in `Cow::Owned`. No interner needed — the blob stores
    ///   raw `InternedKey` u64 hashes directly.
    pub fn get(&self, edge_idx: u32) -> Option<Cow<'_, [(InternedKey, Value)]>> {
        // Overlay first. `Some(None)` = explicit tombstone, hide base.
        if let Some(entry) = self.overlay.get(&edge_idx) {
            return entry.as_ref().map(|v| Cow::Borrowed(v.as_slice()));
        }
        let base = self.base.as_ref()?;
        let bytes = base.slot(edge_idx)?;
        if bytes.is_empty() {
            return None;
        }
        let decoded = decode_props(base.codec, bytes)?;
        if decoded.is_empty() {
            None
        } else {
            Some(Cow::Owned(decoded))
        }
    }

    /// Replace an edge's properties in the overlay.
    pub fn insert(&mut self, edge_idx: u32, props: Vec<(InternedKey, Value)>) {
        if props.is_empty() {
            // Normalise to tombstone/absent so `is_empty()` is meaningful.
            self.remove(edge_idx);
            return;
        }
        self.overlay.insert(edge_idx, Some(props));
    }

    /// Remove an edge's properties. Writes a tombstone into the overlay if
    /// the base might contain this edge; otherwise just drops the overlay entry.
    pub fn remove(&mut self, edge_idx: u32) {
        if let Some(base) = self.base.as_ref() {
            if edge_idx < base.len() {
                self.overlay.insert(edge_idx, None);
                return;
            }
        }
        self.overlay.remove(&edge_idx);
    }

    /// Remove and return the current properties, if any. Mirrors
    /// `HashMap::remove`. Used by `compact_edges` which remaps edge_idx.
    pub fn take(&mut self, edge_idx: u32) -> Option<Vec<(InternedKey, Value)>> {
        let current = self
            .get(edge_idx)
            .map(|cow| cow.into_owned())
            .filter(|v| !v.is_empty());
        self.remove(edge_idx);
        current
    }

    /// True when no edge currently has any properties. Conservative: a
    /// loaded graph with all entries tombstoned in the overlay still
    /// reports `false`. That's OK — the save path always writes a valid
    /// (possibly all-zero-length) file and callers gate on the existence
    /// of at least one property-bearing edge.
    pub fn is_empty(&self) -> bool {
        if self.base.as_ref().is_some_and(|b| b.len() > 0) {
            return false;
        }
        self.overlay.values().all(|v| v.is_none())
    }

    /// Smallest exclusive upper bound on edge_idx values that *might*
    /// have properties in this store. Used by callers (like `compact_edges`)
    /// that need to iterate every potentially-populated slot to remap
    /// indices. The returned value is `max(base.len(), max_overlay_idx + 1)`.
    pub fn upper_bound(&self) -> u32 {
        let base_upper = self.base.as_ref().map(|b| b.len()).unwrap_or(0);
        let overlay_upper = self
            .overlay
            .keys()
            .max()
            .copied()
            .map(|k| k.saturating_add(1))
            .unwrap_or(0);
        base_upper.max(overlay_upper)
    }

    /// Fork a transaction overlay while sharing the immutable columnar base.
    /// Only the mutation-sized overlay is copied.
    pub(crate) fn fork_overlay(&self) -> Self {
        Self {
            base: self.base.clone(),
            overlay: self.overlay.clone(),
        }
    }

    /// Write the current merged state (base ∪ overlay, minus tombstones)
    /// to `target_dir` as the columnar format, then clear the overlay.
    ///
    /// `upper_bound` is an exclusive upper bound on edge_idx values that
    /// need slots in the offsets array. Typically `DiskGraph::next_edge_idx`.
    /// Slots past any actual data are represented as zero-length entries.
    ///
    /// After `save_to` returns, `self.base` is *not* automatically
    /// re-opened — callers that want subsequent reads to hit the freshly
    /// written files should `*self = Self::load_from(target_dir, 2, meta, ...)`.
    pub fn save_to(&mut self, target_dir: &Path, upper_bound: u32) -> io::Result<()> {
        let offsets_path = target_dir.join(OFFSETS_FILE);
        let heap_path = target_dir.join(HEAP_FILE);

        // Release any existing mmap over these exact paths before we
        // overwrite them — memmap2 docs are explicit that overlapping
        // writes to a mapped file are UB.
        if let Some(base) = self.base.as_ref() {
            if base.offsets.file_path() == Some(&offsets_path) {
                self.base = None;
            }
        }

        // 0.8.12 phase-1 fast path: skip the O(upper_bound) sweep when
        // no edge has any properties. The pre-phase-1 loop ran 6.7M
        // overlay HashMap lookups and wrote a 54 MB all-zero offsets
        // file on every save of a property-less wiki500m graph — ~1–2 s
        // per save. Now we emit zero-length files and reload resolves
        // to an empty base (see matching guard in `load_from`). Every
        // `get(edge_idx)` on the reloaded store returns `None`, which
        // is exactly what the full-sweep path produced anyway.
        if self.is_empty() {
            std::fs::write(&offsets_path, b"")?;
            std::fs::write(&heap_path, b"")?;
            let legacy = target_dir.join(LEGACY_FILE);
            if legacy.exists() {
                let _ = std::fs::remove_file(&legacy);
            }
            self.overlay.clear();
            return Ok(());
        }

        let mut offsets: Vec<u64> = Vec::with_capacity(upper_bound as usize + 1);
        // Pre-size heap based on overlay content — avoids repeated realloc.
        let heap_hint: usize = self
            .overlay
            .values()
            .map(|v| v.as_ref().map_or(0, |p| 32 + 16 * p.len()))
            .sum();
        let mut heap: Vec<u8> = Vec::with_capacity(heap_hint);

        for edge_idx in 0..upper_bound {
            offsets.push(heap.len() as u64);
            if let Some(cow) = self.get(edge_idx) {
                if !cow.is_empty() {
                    encode_props_into(cow.as_ref(), &mut heap)?;
                }
            }
        }
        offsets.push(heap.len() as u64);

        MmapOrVec::from_vec(offsets).save_to_file(&offsets_path)?;
        std::fs::write(&heap_path, &heap)?;

        // Remove the legacy file if it's lingering from a format=0 load.
        let legacy = target_dir.join(LEGACY_FILE);
        if legacy.exists() {
            let _ = std::fs::remove_file(&legacy);
        }

        // Overlay has been fully absorbed into the new base file.
        self.overlay.clear();
        Ok(())
    }

    /// Open the store from a directory.
    /// - `format_version` comes from `DiskGraphMeta.edge_properties_format`
    ///   (2 = Postcard columnar).
    /// - `meta` provides the file lengths needed to mmap the columnar files.
    /// - `_interner` is retained by the storage boundary; current columnar
    ///   payloads store raw u64 hashes and never touch it.
    pub fn load_from(
        dir: &Path,
        format_version: u8,
        meta: EdgePropertyStoreMeta,
        _interner: &mut StringInterner,
    ) -> io::Result<Self> {
        if format_version < 2 {
            return Err(crate::graph::io::file::pre_014_bincode_error(
                "edge-property store",
            ));
        }
        let offsets_path = dir.join(OFFSETS_FILE);
        let heap_path = dir.join(HEAP_FILE);
        if !offsets_path.exists() {
            return Ok(Self::new());
        }
        // 0.8.12 phase-1: the save path now emits zero-length
        // offsets/heap files for the "no properties anywhere" case.
        // `MmapOrVec::load_mapped` with `len == 0` would try to
        // `map_mut` a zero-byte region, which fails on some platforms.
        // Short-circuit to an empty store — every `get()` on this
        // store will resolve via the overlay check (base is None) and
        // return `None`, matching the semantic of "no edge has props".
        if meta.offsets_len == 0 {
            return Ok(Self::new());
        }
        let offsets = MmapOrVec::<u64>::load_mapped(&offsets_path, meta.offsets_len)?;
        let heap = MmapBytes::load_mapped(&heap_path, meta.heap_len)?;
        let codec = match format_version {
            2 => crate::serde_codec::CodecVersion::PostcardV1,
            other => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    format!("unsupported edge property format {other}"),
                ));
            }
        };
        Ok(Self {
            base: Some(Arc::new(ColumnarBase {
                offsets,
                heap,
                codec,
            })),
            overlay: HashMap::new(),
        })
    }

    /// Compute the on-disk metadata for this store — call after `save_to`
    /// to get the values that belong in `DiskGraphMeta`. The returned
    /// counts reflect what `save_to` just wrote, not the in-memory state.
    pub fn meta_for(dir: &Path) -> EdgePropertyStoreMeta {
        let offsets = dir.join(OFFSETS_FILE);
        let heap = dir.join(HEAP_FILE);
        EdgePropertyStoreMeta {
            offsets_len: std::fs::metadata(&offsets)
                .map(|m| m.len() as usize / std::mem::size_of::<u64>())
                .unwrap_or(0),
            heap_len: std::fs::metadata(&heap)
                .map(|m| m.len() as usize)
                .unwrap_or(0),
        }
    }
}

#[cfg(test)]
#[allow(clippy::approx_constant)]
mod tests {
    use super::*;
    use crate::datatypes::values::Value;
    use crate::graph::schema::StringInterner;
    use tempfile::TempDir;

    fn k(s: &str, interner: &mut StringInterner) -> InternedKey {
        interner.get_or_intern(s)
    }

    #[test]
    fn empty_store_is_empty() {
        let s = EdgePropertyStore::new();
        assert!(s.is_empty());
    }

    #[test]
    fn insert_and_get_overlay_hit() {
        let mut interner = StringInterner::new();
        let mut s = EdgePropertyStore::new();
        let props = vec![(k("weight", &mut interner), Value::Float64(1.5))];
        s.insert(42, props.clone());
        let got = s.get(42).expect("should hit overlay");
        assert_eq!(got.as_ref(), props.as_slice());
        assert!(!s.is_empty());
    }

    #[test]
    fn remove_without_base_drops_entry() {
        let mut interner = StringInterner::new();
        let mut s = EdgePropertyStore::new();
        s.insert(7, vec![(k("x", &mut interner), Value::Int64(1))]);
        s.remove(7);
        assert!(s.get(7).is_none());
        assert!(s.is_empty());
    }

    #[test]
    fn insert_empty_normalises_to_absent() {
        let mut s = EdgePropertyStore::new();
        s.insert(1, vec![]);
        assert!(s.get(1).is_none());
        assert!(s.is_empty());
    }

    #[test]
    fn save_and_load_round_trip() {
        let tmp = TempDir::new().unwrap();
        let mut interner = StringInterner::new();
        let mut s = EdgePropertyStore::new();

        let p0 = vec![
            (k("name", &mut interner), Value::String("alpha".into())),
            (k("rank", &mut interner), Value::Int64(7)),
        ];
        let p1 = vec![(k("weight", &mut interner), Value::Float64(3.14))];
        s.insert(0, p0.clone());
        s.insert(3, p1.clone()); // edges 1, 2 have no props

        s.save_to(tmp.path(), 4).unwrap();

        let meta = EdgePropertyStore::meta_for(tmp.path());
        // 4 slots + 1 trailing offset = 5 offsets
        assert_eq!(meta.offsets_len, 5);
        assert!(meta.heap_len > 0);

        // Reload and verify each edge.
        let reloaded = EdgePropertyStore::load_from(tmp.path(), 2, meta, &mut interner).unwrap();
        assert_eq!(reloaded.get(0).unwrap().as_ref(), p0.as_slice());
        assert!(reloaded.get(1).is_none());
        assert!(reloaded.get(2).is_none());
        assert_eq!(reloaded.get(3).unwrap().as_ref(), p1.as_slice());
    }

    #[test]
    fn overlay_tombstones_hide_base() {
        let tmp = TempDir::new().unwrap();
        let tmp2 = TempDir::new().unwrap();
        let mut interner = StringInterner::new();

        let mut s = EdgePropertyStore::new();
        s.insert(5, vec![(k("a", &mut interner), Value::Int64(99))]);
        s.save_to(tmp.path(), 6).unwrap();

        let meta = EdgePropertyStore::meta_for(tmp.path());
        let mut reloaded =
            EdgePropertyStore::load_from(tmp.path(), 2, meta, &mut interner).unwrap();
        assert!(reloaded.get(5).is_some());

        reloaded.remove(5);
        assert!(reloaded.get(5).is_none());

        // Save+reload persists the tombstone — no trace of edge 5.
        reloaded.save_to(tmp2.path(), 6).unwrap();
        let meta2 = EdgePropertyStore::meta_for(tmp2.path());
        let after = EdgePropertyStore::load_from(tmp2.path(), 2, meta2, &mut interner).unwrap();
        assert!(after.get(5).is_none());
    }

    #[test]
    fn take_returns_and_removes() {
        let mut interner = StringInterner::new();
        let mut s = EdgePropertyStore::new();
        let p = vec![(k("t", &mut interner), Value::Boolean(true))];
        s.insert(11, p.clone());
        let taken = s.take(11).unwrap();
        assert_eq!(taken, p);
        assert!(s.get(11).is_none());
    }

    #[test]
    fn pre_014_combined_store_is_rejected() {
        let tmp = TempDir::new().unwrap();
        let mut interner = StringInterner::new();
        let error = EdgePropertyStore::load_from(
            tmp.path(),
            0,
            EdgePropertyStoreMeta::default(),
            &mut interner,
        )
        .unwrap_err();
        assert!(error.to_string().contains("pre-0.14"));
    }

    #[test]
    fn pre_014_columnar_and_unknown_versions_fail() {
        let tmp = TempDir::new().unwrap();
        let mut interner = StringInterner::new();
        let key = k("legacy-column", &mut interner);
        let raw = vec![(key.as_u64(), Value::Int64(17))];
        let heap =
            crate::serde_codec::encode_versioned(crate::serde_codec::CURRENT_CODEC, &raw, u64::MAX)
                .unwrap();
        MmapOrVec::from_vec(vec![0u64, heap.len() as u64])
            .save_to_file(&tmp.path().join(OFFSETS_FILE))
            .unwrap();
        std::fs::write(tmp.path().join(HEAP_FILE), &heap).unwrap();
        let meta = EdgePropertyStore::meta_for(tmp.path());

        let error = EdgePropertyStore::load_from(tmp.path(), 1, meta, &mut interner).unwrap_err();
        assert!(error.to_string().contains("pre-0.14"));

        let error = EdgePropertyStore::load_from(tmp.path(), 99, meta, &mut interner).unwrap_err();
        assert_eq!(error.kind(), io::ErrorKind::InvalidData);
        assert!(error
            .to_string()
            .contains("unsupported edge property format"));
    }

    #[test]
    fn empty_store_save_emits_zero_length_files_and_reload_preserves_semantics() {
        // 0.8.12 phase-1 fast path: when no edge has properties, save
        // must not sweep `0..upper_bound` or write a 54 MB all-zero
        // offsets file. Verify the files are zero-length *and* that a
        // reload with a large `upper_bound` answers `get()` as None for
        // every edge.
        let tmp = TempDir::new().unwrap();
        let mut interner = StringInterner::new();
        let mut s = EdgePropertyStore::new();

        // upper_bound is deliberately large — pre-phase-1 this would
        // have written 1_000_000 * 8 = 8 MB of zeros.
        s.save_to(tmp.path(), 1_000_000).unwrap();

        let offsets_meta = std::fs::metadata(tmp.path().join(OFFSETS_FILE)).unwrap();
        let heap_meta = std::fs::metadata(tmp.path().join(HEAP_FILE)).unwrap();
        assert_eq!(offsets_meta.len(), 0, "offsets file should be empty");
        assert_eq!(heap_meta.len(), 0, "heap file should be empty");

        let meta = EdgePropertyStore::meta_for(tmp.path());
        assert_eq!(meta.offsets_len, 0);
        assert_eq!(meta.heap_len, 0);

        let reloaded = EdgePropertyStore::load_from(tmp.path(), 2, meta, &mut interner).unwrap();
        assert!(reloaded.is_empty());
        // get() on any edge_idx returns None — matches the pre-phase-1
        // behavior where every edge had an empty slot.
        assert!(reloaded.get(0).is_none());
        assert!(reloaded.get(999_999).is_none());
        assert!(reloaded.get(u32::MAX).is_none());
    }
}