diff --git a/egglog-experimental/src/fresh_macro.rs b/egglog-experimental/src/fresh_macro.rs index 50ce721..f1a3935 100644 --- a/egglog-experimental/src/fresh_macro.rs +++ b/egglog-experimental/src/fresh_macro.rs @@ -88,7 +88,7 @@ fn desugar_fresh_rule( name: constructor_name.clone(), schema: Schema { input: schema, - output: output_sort, + outputs: vec![output_sort], }, cost: first_opts.cost, unextractable: first_opts.unextractable, diff --git a/egglog-experimental/src/primitive.rs b/egglog-experimental/src/primitive.rs index d92704a..93d8cbe 100644 --- a/egglog-experimental/src/primitive.rs +++ b/egglog-experimental/src/primitive.rs @@ -263,6 +263,8 @@ fn required_context(egraph: &mut EGraph, expr: &ResolvedExpr) -> Context { } } ResolvedCall::Primitive(_) => Context::Pure, + // `values` builds/destructures a tuple value; it reads or writes no tables. + ResolvedCall::Values(_) => Context::Pure, ResolvedCall::Func(func) => match func.subtype { FunctionSubtype::Constructor => Context::Write, FunctionSubtype::Custom => Context::Read, diff --git a/egglog-experimental/src/set_cost.rs b/egglog-experimental/src/set_cost.rs index 20f3e5a..9dbc19d 100644 --- a/egglog-experimental/src/set_cost.rs +++ b/egglog-experimental/src/set_cost.rs @@ -114,7 +114,7 @@ impl Macro> for SetCostDeclarations { if !*unextractable { let cost_table_name = get_cost_table_name(name); let mut cost_table_schema = schema.clone(); - cost_table_schema.output = "i64".into(); + cost_table_schema.outputs = vec!["i64".into()]; cost_table_commands.push(Command::Function { span: span.clone(), name: cost_table_name, diff --git a/egglog-experimental/src/table_rows.rs b/egglog-experimental/src/table_rows.rs index dbd6fca..4e2b3bb 100644 --- a/egglog-experimental/src/table_rows.rs +++ b/egglog-experimental/src/table_rows.rs @@ -40,7 +40,7 @@ pub(crate) fn table_layout(egraph: &EGraph, name: &str, span: Span) -> Result (S, Proof)` + (parent + a proof `key = parent`) with a self-referential `:merge`, replacing the former + two-table `@UF_S (S S) -> Proof` relation plus `@UF_Sf` index. Constructor-view congruence + and the UF self-merge both keep the smaller e-class and stage an oriented displaced edge, + so the `single_parent`/`uf_function_index` rules and per-term self-loop seeds are gone. This + mirrors the term (non-proof) encoding's single-table `@UF : (S) -> S`. +- **Tuple-output functions.** A function may declare more than one output sort, e.g. + `(function interval (Math) (i64 i64) :merge (values (max old0 new0) (min old1 new1)))`. Such a + function stores its outputs as separate value columns; the functional dependency is + `keys -> (value0, value1, ...)`. Outputs are destructured in queries with + `(= (values lo hi) (interval x))`, written with `(set (interval x) (values 0 100))`, and merged + with a `(values ...)` clause whose `i`-th element merges column `i` using the bound variables + `old0`, `new0`, `old1`, `new1`, .... Tuple outputs are only allowed for plain functions (not + constructors, relations, or view tables) and are not supported by the term/proof encoding. +- **`:merge` action blocks.** A `:merge` may be a value-producing action block + `:merge (* )`: the actions run first (with `old`/`new` bound), then the + trailing expression is the merged value. Actions may be `let` (bind an intermediate value used by + later actions or the result), `set` (write another function), or `union` (unify two eclasses). A + bare `:merge ` (no actions) is unchanged. +- Built-in keywords (most command, action, and schedule heads such as `function`, `set`, `union`, + `rule`, `run`, ..., plus the tuple constructor `values`) are now reserved and may no longer be + used as user identifiers (function/sort/constructor/relation/variant names or variables). Names + starting with `:` may not be used as identifiers (definition names), since that prefix marks + option keywords (`:merge`, `:cost`, ...); it is still accepted in expression position so command + macros can consume their own option markers (e.g. `:until`). The common-word commands `input` and + `output` are only partially reserved: they remain usable as variables, but not as definition names + or as the head of a call expression. - Add typed `EGraph` extension state that clones with `EGraph` and is restored by `push`/`pop`. - Report full source file paths in egglog span and error messages. - Fix seminaive matching after nested containers rebuild in place by propagating dirty container ids through parent containers. diff --git a/egglog/core-relations/src/dependency_graph.rs b/egglog/core-relations/src/dependency_graph.rs index 3620d02..25d21c6 100644 --- a/egglog/core-relations/src/dependency_graph.rs +++ b/egglog/core-relations/src/dependency_graph.rs @@ -55,6 +55,14 @@ impl DependencyGraph { pub(crate) fn strata(&self) -> impl Iterator> { self.levels.iter().map(|(_, tables)| tables) } + /// Whether this table's merge reads another table (has a read dependency). Tables with read + /// dependencies sit above level 0; tables with only write dependencies (or none) are at level + /// 0. Used to decide whether the dependency-agnostic `merge_simple` fast path is safe. + pub(crate) fn has_read_deps(&self, table: TableId) -> bool { + self.to_level + .get(table) + .is_some_and(|level| *level != LevelId::new(0)) + } pub(crate) fn write_deps(&self, table: TableId) -> impl Iterator + '_ { self.write_deps .get(table) diff --git a/egglog/core-relations/src/free_join/mod.rs b/egglog/core-relations/src/free_join/mod.rs index d5f055d..4e4c476 100644 --- a/egglog/core-relations/src/free_join/mod.rs +++ b/egglog/core-relations/src/free_join/mod.rs @@ -523,7 +523,16 @@ impl Database { loop { to_merge.clear(); let to_merge_vec = self.notification_list.reset(); - if to_merge_vec.len() < 4 { + // `merge_simple` is faster but ignores read/write-dependency ordering, so it is only + // sound when no dirty table's merge reads another table (e.g. a `Construct`/`Function` + // lookup inside a `:merge`). With a read dependency, an unordered merge could observe a + // stale, not-yet-merged target — and a `Construct` could then mint a duplicate id — so + // fall back to the dependency-ordered strata path however few tables are dirty. + if to_merge_vec.len() < 4 + && !to_merge_vec + .iter() + .any(|table| self.deps.has_read_deps(*table)) + { ever_changed |= self.merge_simple(to_merge_vec); break; } @@ -607,7 +616,16 @@ impl Database { while !to_merge.is_empty() { for table_id in to_merge.iter().copied() { let mut info = self.tables.unwrap_val(table_id); - let mut es = ExecutionState::new(self.read_only_view(), Default::default()); + // Pre-seed the table's OWN buffer so a self-referential merge — one that stages a + // write back into its own table (e.g. the term encoder's `@UF` recursive + // parent-union) — can stage it. The table has been `unwrap_val`'d out of + // `self.tables`, so the lazy `new_buffer()` path (which indexes `self.tables`) would + // fail for the self id. Staged rows land in `pending_state` and are picked up on the + // next fixpoint iteration of this loop. Non-self-referential merges get an empty, + // never-touched buffer. (The strata path in `merge_all` handles this via write-deps.) + let mut bufs = DenseIdMap::default(); + bufs.insert(table_id, info.table.new_buffer()); + let mut es = ExecutionState::new(self.read_only_view(), bufs); changed |= info.table.merge(&mut es).added || es.changed; self.tables.insert(table_id, info); } @@ -625,10 +643,14 @@ impl Database { pub fn merge_table(&mut self, table: TableId) -> bool { let mut info = self.tables.unwrap_val(table); self.total_size_estimate = self.total_size_estimate.wrapping_sub(info.table.len()); - let table_changed = info.table.merge(&mut ExecutionState::new( - self.read_only_view(), - Default::default(), - )); + // Pre-seed the table's own write buffer so a self-referential merge can stage a write back + // into itself (the table has been `unwrap_val`'d out, so the lazy `new_buffer()` path would + // fail for the self id). This mirrors `merge_simple`; see the comment there. + let mut bufs = DenseIdMap::default(); + bufs.insert(table, info.table.new_buffer()); + let table_changed = info + .table + .merge(&mut ExecutionState::new(self.read_only_view(), bufs)); self.total_size_estimate = self.total_size_estimate.wrapping_add(info.table.len()); self.tables.insert(table, info); table_changed.added diff --git a/egglog/egglog-ast/src/generic_ast.rs b/egglog/egglog-ast/src/generic_ast.rs index cf8a29c..3ddd518 100644 --- a/egglog/egglog-ast/src/generic_ast.rs +++ b/egglog/egglog-ast/src/generic_ast.rs @@ -42,6 +42,20 @@ pub struct GenericActions>, ); +/// The `:merge` of a function: a *value-producing action block*. On a functional-dependency +/// conflict, the `actions` run (with `old`/`new` — or `old0`/`new0`/... for tuple outputs — bound +/// to the conflicting values), then `result` is evaluated to the merged value(s). A plain +/// `:merge ` has no actions (see [`GenericMerge::result_only`]). +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct GenericMerge +where + Head: Clone + Display, + Leaf: Clone + PartialEq + Eq + Display + Hash, +{ + pub actions: GenericActions, + pub result: GenericExpr, +} + #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub enum GenericAction where diff --git a/egglog/egglog-ast/src/generic_ast_helpers.rs b/egglog/egglog-ast/src/generic_ast_helpers.rs index e3b0b84..966d0c3 100644 --- a/egglog/egglog-ast/src/generic_ast_helpers.rs +++ b/egglog/egglog-ast/src/generic_ast_helpers.rs @@ -331,6 +331,67 @@ where } } +impl GenericMerge +where + Head: Clone + Display, + Leaf: Clone + PartialEq + Eq + Display + Hash, +{ + /// A merge with no actions: just a result expression (the `:merge ` form). + pub fn result_only(result: GenericExpr) -> Self { + Self { + actions: GenericActions::default(), + result, + } + } + + /// Applies `f` to every expression in the merge's actions and result. + pub fn visit_exprs( + self, + f: &mut impl FnMut(GenericExpr) -> GenericExpr, + ) -> Self { + Self { + actions: self.actions.visit_exprs(f), + result: self.result.visit_exprs(f), + } + } + + /// Applies the `head` and `leaf` mappings to every symbol in the merge. + pub fn map_symbols( + self, + head: &mut impl FnMut(Head) -> Head2, + leaf: &mut impl FnMut(Leaf) -> Leaf2, + ) -> GenericMerge + where + Head2: Clone + Display, + Leaf2: Clone + PartialEq + Eq + Display + Hash, + { + GenericMerge { + actions: self.actions.map_symbols(head, leaf), + result: self.result.map_symbols(head, leaf), + } + } +} + +impl Display for GenericMerge +where + Head: Clone + Display, + Leaf: Clone + PartialEq + Eq + Display + Hash, +{ + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + if self.actions.is_empty() { + // Result-only: `:merge `. + write!(f, "{}", self.result) + } else { + // Action block: `:merge (* )`. + write!(f, "(")?; + for action in self.actions.iter() { + write!(f, "{action} ")?; + } + write!(f, "{})", self.result) + } + } +} + impl GenericAction where Head: Clone + Display, diff --git a/egglog/egglog-bridge/examples/ac.rs b/egglog/egglog-bridge/examples/ac.rs index e6a83b1..7cdf69b 100644 --- a/egglog/egglog-bridge/examples/ac.rs +++ b/egglog/egglog-bridge/examples/ac.rs @@ -15,6 +15,8 @@ fn main() { let mut egraph = EGraph::default(); let int_base = egraph.base_values_mut().register_type::(); let num_table = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Base(int_base), ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -22,6 +24,8 @@ fn main() { can_subsume: false, }); let add_table = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id; 3], default: DefaultVal::FreshId, merge: MergeFn::UnionId, diff --git a/egglog/egglog-bridge/examples/math.rs b/egglog/egglog-bridge/examples/math.rs index 8c0eaa1..5055c31 100644 --- a/egglog/egglog-bridge/examples/math.rs +++ b/egglog/egglog-bridge/examples/math.rs @@ -20,6 +20,8 @@ fn main() { let string_ty = egraph.base_values_mut().register_type::<&'static str>(); // tables let diff = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -27,6 +29,8 @@ fn main() { can_subsume: false, }); let integral = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -35,6 +39,8 @@ fn main() { }); let add = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -42,6 +48,8 @@ fn main() { can_subsume: false, }); let sub = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -50,6 +58,8 @@ fn main() { }); let mul = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -58,6 +68,8 @@ fn main() { }); let div = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -66,6 +78,8 @@ fn main() { }); let pow = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -74,6 +88,8 @@ fn main() { }); let ln = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -82,6 +98,8 @@ fn main() { }); let sqrt = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -90,6 +108,8 @@ fn main() { }); let sin = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -98,6 +118,8 @@ fn main() { }); let cos = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -106,6 +128,8 @@ fn main() { }); let rat = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Base(rational_ty), ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -114,6 +138,8 @@ fn main() { }); let var = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Base(string_ty), ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, diff --git a/egglog/egglog-bridge/src/lib.rs b/egglog/egglog-bridge/src/lib.rs index ed42721..775296e 100644 --- a/egglog/egglog-bridge/src/lib.rs +++ b/egglog/egglog-bridge/src/lib.rs @@ -190,8 +190,18 @@ impl Default for EGraph { /// Properties of a function added to an [`EGraph`]. pub struct FunctionConfig { - /// The function's schema. The last column in the schema is the return type. + /// The function's schema: `n_keys` key (input) columns followed by [`FunctionConfig::n_vals`] + /// value (output) columns. pub schema: Vec, + /// The number of value (output) columns; the remaining leading columns are keys. Must be at + /// least 1 and at most `schema.len()`. A tuple-output function has more than one. + pub n_vals: usize, + /// Opt-in identity-column guard. `Some(k)` marks the first `k` value columns (`k` in + /// `1..=n_vals`) as *identity*: when a key collision leaves those columns unchanged the merge is + /// skipped and the existing row is kept (a subsume-flag change still applies). `None` (the + /// default) runs the merge on every collision, preserving arbitrary/non-idempotent merge + /// semantics. Only opt in for merges that are idempotent on equal inputs (`merge(x, x) == x`). + pub n_identity_vals: Option, /// The behavior of the function when lookups are made on keys not currently present. pub default: DefaultVal, /// How to resolve FD conflicts for the function. @@ -332,10 +342,7 @@ impl EGraph { let mut bufs = DenseIdMap::default(); for (func, row) in values.into_iter() { let table_info = &self.funcs[func]; - let schema_math = SchemaMath { - subsume: table_info.can_subsume, - func_cols: table_info.schema.len(), - }; + let schema_math = table_info.schema_math(); let table_id = table_info.table; extended_row.extend_from_slice(&row); schema_math.write_table_row( @@ -362,10 +369,7 @@ impl EGraph { /// This method panics if the values do not match the arity of the function. pub fn add_term(&mut self, func: FunctionId, inputs: &[Value]) -> Value { let info = &self.funcs[func]; - let schema_math = SchemaMath { - subsume: info.can_subsume, - func_cols: info.schema.len(), - }; + let schema_math = info.schema_math(); let mut extended_row = Vec::new(); extended_row.extend_from_slice(inputs); let res = self.fresh_id(); @@ -388,10 +392,7 @@ impl EGraph { /// (`key`). pub fn lookup_id(&self, func: FunctionId, key: &[Value]) -> Option { let info = &self.funcs[func]; - let schema_math = SchemaMath { - subsume: info.can_subsume, - func_cols: info.schema.len(), - }; + let schema_math = info.schema_math(); let table_id = info.table; let table = self.db.get_table(table_id); let row = table.get_row(key)?; @@ -438,10 +439,7 @@ impl EGraph { pub fn for_each_while(&self, table: FunctionId, mut f: impl FnMut(ScanEntry<'_>) -> bool) { let info = &self.funcs[table]; let table = self.funcs[table].table; - let schema_math = SchemaMath { - subsume: info.can_subsume, - func_cols: info.schema.len(), - }; + let schema_math = info.schema_math(); let imp = self.db.get_table(table); let all = imp.all(); let mut cur = Offset::new(0); @@ -504,9 +502,18 @@ impl EGraph { } /// Register a function in this EGraph. + /// The [`FunctionId`] the next [`EGraph::add_table`] will assign. Lets a caller build a + /// self-referential merge (one that names the table it is about to create) before calling + /// `add_table`; the id is deterministic as long as no other function is added in between. + pub fn peek_next_function_id(&self) -> FunctionId { + self.funcs.next_id() + } + pub fn add_table(&mut self, config: FunctionConfig) -> FunctionId { let FunctionConfig { schema, + n_vals, + n_identity_vals, default, merge, name, @@ -516,22 +523,75 @@ impl EGraph { !schema.is_empty(), "must have at least one column in schema" ); + // `n_vals` (value-column count) is declared by the caller, not inferred from the merge, so a + // tuple-output function whose merge isn't a top-level `Columns` can't silently mis-split its + // key and value columns. `to_callback` cross-checks that the merge produces exactly `n_vals` + // columns, and `check_value_col_indices` that its `OldCol`/`NewCol` stay in range. + assert!( + (1..=schema.len()).contains(&n_vals), + "function {name} declares {n_vals} value columns but has {} columns total", + schema.len() + ); + if let Some(k) = n_identity_vals { + assert!( + (1..=n_vals).contains(&k), + "function {name} declares {k} identity columns but has {n_vals} value columns" + ); + } + merge.check_value_col_indices(n_vals, &name); let to_rebuild: Vec = schema .iter() .enumerate() .filter(|(_, ty)| matches!(ty, ColumnTy::Id)) .map(|(i, _)| ColumnId::from_usize(i)) .collect(); + let n_keys = schema.len() - n_vals; let schema_math = SchemaMath { subsume: can_subsume, + n_keys, func_cols: schema.len(), + n_identity_vals, }; let n_args = schema_math.num_keys(); let n_cols = schema_math.table_columns(); let next_func_id = self.funcs.next_id(); + let name: Arc = name.into(); + // Knot-tying for a self-referential merge (e.g. the term encoder's single-table UF `@UF`, + // whose `:merge` does a `TableInsert` into its OWN table): merge resolution (`fill_deps` / + // `to_callback` -> `TableAction::new`) reads `self.funcs[self_id].table`, so the new + // `FunctionInfo` must already be in `self.funcs` before the merge is built. Reserve the + // table id (the next `add_table_named` assigns exactly this id) and push the `FunctionInfo` + // up front, then build the merge and the backing table. Non-self-referential merges resolve + // the same either way. + let table_id = self.db.next_table_id(); + let res = self.funcs.push(FunctionInfo { + table: table_id, + schema: schema.clone(), + n_keys, + n_identity_vals, + incremental_rebuild_rules: Default::default(), + nonincremental_rebuild_rule: RuleId::new(!0), + default_val: default, + can_subsume, + name: name.clone(), + }); + // Not a debug_assert: a mismatch here silently mis-maps every future reference to this + // function's table (see the knot-tying comment above), so pay the one-time check in release. + assert_eq!( + res, next_func_id, + "peek_next_function_id / add_table id reservation is out of sync" + ); + let mut read_deps = IndexSet::::new(); let mut write_deps = IndexSet::::new(); merge.fill_deps(self, &mut read_deps, &mut write_deps); + // A self-referential merge may only *write* back into its own table (the pre-seeded + // write buffer handles that); *reading* self during its own merge is unsupported and would + // otherwise panic deep inside merge execution. + assert!( + !read_deps.contains(&table_id), + "self-referential merge for `{name}` may only write to its own table, not read it" + ); let merge_fn = merge.to_callback(schema_math, &name, self); let table = SortedWritesTable::new( n_args, @@ -540,24 +600,16 @@ impl EGraph { to_rebuild, merge_fn, ); - let name: Arc = name.into(); - let table_id = self.db.add_table_named( + let assigned_table_id = self.db.add_table_named( table, name.clone(), read_deps.iter().copied(), write_deps.iter().copied(), ); - - let res = self.funcs.push(FunctionInfo { - table: table_id, - schema: schema.clone(), - incremental_rebuild_rules: Default::default(), - nonincremental_rebuild_rule: RuleId::new(!0), - default_val: default, - can_subsume, - name, - }); - debug_assert_eq!(res, next_func_id); + assert_eq!( + assigned_table_id, table_id, + "reserved table id did not match the id assigned by add_table_named" + ); let incremental_rebuild_rules = self.incremental_rebuild_rules(res, &schema); let nonincremental_rebuild_rule = self.nonincremental_rebuild(res, &schema); let info = &mut self.funcs[res]; @@ -995,6 +1047,11 @@ struct CachedPlanInfo { struct FunctionInfo { table: TableId, schema: Vec, + /// The number of key (input) columns. The remaining columns of `schema` are value/return + /// columns (one for most functions, more for tuple-output functions). + n_keys: usize, + /// Opt-in identity-column guard (see [`FunctionConfig::n_identity_vals`]). + n_identity_vals: Option, incremental_rebuild_rules: Vec, nonincremental_rebuild_rule: RuleId, default_val: DefaultVal, @@ -1002,6 +1059,17 @@ struct FunctionInfo { name: Arc, } +impl FunctionInfo { + fn schema_math(&self) -> SchemaMath { + SchemaMath { + subsume: self.can_subsume, + n_keys: self.n_keys, + func_cols: self.schema.len(), + n_identity_vals: self.n_identity_vals, + } + } +} + impl FunctionInfo { fn ret_ty(&self) -> ColumnTy { self.schema.last().copied().unwrap() @@ -1031,14 +1099,53 @@ pub enum MergeFn { /// The output of a merge is determined by looking up the value for the given function and the /// given arguments in the egraph. Function(FunctionId, Vec), - /// Always return the old value for the given function. + /// Always return the old value for the given function (this column's old value). Old, - /// Always return the new value for the given function. + /// Always return the new value for the given function (this column's new value). New, + /// The old value of the `i`th value column. Used by tuple-output merges, where a column's + /// merge may reference any output column of the old row. + OldCol(usize), + /// The new value of the `i`th value column. + NewCol(usize), + /// The value bound by a preceding `let` action in a [`MergeFn::Block`], identified by its slot + /// (0-based, assigned in block order). Reads the per-invocation merge environment. + LetVar(usize), /// Always overwrite the new value for the given function with a constant. This is more useful /// as a "base case" in a more complicated merge function (e.g. one that clamps a value between /// 1 and 100) than it is as a standalone merge function. Const(Value), + /// A merge for a tuple-output function: one [`MergeFn`] per value column. Each inner + /// `MergeFn` produces that column's merged value and may reference any column via + /// [`MergeFn::OldCol`] / [`MergeFn::NewCol`]. The length determines the number of value + /// columns of the function. + Columns(Vec), + /// Look up (minting on miss, like a constructor) the value of a function call, as a merge value + /// expression — e.g. building a proof term. `args` evaluate to the key columns; uses + /// `lookup_or_insert` on the target's own merge. + Lookup(FunctionId, Vec), + /// A value-producing action block: run `actions` for their effects, then evaluate `result` to + /// the merged value(s). This is the merge form for merges that perform side effects (e.g. + /// staging a union). Ordinary merges have no actions. Only valid at the top level of a merge. + Block { + actions: Vec, + result: Box, + }, +} + +/// A side effect run by a [`MergeFn::Block`] before its result is evaluated. +pub enum MergeAction { + /// `(set (f keys...) vals...)` inside a merge: stage a full row (`args` = keys ++ values) into + /// the function's table, respecting that table's own merge. + Set(FunctionId, Vec), + /// `(let x )` inside a merge: evaluate `value` and bind it to `slot` of the merge + /// environment (slots are assigned in block order), so later actions and the result can read it + /// via [`MergeFn::LetVar`]. + Let { slot: usize, value: MergeFn }, + /// `(union a b)` inside a merge: union the two eclasses (stages into the union-find). Only + /// meaningful for eq-sort outputs; action-block merges are rejected under proofs, so this never + /// runs against the proof union-find. + Union(MergeFn, MergeFn), } impl MergeFn { @@ -1064,7 +1171,50 @@ impl MergeFn { UnionId => { write_deps.insert(egraph.uf_table); } - AssertEq | Old | New | Const(..) => {} + Columns(cols) => { + cols.iter() + .for_each(|col| col.fill_deps(egraph, read_deps, write_deps)); + } + Lookup(func, args) => { + 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)); + } + Block { actions, result } => { + actions + .iter() + .for_each(|a| a.fill_deps(egraph, read_deps, write_deps)); + result.fill_deps(egraph, read_deps, write_deps); + } + AssertEq | Old | New | OldCol(..) | NewCol(..) | LetVar(..) | Const(..) => {} + } + } + + /// Assert that every `OldCol(i)`/`NewCol(i)` references a real value column (`i < n_vals`). + /// An out-of-range index would otherwise read a timestamp/subsume column or panic during merge + /// execution; this surfaces the mistake at table-creation time with a clear message. + fn check_value_col_indices(&self, n_vals: usize, name: &str) { + use MergeFn::*; + let check = |i: usize, kind: &str| { + assert!( + i < n_vals, + "merge for `{name}` references {kind}({i}), but the function has only {n_vals} value column(s)" + ); + }; + match self { + OldCol(i) => check(*i, "OldCol"), + NewCol(i) => check(*i, "NewCol"), + Primitive(_, args) | Function(_, args) | Columns(args) | Lookup(_, args) => args + .iter() + .for_each(|a| a.check_value_col_indices(n_vals, name)), + Block { actions, result } => { + actions + .iter() + .for_each(|a| a.check_value_col_indices(n_vals, name)); + result.check_value_col_indices(n_vals, name); + } + AssertEq | Old | New | UnionId | LetVar(..) | Const(..) => {} } } @@ -1074,20 +1224,68 @@ impl MergeFn { function_name: &str, egraph: &mut EGraph, ) -> Box { - let resolved = self.resolve(function_name, egraph); + // Split the merge into side-effect actions (run once) and one value expression per value + // column. Only a top-level `Block` carries actions; `Columns` gives the per-column values. + let (action_srcs, col_srcs): (&[MergeAction], Vec<&MergeFn>) = match self { + MergeFn::Block { actions, result } => (actions.as_slice(), result.columns()), + other => (&[], other.columns()), + }; + let actions: Vec = action_srcs + .iter() + .map(|a| a.resolve(function_name, egraph)) + .collect(); + let resolved: Vec = col_srcs + .iter() + .map(|c| c.resolve(function_name, egraph)) + .collect(); + assert_eq!( + resolved.len(), + schema_math.n_vals(), + "merge for {function_name} must have one entry per value column" + ); Box::new(move |state, cur, new, out| { + // Identity/payload columns: when a collision leaves every identity value column + // unchanged it is not a real value conflict, so keep the existing value columns (and + // skip the actions) instead of running the merge — which would, e.g., stage a spurious + // union. A subsume-flag change is still applied below, so this stays correct for + // subsumable functions. Inert unless the function declares payload columns. + let identity_unchanged = match schema_math.n_identity_vals { + Some(k) => { + let id_lo = schema_math.n_keys; + cur[id_lo..id_lo + k] == new[id_lo..id_lo + k] + } + None => false, + }; + let timestamp = new[schema_math.ts_col()]; + // Environment for `let`-bound values, filled by `Let` actions in slot order and read by + // `LetVar` in later actions / the result. Empty (no allocation) unless the block uses + // `let`. + let mut env = SmallVec::<[Value; 4]>::new(); + + // Run the block's side effects once, before computing the merged values. + if !identity_unchanged { + for action in &actions { + action.run(state, cur, new, schema_math.n_keys, timestamp, &mut env); + } + } + let mut changed = false; - let ret_val = { - let cur = cur[schema_math.ret_val_col()]; - let new = new[schema_math.ret_val_col()]; - let out = resolved.run(state, cur, new, timestamp); - changed |= cur != out; - out - }; + // Compute each merged value column. `cur` and `new` are full rows, so a tuple-output + // column's merge may reference any output column via `OldCol`/`NewCol`. + let mut merged_vals = SmallVec::<[Value; 4]>::new(); + for (i, col_merge) in resolved.iter().enumerate() { + let out_val = if identity_unchanged { + cur[schema_math.val_col(i)] + } else { + col_merge.run(state, cur, new, schema_math.n_keys, i, timestamp, &env) + }; + changed |= cur[schema_math.val_col(i)] != out_val; + merged_vals.push(out_val); + } let subsume = schema_math.subsume.then(|| { let cur = cur[schema_math.subsume_col()]; @@ -1098,12 +1296,15 @@ impl MergeFn { }); if changed { out.extend_from_slice(new); + for (i, val) in merged_vals.iter().enumerate() { + out[schema_math.val_col(i)] = *val; + } schema_math.write_table_row( out, RowVals { timestamp, subsume, - ret_val: Some(ret_val), + ret_val: None, }, ); } @@ -1112,11 +1313,29 @@ impl MergeFn { }) } + /// Resolve this merge into one [`ResolvedMergeFn`] per value column. Single-value merges + /// (the common case) resolve to a one-element vector; [`MergeFn::Columns`] resolves to one + /// entry per column. + /// The per-value-column value expressions of this merge (`Columns` → its entries; any other + /// merge is a single column). `Block` is peeled off by `to_callback` and never reaches here. + fn columns(&self) -> Vec<&MergeFn> { + match self { + MergeFn::Columns(cols) => cols.iter().collect(), + other => vec![other], + } + } + fn resolve(&self, function_name: &str, egraph: &mut EGraph) -> ResolvedMergeFn { match self { MergeFn::Const(v) => ResolvedMergeFn::Const(*v), MergeFn::Old => ResolvedMergeFn::Old, MergeFn::New => ResolvedMergeFn::New, + MergeFn::OldCol(i) => ResolvedMergeFn::OldCol(*i), + MergeFn::NewCol(i) => ResolvedMergeFn::NewCol(*i), + MergeFn::LetVar(slot) => ResolvedMergeFn::LetVar(*slot), + MergeFn::Columns(_) => { + panic!("nested Columns merge is not supported (Columns must be top-level)") + } MergeFn::AssertEq => ResolvedMergeFn::AssertEq { panic: egraph.new_panic(format!( "Illegal merge attempted for function {function_name}" @@ -1159,6 +1378,16 @@ impl MergeFn { .collect::>(), } } + MergeFn::Lookup(func, args) => ResolvedMergeFn::Lookup { + table: TableAction::new(egraph, *func), + args: args + .iter() + .map(|arg| arg.resolve(function_name, egraph)) + .collect::>(), + }, + MergeFn::Block { .. } => { + panic!("nested Block merge is not supported (Block must be top-level)") + } } } } @@ -1171,6 +1400,10 @@ enum ResolvedMergeFn { Const(Value), Old, New, + OldCol(usize), + NewCol(usize), + /// Read slot `.0` of the per-invocation merge environment (a preceding `let`). + LetVar(usize), AssertEq { panic: ExternalFunctionId, }, @@ -1187,15 +1420,157 @@ enum ResolvedMergeFn { args: Vec, panic: ExternalFunctionId, }, + /// Look up / mint a function value inside a merge value expression (mint on miss, like a + /// constructor), via `lookup_or_insert`. + Lookup { + table: TableAction, + args: Vec, + }, +} + +/// A resolved side effect run by a [`MergeFn::Block`] before its result is evaluated. +enum ResolvedMergeAction { + /// `(set (f keys...) vals...)`: stage the full row `args` into `table`, respecting its merge. + Set { + table: TableAction, + args: Vec, + }, + /// `(let x )`: evaluate `value` and push it onto the environment (its slot equals the + /// current environment length, since `let`s run in slot order). + Let { + slot: usize, + value: ResolvedMergeFn, + }, + /// `(union a b)`: stage a union of the two eclasses into the union-find. + Union { + a: ResolvedMergeFn, + b: ResolvedMergeFn, + uf_table: TableId, + }, +} + +impl MergeAction { + fn fill_deps( + &self, + egraph: &EGraph, + read_deps: &mut IndexSet, + write_deps: &mut IndexSet, + ) { + match self { + MergeAction::Set(func, args) => { + write_deps.insert(egraph.funcs[*func].table); + args.iter() + .for_each(|arg| arg.fill_deps(egraph, read_deps, write_deps)); + } + MergeAction::Let { value, .. } => value.fill_deps(egraph, read_deps, write_deps), + MergeAction::Union(a, b) => { + a.fill_deps(egraph, read_deps, write_deps); + b.fill_deps(egraph, read_deps, write_deps); + write_deps.insert(egraph.uf_table); + } + } + } + + fn check_value_col_indices(&self, n_vals: usize, name: &str) { + match self { + MergeAction::Set(_, args) => args + .iter() + .for_each(|a| a.check_value_col_indices(n_vals, name)), + MergeAction::Let { value, .. } => value.check_value_col_indices(n_vals, name), + MergeAction::Union(a, b) => { + a.check_value_col_indices(n_vals, name); + b.check_value_col_indices(n_vals, name); + } + } + } + + fn resolve(&self, function_name: &str, egraph: &mut EGraph) -> ResolvedMergeAction { + match self { + MergeAction::Set(func, args) => ResolvedMergeAction::Set { + table: TableAction::new(egraph, *func), + args: args + .iter() + .map(|arg| arg.resolve(function_name, egraph)) + .collect::>(), + }, + MergeAction::Let { slot, value } => ResolvedMergeAction::Let { + slot: *slot, + value: value.resolve(function_name, egraph), + }, + MergeAction::Union(a, b) => ResolvedMergeAction::Union { + a: a.resolve(function_name, egraph), + b: b.resolve(function_name, egraph), + uf_table: egraph.uf_table, + }, + } + } +} + +impl ResolvedMergeAction { + fn run( + &self, + state: &mut ExecutionState, + cur: &[Value], + new: &[Value], + n_keys: usize, + ts: Value, + env: &mut SmallVec<[Value; 4]>, + ) { + // Action value expressions use explicit `OldCol`/`NewCol`, so `self_col` is irrelevant + // (pass 0). + match self { + ResolvedMergeAction::Set { table, args } => { + // `insert` respects the target table's own merge. + let row = args + .iter() + .map(|arg| arg.run(state, cur, new, n_keys, 0, ts, env)) + .collect::>(); + table.insert(state, row.into_iter()); + } + ResolvedMergeAction::Let { slot, value } => { + let v = value.run(state, cur, new, n_keys, 0, ts, env); + // `let`s run in slot order, so the new binding lands at `slot`. + debug_assert_eq!(*slot, env.len()); + env.push(v); + } + ResolvedMergeAction::Union { a, b, uf_table } => { + let av = a.run(state, cur, new, n_keys, 0, ts, env); + let bv = b.run(state, cur, new, n_keys, 0, ts, env); + if av != bv { + state.stage_insert(*uf_table, &[av, bv, ts]); + } + } + } + } } impl ResolvedMergeFn { - fn run(&self, state: &mut ExecutionState, cur: Value, new: Value, ts: Value) -> Value { + /// Compute the merged value for value column `self_col`. + /// + /// `cur` and `new` are the full conflicting rows. `n_keys` is the number of key columns, so + /// value column `i` lives at `cur[n_keys + i]`. `Old`/`New` refer to `self_col`'s value, while + /// `OldCol`/`NewCol` reference an explicit column — this lets a tuple-output column's merge + /// read any other output column. `env` holds `let`-bound values (read by `LetVar`). + #[allow(clippy::too_many_arguments)] + fn run( + &self, + state: &mut ExecutionState, + cur: &[Value], + new: &[Value], + n_keys: usize, + self_col: usize, + ts: Value, + env: &[Value], + ) -> Value { match self { ResolvedMergeFn::Const(v) => *v, - ResolvedMergeFn::Old => cur, - ResolvedMergeFn::New => new, + ResolvedMergeFn::Old => cur[n_keys + self_col], + ResolvedMergeFn::New => new[n_keys + self_col], + ResolvedMergeFn::OldCol(i) => cur[n_keys + i], + ResolvedMergeFn::NewCol(i) => new[n_keys + i], + ResolvedMergeFn::LetVar(slot) => env[*slot], ResolvedMergeFn::AssertEq { panic } => { + let (cur, new) = (cur[n_keys + self_col], new[n_keys + self_col]); if cur != new { let res = state.call_external_func(*panic, &[]); assert_eq!(res, None); @@ -1203,6 +1578,7 @@ impl ResolvedMergeFn { cur } ResolvedMergeFn::UnionId { uf_table } => { + let (cur, new) = (cur[n_keys + self_col], new[n_keys + self_col]); if cur != new { state.stage_insert(*uf_table, &[cur, new, ts]); // We pick the minimum when unioning. This matches the original egglog @@ -1219,7 +1595,7 @@ impl ResolvedMergeFn { ResolvedMergeFn::Primitive { prim, args, panic } => { let args = args .iter() - .map(|arg| arg.run(state, cur, new, ts)) + .map(|arg| arg.run(state, cur, new, n_keys, self_col, ts, env)) .collect::>(); match state.call_external_func(*prim, &args) { @@ -1227,19 +1603,19 @@ impl ResolvedMergeFn { None => { let res = state.call_external_func(*panic, &[]); assert_eq!(res, None); - cur + cur[n_keys + self_col] } } } ResolvedMergeFn::Function { func, args, panic } => { // see github.com/egraphs-good/egglog/pull/287 - if cur == new { - return cur; + if cur[n_keys + self_col] == new[n_keys + self_col] { + return cur[n_keys + self_col]; } let args = args .iter() - .map(|arg| arg.run(state, cur, new, ts)) + .map(|arg| arg.run(state, cur, new, n_keys, self_col, ts, env)) .collect::>(); // Merge functions dispatch to another function that may be @@ -1250,9 +1626,19 @@ impl ResolvedMergeFn { func.lookup_or_insert(state, &args).unwrap_or_else(|| { let res = state.call_external_func(*panic, &[]); assert_eq!(res, None); - cur + cur[n_keys + self_col] }) } + ResolvedMergeFn::Lookup { table, args } => { + let key = args + .iter() + .map(|arg| arg.run(state, cur, new, n_keys, self_col, ts, env)) + .collect::>(); + // A constructor mints its (single) output on miss; no extra value columns. + table + .lookup_or_insert_multi(state, &key, &[]) + .unwrap_or(cur[n_keys + self_col]) + } } } } @@ -1290,10 +1676,7 @@ impl TableAction { }; TableAction { table: func_info.table, - table_math: SchemaMath { - func_cols: func_info.schema.len(), - subsume: func_info.can_subsume, - }, + table_math: func_info.schema_math(), default: match &func_info.default_val { DefaultVal::FreshId => Some(MergeVal::Counter(egraph.id_counter)), DefaultVal::Fail => None, @@ -1310,10 +1693,9 @@ impl TableAction { self.kind } - /// Number of input columns (schema minus the trailing output - /// column). + /// Number of input (key) columns. pub fn input_arity(&self) -> usize { - self.table_math.func_cols - 1 + self.table_math.num_keys() } /// Look up a row and return its return-value column, or `None` if the @@ -1388,7 +1770,10 @@ impl TableAction { let timestamp = MergeVal::Constant(Value::from_usize(state.read_counter(self.timestamp))); let mut merge_vals = SmallVec::<[MergeVal; 3]>::new(); + // Build just the non-key portion (single value + ts + subsume) of the row. + // `lookup_or_insert` is only used for single-value functions/constructors. SchemaMath { + n_keys: 0, func_cols: 1, ..self.table_math } @@ -1412,6 +1797,47 @@ impl TableAction { } } + /// Multi-value variant of [`TableAction::lookup_or_insert`] for a value-tuple constructor + /// `(children) -> (output, extra...)`: the first value column (`output`) is minted (the + /// configured `FreshId` default) and the rest are written from `provided_vals` (e.g. a proof). + /// Returns the minted `output`. + /// + /// Idempotent: an already-present key returns its existing `output` and writes nothing. Write + /// operation, only safe in action/merge contexts. + pub fn lookup_or_insert_multi( + &self, + state: &mut ExecutionState, + key: &[Value], + provided_vals: &[Value], + ) -> Option { + match self.default { + Some(default) => { + debug_assert_eq!( + self.table_math.n_vals(), + 1 + provided_vals.len(), + "lookup_or_insert_multi: provided_vals must fill every value \ + column except the minted first one" + ); + let timestamp = + MergeVal::Constant(Value::from_usize(state.read_counter(self.timestamp))); + // Non-key columns, in order: [output (minted), provided.., ts, subsume?]. + let mut merge_vals = SmallVec::<[MergeVal; 4]>::new(); + merge_vals.push(default); + merge_vals.extend(provided_vals.iter().map(|v| MergeVal::Constant(*v))); + merge_vals.push(timestamp); + if self.table_math.subsume { + merge_vals.push(MergeVal::Constant(NOT_SUBSUMED)); + } + // The first value column (the minted output) is at `ret_val_col()`. + Some( + state.predict_val(self.table, key, merge_vals.iter().copied()) + [self.table_math.ret_val_col()], + ) + } + None => self.lookup(state, key), + } + } + /// Insert a row into this table. pub fn insert(&self, state: &mut ExecutionState, row: impl Iterator) { let ts = Value::from_usize(state.read_counter(self.timestamp)); @@ -1617,16 +2043,24 @@ fn combine_subsumed(v1: Value, v2: Value) -> Value { /// Functions can have multiple "output columns" in the underlying core-relations layer depending /// on whether different features are enabled. Roughly, tables are laid out as: /// -/// > `[key0, ..., keyn, return value, timestamp, subsume?]` +/// > `[key0, ..., keyn, value0, ..., valuem, timestamp, subsume?]` /// -/// Where there are `n+1` key columns and columns marked with a question mark are optional, -/// depending on the egraph and table-level configuration. +/// Where there are `n+1` key columns, `m+1` value (return) columns, and columns marked with a +/// question mark are optional, depending on the egraph and table-level configuration. +/// +/// Most functions have a single value column (`m == 0`); tuple-output functions (declared with a +/// parenthesized list of output sorts) have more than one. The functional dependency is always +/// from the key columns to the value columns. #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] struct SchemaMath { /// Whether or not the table is enabled for subsumption. subsume: bool, - /// The number of columns in the function (including the return value). + /// The number of key (input) columns. + n_keys: usize, + /// The number of columns in the function (keys plus all value/return columns). func_cols: usize, + /// Opt-in identity-column guard (see [`FunctionConfig::n_identity_vals`]). + n_identity_vals: Option, } /// A struct containing possible non-key portions of a table row. To be used with @@ -1682,15 +2116,27 @@ impl SchemaMath { } fn num_keys(&self) -> usize { - self.func_cols - 1 + self.n_keys + } + + /// The number of value (return) columns. + fn n_vals(&self) -> usize { + self.func_cols - self.n_keys } fn table_columns(&self) -> usize { self.func_cols + 1 /* timestamp */ + if self.subsume { 1 } else { 0 } } + /// The column index of the `i`th value (return) column. + fn val_col(&self, i: usize) -> usize { + self.n_keys + i + } + + /// The first value column. For single-output functions this is *the* return value; + /// for tuple-output functions it is value column 0. fn ret_val_col(&self) -> usize { - self.func_cols - 1 + self.n_keys } fn ts_col(&self) -> usize { diff --git a/egglog/egglog-bridge/src/rule.rs b/egglog/egglog-bridge/src/rule.rs index 8a4a7da..ba4e6c7 100644 --- a/egglog/egglog-bridge/src/rule.rs +++ b/egglog/egglog-bridge/src/rule.rs @@ -357,14 +357,14 @@ impl RuleBuilder<'_> { let schema_math = if let Some(func) = func { let info = &self.egraph.funcs[func]; assert_eq!(info.schema.len(), entries.len()); - SchemaMath { - subsume: info.can_subsume, - func_cols: info.schema.len(), - } + info.schema_math() } else { SchemaMath { subsume: subsume_entry.is_some(), + n_keys: entries.len().saturating_sub(1), func_cols: entries.len(), + // No merge runs off this fallback schema, so no identity guard. + n_identity_vals: None, } }; schema_math.write_table_row( @@ -473,8 +473,8 @@ impl RuleBuilder<'_> { /// /// `entries` should match the number of keys to the function. pub fn subsume(&mut self, func: FunctionId, entries: &[QueryEntry]) { - // First, insert a subsumed value if the tuple is new. - let ret = self.lookup_with_subsumed( + // Ensure the row exists (panics otherwise); its value is re-read per column below. + let _ret = self.lookup_with_subsumed( func, entries, QueryEntry::Const { @@ -484,31 +484,34 @@ impl RuleBuilder<'_> { || "subsumed a nonextestent row!".to_string(), ); let info = &self.egraph.funcs[func]; - let schema_math = SchemaMath { - subsume: info.can_subsume, - func_cols: info.schema.len(), - }; + let schema_math = info.schema_math(); assert!(info.can_subsume); - assert_eq!(entries.len() + 1, info.schema.len()); + assert_eq!(entries.len(), schema_math.num_keys()); + let n_vals = schema_math.n_vals(); let entries = entries.to_vec(); let table = info.table; - let ret: QueryEntry = ret.into(); self.add_callback(move |inner, rb| { // Then, add a tuple subsuming the entry, but only if the entry isn't already subsumed. - // Look up the current subsume value. - let mut dst_entries = inner.convert_all(&entries); + let keys = inner.convert_all(&entries); let cur_subsume_val = rb.lookup( table, - &dst_entries, + &keys, ColumnId::from_usize(schema_math.subsume_col()), )?; + // Re-read every value column so subsumption preserves the whole row (tuple-output views + // carry more than one value, e.g. the e-class and its proof). + let mut dst_entries = keys.clone(); + for i in 0..n_vals { + let v = rb.lookup(table, &keys, ColumnId::from_usize(schema_math.val_col(i)))?; + dst_entries.push(v.into()); + } schema_math.write_table_row( &mut dst_entries, RowVals { timestamp: inner.next_ts(), subsume: Some(SUBSUMED.into()), - ret_val: Some(inner.convert(&ret)), + ret_val: None, }, ); rb.insert_if_eq( @@ -540,10 +543,7 @@ impl RuleBuilder<'_> { .to_var(); let table = info.table; let id_counter = self.query.id_counter; - let schema_math = SchemaMath { - subsume: info.can_subsume, - func_cols: info.schema.len(), - }; + let schema_math = info.schema_math(); let cb: BuildRuleCallback = match info.default_val { DefaultVal::Const(_) | DefaultVal::FreshId => { let wv: WriteVal = match &info.default_val { @@ -673,10 +673,7 @@ impl RuleBuilder<'_> { let info = &self.egraph.funcs[func]; let table = info.table; let entries = entries.to_vec(); - let schema_math = SchemaMath { - subsume: info.can_subsume, - func_cols: info.schema.len(), - }; + let schema_math = info.schema_math(); self.query.add_rule.push(Box::new(move |inner, rb| { let mut dst_vars = inner.convert_all(&entries); schema_math.write_table_row( diff --git a/egglog/egglog-bridge/src/tests.rs b/egglog/egglog-bridge/src/tests.rs index cbee0e8..a321637 100644 --- a/egglog/egglog-bridge/src/tests.rs +++ b/egglog/egglog-bridge/src/tests.rs @@ -19,8 +19,8 @@ use num_rational::Rational64; use once_cell::sync::Lazy; use crate::{ - ColumnTy, DefaultVal, EGraph, FunctionConfig, FunctionId, MergeFn, QueryEntry, add_expressions, - define_rule, + ColumnTy, DefaultVal, EGraph, FunctionConfig, FunctionId, MergeAction, MergeFn, QueryEntry, + add_expressions, define_rule, }; /// Run a simple associativity/commutativity test. @@ -33,6 +33,8 @@ fn ac_test(can_subsume: bool) { let mut egraph = EGraph::default(); let int_base = egraph.base_values_mut().register_type::(); let num_table = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Base(int_base), ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -40,6 +42,8 @@ fn ac_test(can_subsume: bool) { can_subsume, }); let add_table = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id; 3], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -111,6 +115,8 @@ fn ac_fail() { let int_base = egraph.base_values_mut().get_ty::(); let one = egraph.base_value_constant(1i64); let num_table = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Base(int_base), ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -118,6 +124,8 @@ fn ac_fail() { can_subsume: false, }); let add_table = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id; 3], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -207,6 +215,8 @@ fn math_test(mut egraph: EGraph, can_subsume: bool) { let string_ty = egraph.base_values_mut().register_type::<&'static str>(); // tables let diff = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -214,6 +224,8 @@ fn math_test(mut egraph: EGraph, can_subsume: bool) { can_subsume, }); let integral = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -221,6 +233,8 @@ fn math_test(mut egraph: EGraph, can_subsume: bool) { can_subsume, }); let add = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -228,6 +242,8 @@ fn math_test(mut egraph: EGraph, can_subsume: bool) { can_subsume, }); let sub = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -235,6 +251,8 @@ fn math_test(mut egraph: EGraph, can_subsume: bool) { can_subsume, }); let mul = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -242,6 +260,8 @@ fn math_test(mut egraph: EGraph, can_subsume: bool) { can_subsume, }); let div = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -249,6 +269,8 @@ fn math_test(mut egraph: EGraph, can_subsume: bool) { can_subsume, }); let pow = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -257,6 +279,8 @@ fn math_test(mut egraph: EGraph, can_subsume: bool) { }); let ln = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -264,6 +288,8 @@ fn math_test(mut egraph: EGraph, can_subsume: bool) { can_subsume, }); let sqrt = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -271,6 +297,8 @@ fn math_test(mut egraph: EGraph, can_subsume: bool) { can_subsume, }); let sin = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -278,6 +306,8 @@ fn math_test(mut egraph: EGraph, can_subsume: bool) { can_subsume, }); let cos = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -285,6 +315,8 @@ fn math_test(mut egraph: EGraph, can_subsume: bool) { can_subsume, }); let rat = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Base(rational_ty), ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -292,6 +324,8 @@ fn math_test(mut egraph: EGraph, can_subsume: bool) { can_subsume, }); let var = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Base(string_ty), ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -534,6 +568,8 @@ fn container_test() { let mut egraph = EGraph::default(); let int_base = egraph.base_values_mut().register_type::(); let num_table = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Base(int_base), ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -541,6 +577,8 @@ fn container_test() { can_subsume: false, }); let add_table = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id; 3], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -548,6 +586,8 @@ fn container_test() { can_subsume: false, }); let vec_table = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id; 2], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -719,6 +759,8 @@ fn basic_container() { fn run_query_prim_container_match_case(seminaive: bool, seed_canonical: bool) -> bool { let mut egraph = EGraph::default(); let k_table = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -726,6 +768,8 @@ fn run_query_prim_container_match_case(seminaive: bool, seed_canonical: bool) -> can_subsume: false, }); let w_table = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -733,6 +777,8 @@ fn run_query_prim_container_match_case(seminaive: bool, seed_canonical: bool) -> can_subsume: false, }); let l_table = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -821,6 +867,8 @@ fn rhs_only_rule() { let zero = egraph.base_values_mut().get(0i64); let one = egraph.base_values_mut().get(1i64); let num_table = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Base(int_base), ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -909,6 +957,8 @@ fn mergefn_arithmetic() { // Create a function with merge function (+ 1 (* old new)) // This uses nested MergeFn::Primitive with external functions to build the complex merge function let f_table = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Base(int_base), ColumnTy::Base(int_base)], default: DefaultVal::Fail, merge: MergeFn::Primitive( @@ -1005,6 +1055,8 @@ fn mergefn_nested_function() { // Create a function g that will be used in the merge function for f let g_table = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -1015,6 +1067,8 @@ fn mergefn_nested_function() { // 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 let f_table = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Base(int_base), ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::Function( @@ -1127,6 +1181,8 @@ fn constrain_prims_simple() { let int_base = egraph.base_values_mut().register_type::(); let bool_base = egraph.base_values_mut().register_type::(); let f_table = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Base(int_base), ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -1134,6 +1190,8 @@ fn constrain_prims_simple() { can_subsume: false, }); let g_table = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Base(int_base), ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -1220,6 +1278,8 @@ fn constrain_prims_abstract() { let mut egraph = EGraph::default(); let int_base = egraph.base_values_mut().register_type::(); let f_table = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Base(int_base), ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -1227,6 +1287,8 @@ fn constrain_prims_abstract() { can_subsume: false, }); let g_table = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Base(int_base), ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -1317,6 +1379,8 @@ fn basic_subsumption() { let mut egraph = EGraph::default(); let int_base = egraph.base_values_mut().register_type::(); let f_table = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Base(int_base), ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -1324,6 +1388,8 @@ fn basic_subsumption() { can_subsume: true, }); let g_table = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Base(int_base), ColumnTy::Id], default: DefaultVal::FreshId, merge: MergeFn::UnionId, @@ -1396,6 +1462,8 @@ fn basic_subsumption() { fn lookup_failure_panics() { let mut egraph = EGraph::default(); let f = egraph.add_table(FunctionConfig { + n_vals: 1, + n_identity_vals: None, schema: vec![ColumnTy::Id, ColumnTy::Id], default: DefaultVal::Fail, merge: MergeFn::UnionId, @@ -1505,6 +1573,160 @@ fn panic_functions_trigger_early_stop() { assert_eq!(channel.lock().unwrap().as_deref(), Some("lazy panic")); } +#[test] +fn self_referential_merge_union_find() { + // A merge that writes back into its OWN table, like the term encoding's single-table UF. On a + // conflicting parent it keeps the smaller endpoint and re-inserts the displaced edge into + // itself. Exercises `peek_next_function_id`, `MergeFn::TableInsert` into self, `Seq`, and the + // backend's self-write buffer pre-seed. + let mut egraph = EGraph::default(); + let int_base = egraph.base_values_mut().register_type::(); + let min_func = egraph.register_external_func(Box::new(core_relations::make_external_func( + |state, vals| { + let [a, b] = vals else { return None }; + let (a, b) = ( + state.base_values().unwrap::(*a), + state.base_values().unwrap::(*b), + ); + Some(state.base_values().get::(a.min(b))) + }, + ))); + let max_func = egraph.register_external_func(Box::new(core_relations::make_external_func( + |state, vals| { + let [a, b] = vals else { return None }; + let (a, b) = ( + state.base_values().unwrap::(*a), + state.base_values().unwrap::(*b), + ); + Some(state.base_values().get::(a.max(b))) + }, + ))); + + // The merge references the table itself, so reserve its id before creating it. + let uf_id = egraph.peek_next_function_id(); + let min = || MergeFn::Primitive(min_func, vec![MergeFn::Old, MergeFn::New]); + let max = MergeFn::Primitive(max_func, vec![MergeFn::Old, MergeFn::New]); + let uf = egraph.add_table(FunctionConfig { + schema: vec![ColumnTy::Base(int_base), ColumnTy::Base(int_base)], + n_vals: 1, + // The single parent column is the identity: the guard skips the merge when the parent is + // unchanged, so the body stages the displaced edge unconditionally. + n_identity_vals: Some(1), + default: DefaultVal::Fail, + merge: MergeFn::Block { + actions: vec![MergeAction::Set(uf_id, vec![max, min()])], + result: Box::new(min()), + }, + name: "uf".into(), + can_subsume: false, + }); + assert_eq!(uf, uf_id, "peeked id must match the id add_table assigns"); + + let set_parent = |egraph: &mut EGraph, child: i64, parent: i64| { + let (c, p) = ( + egraph.base_value_constant(child), + egraph.base_value_constant(parent), + ); + let r = { + let mut rb = egraph.new_rule("set", true); + rb.set(uf, &[c, p]); + rb.build() + }; + egraph.run_rules(&[r]).unwrap(); + }; + let parent_of = |egraph: &mut EGraph, node: i64| -> Option { + let k = egraph.base_values_mut().get(node); + egraph + .lookup_id(uf, &[k]) + .map(|v| egraph.base_values().unwrap::(v)) + }; + + // uf[5] = 3, then uf[5] = 1: the conflict keeps min (1) and re-inserts the displaced edge 3->1. + set_parent(&mut egraph, 5, 3); + set_parent(&mut egraph, 5, 1); + assert_eq!( + parent_of(&mut egraph, 5), + Some(1), + "5 points at the smaller endpoint" + ); + assert_eq!( + parent_of(&mut egraph, 3), + Some(1), + "displaced edge 3->1 must be re-inserted into uf itself (self-write)" + ); + + // A conflict where the new parent is larger: uf[3] = 2 keeps 1 and re-inserts 2->1. + set_parent(&mut egraph, 3, 2); + assert_eq!(parent_of(&mut egraph, 3), Some(1)); + assert_eq!( + parent_of(&mut egraph, 2), + Some(1), + "displaced edge 2->1 re-inserted" + ); +} + +#[test] +fn identity_column_guard_skips_payload_only_conflicts() { + // A function with an identity column (col 0) and a payload column (col 1): a collision that + // changes only the payload keeps the existing row (the merge is skipped); a collision that + // changes the identity runs the merge. + let mut egraph = EGraph::default(); + let int_base = egraph.base_values_mut().register_type::(); + let f = egraph.add_table(FunctionConfig { + schema: vec![ + ColumnTy::Base(int_base), // key + ColumnTy::Base(int_base), // value col 0 (identity) + ColumnTy::Base(int_base), // value col 1 (payload) + ], + n_vals: 2, + // col 0 is identity, col 1 is payload. + n_identity_vals: Some(1), + default: DefaultVal::Fail, + // On a real (identity) conflict, take the new row. + merge: MergeFn::Columns(vec![MergeFn::NewCol(0), MergeFn::NewCol(1)]), + name: "f".into(), + can_subsume: false, + }); + + let set = |egraph: &mut EGraph, a: i64, b: i64| { + let (k, va, vb) = ( + egraph.base_value_constant(1i64), + egraph.base_value_constant(a), + egraph.base_value_constant(b), + ); + let r = { + let mut rb = egraph.new_rule("s", true); + rb.set(f, &[k, va, vb]); + rb.build() + }; + egraph.run_rules(&[r]).unwrap(); + }; + let read_row = |egraph: &EGraph| -> (i64, i64) { + let mut out = None; + egraph.for_each(f, |row| { + out = Some(( + egraph.base_values().unwrap::(row.vals[1]), + egraph.base_values().unwrap::(row.vals[2]), + )); + }); + out.unwrap() + }; + + set(&mut egraph, 10, 100); + set(&mut egraph, 10, 200); // identity (10) unchanged -> keep old, payload stays 100 + assert_eq!( + read_row(&egraph), + (10, 100), + "payload-only change should keep the existing row" + ); + set(&mut egraph, 20, 300); // identity 10 -> 20 -> merge runs, takes new + assert_eq!( + read_row(&egraph), + (20, 300), + "an identity change should run the merge" + ); +} + const _: () = { const fn assert_send() {} assert_send::() diff --git a/egglog/src/ast/desugar.rs b/egglog/src/ast/desugar.rs index 11614aa..d0bb90e 100644 --- a/egglog/src/ast/desugar.rs +++ b/egglog/src/ast/desugar.rs @@ -71,6 +71,7 @@ pub(crate) fn desugar_command( presort_and_args: None, uf: None, proof_func: None, + proof_ctors: None, unionable: true, }); } @@ -91,6 +92,7 @@ pub(crate) fn desugar_command( presort_and_args: Some((sort, args)), uf: None, proof_func: None, + proof_ctors: None, unionable: true, }); } @@ -106,7 +108,7 @@ pub(crate) fn desugar_command( variant.name, Schema { input: variant.types, - output: datatype.clone(), + outputs: vec![datatype.clone()], }, variant.cost, false, @@ -144,6 +146,7 @@ pub(crate) fn desugar_command( presort_and_args, uf, proof_func, + proof_ctors, unionable, } => vec![NCommand::Sort { span, @@ -151,6 +154,7 @@ pub(crate) fn desugar_command( presort_and_args, uf, proof_func, + proof_ctors, unionable, }], Command::AddRuleset(span, name) => vec![NCommand::AddRuleset(span, name)], @@ -234,6 +238,7 @@ fn desugar_prove(parser: &mut Parser, span: Span, query: Vec) -> Vec) -> Vec) -> Vec) -> Vec, + /// The global proof-constructor names `(congr, trans, sym)`, recorded once on the `Proof` + /// sort by proof desugaring so a desugared program re-parses self-contained (the native + /// congruence merge needs them). See `:internal-proof-names`. + proof_ctors: Option<(String, String, String)>, /// Whether values of this sort can be unioned. /// Defaults to true for user-defined sorts. /// Set to false for relations and term tables that should not allow union. @@ -116,6 +121,7 @@ where presort_and_args, uf, proof_func, + proof_ctors, unionable, } => GenericCommand::Sort { span: span.clone(), @@ -123,6 +129,7 @@ where presort_and_args: presort_and_args.clone(), uf: uf.clone(), proof_func: proof_func.clone(), + proof_ctors: proof_ctors.clone(), unionable: *unionable, }, GenericNCommand::Function(f) => match f.subtype { @@ -242,6 +249,7 @@ where presort_and_args, uf, proof_func, + proof_ctors, unionable, } => GenericNCommand::Sort { span, @@ -249,6 +257,7 @@ where presort_and_args, uf, proof_func, + proof_ctors, unionable, }, GenericNCommand::Function(func) => GenericNCommand::Function(func.visit_exprs(f)), @@ -558,6 +567,10 @@ where /// The name of the proof function for this sort. /// Set by proof desugaring to record where proofs are stored for this sort. proof_func: Option, + /// The global proof-constructor names `(congr, trans, sym)`, recorded once on the `Proof` + /// sort by proof desugaring so a desugared program re-parses self-contained. See + /// `:internal-proof-names`. + proof_ctors: Option<(String, String, String)>, /// Whether values of this sort can be unioned. /// Defaults to true for user-defined sorts. /// Set to false for relations and term tables that should not allow union. @@ -656,8 +669,8 @@ where inputs: Vec, }, - /// The `function` command declare an egglog custom function, which is a database table with a - /// a functional dependency (also called a primary key) on its inputs to one output. + /// The `function` command declares an egglog custom function, which is a database table with a + /// functional dependency (also called a primary key) on its inputs to its output(s). /// /// ```text /// (function @@ -666,6 +679,13 @@ where ///``` /// A function can have a `cost` for extraction. /// + /// The output of a function is usually a single sort, but may be a parenthesized list of + /// sorts, e.g. `(function f (Math) (i64 i64) ...)`, declaring a *tuple-output* function whose + /// functional dependency maps the inputs to a tuple of value columns. Tuple outputs are + /// destructured in queries with `(= (values a b) (f x))`, written with + /// `(set (f x) (values a b))`, and merged with a `(values e0 e1 ...)` clause where `ei` merges + /// column `i` using the bound variables `old0`, `new0`, `old1`, `new1`, .... + /// /// Finally, it can have a `merge` and `on_merge`, which are triggered when /// the function dependency is violated. /// In this case, the merge expression determines which of the two outputs @@ -699,7 +719,7 @@ where span: Span, name: String, schema: Schema, - merge: Option>, + merge: Option>, hidden: bool, let_binding: bool, term_constructor: Option, @@ -956,6 +976,7 @@ where presort_and_args: None, uf, proof_func, + proof_ctors, .. } => { write!(f, "(sort {name}")?; @@ -965,6 +986,9 @@ where if let Some(pf) = proof_func { write!(f, " :internal-proof-func {pf}")?; } + if let Some((congr, trans, sym)) = proof_ctors { + write!(f, " :internal-proof-names {congr} {trans} {sym}")?; + } write!(f, ")") } GenericCommand::Sort { @@ -1249,7 +1273,7 @@ where pub schema: Schema, /// Resolved schema after typechecking is stored here, otherwise "". pub resolved_schema: Head, - pub merge: Option>, + pub merge: Option>, pub cost: Option, pub unextractable: bool, /// Hidden functions are excluded from print-size output. @@ -1289,18 +1313,54 @@ impl Display for Variant { #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct Schema { pub input: Vec, - pub output: String, + /// The output (value-column) sorts, primary first. A tuple-output function (declared with a + /// parenthesized list, e.g. `(function f (Math) (i64 i64) ...)`) has more than one; ordinary + /// functions have exactly one. Always non-empty. + pub outputs: Vec, } impl Display for Schema { fn fmt(&self, f: &mut Formatter) -> std::fmt::Result { - write!(f, "({}) {}", ListDisplay(&self.input, " "), self.output) + if self.is_tuple_output() { + write!( + f, + "({}) ({})", + ListDisplay(&self.input, " "), + ListDisplay(&self.outputs, " ") + ) + } else { + write!(f, "({}) {}", ListDisplay(&self.input, " "), self.output()) + } } } impl Schema { pub fn new(input: Vec, output: String) -> Self { - Self { input, output } + Self { + input, + outputs: vec![output], + } + } + + /// Construct a schema with one or more output sorts. `outputs` must be non-empty. + pub fn new_tuple(input: Vec, outputs: Vec) -> Self { + assert!(!outputs.is_empty(), "schema must have at least one output"); + Self { input, outputs } + } + + /// The primary (first) output sort. + pub fn output(&self) -> &String { + &self.outputs[0] + } + + /// The number of output (value) columns. + pub fn num_outputs(&self) -> usize { + self.outputs.len() + } + + /// Whether this function has more than one output column. + pub fn is_tuple_output(&self) -> bool { + self.outputs.len() > 1 } } @@ -1310,7 +1370,7 @@ impl FunctionDecl { span: Span, name: String, schema: Schema, - merge: Option>, + merge: Option>, ) -> Self { Self { name, @@ -1385,7 +1445,7 @@ pub struct Facts(pub Vec>); impl Facts where - Head: Clone + Display, + Head: Clone + Display + HeadOps, Leaf: Clone + PartialEq + Eq + Display + Hash, { /// Flattens a list of facts into a Query. @@ -1406,21 +1466,66 @@ where for fact in self.0.iter() { match fact { GenericFact::Eq(span, e1, e2) => { - let mut to_equate = vec![]; - let mut process = |expr: &GenericExpr| { - let (child_atoms, expr) = expr.to_query(typeinfo, fresh_gen); - atoms.extend(child_atoms); - to_equate.push(expr.get_corresponding_var_or_lit(typeinfo)); - expr - }; - let e1 = process(e1); - let e2 = process(e2); - atoms.push(GenericAtom { - span: span.clone(), - head: HeadOrEq::Eq, - args: to_equate, - }); - new_body.push(GenericFact::Eq(span.clone(), e1, e2)); + // Tuple destructure: `(= (values v...) (f a...))` (in either order), where `f` + // is a tuple-output function. This lowers directly to the function atom + // `f(a..., v...)` — the `values` targets become the function's output columns. + if let Some(td) = match_tuple_destructure(e1, e2, typeinfo) { + let mut atom_args = vec![]; + let mut mapped_inputs = vec![]; + for arg in td.func_args { + let (child_atoms, mexpr) = arg.to_query(typeinfo, fresh_gen); + atoms.extend(child_atoms); + atom_args.push(mexpr.get_corresponding_var_or_lit(typeinfo)); + mapped_inputs.push(mexpr); + } + let mut mapped_values = vec![]; + for v in td.values_args { + let (child_atoms, mexpr) = v.to_query(typeinfo, fresh_gen); + atoms.extend(child_atoms); + atom_args.push(mexpr.get_corresponding_var_or_lit(typeinfo)); + mapped_values.push(mexpr); + } + atoms.push(GenericAtom { + span: td.func_span.clone(), + head: HeadOrEq::Head(td.func_head.clone()), + args: atom_args, + }); + // Reconstruct a mapped `Eq` fact so type annotation can produce the + // resolved form (which is re-lowered the same way during canonicalization). + let values_mapped = GenericExpr::Call( + td.values_span.clone(), + CorrespondingVar::new( + td.values_head.clone(), + fresh_gen.fresh(td.func_head), + ), + mapped_values, + ); + let func_mapped = GenericExpr::Call( + td.func_span.clone(), + CorrespondingVar::new( + td.func_head.clone(), + fresh_gen.fresh(td.func_head), + ), + mapped_inputs, + ); + new_body.push(GenericFact::Eq(span.clone(), values_mapped, func_mapped)); + } else { + let mut to_equate = vec![]; + let mut process = |expr: &GenericExpr| { + let (child_atoms, expr) = expr.to_query(typeinfo, fresh_gen); + atoms.extend(child_atoms); + to_equate.push(expr.get_corresponding_var_or_lit(typeinfo)); + expr + }; + let e1 = process(e1); + let e2 = process(e2); + atoms.push(GenericAtom { + span: span.clone(), + head: HeadOrEq::Eq, + args: to_equate, + }); + new_body.push(GenericFact::Eq(span.clone(), e1, e2)); + } } GenericFact::Fact(expr) => { let (child_atoms, expr) = expr.to_query(typeinfo, fresh_gen); @@ -1433,6 +1538,47 @@ where } } +/// The pieces of a recognized `(= (values v...) (f a...))` tuple destructure. +struct TupleDestructure<'a, Head, Leaf> { + values_head: &'a Head, + values_span: &'a Span, + values_args: &'a [GenericExpr], + func_head: &'a Head, + func_span: &'a Span, + func_args: &'a [GenericExpr], +} + +/// Recognize `(= (values v...) (f a...))` in either argument order, where `f` is a tuple-output +/// function. The arity (number of `values` targets vs. output columns) is validated later by the +/// usual type-checking arity constraints. +fn match_tuple_destructure<'a, Head, Leaf>( + e1: &'a GenericExpr, + e2: &'a GenericExpr, + typeinfo: &TypeInfo, +) -> Option> +where + Head: Clone + Display + HeadOps, + Leaf: Clone + PartialEq + Eq + Display + Hash, +{ + for (a, b) in [(e1, e2), (e2, e1)] { + if let GenericExpr::Call(values_span, values_head, values_args) = a + && values_head.is_values() + && let GenericExpr::Call(func_span, func_head, func_args) = b + && func_head.is_tuple_output(typeinfo) + { + return Some(TupleDestructure { + values_head, + values_span, + values_args, + func_head, + func_span, + func_args, + }); + } + } + None +} + #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct CorrespondingVar where @@ -1470,6 +1616,9 @@ pub type Actions = GenericActions; pub(crate) type ResolvedActions = GenericActions; pub(crate) type MappedActions = GenericActions, Leaf>; +pub type Merge = GenericMerge; +pub type ResolvedMerge = GenericMerge; + pub type Rule = GenericRule; pub(crate) type ResolvedRule = GenericRule; @@ -1600,6 +1749,7 @@ where presort_and_args, uf, proof_func, + proof_ctors, unionable, } => GenericCommand::Sort { span, @@ -1607,6 +1757,7 @@ where presort_and_args, uf: uf.map(&mut *fun), proof_func: proof_func.map(&mut *fun), + proof_ctors: proof_ctors.map(|(c, t, s)| (fun(c), fun(t), fun(s))), unionable, }, GenericCommand::Datatype { @@ -1674,7 +1825,7 @@ where name: fun(name), schema: Schema { input: schema.input.into_iter().map(&mut *fun).collect(), - output: fun(schema.output), + outputs: schema.outputs.into_iter().map(&mut *fun).collect(), }, cost, unextractable, @@ -1701,7 +1852,7 @@ where name: fun(name), schema: Schema { input: schema.input.into_iter().map(&mut *fun).collect(), - output: fun(schema.output), + outputs: schema.outputs.into_iter().map(&mut *fun).collect(), }, merge, hidden, @@ -1877,6 +2028,7 @@ where presort_and_args, uf, proof_func, + proof_ctors, unionable, } => GenericCommand::Sort { span, @@ -1884,6 +2036,7 @@ where presort_and_args, uf, proof_func, + proof_ctors, unionable, }, GenericCommand::Datatype { diff --git a/egglog/src/ast/parse.rs b/egglog/src/ast/parse.rs index 4d0aebb..6613289 100644 --- a/egglog/src/ast/parse.rs +++ b/egglog/src/ast/parse.rs @@ -38,6 +38,60 @@ macro_rules! error { }; } +/// Names that may not be used as user identifiers (function/sort/constructor/relation/variant +/// names or variables) because they are built-in keywords of the surface syntax: most command, +/// action, and schedule heads, plus `values` (the tuple constructor for tuple-output functions). +/// Keeping these out of the identifier namespace avoids confusing programs where, say, a function +/// named `set` reads like the `set` action. A few common-word commands (`input`, `output`) are +/// only partially reserved — see [`COMMAND_ONLY_KEYWORDS`]. Names starting with `:` are reserved +/// separately (the `:` prefix marks option keywords), see [`Parser::ensure_symbol_not_reserved`]. +const RESERVED_KEYWORDS: &[&str] = &[ + // commands + "sort", + "datatype", + "datatype*", + "function", + "constructor", + "relation", + "ruleset", + "unstable-combined-ruleset", + "rule", + "rewrite", + "birewrite", + "run", + "run-schedule", + "check", + "extract", + "push", + "pop", + "print-function", + "print-size", + "print-stats", + "include", + "fail", + "prove", + "prove-exists", + // actions + "let", + "set", + "union", + "delete", + "subsume", + "panic", + // schedules + "repeat", + "saturate", + "seq", + // tuple-output destructuring/construction + "values", +]; + +/// Commands whose names are common enough to want as ordinary identifiers, so they are only +/// *partially* reserved: usable as variables, but not as the head of a call s-expr (so `(input +/// ...)` always means the command, never a table lookup) nor as a definition (table/sort/etc.) +/// name. +const COMMAND_ONLY_KEYWORDS: &[&str] = &["input", "output"]; + pub enum Sexp { // Will never contain `Literal::Unit`, as this // will be parsed as an empty `Sexp::List`. @@ -168,7 +222,26 @@ impl Default for Parser { impl Parser { fn ensure_symbol_not_reserved(&self, symbol: &str, span: &Span) -> Result<(), ParseError> { - if self.symbol_gen.is_reserved(symbol) && self.ensure_no_reserved_symbols { + // Disabled when re-parsing egglog's own generated programs (e.g. proof/term encoding), + // which may legitimately use internal-prefixed names. + if !self.ensure_no_reserved_symbols { + return Ok(()); + } + if RESERVED_KEYWORDS.contains(&symbol) { + return error!( + span.clone(), + "`{symbol}` is a reserved keyword and cannot be used as a name" + ); + } + // The leading `:` marks option keywords (`:merge`, `:cost`, `:ruleset`, ...); the parser + // uses it to delimit options, so it cannot be the start of a user identifier. + if symbol.starts_with(':') { + return error!( + span.clone(), + "`{symbol}` cannot be used as a name: the `:` prefix is reserved for option keywords" + ); + } + if self.symbol_gen.is_reserved(symbol) { return error!( span.clone(), "symbols starting with '{}' are reserved for egglog internals", @@ -178,6 +251,28 @@ impl Parser { Ok(()) } + /// Check that a name introducing a definition (function/sort/constructor/relation/datatype/ + /// variant) is allowed: neither a reserved keyword nor a command-only keyword like + /// `input`/`output` (which would otherwise create an uncallable table). + fn ensure_definition_name(&self, name: &str, span: &Span) -> Result<(), ParseError> { + self.ensure_symbol_not_reserved(name, span)?; + if self.ensure_no_reserved_symbols && COMMAND_ONLY_KEYWORDS.contains(&name) { + return error!( + span.clone(), + "`{name}` is a command and cannot be used as a definition name" + ); + } + Ok(()) + } + + /// Parse an atom that introduces a new name (e.g. a function, sort, constructor, or relation), + /// rejecting reserved keywords such as `values` and command names like `input`/`output`. + fn parse_name(&self, sexp: &Sexp, what: &'static str) -> Result { + let name = sexp.expect_atom(what)?; + self.ensure_definition_name(&name, &sexp.span())?; + Ok(name) + } + pub fn get_program_from_string( &mut self, filename: Option, @@ -255,38 +350,42 @@ impl Parser { Ok(match head.as_str() { "sort" => { - // Parse sort - :internal-uf/:internal-proof-func and container sorts are mutually exclusive + // Parse sort - the :internal-* annotations and container sorts are mutually exclusive // (sort ) // (sort :internal-uf ) // (sort :internal-proof-func ) + // (sort :internal-proof-names ) // (sort ( *)) match tail { [name] => vec![Command::Sort { span, - name: name.expect_atom("sort name")?, + name: self.parse_name(name, "sort name")?, presort_and_args: None, uf: None, proof_func: None, + proof_ctors: None, unionable: true, }], [name, call @ Sexp::List(..)] => { let (func, args, _) = call.expect_call("container sort declaration")?; vec![Command::Sort { span, - name: name.expect_atom("sort name")?, + name: self.parse_name(name, "sort name")?, presort_and_args: Some(( func, map_fallible(args, self, Self::parse_expr)?, )), uf: None, proof_func: None, + proof_ctors: None, unionable: true, }] } [name, rest @ ..] => { - // Parse :internal-uf and :internal-proof-func annotations + // Parse :internal-uf, :internal-proof-func, and :internal-proof-names annotations let mut uf = None; let mut proof_func = None; + let mut proof_ctors = None; for (key, val) in self.parse_options(rest)? { match (key, val) { (":internal-uf", [uf_func]) => { @@ -296,20 +395,28 @@ impl Parser { proof_func = Some(pf.expect_atom("internal-proof-func function name")?); } + (":internal-proof-names", [congr, trans, sym]) => { + proof_ctors = Some(( + congr.expect_atom("congruence constructor name")?, + trans.expect_atom("transitivity constructor name")?, + sym.expect_atom("symmetry constructor name")?, + )); + } _ => { return error!( span, - "usages:\n(sort )\n(sort :internal-uf )\n(sort :internal-proof-func )\n(sort ( *))" + "usages:\n(sort )\n(sort :internal-uf )\n(sort :internal-proof-func )\n(sort :internal-proof-names )\n(sort ( *))" ); } } } vec![Command::Sort { span, - name: name.expect_atom("sort name")?, + name: self.parse_name(name, "sort name")?, presort_and_args: None, uf, proof_func, + proof_ctors, unionable: true, }] } @@ -324,7 +431,7 @@ impl Parser { "datatype" => match tail { [name, variants @ ..] => vec![Command::Datatype { span, - name: name.expect_atom("sort name")?, + name: self.parse_name(name, "sort name")?, variants: map_fallible(variants, self, Self::variant)?, }], _ => return error!(span, "usage: (datatype *)"), @@ -351,14 +458,44 @@ impl Parser { } merge = Some(None); } - (":merge", [e]) => { + (":merge", args) => { if merge.is_some() { return error!( span, "conflicting merge options: :merge and :no-merge cannot both be specified" ); } - merge = Some(Some(self.parse_expr(e)?)); + // `:merge (* )`: a value-producing action + // block. When the block's first element is itself a list it is an + // action block — the last element is the merged value, the earlier + // ones are actions run first (with old/new bound). Otherwise the + // block is an ordinary result expression: the common back-compatible + // `:merge (max old new)` / `:merge new` (an expression always has an + // atom head, so this is unambiguous). + let [arg] = args else { + return error!( + span, + ":merge takes a single result expression or action block: `:merge (* )`" + ); + }; + let m = match arg { + Sexp::List(items, _) + if matches!(items.first(), Some(Sexp::List(..))) => + { + let (result_sexp, action_sexps) = + items.split_last().unwrap(); + let mut actions = Vec::new(); + for a in action_sexps { + actions.extend(self.parse_action(a)?); + } + GenericMerge { + actions: GenericActions(actions), + result: self.parse_expr(result_sexp)?, + } + } + _ => GenericMerge::result_only(self.parse_expr(arg)?), + }; + merge = Some(Some(m)); } (":internal-hidden", []) => hidden = true, (":internal-let", []) => let_binding = true, @@ -379,7 +516,7 @@ impl Parser { } }; vec![Command::Function { - name: name.expect_atom("function name")?, + name: self.parse_name(name, "function name")?, schema: self.parse_schema(inputs, output)?, merge, hidden, @@ -390,7 +527,8 @@ impl Parser { }] } _ => { - let a = "(function (*) :merge )"; + let a = + "(function (*) :merge * )"; let b = "(function (*) :no-merge)"; return error!(span, "usages:\n{a}\n{b}"); } @@ -419,7 +557,7 @@ impl Parser { vec![Command::Constructor { span, - name: name.expect_atom("constructor name")?, + name: self.parse_name(name, "constructor name")?, schema: self.parse_schema(inputs, output)?, cost, unextractable, @@ -439,7 +577,7 @@ impl Parser { "relation" => match tail { [name, inputs] => vec![Command::Relation { span, - name: name.expect_atom("relation name")?, + name: self.parse_name(name, "relation name")?, inputs: map_fallible(inputs.expect_list("input sorts")?, self, |_, sexp| { sexp.expect_atom("input sort") })?, @@ -918,7 +1056,13 @@ impl Parser { if *symbol == "_" { self.symbol_gen.fresh(symbol) } else { - self.ensure_symbol_not_reserved(symbol, span)?; + // `:`-prefixed atoms are option-keyword markers (e.g. `:until` in a custom + // run-schedule) that command macros consume as `Expr::Var`s; allow them here. + // User variables never start with `:`, so this doesn't weaken the reserved-name + // guarantee for identifiers. + if !symbol.starts_with(':') { + self.ensure_symbol_not_reserved(symbol, span)?; + } symbol.clone() }, ), @@ -931,6 +1075,17 @@ impl Parser { return func.parse(tail, span, self); } + // `input`/`output` are commands, not callable tables, so they may not head an + // expression even though they are allowed as ordinary variable names. + if self.ensure_no_reserved_symbols + && COMMAND_ONLY_KEYWORDS.contains(&head.as_str()) + { + return error!( + span, + "`{head}` is a command and cannot be used as the head of an expression" + ); + } + Expr::Call( span.clone(), head, @@ -950,7 +1105,7 @@ impl Parser { Ok(match head.as_str() { "sort" => match tail { [name, call] => { - let name = name.expect_atom("sort name")?; + let name = self.parse_name(name, "sort name")?; let (func, args, _) = call.expect_call("container sort declaration")?; let args = map_fallible(args, self, Self::parse_expr)?; (span, name, Subdatatypes::NewSort(func, args)) @@ -963,6 +1118,7 @@ impl Parser { } }, _ => { + self.ensure_definition_name(&head, &span)?; let variants = map_fallible(tail, self, Self::variant)?; (span, head, Subdatatypes::Variants(variants)) } @@ -971,6 +1127,7 @@ impl Parser { pub fn variant(&mut self, sexp: &Sexp) -> Result { let (name, tail, span) = sexp.expect_call("datatype variant")?; + self.ensure_definition_name(&name, &span)?; let (types, cost, unextractable) = match tail { [types @ .., Sexp::Atom(o, _)] if *o == ":unextractable" => (types, None, true), @@ -1023,14 +1180,27 @@ impl Parser { } pub fn parse_schema(&self, input: &Sexp, output: &Sexp) -> Result { - Ok(Schema { - input: input - .expect_list("input sorts")? + let input = input + .expect_list("input sorts")? + .iter() + .map(|sexp| sexp.expect_atom("input sort")) + .collect::>()?; + // The output is either a single sort, or a parenthesized list of sorts for a + // tuple-output function (e.g. `(function f (Math) (i64 i64) ...)`). + let outputs: Vec = match output { + Sexp::List(list, _) => list .iter() - .map(|sexp| sexp.expect_atom("input sort")) + .map(|sexp| sexp.expect_atom("output sort")) .collect::>()?, - output: output.expect_atom("output sort")?, - }) + _ => vec![output.expect_atom("output sort")?], + }; + if outputs.is_empty() { + return error!( + output.span(), + "a function must have at least one output sort" + ); + } + Ok(Schema::new_tuple(input, outputs)) } } diff --git a/egglog/src/ast/proof_global_remover.rs b/egglog/src/ast/proof_global_remover.rs index 2cbbb4b..e58e9f0 100644 --- a/egglog/src/ast/proof_global_remover.rs +++ b/egglog/src/ast/proof_global_remover.rs @@ -41,7 +41,7 @@ fn resolved_var_to_call(var: &ResolvedVar) -> ResolvedCall { name: var.name.clone(), subtype: FunctionSubtype::Constructor, input: vec![], - output: var.sort.clone(), + outputs: vec![var.sort.clone()], }) } @@ -79,14 +79,14 @@ fn remove_globals_cmd(cmd: ResolvedNCommand) -> Vec { name: name.name.clone(), subtype: FunctionSubtype::Constructor, input: vec![], - output: ty.clone(), + outputs: vec![ty.clone()], }); let func_decl = ResolvedFunctionDecl { name: name.name, subtype: FunctionSubtype::Constructor, schema: Schema { input: vec![], - output: ty.name().to_owned(), + outputs: vec![ty.name().to_owned()], }, resolved_schema: resolved_call.clone(), merge: None, diff --git a/egglog/src/ast/remove_globals.rs b/egglog/src/ast/remove_globals.rs index c37f908..0be442c 100644 --- a/egglog/src/ast/remove_globals.rs +++ b/egglog/src/ast/remove_globals.rs @@ -61,7 +61,7 @@ fn resolved_var_to_call(var: &ResolvedVar) -> ResolvedCall { name: var.name.clone(), subtype: FunctionSubtype::Custom, input: vec![], - output: var.sort.clone(), + outputs: vec![var.sort.clone()], }) } @@ -97,14 +97,14 @@ impl GlobalRemover<'_> { name: name.name.clone(), subtype: FunctionSubtype::Custom, input: vec![], - output: ty.clone(), + outputs: vec![ty.clone()], }); let func_decl = ResolvedFunctionDecl { name: name.name, subtype: FunctionSubtype::Custom, schema: Schema { input: vec![], - output: ty.name().to_owned(), + outputs: vec![ty.name().to_owned()], }, resolved_schema: resolved_call.clone(), merge: None, diff --git a/egglog/src/constraint.rs b/egglog/src/constraint.rs index 3267a58..319d662 100644 --- a/egglog/src/constraint.rs +++ b/egglog/src/constraint.rs @@ -542,6 +542,21 @@ impl Assignment { .iter() .map(|arg| self.annotate_expr(arg, typeinfo, ctx)) .collect(); + // The `values` tuple constructor resolves to `ResolvedCall::Values` carrying its + // element sorts. A tuple-output function call carries only its input columns here + // (its outputs live on the `values` side of the destructure/set), so it resolves + // from input types alone. + if head.as_str() == "values" { + let sorts = args.iter().map(|arg| arg.output_type()).collect(); + return GenericExpr::Call(span.clone(), ResolvedCall::Values(sorts), args); + } + if let Some(ty) = typeinfo.get_func_type(head).filter(|t| t.is_tuple_output()) { + let input_types: Vec<_> = args.iter().map(|arg| arg.output_type()).collect(); + let resolved_call = + ResolvedCall::from_resolution_func_types(head, &input_types, typeinfo) + .unwrap_or_else(|| ResolvedCall::Func(ty.clone())); + return GenericExpr::Call(span.clone(), resolved_call, args); + } let types: Vec<_> = args .iter() .map(|arg| arg.output_type()) @@ -621,12 +636,23 @@ impl Assignment { .map(|child| self.annotate_expr(child, typeinfo, ctx)) .collect(); let rhs = self.annotate_expr(rhs, typeinfo, ctx); - let types: Vec<_> = children - .iter() - .map(|child| child.output_type()) - .chain(once(rhs.output_type())) - .collect(); - let resolved_call = ResolvedCall::from_resolution(head, &types, typeinfo, ctx); + // For a tuple-output function the `rhs` is a `(values ...)` form, so the function + // is resolved from its input columns alone. + let resolved_call = if let Some(ty) = + typeinfo.get_func_type(head).filter(|t| t.is_tuple_output()) + { + let input_types: Vec<_> = + children.iter().map(|child| child.output_type()).collect(); + ResolvedCall::from_resolution_func_types(head, &input_types, typeinfo) + .unwrap_or_else(|| ResolvedCall::Func(ty.clone())) + } else { + let types: Vec<_> = children + .iter() + .map(|child| child.output_type()) + .chain(once(rhs.output_type())) + .collect(); + ResolvedCall::from_resolution(head, &types, typeinfo, ctx) + }; if !matches!(resolved_call, ResolvedCall::Func(_)) { return Err(TypeError::UnboundFunction(head.clone(), span.clone())); } @@ -835,7 +861,7 @@ impl CoreAction { } let mut args = args.clone(); - args.push(rhs.clone()); + args.extend(rhs.iter().cloned()); Ok(get_literal_and_global_constraints(&args, typeinfo) .chain(get_atom_application_constraints( @@ -845,9 +871,16 @@ impl CoreAction { } CoreAction::Change(span, _change, head, args) => { let mut args = args.clone(); - // Add a dummy last output argument - let var = symbol_gen.fresh(head); - args.push(AtomTerm::Var(span.clone(), var)); + // Add a dummy output argument per output column (tuple-output views have more than + // one), so the atom matches the function's full arity for constraint solving. + let num_outputs = typeinfo + .get_func_type(head) + .map(|t| t.num_outputs()) + .unwrap_or(1); + for _ in 0..num_outputs { + let var = symbol_gen.fresh(head); + args.push(AtomTerm::Var(span.clone(), var)); + } Ok(get_literal_and_global_constraints(&args, typeinfo) .chain(get_atom_application_constraints( @@ -917,11 +950,13 @@ fn get_atom_application_constraints( // `constraint::xor` means one and only one of the instantiation can hold. let mut xor_constraints: Vec>>> = vec![]; - // function atom constraints + // function atom constraints. A function atom carries all input columns followed by all output + // columns (more than one for a tuple-output function). if let Some(typ) = type_info.get_func_type(head) { let mut constraints = vec![]; + let expected = typ.input.len() + typ.num_outputs(); // arity mismatch - if typ.input.len() + 1 != args.len() { + if expected != args.len() { constraints.push(constraint::impossible( ImpossibleConstraint::ArityMismatch { atom: Atom { @@ -929,15 +964,15 @@ fn get_atom_application_constraints( head: head.to_owned(), args: args.to_vec(), }, - expected: typ.input.len() + 1, + expected, }, )); } else { for (arg_typ, arg) in typ .input .iter() + .chain(typ.outputs.iter()) .cloned() - .chain(once(typ.output.clone())) .zip(args.iter().cloned()) { constraints.push(constraint::assign(arg, arg_typ)); @@ -1181,7 +1216,9 @@ pub(crate) fn grounded_check( for atom in body.atoms.iter() { let mut add_global_and_literal = false; match &atom.head { - HeadOrEq::Head(ResolvedCall::Func(_)) => { + // `Values` never appears as a query atom head (tuple destructures are lowered to the + // underlying function atom), but it is grounded like a function if it ever did. + HeadOrEq::Head(ResolvedCall::Func(_) | ResolvedCall::Values(_)) => { for arg in atom.args.iter() { problem.constraints.push(assign(arg.clone(), ())); } diff --git a/egglog/src/core.rs b/egglog/src/core.rs index d0dd46d..34cbecf 100644 --- a/egglog/src/core.rs +++ b/egglog/src/core.rs @@ -111,10 +111,41 @@ impl Hash for SpecializedPrimitive { } } -#[derive(Debug, Clone, PartialEq, Eq, Hash)] +#[derive(Debug, Clone)] pub enum ResolvedCall { Func(FuncType), Primitive(SpecializedPrimitive), + /// The `values` tuple constructor, used to destructure a tuple-output function's outputs in a + /// query (`(= (values a b) (f x))`) or to construct them in a `set` action + /// (`(set (f x) (values a b))`). Carries the output sorts. It never reaches the backend on its + /// own: it is always paired with a tuple-output function call when lowering to core. + Values(Vec), +} + +impl PartialEq for ResolvedCall { + fn eq(&self, other: &Self) -> bool { + match (self, other) { + (ResolvedCall::Func(a), ResolvedCall::Func(b)) => a == b, + (ResolvedCall::Primitive(a), ResolvedCall::Primitive(b)) => a == b, + (ResolvedCall::Values(a), ResolvedCall::Values(b)) => { + a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.name() == y.name()) + } + _ => false, + } + } +} + +impl Eq for ResolvedCall {} + +impl Hash for ResolvedCall { + fn hash(&self, state: &mut H) { + std::mem::discriminant(self).hash(state); + match self { + ResolvedCall::Func(f) => f.hash(state), + ResolvedCall::Primitive(p) => p.hash(state), + ResolvedCall::Values(sorts) => sorts.iter().for_each(|s| s.name().hash(state)), + } + } } impl ResolvedCall { @@ -122,13 +153,18 @@ impl ResolvedCall { match self { ResolvedCall::Func(func) => &func.name, ResolvedCall::Primitive(prim) => prim.name(), + ResolvedCall::Values(_) => "values", } } pub fn output(&self) -> &ArcSort { match self { - ResolvedCall::Func(func) => &func.output, + ResolvedCall::Func(func) => func.output(), ResolvedCall::Primitive(prim) => prim.output(), + // `values` has no single output; its first column is returned only so that callers + // that incidentally ask for "a" sort do not panic. Tuple-output uses are routed + // specially before this is consulted. + ResolvedCall::Values(sorts) => &sorts[0], } } @@ -138,10 +174,11 @@ impl ResolvedCall { match self { ResolvedCall::Func(func) => { let mut types = func.input.clone(); - types.push(func.output.clone()); + types.extend(func.outputs.iter().cloned()); types } ResolvedCall::Primitive(prim) => prim.input().to_vec(), + ResolvedCall::Values(sorts) => sorts.clone(), } } @@ -171,7 +208,7 @@ impl ResolvedCall { ctx: crate::Context, ) -> ResolvedCall { if let Some(ty) = typeinfo.get_func_type(head) { - let expected = ty.input.iter().chain(once(&ty.output)).map(|s| s.name()); + let expected = ty.input.iter().chain(ty.outputs.iter()).map(|s| s.name()); let actual = types.iter().map(|s| s.name()); if expected.eq(actual) { return ResolvedCall::Func(ty.clone()); @@ -206,6 +243,7 @@ impl Display for ResolvedCall { match self { ResolvedCall::Func(func) => write!(f, "{}", func.name), ResolvedCall::Primitive(prim) => write!(f, "{}", prim.name()), + ResolvedCall::Values(_) => write!(f, "values"), } } } @@ -224,6 +262,7 @@ impl IsFunc for ResolvedCall { match self { ResolvedCall::Func(func) => type_info.is_constructor(&func.name), ResolvedCall::Primitive(_) => false, + ResolvedCall::Values(_) => false, } } } @@ -234,6 +273,36 @@ impl IsFunc for String { } } +/// Operations on a call head needed to lower the `values` tuple sugar. Implemented for both the +/// unresolved (`String`) and resolved ([`ResolvedCall`]) head representations so that the same +/// lowering code runs in both the type-checking and canonicalization passes. +pub trait HeadOps { + /// Whether this head is the `values` tuple constructor. + fn is_values(&self) -> bool; + /// Whether this head is a tuple-output function (more than one output column). + fn is_tuple_output(&self, type_info: &TypeInfo) -> bool; +} + +impl HeadOps for String { + fn is_values(&self) -> bool { + self == "values" + } + fn is_tuple_output(&self, type_info: &TypeInfo) -> bool { + type_info + .get_func_type(self) + .is_some_and(|t| t.is_tuple_output()) + } +} + +impl HeadOps for ResolvedCall { + fn is_values(&self) -> bool { + matches!(self, ResolvedCall::Values(_)) + } + fn is_tuple_output(&self, _type_info: &TypeInfo) -> bool { + matches!(self, ResolvedCall::Func(f) if f.is_tuple_output()) + } +} + #[derive(Debug, Clone)] pub enum GenericAtomTerm { Var(Span, Leaf), @@ -462,6 +531,7 @@ impl Query { head: head.clone(), args: atom.args.clone(), }), + ResolvedCall::Values(_) => None, }) } @@ -473,6 +543,7 @@ impl Query { args: atom.args.clone(), }), ResolvedCall::Primitive(_) => None, + ResolvedCall::Values(_) => None, }) } } @@ -481,11 +552,13 @@ impl Query { pub enum GenericCoreAction { Let(Span, Leaf, Head, Vec>), LetAtomTerm(Span, Leaf, GenericAtomTerm), + /// `(set (f keys...) value...)`. The last field holds the function's output (value) columns: + /// one for an ordinary function, more for a tuple-output function. Set( Span, Head, Vec>, - GenericAtomTerm, + Vec>, ), Change(Span, Change, Head, Vec>), Union(Span, GenericAtomTerm, GenericAtomTerm), @@ -554,9 +627,9 @@ where add_from_atom(&mut free_vars, at); free_vars.remove(v); } - GenericCoreAction::Set(_span, _, ats, at) => { + GenericCoreAction::Set(_span, _, ats, vals) => { add_from_atoms(&mut free_vars, ats); - add_from_atom(&mut free_vars, at); + add_from_atoms(&mut free_vars, vals); } GenericCoreAction::Change(_span, _change, _, ats) => { add_from_atoms(&mut free_vars, ats); @@ -610,14 +683,14 @@ pub(crate) trait GenericActionsExt { ctx: &mut CoreActionContext<'_, Head, Leaf, FG>, ) -> Result<(GenericCoreActions, MappedActions), TypeError> where - Head: Clone + Display + IsFunc, + Head: Clone + Display + IsFunc + HeadOps, Leaf: Clone + PartialEq + Eq + Display + Hash, FG: FreshGen; } impl GenericActionsExt for GenericActions where - Head: Clone + Display + IsFunc, + Head: Clone + Display + IsFunc + HeadOps, Leaf: Clone + PartialEq + Eq + Display + Hash, { #[allow(clippy::type_complexity)] @@ -626,7 +699,7 @@ where ctx: &mut CoreActionContext<'_, Head, Leaf, FG>, ) -> Result<(GenericCoreActions, MappedActions), TypeError> where - Head: Clone + Display + IsFunc, + Head: Clone + Display + IsFunc + HeadOps, Leaf: Clone + PartialEq + Eq + Display + Hash, FG: FreshGen, { @@ -663,7 +736,31 @@ where let mapped_arg = arg.to_core_actions(ctx, &mut norm_actions)?; mapped_args.push(mapped_arg); } - let mapped_expr = expr.to_core_actions(ctx, &mut norm_actions)?; + // The value may be a `(values v...)` tuple (for a tuple-output function), which + // contributes one output column per element; otherwise it is a single value. + let (mapped_value, value_terms) = match expr { + GenericExpr::Call(vspan, vhead, vargs) if vhead.is_values() => { + let mut mapped = vec![]; + let mut terms = vec![]; + for v in vargs { + let m = v.to_core_actions(ctx, &mut norm_actions)?; + terms.push(m.get_corresponding_var_or_lit(typeinfo)); + mapped.push(m); + } + let dummy = ctx.fresh_gen.fresh(vhead); + let mapped_call = GenericExpr::Call( + vspan.clone(), + CorrespondingVar::new(vhead.clone(), dummy), + mapped, + ); + (mapped_call, terms) + } + _ => { + let m = expr.to_core_actions(ctx, &mut norm_actions)?; + let term = m.get_corresponding_var_or_lit(typeinfo); + (m, vec![term]) + } + }; norm_actions.push(GenericCoreAction::Set( span.clone(), head.clone(), @@ -671,14 +768,14 @@ where .iter() .map(|e| e.get_corresponding_var_or_lit(typeinfo)) .collect(), - mapped_expr.get_corresponding_var_or_lit(typeinfo), + value_terms, )); let v = ctx.fresh_gen.fresh(head); mapped_actions.0.push(GenericAction::Set( span.clone(), CorrespondingVar::new(head.clone(), v), mapped_args, - mapped_expr, + mapped_value, )); } GenericAction::Change(span, change, head, args) => { @@ -730,7 +827,8 @@ where .iter() .map(|e| e.get_corresponding_var_or_lit(typeinfo)) .collect(), - mapped_expr.get_corresponding_var_or_lit(typeinfo), + // Constructors are single-output, so a single value column. + vec![mapped_expr.get_corresponding_var_or_lit(typeinfo)], )); let v = ctx.fresh_gen.fresh(head); mapped_actions.0.push(GenericAction::Set( @@ -1089,13 +1187,13 @@ pub(crate) trait GenericRuleExt { union_to_set_optimization: bool, ) -> Result, Head, Leaf>, TypeError> where - Head: Clone + Display + IsFunc, + Head: Clone + Display + IsFunc + HeadOps, Leaf: Clone + PartialEq + Eq + Display + Hash + Debug; } impl GenericRuleExt for GenericRule where - Head: Clone + Display + IsFunc, + Head: Clone + Display + IsFunc + HeadOps, Leaf: Clone + PartialEq + Eq + Display + Hash + Debug, { fn to_core_rule( @@ -1105,7 +1203,7 @@ where union_to_set_optimization: bool, ) -> Result, Head, Leaf>, TypeError> where - Head: Clone + Display + IsFunc, + Head: Clone + Display + IsFunc + HeadOps, Leaf: Clone + PartialEq + Eq + Display + Hash + Debug, { let (body, _correspondence) = Facts(self.body.clone()).to_query(typeinfo, fresh_gen); diff --git a/egglog/src/exec_state.rs b/egglog/src/exec_state.rs index ebac11e..10ea82a 100644 --- a/egglog/src/exec_state.rs +++ b/egglog/src/exec_state.rs @@ -290,6 +290,9 @@ pub trait Core<'a, 'db: 'a>: Internal<'a, 'db> { match resolved_call { ResolvedCall::Primitive(primitive) => self.apply_primitive(primitive, &values), ResolvedCall::Func(func) => self.apply_resolved_function(func, &values), + ResolvedCall::Values(_) => { + panic!("`values` cannot be evaluated as a single-valued expression") + } } } } diff --git a/egglog/src/extract.rs b/egglog/src/extract.rs index 3c7aa40..0b77b86 100644 --- a/egglog/src/extract.rs +++ b/egglog/src/extract.rs @@ -562,7 +562,6 @@ impl Extractor { /// Find the canonical representative of a value using the union-find table. /// If no UF is registered for this sort, returns the original value. - /// The UF table stores (value, canonical) pairs - one hop lookup. fn find_canonical(&self, egraph: &EGraph, value: Value, sort: &ArcSort) -> Value { // Check if there's a UF registered for this sort let Some(uf_name) = egraph.proof_state.uf_parent.get(sort.name()) else { @@ -574,12 +573,28 @@ impl Extractor { return value; }; - // Single lookup in UF table - it's guaranteed to be one hop to canonical + // A single-key union-find is a self-referential function keyed by the element (the + // encoding's `UF_`: `(S) -> S` in term mode, `(S) -> (S, Proof)` in proof mode), so + // `lookup_id` returns the parent. Chase to a fixpoint; a miss means the value is its own + // leader. A two-key union-find is a `(child, parent)` relation (e.g. a user-provided + // `:internal-uf`), resolved with a one-hop scan. Dispatching on the key arity keeps both + // correct. + if uf_func.schema.input.len() == 1 { + let mut canonical = value; + loop { + match egraph.backend.lookup_id(uf_func.backend_id, &[canonical]) { + Some(next) if next != canonical => canonical = next, + _ => break, + } + } + return canonical; + } + + // Two-key `(child, parent)` relation: one-hop lookup. let mut canonical = value; egraph .backend .for_each(uf_func.backend_id, |row: egglog_bridge::ScanEntry| { - // UF table has (child, parent) as inputs if row.vals[0] == value { canonical = row.vals[1]; } @@ -715,21 +730,30 @@ impl Function { } } - /// For view tables (with term_constructor), the effective output sort is the last input column. - /// For regular tables, it's the output sort. + /// Whether this is the proof-mode functional-dependency view `(children) -> (eclass, proof)`, + /// where the e-class is the first output column rather than the last input column. + fn is_fd_view(&self) -> bool { + self.decl.term_constructor.is_some() && self.schema.outputs.len() > 1 + } + + /// For view tables (with term_constructor), the effective output sort is the last input column + /// (old form) or the first output column (FD tuple view). For regular tables, it's the output. /// This is used by extraction to determine which sort a table produces values for. pub(crate) fn extraction_output_sort(&self) -> &ArcSort { - if self.decl.term_constructor.is_some() { + if self.is_fd_view() { + self.schema.output() + } else if self.decl.term_constructor.is_some() { self.schema.input.last().unwrap() } else { - &self.schema.output + self.schema.output() } } /// Returns the number of children for extraction purposes. - /// For view tables, this excludes the last column (the e-class). + /// For old-form view tables, this excludes the last input column (the e-class); FD tuple views + /// key on children only, so all inputs are children. pub(crate) fn extraction_num_children(&self) -> usize { - if self.decl.term_constructor.is_some() { + if self.decl.term_constructor.is_some() && !self.is_fd_view() { self.schema.input.len() - 1 } else { self.schema.input.len() @@ -749,13 +773,12 @@ impl Function { /// For view tables, the e-class is the last input column (second-to-last in the row). /// For regular tables, it's the last column (the actual output). pub(crate) fn extraction_output_index(&self) -> usize { - if self.decl.term_constructor.is_some() { - // For view tables: input is [children..., eclass], output is view_sort - // Row is [children..., eclass, view_sort] - // We want eclass which is at index input.len() - 1 + if self.decl.term_constructor.is_some() && !self.is_fd_view() { + // Old-form view: row is [children..., eclass, view_sort]; eclass at input.len() - 1. self.schema.input.len() - 1 } else { - // For regular tables: row is [inputs..., output] + // Regular table: [inputs..., output]. FD view: [children..., eclass, proof]; the eclass + // is the first output column, at index input.len(). self.schema.input.len() } } @@ -812,7 +835,7 @@ impl EGraph { .ok_or(TypeError::UnboundFunction(sym.to_owned(), span!()))?; let mut rootsorts = func.schema.input.clone(); if include_output { - rootsorts.push(func.schema.output.clone()); + rootsorts.extend(func.schema.outputs.iter().cloned()); } let extractor = Extractor::compute_costs_from_rootsorts( Some(rootsorts), @@ -840,11 +863,22 @@ impl EGraph { } inputs.push(termdag.app(sym.to_owned(), children)); if include_output { - let value = row.vals[func.schema.input.len()]; - let sort = &func.schema.output; - let (_, term) = extractor - .extract_best_with_sort(self, &mut termdag, value, sort.clone()) - .unwrap_or_else(|| (0, termdag.var("Unextractable".into()))); + // Extract every output (value) column. A tuple-output function has more than + // one; we display them wrapped in a `(values ...)` term to mirror the surface + // syntax. + let mut out_terms = Vec::new(); + for (i, sort) in func.schema.outputs.iter().enumerate() { + let value = row.vals[func.schema.input.len() + i]; + let (_, term) = extractor + .extract_best_with_sort(self, &mut termdag, value, sort.clone()) + .unwrap_or_else(|| (0, termdag.var("Unextractable".into()))); + out_terms.push(term); + } + let term = if out_terms.len() == 1 { + out_terms.pop().unwrap() + } else { + termdag.app("values".to_owned(), out_terms) + }; output.as_mut().unwrap().push(term); } true diff --git a/egglog/src/lib.rs b/egglog/src/lib.rs index b9d7e20..d56ce42 100644 --- a/egglog/src/lib.rs +++ b/egglog/src/lib.rs @@ -221,7 +221,7 @@ impl std::fmt::Display for CommandOutput { write!(f, "Overall statistics:\n{run_report}") } CommandOutput::PrintFunction(function, termdag, terms_and_outputs, mode) => { - let out_is_unit = function.schema.output.name() == UnitSort.name(); + let out_is_unit = function.schema.output().name() == UnitSort.name(); if *mode == PrintFunctionMode::CSV { let mut wtr = Writer::from_writer(vec![]); for (term_id, output) in terms_and_outputs { @@ -375,17 +375,15 @@ impl Function { #[derive(Clone, Debug)] pub struct ResolvedSchema { pub input: Vec, - pub output: ArcSort, + /// The output (value-column) sorts, primary first. A tuple-output function has more than one; + /// ordinary functions have exactly one. Always non-empty. + pub outputs: Vec, } impl ResolvedSchema { - /// Get the type at position `index`, counting the `output` sort as at position `input.len()`. - pub fn get_by_pos(&self, index: usize) -> Option<&ArcSort> { - if self.input.len() == index { - Some(&self.output) - } else { - self.input.get(index) - } + /// The primary (first) output sort. + pub fn output(&self) -> &ArcSort { + &self.outputs[0] } } @@ -478,6 +476,18 @@ impl Default for EGraph { if a > b { a } else { b } }); + // Orientation helpers for the proof-encoding UF/view merges: given two `(value, proof)` + // pairs `(a, ap)` and `(b, bp)`, return the proof paired with the smaller / larger value + // (same value ordering as `ordering-min`/`ordering-max`). Polymorphic like the ordering + // prims, so each resolves to a single value-level external function and needs no `Proof` + // sort to exist — which keeps the encoding self-contained on a non-proof re-parse. + add_primitive!(&mut eg, "proof-of-min" = |a: #, ap: #, b: #, bp: #| -> # { + if a < b { ap } else { bp } + }); + add_primitive!(&mut eg, "proof-of-max" = |a: #, ap: #, b: #, bp: #| -> # { + if a > b { ap } else { bp } + }); + eg.rulesets .insert("".into(), Ruleset::Rules(Default::default())); @@ -748,22 +758,37 @@ impl EGraph { fn translate_expr_to_mergefn( &self, expr: &ResolvedExpr, + lets: &HashMap, ) -> Result { match expr { GenericExpr::Lit(_, literal) => { let val = literal_to_value(&self.backend, literal); Ok(egglog_bridge::MergeFn::Const(val)) } - GenericExpr::Var(span, resolved_var) => match resolved_var.name.as_str() { - "old" => Ok(egglog_bridge::MergeFn::Old), - "new" => Ok(egglog_bridge::MergeFn::New), - // NB: type-checking should already catch unbound variables here. - _ => Err(TypeError::Unbound(resolved_var.name.clone(), span.clone()).into()), - }, + GenericExpr::Var(span, resolved_var) => { + let name = resolved_var.name.as_str(); + // A `let`-bound variable resolves to its environment slot. Otherwise: single-output + // merges use `old`/`new`; tuple-output merges use `old0`, `new0`, `old1`, ... to + // refer to the old/new value of a specific output column. + if let Some(&slot) = lets.get(name) { + Ok(egglog_bridge::MergeFn::LetVar(slot)) + } else if name == "old" { + Ok(egglog_bridge::MergeFn::Old) + } else if name == "new" { + Ok(egglog_bridge::MergeFn::New) + } else if let Some(i) = name.strip_prefix("old").and_then(|s| s.parse().ok()) { + Ok(egglog_bridge::MergeFn::OldCol(i)) + } else if let Some(i) = name.strip_prefix("new").and_then(|s| s.parse().ok()) { + Ok(egglog_bridge::MergeFn::NewCol(i)) + } else { + // 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) => { let translated_args = args .iter() - .map(|arg| self.translate_expr_to_mergefn(arg)) + .map(|arg| self.translate_expr_to_mergefn(arg, lets)) .collect::, _>>()?; Ok(egglog_bridge::MergeFn::Function( self.functions[&f.name].backend_id, @@ -773,7 +798,7 @@ impl EGraph { GenericExpr::Call(_, ResolvedCall::Primitive(p), args) => { let mut translated_args = args .iter() - .map(|arg| self.translate_expr_to_mergefn(arg)) + .map(|arg| self.translate_expr_to_mergefn(arg, lets)) .collect::, _>>()?; if p.name() == "unstable-fn" { let Some(GenericExpr::Lit(_, Literal::String(name))) = args.first() else { @@ -801,9 +826,289 @@ impl EGraph { translated_args, )) } + // `(values ...)` never legitimately reaches here: a top-level tuple merge is + // destructured per column in `declare_function`, and any other `(values ...)` is + // rejected during type-checking. This arm only keeps the match exhaustive. + GenericExpr::Call(span, ResolvedCall::Values(_), _) => Err(Error::TypeError( + TypeError::TupleMergeNotValues("".to_owned(), span.clone()), + )), } } + /// Lower a resolved `:merge` (a value-producing action block) to a backend [`MergeFn`], keeping + /// the existing merge interpreter. The `result` produces the merged value(s); any `actions` run + /// first as effects. + fn translate_merge_to_mergefn( + &self, + merge: &ResolvedMerge, + ) -> Result { + use egglog_bridge::MergeFn; + // Assign each `let`-bound variable an environment slot, in block order, so `set`/`union` + // args and the result can refer to it via `MergeFn::LetVar`. Built up front because the + // result is lowered before the actions. + let mut lets = HashMap::::default(); + for action in merge.actions.iter() { + if let GenericAction::Let(_, var, _) = action { + let slot = lets.len(); + lets.insert(var.name.as_str().to_owned(), slot); + } + } + // Lower the result value (a `(values ...)` result becomes one column per element). + let result = match &merge.result { + GenericExpr::Call(_, ResolvedCall::Values(_), cols) => MergeFn::Columns( + cols.iter() + .map(|e| self.translate_expr_to_mergefn(e, &lets)) + .collect::, _>>()?, + ), + expr => self.translate_expr_to_mergefn(expr, &lets)?, + }; + if merge.actions.is_empty() { + return Ok(result); + } + // A value-producing action block: run the effects, then evaluate the result value(s). + let actions = merge + .actions + .iter() + .map(|a| self.translate_merge_action(a, &lets)) + .collect::, _>>()?; + Ok(MergeFn::Block { + actions, + result: Box::new(result), + }) + } + + /// Lower a single resolved merge action to a backend [`MergeAction`]. Supports `set`, `let`, and + /// `union`; other actions (`delete`/`panic`/`extract`/...) are not meaningful during a merge. + fn translate_merge_action( + &self, + action: &ResolvedAction, + lets: &HashMap, + ) -> Result { + use egglog_bridge::MergeAction; + match action { + GenericAction::Let(_, var, expr) => Ok(MergeAction::Let { + slot: lets[var.name.as_str()], + value: self.translate_expr_to_mergefn(expr, lets)?, + }), + GenericAction::Union(_, a, b) => Ok(MergeAction::Union( + self.translate_expr_to_mergefn(a, lets)?, + self.translate_expr_to_mergefn(b, lets)?, + )), + GenericAction::Set(_, ResolvedCall::Func(f), keys, val) => { + let backend_id = self + .functions + .get(&f.name) + .ok_or_else(|| { + Error::BackendError(format!( + "merge action sets unknown function `{}`", + f.name + )) + })? + .backend_id; + let mut args = keys + .iter() + .map(|k| self.translate_expr_to_mergefn(k, lets)) + .collect::, _>>()?; + // A tuple-output target is set with `(values ...)`; expand it into value columns. + match val { + GenericExpr::Call(_, ResolvedCall::Values(_), cols) => { + for c in cols { + args.push(self.translate_expr_to_mergefn(c, lets)?); + } + } + _ => args.push(self.translate_expr_to_mergefn(val, lets)?), + } + Ok(MergeAction::Set(backend_id, args)) + } + other => Err(Error::BackendError(format!( + "action `{other}` is not supported inside a :merge block (only `set`, `let`, `union`)" + ))), + } + } + + /// Build the native congruence `:merge` for a term-encoding constructor view + /// `(children) -> (eclass, view_proof)` (proof mode). Returns `Ok(None)` for any function that + /// isn't such a view, leaving the ordinary merge lowering in place. Once the view shape matches, + /// the proof-encoding metadata (`:internal-proof-names`/`:internal-uf`) must resolve; a missing + /// piece is a hard error rather than a silent fallback that would drop congruence. + /// + /// The view proofs are oriented `eclass = f(children)` (eclass on the LEFT). On an FD conflict + /// (two congruent terms share the same canonical children) the merge keeps the SMALLER eclass and + /// stages the oriented union edge `(@UF max_ec) = (values min_ec (Trans max_vp (Sym min_vp)))` + /// into the per-sort UF, so the single-table UF stays acyclic without a `single_parent` rule. + fn native_congruence_merge( + &self, + decl: &ResolvedFunctionDecl, + output: &ArcSort, + num_outputs: usize, + ) -> Result, Error> { + use egglog_bridge::{MergeAction, MergeFn}; + // Detected purely by shape: a term-constructor view with an eq-sort first output and a + // second (proof) output. This shape is only ever emitted by the proof encoder, so we don't + // gate on `proofs_enabled` — that lets a re-parsed desugared program (run in a fresh, non- + // proof egraph) rebuild the same merge, which is what makes the encoding self-contained. + if !(decl.term_constructor.is_some() && num_outputs == 2 && output.is_eq_sort()) { + return Ok(None); + } + // The shape now uniquely identifies a proof-mode congruence view, so every lookup below is + // expected to succeed. If one doesn't, the encoding's `:internal-proof-names`/`:internal-uf` + // metadata is absent or stale; fail loudly instead of silently falling back to a merge that + // never resolves congruence. + let missing = |what: &str| { + Error::BackendError(format!( + "term-encoding congruence view `{}` cannot resolve {what}; the proof encoding's \ + :internal-proof-names/:internal-uf metadata is absent or stale", + decl.name + )) + }; + let resolve = |name: &str| -> Option { + self.type_info + .get_prims(name) + .and_then(|p| p.first()) + .and_then(|p| p.context_ids[crate::Context::Write]) + }; + let min_id = + resolve("ordering-min").ok_or_else(|| missing("the `ordering-min` primitive"))?; + let max_id = + resolve("ordering-max").ok_or_else(|| missing("the `ordering-max` primitive"))?; + let uf_name = self + .proof_state + .uf_parent + .get(decl.schema.output()) + .ok_or_else(|| missing("its per-sort union-find"))?; + let uf_id = self + .functions + .get(uf_name) + .ok_or_else(|| missing("its union-find function"))? + .backend_id; + let trans_id = self + .functions + .get(&self.proof_state.proof_names.eq_trans_constructor) + .ok_or_else(|| missing("the `Trans` proof constructor"))? + .backend_id; + let sym_id = self + .functions + .get(&self.proof_state.proof_names.eq_sym_constructor) + .ok_or_else(|| missing("the `Sym` proof constructor"))? + .backend_id; + let orient_max = + resolve("proof-of-max").ok_or_else(|| missing("the `proof-of-max` primitive"))?; + let orient_min = + resolve("proof-of-min").ok_or_else(|| missing("the `proof-of-min` primitive"))?; + // Column 0 = eclass, column 1 = view_proof (`eclass = f(children)`). + let max_ec = || MergeFn::Primitive(max_id, vec![MergeFn::OldCol(0), MergeFn::NewCol(0)]); + let min_ec = || MergeFn::Primitive(min_id, vec![MergeFn::OldCol(0), MergeFn::NewCol(0)]); + // The view proof paired with the larger / smaller eclass (both prove `eclass = f(children)`), + // selected by eclass order via a pure primitive rather than a conditional merge node. + let orient_args = || { + vec![ + MergeFn::OldCol(0), + MergeFn::OldCol(1), + MergeFn::NewCol(0), + MergeFn::NewCol(1), + ] + }; + let max_pf = || MergeFn::Primitive(orient_max, orient_args()); + let min_pf = || MergeFn::Primitive(orient_min, orient_args()); + // Trans(max_vp, Sym(min_vp)) : max_ec = min_ec (view proofs put the eclass on the left). The + // identity-column guard (col 0) skips this merge when the two eclasses are already equal, so + // the block stages the union unconditionally. + let edge = MergeFn::Lookup( + trans_id, + vec![max_pf(), MergeFn::Lookup(sym_id, vec![min_pf()])], + ); + // Stage the union edge into @UF; the merged value is (smaller eclass, its view proof). + Ok(Some(MergeFn::Block { + actions: vec![MergeAction::Set(uf_id, vec![max_ec(), min_ec(), edge])], + result: Box::new(MergeFn::Columns(vec![min_ec(), min_pf()])), + })) + } + + /// Build the self-referential union-find `:merge` for the encoding's single + /// per-sort UF function `@UF` (the whole union-find). `uf_id` is the + /// function's own backend id (see [`EGraph::peek_next_function_id`]). + /// `num_outputs` selects the encoding: + /// + /// - Term (`num_outputs == 1`), `@UF : (S) -> S`. To union `a, b`: `(set (@UF + /// (ordering-max a b)) (ordering-min a b))`. On an FD conflict the merge + /// keeps `ordering-min(old, new)` AND writes `(set (@UF ordering-max) ordering-min)` + /// into ITS OWN table (recording the displaced edge; strictly decreasing, hence finite). + /// - Proof (`num_outputs == 2`), `@UF : (S) -> (S, Proof)`. Value column 0 is + /// the parent, column 1 a proof `key = parent` (key on the LEFT). On a conflict + /// the merge keeps the smaller parent and stages the oriented displaced edge + /// `(@UF max_ec) = (values min_ec (Trans (Sym max_pf) min_pf))` into itself. + fn build_uf_self_merge( + &self, + uf_id: egglog_bridge::FunctionId, + num_outputs: usize, + ) -> Result { + use egglog_bridge::{MergeAction, MergeFn}; + let resolve = |name: &str| -> Result { + self.type_info + .get_prims(name) + .and_then(|p| p.first()) + .and_then(|p| p.context_ids[crate::Context::Write]) + .ok_or_else(|| { + Error::BackendError(format!( + "UF self-merge: primitive `{name}` not resolvable in write context" + )) + }) + }; + let min_id = resolve("ordering-min")?; + let max_id = resolve("ordering-max")?; + if num_outputs == 1 { + // The identity guard (the single value column, the parent) skips this merge when the + // parent is unchanged, so the body records the displaced edge unconditionally. + let min = || MergeFn::Primitive(min_id, vec![MergeFn::Old, MergeFn::New]); + let max = MergeFn::Primitive(max_id, vec![MergeFn::Old, MergeFn::New]); + return Ok(MergeFn::Block { + actions: vec![MergeAction::Set(uf_id, vec![max, min()])], + result: Box::new(min()), + }); + } + // Proof branch: value col 0 = parent, col 1 = proof (`key = parent`). + let backend_id = |name: &str| -> Result { + self.functions + .get(name) + .map(|f| f.backend_id) + .ok_or_else(|| { + Error::BackendError(format!( + "UF self-merge: proof constructor `{name}` missing" + )) + }) + }; + let trans_id = backend_id(&self.proof_state.proof_names.eq_trans_constructor)?; + let sym_id = backend_id(&self.proof_state.proof_names.eq_sym_constructor)?; + let orient_max = resolve("proof-of-max")?; + let orient_min = resolve("proof-of-min")?; + let max_ec = || MergeFn::Primitive(max_id, vec![MergeFn::OldCol(0), MergeFn::NewCol(0)]); + let min_ec = || MergeFn::Primitive(min_id, vec![MergeFn::OldCol(0), MergeFn::NewCol(0)]); + // The proof paired with the larger / smaller parent (both prove `key = that parent`), + // selected by parent order via a pure primitive rather than a conditional merge node. + let orient_args = || { + vec![ + MergeFn::OldCol(0), + MergeFn::OldCol(1), + MergeFn::NewCol(0), + MergeFn::NewCol(1), + ] + }; + let max_pf = || MergeFn::Primitive(orient_max, orient_args()); + let min_pf = || MergeFn::Primitive(orient_min, orient_args()); + // Trans(Sym(max_pf), min_pf) : max_ec = min_ec (UF proofs put the key on the left). The + // identity-column guard (col 0, the parent) skips this merge when the parent is unchanged, + // so the block stages the displaced edge unconditionally. + let edge = MergeFn::Lookup( + trans_id, + vec![MergeFn::Lookup(sym_id, vec![max_pf()]), min_pf()], + ); + // Stage the displaced edge back into @UF; the merged value is (smaller parent, its proof). + Ok(MergeFn::Block { + actions: vec![MergeAction::Set(uf_id, vec![max_ec(), min_ec(), edge])], + result: Box::new(MergeFn::Columns(vec![min_ec(), min_pf()])), + }) + } + fn declare_function(&mut self, decl: &ResolvedFunctionDecl) -> Result<(), Error> { let get_sort = |name: &String| match self.type_info.get_sort_by_name(name) { Some(sort) => Ok(sort.clone()), @@ -819,7 +1124,13 @@ impl EGraph { .iter() .map(get_sort) .collect::, _>>()?; - let output = get_sort(&decl.schema.output)?; + let outputs = decl + .schema + .outputs + .iter() + .map(get_sort) + .collect::, _>>()?; + let num_outputs = outputs.len(); let can_subsume = match decl.subtype { FunctionSubtype::Constructor => true, @@ -828,30 +1139,61 @@ impl EGraph { }; use egglog_bridge::{DefaultVal, MergeFn}; + // Identity/payload guard: the encoding's UF and congruence-view merges are idempotent on an + // unchanged eclass/parent (value column 0), so they opt into skip-on-unchanged — which is + // what lets them drop their equality guard. Ordinary functions keep the default (no guard). + let mut n_identity_vals = None; + let merge = if self + .proof_state + .self_merge_uf_functions + .contains(&*decl.name) + { + // Single self-referential UF `@UF` (term `(S) -> S`, proof + // `(S) -> (S, Proof)`): override the placeholder source `:merge` with the + // native self-referential merge, dispatched on value-column count. Its own + // backend id is the id `add_table` (below) will assign, peeked deterministically. + n_identity_vals = Some(1); + let uf_id = self.backend.peek_next_function_id(); + self.build_uf_self_merge(uf_id, num_outputs)? + } else if let Some(m) = self.native_congruence_merge(decl, &outputs[0], num_outputs)? { + // Term-encoding constructor view `(children) -> (eclass, proof)`: resolve congruence via + // a native :merge that stages the congruence edge into the per-sort UF table, instead of + // a rule-encoded self-join. + n_identity_vals = Some(1); + m + } else { + match decl.subtype { + FunctionSubtype::Constructor => MergeFn::UnionId, + FunctionSubtype::Custom => match &decl.merge { + Some(merge) => self.translate_merge_to_mergefn(merge)?, + // No merge clause: assert equality per output column. + None if num_outputs > 1 => { + MergeFn::Columns((0..num_outputs).map(|_| MergeFn::AssertEq).collect()) + } + None => MergeFn::AssertEq, + }, + } + }; let backend_id = self.backend.add_table(egglog_bridge::FunctionConfig { schema: input .iter() - .chain([&output]) + .chain(outputs.iter()) .map(|sort| sort.column_ty(&self.backend)) .collect(), + n_vals: num_outputs, + n_identity_vals, default: match decl.subtype { FunctionSubtype::Constructor => DefaultVal::FreshId, FunctionSubtype::Custom => DefaultVal::Fail, }, - merge: match decl.subtype { - FunctionSubtype::Constructor => MergeFn::UnionId, - FunctionSubtype::Custom => match &decl.merge { - None => MergeFn::AssertEq, - Some(expr) => self.translate_expr_to_mergefn(expr)?, - }, - }, + merge, name: decl.name.to_string(), can_subsume, }); let function = Function { decl: decl.clone(), - schema: ResolvedSchema { input, output }, + schema: ResolvedSchema { input, outputs }, can_subsume, backend_id, }; @@ -1594,11 +1936,17 @@ impl EGraph { name, uf, proof_func, + proof_ctors, .. } => { - // If the sort has a :internal-uf field, store the mapping for extraction + // If the sort has a :internal-uf field, store the mapping for extraction and + // register the UF as self-referential, so `declare_function` reattaches the + // native self-merge to a re-parsed desugared program (self-containment). if let Some(uf_name) = uf { - self.proof_state.uf_parent.insert(name.clone(), uf_name); + self.proof_state + .uf_parent + .insert(name.clone(), uf_name.clone()); + self.proof_state.self_merge_uf_functions.insert(uf_name); } // If the sort has a :internal-proof-func field, store the mapping for proof lookup. // This annotation is set by proof instrumentation and consumed here. @@ -1607,6 +1955,15 @@ impl EGraph { .proof_func_parent .insert(name.clone(), proof_func_name); } + // The Proof sort's :internal-proof-names records the global proof-constructor names. + // Repopulating them makes a re-parsed desugared program self-contained (the native + // congruence merge looks Trans/Sym up by name). + if let Some((congr, trans, sym)) = proof_ctors { + let names = &mut self.proof_state.proof_names; + names.congr_constructor = congr; + names.eq_trans_constructor = trans; + names.eq_sym_constructor = sym; + } log::info!("Declared sort {name}.") } ResolvedNCommand::Function(fdecl) => { @@ -1845,9 +2202,11 @@ impl EGraph { } if function_type.subtype != FunctionSubtype::Constructor { - match func.schema.output.name() { - "i64" | "String" | "Unit" => {} - s => panic!("Unsupported type {s} for input"), + for sort in &func.schema.outputs { + match sort.name() { + "i64" | "String" | "Unit" => {} + s => panic!("Unsupported type {s} for input"), + } } } @@ -1861,7 +2220,7 @@ impl EGraph { let mut row_schema = func.schema.input.clone(); if function_type.subtype == FunctionSubtype::Custom { - row_schema.push(func.schema.output.clone()); + row_schema.extend(func.schema.outputs.iter().cloned()); } log::debug!("{row_schema:?}"); @@ -2462,7 +2821,7 @@ fn resolve_function_container_target_with_context( .iter() .zip(&expected_inputs) .all(|(actual, expected)| actual.name() == expected.name()); - if !inputs_match || func_type.output.name() != output.name() { + if !inputs_match || func_type.output().name() != output.name() { let expected_input_names = expected_inputs .iter() .map(|sort| sort.name()) @@ -2479,7 +2838,7 @@ fn resolve_function_container_target_with_context( expected_input_names, output.name(), actual_input_names, - func_type.output.name(), + func_type.output().name(), ))); } @@ -2689,6 +3048,9 @@ impl<'a> BackendRule<'a> { let (p, args, ty) = self.prim(p, &atom.args, ctx); self.rb.query_prim(p, &args, ty).unwrap() } + ResolvedCall::Values(_) => { + unreachable!("`values` is lowered to the underlying function atom before query") + } } } } @@ -2717,6 +3079,9 @@ impl<'a> BackendRule<'a> { format!("{span}: call of primitive {name} failed") }) } + ResolvedCall::Values(_) => { + panic!("`values` cannot be bound as a single value") + } }; self.entries.insert(v, y.into()); } @@ -2725,16 +3090,18 @@ impl<'a> BackendRule<'a> { let x = self.entry(x); self.entries.insert(v, x); } - core::GenericCoreAction::Set(_, f, xs, y) => match f { + core::GenericCoreAction::Set(_, f, xs, ys) => match f { ResolvedCall::Primitive(..) => panic!("runtime primitive set!"), + ResolvedCall::Values(..) => panic!("`values` is not a settable function"), ResolvedCall::Func(f) => { let f = self.func(f); - let args = self.args(xs.iter().chain([y])); + let args = self.args(xs.iter().chain(ys.iter())); self.rb.set(f, &args) } }, core::GenericCoreAction::Change(span, change, f, args) => match f { ResolvedCall::Primitive(..) => panic!("runtime primitive change!"), + ResolvedCall::Values(..) => panic!("`values` is not a changeable function"), ResolvedCall::Func(f) => { let name = f.name.clone(); let can_subsume = self.functions[&f.name].can_subsume; @@ -3097,7 +3464,7 @@ mod tests { ResolvedExpr::Call(_, ResolvedCall::Func(func), children) => { assert_eq!(func.name, "$x"); assert!(children.is_empty()); - assert_eq!(func.output.name(), I64Sort.name()); + assert_eq!(func.output().name(), I64Sort.name()); } other => panic!("expected global function call rewrite, got {other:?}"), } diff --git a/egglog/src/prelude.rs b/egglog/src/prelude.rs index 3e8b6c5..b7ded6b 100644 --- a/egglog/src/prelude.rs +++ b/egglog/src/prelude.rs @@ -758,6 +758,7 @@ pub fn add_sort(egraph: &mut EGraph, name: &str) -> Result, E presort_and_args: None, uf: None, proof_func: None, + proof_ctors: None, unionable: true, }]) } @@ -773,7 +774,7 @@ pub fn add_function( span: span!(), name: name.to_owned(), schema, - merge, + merge: merge.map(GenericMerge::result_only), hidden: false, let_binding: false, term_constructor: None, @@ -824,7 +825,7 @@ macro_rules! datatype { stringify!($name), Schema { input: vec![$(stringify!($args).to_owned()),*], - output: stringify!($sort).to_owned(), + outputs: vec![stringify!($sort).to_owned()], }, [$($cost)*].first().copied(), false, diff --git a/egglog/src/proofs/merge_as_actions_design.md b/egglog/src/proofs/merge_as_actions_design.md new file mode 100644 index 0000000..880b6c6 --- /dev/null +++ b/egglog/src/proofs/merge_as_actions_design.md @@ -0,0 +1,150 @@ +# Design: `:merge` as an action block (retire the `MergeFn` DSL) + +Status: proposal. Not implemented. Intended as a follow-up **after** the term/proof +encoding PR lands on the current `MergeFn` implementation. + +## Problem + +The term/proof encoding needs merges that *perform actions* during an FD conflict — +stage a union edge into the per-sort union-find and compose a `Trans`/`Sym` proof — not +merely pick `old` or `new`. The encoding PR added a bespoke merge DSL in `egglog-bridge`: + +``` +MergeFn::{ Old, New, OldCol, NewCol, Const, Primitive, UnionId, AssertEq, // pick a value + Columns, // per-column + Seq, TableInsert, Construct, IfEq } // perform actions +``` + +The `Seq`/`TableInsert`/`Construct`/`IfEq` group is effectively a *second, weaker action +language* living in the bridge. It has its own ad-hoc lowering (`translate_expr_to_mergefn`) +and no principled type/context checking (e.g. `TupleMergeNotValues`, `TupleMergeArity` are +bespoke error paths). Meanwhile egglog already has a real action IR (`GenericAction`), a +constraint-based typechecker, a typed-capability/context system (`Context::Write`, +`PurePrim`/`ReadPrim`/`WritePrim`), and action lowering used by rules. + +**Goal:** express `:merge` with the existing rule action IR — uniform language, principled +typechecking — and delete the bespoke DSL, while keeping the backend surface small and +*general*. + +## Key insight + +`IfEq` in the encoding does exactly two jobs. Both can be removed: + +- **Role 1 — the guard:** "only stage a union when the eclass columns differ." This is the + *only* genuinely conditional decision in the merge body, and it is precisely "did the + identity column change?". Move it into the table's conflict detection via **identity vs + payload columns** (below). Then the merge body stages *unconditionally*. +- **Role 2 — proof orientation:** "which input proof matches the larger eclass?" This is a + *pure* function of `(old_ec, old_pf, new_ec, new_pf)` — no table access — so it becomes an + ordinary **pure primitive**, not a control-flow node. + +With both gone, the merge body is a straight-line action block terminated by a value, and +`If` is not needed in the IR at all. + +## Design + +### 1. Identity vs payload value columns (backend) + +Split a function's value columns into: +- **identity** columns: participate in FD-conflict detection; +- **payload** columns: carried alongside, not identifying. + +Conflict fires iff an *identity* column differs. A payload-only difference keeps the existing +row (no merge invocation, and the row is not marked dirty — strictly less churn than today's +"fire the merge, `IfEq` keeps old"). Precedent already exists: the subsume and timestamp +columns are non-identity payload. + +For the encoding, `@UF : (S) -> (S, Proof)` and the FD view `(children) -> (eclass, proof)` +declare **col 0 (eclass/parent) identity, col 1 (proof) payload**. + +Backend change: `FunctionConfig` gains an identity-column count/mask; `SortedWritesTable`'s +conflict path compares only identity columns. This is the one *added* backend capability, and +it is general (any "carry a witness/derivation next to the real value" use). + +### 2. `:merge` as an action block ending in an expression (egglog IR) + +A merge body becomes `{ action* ; result_expr }` with `old`/`new` (and `old0`/`new0`/… for +tuple output) bound. `result_expr` produces the merged value(s) (`(values e0 e1)` for tuple). +Actions are the existing set (`Set`/`Union`/constructor application/`Let`), restricted to +`Context::Write` (no queries, no `panic`) — the capability system enforces the restriction, +which is the "principled typechecking" payoff. + +The one genuine IR addition is a **value-producing action block** (a `{ stmts; expr }` form; +actions today are effect-only). Existing `:merge ` is the degenerate case (no actions, +just the result), so `:merge new`, `:merge (max old new)`, `:merge (values …)` keep working. + +### 3. Pure orientation primitives (proof layer) + +Register pure primitives that, given the two `(eclass, proof)` pairs, return the proof/eclass +of the smaller/larger endpoint (same insertion-order comparison as `ordering-min`/`-max`). +Pure ⇒ no dependency obligation ⇒ no `If`. These live in the proof-encoding layer. + +### 4. Lowering & dependency tracking + +`SortedWritesTable` already takes a merge **closure** (`MergeFn::to_callback` compiles to +`Box`). So the egglog layer compiles the typed merge action +block directly to that closure and the `MergeFn` enum is deleted. Read/write dependencies are +**derived** by walking the compiled action block (`set` targets → write deps; lookups / +constructor reads → read deps) and fed to the same strata ordering +(`DependencyGraph`/`has_read_deps`) that exists today — same obligation, better located than +`fill_deps`. Self-referential staging (`set` into the table's own id) still relies on the +core-relations self-write buffer pre-seed; that is fundamental to the single self-referential +UF and unchanged by this proposal. + +## Before / after (proof-mode `@UF` self-merge) + +Today (nested Rust `MergeFn` builders, `build_uf_self_merge`): + +``` +col0 = IfEq { a: OldCol(0), b: NewCol(0), then: OldCol(0), + els: Seq[ TableInsert(@UF, [max_ec, min_ec, + Construct(Trans, [Construct(Sym, [max_pf]), min_pf])]), + min_ec ] } +col1 = IfEq { a: OldCol(0), b: NewCol(0), then: OldCol(1), els: min_pf } +``` + +After (emitted egglog action block; col0 identity so no guard; `pf-of-*` are pure primitives): + +```text +:merge + (set (@UF (ordering-max old0 new0)) + (values (ordering-min old0 new0) + (Trans (Sym (pf-of-max old0 old1 new0 new1)) + (pf-of-min old0 old1 new0 new1)))) + (values (ordering-min old0 new0) + (pf-of-min old0 old1 new0 new1)) +``` + +## Removed / kept + +Removed: `MergeFn::{Seq, TableInsert, Construct, IfEq}` + the `ResolvedMergeFn` machinery; +`translate_expr_to_mergefn` and the bespoke merge typecheck errors; the hand-written +`build_uf_self_merge` / `native_congruence_merge` (become emitted action blocks). + +Kept: the `Columns` idea (subsumed — per-column results are just the result `(values …)`); +the core-relations self-write pre-seed; `ordering-min`/`-max` (plus the new pure orientation +primitives). + +## Risks / open questions + +1. **Perf.** The merge runs in the flush inner loop, once per FD conflict — the PR's headline + path (≈1.8× on proofs). The compiled action-block closure must be as tight as today's + `MergeFn::run`. Bench against current on `tests/math-microbenchmark.egg` proofs; this is the + acceptance gate. +2. **Backend subset-identity conflicts.** Confirm `SortedWritesTable`'s conflict/merge path can + compare a column subset cleanly (keep-old on payload-only diff). +3. **Semantic equivalence.** Prove col-0-identity is exactly today's `IfEq(OldCol0 == NewCol0)` + guard for *both* the view merge and the `@UF` self-merge before trusting it. +4. **Value-producing action block.** The IR addition and its typechecking interaction. +5. **Merge action subset.** Enforce "no query / no panic, `Context::Write`" via the capability + system. +6. **Back-compat.** Existing `:merge ` must keep parsing as a result-only block. + +## Sequencing + +1. Identity columns in the backend (+ test). +2. Value-producing action-block IR + typecheck. +3. Compile merge block → closure + derive deps. +4. Port the encoding's merges + add the pure orientation primitives. +5. Delete the `MergeFn` action DSL. +6. Bench proofs; gate on no regression. diff --git a/egglog/src/proofs/proof_checker.rs b/egglog/src/proofs/proof_checker.rs index 8938026..21bb999 100644 --- a/egglog/src/proofs/proof_checker.rs +++ b/egglog/src/proofs/proof_checker.rs @@ -58,13 +58,14 @@ pub(crate) fn run_merge( if let GenericNCommand::Function(func_decl) = cmd && func_decl.name == func_name { - // run the merge function for this function using eval_expr - let expr = func_decl.merge.as_ref().ok_or_else(|| { + // run the merge function for this function using eval_expr. The proof checker only + // evaluates the merged value; action-block merges aren't used under proofs. + let merge = func_decl.merge.as_ref().ok_or_else(|| { ProofCheckError::from(ProofCheckErrorKind::FunctionNotFound { function_name: func_name.to_string(), }) })?; - return eval_expr_with_subst("merge_function", expr, term_dag, &subst); + return eval_expr_with_subst("merge_function", &merge.result, term_dag, &subst); } } Err(ProofCheckErrorKind::FunctionNotFound { @@ -189,6 +190,7 @@ fn eval_expr_with_subst( }) })? } + ResolvedCall::Values(_) => panic!("`values` is not supported in proofs"), }, }; @@ -1031,6 +1033,7 @@ impl ProofStore { } } } + ResolvedCall::Values(_) => panic!("`values` is not supported in proofs"), } } } diff --git a/egglog/src/proofs/proof_encoding.md b/egglog/src/proofs/proof_encoding.md index c0866b8..7fa3a02 100644 --- a/egglog/src/proofs/proof_encoding.md +++ b/egglog/src/proofs/proof_encoding.md @@ -7,9 +7,11 @@ This makes proof production easier, since all equality reasoning is explicit and can be instrumented with proof tracking. The term encoding adds an explicit union-find structure per sort, and maintains it via rules that run during scheduled maintenance. -To speed up rebuild queries, each sort now uses two UF tables: - a constructor UF table (`UF_`) that stores raw parent edges, and a function UF table - (`UF_f`) that stores the current parent for each term as an index. +Each sort's union-find is a single self-referential function `UF_ : (S) -> S` that maps + each term to its parent; a term with no entry is its own representative (identity-on-miss). +Its native `:merge` (built in `EGraph::build_uf_self_merge`) + keeps the smaller endpoint on a conflict and unions the displaced parent back into `UF_`, + so a single `set` performs a union. For efficiency, every constructor becomes two tables: a term table that stores the actual terms, and a view table storing representative terms along with their e-class (stored as the leader term). The term encoding enables proof tracking, done at the @@ -40,100 +42,69 @@ Lowering the program with the term encoding expands to a bunch of new egglog, wh ```text (ruleset parent) -(ruleset single_parent) -(ruleset uf_function_index) (ruleset rebuilding) (ruleset rebuilding_cleanup) (ruleset delete_subsume_ruleset) ``` -*The new rulesets* orchestrate new rules for per-sort union-find tables (`parent` and `single_parent`), -building a fast function index over UF (`uf_function_index`), +*The new rulesets* orchestrate path compression on the per-sort union-find (`parent`), rebuild-time congruence (`rebuilding` + `rebuilding_cleanup`), and deferred deletions/subsumptions (`delete_subsume_ruleset`). ```text (run-schedule - (saturate - rebuilding_cleanup ;; cleanup merged rows - (saturate single_parent) ;; ensure each term points to single parent - (saturate parent) ;; transitively close parent links - (saturate uf_function_index) ;; mirror UF constructor rows into UF function index - rebuilding) ;; find new equalities via congruence - delete_subsume_ruleset) ;; process deletions/subsumptions + (seq + (saturate + rebuilding_cleanup ;; clean up merged rows + (saturate parent) ;; flatten union-find chains via path compression + rebuilding) ;; find new equalities via congruence + delete_subsume_ruleset)) ;; process deletions/subsumptions ``` *In-between* the original program's commands, the term encoding runs these rulesets to maintain egglog's invariants. ```text -(sort Math) -(function UF_Math (Math Math) Unit :merge old :internal-hidden) -(function UF_Mathf (Math) Math :merge new) +(sort Math :internal-uf UF_Math) +(function UF_Math (Math) Math :merge (ordering-min old new) :unextractable :internal-hidden) +(rule ((= b (UF_Math a)) + (= c (UF_Math b)) + (!= b c)) + ((set (UF_Math a) c)) + :ruleset parent :name "uf_path_compress") ``` -*The union-find* tables for each sort store the equivalence - classes of terms of that sort. -`UF_` remains the source of truth for UF maintenance updates, - while `UF_f` is a function-backed index used by rebuild rules. -`UF_` is always a function whose output type is `Unit` (without proof tracking) - or `Proof` (with proof tracking). -Using `:merge old` ensures that only the first proof/unit value is kept. -When proof tracking is enabled, proofs are stored directly in the UF table - (e.g., `(function UF_Math (Math Math) Proof :merge old :internal-hidden)`). - -```text -(rule ((UF_Math a b) - (UF_Math b c) - (!= b c)) - ((delete (UF_Math a b)) - (set (UF_Math a c) ())) - :ruleset parent :name "uf_update") -(rule ((UF_Math a b) - (UF_Math a c) - (!= b c) - (= (ordering-max b c) b)) - ((delete (UF_Math a b)) - (set (UF_Math b c) ())) - :ruleset single_parent :name "singleparentuf_update") -(rule ((UF_Math a b)) - ((set (UF_Mathf a) b)) - :ruleset uf_function_index :name "uf_function_index_update") -``` +*The union-find* for each sort is the single self-referential function `UF_`, + mapping each term to its parent. +`UF_` is `(S) -> S` without proof tracking, or `(S) -> (S Proof)` with proof tracking + (the extra column carries a proof of the parent edge). +A term with no row is its own representative, so `UF_` acts as an identity-on-miss lookup. +The source `:merge (ordering-min old new)` above is only a placeholder that lets the function + typecheck; `declare_function` replaces it with the native self-referential merge + (see `EGraph::build_uf_self_merge`), which on a conflicting + parent keeps the smaller endpoint and unions the displaced one back into `UF_`. +A single `set` on `UF_` therefore performs a union. *Union-find rules:* -A couple rules ensure the UF function is kept up to date as - equalities are added, and the indexing ruleset mirrors those rows - into the function UF. +The only maintenance rule is path compression (in the `parent` ruleset), which flattens + `a -> b -> c` chains to `a -> c`. We use the `ordering-max` and `ordering-min` egglog primitives to define an arbitrary ordering on terms based on insertion order, so that we can deterministically choose which term becomes the parent in the union-find structure. -**Important invariant:** every representative term must have a self-loop - entry in the constructor union-find table (e.g., `(UF_Math v v)`). -This is because the rebuild rules query the union-find for every - eq-sort column simultaneously, so a missing entry for any column - prevents the rule from firing even when other columns have changed. -Self-loops are added in `add_term_and_view` whenever a constructor - value is created. -The `uf_function_index` ruleset then copies those rows into - `UF_f`, so representatives also satisfy `(= (UF_f v) v)`. -We may want to remove this invariant in the future if we move - to a different encoding, saving some space and time. - ```text (sort view) -(constructor Add (i64 i64) Math) +(constructor Add (i64 i64) Math :unextractable :internal-hidden) (function AddView (i64 i64 Math) Unit :merge old :internal-term-constructor Add) -(constructor to_delete_Add (i64 i64) view) -(constructor to_subsume_Add (i64 i64) view) +(constructor to_delete_Add (i64 i64) view :internal-hidden) +(constructor to_subsume_Add (i64 i64) view :internal-hidden) ``` Each constructor in the original program is expanded to a term table (`Add`), a view table (`AddView`), and helpers for deferred deletion/subsumption (`to_delete_Add`, `to_subsume_Add`). -The view table is always a function whose output type is `Unit` (without proof tracking) +The view table is a function whose output type is `Unit` (without proof tracking) or `Proof` (with proof tracking), with `:merge old`. A view table stores "canonicalized" terms and their e-class representative. A canonicalized term has representative terms for its children. @@ -141,41 +112,41 @@ The last column of the view table is the representative term for the e-class. The view tables are kept up to date during rebuilding. ```text -(rule ((AddView c0 c1 new) - (AddView c0 c1 old) +(rule ((= v (AddView c0 c1 new)) + (= v1 (AddView c0 c1 old)) (!= old new) (= (ordering-max old new) new)) - ((set (UF_Math (ordering-max new old) (ordering-min new old)) ())) - :ruleset rebuilding :name "congruence_rule") -(rule ((= v9 (AddView c0 c1 c2)) - (= c2_leader (UF_Mathf c2)) - (guard - (or (bool-!= c2 c2_leader)))) + ((set (UF_Math (ordering-max new old)) (ordering-min new old))) + :ruleset rebuilding :name "congruence_rule" :internal-include-subsumed) +(rule ((= v2 (AddView c0 c1 c2)) + (= c2_leader (UF_Math c2)) + (!= c2 c2_leader)) ((set (AddView c0 c1 c2_leader) ()) (delete (AddView c0 c1 c2))) - :ruleset rebuilding :name "rebuild_rule") + :ruleset rebuilding :name "rebuild_rule" :internal-include-subsumed) ``` -For each constructor, we add a congruence rule and a rebuild rule. -The congruence rule adds equalities to the union-find table when two constructor applications +For each constructor, we add a congruence rule and rebuild rules. +The congruence rule adds an equality to the union-find when two constructor applications have equal arguments. -The rebuild rule updates view tables so that views - point to representative terms for child e-classes. -Rebuild rules read representatives from `UF_f` (function lookup) - rather than joining directly on `UF_`, - which avoids expensive UF joins during rebuilding. +The rebuild rules keep the view pointing to representative terms. +They are *fanned out*, one per eq-sort column (only `Add`'s `Math` output column here, since its + `i64` children are not eq-sorts): each rule replaces a single column with its `UF_` leader + and re-sets the row. +Because `UF_` has no row for a canonical term (identity-on-miss), a column already at its + leader fails the `(!= c c_leader)` guard, so no self-loops or default lookups are needed. ```text -(function v2 () Math :no-merge) -(set (v2) (Add 1 2)) -(set (AddView 1 2 (v2)) ()) -(set (UF_Math (v2) (v2)) ()) +(function v3 () Math :no-merge :unextractable :internal-let) +(set (v3) (Add 1 2)) +(set (AddView 1 2 (v3)) ()) ``` Above is the desugaring for `(Add 1 2)`. We add to both view and term tables whenever we evaluate a constructor or function application. -The self-loop `(UF_Math (v2) (v2))` initializes the e-class for the new term. +The new term needs no `UF_` entry: with identity-on-miss, a term with no row is already + its own representative. It's straightforward except for global variables. Since global variables are not allowed after this pass, we use functions with no arguments to represent them @@ -183,23 +154,20 @@ Since global variables are not allowed after this pass, ```text -(rule ((= v3 (AddView a b v4))) - ((let v5 (Add a b)) - (set (AddView a b v5) ()) - (set (UF_Math v5 v5) ()) - (let v6 (Add b a)) - (set (AddView b a v6) ()) - (set (UF_Math v6 v6) ()) - (set (UF_Math (ordering-max v5 v6) (ordering-min v5 v6)) ())) +(rule ((= v5 (AddView a b v4))) + ((let v7 (Add a b)) + (set (AddView a b v7) ()) + (let v8 (Add b a)) + (set (AddView b a v8) ()) + (set (UF_Math (ordering-max v7 v8)) (ordering-min v7 v8))) :name "commutativity") ``` Here we have the instrumented commutativity rule. The query uses the view table to find the canonical e-node. -The actions add to the term table, add to the view table, - and add an equality to the union-find table. -We add an equality to the union-find table for the two terms, using the `ordering-max` and - `ordering-min` egglog primitives to correctly choose a parent. +The actions add to the term and view tables, then add an equality to the union-find. +We add the equality with a single `set` on `UF_`, using the `ordering-max` and + `ordering-min` egglog primitives to deterministically choose the parent. @@ -293,7 +261,7 @@ The header defines the proof format corresponding to [`RawProof`](crate::proofs: See the proof header in `proof_encoding_helpers.rs` for details. ```text -(function MathProof (Math) Proof :merge old) +(function MathProof (Math) Proof :merge old :unextractable :internal-hidden) ``` Every sort gets a proof table storing @@ -302,64 +270,63 @@ The proof proves a proposition `t = t` for input term `t`. We store the oldest proof currently. -When proof tracking is enabled, the union-find table's output type is `Proof` instead of `Unit`: +When proof tracking is enabled, the single self-referential union-find carries a +proof in a second value column: ```text -(function UF_Math (Math Math) Proof :merge old :internal-hidden) +(function UF_Math (Math) (Math Proof) :merge (values old0 old1) :unextractable :internal-hidden) ``` -If term `a` has parent `b`, `(UF_Math a b)` returns a - proof of `a = b`. -The path compression and single-parent rules are instrumented to produce - proofs using symmetry (`Sym`) and transitivity (`Trans`) as needed. +If term `k` has parent `p`, `(UF_Math k)` returns `(values p proof)` where `proof` +proves `k = p` (the key on the LEFT). The native self-referential `:merge` (built in +`EGraph::build_uf_self_merge`) keeps the smaller parent on a conflict and stages the oriented +displaced edge back into `UF_Math`, composing proofs with `Trans`/`Sym`. Path compression flattens +cross-key chains via `Trans`. -Similarly, the view table's output type is `Proof` instead of `Unit`: +Similarly, the constructor view is a functional-dependency tuple carrying a proof: ```text -(function AddView (i64 i64 Math) Proof :merge old :internal-term-constructor Add) +(function AddView (i64 i64) (Math Proof) :merge (values old0 old1) :internal-term-constructor Add) ``` -Recall that view tables store a term - along with the e-class representative. -For a term `t` with representative `r`, - the proof (output of the view function) proves that `r = t`. -The direction is important, making - proof production easier later. -We store the earliest proof (`:merge old`). +The view maps a term's canonicalized children to `(eclass, proof)`, where `proof` +proves `eclass = f(children)` (the eclass on the LEFT). Its native congruence `:merge` +(built in `EGraph::native_congruence_merge`) keeps the smaller eclass on a functional-dependency +conflict and stages the oriented union edge into `UF_Math`. ```text -(rule (;; query the view function directly for the proof - (= v9 (AddView a b v8))) +(rule (;; query the view for its eclass and proof (proof that eclass = (Add a b)) + (= (values v11 v12) (AddView a b))) (;; proof list, one per line of the original query - (let v10 (PCons v9 (PNil ))) - - (let v11 (Add a b)) - ;; Proof that Add a b = Add a b - (let v12 (Rule "commutativity" v10 (AstMath v11) (AstMath v11))) - ;; Setting the proof for Add a b - (set (MathProof v11) v12) + (let v13 (PCons v12 (PNil))) - ;; Update the view function (set instead of constructor insertion) - (set (AddView a b v11) v12) + (let v14 (Add a b)) + ;; Proof that Add a b = Add a b + (let v15 (Rule "commutativity" v13 (AstMath v14) (AstMath v14))) + ;; Set the proof for Add a b + (set (MathProof v14) v15) + ;; Update the FD view: children -> (eclass, proof) + (set (AddView a b) (values v14 v15)) - (let v13 (Add b a)) + (let v16 (Add b a)) ;; Proof that Add b a = Add b a - (let v14 (Rule "commutativity" v10 (AstMath v13) (AstMath v13))) - (set (MathProof v13) v14) - (set (AddView b a v13) v14) - - ;; Store a proof that (Add a b) = (Add b a). - (set (UF_Math (ordering-max v11 v13) (ordering-min v11 v13)) - (Rule "commutativity" v10 (AstMath (ordering-max v11 v13)) (AstMath (ordering-min v11 v13))))) + (let v17 (Rule "commutativity" v13 (AstMath v16) (AstMath v16))) + (set (MathProof v16) v17) + (set (AddView b a) (values v16 v17)) + + ;; Union (Add a b) and (Add b a), storing a proof of their equality. + (set (UF_Math (ordering-max v14 v16)) + (values (ordering-min v14 v16) + (Rule "commutativity" v13 (AstMath (ordering-max v14 v16)) (AstMath (ordering-min v14 v16)))))) :name "commutativity") ``` Instrumented rules with proof tracking query the view function directly (since the proof is its output column), then construct proofs for each action. -The structure is the same as term mode — view updates use `set`, UF updates use `set` — - but the values stored are `Proof` terms instead of `()`. +The structure is the same as term mode — view and UF updates both use `set` — + but the values stored carry `Proof` terms instead of `()`. For nested terms, congruence proofs are built to ensure the proof terms match the original queries. diff --git a/egglog/src/proofs/proof_encoding.rs b/egglog/src/proofs/proof_encoding.rs index a51701e..8376ef3 100644 --- a/egglog/src/proofs/proof_encoding.rs +++ b/egglog/src/proofs/proof_encoding.rs @@ -7,7 +7,6 @@ use crate::*; #[derive(Clone)] pub(crate) struct EncodingState { pub uf_parent: HashMap, - pub uf_function: HashMap, /// Maps sort name -> proof function name (set from :internal-proof-func annotation). pub proof_func_parent: HashMap, /// Function name -> (hidden current-value function, input arity). The @@ -22,6 +21,11 @@ pub(crate) struct EncodingState { pub original_typechecking: Option>, pub proofs_enabled: bool, pub proof_testing: bool, + /// Non-proof term encoding: the per-sort single self-referential union-find + /// function names (`@UF : (S) -> S`). `declare_function` (lib.rs) attaches + /// the native self-referential `:merge` (see `EGraph::build_uf_self_merge`) + /// to these, overriding their placeholder source `:merge`. + pub self_merge_uf_functions: HashSet, pub proof_names: EncodingNames, /// Test-only knob: annotate RHS-reading rules `:naive` (the safe /// whole-database baseline) instead of `:unsafe-seminaive`, so tests can @@ -33,7 +37,6 @@ impl EncodingState { pub(crate) fn new(symbol_gen: &mut SymbolGen) -> Self { Self { uf_parent: HashMap::default(), - uf_function: HashMap::default(), proof_func_parent: HashMap::default(), merge_current: HashMap::default(), term_header_added: false, @@ -41,6 +44,7 @@ impl EncodingState { proofs_enabled: false, proof_names: EncodingNames::new(symbol_gen), proof_testing: false, + self_merge_uf_functions: HashSet::default(), force_proof_naive: false, } } @@ -71,6 +75,12 @@ impl<'a> ProofInstrumentor<'a> { let uf_name = self.uf_name(type_name); let smaller = format!("(ordering-min {lhs} {rhs})"); let larger = format!("(ordering-max {lhs} {rhs})"); + // Non-proof: the union-find is a single self-referential function + // `@UF : (S) -> S` keyed by the larger endpoint; the native `:merge` + // resolves conflicts (keeps the min, unions the displaced parent). + if !self.egraph.proof_state.proofs_enabled { + return format!("(set ({uf_name} {larger}) {smaller})"); + } let proof = if self.egraph.proof_state.proofs_enabled { let to_ast_constructor = self .proof_names() @@ -94,7 +104,9 @@ impl<'a> ProofInstrumentor<'a> { } else { "()".to_string() }; - format!("(set ({uf_name} {larger} {smaller}) {proof})") + // Single self-referential UF `@UF : (S) -> (S, Proof)` keyed by the larger + // endpoint; value column 0 = smaller parent, column 1 = `proof` (larger = smaller). + format!("(set ({uf_name} {larger}) (values {smaller} {proof}))") } /// The parent table is the database representation of a union-find datastructure. @@ -102,113 +114,88 @@ impl<'a> ProofInstrumentor<'a> { /// Also, we have a rule that maintains the invariant that each term points to its /// canonical representative. fn declare_sort(&mut self, sort_name: &str) -> Vec { - let pname = self.uf_name(sort_name); - let uf_function_name = self.uf_function_name(sort_name); - let fresh_name = self.egraph.parser.symbol_gen.fresh("uf_update"); - let uf_function_index_name = self.egraph.parser.symbol_gen.fresh("uf_function_index"); + if self.egraph.proof_state.proofs_enabled { + self.declare_sort_proof(sort_name) + } else { + self.declare_sort_non_proof(sort_name) + } + } + /// Proof-mode `declare_sort`: the union-find is a SINGLE self-referential + /// function `@UF : (S) -> (S, Proof)` (native tuple, reusing `uf_name`). Value + /// column 0 is the parent, column 1 a proof `key = parent`. Union `a, b` writes + /// `(set (@UF (ordering-max a b)) (values (ordering-min a b) proof))`; + /// `declare_function` attaches the native self-referential `:merge` (see + /// [`EGraph::build_uf_self_merge`]), overriding the `:merge (values old0 old1)` + /// placeholder. A proof-composing `path_compress` rule flattens cross-key chains + /// via `Trans`. Also emits the per-sort `term_proof` table and AST constructor. + fn declare_sort_proof(&mut self, sort_name: &str) -> Vec { + let uf_name = self.uf_name(sort_name); + let proof_type = self.proof_type_str().to_string(); + let term_proof_name = self.term_proof_name(sort_name); + let add_to_ast_code = self.add_to_ast(sort_name); + let fresh_name = self.egraph.parser.symbol_gen.fresh("uf_path_compress"); let path_compress_ruleset_name = self.proof_names().path_compress_ruleset_name.clone(); - let single_parent_ruleset_name = self.proof_names().single_parent_ruleset_name.clone(); - let uf_function_index_ruleset_name = - self.proof_names().uf_function_index_ruleset_name.clone(); + let trans = self.proof_names().eq_trans_constructor.clone(); - let proof_type = self.proof_type_str().to_string(); + // Mark `@UF` so `declare_function` attaches the native self-referential proof merge. + self.egraph + .proof_state + .self_merge_uf_functions + .insert(uf_name.clone()); + + let code = format!( + "{add_to_ast_code} + (function {term_proof_name} ({sort_name}) {proof_type} :merge old :internal-hidden) + (function {uf_name} ({sort_name}) ({sort_name} {proof_type}) :merge (values old0 old1) :unextractable :internal-hidden) + ;; path compression: a->b (pb: a=b), b->c (pc: b=c) => a->c (Trans pb pc: a=c) + (rule ((= (values b pb) ({uf_name} a)) + (= (values c pc) ({uf_name} b)) + (!= b c)) + ((set ({uf_name} a) (values c ({trans} pb pc)))) + :ruleset {path_compress_ruleset_name} + :name \"{fresh_name}\") + " + ); - // In proof mode, path compression composes proofs via Trans/Sym. - // In term mode, the proof output is Unit and we just write (). - let (path_compress_query, path_compress_action, single_parent_query, single_parent_action) = - if self.egraph.proof_state.proofs_enabled { - let p1_fresh = self.egraph.parser.symbol_gen.fresh("p1"); - let p2_fresh = self.egraph.parser.symbol_gen.fresh("p2"); - let trans = self.proof_names().eq_trans_constructor.clone(); - let sym = self.proof_names().eq_sym_constructor.clone(); - ( - format!( - "(= {p1_fresh} ({pname} a b)) - (= {p2_fresh} ({pname} b c))" - ), - format!( - "(delete ({pname} a b)) - (set ({pname} a c) ({trans} {p1_fresh} {p2_fresh}))" - ), - format!( - "(= {p1_fresh} ({pname} a b)) - (= {p2_fresh} ({pname} a c))" - ), - format!( - "(delete ({pname} a b)) - (set ({pname} b c) ({trans} ({sym} {p1_fresh}) {p2_fresh}))" - ), - ) - } else { - ( - format!("({pname} a b)\n ({pname} b c)"), - format!( - "(delete ({pname} a b))\n (set ({pname} a c) ())" - ), - format!("({pname} a b)\n ({pname} a c)"), - format!( - "(delete ({pname} a b))\n (set ({pname} b c) ())" - ), - ) - }; + self.parse_program(&code) + } - // In proof mode, UF function index stores (leader, proof) pairs. - // In term mode, it just stores the leader. - let (uf_function_output_type, uf_pair_sort_decl, uf_index_query, uf_index_action) = - if self.egraph.proof_state.proofs_enabled { - let pair_sort = self.uf_pair_sort_name(sort_name); - let proof_fresh = self.egraph.parser.symbol_gen.fresh("uf_idx_proof"); - ( - pair_sort.clone(), - format!("(sort {pair_sort} (Pair {sort_name} {proof_type}))"), - format!("(= {proof_fresh} ({pname} a b))"), - format!("(set ({uf_function_name} a) (pair b {proof_fresh}))"), - ) - } else { - ( - sort_name.to_string(), - "".to_string(), - format!("({pname} a b)"), - format!("(set ({uf_function_name} a) b)"), - ) - }; + /// Non-proof `declare_sort`: the whole union-find is a SINGLE + /// self-referential function `@UF : (S) -> S` (reusing `uf_name`), with a + /// native `:merge` (built in [`EGraph::build_uf_self_merge`], attached in + /// `declare_function`). Union `a, b` writes `(set (@UF (ordering-max a b)) + /// (ordering-min a b))`; on an FD conflict the merge keeps the min and unions + /// the displaced parent back into `@UF`, so the `single_parent` / + /// `uf_function_index` rulesets are unnecessary. `path_compress` is kept to + /// flatten cross-key chains (the encoder reads `@UF` as a single + /// identity-on-miss lookup). + fn declare_sort_non_proof(&mut self, sort_name: &str) -> Vec { + let uf_name = self.uf_name(sort_name); + let fresh_name = self.egraph.parser.symbol_gen.fresh("uf_path_compress"); + let path_compress_ruleset_name = self.proof_names().path_compress_ruleset_name.clone(); - let mut code = format!( - "{uf_pair_sort_decl} - (function {pname} ({sort_name} {sort_name}) {proof_type} :merge old :internal-hidden) - (function {uf_function_name} ({sort_name}) {uf_function_output_type} :merge new :unextractable :internal-hidden) - ;; performs path compression, ensuring each term points to the representative - (rule ({path_compress_query} + // Mark `@UF` so `declare_function` attaches the native self-referential + // merge instead of lowering the placeholder `:merge (ordering-min old new)`. + self.egraph + .proof_state + .self_merge_uf_functions + .insert(uf_name.clone()); + + // The source `:merge (ordering-min old new)` is a valid placeholder so the + // function typechecks; `declare_function` overrides it with the native + // self-referential merge (recursive parent-union). + let code = format!( + "(function {uf_name} ({sort_name}) {sort_name} :merge (ordering-min old new) :unextractable :internal-hidden) + (rule ((= b ({uf_name} a)) + (= c ({uf_name} b)) (!= b c)) - ({path_compress_action}) + ((set ({uf_name} a) c)) :ruleset {path_compress_ruleset_name} :name \"{fresh_name}\") - ;; ensures each term has only one parent - (rule ({single_parent_query} - (!= b c) - (= (ordering-max b c) b)) - ({single_parent_action}) - :ruleset {single_parent_ruleset_name} - :name \"singleparent{fresh_name}\") - ;; mirrors UF rows into a function-backed UF index for faster rebuild lookups - (rule ({uf_index_query}) - ({uf_index_action}) - :ruleset {uf_function_index_ruleset_name} - :name \"{uf_function_index_name}\") " ); - if self.egraph.proof_state.proofs_enabled { - let term_proof_name = self.term_proof_name(sort_name); - let add_to_ast_code = self.add_to_ast(sort_name); - code = format!( - "{add_to_ast_code} - (function {term_proof_name} ({sort_name}) {proof_type} :merge old :internal-hidden) - {code}" - ); - } - self.parse_program(&code) } @@ -228,6 +215,30 @@ impl<'a> ProofInstrumentor<'a> { let delete_subsume_ruleset = self.proof_names().delete_subsume_ruleset_name.clone(); let fresh_name = self.egraph.parser.symbol_gen.fresh("delete_rule"); + // The FD tuple view is keyed by children only, so match its value tuple to delete/subsume + // by key (the bridge re-reads every value column when subsuming a tuple-output view). + if self.native_merge_view(fdecl) { + let e = self.fresh_var(); + let pf = self.fresh_var(); + let e2 = self.fresh_var(); + let pf2 = self.fresh_var(); + // Deletion removes the row by key; subsumption marks it subsumed (kept for size/proofs + // but excluded from matching). + return format!( + "(rule (({to_delete_name} {child_names}) + (= (values {e} {pf}) ({view_name} {child_names}))) + ((delete ({view_name} {child_names})) + (delete ({to_delete_name} {child_names}))) + :ruleset {delete_subsume_ruleset} + :name \"{fresh_name}\") + (rule (({subsumed_name} {child_names}) + (= (values {e2} {pf2}) ({view_name} {child_names}))) + ((subsume ({view_name} {child_names}))) + :ruleset {delete_subsume_ruleset} + :name \"{fresh_name}_subsume\")" + ); + } + format!( "(rule (({to_delete_name} {child_names}) ({view_name} {child_names} out)) @@ -297,8 +308,10 @@ impl<'a> ProofInstrumentor<'a> { "()".to_string() }; let mut merge_fn_code = vec![]; + // Proof instrumentation tracks the merged *value*; a `:merge` action block's effects are + // not proof-tracked (action-block merges under proofs are unsupported). let merge_fn_var = self.instrument_action_expr( - merge_fn, + &merge_fn.result, &mut merge_fn_code, &Justification::Merge(name.clone(), p1_fresh.clone(), p2_fresh.clone()), ); @@ -323,7 +336,7 @@ impl<'a> ProofInstrumentor<'a> { let term_and_proof = self.update_view(name, &updated, &proof_var); let cleanup_constructor = self.egraph.parser.symbol_gen.fresh("mergecleanup"); let fresh_sort = self.egraph.parser.symbol_gen.fresh("mergecleanupsort"); - let output_sort = fdecl.schema.output.clone(); + let output_sort = fdecl.schema.output().clone(); // The first runs the merge function adding a new row. // The second deletes rows with old values for the old variable, while the third deletes rows with new values for the new variable. @@ -391,7 +404,7 @@ impl<'a> ProofInstrumentor<'a> { // We also have a proof that other eclass representative r_2 = f(c_1,...,c_n), the same term. // We want a proof that r1 = r2, which we get by transitivity. let union_code = self.union( - &fdecl.schema.output, + fdecl.schema.output(), "new", "old", &Justification::Proof(format!("({trans} {prf1} ({sym} {prf2}))",)), @@ -433,6 +446,9 @@ impl<'a> ProofInstrumentor<'a> { &view_name, &rebuilding_ruleset, ) + } else if self.native_merge_view(fdecl) { + // Congruence is resolved by the view's native `:merge`; no rule needed. + String::new() } else { self.handle_congruence(fdecl, &child_names, &rebuilding_ruleset) } @@ -443,9 +459,15 @@ impl<'a> ProofInstrumentor<'a> { /// The view table stores child terms and their eclass. /// The view table is mutated using delete, but we never delete from term tables. /// We re-use the original name of the function for the term table. + /// Whether this function's view is the proof-mode functional-dependency tuple view + /// `(children) -> (eclass, proof)` whose native `:merge` resolves congruence. + fn native_merge_view(&self, fdecl: &ResolvedFunctionDecl) -> bool { + fdecl.subtype == FunctionSubtype::Constructor && self.egraph.proof_state.proofs_enabled + } + fn term_and_view(&mut self, fdecl: &ResolvedFunctionDecl) -> Vec { let schema = &fdecl.schema; - let out_type = schema.output.clone(); + let out_type = schema.output().clone(); let name = &fdecl.name; let view_name = self.view_name(&fdecl.name); @@ -459,14 +481,14 @@ impl<'a> ProofInstrumentor<'a> { if fdecl.subtype == FunctionSubtype::Constructor { "".to_string() } else { - schema.output.to_string() + schema.output().to_string() } ); let view_sorts = format!("{in_sorts} {out_type}"); let proof_constructors = self.proof_functions(fdecl, &view_sorts); let view_sort = if fdecl.subtype == FunctionSubtype::Constructor { - schema.output.clone() + schema.output().clone() } else { fresh_sort.clone() }; @@ -499,9 +521,20 @@ impl<'a> ProofInstrumentor<'a> { if fdecl.internal_let { view_flags.push_str(" :internal-let"); } - let view_decl = format!( - "(function {view_name} ({view_sorts}) {proof_type} :merge old :internal-term-constructor {name}{view_flags})" - ); + // In proof mode, a constructor's view is a functional-dependency tuple + // `(children) -> (eclass, proof)` whose native `:merge` resolves congruence (see + // `EGraph::native_congruence_merge`). Other functions keep the `(children eclass) -> proof` + // form with a congruence/merge rule. + let fd_view = self.native_merge_view(fdecl); + let view_decl = if fd_view { + format!( + "(function {view_name} ({in_sorts}) ({out_type} {proof_type}) :merge (values old0 old1) :internal-term-constructor {name}{view_flags})" + ) + } else { + format!( + "(function {view_name} ({view_sorts}) {proof_type} :merge old :internal-term-constructor {name}{view_flags})" + ) + }; self.parse_program(&format!( " (sort {fresh_sort}) @@ -523,6 +556,12 @@ impl<'a> ProofInstrumentor<'a> { /// Rules that update the views when children change. fn rebuilding_rules(&mut self, fdecl: &ResolvedFunctionDecl) -> Vec { + if self.native_merge_view(fdecl) { + return self.rebuilding_rules_fd(fdecl); + } + if !self.egraph.proof_state.proofs_enabled { + return self.rebuilding_rules_fanout(fdecl); + } let types = fdecl.resolved_schema.view_types(); // Check if there are any eq-sort columns at all; if not, no rebuild rule needed. @@ -530,125 +569,182 @@ impl<'a> ProofInstrumentor<'a> { return vec![]; } + // Proof-mode fanout: the single-key `@UF` has no row for a canonical column + // (identity-on-miss), so a combined all-columns join would stall whenever any + // column is already canonical. Emit one rule per eq-sort view column, each with a + // single `@UF` lookup, mirroring [`Self::rebuilding_rules_fanout`] but composing a + // proof: `Congr` for a child update, `Trans(Sym uf_pf, view_prf)` for the + // representative (last) column of a constructor view. let view_name = self.view_name(&fdecl.name); let child = |i: usize| format!("c{i}_"); let children_vec: Vec = (0..types.len()).map(child).collect(); let children = format!("{}", ListDisplay(&children_vec, " ")); + let rebuilding_ruleset = self.proof_names().rebuilding_ruleset_name.clone(); + let trans = self.proof_names().eq_trans_constructor.clone(); + let sym = self.proof_names().eq_sym_constructor.clone(); + let congr = self.proof_names().congr_constructor.clone(); - // For each eq-sort column, look up its leader via the UF table. - // For non-eq-sort columns, the leader is the same as the original. - let mut uf_queries = vec![]; - let mut leader_vars: Vec = vec![]; - let mut bool_neq_exprs = vec![]; - let mut uf_proof_vars: Vec> = vec![]; - + let mut rules = String::new(); for (i, ty) in types.iter().enumerate() { - if ty.is_eq_sort() { - let leader_var = format!("c{i}_leader_"); - let uf_function_name = self.uf_function_name(ty.name()); - let ci = child(i); - - if self.egraph.proof_state.proofs_enabled { - // UF function index returns a Pair(leader, proof); one lookup gives both - let pair_var = self.fresh_var(); - let proof_var = format!("(pair-second {pair_var})"); - uf_queries.push(format!( - "(= {pair_var} ({uf_function_name} {ci})) - (= {leader_var} (pair-first {pair_var}))" - )); - uf_proof_vars.push(Some(proof_var)); - } else { - uf_queries.push(format!("(= {leader_var} ({uf_function_name} {ci}))")); - uf_proof_vars.push(None); - } - - bool_neq_exprs.push(format!("(bool-!= {ci} {leader_var})")); - leader_vars.push(leader_var); - } else { - leader_vars.push(child(i)); - uf_proof_vars.push(None); + if !ty.is_eq_sort() { + continue; } + let ci = child(i); + let leader = format!("c{i}_leader_"); + let uf_name = self.uf_name(ty.name()); + let (query_view, view_prf) = self.query_view_and_get_proof(&fdecl.name, &children_vec); + let uf_prf = self.fresh_var(); + let new_pf = self.fresh_var(); + let proof_code = + if fdecl.subtype == FunctionSubtype::Constructor && i == types.len() - 1 { + format!("(let {new_pf} ({trans} ({sym} {uf_prf}) {view_prf}))") + } else { + format!("(let {new_pf} ({congr} {view_prf} {i} {uf_prf}))") + }; + let mut updated = children_vec.clone(); + updated[i] = leader.clone(); + let updated_view = self.update_view(&fdecl.name, &updated, &new_pf); + let fresh_name = self.egraph.parser.symbol_gen.fresh("rebuild_rule"); + rules.push_str(&format!( + "(rule ({query_view} + (= (values {leader} {uf_prf}) ({uf_name} {ci})) + (!= {ci} {leader})) + ( + {proof_code} + {updated_view} + (delete ({view_name} {children})) + ) + :ruleset {rebuilding_ruleset} :name \"{fresh_name}\" :internal-include-subsumed)\n" + )); } + self.parse_program(&rules) + } - let uf_query_str = uf_queries.join("\n "); - let or_expr = format!("(or {})", bool_neq_exprs.join("\n ")); - let filter_query = format!("(guard {or_expr})"); - - // Build the updated children: use leader_var for eq-sort columns, original for others. - let children_updated: Vec = leader_vars.clone(); - - let fresh_name = self.egraph.parser.symbol_gen.fresh("rebuild_rule"); - let (query_view, view_prf) = self.query_view_and_get_proof(&fdecl.name, &children_vec); - - // Build proof code if proofs are enabled. - // We chain congruence proofs for each updated child and a transitivity proof - // for the representative (last column) update. - let (pf_code, pf_var) = if self.egraph.proof_state.proofs_enabled { - let eq_trans_constructor = self.proof_names().eq_trans_constructor.clone(); - let congr_constructor = self.proof_names().congr_constructor.clone(); - let sym_constructor = self.proof_names().eq_sym_constructor.clone(); - - // Start with the view proof and apply congruence for each eq-sort child - // (excluding the last column if this is a constructor, since that's the representative). - let mut current_proof = view_prf.clone(); - let mut proof_code_parts = vec![]; - - for (i, ty) in types.iter().enumerate() { - if !ty.is_eq_sort() { - continue; - } - - let uf_prf = uf_proof_vars[i].as_ref().unwrap(); + /// Non-proof rebuild: one rule per eq-sort view column (children and the + /// constructor's eclass/output column alike). The single-key `@UF` has no row + /// for a canonical (root) node, so a column that is already canonical simply + /// doesn't match — no self-loops or default-lookups needed. Each rule replaces + /// ONE column with its leader, re-setting the view row and deleting the stale + /// one; a row with several stale columns is canonicalized across rebuild + /// iterations. + fn rebuilding_rules_fanout(&mut self, fdecl: &ResolvedFunctionDecl) -> Vec { + let types = fdecl.resolved_schema.view_types(); + if !types.iter().any(|t| t.is_eq_sort()) { + return vec![]; + } + let view_name = self.view_name(&fdecl.name); + let child = |i: usize| format!("c{i}_"); + let children_vec: Vec = (0..types.len()).map(child).collect(); + let children = format!("{}", ListDisplay(&children_vec, " ")); + let rebuilding_ruleset = self.proof_names().rebuilding_ruleset_name.clone(); - if fdecl.subtype == FunctionSubtype::Constructor && i == types.len() - 1 { - // Updating the representative term (last column of constructor): - // use transitivity with sym of the UF proof - let new_proof = self.fresh_var(); - proof_code_parts.push(format!( - "(let {new_proof} - ({eq_trans_constructor} - ({sym_constructor} {uf_prf}) - {current_proof}))", - )); - current_proof = new_proof; - } else { - // Updating a child via congruence - let new_proof = self.fresh_var(); - proof_code_parts.push(format!( - "(let {new_proof} - ({congr_constructor} {current_proof} {i} - {uf_prf}))", - )); - current_proof = new_proof; - } + let mut rules = String::new(); + for (i, ty) in types.iter().enumerate() { + if !ty.is_eq_sort() { + continue; } + let ci = child(i); + let leader = format!("c{i}_leader_"); + let uf_name = self.uf_name(ty.name()); + let (query_view, _pf) = self.query_view_and_get_proof(&fdecl.name, &children_vec); + let mut updated = children_vec.clone(); + updated[i] = leader.clone(); + let updated_view = self.update_view(&fdecl.name, &updated, "()"); + let fresh_name = self.egraph.parser.symbol_gen.fresh("rebuild_rule"); + rules.push_str(&format!( + "(rule ({query_view} + (= {leader} ({uf_name} {ci})) + (!= {ci} {leader})) + ( + {updated_view} + (delete ({view_name} {children})) + ) + :ruleset {rebuilding_ruleset} :name \"{fresh_name}\" :internal-include-subsumed)\n" + )); + } + self.parse_program(&rules) + } - (proof_code_parts.join("\n"), current_proof) - } else { - ("".to_string(), "()".to_string()) - }; - - let updated_view = self.update_view(&fdecl.name, &children_updated, &pf_var); + /// Rebuild rules for a proof-mode functional-dependency constructor view + /// `(children) -> (eclass, proof)`, fanned out per eq-sort column (like the term + /// path): the single-key `@UF` has no row for a canonical column, so a combined + /// all-columns join would stall when any column is already canonical. + /// + /// - A CHILD update (`ci -> leader`, `uf_pf : ci = leader`) re-keys the row: `set` + /// at the canonicalized children (triggering the congruence `:merge` on collision) + /// and `delete` the stale row, with proof `Congr(view_prf, i, uf_pf)`. + /// - The ECLASS value update (`eclass -> leader`, `uf_pf : eclass = leader`) keeps + /// the same key, so it just `set`s the smaller eclass (the view `:merge` keeps the + /// min) with proof `Trans(Sym uf_pf, view_prf) : leader = f(children)` — no delete. + fn rebuilding_rules_fd(&mut self, fdecl: &ResolvedFunctionDecl) -> Vec { + let types = fdecl.resolved_schema.view_types(); + // The last column is the eclass (output); the rest are children (keys). + let n = types.len(); + let child = |i: usize| format!("c{i}_"); + let children_vars: Vec = (0..n - 1).map(child).collect(); + let child_types = &types[..n - 1]; + let eclass_type = &types[n - 1]; - // Make a single rule that updates the view when any child's leader differs. - let rule = format!( + let view_name = self.view_name(&fdecl.name); + let children_str = ListDisplay(&children_vars, " "); + let rebuilding_ruleset = self.proof_names().rebuilding_ruleset_name.clone(); + let trans = self.proof_names().eq_trans_constructor.clone(); + let congr = self.proof_names().congr_constructor.clone(); + let sym = self.proof_names().eq_sym_constructor.clone(); + + let mut rules = String::new(); + // One rule per eq-sort child (re-keys the row via delete + set). + for (i, ty) in child_types.iter().enumerate() { + if !ty.is_eq_sort() { + continue; + } + let ci = child(i); + let leader = format!("c{i}_leader_"); + let uf_name = self.uf_name(ty.name()); + let (query_view, eclass_var, view_prf) = + self.query_fd_view(&fdecl.name, &children_vars); + let uf_prf = self.fresh_var(); + let new_pf = self.fresh_var(); + let mut updated = children_vars.clone(); + updated[i] = leader.clone(); + let updated_view = self.update_fd_view(&fdecl.name, &updated, &eclass_var, &new_pf); + let fresh_name = self.egraph.parser.symbol_gen.fresh("rebuild_rule"); + rules.push_str(&format!( + "(rule ({query_view} + (= (values {leader} {uf_prf}) ({uf_name} {ci})) + (!= {ci} {leader})) + ( + (let {new_pf} ({congr} {view_prf} {i} {uf_prf})) + {updated_view} + (delete ({view_name} {children_str})) + ) + :ruleset {rebuilding_ruleset} :name \"{fresh_name}\" :internal-include-subsumed)\n" + )); + } + // One rule for the eclass value (same key, `set` only; the view merge keeps the min). + let eclass_uf_name = self.uf_name(eclass_type.name()); + let (query_view, eclass_var, view_prf) = self.query_fd_view(&fdecl.name, &children_vars); + let eclass_leader = self.fresh_var(); + let uf_prf = self.fresh_var(); + let new_pf = self.fresh_var(); + let updated_view = + self.update_fd_view(&fdecl.name, &children_vars, &eclass_leader, &new_pf); + let fresh_name = self.egraph.parser.symbol_gen.fresh("rebuild_rule"); + rules.push_str(&format!( "(rule ({query_view} - {uf_query_str} - {filter_query} - ) + (= (values {eclass_leader} {uf_prf}) ({eclass_uf_name} {eclass_var})) + (!= {eclass_var} {eclass_leader})) ( - {pf_code} + (let {new_pf} ({trans} ({sym} {uf_prf}) {view_prf})) {updated_view} - (delete ({view_name} {children})) ) - :ruleset {} :name \"{fresh_name}\" :internal-include-subsumed)", - self.proof_names().rebuilding_ruleset_name - ); - self.parse_program(&rule) + :ruleset {rebuilding_ruleset} :name \"{fresh_name}\" :internal-include-subsumed)\n" + )); + self.parse_program(&rules) } - /// Rules that update the to_subsume tables when children change. - /// copied from above and changed to remove last param since we dont deal with output value in to subsumed rows, removed proof flags since we dont need proofs for this, and changed from function returning unit to constructor for to subsume + /// Rules that update the to_subsume tables when children change. One rule per + /// eq-sort child (no proof needed for subsumed rows). fn rebuilding_subsumed_rules(&mut self, fdecl: &ResolvedFunctionDecl) -> Vec { let ResolvedCall::Func(FuncType { input, .. }) = &fdecl.resolved_schema else { panic!("cannot create subsumed rules for primitives") @@ -659,70 +755,59 @@ impl<'a> ProofInstrumentor<'a> { return vec![]; } + self.rebuilding_subsumed_rules_fanout(fdecl, input.clone()) + } + + /// Subsumed-table rebuild: one rule per eq-sort column, mirroring + /// [`Self::rebuilding_rules_fanout`] (the single-key `@UF` has no row for a + /// canonical node, so a per-column lookup only fires when there is work). In proof + /// mode `@UF` is a `(S) -> (S, Proof)` tuple; its proof is unused for subsumed rows. + fn rebuilding_subsumed_rules_fanout( + &mut self, + fdecl: &ResolvedFunctionDecl, + input: Vec, + ) -> Vec { let subsumed_name = self.subsumed_name(&fdecl.name); let child = |i: usize| format!("c{i}_"); let children_vec: Vec = (0..input.len()).map(child).collect(); let children = format!("{}", ListDisplay(&children_vec, " ")); + let rebuilding_ruleset = self.proof_names().rebuilding_ruleset_name.clone(); + let proofs_enabled = self.egraph.proof_state.proofs_enabled; - // For each eq-sort column, look up its leader via the UF table. - // For non-eq-sort columns, the leader is the same as the original. - let mut uf_queries = vec![]; - let mut leader_vars: Vec = vec![]; - let mut bool_neq_exprs = vec![]; - + let mut rules = String::new(); for (i, ty) in input.iter().enumerate() { - if ty.is_eq_sort() { - let leader_var = format!("c{i}_leader_"); - let uf_function_name = self.uf_function_name(ty.name()); - let ci = child(i); - - if self.egraph.proof_state.proofs_enabled { - // UF function index returns a Pair(leader, proof); one lookup gives both - let pair_var = self.fresh_var(); - uf_queries.push(format!( - "(= {pair_var} ({uf_function_name} {ci})) - (= {leader_var} (pair-first {pair_var}))" - )); - } else { - uf_queries.push(format!("(= {leader_var} ({uf_function_name} {ci}))")); - } - - bool_neq_exprs.push(format!("(bool-!= {ci} {leader_var})")); - leader_vars.push(leader_var); - } else { - leader_vars.push(child(i)); + if !ty.is_eq_sort() { + continue; } + let ci = child(i); + let leader = format!("c{i}_leader_"); + let uf_name = self.uf_name(ty.name()); + let uf_lookup = if proofs_enabled { + let proof_var = self.fresh_var(); + format!("(= (values {leader} {proof_var}) ({uf_name} {ci}))") + } else { + format!("(= {leader} ({uf_name} {ci}))") + }; + let mut updated = children_vec.clone(); + updated[i] = leader.clone(); + let updated_view = ListDisplay(&updated, " "); + let fresh_name = self + .egraph + .parser + .symbol_gen + .fresh("rebuild_to_subsume_rule"); + rules.push_str(&format!( + "(rule (({subsumed_name} {children}) + {uf_lookup} + (!= {ci} {leader})) + ( + ({subsumed_name} {updated_view}) + (delete ({subsumed_name} {children})) + ) + :ruleset {rebuilding_ruleset} :name \"{fresh_name}\" :internal-include-subsumed)\n" + )); } - - let uf_query_str = uf_queries.join("\n "); - let or_expr = format!("(or {})", bool_neq_exprs.join("\n ")); - let filter_query = format!("(guard {or_expr})"); - - // Build the updated children: use leader_var for eq-sort columns, original for others. - let children_updated: Vec = leader_vars.clone(); - - let fresh_name = self - .egraph - .parser - .symbol_gen - .fresh("rebuild_to_subsume_rule"); - - let updated_children_view = ListDisplay(children_updated, " "); - - // Make a single rule that updates the view when any child's leader differs. - let rule = format!( - "(rule (({subsumed_name} {children}) - {uf_query_str} - {filter_query} - ) - ( - ({subsumed_name} {updated_children_view}) - (delete ({subsumed_name} {children})) - ) - :ruleset {} :name \"{fresh_name}\" :internal-include-subsumed)", - self.proof_names().rebuilding_ruleset_name - ); - self.parse_program(&rule) + self.parse_program(&rules) } /// Instrument fact replaces terms with looking up @@ -883,11 +968,19 @@ impl<'a> ProofInstrumentor<'a> { let args_str = ListDisplay(new_args, " "); let proof = { - // View is always a function; query it and bind the output let view_proof_var = self.fresh_var(); - res.push(format!( - "(= {view_proof_var} ({view_name} {args_str} {fv}))" - )); + // In proof mode a constructor view is the FD tuple + // `(children) -> (eclass, proof)`: bind the eclass (`fv`) and proof from + // the tuple. Without proofs it's the `(children eclass) -> Unit` form. + if self.egraph.proof_state.proofs_enabled { + res.push(format!( + "(= (values {fv} {view_proof_var}) ({view_name} {args_str}))" + )); + } else { + res.push(format!( + "(= {view_proof_var} ({view_name} {args_str} {fv}))" + )); + } if self.proofs_enabled() { let mut proof = view_proof_var; for (i, arg_proof) in arg_proofs.into_iter().enumerate() { @@ -938,6 +1031,9 @@ impl<'a> ProofInstrumentor<'a> { (fv.clone(), proof) } + ResolvedCall::Values(_) => { + panic!("tuple-output (`values`) functions are not supported in proofs") + } } } } @@ -1046,6 +1142,23 @@ impl<'a> ProofInstrumentor<'a> { view_update } + /// Write a row into the proof-mode functional-dependency view + /// `(set (@FView children) (values eclass proof))`. Re-setting an existing `children` key with a + /// different `eclass` triggers the view's native congruence `:merge`. + fn update_fd_view( + &mut self, + fname: &str, + children: &[String], + eclass: &str, + proof: &str, + ) -> String { + let view_name = self.view_name(fname); + format!( + "(set ({view_name} {}) (values {eclass} {proof}))", + ListDisplay(children, " ") + ) + } + /// Return some code adding to the view and term tables. /// For constructors, `args` should not include the eclass of the resulting term (since it may not exist yet). /// For custom functions, `args` should include all arguments (including the output for the function). @@ -1100,7 +1213,7 @@ impl<'a> ProofInstrumentor<'a> { let proof_var = self.fresh_var(); // add a proof for the constructor if needed let term_proof = if func_type.subtype == FunctionSubtype::Constructor { - let term_proof_constructor = self.term_proof_name(func_type.output.name()); + let term_proof_constructor = self.term_proof_name(func_type.output().name()); format!("(set ({term_proof_constructor} {fv}) {proof_var})") } else { "".to_string() @@ -1118,18 +1231,18 @@ impl<'a> ProofInstrumentor<'a> { }; res.push(proof_str); - res.push(self.update_view(&func_type.name, &args_with_fv, &view_proof_var)); - - // add to uf table to initialize eclass for constructors - if func_type.subtype == FunctionSubtype::Constructor { - res.push(self.union( - func_type.output.name(), - &fv, - &fv, - &Justification::Proof(view_proof_var), - )); + if func_type.subtype == FunctionSubtype::Constructor + && self.egraph.proof_state.proofs_enabled + { + // FD view: children are the key, the fresh term is the eclass value. + res.push(self.update_fd_view(&func_type.name, args, &fv, &view_proof_var)); + } else { + res.push(self.update_view(&func_type.name, &args_with_fv, &view_proof_var)); } + // No self-loop seed: the single-key `@UF` needs none — a root term simply has + // no `@UF` row (identity-on-miss), in both term and proof mode. + (res, fv) } @@ -1142,6 +1255,20 @@ impl<'a> ProofInstrumentor<'a> { (query, pf_var) } + /// Query the proof-mode functional-dependency view by its `children` key, binding fresh + /// variables for the `eclass` and `proof` output columns: `(= (values e pf) (@FView children))`. + /// Returns `(query, eclass_var, proof_var)`. + fn query_fd_view(&mut self, fname: &str, children: &[String]) -> (String, String, String) { + let view_name = self.view_name(fname); + let eclass_var = self.fresh_var(); + let pf_var = self.fresh_var(); + let query = format!( + "(= (values {eclass_var} {pf_var}) ({view_name} {}))", + ListDisplay(children, " ") + ); + (query, eclass_var, pf_var) + } + // Add to view and term tables, returning a variable for the created term. fn instrument_action_expr( &mut self, @@ -1184,6 +1311,9 @@ impl<'a> ProofInstrumentor<'a> { )); fv } + ResolvedCall::Values(_) => { + panic!("tuple-output (`values`) functions are not supported in proofs") + } } } } @@ -1261,18 +1391,17 @@ impl<'a> ProofInstrumentor<'a> { /// Any schedule should be sound as long as we saturate. fn rebuild(&mut self) -> Schedule { let path_compress_ruleset = self.proof_names().path_compress_ruleset_name.clone(); - let single_parent = self.proof_names().single_parent_ruleset_name.clone(); - let uf_function_index = self.proof_names().uf_function_index_ruleset_name.clone(); let rebuilding_cleanup_ruleset = self.proof_names().rebuilding_cleanup_ruleset_name.clone(); let rebuilding_ruleset = self.proof_names().rebuilding_ruleset_name.clone(); let delete_ruleset = self.proof_names().delete_subsume_ruleset_name.clone(); + // Both term and proof modes use a single self-referential `@UF` whose native + // `:merge` subsumes `single_parent` and makes `uf_function_index` dead, so only + // `path_compress` (flattening cross-key chains) remains as UF maintenance. self.parse_schedule(format!( "(seq (saturate {rebuilding_cleanup_ruleset} - (saturate {single_parent}) (saturate {path_compress_ruleset}) - (saturate {uf_function_index}) {rebuilding_ruleset}) {delete_ruleset})" )) @@ -1343,6 +1472,7 @@ impl<'a> ProofInstrumentor<'a> { presort_and_args: presort_and_args.clone(), uf: Some(uf_name), proof_func, + proof_ctors: None, unionable: *unionable, }); res.extend(self.declare_sort(name)); diff --git a/egglog/src/proofs/proof_encoding_helpers.rs b/egglog/src/proofs/proof_encoding_helpers.rs index 778476c..2eecc8c 100644 --- a/egglog/src/proofs/proof_encoding_helpers.rs +++ b/egglog/src/proofs/proof_encoding_helpers.rs @@ -31,8 +31,6 @@ pub(crate) struct EncodingNames { /// For a given function symbol, the name of the function that converts to the AST type. pub(crate) sort_to_ast_constructor: HashMap, pub(crate) fn_to_term_sort: HashMap, - pub(crate) single_parent_ruleset_name: String, - pub(crate) uf_function_index_ruleset_name: String, pub(crate) pcons: String, pub(crate) pnil: String, // Ruleset names @@ -71,8 +69,6 @@ impl EncodingNames { congr_constructor: symbol_gen.fresh("Congr"), sort_to_ast_constructor: HashMap::default(), fn_to_term_sort: HashMap::default(), - single_parent_ruleset_name: symbol_gen.fresh("single_parent"), - uf_function_index_ruleset_name: symbol_gen.fresh("uf_function_index"), pcons: symbol_gen.fresh("PCons"), pnil: symbol_gen.fresh("PNil"), path_compress_ruleset_name: symbol_gen.fresh("parent"), @@ -101,28 +97,6 @@ impl ProofInstrumentor<'_> { } } - pub(crate) fn uf_function_name(&mut self, sort: &str) -> String { - if let Some(name) = self.egraph.proof_state.uf_function.get(sort) { - name.clone() - } else { - let fresh_name = self.egraph.parser.symbol_gen.fresh(&format!("UF_{sort}f")); - self.egraph - .proof_state - .uf_function - .insert(sort.to_string(), fresh_name.clone()); - fresh_name - } - } - - /// Returns the name of the Pair sort used to bundle (leader, proof) in the UF function index. - /// Only used in proof mode. - pub(crate) fn uf_pair_sort_name(&mut self, sort: &str) -> String { - self.egraph - .parser - .symbol_gen - .fresh(&format!("UFPair_{sort}")) - } - pub(crate) fn parse_program(&mut self, input: &str) -> Vec { self.egraph.parser.ensure_no_reserved_symbols = false; let res = self.egraph.parser.get_program_from_string(None, input); @@ -146,14 +120,10 @@ impl ProofInstrumentor<'_> { pub(crate) fn term_header(&mut self) -> Vec { let str = format!( "(ruleset {}) - (ruleset {}) - (ruleset {}) (ruleset {}) (ruleset {}) (ruleset {})", self.proof_names().path_compress_ruleset_name, - self.proof_names().single_parent_ruleset_name, - self.proof_names().uf_function_index_ruleset_name, self.proof_names().rebuilding_ruleset_name, self.proof_names().rebuilding_cleanup_ruleset_name, self.proof_names().delete_subsume_ruleset_name @@ -373,7 +343,7 @@ impl ProofInstrumentor<'_> { " (sort {proof_list_sort}) (sort {ast_sort}) ;; wrap sorts in this for proofs -(sort {proof_datatype}) +(sort {proof_datatype} :internal-proof-names {congr_constructor} {eq_trans_constructor} {eq_sym_constructor}) (constructor {pcons} ({proof_datatype} {proof_list_sort}) {proof_list_sort} :internal-hidden) (constructor {pnil} () {proof_list_sort} :internal-hidden) @@ -465,6 +435,12 @@ pub enum ProofEncodingUnsupportedReason { "rule uses `:naive` with an eq-sort primitive in the body. Proof encoding can only look up proofs for primitive eq-sort fact results under seminaive-safe query evaluation." )] NaiveEqSortPrimitiveFact, + #[error("tuple-output functions are not supported by the term/proof encoding.")] + TupleOutputFunction, + #[error( + "a `:merge` action block (actions before the result value) is not supported by the term/proof encoding." + )] + MergeActionBlock, } /// Checks whether a desugared program supports proof encoding. @@ -542,6 +518,24 @@ pub(crate) fn command_supports_proof_encoding( return Err(ProofEncodingUnsupportedReason::NaiveEqSortPrimitiveFact); } + // Tuple-output functions store multiple value columns, which the term/proof encoding (built + // around single-output constructor views) does not model. + if let crate::ast::GenericCommand::Function { schema, .. } = command + && schema.is_tuple_output() + { + return Err(ProofEncodingUnsupportedReason::TupleOutputFunction); + } + + // A `:merge` action block runs actions before its result; the proof encoding only instruments + // the merged value, so mark it unsupported rather than emit silently-incomplete proofs. + if let crate::ast::GenericCommand::Function { + merge: Some(merge), .. + } = command + && !merge.actions.is_empty() + { + return Err(ProofEncodingUnsupportedReason::MergeActionBlock); + } + // Check all expressions for primitives without validators let mut all_primitives_have_validators = true; command.clone().visit_exprs(&mut |expr| { diff --git a/egglog/src/proofs/proof_extraction.rs b/egglog/src/proofs/proof_extraction.rs index c05c9c2..6958a38 100644 --- a/egglog/src/proofs/proof_extraction.rs +++ b/egglog/src/proofs/proof_extraction.rs @@ -30,6 +30,9 @@ impl ProofInstrumentor<'_> { ResolvedCall::Primitive(_) => { return Err(ProveExistsError::PrimitivesUnsupported); } + ResolvedCall::Values(_) => { + return Err(ProveExistsError::RequiresConstructor); + } }; let function = self @@ -39,7 +42,7 @@ impl ProofInstrumentor<'_> { .unwrap_or_else(|| panic!("constructor {} is not declared", func.name)); let backend_id = function.backend_id; - let output_sort = function.schema.output.clone(); + let output_sort = function.schema.output().clone(); let mut termdag = TermDag::default(); let mut witness_value = None; @@ -81,7 +84,7 @@ impl ProofInstrumentor<'_> { ) }) .schema - .output + .output() .clone(); let proof_value = self .egraph diff --git a/egglog/src/proofs/proof_normal_form.rs b/egglog/src/proofs/proof_normal_form.rs index 7955468..500b8d0 100644 --- a/egglog/src/proofs/proof_normal_form.rs +++ b/egglog/src/proofs/proof_normal_form.rs @@ -81,7 +81,7 @@ fn proof_form_expr( ref span, ref head @ ResolvedCall::Func(FuncType { subtype: FunctionSubtype::Custom, - ref output, + ref outputs, .. }), ref args, @@ -95,7 +95,7 @@ fn proof_form_expr( span.clone(), ResolvedVar { name: fresh.fresh("n"), - sort: output.clone(), + sort: outputs[0].clone(), is_global_ref: false, }, ); @@ -126,7 +126,7 @@ fn proof_form_expr( // (but allow other primitives to stay inline) ref arg_expr @ ResolvedExpr::Call( ref arg_span, - ResolvedCall::Func(FuncType { ref output, .. }), + ResolvedCall::Func(FuncType { ref outputs, .. }), ref inner_args, ) => { // First recursively normalize the inner arguments @@ -140,7 +140,7 @@ fn proof_form_expr( arg_span.clone(), ResolvedVar { name: fresh.fresh("v"), - sort: output.clone(), + sort: outputs[0].clone(), is_global_ref: false, }, ); diff --git a/egglog/src/proofs/snapshots/egglog__proofs__proof_tests__tests__doc_example_add_function1.snap b/egglog/src/proofs/snapshots/egglog__proofs__proof_tests__tests__doc_example_add_function1.snap index de728f7..a8752e8 100644 --- a/egglog/src/proofs/snapshots/egglog__proofs__proof_tests__tests__doc_example_add_function1.snap +++ b/egglog/src/proofs/snapshots/egglog__proofs__proof_tests__tests__doc_example_add_function1.snap @@ -3,30 +3,16 @@ source: src/proofs/proof_tests.rs expression: snapshot --- (ruleset __parent) -(ruleset __single_parent) -(ruleset __uf_function_index) (ruleset __rebuilding) (ruleset __rebuilding_cleanup) (ruleset __delete_subsume_ruleset) (sort Math :internal-uf __UF_Math) -(function __UF_Math (Math Math) Unit :merge old :unextractable :internal-hidden) -(function __UF_Mathf (Math) Math :merge new :unextractable :internal-hidden) -(rule ((__UF_Math a b) - (__UF_Math b c) +(function __UF_Math (Math) Math :merge (ordering-min old new) :unextractable :internal-hidden) +(rule ((= b (__UF_Math a)) + (= c (__UF_Math b)) (!= b c)) - ((delete (__UF_Math a b)) - (set (__UF_Math a c) ())) - :ruleset __parent :name "__uf_update") -(rule ((__UF_Math a b) - (__UF_Math a c) - (!= b c) - (= (ordering-max b c) b)) - ((delete (__UF_Math a b)) - (set (__UF_Math b c) ())) - :ruleset __single_parent :name "singleparent__uf_update") -(rule ((__UF_Math a b)) - ((set (__UF_Mathf a) b)) - :ruleset __uf_function_index :name "__uf_function_index1") + ((set (__UF_Math a) c)) + :ruleset __parent :name "__uf_path_compress") (sort __view) (constructor Add (i64 i64) Math :unextractable :internal-hidden) (function __AddView (i64 i64 Math) Unit :merge old :internal-term-constructor Add) @@ -36,7 +22,7 @@ expression: snapshot (= __v1 (__AddView c0_ c1_ old)) (!= old new) (= (ordering-max old new) new)) - ((set (__UF_Math (ordering-max new old) (ordering-min new old)) ())) + ((set (__UF_Math (ordering-max new old)) (ordering-min new old))) :ruleset __rebuilding :name "__congruence_rule" :internal-include-subsumed) (rule ((__to_delete_Add c0_ c1_) (__AddView c0_ c1_ out)) @@ -48,26 +34,23 @@ expression: snapshot ((subsume (__AddView c0_ c1_ out))) :ruleset __delete_subsume_ruleset :name "__delete_rule_subsume") (rule ((= __v2 (__AddView c0_ c1_ c2_)) - (= c2_leader_ (__UF_Mathf c2_)) - (guard (or (bool-!= c2_ c2_leader_)))) + (= c2_leader_ (__UF_Math c2_)) + (!= c2_ c2_leader_)) ((set (__AddView c0_ c1_ c2_leader_) ()) (delete (__AddView c0_ c1_ c2_))) :ruleset __rebuilding :name "__rebuild_rule" :internal-include-subsumed) (function __v3 () Math :no-merge :unextractable :internal-let) (set (__v3) (Add 1 2)) (set (__AddView 1 2 (__v3)) ()) -(set (__UF_Math (ordering-max (__v3) (__v3)) (ordering-min (__v3) (__v3))) ()) -(run-schedule (seq (saturate (seq (run __rebuilding_cleanup) (saturate (run __single_parent)) (saturate (run __parent)) (saturate (run __uf_function_index)) (run __rebuilding))) (run __delete_subsume_ruleset))) +(run-schedule (seq (saturate (seq (run __rebuilding_cleanup) (saturate (run __parent)) (run __rebuilding))) (run __delete_subsume_ruleset))) (rule ((= __v5 (__AddView a b __v4))) ((let __v7 (Add a b)) (set (__AddView a b __v7) ()) - (set (__UF_Math (ordering-max __v7 __v7) (ordering-min __v7 __v7)) ()) (let __v8 (Add b a)) (set (__AddView b a __v8) ()) - (set (__UF_Math (ordering-max __v8 __v8) (ordering-min __v8 __v8)) ()) - (set (__UF_Math (ordering-max __v7 __v8) (ordering-min __v7 __v8)) ())) + (set (__UF_Math (ordering-max __v7 __v8)) (ordering-min __v7 __v8))) :name "commutativity") (check (= __v10 (__AddView 1 2 __v9)) (= __v12 (__AddView 2 1 __v11)) (= __v9 __v11)) -(run-schedule (seq (saturate (seq (run __rebuilding_cleanup) (saturate (run __single_parent)) (saturate (run __parent)) (saturate (run __uf_function_index)) (run __rebuilding))) (run __delete_subsume_ruleset))) +(run-schedule (seq (saturate (seq (run __rebuilding_cleanup) (saturate (run __parent)) (run __rebuilding))) (run __delete_subsume_ruleset))) diff --git a/egglog/src/proofs/snapshots/egglog__proofs__proof_tests__tests__doc_example_add_function2.snap b/egglog/src/proofs/snapshots/egglog__proofs__proof_tests__tests__doc_example_add_function2.snap index e664c5f..4591a0e 100644 --- a/egglog/src/proofs/snapshots/egglog__proofs__proof_tests__tests__doc_example_add_function2.snap +++ b/egglog/src/proofs/snapshots/egglog__proofs__proof_tests__tests__doc_example_add_function2.snap @@ -3,8 +3,6 @@ source: src/proofs/proof_tests.rs expression: snapshot --- (ruleset __parent) -(ruleset __single_parent) -(ruleset __uf_function_index) (ruleset __rebuilding) (ruleset __rebuilding_cleanup) (ruleset __delete_subsume_ruleset) @@ -48,4 +46,4 @@ expression: snapshot :ruleset __delete_subsume_ruleset :name "__delete_rule_subsume") (check (= __v (__addView 0 0 __n)) (= __n 0)) -(run-schedule (seq (saturate (seq (run __rebuilding_cleanup) (saturate (run __single_parent)) (saturate (run __parent)) (saturate (run __uf_function_index)) (run __rebuilding))) (run __delete_subsume_ruleset))) +(run-schedule (seq (saturate (seq (run __rebuilding_cleanup) (saturate (run __parent)) (run __rebuilding))) (run __delete_subsume_ruleset))) diff --git a/egglog/src/scheduler.rs b/egglog/src/scheduler.rs index eb5436f..6691d0c 100644 --- a/egglog/src/scheduler.rs +++ b/egglog/src/scheduler.rs @@ -338,6 +338,8 @@ impl SchedulerRuleInfo { .collect(); let decided = egraph.backend.add_table(FunctionConfig { schema, + n_vals: 1, + n_identity_vals: None, default: DefaultVal::Const(unit), merge: MergeFn::AssertEq, name: "backend".to_string(), diff --git a/egglog/src/serialize.rs b/egglog/src/serialize.rs index 83ab829..e185566 100644 --- a/egglog/src/serialize.rs +++ b/egglog/src/serialize.rs @@ -148,8 +148,14 @@ impl EGraph { truncated_functions.push(name.clone()); return false; } - let (out, inps) = row.vals.split_last().unwrap(); - let class_id = self.value_to_class_id(&function.schema.output, *out); + // Split the row into key columns and value columns by the function's input arity. + // Tuple-output functions (e.g. the proof-mode `(children) -> (eclass, proof)` view) + // have more than one value column; the first value column is the e-class / output + // that serialization tracks, and any extra columns (a proof) are not serialized. + let n_keys = function.schema.input.len(); + let inps = &row.vals[..n_keys]; + let out = &row.vals[n_keys]; + let class_id = self.value_to_class_id(function.schema.output(), *out); if function.decl.internal_let { let_bindings .entry(class_id.clone()) @@ -186,7 +192,7 @@ impl EGraph { let node_ids: NodeIDs = all_calls.iter().fold( HashMap::default(), |mut acc, (func, _input, _output, _subsumed, class_id, node_id)| { - if func.schema.output.is_eq_sort() { + if func.schema.output().is_eq_sort() { acc.entry(class_id.clone()) .or_default() .push_back(node_id.clone()); @@ -202,7 +208,7 @@ impl EGraph { }; for (func, input, output, subsumed, class_id, node_id) in all_calls { - self.serialize_value(&mut serializer, &func.schema.output, output, &class_id); + self.serialize_value(&mut serializer, func.schema.output(), output, &class_id); assert_eq!(input.len(), func.schema.input.len()); let children: Vec<_> = input diff --git a/egglog/src/sort/fn.rs b/egglog/src/sort/fn.rs index e7f3f8e..20271c7 100644 --- a/egglog/src/sort/fn.rs +++ b/egglog/src/sort/fn.rs @@ -293,7 +293,7 @@ impl TypeConstraint for FunctionCTorTypeConstraint { // the output type and input types (starting after the partial args) must match between these functions let expected_output = self.function.output.clone(); let expected_input = self.function.inputs.clone(); - let actual_output = func_type.output.clone(); + let actual_output = func_type.output().clone(); let actual_input: Vec = func_type .input .iter() diff --git a/egglog/src/typechecking.rs b/egglog/src/typechecking.rs index 53cf2b4..bc8c933 100644 --- a/egglog/src/typechecking.rs +++ b/egglog/src/typechecking.rs @@ -115,14 +115,38 @@ pub struct FuncType { pub name: String, pub subtype: FunctionSubtype, pub input: Vec, - pub output: ArcSort, + /// The output (value-column) sorts, primary first. A tuple-output function has more than one; + /// ordinary functions have exactly one. Always non-empty. + pub outputs: Vec, +} + +impl FuncType { + /// The primary (first) output sort. + pub fn output(&self) -> &ArcSort { + &self.outputs[0] + } + + /// The number of output (value) columns. + pub fn num_outputs(&self) -> usize { + self.outputs.len() + } + + /// Whether this function has more than one output column. + pub fn is_tuple_output(&self) -> bool { + self.outputs.len() > 1 + } } impl PartialEq for FuncType { fn eq(&self, other: &Self) -> bool { if self.name == other.name && self.subtype == other.subtype - && self.output.name() == other.output.name() + && self.num_outputs() == other.num_outputs() + && self + .outputs + .iter() + .zip(other.outputs.iter()) + .all(|(a, b)| a.name() == b.name()) { if self.input.len() != other.input.len() { return false; @@ -145,7 +169,9 @@ impl Hash for FuncType { fn hash(&self, state: &mut H) { self.name.hash(state); self.subtype.hash(state); - self.output.name().hash(state); + for out in &self.outputs { + out.name().hash(state); + } for inp in &self.input { inp.name().hash(state); } @@ -397,7 +423,7 @@ impl EGraph { // If this is a let binding, add it to global_sorts // This preserves bahavior for lets after desugaring if resolved.internal_let { - let output_sort = self.type_info.sorts.get(&fdecl.schema.output).unwrap(); + let output_sort = self.type_info.sorts.get(fdecl.schema.output()).unwrap(); self.type_info .global_sorts .insert(fdecl.name.clone(), output_sort.clone()); @@ -415,6 +441,7 @@ impl EGraph { presort_and_args, uf, proof_func, + proof_ctors, unionable, } => { // Note this is bad since typechecking should be pure and idempotent @@ -430,6 +457,7 @@ impl EGraph { presort_and_args: presort_and_args.clone(), uf: uf.clone(), proof_func: proof_func.clone(), + proof_ctors: proof_ctors.clone(), unionable: *unionable, } } @@ -458,6 +486,19 @@ impl EGraph { )?) } NCommand::Extract(span, expr, variants) => { + // A tuple-output function returns more than one value, so it can't be extracted as a + // single term; surface a clear error instead of a confusing arity mismatch. + if let GenericExpr::Call(_, head, _) = expr + && self + .type_info + .get_func_type(head) + .is_some_and(|t| t.is_tuple_output()) + { + return Err(TypeError::CannotExtractTupleOutput( + head.clone(), + span.clone(), + )); + } let res_expr = self.type_info.typecheck_standalone_expr( symbol_gen, expr, @@ -688,32 +729,30 @@ impl TypeInfo { } fn function_to_functype(&self, func: &FunctionDecl) -> Result { + let resolve = |name: &String| -> Result { + self.sorts + .get(name) + .cloned() + .ok_or_else(|| TypeError::UndefinedSort(name.clone(), func.span.clone())) + }; let input = func .schema .input .iter() - .map(|name| { - if let Some(sort) = self.sorts.get(name) { - Ok(sort.clone()) - } else { - Err(TypeError::UndefinedSort(name.clone(), func.span.clone())) - } - }) + .map(&resolve) + .collect::, _>>()?; + let outputs = func + .schema + .outputs + .iter() + .map(&resolve) .collect::, _>>()?; - let output = if let Some(sort) = self.sorts.get(&func.schema.output) { - Ok(sort.clone()) - } else { - Err(TypeError::UndefinedSort( - func.schema.output.clone(), - func.span.clone(), - )) - }?; Ok(FuncType { name: func.name.clone(), subtype: func.subtype, input, - output: output.clone(), + outputs, }) } @@ -734,8 +773,13 @@ impl TypeInfo { fdecl.span.clone(), )); } - // View tables (with term_constructor) must have at least one input (the e-class) - if fdecl.term_constructor.is_some() && fdecl.schema.input.is_empty() { + // View tables (with term_constructor) must have at least one input (the e-class), except the + // proof-mode functional-dependency tuple view `(children) -> (eclass, proof)`, which keys on + // children only (a 0-arg constructor's view then has no inputs). + if fdecl.term_constructor.is_some() + && fdecl.schema.input.is_empty() + && !fdecl.schema.is_tuple_output() + { return Err(TypeError::TermConstructorNoInputs( fdecl.name.clone(), fdecl.span.clone(), @@ -748,34 +792,105 @@ impl TypeInfo { fdecl.span.clone(), )); } - 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() { + let outputs: Vec = fdecl + .schema + .outputs + .iter() + .map(|name| self.sorts.get(name).unwrap().clone()) + .collect(); + let is_tuple = fdecl.schema.is_tuple_output(); + + // Tuple outputs are only meaningful for custom functions (which carry a functional + // dependency from keys to a tuple of values). Constructors mint a single e-class id, so they + // may not be tuple-output. Term-constructor *views* may be tuple-output: the proof-mode + // encoder emits `(children) -> (eclass, proof)` views (an internal-only annotation, so this + // can't be reached by user input). + if is_tuple && fdecl.subtype == FunctionSubtype::Constructor { + return Err(TypeError::TupleOutputNotAllowed( + fdecl.name.clone(), + fdecl.span.clone(), + )); + } + if fdecl.subtype == FunctionSubtype::Constructor && !outputs[0].is_eq_sort() { return Err(TypeError::ConstructorOutputNotSort( fdecl.name.clone(), fdecl.span.clone(), )); } - bound_vars.insert("old", (fdecl.span.clone(), output_type.clone())); - bound_vars.insert("new", (fdecl.span.clone(), output_type.clone())); + + // For single-output functions the merge expression refers to `old`/`new`. For + // tuple-output functions it refers to `old0`, `new0`, `old1`, `new1`, ... (one pair per + // output column), and the whole merge is a `(values ...)` form. + let mut bound_vars = IndexMap::default(); + let tuple_var_names: Vec<(String, String)> = (0..outputs.len()) + .map(|i| (format!("old{i}"), format!("new{i}"))) + .collect(); + if is_tuple { + for (i, (old_name, new_name)) in tuple_var_names.iter().enumerate() { + bound_vars.insert(old_name.as_str(), (fdecl.span.clone(), outputs[i].clone())); + bound_vars.insert(new_name.as_str(), (fdecl.span.clone(), outputs[i].clone())); + } + } else { + bound_vars.insert("old", (fdecl.span.clone(), outputs[0].clone())); + bound_vars.insert("new", (fdecl.span.clone(), outputs[0].clone())); + } + + // A `:merge` is a value-producing action block: the `actions` run (writes are allowed, but + // live DB reads would be untracked by seminaive rule execution), then `result` produces the + // merged value(s). Both are typechecked with `old`/`new` (`old0`/`new0`/... for tuple) + // bound; the actions go through the same action typechecker as rule bodies. + let merge = match &fdecl.merge { + Some(merge) => { + let actions = self.typecheck_standalone_actions( + symbol_gen, + &merge.actions, + &bound_vars, + Context::Write, + )?; + // The result is evaluated after the actions, so any `let`-bound variable is in + // scope for it. Extend the binding with each `let`'s solved type before checking + // the result. `let_bindings` owns the names so `result_scope` can borrow them. + let let_bindings: Vec<(String, Span, ArcSort)> = actions + .0 + .iter() + .filter_map(|a| match a { + GenericAction::Let(span, var, _) => { + Some((var.name.as_str().to_owned(), span.clone(), var.sort.clone())) + } + _ => None, + }) + .collect(); + let mut result_scope = bound_vars.clone(); + for (name, span, sort) in &let_bindings { + result_scope.insert(name.as_str(), (span.clone(), sort.clone())); + } + let result = if is_tuple { + self.typecheck_tuple_merge( + symbol_gen, + fdecl, + &merge.result, + &outputs, + &result_scope, + )? + } else { + self.typecheck_standalone_expr( + symbol_gen, + &merge.result, + &result_scope, + Context::Write, + )? + }; + Some(ResolvedMerge { actions, result }) + } + None => None, + }; Ok(ResolvedFunctionDecl { name: fdecl.name.clone(), subtype: fdecl.subtype, schema: fdecl.schema.clone(), resolved_schema: ResolvedCall::Func(self.func_types.get(&fdecl.name).unwrap().clone()), - merge: match &fdecl.merge { - // Merge expressions run as part of action-side table updates: - // writes are allowed, but live DB reads would be untracked by - // seminaive rule execution. - Some(merge) => Some(self.typecheck_standalone_expr( - symbol_gen, - merge, - &bound_vars, - Context::Write, - )?), - None => None, - }, + merge, cost: fdecl.cost, unextractable: fdecl.unextractable, internal_hidden: fdecl.internal_hidden, @@ -785,6 +900,55 @@ impl TypeInfo { }) } + /// Typecheck the `(values e0 e1 ...)` merge of a tuple-output function. Each `ei` is checked + /// with `old0`/`new0`/... bound to the corresponding output columns, and must have the type of + /// output column `i`. The result is a resolved `values` call carrying the output sorts. + fn typecheck_tuple_merge( + &self, + symbol_gen: &mut SymbolGen, + fdecl: &FunctionDecl, + merge: &Expr, + outputs: &[ArcSort], + bound_vars: &IndexMap<&str, (Span, ArcSort)>, + ) -> Result { + let args = match merge { + GenericExpr::Call(_, head, args) if head.as_str() == "values" => args, + _ => { + return Err(TypeError::TupleMergeNotValues( + fdecl.name.clone(), + fdecl.span.clone(), + )); + } + }; + if args.len() != outputs.len() { + return Err(TypeError::TupleMergeArity { + name: fdecl.name.clone(), + expected: outputs.len(), + actual: args.len(), + span: fdecl.span.clone(), + }); + } + let mut resolved_args = Vec::with_capacity(args.len()); + for (arg, expected) in args.iter().zip(outputs) { + let resolved = + self.typecheck_standalone_expr(symbol_gen, arg, bound_vars, Context::Write)?; + let actual = resolved.output_type(); + if actual.name() != expected.name() { + return Err(TypeError::Mismatch { + expr: arg.clone(), + expected: expected.clone(), + actual, + }); + } + resolved_args.push(resolved); + } + Ok(GenericExpr::Call( + merge.span(), + ResolvedCall::Values(outputs.to_vec()), + resolved_args, + )) + } + fn typecheck_schedule( &self, symbol_gen: &mut SymbolGen, @@ -1197,6 +1361,27 @@ pub enum TypeError { crate::GLOBAL_NAME_PREFIX )] GlobalMissingPrefix { name: String, span: Span }, + #[error( + "{1}\nFunction {0} has a tuple output, which is only allowed for plain functions (not constructors, relations, or view tables)." + )] + TupleOutputNotAllowed(String, Span), + #[error( + "{1}\nThe :merge of tuple-output function {0} must be a `(values ...)` form with one expression per output column." + )] + TupleMergeNotValues(String, Span), + #[error( + "{span}\nThe :merge of tuple-output function {name} has {actual} columns but the function has {expected} output columns." + )] + TupleMergeArity { + name: String, + expected: usize, + actual: usize, + span: Span, + }, + #[error( + "{1}\nCannot extract tuple-output function {0}: extraction yields a single term, but a tuple-output function has more than one output column. Read its columns in a rule with `(= (values ...) ({0} ...))` instead." + )] + CannotExtractTupleOutput(String, Span), } #[cfg(test)] diff --git a/egglog/src/util.rs b/egglog/src/util.rs index 8c54a70..9210643 100644 --- a/egglog/src/util.rs +++ b/egglog/src/util.rs @@ -100,8 +100,9 @@ impl FreshGen for SymbolGen { } ); let sort = match name_hint { - ResolvedCall::Func(f) => f.output.clone(), + ResolvedCall::Func(f) => f.output().clone(), ResolvedCall::Primitive(prim) => prim.output().clone(), + ResolvedCall::Values(sorts) => sorts[0].clone(), }; ResolvedVar { name, diff --git a/egglog/tests/merge-action-block.egg b/egglog/tests/merge-action-block.egg new file mode 100644 index 0000000..8d6786a --- /dev/null +++ b/egglog/tests/merge-action-block.egg @@ -0,0 +1,27 @@ +; A value-producing `:merge` action block. Syntax: `:merge (* )` — the +; actions run first (with `old`/`new` bound), then the trailing expression is the merged value. +; Actions may be `let`, `set`, or `union`. +; +; Here, on a conflict `f` binds the smaller/larger values with `let`, records the discarded +; (smaller) value into a side function with `set`, and keeps the larger as the merged value. +(function discarded (i64) i64 :merge new) +(function f (i64) i64 + :merge ((let smaller (min old new)) + (let larger (max old new)) + (set (discarded smaller) 1) + larger)) + +(set (f 0) 5) +(set (f 0) 3) + +(check (= (f 0) 5)) ; the result kept the larger (`let`-bound) value +(check (= (discarded 3) 1)) ; the `set` action recorded the discarded (smaller) value +(fail (check (= (discarded 5) 1))) + +; A merge that `union`s the two conflicting outputs of an eq-sort function: on a conflict the two +; eclasses are unified and the old value is kept. +(datatype Color (Red) (Green)) +(function pick (i64) Color :merge ((union old new) old)) +(set (pick 0) (Red)) +(set (pick 0) (Green)) +(check (= (Red) (Green))) ; the `union` action unified the conflicting values diff --git a/egglog/tests/merge_action_block.rs b/egglog/tests/merge_action_block.rs new file mode 100644 index 0000000..5beecc2 --- /dev/null +++ b/egglog/tests/merge_action_block.rs @@ -0,0 +1,95 @@ +//! Integration tests for value-producing `:merge` action blocks: a `:merge` may run actions +//! (`let`, `set`, `union`) before its result expression, with `old`/`new` bound. Syntax: +//! `:merge (* )` (a bare `:merge ` is the no-actions case). + +use egglog::EGraph; + +fn run(prog: &str) -> Result<(), egglog::Error> { + EGraph::default() + .parse_and_run_program(None, prog) + .map(|_| ()) +} + +#[test] +fn merge_action_block_runs_effect_and_returns_value() { + // On a conflict `f`'s merge records the discarded (smaller) value into `discarded`, then keeps + // the larger value. Asserts both the side effect and the merged result. + run(r#" + (function discarded (i64) i64 :merge new) + (function f (i64) i64 :merge ((set (discarded (min old new)) 1) (max old new))) + (set (f 0) 5) + (set (f 0) 3) + (check (= (f 0) 5)) ; result kept the larger value + (check (= (discarded 3) 1)) ; the action recorded the discarded (smaller) value + (fail (check (= (discarded 5) 1))) + "#) + .unwrap(); +} + +#[test] +fn merge_action_block_plain_expr_still_works() { + // A bare `:merge ` (no actions) is unchanged. + run(r#" + (function g (i64) i64 :merge (max old new)) + (set (g 0) 2) + (set (g 0) 7) + (check (= (g 0) 7)) + "#) + .unwrap(); +} + +#[test] +fn merge_action_block_let_binding() { + // A `let` binds an intermediate value that later actions and the result can reference. + run(r#" + (function discarded (i64) i64 :merge new) + (function f (i64) i64 + :merge ((let smaller (min old new)) + (let larger (max old new)) + (set (discarded smaller) 1) + larger)) + (set (f 0) 5) + (set (f 0) 3) + (check (= (f 0) 5)) ; result is the `let`-bound larger value + (check (= (discarded 3) 1)) ; the `set` used the `let`-bound smaller value + "#) + .unwrap(); +} + +#[test] +fn merge_action_block_union() { + // A `union` action unifies the two conflicting eclasses of an eq-sort function. + run(r#" + (datatype Color (Red) (Green)) + (function pick (i64) Color :merge ((union old new) old)) + (set (pick 0) (Red)) + (set (pick 0) (Green)) + (check (= (Red) (Green))) + "#) + .unwrap(); +} + +#[test] +fn merge_action_block_let_in_tuple_merge() { + // A `let` is in scope for a tuple-output `(values ...)` result too. + run(r#" + (function iv (i64) (i64 i64) + :merge ((let lo (max old0 new0)) + (values lo (min old1 new1)))) + (set (iv 0) (values 1 9)) + (set (iv 0) (values 4 2)) + (check (= (values 4 2) (iv 0))) ; lo = max(1,4) = 4, hi = min(9,2) = 2 + "#) + .unwrap(); +} + +#[test] +fn merge_action_block_rejects_unsupported_action() { + // `set`/`let`/`union` are supported; other actions (here `panic`) are not meaningful during a + // merge and are a clear error. + let err = run(r#"(function g (i64) i64 :merge ((panic "boom") old))"#).unwrap_err(); + assert!( + err.to_string().contains("not supported inside a :merge"), + "expected an unsupported-action error, got: {err}" + ); +} diff --git a/egglog/tests/snapshots/files__proof_unsupported_files.snap b/egglog/tests/snapshots/files__proof_unsupported_files.snap index 190f045..9b6b0cc 100644 --- a/egglog/tests/snapshots/files__proof_unsupported_files.snap +++ b/egglog/tests/snapshots/files__proof_unsupported_files.snap @@ -1,5 +1,5 @@ --- -source: tests/files.rs +source: egglog/tests/files.rs expression: snapshot --- array.egg @@ -35,6 +35,7 @@ looking_up_nonconstructor_in_rewrite_good.egg luminal-llama.egg map.egg math.egg +merge-action-block.egg merge_read.egg multiset.egg naive-action-lookup.egg @@ -59,6 +60,7 @@ string_quotes.egg taylor51.egg towers-of-hanoi.egg tricky-type-checking.egg +tuple-output.egg type-constraints-tests.egg typed_primitive_unstable_app.egg typeinfer.egg diff --git a/egglog/tests/snapshots/files__shared_snapshot_eqsat_basic_proof.snap b/egglog/tests/snapshots/files__shared_snapshot_eqsat_basic_proof.snap index 72d52cf..926b1c7 100644 --- a/egglog/tests/snapshots/files__shared_snapshot_eqsat_basic_proof.snap +++ b/egglog/tests/snapshots/files__shared_snapshot_eqsat_basic_proof.snap @@ -5,43 +5,43 @@ expression: snapshot_content_across_treatments (let t0 (Mul (Num 2) (Add (Var "x") (Num 3)))) (let t1 (Mul (Num 2) (Var "x"))) (let t2 (Add (Num 6) t1)) -(let t3 (Mul (Num 2) (Num 3))) -(let t4 (Add t3 t1)) -(let t5 (Add t1 t3)) +(let t3 (Add t1 (Num 6))) +(let t4 (Mul (Num 2) (Num 3))) +(let t5 (Add t1 t4)) (let t6 (premises (Fiat (= t0 t0)))) (let t7 (substitution - (@rewrite_var__1 t0) - (a (Num 2)) - (b (Var "x")) - (c (Num 3)))) -(let t8 - (premises - (Sym - (= t0 t5) - (Rule - (= t5 t0) - (name "(rewrite (Mul a (Add b c)) (Add (Mul a b) (Mul a c)))") - t6 - t7)))) -(let t9 (substitution (a t1) (b t3) (@rewrite_var__ t0))) + (@rewrite_var__1 t0) + (a (Num 2)) + (b (Var "x")) + (c (Num 3)))) (Trans (= t0 t2) - (Sym (= t0 t4) (Rule (= t4 t0) (name "(rewrite (Add a b) (Add b a))") t8 t9)) (Congr - (= t4 t2) - (Rule (= t4 t4) (name "(rewrite (Add a b) (Add b a))") t8 t9) + (= t0 t3) + (Sym + (= t0 t5) + (Rule + (= t5 t0) + (name "(rewrite (Mul a (Add b c)) (Add (Mul a b) (Mul a c)))") + t6 + t7)) (Rule - (= t3 (Num 6)) + (= t4 (Num 6)) (name "(rewrite (Mul (Num a) (Num b)) (Num (* a b)))") (premises (Rule - (= t3 t3) + (= t4 t4) (name "(rewrite (Mul a (Add b c)) (Add (Mul a b) (Mul a c)))") t6 t7)) - (substitution (a 2) (b 3) (@rewrite_var__3 t3))) - 0)) + (substitution (a 2) (b 3) (@rewrite_var__3 t4))) + 1) + (Rule + (= t3 t2) + (name "(rewrite (Add a b) (Add b a))") + (premises (Fiat (= t2 t2))) + (substitution (a (Num 6)) (b t1) (@rewrite_var__ t2)))) ((Add 4) (Mul 3) (Num 3) diff --git a/egglog/tests/snapshots/files__shared_snapshot_merge_action_block.snap b/egglog/tests/snapshots/files__shared_snapshot_merge_action_block.snap new file mode 100644 index 0000000..5d4244e --- /dev/null +++ b/egglog/tests/snapshots/files__shared_snapshot_merge_action_block.snap @@ -0,0 +1,9 @@ +--- +source: egglog/tests/files.rs +expression: snapshot_content_across_treatments +--- +((Green 1) + (Red 1) + (discarded 1) + (f 1) + (pick 1)) diff --git a/egglog/tests/snapshots/files__shared_snapshot_tuple_output.snap b/egglog/tests/snapshots/files__shared_snapshot_tuple_output.snap new file mode 100644 index 0000000..9a55f89 --- /dev/null +++ b/egglog/tests/snapshots/files__shared_snapshot_tuple_output.snap @@ -0,0 +1,7 @@ +--- +source: egglog/tests/files.rs +expression: snapshot_content_across_treatments +--- +((Add 1) + (Num 2) + (interval 3)) diff --git a/egglog/tests/tuple-output.egg b/egglog/tests/tuple-output.egg new file mode 100644 index 0000000..f0f62e3 --- /dev/null +++ b/egglog/tests/tuple-output.egg @@ -0,0 +1,20 @@ +; Multi-output (tuple) function: an interval analysis whose lower bound merges with `max` and +; whose upper bound merges with `min`, expressed as a single two-output function +; `(Math) -> (i64 i64)`. The value columns are written with `(set (interval e) (values lo hi))` +; and destructured with `(= (values lo hi) (interval e))`. +(datatype Math (Num i64) (Add Math Math)) + +(function interval (Math) (i64 i64) + :merge (values (max old0 new0) (min old1 new1))) + +(rule ((= e (Num n))) + ((set (interval e) (values n n)))) +(rule ((= e (Add a b)) + (= (values alo ahi) (interval a)) + (= (values blo bhi) (interval b))) + ((set (interval e) (values (+ alo blo) (+ ahi bhi))))) + +(let expr (Add (Num 1) (Num 2))) +(run 10) +(check (= (values 1 1) (interval (Num 1)))) +(check (= (values 3 3) (interval expr))) diff --git a/egglog/tests/tuple_outputs.rs b/egglog/tests/tuple_outputs.rs new file mode 100644 index 0000000..a06b6ba --- /dev/null +++ b/egglog/tests/tuple_outputs.rs @@ -0,0 +1,315 @@ +//! Integration tests for tuple-output functions: functions declared with more than one output +//! sort, e.g. `(function f (Math) (i64 i64) :merge (values ...))`. Their value columns are +//! destructured with `(= (values a b) (f x))` and written with `(set (f x) (values a b))`. + +use egglog::{EGraph, Error, TypeError}; + +fn run(prog: &str) -> Result<(), Error> { + EGraph::default() + .parse_and_run_program(None, prog) + .map(|_| ()) +} + +#[test] +fn tuple_merge_interval_analysis() { + // The motivating example: an interval analysis whose lower bound merges with `max` and whose + // upper bound merges with `min`, expressed as a single two-output function. + run(r#" + (datatype Math (Num i64) (Add Math Math)) + (function interval (Math) (i64 i64) + :merge (values (max old0 new0) (min old1 new1))) + + (rule ((= e (Num n))) + ((set (interval e) (values n n)))) + (rule ((= e (Add a b)) + (= (values alo ahi) (interval a)) + (= (values blo bhi) (interval b))) + ((set (interval e) (values (+ alo blo) (+ ahi bhi))))) + + (let expr (Add (Num 1) (Num 2))) + (run 10) + (check (= (values 3 3) (interval expr))) + "#) + .unwrap(); +} + +#[test] +fn tuple_merge_combines_columns_independently() { + // Two `set`s on the same key are resolved column-by-column: lo with max, hi with min. + run(r#" + (datatype M (V i64)) + (function iv (M) (i64 i64) :merge (values (max old0 new0) (min old1 new1))) + (set (iv (V 0)) (values 1 10)) + (set (iv (V 0)) (values 3 7)) + (check (= (values 3 7) (iv (V 0)))) + "#) + .unwrap(); +} + +#[test] +fn tuple_check_distinguishes_columns() { + // A `(values ...)` check must match every column, not just the first. + let bad = run(r#" + (datatype M (V i64)) + (function iv (M) (i64 i64) :merge (values (max old0 new0) (min old1 new1))) + (set (iv (V 0)) (values 3 7)) + (check (= (values 3 99) (iv (V 0)))) + "#); + assert!(bad.is_err(), "check on a wrong second column should fail"); +} + +#[test] +fn tuple_destructure_binds_each_column_in_rule() { + run(r#" + (datatype M (V i64)) + (function iv (M) (i64 i64) :merge (values (max old0 new0) (min old1 new1))) + (relation lo-is (M i64)) + (relation hi-is (M i64)) + (set (iv (V 0)) (values 3 7)) + (rule ((= (values l h) (iv x))) ((lo-is x l) (hi-is x h))) + (run 1) + (check (lo-is (V 0) 3)) + (check (hi-is (V 0) 7)) + (fail (check (lo-is (V 0) 7))) + "#) + .unwrap(); +} + +#[test] +fn tuple_output_eq_sorts() { + // Output columns may be e-classes, and are kept canonical across unions like any other column. + run(r#" + (datatype N (Nn i64)) + (function two (i64) (N N) :merge (values new0 new1)) + (set (two 0) (values (Nn 1) (Nn 2))) + (relation got (N N)) + (rule ((= (values a b) (two k))) ((got a b))) + (run 1) + (check (got (Nn 1) (Nn 2))) + + ; unioning an output e-class is reflected after rebuilding + (union (Nn 1) (Nn 5)) + (run 1) + (check (= (values (Nn 5) (Nn 2)) (two 0))) + "#) + .unwrap(); +} + +#[test] +fn tuple_no_merge_asserts_equality() { + // With no `:merge`, each output column asserts equality; a conflicting write is an error. + run(r#" + (datatype M (V i64)) + (function iv (M) (i64 i64)) + (set (iv (V 0)) (values 1 2)) + (set (iv (V 0)) (values 1 3)) + "#) + .expect_err("conflicting write under default (assert-eq) merge should fail"); +} + +#[test] +fn constructor_tuple_output_rejected() { + let err = EGraph::default() + .parse_and_run_program(None, "(constructor c (i64) (i64 i64))") + .unwrap_err(); + assert!( + matches!(err, Error::TypeError(TypeError::TupleOutputNotAllowed(..))), + "expected TupleOutputNotAllowed, got {err:?}" + ); +} + +#[test] +fn tuple_merge_arity_mismatch_rejected() { + let err = EGraph::default() + .parse_and_run_program( + None, + "(function f (i64) (i64 i64) :merge (values (max old0 new0)))", + ) + .unwrap_err(); + assert!( + matches!(err, Error::TypeError(TypeError::TupleMergeArity { .. })), + "expected TupleMergeArity, got {err:?}" + ); +} + +#[test] +fn builtin_keywords_are_reserved() { + // Built-in keywords (command/action/schedule heads, plus the tuple constructor `values`) may + // not be used as function/constructor/relation/variant names or as variables. + for prog in [ + // `values` (the tuple constructor) in each identifier position + "(function values (i64) i64 :merge new)", + "(sort S) (constructor values (i64) S)", + "(relation values (i64))", + "(datatype D (values i64))", + "(relation r (i64)) (rule ((r values)) ())", + // a sampling of command/action/schedule keywords + "(function set (i64) i64 :merge new)", + "(relation union (i64))", + "(sort run)", + "(datatype D (function i64))", + "(relation r (i64)) (rule ((r let)) ())", + ] { + let err = run(prog).expect_err(&format!("keyword should be rejected in: {prog}")); + assert!( + err.to_string().contains("reserved keyword"), + "expected a reserved-keyword error for {prog:?}, got: {err}" + ); + } + + // The `:` prefix is reserved for option keywords, so it cannot start an identifier. + let err = run("(relation :merge (i64))").unwrap_err(); + assert!( + err.to_string().contains("`:` prefix is reserved"), + "expected a colon-prefix error, got: {err}" + ); +} + +#[test] +fn input_output_are_only_partially_reserved() { + // `input`/`output` are usable as ordinary variables... + run("(relation r (i64)) (r 0) (rule ((r input)) ((r (+ input 1)))) (run 1) (check (r 1))") + .unwrap(); + + // ...but not as definition names... + for prog in [ + "(function input (i64) i64 :merge new)", + "(relation output (i64))", + "(sort input)", + "(datatype D (output i64))", + ] { + let err = run(prog).expect_err(&format!("`input`/`output` should be rejected in: {prog}")); + assert!( + err.to_string() + .contains("cannot be used as a definition name"), + "expected a definition-name error for {prog:?}, got: {err}" + ); + } + + // ...and not as the head of a call expression. + let err = run("(relation r (i64)) (rule ((r x)) ((r (input x))))").unwrap_err(); + assert!( + err.to_string() + .contains("cannot be used as the head of an expression"), + "expected a call-head error, got: {err}" + ); +} + +#[test] +fn tuple_merge_must_be_values_form() { + let err = EGraph::default() + .parse_and_run_program(None, "(function f (i64) (i64 i64) :merge old0)") + .unwrap_err(); + assert!( + matches!(err, Error::TypeError(TypeError::TupleMergeNotValues(..))), + "expected TupleMergeNotValues, got {err:?}" + ); +} + +#[test] +fn tuple_output_input_command_loads_all_columns() { + // `(input f file)` for a tuple-output function must load every output column, not just + // the first, so a TSV row has `input + all outputs` fields. + let dir = std::env::temp_dir().join("egglog_tuple_output_input_test"); + std::fs::create_dir_all(&dir).unwrap(); + std::fs::write(dir.join("iv.tsv"), "1\t10\t20\n2\t30\t40\n").unwrap(); + + let mut egraph = EGraph::default(); + egraph.fact_directory = Some(dir.clone()); + let result = egraph.parse_and_run_program( + None, + r#" + (function iv (i64) (i64 i64) :merge (values old0 old1)) + (input iv "iv.tsv") + (check (= (values 10 20) (iv 1))) + (check (= (values 30 40) (iv 2))) + "#, + ); + + std::fs::remove_dir_all(&dir).ok(); + result.unwrap(); +} + +#[test] +fn extract_tuple_output_rejected() { + // A tuple-output function returns more than one value, so it can't be extracted as a single + // term; this should be a clear error rather than a confusing arity mismatch. + let err = EGraph::default() + .parse_and_run_program( + None, + r#" + (datatype M (V i64)) + (function iv (M) (i64 i64) :merge (values (max old0 new0) (min old1 new1))) + (set (iv (V 0)) (values 3 7)) + (extract (iv (V 0))) + "#, + ) + .unwrap_err(); + assert!( + matches!( + err, + Error::TypeError(TypeError::CannotExtractTupleOutput(..)) + ), + "expected CannotExtractTupleOutput, got {err:?}" + ); +} + +#[test] +fn zero_input_tuple_output() { + // A tuple-output function may have no inputs (a single global tuple cell). + run(r#" + (function pt () (i64 i64) :merge (values (max old0 new0) (min old1 new1))) + (set (pt) (values 1 10)) + (set (pt) (values 5 3)) + (check (= (values 5 3) (pt))) + "#) + .unwrap(); +} + +#[test] +fn tuple_output_delete() { + // `delete` removes a tuple row by key. + run(r#" + (datatype M (V i64)) + (function iv (M) (i64 i64) :merge (values (max old0 new0) (min old1 new1))) + (set (iv (V 0)) (values 3 7)) + (check (= (values 3 7) (iv (V 0)))) + (delete (iv (V 0))) + (fail (check (= (values 3 7) (iv (V 0))))) + "#) + .unwrap(); +} + +#[test] +fn proof_mode_rejects_tuple_output() { + // The term/proof encoding is built around single-output constructor views and does not model + // user tuple-output functions, so declaring one under proofs is rejected. + let err = EGraph::new_with_proofs() + .parse_and_run_program( + None, + "(function iv (i64) (i64 i64) :merge (values old0 old1))", + ) + .unwrap_err(); + assert!( + matches!(err, Error::UnsupportedProofCommand { .. }) + && err + .to_string() + .contains("tuple-output functions are not supported"), + "expected UnsupportedProofCommand for a tuple output, got {err:?}" + ); +} + +#[test] +fn tuple_merge_unnumbered_old_new_rejected() { + // A tuple merge must address columns as `old0`/`new0`/`old1`/...; bare `old`/`new` don't bind. + let err = EGraph::default() + .parse_and_run_program( + None, + "(function iv (i64) (i64 i64) :merge (values old new))", + ) + .unwrap_err(); + assert!( + matches!(err, Error::TypeError(_)), + "expected a type error for bare old/new in a tuple merge, got {err:?}" + ); +}