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,
};
#[derive(Debug, Clone)]
pub struct Traversal {
pub depth: u32,
pub as_of: Option<TimestampMs>,
pub as_at: Option<TimestampMs>,
pub group_id: Option<String>,
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 {
pub fn depth(mut self, depth: u32) -> Self {
self.depth = depth;
self
}
pub fn as_of(mut self, t: TimestampMs) -> Self {
self.as_of = Some(t);
self
}
pub fn as_at(mut self, t: TimestampMs) -> Self {
self.as_at = Some(t);
self
}
pub fn group(mut self, group_id: impl Into<String>) -> Self {
self.group_id = Some(group_id.into());
self
}
pub fn max_nodes(mut self, n: usize) -> Self {
self.max_nodes = n;
self
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct Stats {
pub nodes: usize,
pub edges: usize,
pub episodes: usize,
pub mentions: usize,
pub oplog: usize,
pub purged: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeHistory {
pub node: Node,
pub edges: Vec<Edge>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MergeCandidate {
pub node: Node,
pub name_score: f64,
pub vector_score: Option<f64>,
pub score: f64,
}
#[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)";
#[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)";
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 {
pub fn stats(&self) -> Result<Stats> {
let conn = self.read();
let tx = conn.unchecked_transaction()?;
let count = |table: &str| -> Result<usize> {
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")?,
})
}
pub fn node(&self, id: Uuid) -> Result<Option<Node>> {
node_on(&self.read(), 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())
}
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())
}
pub fn nodes_in_group(&self, group_id: &str) -> Result<Vec<Node>> {
let conn = self.read();
let mut stmt = conn.prepare_cached(&format!(
"SELECT {NODE_COLS} FROM nodes WHERE group_id = ?1 ORDER BY id"
))?;
collect(stmt.query_map(params![group_id], Node::from_row)?)
}
pub fn edges_in_group(&self, group_id: &str) -> Result<Vec<Edge>> {
let conn = self.read();
let mut stmt = conn.prepare_cached(&format!(
"SELECT {EDGE_COLS} FROM edges WHERE group_id = ?1 ORDER BY id"
))?;
collect(stmt.query_map(params![group_id], Edge::from_row)?)
}
pub fn episodes_in_group(&self, group_id: &str) -> Result<Vec<Episode>> {
let conn = self.read();
let mut stmt = conn.prepare_cached(&format!(
"SELECT {EPISODE_COLS} FROM episodes WHERE group_id = ?1
ORDER BY occurred_at, id"
))?;
collect(stmt.query_map(params![group_id], Episode::from_row)?)
}
pub fn get_node_embedding(&self, id: Uuid) -> Result<Option<Vec<f32>>> {
self.get_embedding("vec_nodes", id)
}
pub fn get_edge_embedding(&self, id: Uuid) -> Result<Option<Vec<f32>>> {
self.get_embedding("vec_edges", id)
}
fn get_embedding(&self, table: &str, id: Uuid) -> Result<Option<Vec<f32>>> {
let conn = self.read();
let has_table: i64 = conn.query_row(
"SELECT COUNT(*) FROM sqlite_master WHERE name = ?1",
params![table],
|r| r.get(0),
)?;
if has_table == 0 {
return Ok(None);
}
let mut stmt =
conn.prepare_cached(&format!("SELECT embedding FROM {table} WHERE id = ?1"))?;
let mut rows =
collect(stmt.query_map(params![id.to_string()], |r| r.get::<_, Vec<u8>>(0))?)?;
Ok(rows.pop().map(|bytes| crate::vecext::bytes_as_f32s(&bytes)))
}
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())
}
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);
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();
let tx = conn.unchecked_transaction()?;
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();
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)?)?
};
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 })
}
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 })
}
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, &target.group_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)
}
}
fn node_vector_neighbors(
conn: &rusqlite::Connection,
id: Uuid,
group: &str,
) -> 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);
};
let mut stmt = conn.prepare_cached(
"SELECT id, distance FROM vec_nodes
WHERE embedding MATCH ?1 AND k = 64 AND group_id = ?2",
)?;
let rows = stmt.query_map(params![embedding, group], |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)
}
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")
}
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)
}
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)
}
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);
}
}
}
}