1use rusqlite::params;
5use serde::{Deserialize, Serialize};
6use uuid::Uuid;
7
8use crate::Grit;
9use crate::clock::TimestampMs;
10use crate::error::Result;
11use crate::model::{
12 self, EDGE_COLS, EPISODE_COLS, Edge, Episode, NODE_COLS, Node, Subgraph, collect,
13};
14
15#[derive(Debug, Clone)]
18pub struct Traversal {
19 pub depth: u32,
21 pub as_of: Option<TimestampMs>,
24 pub as_at: Option<TimestampMs>,
27 pub group_id: Option<String>,
29 pub max_nodes: usize,
34}
35
36impl Default for Traversal {
37 fn default() -> Self {
38 Self {
39 depth: 3,
40 as_of: None,
41 as_at: None,
42 group_id: None,
43 max_nodes: 256,
44 }
45 }
46}
47
48impl Traversal {
49 pub fn depth(mut self, depth: u32) -> Self {
51 self.depth = depth;
52 self
53 }
54
55 pub fn as_of(mut self, t: TimestampMs) -> Self {
57 self.as_of = Some(t);
58 self
59 }
60
61 pub fn as_at(mut self, t: TimestampMs) -> Self {
63 self.as_at = Some(t);
64 self
65 }
66
67 pub fn group(mut self, group_id: impl Into<String>) -> Self {
69 self.group_id = Some(group_id.into());
70 self
71 }
72
73 pub fn max_nodes(mut self, n: usize) -> Self {
75 self.max_nodes = n;
76 self
77 }
78}
79
80#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
82pub struct Stats {
83 pub nodes: usize,
85 pub edges: usize,
87 pub episodes: usize,
89 pub mentions: usize,
91 pub oplog: usize,
93 pub purged: usize,
95}
96
97#[derive(Debug, Clone, Serialize, Deserialize)]
100pub struct NodeHistory {
101 pub node: Node,
103 pub edges: Vec<Edge>,
105}
106
107#[derive(Debug, Clone, Serialize, Deserialize)]
110pub struct MergeCandidate {
111 pub node: Node,
113 pub name_score: f64,
115 pub vector_score: Option<f64>,
117 pub score: f64,
119}
120
121#[doc(hidden)]
129pub const EXPAND_SQL_CURRENT: &str = "
130SELECT DISTINCT CASE WHEN e.src = f.value THEN e.dst ELSE e.src END
131FROM json_each(?1) AS f
132JOIN edges AS e ON e.src = f.value OR e.dst = f.value
133WHERE (?2 IS NULL OR e.group_id = ?2)
134 AND e.created_at <= ?3 AND (e.expired_at IS NULL OR e.expired_at > ?3)
135 AND (e.valid_at IS NULL OR e.valid_at <= ?4)
136 AND (e.invalid_at IS NULL OR e.invalid_at > ?4)";
137
138#[doc(hidden)]
143pub const EXPAND_SQL_AS_AT: &str = "
144SELECT DISTINCT CASE WHEN e.src = f.value THEN e.dst ELSE e.src END
145FROM json_each(?1) AS f
146JOIN edges AS e ON e.src = f.value OR e.dst = f.value
147WHERE (?2 IS NULL OR e.group_id = ?2)
148 AND e.created_at <= ?3 AND (e.expired_at IS NULL OR e.expired_at > ?3)
149 AND (e.valid_at IS NULL OR e.valid_at <= ?4)
150 AND NOT EXISTS (SELECT 1 FROM edge_invalidations AS iv
151 WHERE iv.edge_id = e.id
152 AND iv.recorded_at <= ?3 AND iv.invalid_at <= ?4)";
153
154const BELIEVED_NODES_SQL: &str = "
158SELECT id FROM nodes
159WHERE id IN (SELECT value FROM json_each(?1))
160 AND created_at <= ?2 AND (expired_at IS NULL OR expired_at > ?2)
161ORDER BY id";
162
163impl Grit {
164 pub fn stats(&self) -> Result<Stats> {
166 let conn = self.read();
167 let tx = conn.unchecked_transaction()?;
168 let count = |table: &str| -> Result<usize> {
169 let n: i64 =
171 tx.query_row(&format!("SELECT COUNT(*) FROM {table}"), [], |r| r.get(0))?;
172 Ok(n as usize)
173 };
174 Ok(Stats {
175 nodes: count("nodes")?,
176 edges: count("edges")?,
177 episodes: count("episodes")?,
178 mentions: count("mentions")?,
179 oplog: count("oplog")?,
180 purged: count("purged")?,
181 })
182 }
183
184 pub fn node(&self, id: Uuid) -> Result<Option<Node>> {
187 node_on(&self.read(), id)
188 }
189
190 pub fn edge(&self, id: Uuid) -> Result<Option<Edge>> {
192 let conn = self.read();
193 let mut stmt =
194 conn.prepare_cached(&format!("SELECT {EDGE_COLS} FROM edges WHERE id = ?1"))?;
195 let mut rows = collect(stmt.query_map(params![id.to_string()], Edge::from_row)?)?;
196 Ok(rows.pop())
197 }
198
199 pub fn episode(&self, id: Uuid) -> Result<Option<Episode>> {
201 let conn = self.read();
202 let mut stmt = conn.prepare_cached(&format!(
203 "SELECT {EPISODE_COLS} FROM episodes WHERE id = ?1"
204 ))?;
205 let mut rows = collect(stmt.query_map(params![id.to_string()], Episode::from_row)?)?;
206 Ok(rows.pop())
207 }
208
209 pub fn nodes_in_group(&self, group_id: &str) -> Result<Vec<Node>> {
225 let conn = self.read();
226 let mut stmt = conn.prepare_cached(&format!(
227 "SELECT {NODE_COLS} FROM nodes WHERE group_id = ?1 ORDER BY id"
228 ))?;
229 collect(stmt.query_map(params![group_id], Node::from_row)?)
230 }
231
232 pub fn edges_in_group(&self, group_id: &str) -> Result<Vec<Edge>> {
244 let conn = self.read();
245 let mut stmt = conn.prepare_cached(&format!(
246 "SELECT {EDGE_COLS} FROM edges WHERE group_id = ?1 ORDER BY id"
247 ))?;
248 collect(stmt.query_map(params![group_id], Edge::from_row)?)
249 }
250
251 pub fn episodes_in_group(&self, group_id: &str) -> Result<Vec<Episode>> {
264 let conn = self.read();
265 let mut stmt = conn.prepare_cached(&format!(
266 "SELECT {EPISODE_COLS} FROM episodes WHERE group_id = ?1
267 ORDER BY occurred_at, id"
268 ))?;
269 collect(stmt.query_map(params![group_id], Episode::from_row)?)
270 }
271
272 pub fn get_node_embedding(&self, id: Uuid) -> Result<Option<Vec<f32>>> {
286 self.get_embedding("vec_nodes", id)
287 }
288
289 pub fn get_edge_embedding(&self, id: Uuid) -> Result<Option<Vec<f32>>> {
292 self.get_embedding("vec_edges", id)
293 }
294
295 fn get_embedding(&self, table: &str, id: Uuid) -> Result<Option<Vec<f32>>> {
296 let conn = self.read();
297 let has_table: i64 = conn.query_row(
299 "SELECT COUNT(*) FROM sqlite_master WHERE name = ?1",
300 params![table],
301 |r| r.get(0),
302 )?;
303 if has_table == 0 {
304 return Ok(None);
305 }
306 let mut stmt =
308 conn.prepare_cached(&format!("SELECT embedding FROM {table} WHERE id = ?1"))?;
309 let mut rows =
310 collect(stmt.query_map(params![id.to_string()], |r| r.get::<_, Vec<u8>>(0))?)?;
311 Ok(rows.pop().map(|bytes| crate::vecext::bytes_as_f32s(&bytes)))
312 }
313
314 pub fn mentions_of(&self, target: Uuid) -> Result<Vec<Uuid>> {
318 let conn = self.read();
319 let tx = conn.unchecked_transaction()?;
320 episodes_mentioning(&tx, &target.to_string())
321 }
322
323 pub fn traverse(&self, from: &[Uuid], spec: &Traversal) -> Result<Subgraph> {
337 use std::collections::HashSet;
338
339 let now = self.now_ms();
340 let as_of = spec.as_of.unwrap_or(now);
341 let as_at = spec.as_at.unwrap_or(now);
342 let expand_sql = if spec.as_at.is_none_or(|t| t >= now) {
345 EXPAND_SQL_CURRENT
346 } else {
347 EXPAND_SQL_AS_AT
348 };
349
350 let conn = self.read();
351 let tx = conn.unchecked_transaction()?;
355
356 let mut in_visited: HashSet<String> = HashSet::new();
361 let mut seeds: Vec<String> = Vec::with_capacity(from.len());
362 for id in from {
363 let canon = resolve_read(&tx, &id.to_string())?;
364 if in_visited.insert(canon.clone()) {
365 seeds.push(canon);
366 }
367 }
368 let believed: HashSet<String> = believed_nodes(&tx, &seeds, as_at)?.into_iter().collect();
369 let mut visited: Vec<String> = seeds
370 .into_iter()
371 .filter(|s| believed.contains(s))
372 .take(spec.max_nodes)
373 .collect();
374 in_visited = visited.iter().cloned().collect();
375
376 let mut frontier = visited.clone();
381 for _ in 0..spec.depth {
382 if frontier.is_empty() || visited.len() >= spec.max_nodes {
383 break;
384 }
385 let neighbors: Vec<String> = {
386 let mut stmt = tx.prepare_cached(expand_sql)?;
387 let rows = stmt.query_map(
388 params![
389 serde_json::to_string(&frontier)?,
390 spec.group_id,
391 as_at,
392 as_of
393 ],
394 |r| r.get::<_, String>(0),
395 )?;
396 collect(rows)?
397 };
398 let fresh: Vec<String> = neighbors
399 .into_iter()
400 .filter(|id| !in_visited.contains(id))
401 .collect();
402 let fresh: Vec<String> = believed_nodes(&tx, &fresh, as_at)?
403 .into_iter()
404 .filter(|id| !in_visited.contains(id))
405 .take(spec.max_nodes - visited.len())
406 .collect();
407 for id in &fresh {
408 in_visited.insert(id.clone());
409 }
410 visited.extend(fresh.iter().cloned());
411 frontier = fresh;
412 }
413 let visited_json = serde_json::to_string(&visited)?;
414
415 let nodes = {
416 let mut stmt = tx.prepare_cached(&format!(
417 "SELECT {NODE_COLS} FROM nodes
418 WHERE id IN (SELECT value FROM json_each(?1))
419 ORDER BY id"
420 ))?;
421 collect(stmt.query_map(params![visited_json], Node::from_row)?)?
422 };
423 let invalid_clause = if spec.as_at.is_none_or(|t| t >= now) {
427 "(e.invalid_at IS NULL OR e.invalid_at > ?4)"
428 } else {
429 "NOT EXISTS (SELECT 1 FROM edge_invalidations AS iv
430 WHERE iv.edge_id = e.id
431 AND iv.recorded_at <= ?3 AND iv.invalid_at <= ?4)"
432 };
433 let edges = {
434 let mut stmt = tx.prepare_cached(&format!(
435 "SELECT {EDGE_COLS} FROM edges AS e
436 WHERE e.src IN (SELECT value FROM json_each(?1))
437 AND e.dst IN (SELECT value FROM json_each(?1))
438 AND (?2 IS NULL OR e.group_id = ?2)
439 AND e.created_at <= ?3 AND (e.expired_at IS NULL OR e.expired_at > ?3)
440 AND (e.valid_at IS NULL OR e.valid_at <= ?4)
441 AND {invalid_clause}
442 ORDER BY e.id"
443 ))?;
444 collect(stmt.query_map(
445 params![visited_json, spec.group_id, as_at, as_of],
446 Edge::from_row,
447 )?)?
448 };
449 Ok(Subgraph { nodes, edges })
450 }
451
452 pub fn node_history(&self, id: Uuid) -> Result<NodeHistory> {
455 let conn = self.read();
456 let tx = conn.unchecked_transaction()?;
457 let node = node_on(&tx, id)?.ok_or(crate::Error::NotFound(id))?;
458 let mut stmt = tx.prepare_cached(&format!(
459 "SELECT {EDGE_COLS} FROM edges WHERE src = ?1 OR dst = ?1
460 ORDER BY created_at, id"
461 ))?;
462 let edges = collect(stmt.query_map(params![id.to_string()], Edge::from_row)?)?;
463 Ok(NodeHistory { node, edges })
464 }
465
466 pub fn find_merge_candidates(&self, id: Uuid, min_score: f64) -> Result<Vec<MergeCandidate>> {
471 let conn = self.read();
472 let tx = conn.unchecked_transaction()?;
473 let target = node_on(&tx, id)?.ok_or(crate::Error::NotFound(id))?;
474 let vector_scores = node_vector_neighbors(&tx, id, &target.group_id)?;
475
476 let others = {
477 let mut stmt = tx.prepare_cached(&format!(
478 "SELECT {NODE_COLS} FROM nodes
479 WHERE group_id = ?1 AND id != ?2 AND expired_at IS NULL"
480 ))?;
481 collect(stmt.query_map(params![target.group_id, id.to_string()], Node::from_row)?)?
482 };
483 drop(tx);
484 drop(conn);
485
486 let target_name = target.name.to_lowercase();
487 let mut out: Vec<MergeCandidate> = others
488 .into_iter()
489 .map(|node| {
490 let name_score = strsim::jaro_winkler(&target_name, &node.name.to_lowercase());
491 let vector_score = vector_scores.get(&node.id).copied();
492 let score = vector_score.map_or(name_score, |v| name_score.max(v));
493 MergeCandidate {
494 node,
495 name_score,
496 vector_score,
497 score,
498 }
499 })
500 .filter(|c| c.score >= min_score)
501 .collect();
502 out.sort_by(|a, b| b.score.total_cmp(&a.score));
503 Ok(out)
504 }
505}
506
507fn node_vector_neighbors(
512 conn: &rusqlite::Connection,
513 id: Uuid,
514 group: &str,
515) -> Result<std::collections::HashMap<Uuid, f64>> {
516 let has_vec: i64 = conn.query_row(
517 "SELECT COUNT(*) FROM sqlite_master WHERE name = 'vec_nodes'",
518 [],
519 |r| r.get(0),
520 )?;
521 let mut out = std::collections::HashMap::new();
522 if has_vec == 0 {
523 return Ok(out);
524 }
525 let embedding: Option<Vec<u8>> = {
526 let mut stmt = conn.prepare_cached("SELECT embedding FROM vec_nodes WHERE id = ?1")?;
527 let mut rows =
528 collect(stmt.query_map(params![id.to_string()], |r| r.get::<_, Vec<u8>>(0))?)?;
529 rows.pop()
530 };
531 let Some(embedding) = embedding else {
532 return Ok(out);
533 };
534 let mut stmt = conn.prepare_cached(
539 "SELECT id, distance FROM vec_nodes
540 WHERE embedding MATCH ?1 AND k = 64 AND group_id = ?2",
541 )?;
542 let rows = stmt.query_map(params![embedding, group], |r| {
543 Ok((r.get::<_, String>(0)?, r.get::<_, f64>(1)?))
544 })?;
545 for row in rows {
546 let (nid, distance) = row?;
547 if let Ok(nid) = Uuid::parse_str(&nid)
548 && nid != id
549 {
550 out.insert(nid, 1.0 - distance);
551 }
552 }
553 Ok(out)
554}
555
556fn node_on(conn: &rusqlite::Connection, id: Uuid) -> Result<Option<Node>> {
558 let mut stmt = conn.prepare_cached(&format!("SELECT {NODE_COLS} FROM nodes WHERE id = ?1"))?;
559 let mut rows = collect(stmt.query_map(params![id.to_string()], Node::from_row)?)?;
560 Ok(rows.pop())
561}
562
563pub(crate) fn ids_json<'a>(ids: impl Iterator<Item = &'a Uuid>) -> String {
564 serde_json::to_string(&ids.map(Uuid::to_string).collect::<Vec<_>>())
565 .expect("Vec<String> serialization cannot fail")
566}
567
568fn believed_nodes(
571 conn: &rusqlite::Connection,
572 ids: &[String],
573 as_at: TimestampMs,
574) -> Result<Vec<String>> {
575 if ids.is_empty() {
576 return Ok(Vec::new());
577 }
578 let mut stmt = conn.prepare_cached(BELIEVED_NODES_SQL)?;
579 let rows = stmt.query_map(params![serde_json::to_string(ids)?, as_at], |r| {
580 r.get::<_, String>(0)
581 })?;
582 collect(rows)
583}
584
585pub(crate) fn episodes_mentioning(
589 conn: &rusqlite::Connection,
590 id: &str,
591) -> Result<Vec<uuid::Uuid>> {
592 let canon = resolve_read(conn, id)?;
593 let mut stmt = conn.prepare_cached(
594 "WITH RECURSIVE aliases (id) AS (
595 SELECT ?1
596 UNION
597 SELECT n.id FROM nodes AS n JOIN aliases AS a ON n.merged_into = a.id
598 UNION
599 SELECT p.id FROM purged AS p JOIN aliases AS a ON p.merged_into = a.id
600 )
601 SELECT DISTINCT episode_id FROM mentions
602 WHERE target_id IN (SELECT id FROM aliases)
603 ORDER BY episode_id",
604 )?;
605 let rows = stmt.query_map(params![canon], |r| model::uuid_col(r, 0))?;
606 collect(rows)
607}
608
609fn resolve_read(conn: &rusqlite::Connection, id: &str) -> Result<String> {
612 use rusqlite::OptionalExtension;
613 let mut chain: Vec<String> = vec![id.to_owned()];
614 let mut seen: std::collections::HashSet<String> = chain.iter().cloned().collect();
615 loop {
616 let current = chain.last().expect("chain never empty");
617 let mut stmt = conn.prepare_cached(
618 "SELECT merged_into FROM nodes WHERE id = ?1
619 UNION ALL
620 SELECT merged_into FROM purged WHERE id = ?1
621 LIMIT 1",
622 )?;
623 let next: Option<Option<String>> =
624 stmt.query_row(params![current], |r| r.get(0)).optional()?;
625 match next {
626 None | Some(None) => return Ok(chain.pop().expect("chain never empty")),
627 Some(Some(next)) => {
628 if seen.contains(&next) {
629 let pos = chain
630 .iter()
631 .position(|x| x == &next)
632 .expect("seen implies present");
633 return Ok(chain[pos..]
634 .iter()
635 .min()
636 .expect("cycle slice non-empty")
637 .clone());
638 }
639 seen.insert(next.clone());
640 chain.push(next);
641 }
642 }
643 }
644}