engram-retired/crates/engram-core/src/storage.rs

/// Low-level sled key/value operations for nodes, edges, vectors, and salience.
///
/// Key schema:
///   nodes:{uuid}             → bincode-encoded Node
///   edges:from:{from}:{to}  → bincode-encoded Edge
///   edges:to:{to}:{from}    → reverse index (same Edge bytes)
///   vectors:{uuid}           → raw little-endian f32 bytes
///   salience:{uuid}          → 4-byte little-endian f32
use crate::error::{EngramError, EngramResult};
use crate::types::{Edge, Node};
use sled::Db;
use uuid::Uuid;

// ── Key constructors ──────────────────────────────────────────────────────────

pub fn node_key(id: Uuid) -> Vec<u8> {
    format!("nodes:{}", id).into_bytes()
}

pub fn edge_from_key(from: Uuid, to: Uuid) -> Vec<u8> {
    format!("edges:from:{}:{}", from, to).into_bytes()
}

pub fn edge_to_key(to: Uuid, from: Uuid) -> Vec<u8> {
    format!("edges:to:{}:{}", to, from).into_bytes()
}

pub fn vector_key(id: Uuid) -> Vec<u8> {
    format!("vectors:{}", id).into_bytes()
}

pub fn salience_key(id: Uuid) -> Vec<u8> {
    format!("salience:{}", id).into_bytes()
}

// ── Node storage ─────────────────────────────────────────────────────────────

pub fn write_node(db: &Db, node: &Node) -> EngramResult<()> {
    let key = node_key(node.id);
    let val = bincode::serialize(node)?;
    db.insert(key, val)?;

    // Store the embedding separately for fast vector scan
    let vkey = vector_key(node.id);
    let vbytes = floats_to_bytes(&node.embedding);
    db.insert(vkey, vbytes)?;

    // Store salience separately so the decay pass can update it cheaply
    let skey = salience_key(node.id);
    db.insert(skey, f32_to_bytes(node.salience))?;

    Ok(())
}

pub fn read_node(db: &Db, id: Uuid) -> EngramResult<Option<Node>> {
    match db.get(node_key(id))? {
        Some(bytes) => Ok(Some(bincode::deserialize(&bytes)?)),
        None => Ok(None),
    }
}

/// Iterate over every node in the store.
pub fn scan_nodes(db: &Db) -> EngramResult<Vec<Node>> {
    let prefix = b"nodes:";
    let mut nodes = Vec::new();
    for result in db.scan_prefix(prefix) {
        let (_k, v) = result?;
        let node: Node = bincode::deserialize(&v)?;
        nodes.push(node);
    }
    Ok(nodes)
}

// ── Edge storage ─────────────────────────────────────────────────────────────

pub fn write_edge(db: &Db, edge: &Edge) -> EngramResult<()> {
    let bytes = bincode::serialize(edge)?;
    // Forward index: from → to
    db.insert(edge_from_key(edge.from_id, edge.to_id), bytes.clone())?;
    // Reverse index: to → from
    db.insert(edge_to_key(edge.to_id, edge.from_id), bytes)?;
    Ok(())
}

pub fn read_edges_from(db: &Db, from_id: Uuid) -> EngramResult<Vec<Edge>> {
    let prefix = format!("edges:from:{}:", from_id).into_bytes();
    read_edges_with_prefix(db, &prefix)
}

pub fn read_edges_to(db: &Db, to_id: Uuid) -> EngramResult<Vec<Edge>> {
    let prefix = format!("edges:to:{}:", to_id).into_bytes();
    read_edges_with_prefix(db, &prefix)
}

fn read_edges_with_prefix(db: &Db, prefix: &[u8]) -> EngramResult<Vec<Edge>> {
    let mut edges = Vec::new();
    for result in db.scan_prefix(prefix) {
        let (_k, v) = result?;
        let edge: Edge = bincode::deserialize(&v)?;
        edges.push(edge);
    }
    Ok(edges)
}

// ── Vector scan ───────────────────────────────────────────────────────────────

/// Read all stored (uuid, embedding) pairs. Used for flat cosine search.
pub fn scan_vectors(db: &Db) -> EngramResult<Vec<(Uuid, Vec<f32>)>> {
    let prefix = b"vectors:";
    let mut out = Vec::new();
    for result in db.scan_prefix(prefix) {
        let (k, v) = result?;
        // key = "vectors:{uuid}" — slice off the prefix
        let id_str = std::str::from_utf8(&k[prefix.len()..])
            .map_err(|e| EngramError::InvalidParam(e.to_string()))?;
        let id = id_str
            .parse::<Uuid>()
            .map_err(|e| EngramError::InvalidParam(e.to_string()))?;
        let floats = bytes_to_floats(&v);
        out.push((id, floats));
    }
    Ok(out)
}

// ── Salience update ───────────────────────────────────────────────────────────

/// Overwrite the salience entry for a node without rewriting the full node blob.
pub fn write_salience(db: &Db, id: Uuid, salience: f32) -> EngramResult<()> {
    db.insert(salience_key(id), f32_to_bytes(salience))?;
    Ok(())
}

pub fn read_salience(db: &Db, id: Uuid) -> EngramResult<Option<f32>> {
    match db.get(salience_key(id))? {
        Some(b) => Ok(Some(bytes_to_f32(&b))),
        None => Ok(None),
    }
}

/// Count entries matching a key prefix.
pub fn count_prefix(db: &Db, prefix: &[u8]) -> EngramResult<usize> {
    let mut n = 0usize;
    for result in db.scan_prefix(prefix) {
        result?;
        n += 1;
    }
    Ok(n)
}

// ── Byte encoding helpers ─────────────────────────────────────────────────────

fn f32_to_bytes(v: f32) -> Vec<u8> {
    v.to_le_bytes().to_vec()
}

fn bytes_to_f32(b: &[u8]) -> f32 {
    let arr: [u8; 4] = b[..4].try_into().unwrap_or([0u8; 4]);
    f32::from_le_bytes(arr)
}

fn floats_to_bytes(floats: &[f32]) -> Vec<u8> {
    let mut out = Vec::with_capacity(floats.len() * 4);
    for f in floats {
        out.extend_from_slice(&f.to_le_bytes());
    }
    out
}

fn bytes_to_floats(bytes: &[u8]) -> Vec<f32> {
    bytes
        .chunks_exact(4)
        .map(|c| f32::from_le_bytes(c.try_into().unwrap()))
        .collect()
}