From b463f740c5f71e34fedb4958e3e82926e455fe12 Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Mon, 5 Jan 2026 20:36:08 -0800 Subject: [PATCH 1/6] optional union find notification --- core-relations/src/lib.rs | 2 +- core-relations/src/uf/mod.rs | 83 +++++++++++++++++++++++++++++++--- core-relations/src/uf/tests.rs | 48 +++++++++++++++++++- 3 files changed, 124 insertions(+), 9 deletions(-) diff --git a/core-relations/src/lib.rs b/core-relations/src/lib.rs index 724d7798f..009eea344 100644 --- a/core-relations/src/lib.rs +++ b/core-relations/src/lib.rs @@ -42,7 +42,7 @@ pub use table_spec::{ ColumnId, Constraint, MutationBuffer, Offset, Rebuilder, Row, Table, TableChange, TableSpec, TableVersion, WrappedTable, }; -pub use uf::{DisplacedTable, DisplacedTableWithProvenance, ProofReason, ProofStep}; +pub use uf::{DisplacedTable, DisplacedTableWithProvenance, LeaderChange, ProofReason, ProofStep}; use egglog_numeric_id as numeric_id; use egglog_union_find as union_find; diff --git a/core-relations/src/uf/mod.rs b/core-relations/src/uf/mod.rs index 7e8095044..884748afc 100644 --- a/core-relations/src/uf/mod.rs +++ b/core-relations/src/uf/mod.rs @@ -29,6 +29,31 @@ mod tests; type UnionFind = crate::union_find::UnionFind; +/// A callback that runs every time a leader change takes effect. See the documentation for +/// [`LeaderChange`] for the information that is provided. +type LeaderChangeCallback = Arc; + +/// Details for a leader change caused by a union. +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +pub struct LeaderChange { + /// The lhs value provided to the write. + pub write_lhs: Value, + /// The leader of the lhs equivalence class before the union. + pub lhs_leader: Value, + /// The rhs value provided to the write. + pub write_rhs: Value, + /// The leader of the rhs equivalence class before the union. + pub rhs_leader: Value, + /// The timestamp associated with the write that triggered the union. + pub ts: Value, +} + +impl LeaderChange { + pub fn new_leader(&self) -> Value { + std::cmp::min(self.lhs_leader, self.rhs_leader) + } +} + /// A special table backed by a union-find used to efficiently implement /// egglog-style canonicaliztion. /// @@ -59,6 +84,8 @@ pub struct DisplacedTable { changed: bool, lookup_table: HashMap, buffered_writes: Arc>, + /// Stored as Arc so DisplacedTable cloning preserves the callback. + on_leader_change: Option, } struct Canonicalizer<'a> { @@ -205,6 +232,7 @@ impl Default for DisplacedTable { changed: false, lookup_table: HashMap::default(), buffered_writes: Arc::new(SegQueue::new()), + on_leader_change: None, } } } @@ -217,6 +245,7 @@ impl Clone for DisplacedTable { changed: self.changed, lookup_table: self.lookup_table.clone(), buffered_writes: Default::default(), + on_leader_change: self.on_leader_change.clone(), } } } @@ -445,10 +474,10 @@ impl Table for DisplacedTable { }) } - fn merge(&mut self, _: &mut ExecutionState) -> TableChange { + fn merge(&mut self, exec_state: &mut ExecutionState) -> TableChange { while let Some(rowbuf) = self.buffered_writes.pop() { for row in rowbuf.iter() { - self.changed |= self.insert_impl(row).is_some(); + self.changed |= self.insert_impl(row, exec_state).is_some(); } } let changed = mem::take(&mut self.changed); @@ -462,6 +491,28 @@ impl Table for DisplacedTable { } impl DisplacedTable { + /// Construct with a leader-change callback. + pub fn with_leader_change_callback(callback: F) -> Self + where + F: Fn(&mut ExecutionState, LeaderChange) + Send + Sync + 'static, + { + Self { + on_leader_change: Some(Arc::new(callback)), + ..Self::default() + } + } + + pub fn set_leader_change_callback(&mut self, callback: F) + where + F: Fn(&mut ExecutionState, LeaderChange) + Send + Sync + 'static, + { + self.on_leader_change = Some(Arc::new(callback)); + } + + pub fn clear_leader_change_callback(&mut self) { + self.on_leader_change = None; + } + pub fn underlying_uf(&self) -> &UnionFind { &self.uf } @@ -488,9 +539,15 @@ impl DisplacedTable { let vals = self.expand(row); eval_constraint(&vals, constraint) } - fn insert_impl(&mut self, row: &[Value]) -> Option<(Value, Value)> { + fn insert_impl( + &mut self, + row: &[Value], + exec_state: &mut ExecutionState, + ) -> Option<(Value, Value)> { assert_eq!(row.len(), 3, "attempt to insert a row with the wrong arity"); - if self.uf.find(row[0]) == self.uf.find(row[1]) { + let lhs_leader = self.uf.find(row[0]); + let rhs_leader = self.uf.find(row[1]); + if lhs_leader == rhs_leader { return None; } let (parent, child) = self.uf.union(row[0], row[1]); @@ -499,6 +556,18 @@ impl DisplacedTable { let _ = self.uf.find(parent); let _ = self.uf.find(child); let ts = row[2]; + if let Some(callback) = &self.on_leader_change { + callback( + exec_state, + LeaderChange { + write_lhs: row[0], + lhs_leader, + write_rhs: row[1], + rhs_leader, + ts, + }, + ); + } if let Some((_, highest)) = self.displaced.last() { assert!( *highest <= ts, @@ -723,11 +792,11 @@ impl DisplacedTableWithProvenance { .or_insert_with(|| self.proof_graph.add_node(val)) } - fn insert_impl(&mut self, row: &[Value]) { + fn insert_impl(&mut self, row: &[Value], exec_state: &mut ExecutionState) { let [a, b, ts, reason] = row else { panic!("attempt to insert a row with the wrong arity ({row:?})"); }; - match self.base.insert_impl(&[*a, *b, *ts]) { + match self.base.insert_impl(&[*a, *b, *ts], exec_state) { Some((parent, child)) => { self.displaced.push((child, parent)); self.context @@ -810,7 +879,7 @@ impl Table for DisplacedTableWithProvenance { fn merge(&mut self, exec_state: &mut ExecutionState) -> TableChange { while let Some(rowbuf) = self.buffered_writes.pop() { for row in rowbuf.iter() { - self.insert_impl(row); + self.insert_impl(row, exec_state); } } diff --git a/core-relations/src/uf/tests.rs b/core-relations/src/uf/tests.rs index e1d4f29c5..6bac1855c 100644 --- a/core-relations/src/uf/tests.rs +++ b/core-relations/src/uf/tests.rs @@ -1,3 +1,5 @@ +use std::sync::{Arc, Mutex}; + use crate::numeric_id::NumericId; use crate::{ @@ -7,7 +9,7 @@ use crate::{ uf::ProofReason, }; -use super::DisplacedTable; +use super::{DisplacedTable, LeaderChange}; fn v(x: usize) -> Value { Value::from_usize(x) @@ -98,3 +100,47 @@ fn displaced_proof() { ] ) } + +#[test] +fn displaced_leader_change_callback() { + empty_execution_state!(e); + let changes: Arc>> = Arc::new(Mutex::new(Vec::new())); + let changes_ref = Arc::clone(&changes); + let mut d = DisplacedTable::with_leader_change_callback(move |_, change| { + changes_ref.lock().unwrap().push(change); + }); + { + let mut buf = d.new_buffer(); + buf.stage_insert(&[v(5), v(3), v(0)]); + buf.stage_insert(&[v(5), v(3), v(1)]); + } + d.merge(&mut e); + + { + let changes = changes.lock().unwrap(); + assert_eq!(changes.len(), 1); + let change = changes[0]; + assert_eq!(change.write_lhs, v(5)); + assert_eq!(change.lhs_leader, v(5)); + assert_eq!(change.write_rhs, v(3)); + assert_eq!(change.rhs_leader, v(3)); + assert_eq!(change.ts, v(0)); + assert_eq!(change.new_leader(), v(3)); + } + + { + let mut buf = d.new_buffer(); + buf.stage_insert(&[v(5), v(2), v(2)]); + } + d.merge(&mut e); + + let changes = changes.lock().unwrap(); + assert_eq!(changes.len(), 2); + let change = changes[1]; + assert_eq!(change.write_lhs, v(5)); + assert_eq!(change.lhs_leader, v(3)); + assert_eq!(change.write_rhs, v(2)); + assert_eq!(change.rhs_leader, v(2)); + assert_eq!(change.ts, v(2)); + assert_eq!(change.new_leader(), v(2)); +} From d89765884b9f677610edb73d2f9770651aae23c9 Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Mon, 5 Jan 2026 21:08:59 -0800 Subject: [PATCH 2/6] initial UF-specific function info --- egglog-bridge/src/lib.rs | 81 +++++++++++++++++++++++++++++++++++++- egglog-bridge/src/tests.rs | 72 ++++++++++++++++++++++++++++++++- 2 files changed, 150 insertions(+), 3 deletions(-) diff --git a/egglog-bridge/src/lib.rs b/egglog-bridge/src/lib.rs index afab12fd3..5a3d0993e 100644 --- a/egglog-bridge/src/lib.rs +++ b/egglog-bridge/src/lib.rs @@ -20,8 +20,8 @@ use std::{ use crate::core_relations::{ BaseValue, BaseValueId, BaseValues, ColumnId, Constraint, ContainerValue, ContainerValues, CounterId, Database, DisplacedTable, DisplacedTableWithProvenance, ExecutionState, - ExternalFunction, ExternalFunctionId, MergeVal, Offset, PlanStrategy, SortedWritesTable, - TableId, TaggedRowBuffer, Value, WrappedTable, + ExternalFunction, ExternalFunctionId, LeaderChange, MergeVal, Offset, PlanStrategy, + SortedWritesTable, TableId, TaggedRowBuffer, Value, WrappedTable, }; use crate::numeric_id::{DenseIdMap, DenseIdMapWithReuse, IdVec, NumericId, define_id}; use egglog_core_relations as core_relations; @@ -163,6 +163,15 @@ pub struct FunctionConfig { pub can_subsume: bool, } +/// Configuration for a UF-backed function. +pub struct UfFunctionConfig { + pub name: String, + pub on_leader_change: + Option>, + pub read_deps: Vec, + pub write_deps: Vec, +} + impl EGraph { /// Create a new EGraph with tracing (aka 'proofs') enabled. /// @@ -768,6 +777,7 @@ impl EGraph { default_val: default, can_subsume, name, + kind: FunctionKind::Table, }); debug_assert_eq!(res, next_func_id); let incremental_rebuild_rules = self.incremental_rebuild_rules(res, &schema); @@ -778,6 +788,41 @@ impl EGraph { res } + /// Register a UF-backed function in this EGraph. + pub fn add_uf_function(&mut self, config: UfFunctionConfig) -> FunctionId { + if self.tracing { + panic!("UF functions are not supported when tracing is enabled"); + } + let UfFunctionConfig { + name, + on_leader_change, + read_deps, + write_deps, + } = config; + let mut table = DisplacedTable::default(); + if let Some(callback) = on_leader_change { + table.set_leader_change_callback(move |state, change| (callback)(state, change)); + } + let name: Arc = name.into(); + let table_id = self.db.add_table_named( + table, + name.clone(), + read_deps.iter().copied(), + write_deps.iter().copied(), + ); + let schema = vec![ColumnTy::Id, ColumnTy::Id]; + self.funcs.push(FunctionInfo { + table: table_id, + schema, + incremental_rebuild_rules: Vec::new(), + nonincremental_rebuild_rule: RuleId::new(!0), + default_val: DefaultVal::Fail, + can_subsume: false, + name, + kind: FunctionKind::Uf, + }) + } + /// Run the given rules, returning whether the database changed. /// /// If the given rules are malformed, this method can return an error. @@ -825,6 +870,9 @@ impl EGraph { // The UF implementation supports "native" rebuilding. let mut tables = Vec::with_capacity(self.funcs.next_id().index()); for (_, func) in self.funcs.iter() { + if func.is_uf() { + continue; + } tables.push(func.table); } loop { @@ -884,6 +932,9 @@ impl EGraph { self.inc_ts(); let ts = self.next_ts(); for (_, info) in self.funcs.iter_mut() { + if info.is_uf() { + continue; + } let last_rebuilt_at = self.rules[info.nonincremental_rebuild_rule].last_run_at; let table_size = self.db.estimate_size(info.table, None); let uf_size = self.db.estimate_size( @@ -964,6 +1015,9 @@ impl EGraph { // First, figure out which functions will be rebuilt nonincrementally, // vs. incrementally. Group them together. for (func, info) in self.funcs.iter_mut() { + if info.is_uf() { + continue; + } let last_rebuilt_at = self.rules[info.nonincremental_rebuild_rule].last_run_at; let table_size = self.db.estimate_size(info.table, None); let uf_size = self.db.estimate_size( @@ -1154,6 +1208,12 @@ struct CachedPlanInfo { atom_mapping: Vec, } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum FunctionKind { + Table, + Uf, +} + #[derive(Clone)] struct FunctionInfo { table: TableId, @@ -1163,12 +1223,17 @@ struct FunctionInfo { default_val: DefaultVal, can_subsume: bool, name: Arc, + kind: FunctionKind, } impl FunctionInfo { fn ret_ty(&self) -> ColumnTy { self.schema.last().copied().unwrap() } + + fn is_uf(&self) -> bool { + matches!(self.kind, FunctionKind::Uf) + } } /// How defaults are computed for the given function. @@ -1218,6 +1283,10 @@ impl MergeFn { .for_each(|arg| arg.fill_deps(egraph, read_deps, write_deps)); } Function(func, args) => { + assert!( + !egraph.funcs[*func].is_uf(), + "Merge functions cannot call UF-backed functions" + ); read_deps.insert(egraph.funcs[*func].table); write_deps.insert(egraph.funcs[*func].table); args.iter() @@ -1318,6 +1387,10 @@ impl MergeFn { }, MergeFn::Function(func, args) => { let func_info = &egraph.funcs[*func]; + assert!( + !func_info.is_uf(), + "Merge functions cannot call UF-backed functions" + ); assert_eq!( func_info.schema.len(), args.len() + 1, @@ -1463,6 +1536,10 @@ impl TableAction { assert!(!egraph.tracing, "proofs not supported yet"); let func_info = &egraph.funcs[func]; + assert!( + !func_info.is_uf(), + "TableAction does not support UF-backed functions" + ); TableAction { table: func_info.table, table_math: SchemaMath { diff --git a/egglog-bridge/src/tests.rs b/egglog-bridge/src/tests.rs index d670d7f63..f6454ca9c 100644 --- a/egglog-bridge/src/tests.rs +++ b/egglog-bridge/src/tests.rs @@ -19,7 +19,7 @@ use num_rational::Rational64; use crate::{ ColumnTy, DefaultVal, EGraph, FunctionConfig, FunctionId, MergeFn, ProofStore, QueryEntry, - add_expressions, define_rule, + UfFunctionConfig, add_expressions, define_rule, }; /// Run a simple associativity/commutativity test. In addition to testing that the rules properly @@ -117,6 +117,76 @@ fn ac_test(tracing: bool, can_subsume: bool) { } } +#[test] +fn uf_function_callback_inserts() { + let mut egraph = EGraph::default(); + let log_table = egraph.add_table(FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Id], + default: DefaultVal::Fail, + merge: MergeFn::AssertEq, + name: "log".into(), + can_subsume: false, + }); + let log_table_id = egraph.funcs[log_table].table; + let uf_func = egraph.add_uf_function(UfFunctionConfig { + name: "uf_cb".into(), + on_leader_change: Some(Box::new(move |state, change| { + state.stage_insert( + log_table_id, + &[change.write_lhs, change.new_leader(), change.ts], + ); + })), + read_deps: Vec::new(), + write_deps: vec![log_table_id], + }); + + let lhs = Value::from_usize(5); + let rhs = Value::from_usize(3); + egraph.add_values([(uf_func, vec![lhs, rhs])]); + + let leader = egraph.lookup_id(uf_func, &[lhs]).unwrap(); + assert_eq!(leader, rhs); + + let mut rows = Vec::new(); + egraph.for_each(log_table, |row| rows.push(row.vals.to_vec())); + assert_eq!(rows, vec![vec![lhs, rhs]]); +} + +#[test] +fn uf_function_disallowed_in_merge() { + let mut egraph = EGraph::default(); + let uf_func = egraph.add_uf_function(UfFunctionConfig { + name: "uf_func".into(), + on_leader_change: None, + read_deps: Vec::new(), + write_deps: Vec::new(), + }); + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + egraph.add_table(FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Id], + default: DefaultVal::Fail, + merge: MergeFn::Function(uf_func, vec![MergeFn::New]), + name: "bad_merge".into(), + can_subsume: false, + }); + })); + assert!(result.is_err()); +} + +#[test] +fn uf_function_disallowed_with_tracing() { + let mut egraph = EGraph::with_tracing(); + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + egraph.add_uf_function(UfFunctionConfig { + name: "uf_tracing".into(), + on_leader_change: None, + read_deps: Vec::new(), + write_deps: Vec::new(), + }); + })); + assert!(result.is_err()); +} + #[test] fn ac_tracing_subsume() { ac_test(true, true); From b3fed831cc87d1cbd5c2fc5b6305f18a5a463f93 Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Mon, 5 Jan 2026 22:30:58 -0800 Subject: [PATCH 3/6] Expose 'UF-backed' functions from egglog-bridge --- egglog-bridge/examples/ac.rs | 42 +++-- egglog-bridge/examples/ac_tracing.rs | 42 +++-- egglog-bridge/examples/math.rs | 227 ++++++++++++++++----------- egglog-bridge/src/lib.rs | 107 +++++++------ egglog-bridge/src/tests.rs | 198 ++++++++++++----------- src/lib.rs | 87 +++++----- src/prelude.rs | 9 +- src/scheduler.rs | 78 +++++---- 8 files changed, 457 insertions(+), 333 deletions(-) diff --git a/egglog-bridge/examples/ac.rs b/egglog-bridge/examples/ac.rs index e6a83b16e..688e48fcd 100644 --- a/egglog-bridge/examples/ac.rs +++ b/egglog-bridge/examples/ac.rs @@ -1,10 +1,16 @@ -use egglog_bridge::{ColumnTy, DefaultVal, EGraph, FunctionConfig, MergeFn, define_rule}; +use egglog_bridge::{ + ColumnTy, DefaultVal, EGraph, FunctionConfig, FunctionId, MergeFn, define_rule, +}; use mimalloc::MiMalloc; #[global_allocator] static GLOBAL: MiMalloc = MiMalloc; +fn add_table(egraph: &mut EGraph, config: FunctionConfig) -> FunctionId { + egraph.add_table(config).unwrap() +} + #[allow(clippy::disallowed_macros)] fn main() { const N: usize = 13; @@ -14,20 +20,26 @@ fn main() { let start = web_time::Instant::now(); let mut egraph = EGraph::default(); let int_base = egraph.base_values_mut().register_type::(); - let num_table = egraph.add_table(FunctionConfig { - schema: vec![ColumnTy::Base(int_base), ColumnTy::Id], - default: DefaultVal::FreshId, - merge: MergeFn::UnionId, - name: "num".into(), - can_subsume: false, - }); - let add_table = egraph.add_table(FunctionConfig { - schema: vec![ColumnTy::Id; 3], - default: DefaultVal::FreshId, - merge: MergeFn::UnionId, - name: "add".into(), - can_subsume: false, - }); + let num_table = add_table( + &mut egraph, + FunctionConfig { + schema: vec![ColumnTy::Base(int_base), ColumnTy::Id], + default: DefaultVal::FreshId, + merge: MergeFn::UnionId, + name: "num".into(), + can_subsume: false, + }, + ); + let add_table = add_table( + &mut egraph, + FunctionConfig { + schema: vec![ColumnTy::Id; 3], + default: DefaultVal::FreshId, + merge: MergeFn::UnionId, + name: "add".into(), + can_subsume: false, + }, + ); let add_comm = define_rule! { [egraph] ((-> (add_table x y) id)) diff --git a/egglog-bridge/examples/ac_tracing.rs b/egglog-bridge/examples/ac_tracing.rs index c15ba0ab0..68f11203c 100644 --- a/egglog-bridge/examples/ac_tracing.rs +++ b/egglog-bridge/examples/ac_tracing.rs @@ -1,26 +1,38 @@ use std::mem; -use egglog_bridge::{ColumnTy, DefaultVal, EGraph, FunctionConfig, MergeFn, define_rule}; +use egglog_bridge::{ + ColumnTy, DefaultVal, EGraph, FunctionConfig, FunctionId, MergeFn, define_rule, +}; + +fn add_table(egraph: &mut EGraph, config: FunctionConfig) -> FunctionId { + egraph.add_table(config).unwrap() +} fn main() { const N: usize = 12; env_logger::init(); let mut egraph = EGraph::with_tracing(); let int_base = egraph.base_values_mut().register_type::(); - let num_table = egraph.add_table(FunctionConfig { - schema: vec![ColumnTy::Base(int_base), ColumnTy::Id], - default: DefaultVal::FreshId, - merge: MergeFn::UnionId, - name: "num".into(), - can_subsume: false, - }); - let add_table = egraph.add_table(FunctionConfig { - schema: vec![ColumnTy::Id; 3], - default: DefaultVal::FreshId, - merge: MergeFn::UnionId, - name: "add".into(), - can_subsume: false, - }); + let num_table = add_table( + &mut egraph, + FunctionConfig { + schema: vec![ColumnTy::Base(int_base), ColumnTy::Id], + default: DefaultVal::FreshId, + merge: MergeFn::UnionId, + name: "num".into(), + can_subsume: false, + }, + ); + let add_table = add_table( + &mut egraph, + FunctionConfig { + schema: vec![ColumnTy::Id; 3], + default: DefaultVal::FreshId, + merge: MergeFn::UnionId, + name: "add".into(), + can_subsume: false, + }, + ); let add_comm = define_rule! { [egraph] ((-> (add_table x y) id)) diff --git a/egglog-bridge/examples/math.rs b/egglog-bridge/examples/math.rs index 8c0eaa13b..49da49c63 100644 --- a/egglog-bridge/examples/math.rs +++ b/egglog-bridge/examples/math.rs @@ -1,5 +1,5 @@ use egglog_bridge::{ - ColumnTy, DefaultVal, EGraph, FunctionConfig, MergeFn, add_expressions, define_rule, + ColumnTy, DefaultVal, EGraph, FunctionConfig, FunctionId, MergeFn, add_expressions, define_rule, }; use mimalloc::MiMalloc; use num_rational::Rational64; @@ -8,6 +8,10 @@ use web_time::Instant; #[global_allocator] static GLOBAL: MiMalloc = MiMalloc; +fn add_table(egraph: &mut EGraph, config: FunctionConfig) -> FunctionId { + egraph.add_table(config).unwrap() +} + #[allow(clippy::disallowed_macros)] fn main() { env_logger::init(); @@ -19,107 +23,146 @@ fn main() { let rational_ty = egraph.base_values_mut().register_type::(); let string_ty = egraph.base_values_mut().register_type::<&'static str>(); // tables - let diff = egraph.add_table(FunctionConfig { - schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], - default: DefaultVal::FreshId, - merge: MergeFn::UnionId, - name: "diff".into(), - can_subsume: false, - }); - let integral = egraph.add_table(FunctionConfig { - schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], - default: DefaultVal::FreshId, - merge: MergeFn::UnionId, - name: "integral".into(), - can_subsume: false, - }); + let diff = add_table( + &mut egraph, + FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], + default: DefaultVal::FreshId, + merge: MergeFn::UnionId, + name: "diff".into(), + can_subsume: false, + }, + ); + let integral = add_table( + &mut egraph, + FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], + default: DefaultVal::FreshId, + merge: MergeFn::UnionId, + name: "integral".into(), + can_subsume: false, + }, + ); - let add = egraph.add_table(FunctionConfig { - schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], - default: DefaultVal::FreshId, - merge: MergeFn::UnionId, - name: "add".into(), - can_subsume: false, - }); - let sub = egraph.add_table(FunctionConfig { - schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], - default: DefaultVal::FreshId, - merge: MergeFn::UnionId, - name: "sub".into(), - can_subsume: false, - }); + let add = add_table( + &mut egraph, + FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], + default: DefaultVal::FreshId, + merge: MergeFn::UnionId, + name: "add".into(), + can_subsume: false, + }, + ); + let sub = add_table( + &mut egraph, + FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], + default: DefaultVal::FreshId, + merge: MergeFn::UnionId, + name: "sub".into(), + can_subsume: false, + }, + ); - let mul = egraph.add_table(FunctionConfig { - schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], - default: DefaultVal::FreshId, - merge: MergeFn::UnionId, - name: "mul".into(), - can_subsume: false, - }); + let mul = add_table( + &mut egraph, + FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], + default: DefaultVal::FreshId, + merge: MergeFn::UnionId, + name: "mul".into(), + can_subsume: false, + }, + ); - let div = egraph.add_table(FunctionConfig { - schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], - default: DefaultVal::FreshId, - merge: MergeFn::UnionId, - name: "div".into(), - can_subsume: false, - }); + let div = add_table( + &mut egraph, + FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], + default: DefaultVal::FreshId, + merge: MergeFn::UnionId, + name: "div".into(), + can_subsume: false, + }, + ); - let pow = egraph.add_table(FunctionConfig { - schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], - default: DefaultVal::FreshId, - merge: MergeFn::UnionId, - name: "pow".into(), - can_subsume: false, - }); + let pow = add_table( + &mut egraph, + FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], + default: DefaultVal::FreshId, + merge: MergeFn::UnionId, + name: "pow".into(), + can_subsume: false, + }, + ); - let ln = egraph.add_table(FunctionConfig { - schema: vec![ColumnTy::Id, ColumnTy::Id], - default: DefaultVal::FreshId, - merge: MergeFn::UnionId, - name: "ln".into(), - can_subsume: false, - }); + let ln = add_table( + &mut egraph, + FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Id], + default: DefaultVal::FreshId, + merge: MergeFn::UnionId, + name: "ln".into(), + can_subsume: false, + }, + ); - let sqrt = egraph.add_table(FunctionConfig { - schema: vec![ColumnTy::Id, ColumnTy::Id], - default: DefaultVal::FreshId, - merge: MergeFn::UnionId, - name: "sqrt".into(), - can_subsume: false, - }); + let sqrt = add_table( + &mut egraph, + FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Id], + default: DefaultVal::FreshId, + merge: MergeFn::UnionId, + name: "sqrt".into(), + can_subsume: false, + }, + ); - let sin = egraph.add_table(FunctionConfig { - schema: vec![ColumnTy::Id, ColumnTy::Id], - default: DefaultVal::FreshId, - merge: MergeFn::UnionId, - name: "sin".into(), - can_subsume: false, - }); + let sin = add_table( + &mut egraph, + FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Id], + default: DefaultVal::FreshId, + merge: MergeFn::UnionId, + name: "sin".into(), + can_subsume: false, + }, + ); - let cos = egraph.add_table(FunctionConfig { - schema: vec![ColumnTy::Id, ColumnTy::Id], - default: DefaultVal::FreshId, - merge: MergeFn::UnionId, - name: "cos".into(), - can_subsume: false, - }); + let cos = add_table( + &mut egraph, + FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Id], + default: DefaultVal::FreshId, + merge: MergeFn::UnionId, + name: "cos".into(), + can_subsume: false, + }, + ); - let rat = egraph.add_table(FunctionConfig { - schema: vec![ColumnTy::Base(rational_ty), ColumnTy::Id], - default: DefaultVal::FreshId, - merge: MergeFn::UnionId, - name: "rat".into(), - can_subsume: false, - }); + let rat = add_table( + &mut egraph, + FunctionConfig { + schema: vec![ColumnTy::Base(rational_ty), ColumnTy::Id], + default: DefaultVal::FreshId, + merge: MergeFn::UnionId, + name: "rat".into(), + can_subsume: false, + }, + ); - let var = egraph.add_table(FunctionConfig { - schema: vec![ColumnTy::Base(string_ty), ColumnTy::Id], - default: DefaultVal::FreshId, - merge: MergeFn::UnionId, - name: "var".into(), - can_subsume: false, - }); + let var = add_table( + &mut egraph, + FunctionConfig { + schema: vec![ColumnTy::Base(string_ty), ColumnTy::Id], + default: DefaultVal::FreshId, + merge: MergeFn::UnionId, + name: "var".into(), + can_subsume: false, + }, + ); let zero = egraph.base_value_constant(Rational64::new(0, 1)); let one = egraph.base_value_constant(Rational64::new(1, 1)); diff --git a/egglog-bridge/src/lib.rs b/egglog-bridge/src/lib.rs index 5a3d0993e..a278ade12 100644 --- a/egglog-bridge/src/lib.rs +++ b/egglog-bridge/src/lib.rs @@ -136,6 +136,16 @@ pub struct EGraph { pub type Result = std::result::Result; +#[derive(Error, Debug)] +pub enum FunctionConfigError { + #[error("UF functions are not supported when tracing is enabled")] + UfWithTracing, + #[error("Merge functions cannot call UF-backed function {0}")] + MergeUsesUf(String), + #[error("TableAction does not support UF-backed function {0}")] + TableActionUsesUf(String), +} + impl Default for EGraph { fn default() -> Self { let mut db = Database::new(); @@ -724,7 +734,7 @@ impl EGraph { } /// Register a function in this EGraph. - pub fn add_table(&mut self, config: FunctionConfig) -> FunctionId { + pub fn add_table(&mut self, config: FunctionConfig) -> Result { let FunctionConfig { schema, default, @@ -752,8 +762,8 @@ impl EGraph { let next_func_id = self.funcs.next_id(); let mut read_deps = IndexSet::::new(); let mut write_deps = IndexSet::::new(); - merge.fill_deps(self, &mut read_deps, &mut write_deps); - let merge_fn = merge.to_callback(schema_math, &name, self); + merge.fill_deps(self, &mut read_deps, &mut write_deps)?; + let merge_fn = merge.to_callback(schema_math, &name, self)?; let table = SortedWritesTable::new( n_args, n_cols, @@ -785,13 +795,13 @@ impl EGraph { let info = &mut self.funcs[res]; info.incremental_rebuild_rules = incremental_rebuild_rules; info.nonincremental_rebuild_rule = nonincremental_rebuild_rule; - res + Ok(res) } /// Register a UF-backed function in this EGraph. - pub fn add_uf_function(&mut self, config: UfFunctionConfig) -> FunctionId { + pub fn add_uf_function(&mut self, config: UfFunctionConfig) -> Result { if self.tracing { - panic!("UF functions are not supported when tracing is enabled"); + return Err(FunctionConfigError::UfWithTracing.into()); } let UfFunctionConfig { name, @@ -811,7 +821,7 @@ impl EGraph { write_deps.iter().copied(), ); let schema = vec![ColumnTy::Id, ColumnTy::Id]; - self.funcs.push(FunctionInfo { + Ok(self.funcs.push(FunctionInfo { table: table_id, schema, incremental_rebuild_rules: Vec::new(), @@ -820,7 +830,7 @@ impl EGraph { can_subsume: false, name, kind: FunctionKind::Uf, - }) + })) } /// Run the given rules, returning whether the database changed. @@ -1275,28 +1285,33 @@ impl MergeFn { egraph: &EGraph, read_deps: &mut IndexSet, write_deps: &mut IndexSet, - ) { + ) -> Result<()> { use MergeFn::*; match self { Primitive(_, args) => { - args.iter() - .for_each(|arg| arg.fill_deps(egraph, read_deps, write_deps)); + for arg in args { + arg.fill_deps(egraph, read_deps, write_deps)?; + } } Function(func, args) => { - assert!( - !egraph.funcs[*func].is_uf(), - "Merge functions cannot call UF-backed functions" - ); + if egraph.funcs[*func].is_uf() { + return Err(FunctionConfigError::MergeUsesUf( + egraph.funcs[*func].name.to_string(), + ) + .into()); + } read_deps.insert(egraph.funcs[*func].table); write_deps.insert(egraph.funcs[*func].table); - args.iter() - .for_each(|arg| arg.fill_deps(egraph, read_deps, write_deps)); + for arg in args { + arg.fill_deps(egraph, read_deps, write_deps)?; + } } UnionId if !egraph.tracing => { write_deps.insert(egraph.uf_table); } UnionId | AssertEq | Old | New | Const(..) => {} } + Ok(()) } fn to_callback( @@ -1304,15 +1319,15 @@ impl MergeFn { schema_math: SchemaMath, function_name: &str, egraph: &mut EGraph, - ) -> Box { + ) -> Result> { assert!( !egraph.tracing || matches!(self, MergeFn::UnionId), "proofs aren't supported for non-union merge functions" ); - let resolved = self.resolve(function_name, egraph); + let resolved = self.resolve(function_name, egraph)?; - Box::new(move |state, cur, new, out| { + Ok(Box::new(move |state, cur, new, out| { let timestamp = new[schema_math.ts_col()]; let mut changed = false; @@ -1354,51 +1369,50 @@ impl MergeFn { } changed - }) + })) } - fn resolve(&self, function_name: &str, egraph: &mut EGraph) -> ResolvedMergeFn { + fn resolve(&self, function_name: &str, egraph: &mut EGraph) -> Result { match self { - MergeFn::Const(v) => ResolvedMergeFn::Const(*v), - MergeFn::Old => ResolvedMergeFn::Old, - MergeFn::New => ResolvedMergeFn::New, - MergeFn::AssertEq => ResolvedMergeFn::AssertEq { + MergeFn::Const(v) => Ok(ResolvedMergeFn::Const(*v)), + MergeFn::Old => Ok(ResolvedMergeFn::Old), + MergeFn::New => Ok(ResolvedMergeFn::New), + MergeFn::AssertEq => Ok(ResolvedMergeFn::AssertEq { panic: egraph.new_panic(format!( "Illegal merge attempted for function {function_name}" )), - }, - MergeFn::UnionId => ResolvedMergeFn::UnionId { + }), + MergeFn::UnionId => Ok(ResolvedMergeFn::UnionId { uf_table: egraph.uf_table, tracing: egraph.tracing, - }, + }), // NB: The primitive and function-based merge functions heap allocate a single callback // for each layer of nesting. This introduces a bit of overhead, particularly for cases // that look like `(f old new)` or `(f new old)`. We could special-case common cases in // this function if that overhead shows up. - MergeFn::Primitive(prim, args) => ResolvedMergeFn::Primitive { + MergeFn::Primitive(prim, args) => Ok(ResolvedMergeFn::Primitive { prim: *prim, args: args .iter() .map(|arg| arg.resolve(function_name, egraph)) - .collect::>(), + .collect::>>()?, panic: egraph.new_panic(format!( "Merge function for {function_name} primitive call failed" )), - }, + }), MergeFn::Function(func, args) => { let func_info = &egraph.funcs[*func]; - assert!( - !func_info.is_uf(), - "Merge functions cannot call UF-backed functions" - ); + if func_info.is_uf() { + return Err(FunctionConfigError::MergeUsesUf(func_info.name.to_string()).into()); + } assert_eq!( func_info.schema.len(), args.len() + 1, "Merge function for {function_name} must match function arity for {}", func_info.name ); - ResolvedMergeFn::Function { - func: TableAction::new(egraph, *func), + Ok(ResolvedMergeFn::Function { + func: TableAction::new(egraph, *func)?, panic: egraph.new_panic(format!( "Lookup on {} failed in the merge function for {function_name}", func_info.name @@ -1406,8 +1420,8 @@ impl MergeFn { args: args .iter() .map(|arg| arg.resolve(function_name, egraph)) - .collect::>(), - } + .collect::>>()?, + }) } } } @@ -1532,15 +1546,14 @@ impl Clone for TableAction { impl TableAction { /// Create a new `TableAction` to be used later. /// This requires access to the `egglog_bridge::EGraph`. - pub fn new(egraph: &EGraph, func: FunctionId) -> TableAction { + pub fn new(egraph: &EGraph, func: FunctionId) -> Result { assert!(!egraph.tracing, "proofs not supported yet"); let func_info = &egraph.funcs[func]; - assert!( - !func_info.is_uf(), - "TableAction does not support UF-backed functions" - ); - TableAction { + if func_info.is_uf() { + return Err(FunctionConfigError::TableActionUsesUf(func_info.name.to_string()).into()); + } + Ok(TableAction { table: func_info.table, table_math: SchemaMath { func_cols: func_info.schema.len(), @@ -1554,7 +1567,7 @@ impl TableAction { }, timestamp: egraph.timestamp_counter, scratch: Vec::new(), - } + }) } /// A "table lookup" is not a read-only operation. It will insert a row when diff --git a/egglog-bridge/src/tests.rs b/egglog-bridge/src/tests.rs index f6454ca9c..ceb8dc14c 100644 --- a/egglog-bridge/src/tests.rs +++ b/egglog-bridge/src/tests.rs @@ -19,7 +19,7 @@ use num_rational::Rational64; use crate::{ ColumnTy, DefaultVal, EGraph, FunctionConfig, FunctionId, MergeFn, ProofStore, QueryEntry, - UfFunctionConfig, add_expressions, define_rule, + Result, UfFunctionConfig, add_expressions, define_rule, }; /// Run a simple associativity/commutativity test. In addition to testing that the rules properly @@ -29,7 +29,7 @@ use crate::{ /// The `can_subsume` argument is only used to enable subsumption on the underlying tables created /// during this test, and exercise the different column handling caused by enabling subsumption. /// Subsumption itself is not used. -fn ac_test(tracing: bool, can_subsume: bool) { +fn ac_test(tracing: bool, can_subsume: bool) -> Result<()> { const N: usize = 5; let mut egraph = if tracing { EGraph::with_tracing() @@ -43,14 +43,14 @@ fn ac_test(tracing: bool, can_subsume: bool) { merge: MergeFn::UnionId, name: "num".into(), can_subsume, - }); + })?; let add_table = egraph.add_table(FunctionConfig { schema: vec![ColumnTy::Id; 3], default: DefaultVal::FreshId, merge: MergeFn::UnionId, name: "add".into(), can_subsume, - }); + })?; let add_comm = define_rule! { [egraph] ((-> (add_table x y) id)) @@ -115,10 +115,11 @@ fn ac_test(tracing: bool, can_subsume: bool) { // .print_eq_proof(_eq_explanation, &mut std::io::stderr()) // .unwrap(); } + Ok(()) } #[test] -fn uf_function_callback_inserts() { +fn uf_function_callback_inserts() -> Result<()> { let mut egraph = EGraph::default(); let log_table = egraph.add_table(FunctionConfig { schema: vec![ColumnTy::Id, ColumnTy::Id], @@ -126,7 +127,7 @@ fn uf_function_callback_inserts() { merge: MergeFn::AssertEq, name: "log".into(), can_subsume: false, - }); + })?; let log_table_id = egraph.funcs[log_table].table; let uf_func = egraph.add_uf_function(UfFunctionConfig { name: "uf_cb".into(), @@ -138,7 +139,7 @@ fn uf_function_callback_inserts() { })), read_deps: Vec::new(), write_deps: vec![log_table_id], - }); + })?; let lhs = Value::from_usize(5); let rhs = Value::from_usize(3); @@ -150,65 +151,64 @@ fn uf_function_callback_inserts() { let mut rows = Vec::new(); egraph.for_each(log_table, |row| rows.push(row.vals.to_vec())); assert_eq!(rows, vec![vec![lhs, rhs]]); + Ok(()) } #[test] -fn uf_function_disallowed_in_merge() { +fn uf_function_disallowed_in_merge() -> Result<()> { let mut egraph = EGraph::default(); let uf_func = egraph.add_uf_function(UfFunctionConfig { name: "uf_func".into(), on_leader_change: None, read_deps: Vec::new(), write_deps: Vec::new(), + })?; + let result = egraph.add_table(FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Id], + default: DefaultVal::Fail, + merge: MergeFn::Function(uf_func, vec![MergeFn::New]), + name: "bad_merge".into(), + can_subsume: false, }); - let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - egraph.add_table(FunctionConfig { - schema: vec![ColumnTy::Id, ColumnTy::Id], - default: DefaultVal::Fail, - merge: MergeFn::Function(uf_func, vec![MergeFn::New]), - name: "bad_merge".into(), - can_subsume: false, - }); - })); assert!(result.is_err()); + Ok(()) } #[test] -fn uf_function_disallowed_with_tracing() { +fn uf_function_disallowed_with_tracing() -> Result<()> { let mut egraph = EGraph::with_tracing(); - let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - egraph.add_uf_function(UfFunctionConfig { - name: "uf_tracing".into(), - on_leader_change: None, - read_deps: Vec::new(), - write_deps: Vec::new(), - }); - })); + let result = egraph.add_uf_function(UfFunctionConfig { + name: "uf_tracing".into(), + on_leader_change: None, + read_deps: Vec::new(), + write_deps: Vec::new(), + }); assert!(result.is_err()); + Ok(()) } #[test] -fn ac_tracing_subsume() { - ac_test(true, true); +fn ac_tracing_subsume() -> Result<()> { + ac_test(true, true) } #[test] -fn ac_tracing() { - ac_test(true, false); +fn ac_tracing() -> Result<()> { + ac_test(true, false) } #[test] -fn ac() { - ac_test(false, false); +fn ac() -> Result<()> { + ac_test(false, false) } #[test] -fn ac_subsume() { - ac_test(false, true); +fn ac_subsume() -> Result<()> { + ac_test(false, true) } #[test] -fn ac_fail() { +fn ac_fail() -> Result<()> { const N: usize = 5; let mut egraph = EGraph::default(); egraph.base_values_mut().register_type::(); @@ -220,14 +220,14 @@ fn ac_fail() { merge: MergeFn::UnionId, name: "num".into(), can_subsume: false, - }); + })?; let add_table = egraph.add_table(FunctionConfig { schema: vec![ColumnTy::Id; 3], default: DefaultVal::FreshId, merge: MergeFn::UnionId, name: "add".into(), can_subsume: false, - }); + })?; let add_comm = define_rule! { [egraph] ((-> (add_table x y) id) (-> (num_table {one}) x)) @@ -282,28 +282,35 @@ fn ac_fail() { let canon_left = egraph.get_canon_in_uf(left_root); let canon_right = egraph.get_canon_in_uf(right_root); assert_ne!(canon_left, canon_right); + Ok(()) } #[test] -fn math() { +fn math() -> Result<()> { let handles = Vec::from_iter((0..2).map(|_| thread::spawn(|| math_test(EGraph::default(), false)))); - handles.into_iter().for_each(|h| h.join().unwrap()); + for handle in handles { + handle.join().unwrap()?; + } + Ok(()) } #[test] -fn math_subsume() { +fn math_subsume() -> Result<()> { let handles = Vec::from_iter((0..2).map(|_| thread::spawn(|| math_test(EGraph::default(), true)))); - handles.into_iter().for_each(|h| h.join().unwrap()); + for handle in handles { + handle.join().unwrap()?; + } + Ok(()) } #[test] -fn math_tracing() { +fn math_tracing() -> Result<()> { math_test(EGraph::with_tracing(), false) } #[test] -fn math_tracing_subsume() { +fn math_tracing_subsume() -> Result<()> { math_test(EGraph::with_tracing(), true) } @@ -314,7 +321,7 @@ fn math_tracing_subsume() { /// As in `ac_test` the `can_subsume` argument is only used to enable subsumption on the underlying /// tables created during this test, and exercise the different column handling caused by enabling /// subsumption. Subsumption itself is not used. -fn math_test(mut egraph: EGraph, can_subsume: bool) { +fn math_test(mut egraph: EGraph, can_subsume: bool) -> Result<()> { const N: usize = 8; let rational_ty = egraph.base_values_mut().register_type::(); let string_ty = egraph.base_values_mut().register_type::<&'static str>(); @@ -325,49 +332,49 @@ fn math_test(mut egraph: EGraph, can_subsume: bool) { merge: MergeFn::UnionId, name: "diff".into(), can_subsume, - }); + })?; let integral = egraph.add_table(FunctionConfig { schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, name: "integral".into(), can_subsume, - }); + })?; let add = egraph.add_table(FunctionConfig { schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, name: "add".into(), can_subsume, - }); + })?; let sub = egraph.add_table(FunctionConfig { schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, name: "sub".into(), can_subsume, - }); + })?; let mul = egraph.add_table(FunctionConfig { schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, name: "mul".into(), can_subsume, - }); + })?; let div = egraph.add_table(FunctionConfig { schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, name: "div".into(), can_subsume, - }); + })?; let pow = egraph.add_table(FunctionConfig { schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, name: "pow".into(), can_subsume, - }); + })?; let ln = egraph.add_table(FunctionConfig { schema: vec![ColumnTy::Id, ColumnTy::Id], @@ -375,42 +382,42 @@ fn math_test(mut egraph: EGraph, can_subsume: bool) { merge: MergeFn::UnionId, name: "ln".into(), can_subsume, - }); + })?; let sqrt = egraph.add_table(FunctionConfig { schema: vec![ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, name: "sqrt".into(), can_subsume, - }); + })?; let sin = egraph.add_table(FunctionConfig { schema: vec![ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, name: "sin".into(), can_subsume, - }); + })?; let cos = egraph.add_table(FunctionConfig { schema: vec![ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, name: "cos".into(), can_subsume, - }); + })?; let rat = egraph.add_table(FunctionConfig { schema: vec![ColumnTy::Base(rational_ty), ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, name: "rat".into(), can_subsume, - }); + })?; let var = egraph.add_table(FunctionConfig { schema: vec![ColumnTy::Base(string_ty), ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, name: "var".into(), can_subsume, - }); + })?; let zero = egraph.base_value_constant(Rational64::new(0, 1)); let one = egraph.base_value_constant(Rational64::new(1, 1)); @@ -574,6 +581,7 @@ fn math_test(mut egraph: EGraph, can_subsume: bool) { let mut proof_store = ProofStore::default(); let _explain = egraph.explain_term(term_id, &mut proof_store).unwrap(); } + Ok(()) } #[derive(Clone, Debug, Hash, Eq, PartialEq)] @@ -635,7 +643,7 @@ fn assert_unordered_eq(mut a: Vec, mut b: Vec) { assert_eq!(a, b); } -fn container_test() { +fn container_test() -> Result<()> { // Test for containers: // * Basic math setup: (num i64), (add math math), (Vec (vec math)) // * start with: @@ -669,21 +677,21 @@ fn container_test() { merge: MergeFn::UnionId, name: "num".into(), can_subsume: false, - }); + })?; let add_table = egraph.add_table(FunctionConfig { schema: vec![ColumnTy::Id; 3], default: DefaultVal::FreshId, merge: MergeFn::UnionId, name: "add".into(), can_subsume: false, - }); + })?; let vec_table = egraph.add_table(FunctionConfig { schema: vec![ColumnTy::Id; 2], default: DefaultVal::FreshId, merge: MergeFn::UnionId, name: "vec".into(), can_subsume: false, - }); + })?; let int_add = egraph.register_external_func(Box::new(make_external_func(|exec_state, args| { let [x, y] = args else { panic!() }; @@ -834,18 +842,20 @@ fn container_test() { vec![one_id; 4], ], ); + Ok(()) } #[test] -fn basic_container() { +fn basic_container() -> Result<()> { // Run the test 8 times to get a decent sample of incremental/nonincremental, parallel/serial. for _ in 0..8 { - container_test() + container_test()?; } + Ok(()) } #[test] -fn rhs_only_rule() { +fn rhs_only_rule() -> Result<()> { let mut egraph = EGraph::default(); let int_base = egraph.base_values_mut().register_type::(); let zero = egraph.base_values_mut().get(0i64); @@ -856,7 +866,7 @@ fn rhs_only_rule() { merge: MergeFn::UnionId, name: "num".into(), can_subsume: false, - }); + })?; let add_data = { let zero = egraph.base_value_constant(0i64); let one = egraph.base_value_constant(1i64); @@ -880,10 +890,11 @@ fn rhs_only_rule() { contents, vec![vec![zero, Value::new(0)], vec![one, Value::new(1)]] ); + Ok(()) } #[test] -fn rhs_only_rule_only_runs_once() { +fn rhs_only_rule_only_runs_once() -> Result<()> { let mut egraph = EGraph::default(); let counter = Arc::new(AtomicUsize::new(0)); let inner = counter.clone(); @@ -902,10 +913,11 @@ fn rhs_only_rule_only_runs_once() { assert_eq!(counter.load(Ordering::SeqCst), 1); assert!(!egraph.run_rules(&[inc_counter_rule]).unwrap().changed()); assert_eq!(counter.load(Ordering::SeqCst), 1); + Ok(()) } #[test] -fn mergefn_arithmetic() { +fn mergefn_arithmetic() -> Result<()> { let mut egraph = EGraph::default(); let int_base = egraph.base_values_mut().register_type::(); @@ -950,7 +962,7 @@ fn mergefn_arithmetic() { ), name: "f".into(), can_subsume: false, - }); + })?; let value_0 = egraph.base_value_constant(0i64); let value_1 = egraph.base_value_constant(1i64); @@ -1026,10 +1038,11 @@ fn mergefn_arithmetic() { }); contents.sort(); assert_eq!(contents, vec![(1, 4), (2, 29)]); + Ok(()) } #[test] -fn mergefn_nested_function() { +fn mergefn_nested_function() -> Result<()> { let mut egraph = EGraph::default(); let int_base = egraph.base_values_mut().register_type::(); @@ -1040,7 +1053,7 @@ fn mergefn_nested_function() { merge: MergeFn::UnionId, name: "g".into(), can_subsume: true, - }); + })?; // Create a function f whose merge function is (g (g new new) (g old old)) // This uses nested MergeFn::Function to build the complex merge function @@ -1056,7 +1069,7 @@ fn mergefn_nested_function() { ), name: "f".into(), can_subsume: true, - }); + })?; let value_1 = egraph.base_value_constant(1i64); let value_2 = egraph.base_value_constant(2i64); @@ -1147,10 +1160,11 @@ fn mergefn_nested_function() { assert_eq!(base_l1, base_2); assert_eq!(base_r1, base_r2); assert_eq!(base_r1, base_1); + Ok(()) } #[test] -fn constrain_prims_simple() { +fn constrain_prims_simple() -> Result<()> { // Take two functions, f and g. Fill f with (f 1) (f 2) (f 3), then filter for even numbers // when adding to 'g'. This should only add 2 to g. let mut egraph = EGraph::default(); @@ -1162,14 +1176,14 @@ fn constrain_prims_simple() { merge: MergeFn::UnionId, name: "f".into(), can_subsume: false, - }); + })?; let g_table = egraph.add_table(FunctionConfig { schema: vec![ColumnTy::Base(int_base), ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, name: "g".into(), can_subsume: false, - }); + })?; let is_even = egraph.register_external_func(Box::new(core_relations::make_external_func( |state, vals| -> Option { @@ -1230,11 +1244,12 @@ fn constrain_prims_simple() { egraph.run_rules(&[copy_to_g]).unwrap(); let g = get_entries(&egraph, g_table); assert_eq!(g.len(), 1); - assert_eq!(g[0], f[1]) + assert_eq!(g[0], f[1]); + Ok(()) } #[test] -fn constrain_prims_abstract() { +fn constrain_prims_abstract() -> Result<()> { // Take two functions, f and g. Fill f with (f -1) (f 0) (f 1), then filter for numbers where // (neg x) = (abs x) when adding to 'g'. This adds only -1 and 0 to g let mut egraph = EGraph::default(); @@ -1245,14 +1260,14 @@ fn constrain_prims_abstract() { merge: MergeFn::UnionId, name: "f".into(), can_subsume: false, - }); + })?; let g_table = egraph.add_table(FunctionConfig { schema: vec![ColumnTy::Base(int_base), ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, name: "g".into(), can_subsume: false, - }); + })?; let neg = egraph.register_external_func(Box::new(core_relations::make_external_func( |state, vals| -> Option { @@ -1327,11 +1342,12 @@ fn constrain_prims_abstract() { egraph.run_rules(&[copy_to_g]).unwrap(); let g = get_entries(&egraph, g_table); assert_eq!(g.len(), 2); - assert_eq!(g, f[0..2]) + assert_eq!(g, f[0..2]); + Ok(()) } #[test] -fn basic_subsumption() { +fn basic_subsumption() -> Result<()> { // fill (f 1) (f 2). Subsume (f 3) (f 2). Copy (f to g). Should only see (g 1) let mut egraph = EGraph::default(); @@ -1342,14 +1358,14 @@ fn basic_subsumption() { merge: MergeFn::UnionId, name: "f".into(), can_subsume: true, - }); + })?; let g_table = egraph.add_table(FunctionConfig { schema: vec![ColumnTy::Base(int_base), ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, name: "g".into(), can_subsume: false, - }); + })?; let value_1 = egraph.base_value_constant(1i64); let value_2 = egraph.base_value_constant(2i64); @@ -1409,11 +1425,12 @@ fn basic_subsumption() { egraph.run_rules(&[copy_to_g]).unwrap(); let g = get_entries(&egraph, g_table); assert_eq!((g.0.len(), g.1), (1, 0)); - assert_eq!(g.0[0], f.0[0]) + assert_eq!(g.0[0], f.0[0]); + Ok(()) } #[test] -fn lookup_failure_panics() { +fn lookup_failure_panics() -> Result<()> { let mut egraph = EGraph::default(); let f = egraph.add_table(FunctionConfig { schema: vec![ColumnTy::Id, ColumnTy::Id], @@ -1421,7 +1438,7 @@ fn lookup_failure_panics() { merge: MergeFn::UnionId, name: "test".into(), can_subsume: false, - }); + })?; let to_entry = |val: u32| QueryEntry::Const { val: Value::new(val), @@ -1452,10 +1469,11 @@ fn lookup_failure_panics() { rb.build() }; egraph.run_rules(&[lookup_failure]).err().unwrap(); + Ok(()) } #[test] -fn primitive_failure_panics() { +fn primitive_failure_panics() -> Result<()> { let mut egraph = EGraph::default(); let _int_base = egraph.base_values_mut().register_type::(); let unit_base = egraph.base_values_mut().register_type::<()>(); @@ -1495,10 +1513,11 @@ fn primitive_failure_panics() { }; egraph.run_rules(&[assert_odd_rule]).err().unwrap(); + Ok(()) } #[test] -fn test_simple_rule_proof_format() { +fn test_simple_rule_proof_format() -> Result<()> { use crate::proof_format::*; // Setup EGraph with tracing let mut egraph = EGraph::with_tracing(); @@ -1513,7 +1532,7 @@ fn test_simple_rule_proof_format() { merge: MergeFn::UnionId, name: "bool".into(), can_subsume: false, - }); + })?; // Add table for not function let not_table = egraph.add_table(FunctionConfig { schema: vec![ColumnTy::Id, ColumnTy::Id], @@ -1521,7 +1540,7 @@ fn test_simple_rule_proof_format() { merge: MergeFn::UnionId, name: "not".into(), can_subsume: false, - }); + })?; // Add true/false wrapped terms let true_id = egraph.add_term(bool_table, &[true_val], "true"); let false_id = egraph.add_term(bool_table, &[false_val], "false"); @@ -1543,6 +1562,7 @@ fn test_simple_rule_proof_format() { egraph .explain_terms_equal(not_true_id, false_id, &mut proof_store) .unwrap(); + Ok(()) } const _: () = { diff --git a/src/lib.rs b/src/lib.rs index a7013de13..ce9254726 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -550,28 +550,33 @@ impl EGraph { }; use egglog_bridge::{DefaultVal, MergeFn}; - let backend_id = self.backend.add_table(egglog_bridge::FunctionConfig { - schema: input - .iter() - .chain([&output]) - .map(|sort| sort.column_ty(&self.backend)) - .collect(), - default: match decl.subtype { - FunctionSubtype::Constructor => DefaultVal::FreshId, - FunctionSubtype::Custom => DefaultVal::Fail, - FunctionSubtype::Relation => DefaultVal::Const(self.backend.base_values().get(())), - }, - merge: match decl.subtype { - FunctionSubtype::Constructor => MergeFn::UnionId, - FunctionSubtype::Relation => MergeFn::AssertEq, - FunctionSubtype::Custom => match &decl.merge { - None => MergeFn::AssertEq, - Some(expr) => self.translate_expr_to_mergefn(expr)?, + let backend_id = self + .backend + .add_table(egglog_bridge::FunctionConfig { + schema: input + .iter() + .chain([&output]) + .map(|sort| sort.column_ty(&self.backend)) + .collect(), + default: match decl.subtype { + FunctionSubtype::Constructor => DefaultVal::FreshId, + FunctionSubtype::Custom => DefaultVal::Fail, + FunctionSubtype::Relation => { + DefaultVal::Const(self.backend.base_values().get(())) + } }, - }, - name: decl.name.to_string(), - can_subsume, - }); + merge: match decl.subtype { + FunctionSubtype::Constructor => MergeFn::UnionId, + FunctionSubtype::Relation => MergeFn::AssertEq, + FunctionSubtype::Custom => match &decl.merge { + None => MergeFn::AssertEq, + Some(expr) => self.translate_expr_to_mergefn(expr)?, + }, + }, + name: decl.name.to_string(), + can_subsume, + }) + .map_err(|err| Error::BackendError(err.to_string()))?; let function = Function { decl: decl.clone(), @@ -876,7 +881,7 @@ impl EGraph { &self.functions, &self.type_info, ); - translator.query(query, false); + translator.query(query, false)?; translator.actions(actions)?; translator.build() }; @@ -1047,7 +1052,7 @@ impl EGraph { &self.functions, &self.type_info, ); - translator.query(&query, true); + translator.query(&query, true)?; translator .rb .call_external_func(ext_id, &[], egglog_bridge::ColumnTy::Id, || { @@ -1366,7 +1371,8 @@ impl EGraph { let num_facts = parsed_contents.len(); - let mut table_action = egglog_bridge::TableAction::new(&self.backend, func.backend_id); + let mut table_action = egglog_bridge::TableAction::new(&self.backend, func.backend_id) + .map_err(|err| Error::BackendError(err.to_string()))?; if function_type.subtype != FunctionSubtype::Constructor { self.backend.with_execution_state(|es| { @@ -1695,7 +1701,7 @@ impl<'a> BackendRule<'a> { &mut self, prim: &core::SpecializedPrimitive, args: &[core::ResolvedAtomTerm], - ) -> (ExternalFunctionId, Vec, ColumnTy) { + ) -> Result<(ExternalFunctionId, Vec, ColumnTy), Error> { let mut qe_args = self.args(args); if prim.name() == "unstable-fn" { @@ -1703,10 +1709,10 @@ impl<'a> BackendRule<'a> { panic!("expected string literal after `unstable-fn`") }; let id = if let Some(f) = self.type_info.get_func_type(name) { - ResolvedFunctionId::Lookup(egglog_bridge::TableAction::new( - self.rb.egraph(), - self.func(f), - )) + ResolvedFunctionId::Lookup( + egglog_bridge::TableAction::new(self.rb.egraph(), self.func(f)) + .map_err(|err| Error::BackendError(err.to_string()))?, + ) } else if let Some(possible) = self.type_info.get_prims(name) { let mut ps: Vec<_> = possible.iter().collect(); ps.retain(|p| { @@ -1739,11 +1745,11 @@ impl<'a> BackendRule<'a> { }); } - ( + Ok(( prim.external_id(), qe_args, prim.output().column_ty(self.rb.egraph()), - ) + )) } fn args<'b>( @@ -1753,7 +1759,11 @@ impl<'a> BackendRule<'a> { args.into_iter().map(|x| self.entry(x)).collect() } - fn query(&mut self, query: &core::Query, include_subsumed: bool) { + fn query( + &mut self, + query: &core::Query, + include_subsumed: bool, + ) -> Result<(), Error> { for atom in &query.atoms { match &atom.head { ResolvedCall::Func(f) => { @@ -1763,14 +1773,19 @@ impl<'a> BackendRule<'a> { true => None, false => Some(false), }; - self.rb.query_table(f, &args, is_subsumed).unwrap(); + self.rb + .query_table(f, &args, is_subsumed) + .map_err(|e| Error::BackendError(e.to_string()))?; } ResolvedCall::Primitive(p) => { - let (p, args, ty) = self.prim(p, &atom.args); - self.rb.query_prim(p, &args, ty).unwrap() + let (p, args, ty) = self.prim(p, &atom.args)?; + self.rb + .query_prim(p, &args, ty) + .map_err(|e| Error::BackendError(e.to_string()))? } } } + Ok(()) } fn actions(&mut self, actions: &core::ResolvedCoreActions) -> Result<(), Error> { @@ -1790,7 +1805,7 @@ impl<'a> BackendRule<'a> { } ResolvedCall::Primitive(p) => { let name = p.name().to_owned(); - let (p, args, ty) = self.prim(p, args); + let (p, args, ty) = self.prim(p, args)?; let span = span.clone(); self.rb.call_external_func(p, &args, ty, move || { format!("{span}: call of primitive {name} failed") diff --git a/src/prelude.rs b/src/prelude.rs index 48af1367b..74dc0d31b 100644 --- a/src/prelude.rs +++ b/src/prelude.rs @@ -454,12 +454,11 @@ pub fn rust_rule( .functions .iter() .map(|(k, v)| { - ( - k.clone(), - egglog_bridge::TableAction::new(&egraph.backend, v.backend_id), - ) + egglog_bridge::TableAction::new(&egraph.backend, v.backend_id) + .map(|action| (k.clone(), action)) }) - .collect(), + .collect::>() + .map_err(|err| Error::BackendError(err.to_string()))?, panic_id, func, }); diff --git a/src/scheduler.rs b/src/scheduler.rs index c6924dac4..d1b81b323 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -195,12 +195,12 @@ impl EGraph { // Step 1: build all the query/action rules and worklist if have not already let record = &mut schedulers[scheduler_id]; - rules.iter().for_each(|(id, rule)| { - record - .rule_info - .entry((*id).to_owned()) - .or_insert_with(|| SchedulerRuleInfo::new(self, rule, id)); - }); + for (id, rule) in &rules { + if !record.rule_info.contains_key(id) { + let info = SchedulerRuleInfo::new(self, rule, id)?; + record.rule_info.insert(id.to_owned(), info); + } + } // Step 2: run all the queries for one iteration let query_rules = rules @@ -222,21 +222,24 @@ impl EGraph { .map_err(|e| Error::BackendError(e.to_string()))?; // Step 3: let the scheduler decide which matches need to be kept - self.backend.with_execution_state(|state| { - for (rule_id, _rule) in rules.iter() { - let rule_info = record.rule_info.get_mut(rule_id).unwrap(); - - let matches: Vec = - std::mem::take(rule_info.matches.lock().unwrap().as_mut()); - let mut matches = Matches::new(matches, rule_info.free_vars.clone()); - rule_info.should_seek = - record - .scheduler - .filter_matches(rule_id, ruleset, &mut matches); - let table_action = TableAction::new(&self.backend, rule_info.decided); - *rule_info.matches.lock().unwrap() = matches.instantiate(state, table_action); - } - }); + self.backend + .with_execution_state(|state| -> Result<(), Error> { + for (rule_id, _rule) in rules.iter() { + let rule_info = record.rule_info.get_mut(rule_id).unwrap(); + + let matches: Vec = + std::mem::take(rule_info.matches.lock().unwrap().as_mut()); + let mut matches = Matches::new(matches, rule_info.free_vars.clone()); + rule_info.should_seek = + record + .scheduler + .filter_matches(rule_id, ruleset, &mut matches); + let table_action = TableAction::new(&self.backend, rule_info.decided) + .map_err(|err| Error::BackendError(err.to_string()))?; + *rule_info.matches.lock().unwrap() = matches.instantiate(state, table_action); + } + Ok(()) + })?; self.backend.flush_updates(); // Step 4: run the action rules @@ -320,7 +323,11 @@ impl ExternalFunction for CollectMatches { } impl SchedulerRuleInfo { - fn new(egraph: &mut EGraph, rule: &ResolvedCoreRule, name: &str) -> SchedulerRuleInfo { + fn new( + egraph: &mut EGraph, + rule: &ResolvedCoreRule, + name: &str, + ) -> Result { let free_vars = rule.head.get_free_vars().into_iter().collect::>(); let unit_type = egraph.backend.base_values().get_ty::<()>(); let unit = egraph.backend.base_values().get(()); @@ -335,13 +342,16 @@ impl SchedulerRuleInfo { .map(|v| v.sort.column_ty(&egraph.backend)) .chain(std::iter::once(ColumnTy::Base(unit_type))) .collect(); - let decided = egraph.backend.add_table(FunctionConfig { - schema, - default: DefaultVal::Const(unit), - merge: MergeFn::AssertEq, - name: "backend".to_string(), - can_subsume: false, - }); + let decided = egraph + .backend + .add_table(FunctionConfig { + schema, + default: DefaultVal::Const(unit), + merge: MergeFn::AssertEq, + name: "backend".to_string(), + can_subsume: false, + }) + .map_err(|err| Error::BackendError(err.to_string()))?; // Step 1: build the query rule let mut qrule_builder = BackendRule::new( @@ -349,7 +359,7 @@ impl SchedulerRuleInfo { &egraph.functions, &egraph.type_info, ); - qrule_builder.query(&rule.body, true); + qrule_builder.query(&rule.body, true)?; let entries = free_vars .iter() .map(|fv| qrule_builder.entry(&GenericAtomTerm::Var(span!(), fv.clone()))) @@ -376,21 +386,21 @@ impl SchedulerRuleInfo { arule_builder .rb .query_table(decided, &entries, None) - .unwrap(); - arule_builder.actions(&rule.head).unwrap(); + .map_err(|e| Error::BackendError(e.to_string()))?; + arule_builder.actions(&rule.head)?; // Remove the entry as it's now done entries.pop(); arule_builder.rb.remove(decided, &entries); let arule_id = arule_builder.build(); - SchedulerRuleInfo { + Ok(SchedulerRuleInfo { free_vars, query_rule: qrule_id, action_rule: arule_id, matches, decided, should_seek: true, - } + }) } } From c5664fb622aba38218244a6c5c7de739379baf54 Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Mon, 5 Jan 2026 23:19:37 -0800 Subject: [PATCH 4/6] Add rebuilding for custom uf tables. --- egglog-bridge/src/lib.rs | 76 ++++++++++++++++++++++++++++++++++---- egglog-bridge/src/tests.rs | 27 +++++++++++++- 2 files changed, 94 insertions(+), 9 deletions(-) diff --git a/egglog-bridge/src/lib.rs b/egglog-bridge/src/lib.rs index a278ade12..3a664763d 100644 --- a/egglog-bridge/src/lib.rs +++ b/egglog-bridge/src/lib.rs @@ -173,11 +173,13 @@ pub struct FunctionConfig { pub can_subsume: bool, } +pub type LeaderChangeCallback = + Box; + /// Configuration for a UF-backed function. pub struct UfFunctionConfig { pub name: String, - pub on_leader_change: - Option>, + pub on_leader_change: Option, pub read_deps: Vec, pub write_deps: Vec, } @@ -788,6 +790,7 @@ impl EGraph { can_subsume, name, kind: FunctionKind::Table, + uf_rebuild_rule: None, }); debug_assert_eq!(res, next_func_id); let incremental_rebuild_rules = self.incremental_rebuild_rules(res, &schema); @@ -821,7 +824,7 @@ impl EGraph { write_deps.iter().copied(), ); let schema = vec![ColumnTy::Id, ColumnTy::Id]; - Ok(self.funcs.push(FunctionInfo { + let res = self.funcs.push(FunctionInfo { table: table_id, schema, incremental_rebuild_rules: Vec::new(), @@ -830,7 +833,11 @@ impl EGraph { can_subsume: false, name, kind: FunctionKind::Uf, - })) + uf_rebuild_rule: None, + }); + let uf_rebuild_rule = self.uf_rebuild_rule(res); + self.funcs[res].uf_rebuild_rule = Some(uf_rebuild_rule); + Ok(res) } /// Run the given rules, returning whether the database changed. @@ -865,6 +872,7 @@ impl EGraph { } fn rebuild(&mut self) -> Result<()> { + let uf_rules = self.uf_rebuild_rules(); fn do_parallel() -> bool { #[cfg(test)] { @@ -886,6 +894,7 @@ impl EGraph { tables.push(func.table); } loop { + let ts = self.next_ts(); // Order matters here: we need to rebuild containers first and then rebuild the // tables. Why? // @@ -918,11 +927,21 @@ impl EGraph { // Rebuilding containers first will find that v3 and v2 are equal, and the rest of // the rules can proceed. let container_rebuild = self.db.rebuild_containers(self.uf_table); - let table_rebuild = - self.db - .apply_rebuild(self.uf_table, &tables, self.next_ts().to_value()); + let table_rebuild = self.db.apply_rebuild(self.uf_table, &tables, ts.to_value()); + let uf_rebuild = if uf_rules.is_empty() { + false + } else { + run_rules_impl( + &mut self.db, + &mut self.rules, + &uf_rules, + ts, + ReportLevel::TimeOnly, + )? + .changed + }; self.inc_ts(); - if !table_rebuild && !container_rebuild { + if !table_rebuild && !container_rebuild && !uf_rebuild { break; } } @@ -991,6 +1010,16 @@ impl EGraph { })?; } } + if !uf_rules.is_empty() { + changed |= run_rules_impl( + &mut self.db, + &mut self.rules, + &uf_rules, + ts, + ReportLevel::TimeOnly, + )? + .changed; + } } log::info!("rebuild took {:?}", start.elapsed()); Ok(()) @@ -1002,6 +1031,7 @@ impl EGraph { /// when the number of active threads is greater than 1. fn rebuild_parallel(&mut self) -> Result<()> { let start = Instant::now(); + let uf_rules = self.uf_rebuild_rules(); #[derive(Default)] struct RebuildState { nonincremental: Vec, @@ -1078,6 +1108,16 @@ impl EGraph { .changed; scratch.clear(); } + if !uf_rules.is_empty() { + changed |= run_rules_impl( + &mut self.db, + &mut self.rules, + &uf_rules, + ts, + ReportLevel::TimeOnly, + )? + .changed; + } } log::info!("rebuild took {:?}", start.elapsed()); Ok(()) @@ -1179,6 +1219,25 @@ impl EGraph { rb.build_internal(None) // skip the syntax check } + /// Build a rule that forwards global uf_table unions into a UF-backed table so it stays + /// consistent with the main union-find. + fn uf_rebuild_rule(&mut self, table: FunctionId) -> RuleId { + let uf_table = self.uf_table; + let mut rb = self.new_rule(&format!("uf rebuild {table:?}"), true); + let lhs: QueryEntry = rb.new_var(ColumnTy::Id).into(); + let rhs: QueryEntry = rb.new_var(ColumnTy::Id).into(); + rb.add_atom_with_timestamp_and_func(uf_table, None, None, &[lhs.clone(), rhs.clone()]); + rb.set(table, &[lhs, rhs]); + rb.build_internal(None) + } + + fn uf_rebuild_rules(&self) -> Vec { + self.funcs + .iter() + .filter_map(|(_, info)| info.uf_rebuild_rule) + .collect() + } + /// Gives the user a handle to the underlying ExecutionState. Useful for staging updates /// to the database. /// @@ -1234,6 +1293,7 @@ struct FunctionInfo { can_subsume: bool, name: Arc, kind: FunctionKind, + uf_rebuild_rule: Option, } impl FunctionInfo { diff --git a/egglog-bridge/src/tests.rs b/egglog-bridge/src/tests.rs index ceb8dc14c..e6ac2b84c 100644 --- a/egglog-bridge/src/tests.rs +++ b/egglog-bridge/src/tests.rs @@ -19,7 +19,7 @@ use num_rational::Rational64; use crate::{ ColumnTy, DefaultVal, EGraph, FunctionConfig, FunctionId, MergeFn, ProofStore, QueryEntry, - Result, UfFunctionConfig, add_expressions, define_rule, + Result, UfFunctionConfig, UnionAction, add_expressions, define_rule, }; /// Run a simple associativity/commutativity test. In addition to testing that the rules properly @@ -154,6 +154,31 @@ fn uf_function_callback_inserts() -> Result<()> { Ok(()) } +#[test] +fn uf_function_rebuilds_from_global_uf() -> Result<()> { + let mut egraph = EGraph::default(); + let uf_func = egraph.add_uf_function(UfFunctionConfig { + name: "uf_rebuild".into(), + on_leader_change: None, + read_deps: Vec::new(), + write_deps: Vec::new(), + })?; + let union = UnionAction::new(&egraph); + let lhs = Value::from_usize(10); + let rhs = Value::from_usize(3); + egraph.with_execution_state(|state| { + union.union(state, lhs, rhs); + }); + egraph.flush_updates(); + + let canon_lhs = egraph.get_canon_in_uf(lhs); + let key = if canon_lhs != lhs { lhs } else { rhs }; + let expected = egraph.get_canon_in_uf(key); + let leader = egraph.lookup_id(uf_func, &[key]).unwrap(); + assert_eq!(leader, expected); + Ok(()) +} + #[test] fn uf_function_disallowed_in_merge() -> Result<()> { let mut egraph = EGraph::default(); From 8f58961faca6a5a0f8ca4359348ce03d828026cb Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Mon, 5 Jan 2026 23:36:00 -0800 Subject: [PATCH 5/6] Better tests, and expose an external function --- egglog-bridge/src/lib.rs | 37 ++++++++++++++++++++--- egglog-bridge/src/tests.rs | 61 ++++++++++++++++++++++++++++++++++++-- 2 files changed, 91 insertions(+), 7 deletions(-) diff --git a/egglog-bridge/src/lib.rs b/egglog-bridge/src/lib.rs index 3a664763d..25f35dec4 100644 --- a/egglog-bridge/src/lib.rs +++ b/egglog-bridge/src/lib.rs @@ -21,7 +21,7 @@ use crate::core_relations::{ BaseValue, BaseValueId, BaseValues, ColumnId, Constraint, ContainerValue, ContainerValues, CounterId, Database, DisplacedTable, DisplacedTableWithProvenance, ExecutionState, ExternalFunction, ExternalFunctionId, LeaderChange, MergeVal, Offset, PlanStrategy, - SortedWritesTable, TableId, TaggedRowBuffer, Value, WrappedTable, + SortedWritesTable, TableId, TaggedRowBuffer, Value, WrappedTable, make_external_func, }; use crate::numeric_id::{DenseIdMap, DenseIdMapWithReuse, IdVec, NumericId, define_id}; use egglog_core_relations as core_relations; @@ -802,7 +802,25 @@ impl EGraph { } /// Register a UF-backed function in this EGraph. - pub fn add_uf_function(&mut self, config: UfFunctionConfig) -> Result { + /// + /// UF Functions record the series of "leader changes" in the e-graph, where the leader is + /// always 'up to date' with the latest value. It has two columns: one key (the "displaced" + /// value), mapping to its current leader. Queries of this such a function, coupled with + /// semi-naive evaluation, allow for rules to "listen" on 'leader changes' in the underlying + /// union-find. This makes it more efficient, if less convenient, than a standard + /// quadratic-sized equivalence relation. + /// + /// The returned external function canonicalizes its single argument against the underlying + /// union-find (the more naive "equivalence relation" encoding, wrapped as a primitive call). + /// + /// Values in this union find are "coherent" with those of the underlying union-find in the + /// e-graph: rebuilding the e-graph also canonicalizes these union-finds with respect to the + /// global e-graph. What makes these union-finds useful is that they can store _more_ + /// equalities that are themselves "visible" to egglog rules: we use this in the term encoding. + pub fn add_uf_function( + &mut self, + config: UfFunctionConfig, + ) -> Result<(FunctionId, ExternalFunctionId)> { if self.tracing { return Err(FunctionConfigError::UfWithTracing.into()); } @@ -814,7 +832,7 @@ impl EGraph { } = config; let mut table = DisplacedTable::default(); if let Some(callback) = on_leader_change { - table.set_leader_change_callback(move |state, change| (callback)(state, change)); + table.set_leader_change_callback(move |state, change| callback(state, change)); } let name: Arc = name.into(); let table_id = self.db.add_table_named( @@ -823,6 +841,17 @@ impl EGraph { read_deps.iter().copied(), write_deps.iter().copied(), ); + let canon = + self.register_external_func(Box::new(make_external_func(move |state, vals| { + let [val] = vals else { + panic!("uf canonicalizer expected 1 value, got {vals:?}") + }; + let table = state.get_table(table_id); + let canon = table + .get_row_column(&[*val], ColumnId::new(1)) + .unwrap_or(*val); + Some(canon) + }))); let schema = vec![ColumnTy::Id, ColumnTy::Id]; let res = self.funcs.push(FunctionInfo { table: table_id, @@ -837,7 +866,7 @@ impl EGraph { }); let uf_rebuild_rule = self.uf_rebuild_rule(res); self.funcs[res].uf_rebuild_rule = Some(uf_rebuild_rule); - Ok(res) + Ok((res, canon)) } /// Run the given rules, returning whether the database changed. diff --git a/egglog-bridge/src/tests.rs b/egglog-bridge/src/tests.rs index e6ac2b84c..2e6eb439e 100644 --- a/egglog-bridge/src/tests.rs +++ b/egglog-bridge/src/tests.rs @@ -129,7 +129,7 @@ fn uf_function_callback_inserts() -> Result<()> { can_subsume: false, })?; let log_table_id = egraph.funcs[log_table].table; - let uf_func = egraph.add_uf_function(UfFunctionConfig { + let (uf_func, _uf_canon) = egraph.add_uf_function(UfFunctionConfig { name: "uf_cb".into(), on_leader_change: Some(Box::new(move |state, change| { state.stage_insert( @@ -157,7 +157,7 @@ fn uf_function_callback_inserts() -> Result<()> { #[test] fn uf_function_rebuilds_from_global_uf() -> Result<()> { let mut egraph = EGraph::default(); - let uf_func = egraph.add_uf_function(UfFunctionConfig { + let (uf_func, _uf_canon) = egraph.add_uf_function(UfFunctionConfig { name: "uf_rebuild".into(), on_leader_change: None, read_deps: Vec::new(), @@ -179,10 +179,65 @@ fn uf_function_rebuilds_from_global_uf() -> Result<()> { Ok(()) } +#[test] +fn uf_function_canon_prim_query() -> Result<()> { + let mut egraph = EGraph::default(); + let (uf_func, uf_canon) = egraph.add_uf_function(UfFunctionConfig { + name: "uf_canon".into(), + on_leader_change: None, + read_deps: Vec::new(), + write_deps: Vec::new(), + })?; + let input_table = egraph.add_table(FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Id], + default: DefaultVal::Fail, + merge: MergeFn::AssertEq, + name: "input".into(), + can_subsume: false, + })?; + let log_table = egraph.add_table(FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Id], + default: DefaultVal::Fail, + merge: MergeFn::AssertEq, + name: "log".into(), + can_subsume: false, + })?; + + let canon_rule = { + let mut rb = egraph.new_rule("uf canon prim", true); + let input: QueryEntry = rb.new_var(ColumnTy::Id).into(); + let input_val: QueryEntry = rb.new_var(ColumnTy::Id).into(); + let canon: QueryEntry = rb.new_var(ColumnTy::Id).into(); + rb.query_table(input_table, &[input.clone(), input_val], Some(false))?; + rb.query_prim(uf_canon, &[input.clone(), canon.clone()], ColumnTy::Id)?; + rb.set(log_table, &[input, canon]); + rb.build() + }; + + let lhs = Value::from_usize(10); + let rhs = Value::from_usize(3); + let untouched = Value::from_usize(77); + egraph.add_values(vec![ + (uf_func, vec![lhs, rhs]), + (input_table, vec![lhs, lhs]), + (input_table, vec![rhs, rhs]), + (input_table, vec![untouched, untouched]), + ]); + egraph.run_rules(&[canon_rule])?; + + let mut rows = Vec::new(); + egraph.for_each(log_table, |row| rows.push(row.vals.to_vec())); + assert_unordered_eq( + rows, + vec![vec![lhs, rhs], vec![rhs, rhs], vec![untouched, untouched]], + ); + Ok(()) +} + #[test] fn uf_function_disallowed_in_merge() -> Result<()> { let mut egraph = EGraph::default(); - let uf_func = egraph.add_uf_function(UfFunctionConfig { + let (uf_func, _uf_canon) = egraph.add_uf_function(UfFunctionConfig { name: "uf_func".into(), on_leader_change: None, read_deps: Vec::new(), From 8ff50f47f3d0db1f42c032f9ec05c4298c3726c1 Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Tue, 6 Jan 2026 09:04:46 -0800 Subject: [PATCH 6/6] initial draft of "custom UF tables" --- egglog-bridge/src/lib.rs | 4 + src/ast/desugar.rs | 3 +- src/ast/mod.rs | 55 +++++++- src/ast/parse.rs | 126 ++++++++++++++--- src/ast/proof_global_remover.rs | 1 + src/ast/remove_globals.rs | 1 + src/lib.rs | 144 +++++++++++++++----- src/prelude.rs | 1 + src/term_encoding.rs | 11 +- src/typechecking.rs | 141 ++++++++++++++++++- tests/fail-typecheck/uf_onchange_schema.egg | 5 + tests/fail-typecheck/uf_schema.egg | 3 + tests/uf_function.egg | 18 +++ 13 files changed, 452 insertions(+), 61 deletions(-) create mode 100644 tests/fail-typecheck/uf_onchange_schema.egg create mode 100644 tests/fail-typecheck/uf_schema.egg create mode 100644 tests/uf_function.egg diff --git a/egglog-bridge/src/lib.rs b/egglog-bridge/src/lib.rs index 25f35dec4..cf5ea203a 100644 --- a/egglog-bridge/src/lib.rs +++ b/egglog-bridge/src/lib.rs @@ -584,6 +584,10 @@ impl EGraph { self.db.get_table(self.funcs[table].table).len() } + pub fn table_id(&self, table: FunctionId) -> TableId { + self.funcs[table].table + } + /// Generate a proof explaining why a given term is in the database. /// /// # Errors diff --git a/src/ast/desugar.rs b/src/ast/desugar.rs index 01c60a470..a7eb21c2a 100644 --- a/src/ast/desugar.rs +++ b/src/ast/desugar.rs @@ -15,8 +15,9 @@ pub(crate) fn desugar_command( name, schema, merge, + impl_kind, } => vec![NCommand::Function(FunctionDecl::function( - span, name, schema, merge, + span, name, schema, merge, impl_kind, ))], Command::Constructor { span, diff --git a/src/ast/mod.rs b/src/ast/mod.rs index a54fb8efe..0661b3b61 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -118,6 +118,7 @@ where schema: f.schema.clone(), name: f.name.clone(), merge: f.merge.clone(), + impl_kind: f.impl_kind.clone(), }, }, GenericNCommand::AddRuleset(span, name) => { @@ -432,9 +433,8 @@ where /// a functional dependency (also called a primary key) on its inputs to one output. /// /// ```text - /// (function - /// (:on_merge >)? - /// (:merge )?) + /// (function + /// (:merge | :no-merge | (:impl displaced-union-find (:onchange )? (:canon-prim )?)) ) ///``` /// A function can have a `cost` for extraction. /// @@ -472,6 +472,7 @@ where name: String, schema: Schema, merge: Option>, + impl_kind: FunctionImpl, }, /// Using the `ruleset` command, defines a new @@ -729,13 +730,30 @@ where name, schema, merge, + impl_kind, } => { let name = sanitize_internal_name(name); write!(f, "(function {name} {schema}")?; - if let Some(merge) = &merge { - write!(f, " :merge {merge}")?; - } else { - write!(f, " :no-merge")?; + match impl_kind { + FunctionImpl::Default => { + if let Some(merge) = &merge { + write!(f, " :merge {merge}")?; + } else { + write!(f, " :no-merge")?; + } + } + FunctionImpl::DisplacedUnionFind { + onchange, + canon_prim, + } => { + write!(f, " :impl displaced-union-find")?; + if let Some(onchange) = onchange { + write!(f, " :onchange {onchange}")?; + } + if let Some(canon_prim) = canon_prim { + write!(f, " :canon-prim {canon_prim}")?; + } + } } write!(f, ")") } @@ -941,6 +959,21 @@ impl Display for FunctionSubtype { } } +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub enum FunctionImpl { + Default, + DisplacedUnionFind { + onchange: Option, + canon_prim: Option, + }, +} + +impl FunctionImpl { + pub fn is_uf(&self) -> bool { + matches!(self, FunctionImpl::DisplacedUnionFind { .. }) + } +} + /// Represents the declaration of a function /// directly parsed from source syntax. #[derive(Clone, Debug, PartialEq, Eq, Hash)] @@ -951,6 +984,7 @@ where { pub name: String, pub subtype: FunctionSubtype, + pub impl_kind: FunctionImpl, /// Untyped schema pub schema: Schema, /// Resolved schema after typechecking is stored here, otherwise "". @@ -1012,10 +1046,12 @@ impl FunctionDecl { name: String, schema: Schema, merge: Option>, + impl_kind: FunctionImpl, ) -> Self { Self { name, subtype: FunctionSubtype::Custom, + impl_kind, schema, resolved_schema: String::new(), merge, @@ -1037,6 +1073,7 @@ impl FunctionDecl { Self { name, subtype: FunctionSubtype::Constructor, + impl_kind: FunctionImpl::Default, resolved_schema: String::new(), schema, merge: None, @@ -1052,6 +1089,7 @@ impl FunctionDecl { Self { name, subtype: FunctionSubtype::Relation, + impl_kind: FunctionImpl::Default, schema: Schema { input, output: String::from("Unit"), @@ -1078,6 +1116,7 @@ where GenericFunctionDecl { name: self.name, subtype: self.subtype, + impl_kind: self.impl_kind, schema: self.schema, resolved_schema: self.resolved_schema, merge: self.merge.map(|expr| expr.visit_exprs(f)), @@ -1315,11 +1354,13 @@ where name, schema, merge, + impl_kind, } => GenericCommand::Function { span, name, schema, merge: merge.map(|expr| expr.make_unresolved()), + impl_kind, }, GenericCommand::AddRuleset(span, name) => GenericCommand::AddRuleset(span, name), GenericCommand::UnstableCombinedRuleset(span, name, others) => { diff --git a/src/ast/parse.rs b/src/ast/parse.rs index 87ebfd282..74c9064c1 100644 --- a/src/ast/parse.rs +++ b/src/ast/parse.rs @@ -1,5 +1,6 @@ //! Parse a string into egglog. +use super::FunctionImpl; use crate::util::INTERNAL_SYMBOL_PREFIX; use crate::*; use egglog_ast::generic_ast::*; @@ -285,26 +286,119 @@ impl Parser { datatypes: map_fallible(tail, self, Self::rec_datatype)?, }], "function" => match tail { - [name, inputs, output, rest @ ..] => vec![Command::Function { - name: name.expect_atom("function name")?, - schema: self.parse_schema(inputs, output)?, - merge: match self.parse_options(rest)?.as_slice() { - [(":no-merge", [])] => None, - [(":merge", [e])] => Some(self.parse_expr(e)?), - [] => { - return error!( - span, - "functions are required to specify merge behaviour" - ); + [name, inputs, output, rest @ ..] => { + let mut merge = None; + let mut merge_specified = false; + let mut onchange = None; + let mut canon_prim = None; + let mut impl_kind = FunctionImpl::Default; + + for (opt, args) in self.parse_options(rest)? { + match opt { + ":merge" => { + if merge_specified { + return error!(span, "merge option specified more than once"); + } + let [expr] = args else { + return error!(span, "usage: :merge "); + }; + merge = Some(self.parse_expr(expr)?); + merge_specified = true; + } + ":no-merge" => { + if merge_specified { + return error!(span, "merge option specified more than once"); + } + if !args.is_empty() { + return error!(span, "usage: :no-merge"); + } + merge = None; + merge_specified = true; + } + ":impl" => { + if !matches!(impl_kind, FunctionImpl::Default) { + return error!(span, "impl option specified more than once"); + } + let [impl_name] = args else { + return error!(span, "usage: :impl displaced-union-find"); + }; + let impl_name = impl_name.expect_atom("function implementation")?; + if impl_name != "displaced-union-find" { + return error!( + span, + "unsupported function implementation {impl_name}" + ); + } + impl_kind = FunctionImpl::DisplacedUnionFind { + onchange: None, + canon_prim: None, + }; + } + ":onchange" => { + let [rel] = args else { + return error!(span, "usage: :onchange "); + }; + if onchange.is_some() { + return error!(span, ":onchange specified more than once"); + } + onchange = Some(rel.expect_atom("onchange relation")?); + } + ":canon-prim" => { + let [prim] = args else { + return error!(span, "usage: :canon-prim "); + }; + if canon_prim.is_some() { + return error!(span, ":canon-prim specified more than once"); + } + canon_prim = Some(prim.expect_atom("canon-prim name")?); + } + _ => return error!(span, "could not parse function options"), } - _ => return error!(span, "could not parse function options"), - }, - span, - }], + } + + let impl_kind = match impl_kind { + FunctionImpl::Default => { + if onchange.is_some() || canon_prim.is_some() { + return error!( + span, + ":onchange and :canon-prim require :impl displaced-union-find" + ); + } + if !merge_specified { + return error!( + span, + "functions are required to specify merge behaviour" + ); + } + FunctionImpl::Default + } + FunctionImpl::DisplacedUnionFind { .. } => { + if merge_specified { + return error!( + span, + "displaced-union-find functions cannot specify merge behaviour" + ); + } + FunctionImpl::DisplacedUnionFind { + onchange, + canon_prim, + } + } + }; + + vec![Command::Function { + name: name.expect_atom("function name")?, + schema: self.parse_schema(inputs, output)?, + merge, + impl_kind, + span, + }] + } _ => { let a = "(function (*) :merge )"; let b = "(function (*) :no-merge)"; - return error!(span, "usages:\n{a}\n{b}"); + let c = "(function (*) :impl displaced-union-find (:onchange )? (:canon-prim )?)"; + return error!(span, "usages:\n{a}\n{b}\n{c}"); } }, "constructor" => match tail { diff --git a/src/ast/proof_global_remover.rs b/src/ast/proof_global_remover.rs index 89335e548..eba1b980a 100644 --- a/src/ast/proof_global_remover.rs +++ b/src/ast/proof_global_remover.rs @@ -84,6 +84,7 @@ fn remove_globals_cmd(cmd: ResolvedNCommand) -> Vec { let func_decl = ResolvedFunctionDecl { name: name.name, subtype: FunctionSubtype::Constructor, + impl_kind: FunctionImpl::Default, schema: Schema { input: vec![], output: ty.name().to_owned(), diff --git a/src/ast/remove_globals.rs b/src/ast/remove_globals.rs index cf04d27af..41f8695a7 100644 --- a/src/ast/remove_globals.rs +++ b/src/ast/remove_globals.rs @@ -102,6 +102,7 @@ impl GlobalRemover<'_> { let func_decl = ResolvedFunctionDecl { name: name.name, subtype: FunctionSubtype::Custom, + impl_kind: FunctionImpl::Default, schema: Schema { input: vec![], output: ty.name().to_owned(), diff --git a/src/lib.rs b/src/lib.rs index ce9254726..b6a1e00ea 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -503,7 +503,10 @@ impl EGraph { // NB: type-checking should already catch unbound variables here. _ => Err(TypeError::Unbound(resolved_var.name.clone(), span.clone()).into()), }, - GenericExpr::Call(_, ResolvedCall::Func(f), args) => { + GenericExpr::Call(span, ResolvedCall::Func(f), args) => { + if self.functions[&f.name].decl.impl_kind.is_uf() { + return Err(TypeError::UfFunctionInMerge(f.name.clone(), span.clone()).into()); + } let translated_args = args .iter() .map(|arg| self.translate_expr_to_mergefn(arg)) @@ -543,40 +546,115 @@ impl EGraph { .collect::, _>>()?; let output = get_sort(&decl.schema.output)?; - let can_subsume = match decl.subtype { - FunctionSubtype::Constructor => true, - FunctionSubtype::Relation => true, - FunctionSubtype::Custom => false, - }; + let (backend_id, can_subsume) = match &decl.impl_kind { + FunctionImpl::Default => { + let can_subsume = match decl.subtype { + FunctionSubtype::Constructor => true, + FunctionSubtype::Relation => true, + FunctionSubtype::Custom => false, + }; - use egglog_bridge::{DefaultVal, MergeFn}; - let backend_id = self - .backend - .add_table(egglog_bridge::FunctionConfig { - schema: input - .iter() - .chain([&output]) - .map(|sort| sort.column_ty(&self.backend)) - .collect(), - default: match decl.subtype { - FunctionSubtype::Constructor => DefaultVal::FreshId, - FunctionSubtype::Custom => DefaultVal::Fail, - FunctionSubtype::Relation => { - DefaultVal::Const(self.backend.base_values().get(())) + use egglog_bridge::{DefaultVal, MergeFn}; + let backend_id = self + .backend + .add_table(egglog_bridge::FunctionConfig { + schema: input + .iter() + .chain([&output]) + .map(|sort| sort.column_ty(&self.backend)) + .collect(), + default: match decl.subtype { + FunctionSubtype::Constructor => DefaultVal::FreshId, + FunctionSubtype::Custom => DefaultVal::Fail, + FunctionSubtype::Relation => { + DefaultVal::Const(self.backend.base_values().get(())) + } + }, + merge: match decl.subtype { + FunctionSubtype::Constructor => MergeFn::UnionId, + FunctionSubtype::Relation => MergeFn::AssertEq, + FunctionSubtype::Custom => match &decl.merge { + None => MergeFn::AssertEq, + Some(expr) => self.translate_expr_to_mergefn(expr)?, + }, + }, + name: decl.name.to_string(), + can_subsume, + }) + .map_err(|err| Error::BackendError(err.to_string()))?; + (backend_id, can_subsume) + } + FunctionImpl::DisplacedUnionFind { + onchange, + canon_prim, + } => { + let mut on_leader_change: Option = None; + let mut write_deps = Vec::new(); + if let Some(onchange) = onchange { + let onchange_func = self.functions.get(onchange).ok_or_else(|| { + Error::TypeError(TypeError::UnboundFunction( + onchange.clone(), + decl.span.clone(), + )) + })?; + let onchange_backend_id = onchange_func.backend_id; + write_deps.push(self.backend.table_id(onchange_backend_id)); + let table_action = + egglog_bridge::TableAction::new(&self.backend, onchange_backend_id) + .map_err(|err| Error::BackendError(err.to_string()))?; + let callback: egglog_bridge::LeaderChangeCallback = Box::new( + move |state: &mut ExecutionState, change: core_relations::LeaderChange| { + let mut action = table_action.clone(); + let unit = state.base_values().get(()); + let new_leader = change.new_leader(); + action.insert( + state, + [ + change.write_lhs, + change.write_rhs, + change.lhs_leader, + change.rhs_leader, + new_leader, + unit, + ] + .into_iter(), + ); + }, + ); + on_leader_change = Some(callback); + } + let (backend_id, canon_prim_id) = self + .backend + .add_uf_function(egglog_bridge::UfFunctionConfig { + name: decl.name.to_string(), + on_leader_change, + read_deps: Vec::new(), + write_deps, + }) + .map_err(|err| Error::BackendError(err.to_string()))?; + + match canon_prim { + Some(canon_prim) => { + let old = self + .type_info + .replace_primitive_external_id(canon_prim, canon_prim_id) + .ok_or_else(|| { + Error::TypeError(TypeError::Unbound( + canon_prim.clone(), + decl.span.clone(), + )) + })?; + if old != canon_prim_id { + self.backend.free_external_func(old); + } } - }, - merge: match decl.subtype { - FunctionSubtype::Constructor => MergeFn::UnionId, - FunctionSubtype::Relation => MergeFn::AssertEq, - FunctionSubtype::Custom => match &decl.merge { - None => MergeFn::AssertEq, - Some(expr) => self.translate_expr_to_mergefn(expr)?, - }, - }, - name: decl.name.to_string(), - can_subsume, - }) - .map_err(|err| Error::BackendError(err.to_string()))?; + None => { + self.backend.free_external_func(canon_prim_id); + } + } + (backend_id, false) + } + }; let function = Function { decl: decl.clone(), diff --git a/src/prelude.rs b/src/prelude.rs index 74dc0d31b..8191a900e 100644 --- a/src/prelude.rs +++ b/src/prelude.rs @@ -599,6 +599,7 @@ pub fn add_function( name: name.to_owned(), schema, merge, + impl_kind: FunctionImpl::Default, }]) } diff --git a/src/term_encoding.rs b/src/term_encoding.rs index 90828c8fa..809368ce1 100644 --- a/src/term_encoding.rs +++ b/src/term_encoding.rs @@ -1184,8 +1184,15 @@ pub fn command_supports_proof_encoding(command: &ResolvedCommand) -> bool { | GenericCommand::Relation { .. } | GenericCommand::Input { .. } => false, ResolvedCommand::Action(ResolvedAction::Let(_, _, expr)) => expr.output_type().is_eq_sort(), - // no-merge isn't supported right now - ResolvedCommand::Function { merge: None, .. } => false, + ResolvedCommand::Function { + impl_kind, merge, .. + } => { + if impl_kind.is_uf() { + return false; + } + // no-merge isn't supported right now + merge.is_some() + } // delete or subsume on custom functions isn't supported _ => true, } diff --git a/src/typechecking.rs b/src/typechecking.rs index 2ba7fcc98..c10fc4342 100644 --- a/src/typechecking.rs +++ b/src/typechecking.rs @@ -1,5 +1,7 @@ use crate::{ + constraint::AllEqualTypeConstraint, core::{CoreActionContext, CoreRule, GenericActionsExt}, + sort::UnitSort, *, }; use ast::{ResolvedAction, ResolvedExpr, ResolvedFact, ResolvedRule, ResolvedVar, Rule}; @@ -47,6 +49,29 @@ impl Debug for PrimitiveWithId { } } +#[derive(Clone)] +struct UfCanonPrimitive { + name: String, + sort: ArcSort, +} + +impl Primitive for UfCanonPrimitive { + fn name(&self) -> &str { + &self.name + } + + fn get_type_constraints(&self, span: &Span) -> Box { + AllEqualTypeConstraint::new(&self.name, span.clone()) + .with_all_arguments_sort(self.sort.clone()) + .with_exact_length(2) + .into_box() + } + + fn apply(&self, _exec_state: &mut ExecutionState, _args: &[Value]) -> Option { + panic!("uf canon primitive should be bound to an external function"); + } +} + /// Stores resolved typechecking information. #[derive(Clone, Default)] pub struct TypeInfo { @@ -154,7 +179,26 @@ impl EGraph { let command: ResolvedNCommand = match command { NCommand::Function(fdecl) => { - ResolvedNCommand::Function(self.type_info.typecheck_function(symbol_gen, fdecl)?) + let resolved = self.type_info.typecheck_function(symbol_gen, fdecl)?; + if let FunctionImpl::DisplacedUnionFind { + canon_prim: Some(canon_prim), + .. + } = &resolved.impl_kind + { + let Some(sort) = self.type_info.get_sort_by_name(&resolved.schema.output) + else { + return Err(TypeError::UndefinedSort( + resolved.schema.output.clone(), + resolved.span.clone(), + )); + }; + let prim = UfCanonPrimitive { + name: canon_prim.clone(), + sort: sort.clone(), + }; + self.add_primitive(prim); + } + ResolvedNCommand::Function(resolved) } NCommand::NormRule { rule } => ResolvedNCommand::NormRule { rule: self.type_info.typecheck_rule(symbol_gen, rule)?, @@ -422,13 +466,81 @@ impl TypeInfo { fdecl.span.clone(), )); } + if let FunctionImpl::DisplacedUnionFind { + canon_prim: Some(canon_prim), + .. + } = &fdecl.impl_kind + { + if self.sorts.contains_key(canon_prim) { + return Err(TypeError::SortAlreadyBound( + canon_prim.clone(), + fdecl.span.clone(), + )); + } + if self.func_types.contains_key(canon_prim) { + return Err(TypeError::FunctionAlreadyBound( + canon_prim.clone(), + fdecl.span.clone(), + )); + } + if self.is_primitive(canon_prim) { + return Err(TypeError::PrimitiveAlreadyBound( + canon_prim.clone(), + fdecl.span.clone(), + )); + } + } let ftype = self.function_to_functype(fdecl)?; - if self.func_types.insert(fdecl.name.clone(), ftype).is_some() { + if self.func_types.contains_key(&fdecl.name) { return Err(TypeError::FunctionAlreadyBound( fdecl.name.clone(), fdecl.span.clone(), )); } + if let FunctionImpl::DisplacedUnionFind { onchange, .. } = &fdecl.impl_kind { + if fdecl.merge.is_some() { + return Err(TypeError::UfFunctionMerge( + fdecl.name.clone(), + fdecl.span.clone(), + )); + } + if fdecl.schema.input.len() != 1 + || ftype.input.len() != 1 + || ftype.input[0].name() != ftype.output.name() + || !ftype.output.is_eq_sort() + { + return Err(TypeError::UfFunctionSchema( + fdecl.name.clone(), + fdecl.span.clone(), + )); + } + if let Some(onchange) = onchange { + let Some(rel_type) = self.func_types.get(onchange) else { + return Err(TypeError::UnboundFunction( + onchange.clone(), + fdecl.span.clone(), + )); + }; + if rel_type.subtype != FunctionSubtype::Relation { + return Err(TypeError::UfOnChangeNotRelation( + onchange.clone(), + fdecl.span.clone(), + )); + } + let out_name = ftype.output.name().to_owned(); + let valid_schema = rel_type.input.len() == 5 + && rel_type.output.name() == UnitSort.name() + && rel_type.input.iter().all(|sort| sort.name() == out_name); + if !valid_schema { + return Err(TypeError::UfOnChangeSchema( + onchange.clone(), + out_name, + fdecl.span.clone(), + )); + } + } + } + self.func_types.insert(fdecl.name.clone(), ftype); let mut bound_vars = IndexMap::default(); let output_type = self.sorts.get(&fdecl.schema.output).unwrap(); if fdecl.subtype == FunctionSubtype::Constructor && !output_type.is_eq_sort() { @@ -443,6 +555,7 @@ impl TypeInfo { Ok(ResolvedFunctionDecl { name: fdecl.name.clone(), subtype: fdecl.subtype, + impl_kind: fdecl.impl_kind.clone(), schema: fdecl.schema.clone(), resolved_schema: ResolvedCall::Func(self.func_types.get(&fdecl.name).unwrap().clone()), merge: match &fdecl.merge { @@ -683,6 +796,20 @@ impl TypeInfo { self.primitives.get(sym).map(Vec::as_slice) } + pub(crate) fn replace_primitive_external_id( + &mut self, + sym: &str, + new_id: ExternalFunctionId, + ) -> Option { + let prims = self.primitives.get_mut(sym)?; + if prims.len() != 1 { + return None; + } + let old = prims[0].1; + prims[0].1 = new_id; + Some(old) + } + pub fn is_primitive(&self, sym: &str) -> bool { self.primitives.contains_key(sym) || self.reserved_primitives.contains(sym) } @@ -745,6 +872,16 @@ pub enum TypeError { AlreadyDefined(String, Span), #[error("{1}\nThe output type of constructor function {0} must be sort")] ConstructorOutputNotSort(String, Span), + #[error("{1}\nDisplaced-union-find function {0} cannot specify merge behaviour.")] + UfFunctionMerge(String, Span), + #[error("{1}\nDisplaced-union-find function {0} must have schema (S) S for an EqSort.")] + UfFunctionSchema(String, Span), + #[error("{1}\n:onchange relation {0} must be a relation.")] + UfOnChangeNotRelation(String, Span), + #[error("{2}\n:onchange relation {0} must have schema ({1} {1} {1} {1} {1}) -> Unit.")] + UfOnChangeSchema(String, String, Span), + #[error("{1}\nMerge expressions cannot call displaced-union-find function {0}.")] + UfFunctionInMerge(String, Span), #[error("{1}\nValue lookup of non-constructor function {0} in rule is disallowed.")] LookupInRuleDisallowed(String, Span), #[error("All alternative definitions considered failed\n{}", .0.iter().map(|e| format!(" {e}\n")).collect::>().join(""))] diff --git a/tests/fail-typecheck/uf_onchange_schema.egg b/tests/fail-typecheck/uf_onchange_schema.egg new file mode 100644 index 000000000..b67b057bb --- /dev/null +++ b/tests/fail-typecheck/uf_onchange_schema.egg @@ -0,0 +1,5 @@ +(datatype Math (A)) + +(relation BadOnChange (Math Math Math Math)) + +(function __UF_Math (Math) Math :impl displaced-union-find :onchange BadOnChange) diff --git a/tests/fail-typecheck/uf_schema.egg b/tests/fail-typecheck/uf_schema.egg new file mode 100644 index 000000000..8d0009ae3 --- /dev/null +++ b/tests/fail-typecheck/uf_schema.egg @@ -0,0 +1,3 @@ +(datatype Math (A) (B)) + +(function __UF_Bad (Math Math) Math :impl displaced-union-find) diff --git a/tests/uf_function.egg b/tests/uf_function.egg new file mode 100644 index 000000000..df3581475 --- /dev/null +++ b/tests/uf_function.egg @@ -0,0 +1,18 @@ +(datatype Math + (A) + (B)) + +(relation __MathUfChange (Math Math Math Math Math)) + +(function __UF_Math (Math) Math :impl displaced-union-find + :onchange __MathUfChange + :canon-prim canon-uf-math) + +(let $a (A)) +(let $b (B)) + +(set (__UF_Math $a) $b) + +(check (= (__UF_Math (ordering-max $a $b)) (ordering-min $a $b))) +(check (= (canon-uf-math $a) (canon-uf-math $b))) +(check (__MathUfChange $a $b $a $b (ordering-min $a $b)))