use std::collections::BTreeMap;
use crate::connection::{Annotation, Database};
use crate::error::{DbError, Result};
const NO_EDGES: &[EdgeRef] = &[];
#[derive(Debug, Clone, Default)]
pub struct Subgraph {
nodes: BTreeMap<String, NodeData>,
out_adj: BTreeMap<String, Vec<EdgeRef>>,
in_adj: BTreeMap<String, Vec<EdgeRef>>,
pool: Interner,
}
#[derive(Debug, Clone, PartialEq)]
pub struct NodeData {
title: String,
content: Option<String>,
embedding_model: Option<String>,
valid_from: String,
valid_to: String,
}
impl NodeData {
pub fn new(
title: impl Into<String>,
valid_from: impl Into<String>,
valid_to: impl Into<String>,
) -> Self {
Self {
title: title.into(),
content: None,
embedding_model: None,
valid_from: valid_from.into(),
valid_to: valid_to.into(),
}
}
#[must_use]
pub fn with_content(mut self, content: impl Into<String>) -> Self {
self.content = Some(content.into());
self
}
#[must_use]
pub fn with_embedding_model(mut self, model: Option<String>) -> Self {
self.embedding_model = model;
self
}
pub fn title(&self) -> &str {
&self.title
}
pub fn content(&self) -> Option<&str> {
self.content.as_deref()
}
pub fn embedding_model(&self) -> Option<&str> {
self.embedding_model.as_deref()
}
pub fn valid_from(&self) -> &str {
&self.valid_from
}
pub fn valid_to(&self) -> &str {
&self.valid_to
}
}
#[derive(Debug, Clone, Default)]
struct Interner {
strings: Vec<String>,
index: BTreeMap<String, u32>,
bytes: usize,
}
impl Interner {
fn intern(&mut self, s: &str) -> (u32, usize) {
if let Some(&i) = self.index.get(s) {
return (i, 0);
}
let i = u32::try_from(self.strings.len())
.expect("a subgraph cannot hold 2^32 distinct strings within any byte budget");
self.strings.push(s.to_string());
self.index.insert(s.to_string(), i);
let cost = Self::entry_bytes(s);
self.bytes += cost;
(i, cost)
}
fn entry_bytes(s: &str) -> usize {
2 * s.len() + std::mem::size_of::<String>() + std::mem::size_of::<u32>()
}
fn get(&self, i: u32) -> &str {
&self.strings[i as usize]
}
fn estimated_bytes(&self) -> usize {
self.bytes
}
}
#[derive(Clone, Copy, PartialEq)]
pub struct EdgeRef {
node: u32,
edge_type: u32,
weight: f64,
valid_from: u32,
valid_to: u32,
}
impl std::fmt::Debug for EdgeRef {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"EdgeRef(node=#{}, type=#{}, w={}, from=#{}, to=#{})",
self.node, self.edge_type, self.weight, self.valid_from, self.valid_to
)
}
}
impl EdgeRef {
pub fn node<'a>(&self, graph: &'a Subgraph) -> &'a str {
graph.pool.get(self.node)
}
pub fn edge_type<'a>(&self, graph: &'a Subgraph) -> &'a str {
graph.pool.get(self.edge_type)
}
pub fn weight(&self) -> f64 {
self.weight
}
pub fn valid_from<'a>(&self, graph: &'a Subgraph) -> &'a str {
graph.pool.get(self.valid_from)
}
pub fn valid_to<'a>(&self, graph: &'a Subgraph) -> &'a str {
graph.pool.get(self.valid_to)
}
}
impl Subgraph {
pub fn contains_node(&self, id: &str) -> bool {
self.nodes.contains_key(id)
}
pub fn node(&self, id: &str) -> Option<&NodeData> {
self.nodes.get(id)
}
pub fn node_ids(&self) -> impl ExactSizeIterator<Item = &str> + '_ {
self.nodes.keys().map(String::as_str)
}
pub fn node_count(&self) -> usize {
self.nodes.len()
}
pub fn nodes(&self) -> impl ExactSizeIterator<Item = (&str, &NodeData)> + '_ {
self.nodes.iter().map(|(id, d)| (id.as_str(), d))
}
pub fn out_adjacency(&self) -> impl Iterator<Item = (&str, &[EdgeRef])> + '_ {
self.out_adj.iter().map(|(id, e)| (id.as_str(), e.as_slice()))
}
pub fn in_adjacency(&self) -> impl Iterator<Item = (&str, &[EdgeRef])> + '_ {
self.in_adj.iter().map(|(id, e)| (id.as_str(), e.as_slice()))
}
pub fn insert_node(&mut self, id: impl Into<String>, data: NodeData) -> Option<NodeData> {
self.nodes.insert(id.into(), data)
}
pub fn out_edges(&self, node: &str) -> &[EdgeRef] {
self.out_adj.get(node).map_or(NO_EDGES, Vec::as_slice)
}
pub fn in_edges(&self, node: &str) -> &[EdgeRef] {
self.in_adj.get(node).map_or(NO_EDGES, Vec::as_slice)
}
pub fn degree(&self, node: &str) -> usize {
self.out_edges(node).len() + self.in_edges(node).len()
}
pub fn weighted_degree(&self, node: &str) -> f64 {
self.out_edges(node).iter().map(|e| e.weight).sum::<f64>()
+ self.in_edges(node).iter().map(|e| e.weight).sum::<f64>()
}
pub fn total_weight(&self) -> f64 {
self.out_adj
.values()
.flat_map(|edges| edges.iter().map(|e| e.weight))
.sum()
}
pub fn edge_count(&self) -> usize {
self.out_adj.values().map(Vec::len).sum()
}
fn drop_dangling_adjacency(&mut self) {
let Subgraph {
nodes,
out_adj,
in_adj,
pool,
} = self;
for adj in [out_adj, in_adj] {
adj.retain(|id, edges| {
if !nodes.contains_key(id) {
return false;
}
edges.retain(|e| nodes.contains_key(pool.get(e.node)));
!edges.is_empty()
});
}
}
pub fn is_closed(&self) -> bool {
self.out_adj
.iter()
.chain(self.in_adj.iter())
.all(|(id, edges)| {
self.nodes.contains_key(id)
&& edges
.iter()
.all(|e| self.nodes.contains_key(self.pool.get(e.node)))
})
}
pub fn add_edge(
&mut self,
source: &str,
target: &str,
edge_type: &str,
weight: f64,
valid_from: &str,
valid_to: &str,
) -> usize {
let (src, b1) = self.pool.intern(source);
let (tgt, b2) = self.pool.intern(target);
let (ty, b3) = self.pool.intern(edge_type);
let (from, b4) = self.pool.intern(valid_from);
let (to, b5) = self.pool.intern(valid_to);
let pooled = b1 + b2 + b3 + b4 + b5;
self.out_adj.entry(source.to_string()).or_default().push(EdgeRef {
node: tgt,
edge_type: ty,
weight,
valid_from: from,
valid_to: to,
});
self.in_adj.entry(target.to_string()).or_default().push(EdgeRef {
node: src,
edge_type: ty,
weight,
valid_from: from,
valid_to: to,
});
2 * std::mem::size_of::<EdgeRef>() + pooled
}
fn node_bytes(id: &str, d: &NodeData) -> usize {
id.len()
+ d.title.len()
+ d.content.as_ref().map_or(0, String::len)
+ d.embedding_model.as_ref().map_or(0, String::len)
+ d.valid_from.len()
+ d.valid_to.len()
+ std::mem::size_of::<NodeData>()
}
fn edge_bytes(_e: &EdgeRef) -> usize {
std::mem::size_of::<EdgeRef>()
}
pub fn estimated_bytes(&self) -> usize {
let nodes: usize = self
.nodes
.iter()
.map(|(id, d)| Self::node_bytes(id, d))
.sum();
let edges: usize = self
.out_adj
.values()
.chain(self.in_adj.values())
.flat_map(|v| v.iter())
.map(Self::edge_bytes)
.sum();
nodes + edges + self.pool.estimated_bytes()
}
pub async fn write_back_annotations(
&self,
db: &Database,
label: &str,
values: &BTreeMap<String, String>,
) -> Result<usize> {
let rows: Vec<Annotation> = self
.nodes
.keys()
.filter_map(|id| {
values
.get(id)
.map(|value| Annotation::new(id.clone(), label, value.clone()))
})
.collect();
db.write_analytics_annotations(rows).await
}
}
impl Database {
pub async fn load_subgraph(
&self,
start_node: &str,
max_hops: u32,
now_ts: &str,
byte_budget: usize,
) -> Result<Subgraph> {
self.load_subgraph_with(
&super::TraversalBuilder::new(start_node)
.max_depth(max_hops as usize)
.min_weight(f64::NEG_INFINITY),
now_ts,
byte_budget,
)
.await
}
pub async fn load_subgraph_with(
&self,
traversal: &super::TraversalBuilder,
now_ts: &str,
byte_budget: usize,
) -> Result<Subgraph> {
let start_node = traversal.start_node.as_str();
let max_hops = traversal.max_depth as u32;
let conn = self.read_conn();
let mut graph = Subgraph::default();
let mut bytes = 0usize;
let edge_filter = traversal.edge_filter_sql();
let sql = format!(
"{}{}",
traversal.walk_cte(),
format_args!(
r#"
-- **The `DISTINCT` is why this query is superlinear, and it is not removable.**
--
-- Wave 3 measured `load_subgraph` at 12.5x for 10x the nodes and could not say
-- why; Wave 4 answered it from the plan. `EXPLAIN` reports
-- `USE TEMP B-TREE FOR DISTINCT`: an O(E log E) sort over the output, and
-- n log n predicts ~13.3x for 10x, against the 12.5x measured. That is the term.
--
-- It is load-bearing: two branches can reach the same node, so a node appears in
-- `walk` at more than one depth and the join would otherwise emit its edges once
-- per depth. Without `DISTINCT` a caller gets duplicate edges.
--
-- **Corrected in 0.6.0 (T0.1), and the correction is not that the analysis was
-- wrong.** Everything above holds, and D-070's two rejected fixes were measured
-- honestly. What was wrong was the fixture: `benches/` seeds a chain of stars,
-- which is a *tree*, and in a tree there is exactly one path to each node — so
-- the term that actually dominated was identically 1 and invisible. D-070
-- concluded the growth was "inherent to producing a deduplicated result", which
-- is true of trees and false of graphs. The real cost was the walk enumerating
-- **paths** rather than nodes; see `walk_cte`. On a 328-edge layered graph at
-- depth 6 that was 299,593 walk rows and 428 ms, against 49 rows and 0.1 ms now.
-- The `DISTINCT` stays, and it is no longer the leading term.
--
-- The filters appear **twice**, and that is the contract (D-073). The walk uses
-- them to bound which nodes are reached; the projection uses them to decide
-- which adjacency lands in the result. Filtering only the walk would hand a
-- caller who asked for `CITES` a graph reached via `CITES` and populated with
-- every other edge type those nodes happen to have.
SELECT DISTINCT l.source_id, l.target_id, l.edge_type, l.weight, l.valid_from, l.valid_to
FROM walk w
JOIN links_current l ON l.source_id = w.node_id
WHERE l.valid_from <= ?3 AND ?3 < l.valid_to
AND l.weight >= ?4
{edge_filter}
ORDER BY l.source_id, l.target_id, l.edge_type
"#
)
);
let mut params: Vec<libsql::Value> = vec![
start_node.into(),
(max_hops as i64).into(),
now_ts.into(),
traversal.min_weight.into(),
];
params.extend(traversal.edge_types.iter().map(|t| t.as_str().into()));
let mut rows = conn.query(&sql, params).await?;
while let Some(row) = rows.next().await? {
let source: String = row.get(0)?;
let target: String = row.get(1)?;
let weight: f64 = row.get(3)?;
if weight < 0.0 || weight.is_nan() {
return Err(DbError::NegativeEdgeWeight {
source_id: source,
target_id: target,
weight,
});
}
let edge_type: String = row.get(2)?;
let valid_from: String = row.get(4)?;
let valid_to: String = row.get(5)?;
bytes += graph.add_edge(&source, &target, &edge_type, weight, &valid_from, &valid_to);
if bytes > byte_budget {
return Err(DbError::SubgraphTooLarge {
n: bytes,
budget: byte_budget,
});
}
}
let mut ids: Vec<String> = graph
.out_adj
.keys()
.chain(graph.in_adj.keys())
.cloned()
.collect();
ids.push(start_node.to_string());
ids.sort();
ids.dedup();
hydrate(conn, &mut graph, &ids, bytes, byte_budget, traversal.content).await?;
graph.drop_dangling_adjacency();
Ok(graph)
}
}
use crate::util::limits::HYDRATE_CHUNK;
async fn hydrate(
conn: &libsql::Connection,
graph: &mut Subgraph,
ids: &[String],
bytes_so_far: usize,
byte_budget: usize,
with_content: bool,
) -> Result<()> {
let mut bytes = bytes_so_far;
for chunk in ids.chunks(HYDRATE_CHUNK) {
let list = (1..=chunk.len())
.map(|i| format!("?{i}"))
.collect::<Vec<_>>()
.join(", ");
let sql = format!(
"SELECT id, title, content, embedding_model, valid_from, valid_to \
FROM concepts WHERE retired = 0 AND id IN ({list})"
);
let params: Vec<libsql::Value> = chunk
.iter()
.map(|id| libsql::Value::Text(id.clone()))
.collect();
let mut rows = conn.query(&sql, params).await?;
while let Some(row) = rows.next().await? {
let id: String = row.get(0)?;
let data = NodeData {
title: row.get(1)?,
content: if with_content { row.get(2).ok() } else { None },
embedding_model: row.get(3).ok(),
valid_from: row.get(4)?,
valid_to: row.get(5)?,
};
bytes += Subgraph::node_bytes(&id, &data);
graph.nodes.insert(id, data);
if bytes > byte_budget {
return Err(DbError::SubgraphTooLarge {
n: bytes,
budget: byte_budget,
});
}
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn adding_an_edge_indexes_it_in_both_directions() {
let mut g = Subgraph::default();
g.add_edge("A", "B", "KNOWS", 0.5, "2026-01-01T00:00:00.000000Z", "9999-12-31T23:59:59.999999Z");
assert_eq!(g.out_edges("A").len(), 1);
assert_eq!(g.out_edges("A")[0].node(&g), "B");
assert_eq!(g.in_edges("B").len(), 1);
assert_eq!(g.in_edges("B")[0].node(&g), "A", "in_adj holds the source");
assert_eq!(g.degree("A") + g.degree("B"), 2);
assert_eq!(g.weighted_degree("A") + g.weighted_degree("B"), 1.0);
assert_eq!(g.total_weight(), 0.5);
}
#[test]
fn a_missing_node_has_no_edges_rather_than_panicking() {
let g = Subgraph::default();
assert!(g.out_edges("nobody").is_empty());
assert!(g.in_edges("nobody").is_empty());
assert_eq!(g.degree("nobody"), 0);
assert_eq!(g.weighted_degree("nobody"), 0.0);
}
}