grit_core/

query.rs

//! Read path: point lookups, validity-filtered traversal, bi-temporal history,
//! and merge-candidate scoring.

use rusqlite::params;
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use crate::Grit;
use crate::clock::TimestampMs;
use crate::error::Result;
use crate::model::{
    self, EDGE_COLS, EPISODE_COLS, Edge, Episode, NODE_COLS, Node, Subgraph, collect,
};

/// Parameters for [`Grit::traverse`]. `Default` gives depth 3 (the AGENTS.md
/// default bound), both timelines evaluated "now", no group filter.
#[derive(Debug, Clone)]
pub struct Traversal {
    /// Maximum hops from the seed set (inclusive bound; default 3).
    pub depth: u32,
    /// Event-time instant: only walk edges whose `[valid_at, invalid_at)`
    /// interval contains this. `None` = clock now.
    pub as_of: Option<TimestampMs>,
    /// System-time instant ("what did I believe at t?"): only walk edges whose
    /// `[created_at, expired_at)` interval contains this. `None` = clock now.
    pub as_at: Option<TimestampMs>,
    /// Restrict the walk to one namespace.
    pub group_id: Option<String>,
    /// Node budget: the walk halts (breadth-first, nearest wins) once this
    /// many nodes are reached. Keeps traversal latency independent of
    /// neighborhood size — Layer 3 asks for "context that fits", never the
    /// 7,000-node ball around a hub. Default 256.
    pub max_nodes: usize,
}

impl Default for Traversal {
    fn default() -> Self {
        Self {
            depth: 3,
            as_of: None,
            as_at: None,
            group_id: None,
            max_nodes: 256,
        }
    }
}

impl Traversal {
    /// Set the hop bound.
    pub fn depth(mut self, depth: u32) -> Self {
        self.depth = depth;
        self
    }

    /// Set the event-time instant.
    pub fn as_of(mut self, t: TimestampMs) -> Self {
        self.as_of = Some(t);
        self
    }

    /// Set the system-time (belief) instant — this is the time-travel knob.
    pub fn as_at(mut self, t: TimestampMs) -> Self {
        self.as_at = Some(t);
        self
    }

    /// Restrict to a namespace.
    pub fn group(mut self, group_id: impl Into<String>) -> Self {
        self.group_id = Some(group_id.into());
        self
    }

    /// Set the node budget (`usize::MAX` for an unbounded walk).
    pub fn max_nodes(mut self, n: usize) -> Self {
        self.max_nodes = n;
        self
    }
}

/// Row counts, as reported by [`Grit::stats`].
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct Stats {
    /// Total node rows (including expired/merged-away ones).
    pub nodes: usize,
    /// Total edge rows.
    pub edges: usize,
    /// Total episodes.
    pub episodes: usize,
    /// Total mention links.
    pub mentions: usize,
    /// Total oplog entries.
    pub oplog: usize,
    /// Purged (tombstoned) ids.
    pub purged: usize,
}

/// Bi-temporal audit trail for one node: every incident edge row ever
/// believed, including invalidated and expired ones, oldest belief first.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeHistory {
    /// The node row (even if expired/merged away).
    pub node: Node,
    /// All incident edges, unfiltered, ordered by `created_at`.
    pub edges: Vec<Edge>,
}

/// A scored suggestion from [`Grit::find_merge_candidates`]. Grit never merges
/// on its own — deciding is Layer 2's LLM-judgment call; grit only scores.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MergeCandidate {
    /// The candidate duplicate.
    pub node: Node,
    /// Jaro-Winkler similarity of names, in `[0, 1]`.
    pub name_score: f64,
    /// Cosine similarity of embeddings, in `[-1, 1]`, if both are embedded.
    pub vector_score: Option<f64>,
    /// max(name_score, vector_score) — the ranking key.
    pub score: f64,
}

/// One BFS level: expand every frontier node (bound as a JSON array — no SQL
/// built from user input) across valid edges to the DISTINCT set of neighbor
/// ids. ?1 frontier, ?2 group, ?3 as_at, ?4 as_of.
///
/// This is the current-belief variant: `as_at` is now, so the denormalized
/// `edges.invalid_at` (maintained as MIN over all invalidation records) IS
/// the belief-current value and the per-edge subquery is unnecessary.
#[doc(hidden)]
pub const EXPAND_SQL_CURRENT: &str = "
SELECT DISTINCT CASE WHEN e.src = f.value THEN e.dst ELSE e.src END
FROM json_each(?1) AS f
JOIN edges AS e ON e.src = f.value OR e.dst = f.value
WHERE (?2 IS NULL OR e.group_id = ?2)
  AND e.created_at <= ?3 AND (e.expired_at IS NULL OR e.expired_at > ?3)
  AND (e.valid_at IS NULL OR e.valid_at <= ?4)
  AND (e.invalid_at IS NULL OR e.invalid_at > ?4)";

/// Time-travel variant of [`EXPAND_SQL_CURRENT`]: a historical `as_at` must
/// ignore invalidations recorded after it, so validity consults the
/// belief-versioned `edge_invalidations` records instead of the folded
/// `edges.invalid_at`.
#[doc(hidden)]
pub const EXPAND_SQL_AS_AT: &str = "
SELECT DISTINCT CASE WHEN e.src = f.value THEN e.dst ELSE e.src END
FROM json_each(?1) AS f
JOIN edges AS e ON e.src = f.value OR e.dst = f.value
WHERE (?2 IS NULL OR e.group_id = ?2)
  AND e.created_at <= ?3 AND (e.expired_at IS NULL OR e.expired_at > ?3)
  AND (e.valid_at IS NULL OR e.valid_at <= ?4)
  AND NOT EXISTS (SELECT 1 FROM edge_invalidations AS iv
                  WHERE iv.edge_id = e.id
                    AND iv.recorded_at <= ?3 AND iv.invalid_at <= ?4)";

/// Keep only ids that are believed nodes at ?2 (purged and merged-away nodes
/// fail this — the walk never passes through ghosts). ORDER BY id makes the
/// budget cut deterministic.
const BELIEVED_NODES_SQL: &str = "
SELECT id FROM nodes
WHERE id IN (SELECT value FROM json_each(?1))
  AND created_at <= ?2 AND (expired_at IS NULL OR expired_at > ?2)
ORDER BY id";

impl Grit {
    /// Row counts across the whole file, from one consistent snapshot.
    pub fn stats(&self) -> Result<Stats> {
        let conn = self.read();
        let tx = conn.unchecked_transaction()?;
        let count = |table: &str| -> Result<usize> {
            // Table names are compile-time constants below, never user input.
            let n: i64 =
                tx.query_row(&format!("SELECT COUNT(*) FROM {table}"), [], |r| r.get(0))?;
            Ok(n as usize)
        };
        Ok(Stats {
            nodes: count("nodes")?,
            edges: count("edges")?,
            episodes: count("episodes")?,
            mentions: count("mentions")?,
            oplog: count("oplog")?,
            purged: count("purged")?,
        })
    }

    /// Fetch one node by id (regardless of validity — point lookups see
    /// everything; filtering is for traversal/search).
    pub fn node(&self, id: Uuid) -> Result<Option<Node>> {
        node_on(&self.read(), id)
    }

    /// Fetch one edge by id.
    pub fn edge(&self, id: Uuid) -> Result<Option<Edge>> {
        let conn = self.read();
        let mut stmt =
            conn.prepare_cached(&format!("SELECT {EDGE_COLS} FROM edges WHERE id = ?1"))?;
        let mut rows = collect(stmt.query_map(params![id.to_string()], Edge::from_row)?)?;
        Ok(rows.pop())
    }

    /// Fetch one episode by id.
    pub fn episode(&self, id: Uuid) -> Result<Option<Episode>> {
        let conn = self.read();
        let mut stmt = conn.prepare_cached(&format!(
            "SELECT {EPISODE_COLS} FROM episodes WHERE id = ?1"
        ))?;
        let mut rows = collect(stmt.query_map(params![id.to_string()], Episode::from_row)?)?;
        Ok(rows.pop())
    }

    /// Episode ids that mention the given node or edge (provenance lookup).
    /// Mention rows keep their original target ids, so this resolves `target`
    /// to its canonical node and gathers mentions across all merged aliases.
    pub fn mentions_of(&self, target: Uuid) -> Result<Vec<Uuid>> {
        let conn = self.read();
        let tx = conn.unchecked_transaction()?;
        episodes_mentioning(&tx, &target.to_string())
    }

    /// Expand from `from` along edges valid at the requested instant, up to
    /// `spec.depth` hops, and return the induced subgraph. Time-travel is the
    /// public API here (Design Invariant 4): pass `as_at` for "what did I
    /// believe at t?", `as_of` for "what was true at t?".
    ///
    /// # Example
    /// ```no_run
    /// # use grit_core::{Grit, Options, Traversal};
    /// # let g = Grit::open("memory.db", Options::new("laptop"))?;
    /// # let (node_id, march) = (g.new_id(), 0);
    /// let believed_in_march = g.traverse(&[node_id], &Traversal::default().as_at(march))?;
    /// # Ok::<(), grit_core::Error>(())
    /// ```
    pub fn traverse(&self, from: &[Uuid], spec: &Traversal) -> Result<Subgraph> {
        use std::collections::HashSet;

        let now = self.now_ms();
        let as_of = spec.as_of.unwrap_or(now);
        let as_at = spec.as_at.unwrap_or(now);
        // A historical as_at must ignore invalidations recorded after it;
        // otherwise the folded edges.invalid_at is exact and much cheaper.
        let expand_sql = if spec.as_at.is_none_or(|t| t >= now) {
            EXPAND_SQL_CURRENT
        } else {
            EXPAND_SQL_AS_AT
        };

        let conn = self.read();
        // One read transaction for the whole walk: every BFS level and the
        // final fetches see the same committed snapshot — a write landing
        // mid-traversal cannot produce a subgraph no database state ever held.
        let tx = conn.unchecked_transaction()?;

        // Stale handles are fine: seeds follow merged_into to their canonical
        // node before the walk starts, then must themselves be believed at
        // as_at (purged / merged-away nodes are dropped, so the walk never
        // starts inside a ghost).
        let mut in_visited: HashSet<String> = HashSet::new();
        let mut seeds: Vec<String> = Vec::with_capacity(from.len());
        for id in from {
            let canon = resolve_read(&tx, &id.to_string())?;
            if in_visited.insert(canon.clone()) {
                seeds.push(canon);
            }
        }
        let believed: HashSet<String> = believed_nodes(&tx, &seeds, as_at)?.into_iter().collect();
        let mut visited: Vec<String> = seeds
            .into_iter()
            .filter(|s| believed.contains(s))
            .take(spec.max_nodes)
            .collect();
        in_visited = visited.iter().cloned().collect();

        // Level-synchronous BFS: one batched expansion per depth level, each
        // node expanded exactly once, believed-ness checked once per NEW node
        // (not once per edge). The budget cut is deterministic — nearest
        // level first, ordered by id within a level.
        let mut frontier = visited.clone();
        for _ in 0..spec.depth {
            if frontier.is_empty() || visited.len() >= spec.max_nodes {
                break;
            }
            let neighbors: Vec<String> = {
                let mut stmt = tx.prepare_cached(expand_sql)?;
                let rows = stmt.query_map(
                    params![
                        serde_json::to_string(&frontier)?,
                        spec.group_id,
                        as_at,
                        as_of
                    ],
                    |r| r.get::<_, String>(0),
                )?;
                collect(rows)?
            };
            let fresh: Vec<String> = neighbors
                .into_iter()
                .filter(|id| !in_visited.contains(id))
                .collect();
            let fresh: Vec<String> = believed_nodes(&tx, &fresh, as_at)?
                .into_iter()
                .filter(|id| !in_visited.contains(id))
                .take(spec.max_nodes - visited.len())
                .collect();
            for id in &fresh {
                in_visited.insert(id.clone());
            }
            visited.extend(fresh.iter().cloned());
            frontier = fresh;
        }
        let visited_json = serde_json::to_string(&visited)?;

        let nodes = {
            let mut stmt = tx.prepare_cached(&format!(
                "SELECT {NODE_COLS} FROM nodes
                 WHERE id IN (SELECT value FROM json_each(?1))
                 ORDER BY id"
            ))?;
            collect(stmt.query_map(params![visited_json], Node::from_row)?)?
        };
        // Induced subgraph: both endpoints in the visited set, which contains
        // only believed nodes by construction — ghost edges (an endpoint
        // purged or merged away) can never qualify.
        let invalid_clause = if spec.as_at.is_none_or(|t| t >= now) {
            "(e.invalid_at IS NULL OR e.invalid_at > ?4)"
        } else {
            "NOT EXISTS (SELECT 1 FROM edge_invalidations AS iv
                         WHERE iv.edge_id = e.id
                           AND iv.recorded_at <= ?3 AND iv.invalid_at <= ?4)"
        };
        let edges = {
            let mut stmt = tx.prepare_cached(&format!(
                "SELECT {EDGE_COLS} FROM edges AS e
                 WHERE e.src IN (SELECT value FROM json_each(?1))
                   AND e.dst IN (SELECT value FROM json_each(?1))
                   AND (?2 IS NULL OR e.group_id = ?2)
                   AND e.created_at <= ?3 AND (e.expired_at IS NULL OR e.expired_at > ?3)
                   AND (e.valid_at IS NULL OR e.valid_at <= ?4)
                   AND {invalid_clause}
                 ORDER BY e.id"
            ))?;
            collect(stmt.query_map(
                params![visited_json, spec.group_id, as_at, as_of],
                Edge::from_row,
            )?)?
        };
        Ok(Subgraph { nodes, edges })
    }

    /// Full bi-temporal audit for a node: the node row plus every incident
    /// edge ever believed, invalidated and expired included — one snapshot.
    pub fn node_history(&self, id: Uuid) -> Result<NodeHistory> {
        let conn = self.read();
        let tx = conn.unchecked_transaction()?;
        let node = node_on(&tx, id)?.ok_or(crate::Error::NotFound(id))?;
        let mut stmt = tx.prepare_cached(&format!(
            "SELECT {EDGE_COLS} FROM edges WHERE src = ?1 OR dst = ?1
             ORDER BY created_at, id"
        ))?;
        let edges = collect(stmt.query_map(params![id.to_string()], Edge::from_row)?)?;
        Ok(NodeHistory { node, edges })
    }

    /// Score possible duplicates of `id` within its group: Jaro-Winkler on
    /// names, cosine over embeddings when present. Returns candidates with
    /// `score >= min_score`, best first. Grit never acts on these — executing
    /// a [`crate::GraphOp::MergeNodes`] is Layer 2's decision.
    pub fn find_merge_candidates(&self, id: Uuid, min_score: f64) -> Result<Vec<MergeCandidate>> {
        let conn = self.read();
        let tx = conn.unchecked_transaction()?;
        let target = node_on(&tx, id)?.ok_or(crate::Error::NotFound(id))?;
        let vector_scores = node_vector_neighbors(&tx, id)?;

        let others = {
            let mut stmt = tx.prepare_cached(&format!(
                "SELECT {NODE_COLS} FROM nodes
                 WHERE group_id = ?1 AND id != ?2 AND expired_at IS NULL"
            ))?;
            collect(stmt.query_map(params![target.group_id, id.to_string()], Node::from_row)?)?
        };
        drop(tx);
        drop(conn);

        let target_name = target.name.to_lowercase();
        let mut out: Vec<MergeCandidate> = others
            .into_iter()
            .map(|node| {
                let name_score = strsim::jaro_winkler(&target_name, &node.name.to_lowercase());
                let vector_score = vector_scores.get(&node.id).copied();
                let score = vector_score.map_or(name_score, |v| name_score.max(v));
                MergeCandidate {
                    node,
                    name_score,
                    vector_score,
                    score,
                }
            })
            .filter(|c| c.score >= min_score)
            .collect();
        out.sort_by(|a, b| b.score.total_cmp(&a.score));
        Ok(out)
    }
}

/// Cosine similarity of every embedded node against `id`'s embedding,
/// empty if `id` has no embedding (or no model is registered).
fn node_vector_neighbors(
    conn: &rusqlite::Connection,
    id: Uuid,
) -> Result<std::collections::HashMap<Uuid, f64>> {
    let has_vec: i64 = conn.query_row(
        "SELECT COUNT(*) FROM sqlite_master WHERE name = 'vec_nodes'",
        [],
        |r| r.get(0),
    )?;
    let mut out = std::collections::HashMap::new();
    if has_vec == 0 {
        return Ok(out);
    }
    let embedding: Option<Vec<u8>> = {
        let mut stmt = conn.prepare_cached("SELECT embedding FROM vec_nodes WHERE id = ?1")?;
        let mut rows =
            collect(stmt.query_map(params![id.to_string()], |r| r.get::<_, Vec<u8>>(0))?)?;
        rows.pop()
    };
    let Some(embedding) = embedding else {
        return Ok(out);
    };
    // distance_metric=cosine ⇒ distance = 1 - cos_sim.
    // KNN queries only allow the MATCH + k constraints; the self-hit is
    // filtered out here instead.
    let mut stmt = conn
        .prepare_cached("SELECT id, distance FROM vec_nodes WHERE embedding MATCH ?1 AND k = 64")?;
    let rows = stmt.query_map(params![embedding], |r| {
        Ok((r.get::<_, String>(0)?, r.get::<_, f64>(1)?))
    })?;
    for row in rows {
        let (nid, distance) = row?;
        if let Ok(nid) = Uuid::parse_str(&nid)
            && nid != id
        {
            out.insert(nid, 1.0 - distance);
        }
    }
    Ok(out)
}

/// Point lookup on an already-held connection/transaction.
fn node_on(conn: &rusqlite::Connection, id: Uuid) -> Result<Option<Node>> {
    let mut stmt = conn.prepare_cached(&format!("SELECT {NODE_COLS} FROM nodes WHERE id = ?1"))?;
    let mut rows = collect(stmt.query_map(params![id.to_string()], Node::from_row)?)?;
    Ok(rows.pop())
}

pub(crate) fn ids_json<'a>(ids: impl Iterator<Item = &'a Uuid>) -> String {
    serde_json::to_string(&ids.map(Uuid::to_string).collect::<Vec<_>>())
        .expect("Vec<String> serialization cannot fail")
}

/// Filter `ids` down to believed nodes at `as_at`, sorted by id (see
/// [`BELIEVED_NODES_SQL`]).
fn believed_nodes(
    conn: &rusqlite::Connection,
    ids: &[String],
    as_at: TimestampMs,
) -> Result<Vec<String>> {
    if ids.is_empty() {
        return Ok(Vec::new());
    }
    let mut stmt = conn.prepare_cached(BELIEVED_NODES_SQL)?;
    let rows = stmt.query_map(params![serde_json::to_string(ids)?, as_at], |r| {
        r.get::<_, String>(0)
    })?;
    collect(rows)
}

/// Provenance across merges: resolve `id` to its canonical node, then collect
/// mentions of every alias that folded into it (mention rows keep original
/// target ids — the writer never re-points them).
pub(crate) fn episodes_mentioning(
    conn: &rusqlite::Connection,
    id: &str,
) -> Result<Vec<uuid::Uuid>> {
    let canon = resolve_read(conn, id)?;
    let mut stmt = conn.prepare_cached(
        "WITH RECURSIVE aliases (id) AS (
             SELECT ?1
             UNION
             SELECT n.id FROM nodes AS n JOIN aliases AS a ON n.merged_into = a.id
             UNION
             SELECT p.id FROM purged AS p JOIN aliases AS a ON p.merged_into = a.id
         )
         SELECT DISTINCT episode_id FROM mentions
         WHERE target_id IN (SELECT id FROM aliases)
         ORDER BY episode_id",
    )?;
    let rows = stmt.query_map(params![canon], |r| model::uuid_col(r, 0))?;
    collect(rows)
}

/// Read-side twin of the writer's canonical resolution: follow `merged_into`
/// to the surviving node, cycle-safe (min id wins, same rule as apply).
fn resolve_read(conn: &rusqlite::Connection, id: &str) -> Result<String> {
    use rusqlite::OptionalExtension;
    let mut chain: Vec<String> = vec![id.to_owned()];
    let mut seen: std::collections::HashSet<String> = chain.iter().cloned().collect();
    loop {
        let current = chain.last().expect("chain never empty");
        let mut stmt = conn.prepare_cached(
            "SELECT merged_into FROM nodes WHERE id = ?1
             UNION ALL
             SELECT merged_into FROM purged WHERE id = ?1
             LIMIT 1",
        )?;
        let next: Option<Option<String>> =
            stmt.query_row(params![current], |r| r.get(0)).optional()?;
        match next {
            None | Some(None) => return Ok(chain.pop().expect("chain never empty")),
            Some(Some(next)) => {
                if seen.contains(&next) {
                    let pos = chain
                        .iter()
                        .position(|x| x == &next)
                        .expect("seen implies present");
                    return Ok(chain[pos..]
                        .iter()
                        .min()
                        .expect("cycle slice non-empty")
                        .clone());
                }
                seen.insert(next.clone());
                chain.push(next);
            }
        }
    }
}