use crate::def::{evaluate, is_keymatch_rooted, NodeView, Predicate, RuleDef};
use crate::index::{
candidate_spec, candidate_spec_approx_with_k, ivf_drift_rebuild_threshold, CandidateSpec,
RuleIndex,
};
use core_storage::{ColumnStore, EdgeProps, IdMap, Interner, Topology, Value};
use std::collections::{BTreeMap, BTreeSet};
#[derive(Debug, Clone)]
pub struct EngineEdgeDelta {
pub rule: String,
pub src_key: String,
pub dst_key: String,
pub edge_type: String,
pub etype_sym: u32,
pub src_id: u32,
pub dst_id: u32,
pub fired: bool,
}
#[cfg(test)]
pub use crate::index::{with_ivf_drift_rebuild, with_vector_dim_reject, with_vector_early_exit};
pub struct GraphMut<'a> {
pub ids: &'a IdMap,
pub syms: &'a mut Interner,
pub labels: &'a [u32],
pub props: &'a ColumnStore,
pub topo: &'a mut Topology,
pub edge_props: &'a mut EdgeProps,
}
pub const DEFAULT_MAX_EDGES: u64 = 1_000_000;
type Triple = (u32, u32, u32);
type Touch = (u32, u32, u32, u32);
pub type SideIvfExport = (Vec<Vec<f64>>, BTreeMap<u32, usize>, u64);
pub type RuleIvfExport = (SideIvfExport, SideIvfExport);
#[derive(Debug, Default)]
pub struct RuleEngine {
rules: BTreeMap<String, RuleDef>,
indexes: BTreeMap<String, RuleIndex>,
provenance: BTreeMap<String, BTreeSet<Triple>>,
owned: BTreeSet<Triple>,
by_node: BTreeMap<u32, BTreeSet<Touch>>,
rule_intern: BTreeMap<String, u32>,
intern_rule: Vec<String>,
tripped: BTreeMap<String, bool>,
fires: BTreeMap<String, u64>,
pending_deltas: Vec<EngineEdgeDelta>,
emit_deltas: bool,
rebuild_needed: BTreeSet<String>,
}
fn candidate_spec_for(def: &RuleDef) -> CandidateSpec<'_> {
if def.approximate {
let k = def.max_edges.map(|me| me.max(64)).unwrap_or(64) as usize;
candidate_spec_approx_with_k(&def.predicate, k)
} else {
candidate_spec(&def.predicate)
}
}
fn src_lookup_spec_for(def: &RuleDef) -> CandidateSpec<'_> {
if is_keymatch_rooted(&def.predicate) {
let field =
keymatch_field(&def.predicate).expect("keymatch-rooted predicate has a KeyMatch field");
CandidateSpec::Scalar { field }
} else {
candidate_spec_for(def)
}
}
fn predicate_covers_field(p: &Predicate, field: &str) -> bool {
match p {
Predicate::VectorSimilar { field: f, .. } => f == field,
Predicate::All(parts) | Predicate::Any(parts) => {
parts.iter().any(|q| predicate_covers_field(q, field))
}
_ => false,
}
}
fn keymatch_field(p: &Predicate) -> Option<&str> {
match p {
Predicate::KeyMatch { field } => Some(field),
Predicate::All(parts) => parts.first().and_then(keymatch_field),
Predicate::Any(_) => None,
_ => None,
}
}
fn compute_desired(
def: &RuleDef,
index: &RuleIndex,
n: u32,
on_src_side: bool,
g: &GraphMut<'_>,
) -> BTreeMap<(u32, u32), f64> {
let (my_label, other_label) = if on_src_side {
(&def.src_label, &def.dst_label)
} else {
(&def.dst_label, &def.src_label)
};
let Some(my_sym) = g.syms.get(my_label) else {
return BTreeMap::new();
};
if g.labels.get(n as usize).copied() != Some(my_sym) {
return BTreeMap::new();
}
let other_sym = g.syms.get(other_label);
let n_key = match g.ids.key_of(n) {
Some(k) => k,
None => return BTreeMap::new(),
};
let n_get = |f: &str| g.props.get(n, f).cloned();
let spec = candidate_spec_for(def);
let candidates: BTreeSet<u32> = if on_src_side {
if is_keymatch_rooted(&def.predicate) {
let field = keymatch_field(&def.predicate).expect("ByKey always comes from KeyMatch");
match n_get(field) {
Some(Value::Str(ref target_key)) => match g.ids.get(target_key) {
Some(dst_id) => std::iter::once(dst_id).collect(),
None => BTreeSet::new(),
},
_ => BTreeSet::new(),
}
} else {
index.dst_side.candidates(&spec, &n_get)
}
} else {
let src_spec = src_lookup_spec_for(def);
if is_keymatch_rooted(&def.predicate) {
let key_getter = |_: &str| Some(Value::Str(n_key.to_string()));
index.src_side.candidates(&src_spec, &key_getter)
} else {
index.src_side.candidates(&src_spec, &n_get)
}
};
let n_early_exit_hint: Option<(Vec<f64>, f64, [f64; 8])> = if !def.approximate {
if let Predicate::VectorSimilar { field, .. } = &def.predicate {
if crate::index::vector_early_exit_enabled() {
let n_side = if on_src_side {
&index.src_side
} else {
&index.dst_side
};
if let Some(vn_v) = n_get(field) {
if let Some(vn) = crate::index::as_numeric_list(&vn_v) {
if let Some((norm_n, ckpts_n)) = n_side.fresh_ckpts_for(n, &vn) {
Some((vn, norm_n, *ckpts_n))
} else {
None
}
} else {
None
}
} else {
None
}
} else {
None
}
} else {
None
}
} else {
None
};
let mut out = BTreeMap::new();
for m in candidates {
if m == n {
continue; }
if g.labels.get(m as usize).copied() != other_sym {
continue; }
let m_key = match g.ids.key_of(m) {
Some(k) => k,
None => continue,
};
let m_get = |f: &str| g.props.get(m, f).cloned();
let (s_view, d_view, s_id, d_id) = if on_src_side {
(
NodeView {
key: n_key,
props: &n_get,
},
NodeView {
key: m_key,
props: &m_get,
},
n,
m,
)
} else {
(
NodeView {
key: m_key,
props: &m_get,
},
NodeView {
key: n_key,
props: &n_get,
},
m,
n,
)
};
if let (Some((ref vn, norm_n, ckpts_n)), Predicate::VectorSimilar { field, min }) =
(&n_early_exit_hint, &def.predicate)
{
let m_side = if on_src_side {
&index.dst_side
} else {
&index.src_side
};
if let Some(vm_v) = m_get(field) {
if let Some(vm) = crate::index::as_numeric_list(&vm_v) {
if let Some((norm_m, ckpts_m)) = m_side.fresh_ckpts_for(m, &vm) {
let (va, ckpts_a, na, vb, ckpts_b, nb) = if on_src_side {
(
vn.as_slice(),
ckpts_n,
*norm_n,
vm.as_slice(),
ckpts_m,
norm_m,
)
} else {
(
vm.as_slice(),
ckpts_m,
norm_m,
vn.as_slice(),
ckpts_n,
*norm_n,
)
};
match crate::def::cosine_early_exit(va, vb, ckpts_a, ckpts_b, na, nb, *min)
{
None => continue, Some(score) => {
out.insert((s_id, d_id), score);
continue; }
}
}
}
}
}
if let Some(score) = evaluate(&def.predicate, &s_view, &d_view) {
out.insert((s_id, d_id), score);
}
}
out
}
fn compute_desired_via(
def: &RuleDef,
anchor: ViaAnchor,
g: &GraphMut<'_>,
) -> BTreeMap<(u32, u32), f64> {
let via_label = def.via_label.as_deref().unwrap();
let via_edge_str = def.via_edge.as_deref().unwrap();
let via_dir = def.via_dir.unwrap_or(core_storage::Direction::Out);
let src_sym = match g.syms.get(&def.src_label) {
Some(s) => s,
None => return BTreeMap::new(),
};
let via_sym = match g.syms.get(via_label) {
Some(s) => s,
None => return BTreeMap::new(),
};
let dst_sym = match g.syms.get(&def.dst_label) {
Some(s) => s,
None => return BTreeMap::new(),
};
let via_etype = match g.syms.get(via_edge_str) {
Some(e) => e,
None => return BTreeMap::new(),
};
let srcs: Vec<u32> = match anchor {
ViaAnchor::Src(src_id) => {
if g.labels.get(src_id as usize).copied() == Some(src_sym) {
vec![src_id]
} else {
return BTreeMap::new();
}
}
ViaAnchor::Dst(_) => {
(0..g.ids.len() as u32)
.filter(|&id| {
matches!(
g.labels.get(id as usize).copied(),
Some(s) if s != u32::MAX && s == src_sym
)
})
.collect()
}
};
let anchored_dst: Option<u32> = match anchor {
ViaAnchor::Dst(dst_id) => {
if g.labels.get(dst_id as usize).copied() == Some(dst_sym) {
Some(dst_id)
} else {
return BTreeMap::new();
}
}
_ => None,
};
let mut out = BTreeMap::new();
for src in srcs {
let _src_key = match g.ids.key_of(src) {
Some(k) => k,
None => continue,
};
let via_neighbors: Vec<u32> = g
.topo
.neighbors(via_etype, via_dir, src)
.iter()
.copied()
.filter(|&v| g.labels.get(v as usize).copied() == Some(via_sym))
.collect();
if via_neighbors.is_empty() {
continue;
}
let dsts: Vec<u32> = if let Some(dst_id) = anchored_dst {
vec![dst_id]
} else {
(0..g.ids.len() as u32)
.filter(|&id| {
id != src
&& matches!(
g.labels.get(id as usize).copied(),
Some(s) if s != u32::MAX && s == dst_sym
)
})
.collect()
};
for dst in dsts {
if dst == src {
continue; }
let dst_key = match g.ids.key_of(dst) {
Some(k) => k,
None => continue,
};
let dst_get = |f: &str| g.props.get(dst, f).cloned();
let dst_view = NodeView {
key: dst_key,
props: &dst_get,
};
let mut best: Option<f64> = None;
for &via_id in &via_neighbors {
let via_key = match g.ids.key_of(via_id) {
Some(k) => k,
None => continue,
};
let via_get = |f: &str| g.props.get(via_id, f).cloned();
let via_view = NodeView {
key: via_key,
props: &via_get,
};
if let Some(score) = evaluate(&def.predicate, &via_view, &dst_view) {
best = Some(match best {
None => score,
Some(prev) => prev.max(score),
});
}
}
if let Some(score) = best {
out.insert((src, dst), score);
}
}
}
out
}
enum ViaAnchor {
Src(u32),
Dst(u32),
}
fn edge_budget(def: &RuleDef) -> u64 {
def.max_edges.unwrap_or(DEFAULT_MAX_EDGES)
}
pub(crate) fn filter_src_top_k(
per_src: BTreeMap<(u32, u32), f64>,
k: u64,
ids: &core_storage::IdMap,
) -> BTreeMap<(u32, u32), f64> {
if per_src.len() as u64 <= k {
return per_src;
}
let mut candidates: Vec<((u32, u32), f64)> = per_src.into_iter().collect();
candidates.sort_by(|&((_, da), sa), &((_, db), sb)| {
sb.total_cmp(&sa).then_with(|| {
let ka = ids.key_of(da).unwrap_or("");
let kb = ids.key_of(db).unwrap_or("");
ka.cmp(kb)
})
});
candidates.truncate(k as usize);
candidates.into_iter().collect()
}
fn apply_per_src_top_k(
def: &RuleDef,
src: u32,
desired_from_src: BTreeMap<(u32, u32), f64>,
prov: &mut ProvSets<'_>,
g: &mut GraphMut<'_>,
) {
let et = g.syms.intern(&def.edge_type);
let current: Vec<Triple> = {
let rid = prov.rule_intern.get(&def.name).copied();
prov.by_node
.get(&src)
.into_iter()
.flatten()
.filter(|(r, t, s, _d)| Some(*r) == rid && *t == et && *s == src)
.map(|(_, t, s, d)| (*t, *s, *d))
.collect()
};
for (t, s, d) in current {
if !desired_from_src.contains_key(&(s, d)) {
g.topo.remove_edge(t, s, d);
g.edge_props.remove_edge(t, s, d);
prov.remove(&def.name, (t, s, d), g.ids, g.syms);
}
}
for ((s, d), score) in &desired_from_src {
let triple = (et, *s, *d);
let already = prov.contains(&triple);
if !already {
let newly = g.topo.add_edge(et, *s, *d);
if newly {
prov.insert(&def.name, triple, g.ids, g.syms);
}
}
let is_owned = already || prov.contains(&triple);
if is_owned {
if let Some(p) = &def.weight_prop {
g.edge_props.set(et, *s, *d, p, Value::Float(*score));
}
}
}
}
fn intern_rule(intern: &mut BTreeMap<String, u32>, names: &mut Vec<String>, rule: &str) -> u32 {
if let Some(&id) = intern.get(rule) {
return id;
}
let id = names.len() as u32;
intern.insert(rule.to_string(), id);
names.push(rule.to_string());
id
}
type ByNodeRebuild = (
BTreeMap<u32, BTreeSet<Touch>>,
BTreeMap<String, u32>,
Vec<String>,
);
fn rebuild_by_node(provenance: &BTreeMap<String, BTreeSet<Triple>>) -> ByNodeRebuild {
let mut by_node = BTreeMap::new();
let mut intern = BTreeMap::new();
let mut names = Vec::new();
for (rule, set) in provenance {
let rid = intern_rule(&mut intern, &mut names, rule);
for &triple in set {
touch_insert(&mut by_node, rid, triple);
}
}
(by_node, intern, names)
}
fn touch_insert(by_node: &mut BTreeMap<u32, BTreeSet<Touch>>, rid: u32, triple: Triple) {
let (t, s, d) = triple;
let entry = (rid, t, s, d);
by_node.entry(s).or_default().insert(entry);
if s != d {
by_node.entry(d).or_default().insert(entry);
}
}
fn touch_remove(by_node: &mut BTreeMap<u32, BTreeSet<Touch>>, rid: u32, triple: Triple) {
let (t, s, d) = triple;
let entry = (rid, t, s, d);
if let Some(set) = by_node.get_mut(&s) {
set.remove(&entry);
if set.is_empty() {
by_node.remove(&s);
}
}
if s != d {
if let Some(set) = by_node.get_mut(&d) {
set.remove(&entry);
if set.is_empty() {
by_node.remove(&d);
}
}
}
}
#[cfg(test)]
fn resolve_by_node(
by_node: &BTreeMap<u32, BTreeSet<Touch>>,
names: &[String],
) -> BTreeMap<u32, BTreeSet<(String, Triple)>> {
by_node
.iter()
.map(|(&n, set)| {
let resolved = set
.iter()
.map(|&(rid, t, s, d)| (names[rid as usize].clone(), (t, s, d)))
.collect();
(n, resolved)
})
.collect()
}
struct ProvSets<'a> {
set: &'a mut BTreeSet<Triple>,
owned: &'a mut BTreeSet<Triple>,
by_node: &'a mut BTreeMap<u32, BTreeSet<Touch>>,
rule_intern: &'a mut BTreeMap<String, u32>,
intern_rule: &'a mut Vec<String>,
deltas: &'a mut Vec<EngineEdgeDelta>,
emit: bool,
}
impl ProvSets<'_> {
fn insert(&mut self, rule: &str, triple: Triple, ids: &IdMap, syms: &Interner) -> bool {
if !self.set.insert(triple) {
return false;
}
self.owned.insert(triple);
let rid = intern_rule(self.rule_intern, self.intern_rule, rule);
touch_insert(self.by_node, rid, triple);
let (etype, src, dst) = triple;
if self.emit {
if let (Some(sk), Some(dk), Some(et)) =
(ids.key_of(src), ids.key_of(dst), syms.resolve(etype))
{
self.deltas.push(EngineEdgeDelta {
rule: rule.to_string(),
src_key: sk.to_string(),
dst_key: dk.to_string(),
edge_type: et.to_string(),
etype_sym: etype,
src_id: src,
dst_id: dst,
fired: true,
});
}
}
true
}
fn remove(&mut self, rule: &str, triple: Triple, ids: &IdMap, syms: &Interner) -> bool {
if !self.set.remove(&triple) {
return false;
}
self.owned.remove(&triple);
let rid = intern_rule(self.rule_intern, self.intern_rule, rule);
touch_remove(self.by_node, rid, triple);
let (etype, src, dst) = triple;
if self.emit {
if let (Some(sk), Some(dk), Some(et)) =
(ids.key_of(src), ids.key_of(dst), syms.resolve(etype))
{
self.deltas.push(EngineEdgeDelta {
rule: rule.to_string(),
src_key: sk.to_string(),
dst_key: dk.to_string(),
edge_type: et.to_string(),
etype_sym: etype,
src_id: src,
dst_id: dst,
fired: false,
});
}
}
true
}
fn contains(&self, triple: &Triple) -> bool {
self.set.contains(triple)
}
fn len(&self) -> usize {
self.set.len()
}
}
fn apply_desired(
def: &RuleDef,
desired: BTreeMap<(u32, u32), f64>,
retract_touching: Option<u32>,
prov: &mut ProvSets<'_>,
tripped: &mut bool,
g: &mut GraphMut<'_>,
) {
let budget = edge_budget(def);
let et = g.syms.intern(&def.edge_type);
let current: Vec<Triple> = match retract_touching {
None => prov
.set
.iter()
.filter(|(t, _, _)| *t == et)
.copied()
.collect(),
Some(n) => {
let rid = prov.rule_intern.get(&def.name).copied();
prov.by_node
.get(&n)
.into_iter()
.flatten()
.filter(|(r, t, _, _)| Some(*r) == rid && *t == et)
.map(|(_, t, s, d)| (*t, *s, *d))
.collect()
}
};
for (t, s, d) in current {
if !desired.contains_key(&(s, d)) {
g.topo.remove_edge(t, s, d);
g.edge_props.remove_edge(t, s, d);
prov.remove(&def.name, (t, s, d), g.ids, g.syms);
}
}
for ((s, d), score) in desired {
let triple = (et, s, d);
let already = prov.contains(&triple);
if !already {
if *tripped || prov.len() as u64 >= budget {
*tripped = true;
continue;
}
let newly = g.topo.add_edge(et, s, d);
if newly {
prov.insert(&def.name, triple, g.ids, g.syms);
}
}
let is_owned_here = already || prov.contains(&triple);
if is_owned_here {
if let Some(p) = &def.weight_prop {
g.edge_props.set(et, s, d, p, Value::Float(score));
}
}
}
}
#[cfg(test)]
#[allow(dead_code)]
fn compute_full_desired(
def: &RuleDef,
index: &RuleIndex,
g: &GraphMut<'_>,
) -> BTreeMap<(u32, u32), f64> {
let mut desired = BTreeMap::new();
let src_sym = g.syms.get(&def.src_label);
for id in 0..g.ids.len() as u32 {
let label_sym = match g.labels.get(id as usize).copied() {
Some(s) if s != u32::MAX => s,
_ => continue,
};
if src_sym == Some(label_sym) {
desired.extend(compute_desired(def, index, id, true, g));
}
}
desired
}
fn pair_still_desired(def: &RuleDef, s: u32, d: u32, g: &GraphMut<'_>) -> bool {
let src_sym = match g.syms.get(&def.src_label) {
Some(sym) => sym,
None => return false,
};
let dst_sym = match g.syms.get(&def.dst_label) {
Some(sym) => sym,
None => return false,
};
if g.labels.get(s as usize).copied() != Some(src_sym) {
return false;
}
if g.labels.get(d as usize).copied() != Some(dst_sym) {
return false;
}
let s_key = match g.ids.key_of(s) {
Some(k) => k,
None => return false,
};
let d_key = match g.ids.key_of(d) {
Some(k) => k,
None => return false,
};
let s_get = |f: &str| g.props.get(s, f).cloned();
let d_get = |f: &str| g.props.get(d, f).cloned();
evaluate(
&def.predicate,
&NodeView {
key: s_key,
props: &s_get,
},
&NodeView {
key: d_key,
props: &d_get,
},
)
.is_some()
}
fn count_desired_up_to(def: &RuleDef, index: &RuleIndex, limit: u64, g: &GraphMut<'_>) -> u64 {
let mut count = 0u64;
let src_sym = g.syms.get(&def.src_label);
for id in 0..g.ids.len() as u32 {
let label_sym = match g.labels.get(id as usize).copied() {
Some(s) if s != u32::MAX => s,
_ => continue,
};
if src_sym != Some(label_sym) {
continue;
}
count += compute_desired(def, index, id, true, g).len() as u64;
if count > limit {
return count;
}
}
count
}
fn apply_streaming_create(
def: &RuleDef,
index: &RuleIndex,
prov: &mut ProvSets<'_>,
tripped: &mut bool,
g: &mut GraphMut<'_>,
) {
let budget = edge_budget(def);
let et = g.syms.intern(&def.edge_type);
let src_sym = g.syms.get(&def.src_label);
'outer: for id in 0..g.ids.len() as u32 {
let label_sym = match g.labels.get(id as usize).copied() {
Some(s) if s != u32::MAX => s,
_ => continue,
};
if src_sym != Some(label_sym) {
continue;
}
let per_src = compute_desired(def, index, id, true, g);
for ((s, d), score) in per_src {
let triple = (et, s, d);
let already = prov.contains(&triple);
if !already {
if *tripped || prov.len() as u64 >= budget {
*tripped = true;
break 'outer;
}
let newly = g.topo.add_edge(et, s, d);
if newly {
prov.insert(&def.name, triple, g.ids, g.syms);
}
}
let is_owned_here = already || prov.contains(&triple);
if is_owned_here {
if let Some(p) = &def.weight_prop {
g.edge_props.set(et, s, d, p, Value::Float(score));
}
}
}
}
}
fn apply_streaming_create_top_k(
def: &RuleDef,
k: u64,
index: &RuleIndex,
prov: &mut ProvSets<'_>,
g: &mut GraphMut<'_>,
) {
let src_sym = g.syms.get(&def.src_label);
for id in 0..g.ids.len() as u32 {
let label_sym = match g.labels.get(id as usize).copied() {
Some(s) if s != u32::MAX => s,
_ => continue,
};
if src_sym != Some(label_sym) {
continue;
}
let per_src = compute_desired(def, index, id, true, g);
let top_k = filter_src_top_k(per_src, k, g.ids);
apply_per_src_top_k(def, id, top_k, prov, g);
}
}
fn apply_streaming_rebuild_top_k(
def: &RuleDef,
k: u64,
index: &RuleIndex,
prov: &mut ProvSets<'_>,
g: &mut GraphMut<'_>,
) {
let et = g.syms.intern(&def.edge_type);
let existing_srcs: BTreeSet<u32> = prov
.set
.iter()
.filter(|(t, _, _)| *t == et)
.map(|(_, s, _)| *s)
.collect();
let src_sym = g.syms.get(&def.src_label);
let mut all_srcs: BTreeSet<u32> = existing_srcs;
for id in 0..g.ids.len() as u32 {
let label_sym = match g.labels.get(id as usize).copied() {
Some(s) if s != u32::MAX => s,
_ => continue,
};
if src_sym == Some(label_sym) {
all_srcs.insert(id);
}
}
for src in all_srcs {
let desired_src = compute_desired(def, index, src, true, g);
let top_k = filter_src_top_k(desired_src, k, g.ids);
apply_per_src_top_k(def, src, top_k, prov, g);
}
}
fn apply_streaming_rebuild(
def: &RuleDef,
index: &RuleIndex,
prov: &mut ProvSets<'_>,
tripped: &mut bool,
g: &mut GraphMut<'_>,
) {
let budget = edge_budget(def);
let et = g.syms.intern(&def.edge_type);
let total = count_desired_up_to(def, index, budget, g);
if total > budget {
*tripped = true;
return; }
*tripped = false;
let current: Vec<Triple> = prov
.set
.iter()
.filter(|(t, _, _)| *t == et)
.copied()
.collect();
for (t, s, d) in current {
if !pair_still_desired(def, s, d, g) {
g.topo.remove_edge(t, s, d);
g.edge_props.remove_edge(t, s, d);
prov.remove(&def.name, (t, s, d), g.ids, g.syms);
}
}
let src_sym = g.syms.get(&def.src_label);
for id in 0..g.ids.len() as u32 {
let label_sym = match g.labels.get(id as usize).copied() {
Some(s) if s != u32::MAX => s,
_ => continue,
};
if src_sym != Some(label_sym) {
continue;
}
let per_src = compute_desired(def, index, id, true, g);
for ((s, d), score) in per_src {
let triple = (et, s, d);
let already = prov.contains(&triple);
if !already {
let newly = g.topo.add_edge(et, s, d);
if newly {
prov.insert(&def.name, triple, g.ids, g.syms);
}
}
let is_owned_here = already || prov.contains(&triple);
if is_owned_here {
if let Some(p) = &def.weight_prop {
g.edge_props.set(et, s, d, p, Value::Float(score));
}
}
}
}
}
fn bump_fires_for_participants(def: &RuleDef, g: &GraphMut<'_>, fires: &mut u64) {
let src_sym = g.syms.get(&def.src_label);
let dst_sym = g.syms.get(&def.dst_label);
for id in 0..g.ids.len() as u32 {
let label_sym = match g.labels.get(id as usize).copied() {
Some(s) if s != u32::MAX => s,
_ => continue,
};
if src_sym == Some(label_sym) || dst_sym == Some(label_sym) {
*fires += 1;
}
}
}
fn index_node_for_rule(
id: u32,
label_sym: u32,
def: &RuleDef,
index: &mut RuleIndex,
syms: &Interner,
props: &ColumnStore,
) {
let get = |f: &str| props.get(id, f).cloned();
if syms.get(&def.src_label) == Some(label_sym) {
let spec = src_lookup_spec_for(def);
index.src_side.insert(&spec, id, &get);
}
if syms.get(&def.dst_label) == Some(label_sym) {
let spec = candidate_spec_for(def);
index.dst_side.insert(&spec, id, &get);
}
}
impl RuleEngine {
pub fn new() -> Self {
Self::default()
}
pub fn rules(&self) -> impl Iterator<Item = &RuleDef> {
self.rules.values()
}
pub fn is_owned(&self, etype: u32, src: u32, dst: u32) -> bool {
self.owned.contains(&(etype, src, dst))
}
pub fn provenance(&self) -> &BTreeMap<String, BTreeSet<(u32, u32, u32)>> {
&self.provenance
}
pub fn provenance_touching(
&self,
node: u32,
) -> impl Iterator<Item = (&str, u32, u32, u32)> + '_ {
self.by_node
.get(&node)
.into_iter()
.flatten()
.map(|&(rid, t, s, d)| (self.intern_rule[rid as usize].as_str(), t, s, d))
}
pub fn provenance_touching_len(&self, node: u32) -> usize {
self.by_node.get(&node).map_or(0, BTreeSet::len)
}
pub fn is_tripped(&self, name: &str) -> bool {
self.tripped.get(name).copied().unwrap_or(false)
}
pub fn fire_count(&self, name: &str) -> u64 {
self.fires.get(name).copied().unwrap_or(0)
}
pub fn drain_deltas(&mut self) -> Vec<EngineEdgeDelta> {
std::mem::take(&mut self.pending_deltas)
}
pub fn pending_delta_count(&self) -> usize {
self.pending_deltas.len()
}
pub fn pending_deltas_since(&self, cursor: usize) -> &[EngineEdgeDelta] {
&self.pending_deltas[cursor..]
}
#[allow(clippy::type_complexity)]
pub fn to_persist(
&self,
) -> (
Vec<RuleDef>,
BTreeMap<String, BTreeSet<(u32, u32, u32)>>,
BTreeMap<String, bool>,
BTreeMap<String, u64>,
) {
(
self.rules.values().cloned().collect(),
self.provenance.clone(),
self.tripped.clone(),
self.fires.clone(),
)
}
pub fn from_persist(
rules: Vec<RuleDef>,
prov: BTreeMap<String, BTreeSet<(u32, u32, u32)>>,
tripped: BTreeMap<String, bool>,
fires: BTreeMap<String, u64>,
) -> Self {
let mut owned = BTreeSet::new();
for set in prov.values() {
owned.extend(set.iter().copied());
}
let indexes = rules
.iter()
.map(|r| (r.name.clone(), RuleIndex::default()))
.collect();
let rules: BTreeMap<String, RuleDef> =
rules.into_iter().map(|r| (r.name.clone(), r)).collect();
let mut tripped = tripped;
let mut fires = fires;
for name in rules.keys() {
tripped.entry(name.clone()).or_insert(false);
fires.entry(name.clone()).or_insert(0);
}
let (by_node, rule_intern, intern_rule) = rebuild_by_node(&prov);
Self {
rules,
indexes,
provenance: prov,
owned,
by_node,
rule_intern,
intern_rule,
tripped,
fires,
pending_deltas: Vec::new(),
emit_deltas: false,
rebuild_needed: BTreeSet::new(),
}
}
pub fn set_emit_deltas(&mut self, emit: bool) {
self.emit_deltas = emit;
}
pub fn emit_deltas(&self) -> bool {
self.emit_deltas
}
pub fn take_rebuild_needed(&mut self) -> Vec<String> {
std::mem::take(&mut self.rebuild_needed)
.into_iter()
.collect()
}
pub fn queue_rebuild_needed(&mut self, name: String) {
self.rebuild_needed.insert(name);
}
fn maybe_queue_ivf_rebuild(&mut self, rule_name: &str, def: &RuleDef) {
if !def.approximate {
return;
}
let Some(idx) = self.indexes.get(rule_name) else {
return;
};
if idx.dst_side.ivf_drift > ivf_drift_rebuild_threshold() {
self.rebuild_needed.insert(rule_name.to_string());
}
}
pub fn export_ivf_state(&self) -> BTreeMap<String, RuleIvfExport> {
let mut out = BTreeMap::new();
for (name, def) in &self.rules {
if def.approximate {
if let Some(idx) = self.indexes.get(name) {
out.insert(
name.clone(),
(
idx.src_side.export_ivf_state(),
idx.dst_side.export_ivf_state(),
),
);
}
}
}
out
}
pub fn reindex_all(
&mut self,
ids: &IdMap,
syms: &Interner,
labels: &[u32],
props: &ColumnStore,
) {
for idx in self.indexes.values_mut() {
*idx = RuleIndex::default();
}
let rule_names: Vec<String> = self.rules.keys().cloned().collect();
for name in &rule_names {
if self.rules[name].approximate {
let idx = self.indexes.get_mut(name).unwrap();
idx.src_side.init_hnsw(name);
idx.dst_side.init_hnsw(name);
}
}
for id in 0..ids.len() as u32 {
let label_sym = match labels.get(id as usize).copied() {
Some(s) if s != u32::MAX => s,
_ => continue,
};
for name in &rule_names {
let def = self.rules[name].clone();
let idx = self.indexes.get_mut(name).unwrap();
index_node_for_rule(id, label_sym, &def, idx, syms, props);
}
}
for name in &rule_names {
if self.rules[name].approximate {
let idx = self.indexes.get_mut(name).unwrap();
idx.src_side.fit_ivf_clusters(name);
idx.dst_side.fit_ivf_clusters(name);
}
}
}
pub fn reindex_all_load_ivf(
&mut self,
ids: &IdMap,
syms: &Interner,
labels: &[u32],
props: &ColumnStore,
ivf_state: BTreeMap<String, RuleIvfExport>,
) {
for idx in self.indexes.values_mut() {
*idx = RuleIndex::default();
}
let rule_names: Vec<String> = self.rules.keys().cloned().collect();
for name in &rule_names {
if self.rules[name].approximate {
let idx = self.indexes.get_mut(name).unwrap();
idx.src_side.init_hnsw(name);
idx.dst_side.init_hnsw(name);
}
}
for id in 0..ids.len() as u32 {
let label_sym = match labels.get(id as usize).copied() {
Some(s) if s != u32::MAX => s,
_ => continue,
};
for name in &rule_names {
let def = self.rules[name].clone();
let idx = self.indexes.get_mut(name).unwrap();
index_node_for_rule(id, label_sym, &def, idx, syms, props);
}
}
for name in &rule_names {
if !self.rules[name].approximate {
continue;
}
let idx = self.indexes.get_mut(name).unwrap();
if let Some(((sc, sa, sd), (dc, da, dd))) = ivf_state.get(name) {
idx.src_side.load_ivf_state(sc.clone(), sa.clone(), *sd);
idx.dst_side.load_ivf_state(dc.clone(), da.clone(), *dd);
} else {
idx.src_side.fit_ivf_clusters(name);
idx.dst_side.fit_ivf_clusters(name);
}
}
}
pub fn export_hnsw_state(&self) -> BTreeMap<String, (Vec<u8>, Vec<u8>)> {
let mut out = BTreeMap::new();
for (name, def) in &self.rules {
if def.approximate {
if let Some(idx) = self.indexes.get(name) {
out.insert(
name.clone(),
(
idx.src_side.export_hnsw_blob(),
idx.dst_side.export_hnsw_blob(),
),
);
}
}
}
out
}
pub fn load_hnsw_state(&mut self, blobs: BTreeMap<String, (Vec<u8>, Vec<u8>)>) {
for (name, (src_blob, dst_blob)) in blobs {
if let Some(idx) = self.indexes.get_mut(&name) {
if !src_blob.is_empty() {
idx.src_side.load_hnsw_blob(&src_blob);
}
if !dst_blob.is_empty() {
idx.dst_side.load_hnsw_blob(&dst_blob);
}
}
}
}
pub fn hnsw_search_dst(
&self,
field: &str,
dst_label: &str,
q: &[f64],
k: usize,
) -> Option<Vec<(u32, f64)>> {
for (name, def) in &self.rules {
if !def.approximate || def.dst_label != dst_label {
continue;
}
if !predicate_covers_field(&def.predicate, field) {
continue;
}
if let Some(idx) = self.indexes.get(name) {
if let Some(h) = idx.dst_side.hnsw_ref() {
if !h.is_empty() {
return Some(h.search(q, k));
}
}
}
}
None
}
pub fn create_rule(&mut self, def: RuleDef, g: &mut GraphMut<'_>) -> Result<(), String> {
def.validate()?;
if self.rules.contains_key(&def.name) {
return Err(format!("rule {:?} already exists", def.name));
}
let name = def.name.clone();
self.rules.insert(name.clone(), def);
self.indexes.insert(name.clone(), RuleIndex::default());
self.provenance.entry(name.clone()).or_default();
self.tripped.insert(name.clone(), false);
self.fires.insert(name.clone(), 0);
let n_total = g.ids.len() as u32;
let def = self.rules[&name].clone();
if def.approximate {
let idx = self.indexes.get_mut(&name).unwrap();
idx.src_side.init_hnsw(&name);
idx.dst_side.init_hnsw(&name);
}
for id in 0..n_total {
let label_sym = match g.labels.get(id as usize).copied() {
Some(s) if s != u32::MAX => s,
_ => continue,
};
let idx = self.indexes.get_mut(&name).unwrap();
index_node_for_rule(id, label_sym, &def, idx, g.syms, g.props);
}
if def.approximate {
let idx = self.indexes.get_mut(&name).unwrap();
idx.src_side.fit_ivf_clusters(&name);
idx.dst_side.fit_ivf_clusters(&name);
}
let mut prov = ProvSets {
set: self.provenance.get_mut(&name).unwrap(),
owned: &mut self.owned,
by_node: &mut self.by_node,
rule_intern: &mut self.rule_intern,
intern_rule: &mut self.intern_rule,
deltas: &mut self.pending_deltas,
emit: self.emit_deltas,
};
if def.via_label.is_some() {
let budget = edge_budget(&def);
let et = g.syms.intern(&def.edge_type);
let src_sym = g.syms.get(&def.src_label);
let tripped = self.tripped.get_mut(&name).unwrap();
'via_outer: for id in 0..g.ids.len() as u32 {
let label_sym = match g.labels.get(id as usize).copied() {
Some(s) if s != u32::MAX => s,
_ => continue,
};
if src_sym != Some(label_sym) {
continue;
}
let per_src = compute_desired_via(&def, ViaAnchor::Src(id), g);
if let Some(k) = def.max_edges {
let top_k = filter_src_top_k(per_src, k, g.ids);
apply_per_src_top_k(&def, id, top_k, &mut prov, g);
} else {
for ((s, d), score) in per_src {
let triple = (et, s, d);
let already = prov.contains(&triple);
if !already {
if *tripped || prov.len() as u64 >= budget {
*tripped = true;
break 'via_outer;
}
let newly = g.topo.add_edge(et, s, d);
if newly {
prov.insert(&name, triple, g.ids, g.syms);
}
}
let is_owned_here = already || prov.contains(&triple);
if is_owned_here {
if let Some(p) = &def.weight_prop {
g.edge_props.set(et, s, d, p, Value::Float(score));
}
}
}
}
}
} else if let Some(k) = def.max_edges {
apply_streaming_create_top_k(&def, k, &self.indexes[&name], &mut prov, g);
} else {
let tripped = self.tripped.get_mut(&name).unwrap();
apply_streaming_create(&def, &self.indexes[&name], &mut prov, tripped, g);
}
let fires = self.fires.get_mut(&name).unwrap();
bump_fires_for_participants(&def, g, fires);
Ok(())
}
pub fn delete_rule(&mut self, name: &str, g: &mut GraphMut<'_>) -> Result<(), String> {
if !self.rules.contains_key(name) {
return Err(format!("rule {:?} not found", name));
}
let def = self.rules.remove(name).unwrap();
self.indexes.remove(name);
self.tripped.remove(name);
self.fires.remove(name);
let mut leftover = self.provenance.remove(name).unwrap_or_default();
let _et = g.syms.intern(&def.edge_type);
let triples: Vec<Triple> = leftover.iter().copied().collect();
let mut sets = ProvSets {
set: &mut leftover,
owned: &mut self.owned,
by_node: &mut self.by_node,
rule_intern: &mut self.rule_intern,
intern_rule: &mut self.intern_rule,
deltas: &mut self.pending_deltas,
emit: self.emit_deltas,
};
for triple in triples {
let (t, s, d) = triple;
g.topo.remove_edge(t, s, d);
g.edge_props.remove_edge(t, s, d);
sets.remove(name, triple, g.ids, g.syms);
}
let same_etype_survivors: Vec<String> = self
.rules
.values()
.filter(|r| r.edge_type == def.edge_type)
.map(|r| r.name.clone())
.collect();
for survivor in same_etype_survivors {
let _ = self.rebuild(&survivor, g);
}
Ok(())
}
pub fn on_node_changed(
&mut self,
n: u32,
changed: Option<(&str, Option<Value>)>,
g: &mut GraphMut<'_>,
) {
let n_label = g.labels.get(n as usize).copied();
let rule_names: Vec<String> = self.rules.keys().cloned().collect();
for rule_name in rule_names {
let def = self.rules[&rule_name].clone();
if def.via_label.is_some() {
self.on_node_changed_via(&rule_name, &def, n, n_label, changed.clone(), g);
} else {
let src_sym = g.syms.get(&def.src_label);
let dst_sym = g.syms.get(&def.dst_label);
let as_src = src_sym.is_some() && n_label == src_sym;
let as_dst = dst_sym.is_some() && n_label == dst_sym;
let fires = match changed {
None => as_src || as_dst,
Some((field, _)) => def.watched_fields().contains(field) && (as_src || as_dst),
};
if !fires {
continue;
}
*self.fires.entry(rule_name.clone()).or_default() += 1;
if let Some((field, ref old_val)) = changed {
let old_val_cloned = old_val.clone();
let old_getter = |f: &str| {
if f == field {
old_val_cloned.clone()
} else {
g.props.get(n, f).cloned()
}
};
let idx = self.indexes.get_mut(&rule_name).unwrap();
if as_src {
let spec = src_lookup_spec_for(&def);
idx.src_side.remove(&spec, n, &old_getter);
}
if as_dst {
let spec = candidate_spec_for(&def);
idx.dst_side.remove(&spec, n, &old_getter);
}
}
{
let cur_getter = |f: &str| g.props.get(n, f).cloned();
let idx = self.indexes.get_mut(&rule_name).unwrap();
if as_src {
let spec = src_lookup_spec_for(&def);
idx.src_side.insert(&spec, n, &cur_getter);
}
if as_dst {
let spec = candidate_spec_for(&def);
idx.dst_side.insert(&spec, n, &cur_getter);
}
}
self.maybe_queue_ivf_rebuild(&rule_name, &def);
if let Some(k) = def.max_edges {
let et = g.syms.intern(&def.edge_type);
let affected_srcs_for_n_dst: BTreeSet<u32> = if as_dst {
let rid = self.rule_intern.get(&def.name).copied();
self.by_node
.get(&n)
.into_iter()
.flatten()
.filter(|(r, t, _s, d)| Some(*r) == rid && *t == et && *d == n)
.map(|(_, _, s, _)| *s)
.collect()
} else {
BTreeSet::new()
};
let mut prov = ProvSets {
set: self.provenance.entry(rule_name.clone()).or_default(),
owned: &mut self.owned,
by_node: &mut self.by_node,
rule_intern: &mut self.rule_intern,
intern_rule: &mut self.intern_rule,
deltas: &mut self.pending_deltas,
emit: self.emit_deltas,
};
if as_src {
let desired_n_src =
compute_desired(&def, &self.indexes[&rule_name], n, true, g);
let top_k = filter_src_top_k(desired_n_src, k, g.ids);
apply_per_src_top_k(&def, n, top_k, &mut prov, g);
}
if as_dst {
let new_desired =
compute_desired(&def, &self.indexes[&rule_name], n, false, g);
let new_srcs: BTreeSet<u32> = new_desired.keys().map(|(s, _)| *s).collect();
let affected_srcs: BTreeSet<u32> =
affected_srcs_for_n_dst.union(&new_srcs).copied().collect();
for src in affected_srcs {
if src == n {
continue;
}
let desired_src =
compute_desired(&def, &self.indexes[&rule_name], src, true, g);
let top_k = filter_src_top_k(desired_src, k, g.ids);
apply_per_src_top_k(&def, src, top_k, &mut prov, g);
}
}
} else {
let mut desired = BTreeMap::new();
if as_src {
desired.extend(compute_desired(
&def,
&self.indexes[&rule_name],
n,
true,
g,
));
}
if as_dst {
desired.extend(compute_desired(
&def,
&self.indexes[&rule_name],
n,
false,
g,
));
}
let tripped = self.tripped.entry(rule_name.clone()).or_default();
apply_desired(
&def,
desired,
Some(n),
&mut ProvSets {
set: self.provenance.entry(rule_name).or_default(),
owned: &mut self.owned,
by_node: &mut self.by_node,
rule_intern: &mut self.rule_intern,
intern_rule: &mut self.intern_rule,
deltas: &mut self.pending_deltas,
emit: self.emit_deltas,
},
tripped,
g,
);
}
}
}
}
fn on_node_changed_via(
&mut self,
rule_name: &str,
def: &RuleDef,
n: u32,
n_label: Option<u32>,
changed: Option<(&str, Option<Value>)>,
g: &mut GraphMut<'_>,
) {
let src_sym = g.syms.get(&def.src_label);
let dst_sym = g.syms.get(&def.dst_label);
let via_sym = def.via_label.as_deref().and_then(|l| g.syms.get(l));
let as_src = src_sym.is_some() && n_label == src_sym;
let as_dst = dst_sym.is_some() && n_label == dst_sym;
let as_via = via_sym.is_some() && n_label == via_sym;
let fires = match changed {
None => as_src || as_via || as_dst,
Some((field, _)) => {
let wf = def.watched_fields();
(wf.contains(field)) && (as_src || as_via || as_dst)
}
};
if !fires {
return;
}
*self.fires.entry(rule_name.to_string()).or_default() += 1;
let mut affected_srcs: BTreeSet<u32> = BTreeSet::new();
if as_src {
affected_srcs.insert(n);
}
if as_via {
let via_edge_str = def.via_edge.as_deref().unwrap();
let via_dir = def.via_dir.unwrap_or(core_storage::Direction::Out);
let rev_dir = match via_dir {
core_storage::Direction::Out => core_storage::Direction::In,
core_storage::Direction::In => core_storage::Direction::Out,
};
if let (Some(via_etype), Some(s_sym)) = (g.syms.get(via_edge_str), src_sym) {
for &src in g.topo.neighbors(via_etype, rev_dir, n).as_ref() {
if g.labels.get(src as usize).copied() == Some(s_sym) {
affected_srcs.insert(src);
}
}
}
}
if as_dst {
let desired_touching_n = compute_desired_via(def, ViaAnchor::Dst(n), g);
for (src, _dst) in desired_touching_n.keys() {
affected_srcs.insert(*src);
}
let et = g.syms.intern(&def.edge_type);
let rid = self.rule_intern.get(rule_name).copied();
let old_srcs: Vec<u32> = self
.by_node
.get(&n)
.into_iter()
.flatten()
.filter(|(r, t, _s, d)| Some(*r) == rid && *t == et && *d == n)
.map(|(_, _, s, _)| *s)
.collect();
affected_srcs.extend(old_srcs);
}
let affected_srcs: Vec<u32> = affected_srcs.into_iter().collect();
if let Some(k) = def.max_edges {
let mut prov = ProvSets {
set: self.provenance.entry(rule_name.to_string()).or_default(),
owned: &mut self.owned,
by_node: &mut self.by_node,
rule_intern: &mut self.rule_intern,
intern_rule: &mut self.intern_rule,
deltas: &mut self.pending_deltas,
emit: self.emit_deltas,
};
for src in affected_srcs {
let desired_src = compute_desired_via(def, ViaAnchor::Src(src), g);
let top_k = filter_src_top_k(desired_src, k, g.ids);
apply_per_src_top_k(def, src, top_k, &mut prov, g);
}
} else {
let tripped = self.tripped.entry(rule_name.to_string()).or_default();
let budget = edge_budget(def);
for src in affected_srcs {
let desired_src = compute_desired_via(def, ViaAnchor::Src(src), g);
if !*tripped {
let mut prov = ProvSets {
set: self.provenance.entry(rule_name.to_string()).or_default(),
owned: &mut self.owned,
by_node: &mut self.by_node,
rule_intern: &mut self.rule_intern,
intern_rule: &mut self.intern_rule,
deltas: &mut self.pending_deltas,
emit: self.emit_deltas,
};
apply_desired(def, desired_src, Some(src), &mut prov, tripped, g);
}
let _ = budget;
}
}
}
pub fn on_edge_changed(
&mut self,
etype_str: &str,
src_id: u32,
dst_id: u32,
g: &mut GraphMut<'_>,
) {
let rule_names: Vec<String> = self.rules.keys().cloned().collect();
for rule_name in rule_names {
let def = self.rules[&rule_name].clone();
let Some(ref via_edge) = def.via_edge else {
continue; };
if via_edge != etype_str {
continue; }
let src_sym = match g.syms.get(&def.src_label) {
Some(s) => s,
None => continue,
};
let via_sym = match def.via_label.as_deref().and_then(|l| g.syms.get(l)) {
Some(s) => s,
None => continue,
};
let via_dir = def.via_dir.unwrap_or(core_storage::Direction::Out);
let (rule_src, rule_via) = match via_dir {
core_storage::Direction::Out => (src_id, dst_id),
core_storage::Direction::In => (dst_id, src_id),
};
if g.labels.get(rule_src as usize).copied() != Some(src_sym) {
continue;
}
if g.labels.get(rule_via as usize).copied() != Some(via_sym) {
continue;
}
*self.fires.entry(rule_name.clone()).or_default() += 1;
let desired_src = compute_desired_via(&def, ViaAnchor::Src(rule_src), g);
if let Some(k) = def.max_edges {
let mut prov = ProvSets {
set: self.provenance.entry(rule_name).or_default(),
owned: &mut self.owned,
by_node: &mut self.by_node,
rule_intern: &mut self.rule_intern,
intern_rule: &mut self.intern_rule,
deltas: &mut self.pending_deltas,
emit: self.emit_deltas,
};
let top_k = filter_src_top_k(desired_src, k, g.ids);
apply_per_src_top_k(&def, rule_src, top_k, &mut prov, g);
} else {
let tripped = self.tripped.entry(rule_name.clone()).or_default();
let mut prov = ProvSets {
set: self.provenance.entry(rule_name).or_default(),
owned: &mut self.owned,
by_node: &mut self.by_node,
rule_intern: &mut self.rule_intern,
intern_rule: &mut self.intern_rule,
deltas: &mut self.pending_deltas,
emit: self.emit_deltas,
};
apply_desired(&def, desired_src, Some(rule_src), &mut prov, tripped, g);
}
}
}
pub fn on_node_removed(&mut self, n: u32, g: &mut GraphMut<'_>) {
let n_label = g.labels.get(n as usize).copied();
let rule_names: Vec<String> = self.rules.keys().cloned().collect();
for rule_name in rule_names {
let def = self.rules[&rule_name].clone();
let src_sym = g.syms.get(&def.src_label);
let dst_sym = g.syms.get(&def.dst_label);
let as_src = src_sym.is_some() && n_label == src_sym;
let as_dst = dst_sym.is_some() && n_label == dst_sym;
{
let cur_getter = |f: &str| g.props.get(n, f).cloned();
let idx = self.indexes.get_mut(&rule_name).unwrap();
if as_src {
let spec = src_lookup_spec_for(&def);
idx.src_side.remove(&spec, n, &cur_getter);
}
if as_dst {
let spec = candidate_spec_for(&def);
idx.dst_side.remove(&spec, n, &cur_getter);
}
}
self.maybe_queue_ivf_rebuild(&rule_name, &def);
}
let touching: Vec<(String, Triple)> = self
.by_node
.get(&n)
.into_iter()
.flatten()
.map(|&(rid, t, s, d)| (self.intern_rule[rid as usize].clone(), (t, s, d)))
.collect();
let topk_backfill: Vec<(String, u32)> = touching
.iter()
.filter_map(|(rule_name, triple)| {
let &(_, s, d) = triple;
let def = self.rules.get(rule_name)?;
def.max_edges?; if d == n && s != n {
Some((rule_name.clone(), s))
} else {
None
}
})
.collect();
for (rule_name, triple) in touching {
let (t, s, d) = triple;
g.topo.remove_edge(t, s, d);
g.edge_props.remove_edge(t, s, d);
if let Some(set) = self.provenance.get_mut(&rule_name) {
ProvSets {
set,
owned: &mut self.owned,
by_node: &mut self.by_node,
rule_intern: &mut self.rule_intern,
intern_rule: &mut self.intern_rule,
deltas: &mut self.pending_deltas,
emit: self.emit_deltas,
}
.remove(&rule_name, triple, g.ids, g.syms);
}
}
for (rule_name, src) in topk_backfill {
let def = self.rules[&rule_name].clone();
let k = def.max_edges.unwrap(); let desired_src = compute_desired(&def, &self.indexes[&rule_name], src, true, g);
let top_k = filter_src_top_k(desired_src, k, g.ids);
let mut prov = ProvSets {
set: self.provenance.entry(rule_name.clone()).or_default(),
owned: &mut self.owned,
by_node: &mut self.by_node,
rule_intern: &mut self.rule_intern,
intern_rule: &mut self.intern_rule,
deltas: &mut self.pending_deltas,
emit: self.emit_deltas,
};
apply_per_src_top_k(&def, src, top_k, &mut prov, g);
}
}
pub fn rebuild(&mut self, name: &str, g: &mut GraphMut<'_>) -> Result<(), String> {
if !self.rules.contains_key(name) {
return Err(format!("rule {:?} not found", name));
}
self.rebuild_needed.remove(name);
let def = self.rules[name].clone();
*self.indexes.get_mut(name).unwrap() = RuleIndex::default();
if def.approximate {
let idx = self.indexes.get_mut(name).unwrap();
idx.src_side.init_hnsw(name);
idx.dst_side.init_hnsw(name);
}
let n_total = g.ids.len() as u32;
for id in 0..n_total {
let label_sym = match g.labels.get(id as usize).copied() {
Some(s) if s != u32::MAX => s,
_ => continue,
};
let idx = self.indexes.get_mut(name).unwrap();
index_node_for_rule(id, label_sym, &def, idx, g.syms, g.props);
}
if def.approximate {
let idx = self.indexes.get_mut(name).unwrap();
idx.src_side.fit_ivf_clusters(name);
idx.dst_side.fit_ivf_clusters(name);
}
let mut prov = ProvSets {
set: self.provenance.get_mut(name).unwrap(),
owned: &mut self.owned,
by_node: &mut self.by_node,
rule_intern: &mut self.rule_intern,
intern_rule: &mut self.intern_rule,
deltas: &mut self.pending_deltas,
emit: self.emit_deltas,
};
if let Some(k) = def.max_edges {
apply_streaming_rebuild_top_k(&def, k, &self.indexes[name], &mut prov, g);
} else {
let tripped = self.tripped.get_mut(name).unwrap();
apply_streaming_rebuild(&def, &self.indexes[name], &mut prov, tripped, g);
}
let fires = self.fires.entry(name.to_string()).or_default();
bump_fires_for_participants(&def, g, fires);
Ok(())
}
#[cfg(test)]
fn by_node_consistent(&self) -> bool {
let (rebuilt, intern, names) = rebuild_by_node(&self.provenance);
resolve_by_node(&self.by_node, &self.intern_rule) == resolve_by_node(&rebuilt, &names)
&& intern.len() == names.len()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::def::{evaluate, NodeView, Predicate, RuleDef};
use core_storage::{ColumnStore, Direction, EdgeProps, IdMap, Interner, Topology, Value};
struct Fx {
ids: IdMap,
syms: Interner,
labels: Vec<u32>,
props: ColumnStore,
topo: Topology,
eprops: EdgeProps,
}
impl Fx {
fn new() -> Self {
Fx {
ids: IdMap::new(),
syms: Interner::new(),
labels: vec![],
props: ColumnStore::new(),
topo: Topology::new(),
eprops: EdgeProps::new(),
}
}
fn add(&mut self, label: &str, key: &str, props: Vec<(&str, Value)>) -> u32 {
let id = self.ids.get_or_insert(key);
let sym = self.syms.intern(label);
self.labels.resize(id as usize + 1, u32::MAX);
self.labels[id as usize] = sym;
for (f, v) in props {
self.props.set(id, f, v);
}
id
}
fn g(&mut self) -> GraphMut<'_> {
GraphMut {
ids: &self.ids,
syms: &mut self.syms,
labels: &self.labels,
props: &self.props,
topo: &mut self.topo,
edge_props: &mut self.eprops,
}
}
}
fn tags(items: &[&str]) -> Value {
Value::List(items.iter().map(|s| Value::Str((*s).into())).collect())
}
fn overlap_rule() -> RuleDef {
RuleDef {
name: "rel".into(),
src_label: "A".into(),
dst_label: "A".into(),
predicate: Predicate::Overlap {
field: "tags".into(),
min: 0.4,
},
edge_type: "REL".into(),
weight_prop: Some("score".into()),
max_edges: None,
approximate: false,
via_label: None,
via_edge: None,
via_dir: None,
}
}
fn emb(xs: &[f64]) -> Value {
Value::List(xs.iter().copied().map(Value::Float).collect())
}
fn approx_vec_rule() -> RuleDef {
RuleDef {
name: "sim".into(),
src_label: "V".into(),
dst_label: "V".into(),
predicate: Predicate::VectorSimilar {
field: "emb".into(),
min: 0.5,
},
edge_type: "SIM".into(),
weight_prop: None,
max_edges: None,
approximate: true,
via_label: None,
via_edge: None,
via_dir: None,
}
}
#[test]
fn approximate_rule_rebuilds_after_drift_threshold() {
with_ivf_drift_rebuild(1, || {
let mut fx = Fx::new();
let mut ids = Vec::new();
for i in 0..6 {
let x = i as f64 * 0.2;
ids.push(fx.add("V", &format!("v{i}"), vec![("emb", emb(&[x, 1.0 - x]))]));
}
let mut eng = RuleEngine::new();
{
let mut g = fx.g();
eng.create_rule(approx_vec_rule(), &mut g).unwrap();
}
assert!(eng.take_rebuild_needed().is_empty());
{
let mut g = fx.g();
eng.on_node_removed(ids[0], &mut g);
}
assert!(
eng.take_rebuild_needed().is_empty(),
"drift=1 is not > threshold 1"
);
{
let mut g = fx.g();
eng.on_node_removed(ids[1], &mut g);
}
assert_eq!(eng.take_rebuild_needed(), vec!["sim".to_string()]);
{
let mut g = fx.g();
eng.rebuild("sim", &mut g).unwrap();
}
assert!(
eng.take_rebuild_needed().is_empty(),
"rebuild must reset drift and not re-queue itself"
);
let drift = eng
.export_ivf_state()
.get("sim")
.map(|(_, dst)| dst.2)
.unwrap();
assert_eq!(drift, 0, "rebuild resets dst-side IVF drift");
});
}
#[test]
fn backfill_creates_edges_with_scores_and_delete_removes_exactly_them() {
let mut fx = Fx::new();
let a = fx.add("A", "a", vec![("tags", tags(&["x", "y"]))]);
let b = fx.add("A", "b", vec![("tags", tags(&["x", "y"]))]);
let _c = fx.add("A", "c", vec![("tags", tags(&["q"]))]);
let et = fx.syms.intern("REL");
fx.topo.add_edge(et, a, b);
let mut eng = RuleEngine::new();
let mut g = fx.g();
eng.create_rule(overlap_rule(), &mut g).unwrap();
assert!(g.topo.neighbors(et, Direction::Out, b).contains(&a));
assert_eq!(
g.edge_props.get(et, b, a, "score"),
Some(&Value::Float(1.0))
);
assert!(!eng.is_owned(et, a, b));
assert!(eng.is_owned(et, b, a));
eng.delete_rule("rel", &mut g).unwrap();
assert!(g.topo.neighbors(et, Direction::Out, a).contains(&b)); assert!(!g.topo.neighbors(et, Direction::Out, b).contains(&a)); assert_eq!(g.edge_props.get(et, b, a, "score"), None);
}
#[test]
fn incremental_update_adds_and_removes_edges() {
let mut fx = Fx::new();
let a = fx.add("A", "a", vec![("tags", tags(&["x", "y"]))]);
let b = fx.add("A", "b", vec![("tags", tags(&["y", "z"]))]);
let et = fx.syms.intern("REL");
let mut eng = RuleEngine::new();
{
let mut g = fx.g();
eng.create_rule(overlap_rule(), &mut g).unwrap(); assert_eq!(g.topo.edge_count(), 0);
}
let old = fx.props.get(b, "tags").cloned();
fx.props.set(b, "tags", tags(&["x", "y"]));
{
let mut g = fx.g();
eng.on_node_changed(b, Some(("tags", old)), &mut g);
assert!(g.topo.neighbors(et, Direction::Out, a).contains(&b));
assert!(g.topo.neighbors(et, Direction::Out, b).contains(&a));
}
let old = fx.props.get(b, "tags").cloned();
fx.props.set(b, "tags", tags(&["qqq"]));
let mut g = fx.g();
eng.on_node_changed(b, Some(("tags", old)), &mut g);
assert_eq!(g.topo.edge_count(), 0);
assert_eq!(g.edge_props.get(et, a, b, "score"), None);
}
#[test]
fn key_match_new_node_links_and_rebuild_is_noop() {
let mut fx = Fx::new();
fx.add("C", "c1", vec![]);
let mut eng = RuleEngine::new();
{
let mut g = fx.g();
eng.create_rule(
RuleDef {
name: "fk".into(),
src_label: "T".into(),
dst_label: "C".into(),
predicate: Predicate::KeyMatch {
field: "cid".into(),
},
edge_type: "AT".into(),
weight_prop: None,
max_edges: None,
approximate: false,
via_label: None,
via_edge: None,
via_dir: None,
},
&mut g,
)
.unwrap();
}
let t = fx.add("T", "t1", vec![("cid", Value::Str("c1".into()))]);
let (at, c1, count_before) = {
let mut g = fx.g();
eng.on_node_changed(t, None, &mut g);
let at = g.syms.get("AT").unwrap();
let c1 = g.ids.get("c1").unwrap();
assert!(g.topo.neighbors(at, Direction::Out, t).contains(&c1));
(at, c1, g.topo.edge_count())
};
let mut g = fx.g();
eng.rebuild("fk", &mut g).unwrap();
assert_eq!(g.topo.edge_count(), count_before); assert!(g.topo.neighbors(at, Direction::Out, t).contains(&c1));
}
#[test]
fn score_refresh_on_persisting_owned_edge() {
let mut fx = Fx::new();
let a = fx.add("A", "a", vec![("tags", tags(&["x", "y", "z"]))]);
let b = fx.add("A", "b", vec![("tags", tags(&["x", "y", "q"]))]);
let et = fx.syms.intern("SIM");
let mut eng = RuleEngine::new();
{
let mut g = fx.g();
eng.create_rule(
RuleDef {
name: "sim".into(),
src_label: "A".into(),
dst_label: "A".into(),
predicate: Predicate::Overlap {
field: "tags".into(),
min: 0.2,
},
edge_type: "SIM".into(),
weight_prop: Some("score".into()),
max_edges: None,
approximate: false,
via_label: None,
via_edge: None,
via_dir: None,
},
&mut g,
)
.unwrap();
assert!(g.topo.neighbors(et, Direction::Out, a).contains(&b));
assert!(g.topo.neighbors(et, Direction::Out, b).contains(&a));
assert!(eng.is_owned(et, a, b) || eng.is_owned(et, b, a));
let check = |v: Option<&Value>| {
if let Some(Value::Float(f)) = v {
assert!(
(f - 0.5).abs() < 1e-9,
"initial score should be 0.5, got {f}"
);
}
};
check(g.edge_props.get(et, a, b, "score"));
check(g.edge_props.get(et, b, a, "score"));
}
let old = fx.props.get(b, "tags").cloned();
fx.props.set(b, "tags", tags(&["x", "y", "z"]));
{
let mut g = fx.g();
eng.on_node_changed(b, Some(("tags", old)), &mut g);
assert!(g.topo.neighbors(et, Direction::Out, a).contains(&b));
assert!(g.topo.neighbors(et, Direction::Out, b).contains(&a));
assert_eq!(
g.edge_props.get(et, a, b, "score"),
Some(&Value::Float(1.0)),
"score on a→b must refresh to 1.0"
);
assert_eq!(
g.edge_props.get(et, b, a, "score"),
Some(&Value::Float(1.0)),
"score on b→a must refresh to 1.0"
);
}
}
#[test]
fn dst_side_keymatch_links_when_c_node_inserted_after_t() {
let mut fx = Fx::new();
let t = fx.add("T", "t1", vec![("cid", Value::Str("c9".into()))]);
let mut eng = RuleEngine::new();
{
let mut g = fx.g();
eng.create_rule(
RuleDef {
name: "fk".into(),
src_label: "T".into(),
dst_label: "C".into(),
predicate: Predicate::KeyMatch {
field: "cid".into(),
},
edge_type: "AT".into(),
weight_prop: None,
max_edges: None,
approximate: false,
via_label: None,
via_edge: None,
via_dir: None,
},
&mut g,
)
.unwrap();
let at = g.syms.intern("AT");
assert_eq!(g.topo.edge_count(), 0, "no C node yet → no edge");
let _ = at;
}
let c9 = fx.add("C", "c9", vec![]);
{
let mut g = fx.g();
eng.on_node_changed(c9, None, &mut g);
let at = g.syms.get("AT").unwrap();
assert!(
g.topo.neighbors(at, Direction::Out, t).contains(&c9),
"T→C edge must appear when C node is inserted"
);
assert!(eng.is_owned(at, t, c9));
}
}
#[test]
fn on_node_removed_retracts_both_sides_and_deindexes() {
let mut fx = Fx::new();
let a = fx.add("A", "a", vec![("tags", tags(&["x", "y"]))]);
let b = fx.add("A", "b", vec![("tags", tags(&["x", "y"]))]);
let et = fx.syms.intern("REL");
let mut eng = RuleEngine::new();
{
let mut g = fx.g();
eng.create_rule(overlap_rule(), &mut g).unwrap();
assert!(g.topo.neighbors(et, Direction::Out, a).contains(&b));
assert!(g.topo.neighbors(et, Direction::Out, b).contains(&a));
}
{
let mut g = fx.g();
eng.on_node_removed(a, &mut g);
assert!(!g.topo.neighbors(et, Direction::Out, a).contains(&b));
assert!(!g.topo.neighbors(et, Direction::Out, b).contains(&a));
assert_eq!(g.edge_props.get(et, a, b, "score"), None);
assert_eq!(g.edge_props.get(et, b, a, "score"), None);
assert!(!eng.is_owned(et, a, b));
assert!(!eng.is_owned(et, b, a));
}
let c = fx.add("A", "c", vec![("tags", tags(&["x", "y"]))]);
{
let mut g = fx.g();
eng.on_node_changed(c, None, &mut g);
assert!(g.topo.neighbors(et, Direction::Out, b).contains(&c));
assert!(g.topo.neighbors(et, Direction::Out, c).contains(&b));
assert!(!g.topo.neighbors(et, Direction::Out, c).contains(&a));
assert!(!g.topo.neighbors(et, Direction::Out, a).contains(&c));
}
{
let mut g = fx.g();
eng.on_node_removed(a, &mut g);
assert!(g.topo.neighbors(et, Direction::Out, b).contains(&c));
}
}
#[test]
fn duplicate_name_and_unknown_delete_error() {
let mut fx = Fx::new();
let mut eng = RuleEngine::new();
let mut g = fx.g();
eng.create_rule(overlap_rule(), &mut g).unwrap();
assert!(eng.create_rule(overlap_rule(), &mut g).is_err());
assert!(eng.delete_rule("nope", &mut g).is_err());
}
#[test]
fn coowned_edge_type_survives_first_delete_gone_after_second() {
let mut fx = Fx::new();
let a = fx.add("A", "a", vec![("tags", tags(&["x", "y"]))]);
let b = fx.add("A", "b", vec![("tags", tags(&["x", "y"]))]);
let mut eng = RuleEngine::new();
{
let mut g = fx.g();
eng.create_rule(
RuleDef {
name: "r1".into(),
src_label: "A".into(),
dst_label: "A".into(),
predicate: Predicate::Overlap {
field: "tags".into(),
min: 0.1,
},
edge_type: "REL2".into(),
weight_prop: None,
max_edges: None,
approximate: false,
via_label: None,
via_edge: None,
via_dir: None,
},
&mut g,
)
.unwrap();
eng.create_rule(
RuleDef {
name: "r2".into(),
src_label: "A".into(),
dst_label: "A".into(),
predicate: Predicate::Overlap {
field: "tags".into(),
min: 0.2,
},
edge_type: "REL2".into(),
weight_prop: None,
max_edges: None,
approximate: false,
via_label: None,
via_edge: None,
via_dir: None,
},
&mut g,
)
.unwrap();
let et = g.syms.intern("REL2");
assert!(
g.topo.neighbors(et, Direction::Out, a).contains(&b),
"a→b must exist after both rules created"
);
assert!(
g.topo.neighbors(et, Direction::Out, b).contains(&a),
"b→a must exist after both rules created"
);
eng.delete_rule("r1", &mut g).unwrap();
assert!(
g.topo.neighbors(et, Direction::Out, a).contains(&b),
"a→b must survive R1 deletion (R2 rebuilds and claims it)"
);
assert!(
g.topo.neighbors(et, Direction::Out, b).contains(&a),
"b→a must survive R1 deletion (R2 rebuilds and claims it)"
);
assert!(
eng.is_owned(et, a, b),
"a→b must be owned by R2 after rebuild"
);
assert!(
eng.is_owned(et, b, a),
"b→a must be owned by R2 after rebuild"
);
eng.delete_rule("r2", &mut g).unwrap();
assert!(
!g.topo.neighbors(et, Direction::Out, a).contains(&b),
"a→b must be gone after both rules deleted"
);
assert!(
!g.topo.neighbors(et, Direction::Out, b).contains(&a),
"b→a must be gone after both rules deleted"
);
}
}
fn topk_eq_rule(k: u64) -> RuleDef {
RuleDef {
name: "eq".into(),
src_label: "N".into(),
dst_label: "N".into(),
predicate: Predicate::FieldEqual { field: "k".into() },
edge_type: "EQ".into(),
weight_prop: None,
max_edges: Some(k),
approximate: false,
via_label: None,
via_edge: None,
via_dir: None,
}
}
fn prov_pairs(eng: &RuleEngine, name: &str) -> BTreeSet<(u32, u32)> {
eng.provenance()
.get(name)
.map(|s| s.iter().map(|&(_, a, b)| (a, b)).collect())
.unwrap_or_default()
}
#[test]
fn topk_k1_keeps_best_scored_dst() {
let mut fx = Fx::new();
let mut eng = RuleEngine::new();
{
let mut g = fx.g();
eng.create_rule(topk_eq_rule(1), &mut g).unwrap();
}
let mut ids = Vec::new();
for i in 0..4usize {
let id = fx.add(
"N",
&format!("n{i}"),
vec![("k", Value::Str("const".into()))],
);
ids.push(id);
let mut g = fx.g();
eng.on_node_changed(id, None, &mut g);
}
let et = fx.syms.get("EQ").unwrap();
let expected_dsts = [ids[1], ids[0], ids[0], ids[0]];
for (i, (&src, &expected_dst)) in ids.iter().zip(expected_dsts.iter()).enumerate() {
let out: Vec<u32> = fx.topo.neighbors(et, Direction::Out, src).to_vec();
assert_eq!(
out,
vec![expected_dst],
"src n{i} should point only to the best dst"
);
}
assert_eq!(eng.provenance()["eq"].len(), 4);
assert!(!eng.is_tripped("eq"), "top-k rules never trip");
}
#[test]
fn topk_insert_evict() {
let mut fx = Fx::new();
let rule = RuleDef {
name: "nw".into(),
src_label: "S".into(),
dst_label: "D".into(),
predicate: Predicate::NumericWithin {
field: "v".into(),
tolerance: 10.0,
},
edge_type: "NEAR".into(),
weight_prop: Some("score".into()),
max_edges: Some(1),
approximate: false,
via_label: None,
via_edge: None,
via_dir: None,
};
let mut eng = RuleEngine::new();
{
let mut g = fx.g();
eng.create_rule(rule, &mut g).unwrap();
}
let s0 = fx.add("S", "s0", vec![("v", Value::Float(0.0))]);
let d_far = fx.add("D", "d_far", vec![("v", Value::Float(9.0))]);
{
let mut g = fx.g();
eng.on_node_changed(s0, None, &mut g);
eng.on_node_changed(d_far, None, &mut g);
}
let et = fx.syms.get("NEAR").unwrap();
assert!(fx.topo.neighbors(et, Direction::Out, s0).contains(&d_far));
assert_eq!(eng.provenance()["nw"].len(), 1);
let d_close = fx.add("D", "d_close", vec![("v", Value::Float(1.0))]);
{
let mut g = fx.g();
eng.on_node_changed(d_close, None, &mut g);
}
let out: Vec<u32> = fx.topo.neighbors(et, Direction::Out, s0).to_vec();
assert_eq!(out, vec![d_close], "d_close should evict d_far");
assert!(!fx.topo.neighbors(et, Direction::Out, s0).contains(&d_far));
assert_eq!(eng.provenance()["nw"].len(), 1);
assert!(eng.by_node_consistent());
}
#[test]
fn topk_retract_backfill() {
let mut fx = Fx::new();
let rule = RuleDef {
name: "nw".into(),
src_label: "S".into(),
dst_label: "D".into(),
predicate: Predicate::NumericWithin {
field: "v".into(),
tolerance: 10.0,
},
edge_type: "NEAR".into(),
weight_prop: Some("score".into()),
max_edges: Some(1),
approximate: false,
via_label: None,
via_edge: None,
via_dir: None,
};
let mut eng = RuleEngine::new();
let s0 = fx.add("S", "s0", vec![("v", Value::Float(0.0))]);
let d_close = fx.add("D", "d_close", vec![("v", Value::Float(1.0))]); let d_far = fx.add("D", "d_far", vec![("v", Value::Float(8.0))]); {
let mut g = fx.g();
eng.create_rule(rule, &mut g).unwrap();
}
let et = fx.syms.get("NEAR").unwrap();
assert!(fx.topo.neighbors(et, Direction::Out, s0).contains(&d_close));
assert!(!fx.topo.neighbors(et, Direction::Out, s0).contains(&d_far));
assert_eq!(eng.provenance()["nw"].len(), 1);
let old = fx.props.get(d_close, "v").cloned();
fx.props.set(d_close, "v", Value::Float(50.0));
{
let mut g = fx.g();
eng.on_node_changed(d_close, Some(("v", old)), &mut g);
}
assert!(!fx.topo.neighbors(et, Direction::Out, s0).contains(&d_close));
assert!(
fx.topo.neighbors(et, Direction::Out, s0).contains(&d_far),
"d_far should backfill after d_close retracted"
);
assert_eq!(eng.provenance()["nw"].len(), 1);
assert!(eng.by_node_consistent());
}
#[test]
fn topk_tie_broken_by_dst_key() {
let mut fx = Fx::new();
let mut eng = RuleEngine::new();
{
let mut g = fx.g();
eng.create_rule(topk_eq_rule(2), &mut g).unwrap();
}
for name in ["a", "b", "c", "d", "e"] {
let id = fx.add("N", name, vec![("k", Value::Str("x".into()))]);
let mut g = fx.g();
eng.on_node_changed(id, None, &mut g);
}
let et = fx.syms.get("EQ").unwrap();
let get_id = |key: &str| fx.ids.get(key).unwrap();
let a = get_id("a");
let b = get_id("b");
let c = get_id("c");
let out_a: BTreeSet<u32> = fx
.topo
.neighbors(et, Direction::Out, a)
.iter()
.copied()
.collect();
assert!(out_a.contains(&b), "a→b (b is best key after a)");
assert!(out_a.contains(&c), "a→c (c is 2nd best key)");
assert_eq!(out_a.len(), 2);
let e = get_id("e");
let out_e: BTreeSet<u32> = fx
.topo
.neighbors(et, Direction::Out, e)
.iter()
.copied()
.collect();
assert!(out_e.contains(&a), "e→a");
assert!(out_e.contains(&b), "e→b");
assert_eq!(out_e.len(), 2);
assert!(eng.by_node_consistent());
}
#[test]
fn topk_k_larger_than_candidate_count() {
let mut fx = Fx::new();
let mut eng = RuleEngine::new();
{
let mut g = fx.g();
eng.create_rule(topk_eq_rule(100), &mut g).unwrap();
}
for i in 0..4usize {
let id = fx.add("N", &format!("n{i}"), vec![("k", Value::Str("c".into()))]);
let mut g = fx.g();
eng.on_node_changed(id, None, &mut g);
}
assert_eq!(eng.provenance()["eq"].len(), 12);
assert!(!eng.is_tripped("eq"));
}
#[test]
fn topk_rebuild_exact() {
let mut fx = Fx::new();
let mut eng = RuleEngine::new();
{
let mut g = fx.g();
eng.create_rule(topk_eq_rule(1), &mut g).unwrap();
}
let _a = fx.add("N", "a", vec![("k", Value::Str("x".into()))]);
let _b = fx.add("N", "b", vec![("k", Value::Str("x".into()))]);
let _c = fx.add("N", "c", vec![("k", Value::Str("x".into()))]);
{
let mut g = fx.g();
eng.on_node_changed(_a, None, &mut g);
eng.on_node_changed(_b, None, &mut g);
eng.on_node_changed(_c, None, &mut g);
}
assert_eq!(eng.provenance()["eq"].len(), 3);
{
let mut g = fx.g();
eng.rebuild("eq", &mut g).unwrap();
}
assert_eq!(eng.provenance()["eq"].len(), 3);
assert!(!eng.is_tripped("eq"));
assert!(eng.by_node_consistent());
}
#[test]
fn topk_by_node_consistent() {
let mut fx = Fx::new();
let mut eng = RuleEngine::new();
{
let mut g = fx.g();
eng.create_rule(topk_eq_rule(2), &mut g).unwrap();
}
for i in 0..5usize {
let id = fx.add(
"N",
&format!("n{i}"),
vec![("k", Value::Str("const".into()))],
);
let mut g = fx.g();
eng.on_node_changed(id, None, &mut g);
}
assert!(eng.by_node_consistent(), "consistent after insertions");
let id2 = fx.ids.get("n2").unwrap();
let old = fx.props.get(id2, "k").cloned();
fx.props.set(id2, "k", Value::Str("other".into()));
{
let mut g = fx.g();
eng.on_node_changed(id2, Some(("k", old)), &mut g);
}
assert!(eng.by_node_consistent(), "consistent after eviction");
{
let mut g = fx.g();
eng.rebuild("eq", &mut g).unwrap();
}
assert!(eng.by_node_consistent(), "consistent after rebuild");
}
fn numeric_rule() -> RuleDef {
RuleDef {
name: "nw".into(),
src_label: "C".into(),
dst_label: "C".into(),
predicate: Predicate::NumericWithin {
field: "year".into(),
tolerance: 2.0,
},
edge_type: "NEAR".into(),
weight_prop: Some("score".into()),
max_edges: None,
approximate: false,
via_label: None,
via_edge: None,
via_dir: None,
}
}
fn geo_rule() -> RuleDef {
RuleDef {
name: "geo".into(),
src_label: "City".into(),
dst_label: "City".into(),
predicate: Predicate::GeoRadius {
field: "loc".into(),
km: 400.0,
},
edge_type: "NEAR_GEO".into(),
weight_prop: Some("score".into()),
max_edges: None,
approximate: false,
via_label: None,
via_edge: None,
via_dir: None,
}
}
fn vec_rule() -> RuleDef {
RuleDef {
name: "vec".into(),
src_label: "Doc".into(),
dst_label: "Doc".into(),
predicate: Predicate::VectorSimilar {
field: "emb".into(),
min: 0.9,
},
edge_type: "SIM".into(),
weight_prop: Some("score".into()),
max_edges: None,
approximate: false,
via_label: None,
via_edge: None,
via_dir: None,
}
}
fn pair_edges(topo: &Topology, et: u32, a: u32, b: u32) -> bool {
topo.neighbors(et, Direction::Out, a).contains(&b)
&& topo.neighbors(et, Direction::Out, b).contains(&a)
}
#[test]
fn numeric_within_incremental_crosses_bucket_and_clears_old_index() {
let mut fx = Fx::new();
let a = fx.add("C", "a", vec![("year", Value::Float(10.0))]);
let b = fx.add("C", "b", vec![("year", Value::Float(12.0))]);
let et = fx.syms.intern("NEAR");
let mut eng = RuleEngine::new();
{
let mut g = fx.g();
eng.create_rule(numeric_rule(), &mut g).unwrap();
assert!(pair_edges(g.topo, et, a, b));
}
let old = fx.props.get(b, "year").cloned();
fx.props.set(b, "year", Value::Float(16.1));
{
let mut g = fx.g();
eng.on_node_changed(b, Some(("year", old)), &mut g);
assert!(!pair_edges(g.topo, et, a, b));
assert_eq!(g.topo.edge_count(), 0);
}
let def = numeric_rule();
let spec = candidate_spec_for(&def);
let old_map: std::collections::HashMap<_, _> =
[("year".to_string(), Value::Float(12.0))].into();
let old_get = |f: &str| old_map.get(f).cloned();
let src_hits = eng.indexes["nw"].src_side.candidates(&spec, &old_get);
let dst_hits = eng.indexes["nw"].dst_side.candidates(&spec, &old_get);
assert!(!src_hits.contains(&b), "old src bucket must drop b");
assert!(!dst_hits.contains(&b), "old dst bucket must drop b");
assert!(src_hits.contains(&a));
let old = fx.props.get(b, "year").cloned();
fx.props.set(b, "year", Value::Float(11.9));
let mut g = fx.g();
eng.on_node_changed(b, Some(("year", old)), &mut g);
assert!(pair_edges(g.topo, et, a, b));
}
fn loc_val(lat: f64, lon: f64) -> Value {
Value::List(vec![Value::Float(lat), Value::Float(lon)])
}
fn emb_val(vals: &[f64]) -> Value {
Value::List(vals.iter().copied().map(Value::Float).collect())
}
#[test]
fn rebuild_is_noop_for_numeric_geo_and_vector() {
let mut fx = Fx::new();
let ca = fx.add("C", "ca", vec![("year", Value::Int(1998))]);
let cb = fx.add("C", "cb", vec![("year", Value::Float(2000.0))]);
let pa = fx.add("City", "paris", vec![("loc", loc_val(48.8566, 2.3522))]);
let lo = fx.add("City", "london", vec![("loc", loc_val(51.5074, -0.1278))]);
let da = fx.add("Doc", "d1", vec![("emb", emb_val(&[1.0, 0.0]))]);
let db = fx.add("Doc", "d2", vec![("emb", emb_val(&[1.0, 0.0]))]);
let mut eng = RuleEngine::new();
{
let mut g = fx.g();
eng.create_rule(numeric_rule(), &mut g).unwrap();
eng.create_rule(geo_rule(), &mut g).unwrap();
eng.create_rule(vec_rule(), &mut g).unwrap();
}
let (near, ngeo, sim) = (
fx.syms.get("NEAR").unwrap(),
fx.syms.get("NEAR_GEO").unwrap(),
fx.syms.get("SIM").unwrap(),
);
assert!(pair_edges(&fx.topo, near, ca, cb));
assert!(pair_edges(&fx.topo, ngeo, pa, lo));
assert!(pair_edges(&fx.topo, sim, da, db));
let before = fx.topo.edge_count();
{
let mut g = fx.g();
eng.rebuild("nw", &mut g).unwrap();
eng.rebuild("geo", &mut g).unwrap();
eng.rebuild("vec", &mut g).unwrap();
}
assert_eq!(fx.topo.edge_count(), before);
assert!(pair_edges(&fx.topo, near, ca, cb));
assert!(pair_edges(&fx.topo, ngeo, pa, lo));
assert!(pair_edges(&fx.topo, sim, da, db));
}
fn fk_rule() -> RuleDef {
RuleDef {
name: "works_at".into(),
src_label: "T".into(),
dst_label: "C".into(),
predicate: Predicate::KeyMatch {
field: "cid".into(),
},
edge_type: "AT".into(),
weight_prop: None,
max_edges: None,
approximate: false,
via_label: None,
via_edge: None,
via_dir: None,
}
}
#[test]
fn by_node_matches_rebuild_after_mutation_storm() {
let mut fx = Fx::new();
let hub = fx.add("C", "hub", vec![]);
let other = fx.add("C", "other", vec![]);
let mut people = Vec::new();
for i in 0..40 {
let cid = if i < 30 { "hub" } else { "other" };
people.push(fx.add(
"T",
&format!("t{i}"),
vec![("cid", Value::Str(cid.into())), ("tags", tags(&["x", "y"]))],
));
}
let mut overlap = overlap_rule();
overlap.src_label = "T".into();
overlap.dst_label = "T".into();
let mut eng = RuleEngine::new();
{
let mut g = fx.g();
eng.create_rule(fk_rule(), &mut g).unwrap();
eng.create_rule(overlap, &mut g).unwrap();
}
assert!(eng.by_node_consistent());
assert_eq!(eng.provenance_touching_len(hub), 30);
for (i, &id) in people.iter().enumerate().take(15) {
let old = fx.props.get(id, "cid").cloned();
fx.props.set(id, "cid", Value::Str("other".into()));
let mut g = fx.g();
eng.on_node_changed(id, Some(("cid", old)), &mut g);
assert!(
eng.by_node_consistent(),
"inconsistent after cid update {i}"
);
}
for &id in people.iter().take(8) {
let old = fx.props.get(id, "tags").cloned();
fx.props.set(id, "tags", tags(&["q"]));
let mut g = fx.g();
eng.on_node_changed(id, Some(("tags", old)), &mut g);
}
assert!(eng.by_node_consistent());
{
let mut g = fx.g();
eng.on_node_removed(people[0], &mut g);
}
fx.labels[people[0] as usize] = u32::MAX;
assert!(eng.by_node_consistent());
assert_eq!(eng.provenance_touching_len(people[0]), 0);
{
let mut g = fx.g();
eng.rebuild("works_at", &mut g).unwrap();
eng.rebuild("rel", &mut g).unwrap();
}
assert!(eng.by_node_consistent());
{
let mut g = fx.g();
eng.delete_rule("rel", &mut g).unwrap();
}
assert!(eng.by_node_consistent());
assert_eq!(eng.provenance_touching(people[1]).count(), 1);
let (defs, prov, tripped, fires) = eng.to_persist();
let restored = RuleEngine::from_persist(defs, prov, tripped, fires);
assert!(restored.by_node_consistent());
assert_eq!(
restored.provenance_touching_len(hub),
eng.provenance_touching_len(hub)
);
assert_eq!(
restored.provenance_touching_len(other),
eng.provenance_touching_len(other)
);
}
#[test]
fn provenance_touching_high_degree_hub() {
let mut fx = Fx::new();
let hub = fx.add("C", "hub", vec![]);
let mut first = None;
for i in 0..256 {
let id = fx.add(
"T",
&format!("t{i}"),
vec![("cid", Value::Str("hub".into()))],
);
if first.is_none() {
first = Some(id);
}
}
let first = first.unwrap();
let mut eng = RuleEngine::new();
{
let mut g = fx.g();
eng.create_rule(fk_rule(), &mut g).unwrap();
}
assert!(eng.by_node_consistent());
assert_eq!(eng.provenance_touching_len(hub), 256);
assert_eq!(eng.provenance_touching_len(first), 1);
let hits: Vec<_> = eng.provenance_touching(first).collect();
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].0, "works_at");
assert_eq!(hits[0].2, first);
assert_eq!(hits[0].3, hub);
}
#[test]
fn by_node_consistent_across_inserts_and_rebuild() {
let mut fx = Fx::new();
let mut eng = RuleEngine::new();
let rule = RuleDef {
name: "eq".into(),
src_label: "N".into(),
dst_label: "N".into(),
predicate: Predicate::FieldEqual { field: "k".into() },
edge_type: "EQ".into(),
weight_prop: None,
max_edges: None, approximate: false,
via_label: None,
via_edge: None,
via_dir: None,
};
{
let mut g = fx.g();
eng.create_rule(rule, &mut g).unwrap();
}
let mut ids = Vec::new();
for i in 0..6 {
let id = fx.add(
"N",
&format!("n{i}"),
vec![("k", Value::Str("const".into()))],
);
ids.push(id);
let mut g = fx.g();
eng.on_node_changed(id, None, &mut g);
}
assert_eq!(eng.provenance()["eq"].len(), 30);
assert!(!eng.is_tripped("eq"));
assert!(eng.by_node_consistent(), "consistent after insertions");
let old = fx.props.get(ids[3], "k").cloned();
fx.props.set(ids[3], "k", Value::Str("other".into()));
{
let mut g = fx.g();
eng.on_node_changed(ids[3], Some(("k", old)), &mut g);
}
assert!(eng.by_node_consistent(), "consistent after property change");
{
let mut g = fx.g();
eng.rebuild("eq", &mut g).unwrap();
}
assert!(!eng.is_tripped("eq"));
assert!(eng.by_node_consistent(), "consistent after rebuild");
}
fn mix64(mut x: u64) -> u64 {
x = x.wrapping_add(0x9E3779B97F4A7C15);
x = (x ^ (x >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
x = (x ^ (x >> 27)).wrapping_mul(0x94D049BB133111EB);
x ^ (x >> 31)
}
fn rand_emb(seed: u64, i: u32, dim: usize) -> Value {
let vals: Vec<f64> = (0..dim)
.map(|d| {
let bits = mix64(seed ^ ((i as u64 + 1).wrapping_mul(0x100000001)) ^ (d as u64));
let mut f = (bits as f64) / (u64::MAX as f64) * 2.0 - 1.0;
if f == 0.0 {
f = 1.0;
}
f
})
.collect();
emb_val(&vals)
}
fn seed_docs(n: u32, seed: u64) -> (Fx, Vec<u32>) {
let dims = [2usize, 3, 4, 8];
let mut fx = Fx::new();
let mut ids = Vec::new();
for i in 0..n {
let dim = dims[(i as usize) % dims.len()];
ids.push(fx.add(
"Doc",
&format!("d{i}"),
vec![("emb", rand_emb(seed, i, dim))],
));
}
(fx, ids)
}
#[test]
fn vector_dim_reject_matches_unfiltered_and_oracle() {
const N: u32 = 500;
const SEED: u64 = 0xC0FF_EE00_D15C;
let def = vec_rule();
let (mut fx_on, ids) = seed_docs(N, SEED);
let mut eng_on = RuleEngine::new();
{
let mut g = fx_on.g();
eng_on.create_rule(def.clone(), &mut g).unwrap();
}
let on = prov_pairs(&eng_on, "vec");
assert!(!on.is_empty(), "seeded set must produce some edges");
let (mut fx_off, _) = seed_docs(N, SEED);
let mut eng_off = RuleEngine::new();
{
let mut g = fx_off.g();
with_vector_dim_reject(false, || {
eng_off.create_rule(def.clone(), &mut g).unwrap();
});
}
assert_eq!(on, prov_pairs(&eng_off, "vec"), "filter vs no-filter");
let mut brute = BTreeSet::new();
for &s in &ids {
for &d in &ids {
if s == d {
continue;
}
let skey = fx_on.ids.key_of(s).unwrap();
let dkey = fx_on.ids.key_of(d).unwrap();
let sget = |f: &str| fx_on.props.get(s, f).cloned();
let dget = |f: &str| fx_on.props.get(d, f).cloned();
if evaluate(
&def.predicate,
&NodeView {
key: skey,
props: &sget,
},
&NodeView {
key: dkey,
props: &dget,
},
)
.is_some()
{
brute.insert((s, d));
}
}
}
assert_eq!(on, brute, "filter vs brute-force evaluate");
}
#[test]
fn vector_dim_change_updates_cache_and_matches_fresh_build() {
let mut fx = Fx::new();
let a = fx.add("Doc", "a", vec![("emb", emb_val(&[1.0, 0.0]))]);
let b = fx.add("Doc", "b", vec![("emb", emb_val(&[1.0, 0.0]))]);
let c = fx.add("Doc", "c", vec![("emb", emb_val(&[1.0, 0.0, 0.0]))]);
let mut eng = RuleEngine::new();
{
let mut g = fx.g();
eng.create_rule(vec_rule(), &mut g).unwrap();
}
assert_eq!(eng.indexes["vec"].src_side.vec_dim(a), Some(2));
assert_eq!(eng.indexes["vec"].src_side.vec_dim(c), Some(3));
assert_eq!(prov_pairs(&eng, "vec"), BTreeSet::from([(a, b), (b, a)]));
let old = fx.props.get(b, "emb").cloned();
fx.props.set(b, "emb", emb_val(&[1.0, 0.0, 0.0]));
{
let mut g = fx.g();
eng.on_node_changed(b, Some(("emb", old)), &mut g);
}
assert_eq!(eng.indexes["vec"].src_side.vec_dim(b), Some(3));
assert_eq!(eng.indexes["vec"].dst_side.vec_dim(b), Some(3));
let after = prov_pairs(&eng, "vec");
assert_eq!(after, BTreeSet::from([(b, c), (c, b)]));
let mut fresh_fx = Fx::new();
let fa = fresh_fx.add("Doc", "a", vec![("emb", emb_val(&[1.0, 0.0]))]);
let fb = fresh_fx.add("Doc", "b", vec![("emb", emb_val(&[1.0, 0.0, 0.0]))]);
let fc = fresh_fx.add("Doc", "c", vec![("emb", emb_val(&[1.0, 0.0, 0.0]))]);
let mut fresh = RuleEngine::new();
{
let mut g = fresh_fx.g();
fresh.create_rule(vec_rule(), &mut g).unwrap();
}
assert_eq!(
prov_pairs(&fresh, "vec"),
BTreeSet::from([(fb, fc), (fc, fb)])
);
assert_eq!(fresh.indexes["vec"].src_side.vec_dim(fb), Some(3));
assert_eq!(fresh.indexes["vec"].src_side.vec_dim(fa), Some(2));
}
#[test]
fn streaming_topk_order_identity_property_test() {
fn reference_topk(rule: &RuleDef, k: u64, fx: &mut Fx) -> BTreeSet<(u32, u32)> {
let mut idx = RuleIndex::default();
for id in 0..fx.ids.len() as u32 {
let label_sym = match fx.labels.get(id as usize).copied() {
Some(s) if s != u32::MAX => s,
_ => continue,
};
index_node_for_rule(id, label_sym, rule, &mut idx, &fx.syms, &fx.props);
}
let src_sym = fx.syms.get(&rule.src_label);
let mut out = BTreeSet::new();
let ids_snap: Vec<u32> = (0..fx.ids.len() as u32).collect();
for id in ids_snap {
let label_sym = match fx.labels.get(id as usize).copied() {
Some(s) if s != u32::MAX => s,
_ => continue,
};
if src_sym != Some(label_sym) {
continue;
}
let g = GraphMut {
ids: &fx.ids,
syms: &mut fx.syms,
labels: &fx.labels,
props: &fx.props,
topo: &mut fx.topo,
edge_props: &mut fx.eprops,
};
let per_src = compute_desired(rule, &idx, id, true, &g);
let mut candidates: Vec<((u32, u32), f64)> = per_src.into_iter().collect();
candidates.sort_by(|&((_, da), sa), &((_, db), sb)| {
sb.total_cmp(&sa).then_with(|| {
let ka = fx.ids.key_of(da).unwrap_or("");
let kb = fx.ids.key_of(db).unwrap_or("");
ka.cmp(kb)
})
});
candidates.truncate(k as usize);
out.extend(candidates.into_iter().map(|(k, _)| k));
}
out
}
fn streaming_pairs(rule: RuleDef, fx: &mut Fx) -> BTreeSet<(u32, u32)> {
let name = rule.name.clone();
let mut eng = RuleEngine::new();
eng.create_rule(rule, &mut fx.g()).unwrap();
eng.provenance()
.get(&name)
.map(|s| s.iter().map(|&(_, a, b)| (a, b)).collect())
.unwrap_or_default()
}
for seed in [0u64, 1, 42, 0xDEAD_BEEF, 0x1234_5678, 99, 12_648_430, 7] {
for k in [1u64, 2, 3, 5] {
let rule = RuleDef {
name: "eq".into(),
src_label: "N".into(),
dst_label: "N".into(),
predicate: Predicate::FieldEqual { field: "k".into() },
edge_type: "EQ".into(),
weight_prop: None,
max_edges: Some(k),
approximate: false,
via_label: None,
via_edge: None,
via_dir: None,
};
let build = || {
let mut fx = Fx::new();
for i in 0..12u32 {
let h = mix64(seed ^ (i as u64 + 1));
let val = match h % 3 {
0 => "a",
1 => "b",
_ => "c",
};
fx.add(
"N",
&format!("n{i:02}"),
vec![("k", Value::Str(val.into()))],
);
}
fx
};
let expected = reference_topk(&rule, k, &mut build());
let actual = streaming_pairs(rule, &mut build());
assert_eq!(
expected, actual,
"FieldEqual seed={seed} k={k}: streaming top-k must match brute-force top-k"
);
}
}
for seed in [0u64, 1, 42, 7] {
for k in [1u64, 2, 4] {
let rule = RuleDef {
name: "nw".into(),
src_label: "S".into(),
dst_label: "D".into(),
predicate: Predicate::NumericWithin {
field: "v".into(),
tolerance: 10.0,
},
edge_type: "NEAR".into(),
weight_prop: Some("score".into()),
max_edges: Some(k),
approximate: false,
via_label: None,
via_edge: None,
via_dir: None,
};
let build = || {
let mut fx = Fx::new();
for i in 0..6u32 {
let h = mix64(seed ^ (i as u64 + 1));
let v = (h % 20) as f64;
fx.add("S", &format!("s{i}"), vec![("v", Value::Float(v))]);
}
for i in 0..8u32 {
let h = mix64(seed ^ (i as u64 + 101));
let v = (h % 20) as f64;
fx.add("D", &format!("d{i}"), vec![("v", Value::Float(v))]);
}
fx
};
let expected = reference_topk(&rule, k, &mut build());
let actual = streaming_pairs(rule, &mut build());
assert_eq!(
expected, actual,
"NumericWithin seed={seed} k={k}: streaming top-k must match brute-force top-k"
);
}
}
for seed in [0u64, 1, 42, 7] {
for k in [1u64, 2] {
let rule = RuleDef {
name: "fk".into(),
src_label: "T".into(),
dst_label: "C".into(),
predicate: Predicate::KeyMatch {
field: "cid".into(),
},
edge_type: "AT".into(),
weight_prop: None,
max_edges: Some(k),
approximate: false,
via_label: None,
via_edge: None,
via_dir: None,
};
let build = || {
let mut fx = Fx::new();
for i in 0..4u32 {
fx.add("C", &format!("c{i}"), vec![]);
}
for i in 0..8u32 {
let h = mix64(seed ^ (i as u64 + 1));
let cid = format!("c{}", h % 4);
fx.add("T", &format!("t{i}"), vec![("cid", Value::Str(cid))]);
}
fx
};
let expected = reference_topk(&rule, k, &mut build());
let actual = streaming_pairs(rule, &mut build());
assert_eq!(
expected, actual,
"KeyMatch seed={seed} k={k}: streaming top-k must match brute-force top-k"
);
}
}
{
let cluster_a: &[(&str, f64, f64)] = &[
("va0", 1.0_f64, 0.0_f64),
("va1", 0.98_f64, 0.199_f64), ("va2", 0.97_f64, 0.243_f64), ];
let cluster_b: &[(&str, f64, f64)] = &[
("vb0", 0.0_f64, 1.0_f64),
("vb1", 0.1_f64, 0.995_f64),
("vb2", 0.05_f64, 0.999_f64),
];
for k in [1u64, 2] {
let rule = RuleDef {
name: "vsim".into(),
src_label: "V".into(),
dst_label: "V".into(),
predicate: Predicate::VectorSimilar {
field: "emb".into(),
min: 0.9,
},
edge_type: "VSIM".into(),
weight_prop: Some("score".into()),
max_edges: Some(k),
approximate: false,
via_label: None,
via_edge: None,
via_dir: None,
};
let build = || {
let mut fx = Fx::new();
let mut add_v = |key: &str, x: f64, y: f64| {
let norm = (x * x + y * y).sqrt();
let v = Value::List(vec![Value::Float(x / norm), Value::Float(y / norm)]);
fx.add("V", key, vec![("emb", v)]);
};
for &(k, x, y) in cluster_a.iter().chain(cluster_b.iter()) {
add_v(k, x, y);
}
fx
};
let expected = reference_topk(&rule, k, &mut build());
let actual = streaming_pairs(rule, &mut build());
assert_eq!(
expected, actual,
"VectorSimilar/ScanAll k={k}: streaming top-k must match brute-force top-k"
);
}
}
}
#[test]
#[ignore]
fn streaming_peak_transient_bound() {
use std::sync::{
atomic::{AtomicBool, AtomicU64, Ordering},
Arc,
};
fn peak_rss_during<F: FnOnce()>(f: F) -> u64 {
let done = Arc::new(AtomicBool::new(false));
let peak = Arc::new(AtomicU64::new(0));
let done2 = done.clone();
let peak2 = peak.clone();
let pid = std::process::id().to_string();
let handle = std::thread::spawn(move || {
while !done2.load(Ordering::Relaxed) {
let rss = std::process::Command::new("ps")
.args(["-o", "rss=", "-p", &pid])
.output()
.ok()
.and_then(|o| String::from_utf8(o.stdout).ok())
.and_then(|s| s.trim().parse::<u64>().ok())
.unwrap_or(0)
* 1024;
peak2.fetch_max(rss, Ordering::Relaxed);
std::thread::sleep(std::time::Duration::from_millis(1));
}
});
f();
done.store(true, Ordering::Relaxed);
let _ = handle.join();
peak.load(Ordering::Relaxed)
}
let mut fx = Fx::new();
for i in 0..500u32 {
fx.add(
"Talent",
&format!("t{i}"),
vec![("k", Value::Str("same".into()))],
);
}
for i in 0..500u32 {
fx.add(
"Company",
&format!("c{i}"),
vec![("k", Value::Str("same".into()))],
);
}
let rule = RuleDef {
name: "eq_tc".into(),
src_label: "Talent".into(),
dst_label: "Company".into(),
predicate: Predicate::FieldEqual { field: "k".into() },
edge_type: "EQ".into(),
weight_prop: None,
max_edges: Some(2), approximate: false,
via_label: None,
via_edge: None,
via_dir: None,
};
let pid = std::process::id().to_string();
let baseline = std::process::Command::new("ps")
.args(["-o", "rss=", "-p", &pid])
.output()
.ok()
.and_then(|o| String::from_utf8(o.stdout).ok())
.and_then(|s| s.trim().parse::<u64>().ok())
.unwrap_or(0)
* 1024;
let mut eng = RuleEngine::new();
let peak = peak_rss_during(|| {
eng.create_rule(rule, &mut fx.g()).unwrap();
});
let peak_delta = peak.saturating_sub(baseline);
assert!(
peak_delta < 3 * 1024 * 1024,
"peak transient delta {} bytes ({} KiB) exceeded 3 MiB; \
streaming path may be building the full pairs map",
peak_delta,
peak_delta / 1024
);
assert_eq!(eng.provenance()["eq_tc"].len(), 1_000); assert!(!eng.is_tripped("eq_tc")); eprintln!(
"streaming_peak_transient_bound: baseline={baseline} peak={peak} \
delta={peak_delta} bytes ({} KiB)",
peak_delta / 1024
);
}
fn near_threshold_pair(dim: usize, min: f64) -> (Vec<f64>, Vec<f64>) {
let cos_target = min + 1e-6; let sin_small = (1.0 - cos_target * cos_target).sqrt();
let mut a = vec![0.0f64; dim];
a[0] = 1.0;
let mut b = vec![0.0f64; dim];
b[0] = cos_target;
if dim > 1 {
b[1] = sin_small;
}
(a, b)
}
fn emb_val2(xs: &[f64]) -> Value {
Value::List(xs.iter().copied().map(Value::Float).collect())
}
fn make_early_exit_fixture(seed: u64, min: f64) -> (Fx, Vec<u32>, usize, usize) {
let dims = [2usize, 4, 8, 16];
let n = 100u32;
let mut fx = Fx::new();
let mut ids = Vec::new();
for i in 0..n {
let dim = dims[(i as usize) % dims.len()];
let emb = rand_emb(seed, i, dim);
ids.push(fx.add("Doc", &format!("d{i}"), vec![("emb", emb)]));
}
let (va, vb) = near_threshold_pair(8, min);
let nt_a = fx.add("Doc", "nt_a", vec![("emb", emb_val2(&va))]);
let nt_b = fx.add("Doc", "nt_b", vec![("emb", emb_val2(&vb))]);
ids.push(nt_a);
ids.push(nt_b);
(fx, ids, nt_a as usize, nt_b as usize)
}
#[test]
fn vector_early_exit_identity_proof() {
const SEED: u64 = 0xEA_4E_5A;
const MIN: f64 = 0.85;
let def = RuleDef {
name: "vec".into(),
src_label: "Doc".into(),
dst_label: "Doc".into(),
predicate: Predicate::VectorSimilar {
field: "emb".into(),
min: MIN,
},
edge_type: "SIM".into(),
weight_prop: Some("score".into()),
max_edges: None,
approximate: false,
via_label: None,
via_edge: None,
via_dir: None,
};
let (mut fx_on, ids, nt_a, nt_b) = make_early_exit_fixture(SEED, MIN);
let (mut fx_off, _, _, _) = make_early_exit_fixture(SEED, MIN);
let (fx_oracle, _, _, _) = make_early_exit_fixture(SEED, MIN);
let nt_a = nt_a as u32;
let nt_b = nt_b as u32;
let mut eng_on = RuleEngine::new();
{
let mut g = fx_on.g();
eng_on.create_rule(def.clone(), &mut g).unwrap();
}
let edges_on = prov_pairs(&eng_on, "vec");
assert!(!edges_on.is_empty(), "should produce some edges");
assert!(
edges_on.contains(&(nt_a, nt_b)),
"near-threshold pair nt_a→nt_b must match with early-exit ON"
);
assert!(
edges_on.contains(&(nt_b, nt_a)),
"near-threshold pair nt_b→nt_a must match with early-exit ON"
);
let mut eng_off = RuleEngine::new();
{
let mut g = fx_off.g();
with_vector_early_exit(false, || {
eng_off.create_rule(def.clone(), &mut g).unwrap();
});
}
let edges_off = prov_pairs(&eng_off, "vec");
assert_eq!(
edges_on, edges_off,
"early-exit ON vs OFF must produce identical edges"
);
let mut oracle = BTreeSet::new();
for &s in &ids {
for &d in &ids {
if s == d {
continue;
}
let skey = fx_oracle.ids.key_of(s).unwrap();
let dkey = fx_oracle.ids.key_of(d).unwrap();
let sg = |f: &str| fx_oracle.props.get(s, f).cloned();
let dg = |f: &str| fx_oracle.props.get(d, f).cloned();
if evaluate(
&def.predicate,
&NodeView {
key: skey,
props: &sg,
},
&NodeView {
key: dkey,
props: &dg,
},
)
.is_some()
{
oracle.insert((s, d));
}
}
}
assert_eq!(
edges_on, oracle,
"early-exit ON vs brute-force oracle must be identical"
);
}
#[test]
fn vector_early_exit_checkpoint_coherence() {
let mut fx = Fx::new();
let a = fx.add("Doc", "a", vec![("emb", emb_val(&[1.0, 0.0, 0.0, 0.0]))]);
let b = fx.add("Doc", "b", vec![("emb", emb_val(&[1.0, 0.0, 0.0, 0.0]))]);
let c = fx.add(
"Doc",
"c",
vec![("emb", emb_val(&[1.0, 0.0, 0.0, 0.0, 0.0, 0.0]))],
);
let def = RuleDef {
name: "vec".into(),
src_label: "Doc".into(),
dst_label: "Doc".into(),
predicate: Predicate::VectorSimilar {
field: "emb".into(),
min: 0.9,
},
edge_type: "SIM".into(),
weight_prop: None,
max_edges: None,
approximate: false,
via_label: None,
via_edge: None,
via_dir: None,
};
let mut eng = RuleEngine::new();
{
let mut g = fx.g();
eng.create_rule(def.clone(), &mut g).unwrap();
}
assert!(
eng.indexes["vec"].src_side.vec_ckpts(a).is_some(),
"a must have src checkpoints"
);
assert!(
eng.indexes["vec"].dst_side.vec_ckpts(b).is_some(),
"b must have dst checkpoints"
);
assert!(
eng.indexes["vec"].src_side.vec_ckpts(c).is_some(),
"c must have src checkpoints (dim=6)"
);
let ckpts_a = *eng.indexes["vec"].src_side.vec_ckpts(a).unwrap();
let norm_a = eng.indexes["vec"].src_side.vec_meta(a).unwrap().1;
assert!(
(ckpts_a[0] - norm_a).abs() < 1e-12,
"ckpts[0] must equal the full L2 norm"
);
assert_eq!(prov_pairs(&eng, "vec"), BTreeSet::from([(a, b), (b, a)]));
let old_b = fx.props.get(b, "emb").cloned();
fx.props
.set(b, "emb", emb_val(&[1.0, 0.0, 0.0, 0.0, 0.0, 0.0]));
{
let mut g = fx.g();
eng.on_node_changed(b, Some(("emb", old_b)), &mut g);
}
assert_eq!(eng.indexes["vec"].src_side.vec_dim(b), Some(6));
assert_eq!(eng.indexes["vec"].dst_side.vec_dim(b), Some(6));
assert!(eng.indexes["vec"].src_side.vec_ckpts(b).is_some());
assert_eq!(prov_pairs(&eng, "vec"), BTreeSet::from([(b, c), (c, b)]));
let wrong_live = vec![2.0f64, 0.0, 0.0, 0.0, 0.0, 0.0]; let gate_result = eng.indexes["vec"].src_side.fresh_ckpts_for(b, &wrong_live);
assert!(
gate_result.is_none(),
"freshness gate must reject a mismatched-norm live vector"
);
let correct_live = vec![1.0f64, 0.0, 0.0, 0.0, 0.0, 0.0];
let gate_result = eng.indexes["vec"]
.src_side
.fresh_ckpts_for(b, &correct_live);
assert!(
gate_result.is_some(),
"freshness gate must accept the matching live vector"
);
}
#[test]
fn vector_early_exit_razor_dim1536() {
const MIN: f64 = 0.85;
const DIM: usize = 1536;
let target = MIN + 5e-13;
let inv_sqrt = 1.0 / (DIM as f64).sqrt();
let a: Vec<f64> = vec![inv_sqrt; DIM];
let perp_scale = (1.0 - target * target).sqrt() / (2.0f64).sqrt();
let mut b: Vec<f64> = vec![target * inv_sqrt; DIM];
b[0] += perp_scale;
b[1] -= perp_scale;
let def = RuleDef {
name: "razor".into(),
src_label: "Doc".into(),
dst_label: "Doc".into(),
predicate: Predicate::VectorSimilar {
field: "emb".into(),
min: MIN,
},
edge_type: "SIM".into(),
weight_prop: None,
max_edges: None,
approximate: false,
via_label: None,
via_edge: None,
via_dir: None,
};
let build_fx = || {
let mut fx = Fx::new();
let na = fx.add("Doc", "razor_a", vec![("emb", emb_val2(&a))]);
let nb = fx.add("Doc", "razor_b", vec![("emb", emb_val2(&b))]);
(fx, na, nb)
};
let (mut fx_on, na, nb) = build_fx();
let (mut fx_off, _, _) = build_fx();
let (fx_oracle, _, _) = build_fx();
let mut eng_on = RuleEngine::new();
{
let mut g = fx_on.g();
eng_on.create_rule(def.clone(), &mut g).unwrap();
}
let edges_on = prov_pairs(&eng_on, "razor");
assert!(
edges_on.contains(&(na, nb)),
"razor pair razor_a→razor_b must be present with early-exit ON (cos={target:.15}, min={MIN})"
);
assert!(
edges_on.contains(&(nb, na)),
"razor pair razor_b→razor_a must be present with early-exit ON"
);
let mut eng_off = RuleEngine::new();
{
let mut g = fx_off.g();
with_vector_early_exit(false, || {
eng_off.create_rule(def.clone(), &mut g).unwrap();
});
}
let edges_off = prov_pairs(&eng_off, "razor");
assert_eq!(
edges_on, edges_off,
"razor dim=1536: early-exit ON vs OFF must produce identical edges"
);
let ids = [na, nb];
let mut oracle = BTreeSet::new();
for &s in &ids {
for &d in &ids {
if s == d {
continue;
}
let skey = fx_oracle.ids.key_of(s).unwrap();
let dkey = fx_oracle.ids.key_of(d).unwrap();
let sg = |f: &str| fx_oracle.props.get(s, f).cloned();
let dg = |f: &str| fx_oracle.props.get(d, f).cloned();
if evaluate(
&def.predicate,
&NodeView {
key: skey,
props: &sg,
},
&NodeView {
key: dkey,
props: &dg,
},
)
.is_some()
{
oracle.insert((s, d));
}
}
}
assert_eq!(
edges_on, oracle,
"razor dim=1536: early-exit ON vs brute-force oracle must be identical"
);
}
}