diff --git a/.claude/skills/tidy-diff-docs/SKILL.md b/.claude/skills/tidy-diff-docs/SKILL.md index 61c6866d0..68568c45f 100644 --- a/.claude/skills/tidy-diff-docs/SKILL.md +++ b/.claude/skills/tidy-diff-docs/SKILL.md @@ -45,6 +45,17 @@ Two guardrails: need not repeat the shape; if a parameter's meaning is obvious from its name and type, do not spell it out again. +## Also fix in the diff + +- **Inline full paths for types.** Import the type and use the short name. + Write `use crate::util::HashMap;` and then `HashMap`, not + `crate::util::HashMap` in a field or a signature. A qualified + path inline makes the type harder to scan and hides the dependency from the + module's import list. Two exemptions: intra-doc links (``[`crate::Foo`]``) + need the path to resolve, and a path may disambiguate two same-named types in + scope. A qualified one-off *call* such as `std::mem::take(..)` is idiomatic + and not worth an import. + ## Keep these - What the item does and the shape of what it returns. diff --git a/CHANGELOG.md b/CHANGELOG.md index 334a043e3..2d5249da0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ - Fix a crash where bounded table scans with column constraints could yield stale (deleted) rows, leaking the stale value marker into primitives and panicking with an out-of-bounds intern-table read (e.g. `index out of bounds: the len is 0 but the index is 2147483647`). - Fix a build failure when egglog is compiled without default features (as a library dependency). The `egglog-add-primitive` proc macro parses full Rust expressions and now declares `syn`'s `full` feature directly, instead of relying on another crate (`clap_derive`, via the `bin` feature) to unify it onto our `syn`. This surfaced when `clap_derive` moved to `syn` 3.x. A CI job now builds `-p egglog --no-default-features` to catch regressions. - Convert many reachable `panic!`/`unwrap`/`expect`/`todo!` sites into recoverable errors, so malformed or edge-case programs report an error instead of aborting the process. Examples now returning `Error`/`TypeError`/`ParseError`/`ProveExistsError`: malformed sort-constructor declarations like `(sort S (Vec))` or `(sort S (UnstableFn))`; negative `extract` variant counts; duplicate rule names; `unstable-fn` referencing an unknown, non-literal, or mis-typed target; `(fail ...)` wrapping `include` or an empty expansion; subsuming a non-call rewrite; running `prove`/`prove-exists` without proofs enabled; and missing/unreadable files for `input`, `print-function`, `print-overall-statistics`, and the CLI. Several primitives became partial (returning no result instead of panicking) on out-of-range input: `vec-set`/`vec-remove` indices, `multiset-pick` on an empty multiset, count overflow in `multiset` operations, and the numeric primitives `bigint <<`/`>>`, `bigrat`, and `log2`. A few scheduler edge cases (unknown ruleset, rules with no free variables) no longer panic; variable-free rules now correctly apply their actions when scheduled. Primitive resolution now returns `TypeError::AmbiguousPrimitive`/`TypeError::UnresolvedPrimitive` instead of panicking when duplicate same-signature registrations are indistinguishable or nothing resolves; both direct calls and `unstable-fn` primitive targets report the same variants. `step_rules_with_scheduler` now restores its `rulesets`/`schedulers` on every fallible path, so an error during scheduled rule compilation no longer leaves the `EGraph` in a corrupted state. +- **Breaking:** a function's signature is stored once, shared between `TypeInfo` and `Function` instead of copied into both. `TypeInfo::get_func_type` returns `Option>`, `Function::schema()` becomes `Function::func_type() -> &FuncType`, and `ResolvedSchema` is removed (its `get_by_pos` moves to `FuncType`). - **Breaking:** `EGraph::print_function` now takes its output sink as `Option<(File, PathBuf)>` plus a `Span`, so write failures return `Error::IoError` instead of panicking. - Speed up query evaluation by building on-the-fly per-subset column indexes as sorted arrays (`SortedColumnIndex`) instead of hash maps. These indexes are typically iterated once and probed a bounded number of times over high-cardinality columns, so skipping hash-table construction is a large win (e.g. ~33% faster on the `gemma` benchmark). - Share trie roots (and their cached sub-indexes and child nodes) across query plans within a single `run_rule_set` instead of rebuilding a fresh trie per plan. Plans that scan the same table under the same header (fast) constraints reuse one root, so on-the-fly per-subset index builds happen once rather than per plan; only roots that more than one plan uses are shared, so workloads that would not benefit keep the per-plan behavior. Large speedups on transformer workloads (e.g. ~15% faster on `whisper`, ~12% on `gemma`, ~8% on `qwen3_moe`). @@ -27,7 +28,7 @@ - **Typed primitive surface for seminaive safety (#772).** Custom primitives now pick one of `PurePrim` / `ReadPrim` / `WritePrim` / `FullPrim` based on what the body needs, and register via the matching `add_*_primitive`. Rust enforces capability bounds via the state wrapper passed to the body; the egglog typechecker enforces context bounds. See the `egglog::exec_state` module docs and the `*Prim` trait docs for the full picture. Migration: `rust_rule` callbacks now take `&mut WriteState` (replacing `RustRuleContext`); a new `rust_rule_full` gives action callbacks read access. Higher-order primitives over `unstable-fn` values dispatch via `state.apply_function(&fc, args)`. - Expose `Read::table_size(name)` and `Read::table_sizes()` so read-capable primitives can inspect row counts without raw execution-state access, while avoiding an all-table scan when only one table is needed. - **`:naive` and `:unsafe-seminaive` rule options** (mutually exclusive). Both compile a rule under the permissive `Read`/`Full` contexts so its RHS can read the database (read-primitives and function-table lookups). `:naive` matches the whole database every iteration; `:unsafe-seminaive` keeps seminaive (delta) matching, which is faster but **unsafe** — an RHS read observes the database mid-iteration, so results can depend on evaluation order. `:unsafe-seminaive` is rejected by the term/proof encoding. -- **Name-indexed e-graph access from primitives and `rust_rule` callbacks (#745, #751).** New `Read` / `Write` capability traits on the state wrappers let primitive bodies and rule callbacks read/write tables by name (`fs.lookup`, `fs.set`, `fs.add`, `fs.union`, `fs.function_entries`, `fs.constructor_enodes`, etc.) instead of through raw `FunctionId` + `&[Value]`; `EGraph::update(|fs| ...)` gives the same surface outside a rule, and `EGraph::function_entries` / `EGraph::constructor_enodes` expose the table scans directly at the top level. Misuse (wrong subtype, wrong arity, unknown table) surfaces as `Error::ApiError`. +- **Name-indexed e-graph access from primitives and `rust_rule` callbacks (#745, #751).** New `Read` / `Write` capability traits on the state wrappers let primitive bodies and rule callbacks read/write tables by name (`fs.lookup`, `fs.set`, `fs.add`, `fs.union`, `fs.function_entries`, `fs.constructor_enodes`, etc.) instead of through raw `FunctionId` + `&[Value]`; `EGraph::update(|fs| ...)` gives the same surface outside a rule, and `EGraph::function_entries` / `EGraph::constructor_enodes` expose the table scans directly at the top level. Misuse (wrong subtype, wrong arity, unknown table) surfaces as `Error::ApiError`. Also `Read::enodes_for_eclass` (a constructor's rows by output e-class, through the backend's column index rather than a scan), `Read::constructor_schema` / `Read::function_schema` / `Read::table_subtype` (a table's declared signature and subtype, which a primitive body cannot get from `TypeInfo`), and `Core::rebuild_container` (remap a container value's contents and intern the result, for container sorts whose Rust type the caller cannot name). Together these let an out-of-tree primitive walk and rebuild a sub-e-graph — see `unstable-subst` in `egglog-experimental`. - **Container support in the term/proof encoding.** Programs using container sorts (`Vec`, `Set`, `Map`, `MultiSet`, `Pair`) now work under the term/proof encoding (previously rejected), including containers read (`vec-get`, `map-get`, …) or constructed (`vec-of`, `set-of`, …) in a rule body (`set-get` excepted: it indexes an internal runtime order that proofs cannot reproduce). A container built in the body is a *side condition* with no carryable proof: it is marked with an `Eval` proof step and re-evaluated against the typed rule when checked, so it can be read or matched in the query but not carried into an action (that is rejected). Two user-visible extraction changes: container terms extract in a deterministic, reproducible order rather than value-id order, and maps extract in a flat `(map-of k0 v0 …)` form (new `map-of` constructor) instead of nested `map-insert`s. ## [2.0.0] - 2026-02-11 diff --git a/core-relations/src/action/mod.rs b/core-relations/src/action/mod.rs index 2a1d82e0d..f8a92a918 100644 --- a/core-relations/src/action/mod.rs +++ b/core-relations/src/action/mod.rs @@ -12,18 +12,20 @@ use std::{ use crate::{ common::HashMap, - free_join::{invoke_batch, invoke_batch_assign}, + free_join::{get_column_index_from_tableinfo, invoke_batch, invoke_batch_assign}, numeric_id::{DenseIdMap, NumericId}, }; use egglog_concurrency::NotificationList; use smallvec::SmallVec; use crate::{ - BaseValues, ContainerValues, ExternalFunctionId, WrappedTable, + BaseValues, ContainerValues, ExternalFunctionId, Offset, WrappedTable, common::Value, free_join::{CounterId, Counters, ExternalFunctions, TableId, TableInfo, Variable}, + offsets::Subset, pool::{Clear, Pooled, with_pool_set}, - table_spec::{ColumnId, MutationBuffer}, + row_buffer::TaggedRowBuffer, + table_spec::{ColumnId, Constraint, MutationBuffer}, }; use self::mask::{Mask, MaskIter, ValueSource}; @@ -483,6 +485,66 @@ impl<'a> ExecutionState<'a> { &self.db.table_info[table].table } + /// Iterate over visible rows in `table` whose `col` equals `value`. + /// + /// Cacheable columns use the table's lazy column index; uncacheable columns + /// fall back to scanning with the equality constraint. The callback + /// receives each matching row as a full value slice. + pub fn for_each_matching_col( + &self, + table: TableId, + col: ColumnId, + value: Value, + mut f: impl FnMut(&[Value]), + ) { + let table_info = &self.db.table_info[table]; + let constraint = Constraint::EqConst { col, val: value }; + let (mut subset, _fast, mut slow) = table_info + .table + .split_fast_slow(std::slice::from_ref(&constraint)); + + debug_assert!(slow.iter().all(|c| matches!(c, Constraint::EqConst { .. }))); + + if !*table_info + .spec + .uncacheable_columns + .get(col) + .unwrap_or(&false) + { + let index = get_column_index_from_tableinfo(table_info, col); + match index.get().unwrap().get_subset(&value) { + Some(s) => { + with_pool_set(|ps| subset.intersect(s, &ps.get_pool())); + } + None => { + subset = Subset::empty(); + } + } + slow.clear(); + } + + let imp = &table_info.table; + let cols: SmallVec<[_; 8]> = (0..imp.spec().arity()).map(ColumnId::from_usize).collect(); + let mut cur = Offset::new(0); + let mut buf = TaggedRowBuffer::new_inline(imp.spec().arity()); + + macro_rules! drain_buf { + ($buf:expr) => { + for (_, row) in $buf.non_stale() { + f(row); + } + $buf.clear(); + }; + } + + while let Some(next) = imp.scan_project(subset.as_ref(), &cols, cur, 1024, &slow, &mut buf) + { + drain_buf!(buf); + cur = next; + } + drain_buf!(buf); + } + /// Get the human-readable name for a table, if one exists. pub fn table_name(&self, table: TableId) -> Option<&'a str> { self.db.table_info[table].name() diff --git a/core-relations/src/free_join/mod.rs b/core-relations/src/free_join/mod.rs index 97417988f..d1f5bf3b0 100644 --- a/core-relations/src/free_join/mod.rs +++ b/core-relations/src/free_join/mod.rs @@ -863,7 +863,10 @@ fn get_index_from_tableinfo(table_info: &TableInfo, cols: &[ColumnId]) -> HashIn /// The core logic behind getting and updating a column index. /// /// This is the single-column analog to [`get_index_from_tableinfo`]. -fn get_column_index_from_tableinfo(table_info: &TableInfo, col: ColumnId) -> HashColumnIndex { +pub(crate) fn get_column_index_from_tableinfo( + table_info: &TableInfo, + col: ColumnId, +) -> HashColumnIndex { let index: Arc<_> = table_info.column_indexes.get_or_insert(col, || { Arc::new(ResettableOnceLock::new(Index::new( vec![col], diff --git a/egglog-bridge/src/lib.rs b/egglog-bridge/src/lib.rs index bc890c007..e527fc3a1 100644 --- a/egglog-bridge/src/lib.rs +++ b/egglog-bridge/src/lib.rs @@ -1454,6 +1454,28 @@ impl TableAction { drain_buf!(buf); } + /// Iterate over rows whose output column is `value`. + /// + /// This reaches the table through an [`ExecutionState`], so read-capable + /// state wrappers can use the same lazy column indexes as direct backend + /// callers without exposing lower-level scan buffers or projection details. + pub fn for_each_output_value( + &self, + state: &ExecutionState, + value: Value, + mut f: impl FnMut(ScanEntry<'_>), + ) { + let schema_math = self.table_math; + let output_col = ColumnId::from_usize(self.input_arity()); + state.for_each_matching_col(self.table, output_col, value, |row| { + let subsumed = schema_math.subsume && row[schema_math.subsume_col()] == SUBSUMED; + f(ScanEntry { + vals: &row[0..schema_math.func_cols], + subsumed, + }); + }); + } + /// Look up a row, inserting the configured default value if absent. /// For constructor tables this mints a fresh eclass ID; for custom /// functions (no default) this behaves identically to diff --git a/src/core.rs b/src/core.rs index 47a68abb8..eeeea5d0b 100644 --- a/src/core.rs +++ b/src/core.rs @@ -158,7 +158,7 @@ impl ResolvedCall { let expected = ty.input.iter().map(|s| s.name()); let actual = types.iter().map(|s| s.name()); if expected.eq(actual) { - return Some(ResolvedCall::Func(ty.clone())); + return Some(ResolvedCall::Func(ty.as_ref().clone())); } } None @@ -175,7 +175,7 @@ impl ResolvedCall { let expected = ty.input.iter().chain(once(&ty.output)).map(|s| s.name()); let actual = types.iter().map(|s| s.name()); if expected.eq(actual) { - return Ok(ResolvedCall::Func(ty.clone())); + return Ok(ResolvedCall::Func(ty.as_ref().clone())); } } diff --git a/src/exec_state.rs b/src/exec_state.rs index ebac11e1d..3dcbb304a 100644 --- a/src/exec_state.rs +++ b/src/exec_state.rs @@ -30,7 +30,9 @@ //! [`ReadPrim`]: crate::ReadPrim //! [`FullPrim`]: crate::FullPrim +use std::any::TypeId; use std::ops::Deref; +use std::sync::Arc; use crate::Error; use crate::api::{ApiError, IntoValue, IntoValues}; @@ -42,7 +44,7 @@ use crate::{ ast::{FunctionSubtype, Literal, ResolvedExpr}, core::ResolvedCall, sort::{F, S}, - typechecking::FuncType, + typechecking::{FuncType, SharedFuncTypes}, }; use egglog_bridge::{ActionRegistry, TableAction, TableKind}; use smallvec::SmallVec; @@ -142,6 +144,9 @@ pub(crate) trait Internal<'a, 'db: 'a>: 'a { /// name through it. pub(crate) trait RegistrySealed<'a, 'db: 'a>: Internal<'a, 'db> { fn registry(&self) -> &ActionRegistry; + /// The e-graph's signature map, unlocked. The schema accessors lock it + /// only when they are called, so a primitive that never asks pays nothing. + fn func_types(&self) -> &'a SharedFuncTypes; } // ===================================================================== @@ -183,6 +188,29 @@ pub trait Core<'a, 'db: 'a>: Internal<'a, 'db> { cv.register_val(container, es) } + /// Rebuild the contents of a container value through `remap`, + /// returning the interned value of the result (or `value` unchanged + /// if it is not a registered container of `type_id`, or if nothing + /// inside it was remapped). `type_id` comes from + /// [`Sort::value_type`](crate::sort::Sort::value_type). + /// + /// `remap` sees every value the container stores, including ones of + /// sorts it does not care about, so it should return its argument + /// unchanged for those. + fn rebuild_container( + &mut self, + type_id: TypeId, + value: Value, + remap: &(dyn Fn(Value) -> Value + Send + Sync), + ) -> Value { + // Same borrow shape as `register_container`: the `ContainerValues` + // reference is tied to the inner ExecutionState's lifetime, not + // to `&self`, so it survives the `&mut` reborrow. + let cv = self.container_values(); + let es = self.es_mut(); + cv.rebuild_val_with(type_id, value, es, remap) + } + /// Convert an egglog [`Value`] to a Rust base type, assuming that the /// value belongs to `T`. fn value_to_base(&self, x: Value) -> T { @@ -304,8 +332,10 @@ pub trait Core<'a, 'db: 'a>: Internal<'a, 'db> { /// The single-entry methods (`lookup`, `eclass_of`, `contains`) /// return `None` if absent — never insert. The iteration / /// introspection methods (`function_entries`, `constructor_enodes`, -/// `table_size`, `table_sizes`) walk the current contents of the -/// database. +/// `enodes_for_eclass`, `table_size`, `table_sizes`) walk the current +/// contents of the database, while `constructor_schema` / +/// `function_schema` / `table_subtype` report how a table is declared +/// rather than what it holds. /// /// Detectable misuse (wrong table subtype, wrong arity) is reported /// as [`crate::ApiError`] via the method's `Result`. Per-column sort @@ -349,6 +379,32 @@ pub trait Read<'a, 'db: 'a>: Core<'a, 'db> + RegistrySealed<'a, 'db> { Ok(action.lookup(self.es(), &key_values).is_some()) } + /// A constructor's declared signature: its input sorts and the sort of the + /// eclass column. Pair with [`Read::table_sizes`] to walk every table. + /// + /// **Only valid for constructor tables.** Functions error; use + /// [`Read::function_schema`] for those. + fn constructor_schema(&self, name: &str) -> Result, Error> { + func_type_of(self.func_types(), name, FunctionSubtype::Constructor) + } + + /// A function's declared signature: its input sorts and its output sort. + /// + /// **Only valid for `function` tables.** Constructors error; use + /// [`Read::constructor_schema`] for those. + fn function_schema(&self, name: &str) -> Result, Error> { + func_type_of(self.func_types(), name, FunctionSubtype::Custom) + } + + /// Whether the named table is a `constructor` or a `function`, or `None` + /// if no table with that name is registered. Unlike the two schema + /// accessors, a mismatch is not an error, so this is the one to dispatch + /// on when either subtype is acceptable. + fn table_subtype(&self, name: &str) -> Option { + let func_types = self.func_types().read().unwrap(); + func_types.get(name).map(|func_type| func_type.subtype) + } + /// Return the current row count for the named table, or `None` if no table /// with that name is registered. fn table_size(&self, name: &str) -> Option { @@ -395,6 +451,34 @@ pub trait Read<'a, 'db: 'a>: Core<'a, 'db> + RegistrySealed<'a, 'db> { Ok(()) } + /// Call `f` on each [`Enode`] of a constructor / relation table whose + /// output eclass is `eclass`. + /// + /// This uses the backend's indexed output-column lookup instead of scanning + /// the whole constructor table. Errors with `WrongSubtype` if `name` is a + /// function. + fn enodes_for_eclass( + &self, + name: &str, + eclass: Value, + mut f: impl FnMut(Enode<'_>), + ) -> Result<(), Error> { + let action = lookup_action(self.registry(), name)?; + check_subtype(name, &action, TableKind::Constructor, "constructor")?; + action.for_each_output_value(self.es(), eclass, |row| { + let (eclass, children) = row + .vals + .split_last() + .expect("constructor row has at least an eclass column"); + f(Enode { + children, + eclass: *eclass, + subsumed: row.subsumed, + }); + }); + Ok(()) + } + /// Call `f` on each [`FunctionEntry`] of a function table. Errors /// with `WrongSubtype` if `name` is a constructor. To stop early, /// use [`Read::function_entries_while`]. @@ -541,6 +625,36 @@ pub trait Write<'a, 'db: 'a>: Core<'a, 'db> + RegistrySealed<'a, 'db> { } } +fn subtype_label(subtype: FunctionSubtype) -> &'static str { + match subtype { + FunctionSubtype::Constructor => "constructor", + FunctionSubtype::Custom => "function", + } +} + +fn func_type_of( + func_types: &SharedFuncTypes, + name: &str, + expected: FunctionSubtype, +) -> Result, Error> { + let func_type = func_types.read().unwrap().get(name).cloned(); + let Some(func_type) = func_type else { + return Err(ApiError::MissingTable { + name: name.to_string(), + } + .into()); + }; + if func_type.subtype != expected { + return Err(ApiError::WrongSubtype { + name: name.to_string(), + expected: subtype_label(expected), + actual: subtype_label(func_type.subtype), + } + .into()); + } + Ok(func_type) +} + fn lookup_action(registry: &ActionRegistry, name: &str) -> Result { registry.lookup_table(name).cloned().ok_or_else(|| { ApiError::MissingTable { @@ -626,6 +740,7 @@ pub struct PureState<'a, 'db> { pub struct ReadState<'a, 'db> { pub(crate) inner: &'a mut ExecutionState<'db>, pub(crate) registry: &'a ActionRegistry, + pub(crate) func_types: &'a SharedFuncTypes, pub(crate) ctx: Context, } @@ -640,6 +755,7 @@ pub struct ReadState<'a, 'db> { pub struct WriteState<'a, 'db> { pub(crate) inner: &'a mut ExecutionState<'db>, pub(crate) registry: &'a ActionRegistry, + pub(crate) func_types: &'a SharedFuncTypes, pub(crate) ctx: Context, } @@ -654,6 +770,7 @@ pub struct WriteState<'a, 'db> { pub struct FullState<'a, 'db> { pub(crate) inner: &'a mut ExecutionState<'db>, pub(crate) registry: &'a ActionRegistry, + pub(crate) func_types: &'a SharedFuncTypes, pub(crate) ctx: Context, } @@ -670,11 +787,13 @@ impl<'a, 'db: 'a> ReadState<'a, 'db> { pub(crate) fn wrap( es: &'a mut ExecutionState<'db>, registry: &'a ActionRegistry, + func_types: &'a SharedFuncTypes, ctx: Context, ) -> Self { Self { inner: es, registry, + func_types, ctx, } } @@ -687,11 +806,13 @@ impl<'a, 'db: 'a> WriteState<'a, 'db> { pub(crate) fn wrap( es: &'a mut ExecutionState<'db>, registry: &'a ActionRegistry, + func_types: &'a SharedFuncTypes, ctx: Context, ) -> Self { Self { inner: es, registry, + func_types, ctx, } } @@ -704,11 +825,13 @@ impl<'a, 'db: 'a> FullState<'a, 'db> { pub(crate) fn wrap( es: &'a mut ExecutionState<'db>, registry: &'a ActionRegistry, + func_types: &'a SharedFuncTypes, ctx: Context, ) -> Self { Self { inner: es, registry, + func_types, ctx, } } @@ -753,6 +876,9 @@ impl<'a, 'db: 'a> RegistrySealed<'a, 'db> for ReadState<'a, 'db> { fn registry(&self) -> &ActionRegistry { self.registry } + fn func_types(&self) -> &'a SharedFuncTypes { + self.func_types + } } impl<'a, 'db: 'a> Core<'a, 'db> for ReadState<'a, 'db> {} impl<'a, 'db: 'a> Read<'a, 'db> for ReadState<'a, 'db> {} @@ -775,6 +901,9 @@ impl<'a, 'db: 'a> RegistrySealed<'a, 'db> for WriteState<'a, 'db> { fn registry(&self) -> &ActionRegistry { self.registry } + fn func_types(&self) -> &'a SharedFuncTypes { + self.func_types + } } impl<'a, 'db: 'a> Core<'a, 'db> for WriteState<'a, 'db> {} impl<'a, 'db: 'a> Write<'a, 'db> for WriteState<'a, 'db> {} @@ -797,6 +926,9 @@ impl<'a, 'db: 'a> RegistrySealed<'a, 'db> for FullState<'a, 'db> { fn registry(&self) -> &ActionRegistry { self.registry } + fn func_types(&self) -> &'a SharedFuncTypes { + self.func_types + } } impl<'a, 'db: 'a> Core<'a, 'db> for FullState<'a, 'db> {} impl<'a, 'db: 'a> Read<'a, 'db> for FullState<'a, 'db> {} diff --git a/src/extract.rs b/src/extract.rs index 1b7593e28..f8e385db1 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -209,7 +209,7 @@ impl Extractor { let func = egraph.functions.get(h).unwrap(); // For view tables, children are all but the last input (which is the e-class) let num_children = func.extraction_num_children(); - for ch in func.schema.input.iter().take(num_children) { + for ch in func.func_type.input.iter().take(num_children) { let ch_name = ch.name(); if !seen.contains(ch_name) { q.push_back(ch.clone()); @@ -284,7 +284,7 @@ impl Extractor { func: &Function, ) -> Option { let mut ch_costs: Vec = Vec::new(); - let sorts = &func.schema.input; + let sorts = &func.func_type.input; let num_children = func.extraction_num_children(); for (value, sort) in row.vals.iter().take(num_children).zip(sorts.iter()) { ch_costs.push(self.compute_cost_node(egraph, *value, sort)?); @@ -327,7 +327,7 @@ impl Extractor { row: &egglog_bridge::ScanEntry, func: &Function, ) -> usize { - let sorts = &func.schema.input; + let sorts = &func.func_type.input; let num_children = func.extraction_num_children(); row.vals .iter() @@ -489,7 +489,7 @@ impl Extractor { .get(&value) .unwrap(); let func = egraph.functions.get(func_name).unwrap(); - let ch_sorts = &func.schema.input; + let ch_sorts = &func.func_type.input; let num_children = func.extraction_num_children(); let output_name = func.extraction_term_name(); @@ -653,7 +653,7 @@ impl Extractor { for (cost, func_name, hyperedge) in root_variants { let mut ch_terms: Vec = Vec::new(); let func = egraph.functions.get(&func_name).unwrap(); - let ch_sorts = &func.schema.input; + let ch_sorts = &func.func_type.input; let num_children = func.extraction_num_children(); // For view tables, children are all but the last input (which is the e-class) for (value, sort) in hyperedge.iter().zip(ch_sorts.iter()).take(num_children) { @@ -725,9 +725,9 @@ impl Function { /// 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() { - self.schema.input.last().unwrap() + self.func_type.input.last().unwrap() } else { - &self.schema.output + &self.func_type.output } } @@ -735,9 +735,9 @@ impl Function { /// For view tables, this excludes the last column (the e-class). pub(crate) fn extraction_num_children(&self) -> usize { if self.decl.term_constructor.is_some() { - self.schema.input.len() - 1 + self.func_type.input.len() - 1 } else { - self.schema.input.len() + self.func_type.input.len() } } @@ -758,10 +758,10 @@ impl Function { // 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 - self.schema.input.len() - 1 + self.func_type.input.len() - 1 } else { // For regular tables: row is [inputs..., output] - self.schema.input.len() + self.func_type.input.len() } } } @@ -822,9 +822,9 @@ impl EGraph { .functions .get(sym) .ok_or(TypeError::UnboundFunction(sym.to_owned(), span!()))?; - let mut rootsorts = func.schema.input.clone(); + let mut rootsorts = func.func_type.input.clone(); if include_output { - rootsorts.push(func.schema.output.clone()); + rootsorts.push(func.func_type.output.clone()); } let extractor = Extractor::compute_costs_from_rootsorts( Some(rootsorts), @@ -844,7 +844,7 @@ impl EGraph { if inputs.len() < n { // include subsumed rows let mut children: Vec = Vec::new(); - for (value, sort) in row.vals.iter().zip(&func.schema.input) { + for (value, sort) in row.vals.iter().zip(&func.func_type.input) { let (_, term_id) = extractor .extract_best_with_sort(self, &mut termdag, *value, sort.clone()) .unwrap_or_else(|| (0, termdag.var("Unextractable".into()))); @@ -852,8 +852,8 @@ 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 value = row.vals[func.func_type.input.len()]; + let sort = &func.func_type.output; let (_, term) = extractor .extract_best_with_sort(self, &mut termdag, value, sort.clone()) .unwrap_or_else(|| (0, termdag.var("Unextractable".into()))); diff --git a/src/lib.rs b/src/lib.rs index abbd2ed62..5475a2185 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -71,6 +71,7 @@ use std::path::PathBuf; use std::sync::Arc; pub use termdag::{OrdTerm, Term, TermDag, TermId}; use thiserror::Error; +pub use typechecking::FuncType; pub use typechecking::PrimitiveValidator; pub use typechecking::TypeError; pub use typechecking::TypeInfo; @@ -220,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.func_type.output.name() == UnitSort.name(); if *mode == PrintFunctionMode::CSV { let mut wtr = Writer::from_writer(vec![]); for (term_id, output) in terms_and_outputs { @@ -326,7 +327,8 @@ pub trait UserDefinedCommand: Send + Sync { #[derive(Clone)] pub struct Function { decl: ResolvedFunctionDecl, - schema: ResolvedSchema, + /// Shared with [`TypeInfo`], which is where a signature is resolved. + func_type: Arc, can_subsume: bool, backend_id: egglog_bridge::FunctionId, } @@ -337,9 +339,9 @@ impl Function { &self.decl.name } - /// Get the schema of the function. - pub fn schema(&self) -> &ResolvedSchema { - &self.schema + /// The function's resolved signature. + pub fn func_type(&self) -> &FuncType { + &self.func_type } /// Whether this function supports subsumption. @@ -366,28 +368,11 @@ impl Function { } } -#[derive(Clone, Debug)] -pub struct ResolvedSchema { - pub input: Vec, - pub output: ArcSort, -} - -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) - } - } -} - impl Debug for Function { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("Function") .field("decl", &self.decl) - .field("schema", &self.schema) + .field("func_type", &self.func_type) .finish() } } @@ -729,7 +714,14 @@ impl EGraph { // Preserve the symbol generator so that fresh symbols // generated after pop don't collide with ones generated before pop. std::mem::swap(&mut self.parser.symbol_gen, &mut e.parser.symbol_gen); + // Restore the pushed signatures into the cell the registered + // primitives hold, rather than swapping in the saved cell they + // have no handle to. + let live_func_types = self.type_info.func_types_handle().clone(); + let restored = e.type_info.func_types_handle().read().unwrap().clone(); *self = *e; + *live_func_types.write().unwrap() = restored; + self.type_info.set_func_types_handle(live_func_types); Ok(()) } None => Err(Error::Pop(span!())), @@ -799,22 +791,36 @@ impl EGraph { } 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()), - None => Err(Error::TypeError(TypeError::UndefinedSort( - name.to_owned(), - decl.span.clone(), - ))), + // Typechecking resolved this signature already; reuse it rather than + // resolving the sorts again into a second copy. Desugaring generates + // some functions (global bindings, proof tables) without typechecking + // them, so those are resolved and recorded here instead. + let func_type = match self.type_info.get_func_type(&decl.name) { + Some(func_type) => func_type, + None => { + let get_sort = |name: &String| match self.type_info.get_sort_by_name(name) { + Some(sort) => Ok(sort.clone()), + None => Err(Error::TypeError(TypeError::UndefinedSort( + name.to_owned(), + decl.span.clone(), + ))), + }; + let func_type = Arc::new(FuncType { + name: decl.name.clone(), + subtype: decl.subtype, + input: decl + .schema + .input + .iter() + .map(get_sort) + .collect::, _>>()?, + output: get_sort(&decl.schema.output)?, + }); + self.type_info.declare_func_type(func_type.clone()); + func_type + } }; - let input = decl - .schema - .input - .iter() - .map(get_sort) - .collect::, _>>()?; - let output = get_sort(&decl.schema.output)?; - let can_subsume = match decl.subtype { FunctionSubtype::Constructor => true, // View tables (functions with term_constructor) need subsumption support @@ -823,9 +829,10 @@ impl EGraph { use egglog_bridge::{DefaultVal, MergeFn}; let backend_id = self.backend.add_table(egglog_bridge::FunctionConfig { - schema: input + schema: func_type + .input .iter() - .chain([&output]) + .chain([&func_type.output]) .map(|sort| sort.column_ty(&self.backend)) .collect(), default: match decl.subtype { @@ -845,7 +852,7 @@ impl EGraph { let function = Function { decl: decl.clone(), - schema: ResolvedSchema { input, output }, + func_type, can_subsume, backend_id, }; @@ -1231,8 +1238,11 @@ impl EGraph { pub fn read(&self, f: impl FnOnce(ReadState<'_, '_>) -> R) -> R { let registry = self.backend.action_registry().clone(); let guard = registry.read().unwrap(); + let func_types = self.type_info.func_types_handle(); self.backend - .with_execution_state_tracked(|es| f(ReadState::wrap(es, &guard, Context::Read))) + .with_execution_state_tracked(|es| { + f(ReadState::wrap(es, &guard, func_types, Context::Read)) + }) .0 } @@ -1874,7 +1884,7 @@ impl EGraph { // check that the function uses supported types - for t in &func.schema.input { + for t in &func.func_type.input { match t.name() { "i64" | "f64" | "String" => {} s => return Err(Error::UnsupportedInputType(s.to_string(), span.clone())), @@ -1882,7 +1892,7 @@ impl EGraph { } if function_type.subtype != FunctionSubtype::Constructor { - match func.schema.output.name() { + match func.func_type.output.name() { "i64" | "String" | "Unit" => {} s => return Err(Error::UnsupportedInputType(s.to_string(), span.clone())), } @@ -1898,9 +1908,9 @@ impl EGraph { // Can also do a row-major Vec let mut parsed_contents: Vec> = Vec::with_capacity(contents.lines().count()); - let mut row_schema = func.schema.input.clone(); + let mut row_schema = func.func_type.input.clone(); if function_type.subtype == FunctionSubtype::Custom { - row_schema.push(func.schema.output.clone()); + row_schema.push(func.func_type.output.clone()); } log::debug!("{row_schema:?}"); @@ -2379,9 +2389,10 @@ impl EGraph { ) -> Result { let registry = self.backend.action_registry().clone(); let guard = registry.read().unwrap(); - let (result, changed) = self - .backend - .with_execution_state_tracked(|es| f(FullState::wrap(es, &guard, Context::Full))); + let func_types = self.type_info.func_types_handle().clone(); + let (result, changed) = self.backend.with_execution_state_tracked(|es| { + f(FullState::wrap(es, &guard, &func_types, Context::Full)) + }); drop(guard); // A read-only closure stages nothing, so `flush_updates` would only do // a no-op merge plus a spurious timestamp bump and rebuild check. Skip diff --git a/src/proofs/proof_extraction.rs b/src/proofs/proof_extraction.rs index 535686754..8a6ce8eeb 100644 --- a/src/proofs/proof_extraction.rs +++ b/src/proofs/proof_extraction.rs @@ -41,7 +41,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.func_type.output.clone(); let mut termdag = TermDag::default(); let mut witness_value = None; @@ -79,7 +79,7 @@ impl ProofInstrumentor<'_> { func.name ) }) - .schema + .func_type .output .clone(); let proof_value = self diff --git a/src/proofs/proof_extractor.rs b/src/proofs/proof_extractor.rs index a7a58230c..369ed6f55 100644 --- a/src/proofs/proof_extractor.rs +++ b/src/proofs/proof_extractor.rs @@ -115,7 +115,11 @@ impl RootExtractor { let num_children = func.extraction_num_children(); let mut ch_terms = Vec::with_capacity(num_children); let mut valid = true; - for (value, sort) in row.iter().take(num_children).zip(func.schema.input.iter()) { + for (value, sort) in row + .iter() + .take(num_children) + .zip(func.func_type.input.iter()) + { match self.extract(egraph, termdag, *value, sort) { Some(term) => ch_terms.push(term), None => { diff --git a/src/serialize.rs b/src/serialize.rs index 4ff3230e8..b81b1fbe5 100644 --- a/src/serialize.rs +++ b/src/serialize.rs @@ -149,7 +149,7 @@ impl EGraph { return false; } let (out, inps) = row.vals.split_last().unwrap(); - let class_id = self.value_to_class_id(&function.schema.output, *out); + let class_id = self.value_to_class_id(&function.func_type.output, *out); if function.decl.internal_let { let_bindings .entry(class_id.clone()) @@ -186,7 +186,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.func_type.output.is_eq_sort() { acc.entry(class_id.clone()) .or_default() .push_back(node_id.clone()); @@ -202,12 +202,12 @@ 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.func_type.output, output, &class_id); - assert_eq!(input.len(), func.schema.input.len()); + assert_eq!(input.len(), func.func_type.input.len()); let children: Vec<_> = input .iter() - .zip(&func.schema.input) + .zip(&func.func_type.input) .map(|(&v, sort)| { self.serialize_value(&mut serializer, sort, v, &self.value_to_class_id(sort, v)) }) diff --git a/src/typechecking.rs b/src/typechecking.rs index 28031a56b..3c479afe9 100644 --- a/src/typechecking.rs +++ b/src/typechecking.rs @@ -44,6 +44,7 @@ impl ExternalFunction for PurePrimWrapper { struct RegistryPrimWrapper { prim: T, registry: Arc>, + func_types: SharedFuncTypes, /// Stamped onto the state wrapper. ctx: Context, _wrap: std::marker::PhantomData S>, @@ -56,6 +57,7 @@ trait RegistryWrap: Clone + Send + Sync { ctx: Context, args: &[Value], registry: &ActionRegistry, + func_types: &SharedFuncTypes, ) -> Option; } @@ -69,8 +71,9 @@ impl RegistryWrap for WrapRead { ctx: Context, args: &[Value], registry: &ActionRegistry, + func_types: &SharedFuncTypes, ) -> Option { - prim.apply(ReadState::wrap(exec_state, registry, ctx), args) + prim.apply(ReadState::wrap(exec_state, registry, func_types, ctx), args) } } #[derive(Clone)] @@ -83,8 +86,12 @@ impl RegistryWrap for WrapWrite { ctx: Context, args: &[Value], registry: &ActionRegistry, + func_types: &SharedFuncTypes, ) -> Option { - prim.apply(WriteState::wrap(exec_state, registry, ctx), args) + prim.apply( + WriteState::wrap(exec_state, registry, func_types, ctx), + args, + ) } } #[derive(Clone)] @@ -97,8 +104,9 @@ impl RegistryWrap for WrapFull { ctx: Context, args: &[Value], registry: &ActionRegistry, + func_types: &SharedFuncTypes, ) -> Option { - prim.apply(FullState::wrap(exec_state, registry, ctx), args) + prim.apply(FullState::wrap(exec_state, registry, func_types, ctx), args) } } @@ -107,7 +115,14 @@ impl + 'static> ExternalFun { fn invoke(&self, exec_state: &mut ExecutionState, args: &[Value]) -> Option { let registry = self.registry.read().unwrap(); - S::invoke(&self.prim, exec_state, self.ctx, args, ®istry) + S::invoke( + &self.prim, + exec_state, + self.ctx, + args, + ®istry, + &self.func_types, + ) } } @@ -119,6 +134,18 @@ pub struct FuncType { pub output: ArcSort, } +impl FuncType { + /// The sort at column `index`, counting the output column as + /// `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) + } + } +} + impl PartialEq for FuncType { fn eq(&self, other: &Self) -> bool { if self.name == other.name @@ -152,6 +179,15 @@ impl Hash for FuncType { } } } +/// Every declared function's signature, by name. +/// +/// Shared (via `Arc>`) with the state wrappers, which cannot reach +/// [`TypeInfo`] from a primitive body and read it through +/// [`Read::constructor_schema`] / [`Read::function_schema`]. `EGraph`'s +/// `push`/`pop` snapshot and restore its contents so a popped declaration +/// stops resolving. +pub(crate) type SharedFuncTypes = Arc>>>; + /// Validators take a termdag and arguments (as TermIds) and return /// a newly computed TermId if the primitive application is valid, /// or None if it is invalid. @@ -204,19 +240,37 @@ impl Debug for PrimitiveWithId { } /// Stores resolved typechecking information. -#[derive(Clone, Default)] +#[derive(Default)] pub struct TypeInfo { mksorts: HashMap, // TODO(yz): I want to get rid of this as now we have user-defined primitives and constraint based type checking reserved_primitives: HashSet<&'static str>, pub(crate) sorts: HashMap>, primitives: HashMap>, - func_types: HashMap, + func_types: SharedFuncTypes, pub(crate) global_sorts: HashMap, /// Sorts that do not allow union (e.g., from `:no-union` sorts or relations). pub(crate) non_unionable_sorts: HashSet, } +impl Clone for TypeInfo { + /// `func_types` is deep-copied rather than shared. A clone is an + /// independent e-graph — a `push`ed copy, or the parallel typechecking the + /// proof checker keeps — and declaring a function in one must not make it + /// resolve in the other. + fn clone(&self) -> Self { + Self { + mksorts: self.mksorts.clone(), + reserved_primitives: self.reserved_primitives.clone(), + sorts: self.sorts.clone(), + primitives: self.primitives.clone(), + func_types: Arc::new(RwLock::new(self.func_types.read().unwrap().clone())), + global_sorts: self.global_sorts.clone(), + non_unionable_sorts: self.non_unionable_sorts.clone(), + } + } +} + // These methods need to be on the `EGraph` in order to // register sorts and primitives with the backend. impl EGraph { @@ -237,7 +291,13 @@ impl EGraph { span: Span, ) -> Result<(), TypeError> { let name = name.into(); - if self.type_info.func_types.contains_key(&name) { + if self + .type_info + .func_types + .read() + .unwrap() + .contains_key(&name) + { return Err(TypeError::FunctionAlreadyBound(name, span)); } @@ -328,10 +388,12 @@ impl EGraph { S: RegistryWrap + 'static, { let registry = self.backend.action_registry().clone(); + let func_types = self.type_info.func_types_handle().clone(); self.register_per_context(x, validator, valid_ctxs, move |x, ctx| { Box::new(RegistryPrimWrapper:: { prim: x, registry: registry.clone(), + func_types: func_types.clone(), ctx, _wrap: std::marker::PhantomData, }) @@ -572,7 +634,10 @@ impl EGraph { span.clone(), )); } - ResolvedNCommand::ProveExists(span.clone(), ResolvedCall::Func(func_type.clone())) + ResolvedNCommand::ProveExists( + span.clone(), + ResolvedCall::Func(func_type.as_ref().clone()), + ) } NCommand::Output { span, file, exprs } => { let exprs = exprs @@ -784,7 +849,13 @@ impl TypeInfo { )); } let ftype = self.function_to_functype(fdecl)?; - if self.func_types.insert(fdecl.name.clone(), ftype).is_some() { + if self + .func_types + .write() + .unwrap() + .insert(fdecl.name.clone(), Arc::new(ftype)) + .is_some() + { return Err(TypeError::FunctionAlreadyBound( fdecl.name.clone(), fdecl.span.clone(), @@ -805,7 +876,12 @@ impl TypeInfo { name: fdecl.name.clone(), subtype: fdecl.subtype, schema: fdecl.schema.clone(), - resolved_schema: ResolvedCall::Func(self.func_types.get(&fdecl.name).unwrap().clone()), + resolved_schema: ResolvedCall::Func( + self.get_func_type(&fdecl.name) + .expect("just inserted") + .as_ref() + .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 @@ -1138,13 +1214,32 @@ impl TypeInfo { .any(|p| p.context_ids.iter().any(|(_, pid)| *pid == Some(id)) && p.validator.is_some()) } - pub fn get_func_type(&self, sym: &str) -> Option<&FuncType> { - self.func_types.get(sym) + pub fn get_func_type(&self, sym: &str) -> Option> { + self.func_types.read().unwrap().get(sym).cloned() } - pub fn is_constructor(&self, sym: &str) -> bool { + /// Record a signature for a function that did not come through + /// typechecking — desugaring generates some (global bindings, proof + /// tables) directly. + pub(crate) fn declare_func_type(&mut self, func_type: Arc) { self.func_types - .get(sym) + .write() + .unwrap() + .insert(func_type.name.clone(), func_type); + } + + /// The shared signature map, for sharing with the state wrappers. + pub(crate) fn func_types_handle(&self) -> &SharedFuncTypes { + &self.func_types + } + + /// Adopt an existing shared cell as this `TypeInfo`'s signature map. + pub(crate) fn set_func_types_handle(&mut self, func_types: SharedFuncTypes) { + self.func_types = func_types; + } + + pub fn is_constructor(&self, sym: &str) -> bool { + self.get_func_type(sym) .is_some_and(|f| f.subtype == FunctionSubtype::Constructor) }