From 1118c51a327c3ef99a97e60122c43b08f513ccb1 Mon Sep 17 00:00:00 2001 From: Oliver Flatt Date: Thu, 6 Aug 2026 19:44:57 +0000 Subject: [PATCH 1/8] Expose enough e-graph introspection to walk and rebuild a sub-e-graph Three additions to what a primitive body can see, none of them specific to any one extension. Together they are what an out-of-tree primitive needs to walk the term structure under an e-class and build a modified copy of it. Read::enodes_for_eclass(name, eclass, f) walks a constructor's rows by output e-class through the backend's lazy column index, instead of scanning the table and filtering. Cherry-picked from #934 along with the core-relations ExecutionState::for_each_matching_col and egglog-bridge TableAction::for_each_output_value it rests on. Read::table_schema(name) and Read::table_subtype(name) report a table's declared column sorts and its subtype. EGraph::functions_iter already exposes this from &EGraph, but a primitive body only ever sees a state wrapper, and those carried no sort information at all - so a primitive could read rows without being able to tell an e-class column from a base value. Backed by a FunctionSchemas map the e-graph shares with the wrappers exactly as it already shares ActionRegistry, and snapshot/restored across push/pop so a popped table stops resolving. table_subtype also replaces probing a subtype by starting a scan and reading the error, which egglog-experimental does today. Core::rebuild_container(type_id, value, remap) remaps a container value's contents and interns the result. Out-of-tree code cannot go through Core::register_container, which requires naming the container's Rust type. Co-Authored-By: Claude Opus 5 --- CHANGELOG.md | 2 + core-relations/src/action/mod.rs | 68 ++++++++++++++++- core-relations/src/free_join/mod.rs | 5 +- egglog-bridge/src/lib.rs | 22 ++++++ src/exec_state.rs | 114 +++++++++++++++++++++++++++- src/lib.rs | 44 +++++++++-- src/typechecking.rs | 17 ++++- 7 files changed, 255 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 334a043e3..fc7e3ed89 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,8 @@ ## [Unreleased] - ReleaseDate +- **Enough e-graph introspection for a primitive to walk and rebuild a sub-e-graph.** Three additions, together sufficient to implement substitution as an out-of-tree primitive (see `unstable-subst` in `egglog-experimental`): `Read::enodes_for_eclass(name, eclass, f)` walks a constructor's rows by output e-class through the backend's column index instead of scanning the table (with `ExecutionState::for_each_matching_col` and `TableAction::for_each_output_value` behind it); `Read::table_schema(name)` and `Read::table_subtype(name)` expose a table's declared column sorts and subtype to a primitive body, which cannot reach `TypeInfo` (`table_subtype` also replaces probing a table's subtype by starting a scan); and `Core::rebuild_container(type_id, value, remap)` remaps a container value's contents and interns the result, which out-of-tree code cannot do through `register_container` because it cannot name an arbitrary container's Rust type. + - 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. 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/exec_state.rs b/src/exec_state.rs index ebac11e1d..192eebb0c 100644 --- a/src/exec_state.rs +++ b/src/exec_state.rs @@ -39,6 +39,7 @@ use crate::core_relations::{ Value, }; use crate::{ + ResolvedSchema, ast::{FunctionSubtype, Literal, ResolvedExpr}, core::ResolvedCall, sort::{F, S}, @@ -84,6 +85,28 @@ impl Context { pub const ALL: [Context; 4] = [Context::Pure, Context::Write, Context::Read, Context::Full]; } +/// The declared column sorts of every table, by name. +/// +/// The [`ActionRegistry`] resolves a table name to the handle that reads and +/// writes it, but it lives in the backend and so knows nothing about egglog +/// sorts. This is the matching name-indexed view of the *types*, kept by the +/// e-graph and shared with the state wrappers so a primitive body — which +/// cannot reach `TypeInfo` — can still ask what a column holds. +#[derive(Clone, Default)] +pub struct FunctionSchemas { + tables: crate::util::HashMap, +} + +impl FunctionSchemas { + pub(crate) fn declare(&mut self, name: &str, subtype: FunctionSubtype, schema: ResolvedSchema) { + self.tables.insert(name.to_owned(), (subtype, schema)); + } + + fn get(&self, name: &str) -> Option<&(FunctionSubtype, ResolvedSchema)> { + self.tables.get(name) + } +} + // ===================================================================== // Sealed traits. // @@ -142,6 +165,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; + /// Borrowed for `'a` rather than for `&self`, so a caller can hold a + /// column sort while it goes on to write through the same state. + fn schemas(&self) -> &'a FunctionSchemas; } // ===================================================================== @@ -183,6 +209,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: std::any::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 +353,9 @@ 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, and `table_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 +399,20 @@ pub trait Read<'a, 'db: 'a>: Core<'a, 'db> + RegistrySealed<'a, 'db> { Ok(action.lookup(self.es(), &key_values).is_some()) } + /// The declared sorts of a table's columns, or `None` if no table with + /// that name is registered. Pair with [`Read::table_sizes`] to walk every + /// table, and with [`Read::table_subtype`] to tell a constructor's eclass + /// column from a function's output column. + fn table_schema(&self, name: &str) -> Option<&'a ResolvedSchema> { + self.schemas().get(name).map(|(_, schema)| schema) + } + + /// Whether the named table is a `constructor` or a `function`, or `None` + /// if no table with that name is registered. + fn table_subtype(&self, name: &str) -> Option { + self.schemas().get(name).map(|(subtype, _)| *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 +459,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`]. @@ -626,6 +718,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) schemas: &'a FunctionSchemas, pub(crate) ctx: Context, } @@ -640,6 +733,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) schemas: &'a FunctionSchemas, pub(crate) ctx: Context, } @@ -654,6 +748,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) schemas: &'a FunctionSchemas, pub(crate) ctx: Context, } @@ -670,11 +765,13 @@ impl<'a, 'db: 'a> ReadState<'a, 'db> { pub(crate) fn wrap( es: &'a mut ExecutionState<'db>, registry: &'a ActionRegistry, + schemas: &'a FunctionSchemas, ctx: Context, ) -> Self { Self { inner: es, registry, + schemas, ctx, } } @@ -687,11 +784,13 @@ impl<'a, 'db: 'a> WriteState<'a, 'db> { pub(crate) fn wrap( es: &'a mut ExecutionState<'db>, registry: &'a ActionRegistry, + schemas: &'a FunctionSchemas, ctx: Context, ) -> Self { Self { inner: es, registry, + schemas, ctx, } } @@ -704,11 +803,13 @@ impl<'a, 'db: 'a> FullState<'a, 'db> { pub(crate) fn wrap( es: &'a mut ExecutionState<'db>, registry: &'a ActionRegistry, + schemas: &'a FunctionSchemas, ctx: Context, ) -> Self { Self { inner: es, registry, + schemas, ctx, } } @@ -753,6 +854,9 @@ impl<'a, 'db: 'a> RegistrySealed<'a, 'db> for ReadState<'a, 'db> { fn registry(&self) -> &ActionRegistry { self.registry } + fn schemas(&self) -> &'a FunctionSchemas { + self.schemas + } } impl<'a, 'db: 'a> Core<'a, 'db> for ReadState<'a, 'db> {} impl<'a, 'db: 'a> Read<'a, 'db> for ReadState<'a, 'db> {} @@ -775,6 +879,9 @@ impl<'a, 'db: 'a> RegistrySealed<'a, 'db> for WriteState<'a, 'db> { fn registry(&self) -> &ActionRegistry { self.registry } + fn schemas(&self) -> &'a FunctionSchemas { + self.schemas + } } impl<'a, 'db: 'a> Core<'a, 'db> for WriteState<'a, 'db> {} impl<'a, 'db: 'a> Write<'a, 'db> for WriteState<'a, 'db> {} @@ -797,6 +904,9 @@ impl<'a, 'db: 'a> RegistrySealed<'a, 'db> for FullState<'a, 'db> { fn registry(&self) -> &ActionRegistry { self.registry } + fn schemas(&self) -> &'a FunctionSchemas { + self.schemas + } } 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/lib.rs b/src/lib.rs index abbd2ed62..9d26e17ad 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -44,7 +44,8 @@ use egglog_core_relations as core_relations; use egglog_numeric_id as numeric_id; use egglog_reports::{ReportLevel, RunReport}; pub use exec_state::{ - Context, Core, Enode, FullState, FunctionEntry, PureState, Read, ReadState, Write, WriteState, + Context, Core, Enode, FullState, FunctionEntry, FunctionSchemas, PureState, Read, ReadState, + Write, WriteState, }; use extract::{DefaultCost, Extractor, TreeAdditiveCostModel}; use indexmap::map::Entry; @@ -68,7 +69,7 @@ use std::io::{Read as _, Write as _}; use std::iter::once; use std::ops::Deref; use std::path::PathBuf; -use std::sync::Arc; +use std::sync::{Arc, RwLock}; pub use termdag::{OrdTerm, Term, TermDag, TermId}; use thiserror::Error; pub use typechecking::PrimitiveValidator; @@ -307,6 +308,11 @@ pub struct EGraph { proof_state: EncodingState, /// In proof mode, this is the program before proof instrumentation and the version we use for proof checking. proof_check_program: Vec, + /// Declared column sorts by table name. Shared (via `Arc>`) + /// with the state wrappers, which cannot reach [`TypeInfo`] from a + /// primitive body, and read through [`Read::table_schema`]. `push`/`pop` + /// snapshot and restore its contents, so a popped table stops resolving. + pub(crate) function_schemas: Arc>, } /// A user-defined command allows users to inject custom command that can be called @@ -416,6 +422,7 @@ impl Default for EGraph { command_macros: Default::default(), proof_state, proof_check_program: vec![], + function_schemas: Default::default(), }; add_base_sort(&mut eg, UnitSort, span!()).unwrap(); add_base_sort(&mut eg, StringSort, span!()).unwrap(); @@ -714,6 +721,11 @@ impl EGraph { let prev_prev: Option> = self.pushed_egraph.take(); let mut prev = self.clone(); prev.pushed_egraph = prev_prev; + // The live schemas are shared with the registered primitives, so the + // saved copy gets its own cell rather than a handle to the one that + // keeps being appended to. + prev.function_schemas = + Arc::new(RwLock::new(self.function_schemas.read().unwrap().clone())); self.pushed_egraph = Some(Box::new(prev)); } @@ -729,7 +741,13 @@ 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 schemas into the cell the registered + // primitives hold, rather than swapping in the saved cell they + // have no handle to. + let live_schemas = self.function_schemas.clone(); + *live_schemas.write().unwrap() = e.function_schemas.read().unwrap().clone(); *self = *e; + self.function_schemas = live_schemas; Ok(()) } None => Err(Error::Pop(span!())), @@ -843,9 +861,15 @@ impl EGraph { can_subsume, }); + let schema = ResolvedSchema { input, output }; + self.function_schemas + .write() + .unwrap() + .declare(&decl.name, decl.subtype, schema.clone()); + let function = Function { decl: decl.clone(), - schema: ResolvedSchema { input, output }, + schema, can_subsume, backend_id, }; @@ -1231,8 +1255,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 schemas = self.function_schemas.read().unwrap(); self.backend - .with_execution_state_tracked(|es| f(ReadState::wrap(es, &guard, Context::Read))) + .with_execution_state_tracked(|es| { + f(ReadState::wrap(es, &guard, &schemas, Context::Read)) + }) .0 } @@ -2379,9 +2406,12 @@ 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 schemas = self.function_schemas.clone(); + let schemas = schemas.read().unwrap(); + let (result, changed) = self.backend.with_execution_state_tracked(|es| { + f(FullState::wrap(es, &guard, &schemas, Context::Full)) + }); + drop(schemas); 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/typechecking.rs b/src/typechecking.rs index 28031a56b..45c1925a5 100644 --- a/src/typechecking.rs +++ b/src/typechecking.rs @@ -1,6 +1,7 @@ use std::hash::Hasher; use crate::Context; +use crate::FunctionSchemas; use crate::proofs::proof_container_rebuild::register_container_rebuild_from_spec; use crate::{ core::{CoreActionContext, CoreRule, GenericActionsExt, ResolvedCall}, @@ -44,6 +45,7 @@ impl ExternalFunction for PurePrimWrapper { struct RegistryPrimWrapper { prim: T, registry: Arc>, + schemas: Arc>, /// Stamped onto the state wrapper. ctx: Context, _wrap: std::marker::PhantomData S>, @@ -56,6 +58,7 @@ trait RegistryWrap: Clone + Send + Sync { ctx: Context, args: &[Value], registry: &ActionRegistry, + schemas: &FunctionSchemas, ) -> Option; } @@ -69,8 +72,9 @@ impl RegistryWrap for WrapRead { ctx: Context, args: &[Value], registry: &ActionRegistry, + schemas: &FunctionSchemas, ) -> Option { - prim.apply(ReadState::wrap(exec_state, registry, ctx), args) + prim.apply(ReadState::wrap(exec_state, registry, schemas, ctx), args) } } #[derive(Clone)] @@ -83,8 +87,9 @@ impl RegistryWrap for WrapWrite { ctx: Context, args: &[Value], registry: &ActionRegistry, + schemas: &FunctionSchemas, ) -> Option { - prim.apply(WriteState::wrap(exec_state, registry, ctx), args) + prim.apply(WriteState::wrap(exec_state, registry, schemas, ctx), args) } } #[derive(Clone)] @@ -97,8 +102,9 @@ impl RegistryWrap for WrapFull { ctx: Context, args: &[Value], registry: &ActionRegistry, + schemas: &FunctionSchemas, ) -> Option { - prim.apply(FullState::wrap(exec_state, registry, ctx), args) + prim.apply(FullState::wrap(exec_state, registry, schemas, ctx), args) } } @@ -107,7 +113,8 @@ 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) + let schemas = self.schemas.read().unwrap(); + S::invoke(&self.prim, exec_state, self.ctx, args, ®istry, &schemas) } } @@ -328,10 +335,12 @@ impl EGraph { S: RegistryWrap + 'static, { let registry = self.backend.action_registry().clone(); + let schemas = self.function_schemas.clone(); self.register_per_context(x, validator, valid_ctxs, move |x, ctx| { Box::new(RegistryPrimWrapper:: { prim: x, registry: registry.clone(), + schemas: schemas.clone(), ctx, _wrap: std::marker::PhantomData, }) From 946b4303b528f7756f468759c381c7855d39399d Mon Sep 17 00:00:00 2001 From: Oliver Flatt Date: Thu, 6 Aug 2026 20:32:22 +0000 Subject: [PATCH 2/8] Split table_schema by subtype, and tidy the diff's docs table_schema returned a schema for either subtype, which reads past the fact that a constructor's last column is an e-class and a function's is an output. Splitting it matches how the rest of Read and Write already work - lookup / eclass_of, constructor_enodes / function_entries, set / add - so it is now constructor_schema and function_schema, each erroring with WrongSubtype on a mismatch. table_subtype stays as the error-free predicate to dispatch on when either subtype is acceptable, which is what retires the subtype probe in egglog-experimental. Also applies the tidy-diff-docs skill to the comments this branch adds, and records a new rule in that skill: import a type rather than naming it by an inline full path, which is what FunctionSchemas was doing with crate::util::HashMap. Co-Authored-By: Claude Opus 5 --- .claude/skills/tidy-diff-docs/SKILL.md | 11 ++++ CHANGELOG.md | 2 +- src/exec_state.rs | 80 +++++++++++++++++++------- src/lib.rs | 4 +- 4 files changed, 74 insertions(+), 23 deletions(-) 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 fc7e3ed89..ff5baf226 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ ## [Unreleased] - ReleaseDate -- **Enough e-graph introspection for a primitive to walk and rebuild a sub-e-graph.** Three additions, together sufficient to implement substitution as an out-of-tree primitive (see `unstable-subst` in `egglog-experimental`): `Read::enodes_for_eclass(name, eclass, f)` walks a constructor's rows by output e-class through the backend's column index instead of scanning the table (with `ExecutionState::for_each_matching_col` and `TableAction::for_each_output_value` behind it); `Read::table_schema(name)` and `Read::table_subtype(name)` expose a table's declared column sorts and subtype to a primitive body, which cannot reach `TypeInfo` (`table_subtype` also replaces probing a table's subtype by starting a scan); and `Core::rebuild_container(type_id, value, remap)` remaps a container value's contents and interns the result, which out-of-tree code cannot do through `register_container` because it cannot name an arbitrary container's Rust type. +- **Enough e-graph introspection for a primitive to walk and rebuild a sub-e-graph.** Three additions, together sufficient to implement substitution as an out-of-tree primitive (see `unstable-subst` in `egglog-experimental`): `Read::enodes_for_eclass(name, eclass, f)` walks a constructor's rows by output e-class through the backend's column index instead of scanning the table (with `ExecutionState::for_each_matching_col` and `TableAction::for_each_output_value` behind it); `Read::constructor_schema(name)`, `Read::function_schema(name)`, and `Read::table_subtype(name)` expose a table's declared column sorts and subtype to a primitive body, which cannot reach `TypeInfo` (`table_subtype` also replaces probing a table's subtype by starting a scan); and `Core::rebuild_container(type_id, value, remap)` remaps a container value's contents and interns the result, which out-of-tree code cannot do through `register_container` because it cannot name an arbitrary container's Rust type. - 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. diff --git a/src/exec_state.rs b/src/exec_state.rs index 192eebb0c..64fe29205 100644 --- a/src/exec_state.rs +++ b/src/exec_state.rs @@ -30,6 +30,7 @@ //! [`ReadPrim`]: crate::ReadPrim //! [`FullPrim`]: crate::FullPrim +use std::any::TypeId; use std::ops::Deref; use crate::Error; @@ -44,6 +45,7 @@ use crate::{ core::ResolvedCall, sort::{F, S}, typechecking::FuncType, + util::HashMap, }; use egglog_bridge::{ActionRegistry, TableAction, TableKind}; use smallvec::SmallVec; @@ -85,16 +87,12 @@ impl Context { pub const ALL: [Context; 4] = [Context::Pure, Context::Write, Context::Read, Context::Full]; } -/// The declared column sorts of every table, by name. -/// -/// The [`ActionRegistry`] resolves a table name to the handle that reads and -/// writes it, but it lives in the backend and so knows nothing about egglog -/// sorts. This is the matching name-indexed view of the *types*, kept by the -/// e-graph and shared with the state wrappers so a primitive body — which -/// cannot reach `TypeInfo` — can still ask what a column holds. +/// The declared column sorts of every table, by name. Read through +/// [`Read::constructor_schema`], [`Read::function_schema`], and +/// [`Read::table_subtype`]. #[derive(Clone, Default)] pub struct FunctionSchemas { - tables: crate::util::HashMap, + tables: HashMap, } impl FunctionSchemas { @@ -165,8 +163,8 @@ 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; - /// Borrowed for `'a` rather than for `&self`, so a caller can hold a - /// column sort while it goes on to write through the same state. + /// Borrowed for `'a`, not for `&self`, so a caller can hold a column sort + /// across a write through the same state. fn schemas(&self) -> &'a FunctionSchemas; } @@ -220,7 +218,7 @@ pub trait Core<'a, 'db: 'a>: Internal<'a, 'db> { /// unchanged for those. fn rebuild_container( &mut self, - type_id: std::any::TypeId, + type_id: TypeId, value: Value, remap: &(dyn Fn(Value) -> Value + Send + Sync), ) -> Value { @@ -354,8 +352,9 @@ pub trait Core<'a, 'db: 'a>: Internal<'a, 'db> { /// return `None` if absent — never insert. The iteration / /// introspection methods (`function_entries`, `constructor_enodes`, /// `enodes_for_eclass`, `table_size`, `table_sizes`) walk the current -/// contents of the database, and `table_schema` / `table_subtype` -/// report how a table is declared rather than what it holds. +/// 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 @@ -399,16 +398,28 @@ pub trait Read<'a, 'db: 'a>: Core<'a, 'db> + RegistrySealed<'a, 'db> { Ok(action.lookup(self.es(), &key_values).is_some()) } - /// The declared sorts of a table's columns, or `None` if no table with - /// that name is registered. Pair with [`Read::table_sizes`] to walk every - /// table, and with [`Read::table_subtype`] to tell a constructor's eclass - /// column from a function's output column. - fn table_schema(&self, name: &str) -> Option<&'a ResolvedSchema> { - self.schemas().get(name).map(|(_, schema)| schema) + /// The declared sorts of a constructor's columns: its inputs followed by + /// 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<&'a ResolvedSchema, Error> { + schema_of(self.schemas(), name, FunctionSubtype::Constructor) + } + + /// The declared sorts of a function's columns: its inputs followed by the + /// output column. + /// + /// **Only valid for `function` tables.** Constructors error; use + /// [`Read::constructor_schema`] for those. + fn function_schema(&self, name: &str) -> Result<&'a ResolvedSchema, Error> { + schema_of(self.schemas(), name, FunctionSubtype::Custom) } /// Whether the named table is a `constructor` or a `function`, or `None` - /// if no table with that name is registered. + /// 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 { self.schemas().get(name).map(|(subtype, _)| *subtype) } @@ -633,6 +644,35 @@ 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 schema_of<'a>( + schemas: &'a FunctionSchemas, + name: &str, + expected: FunctionSubtype, +) -> Result<&'a ResolvedSchema, Error> { + let Some((subtype, schema)) = schemas.get(name) else { + return Err(ApiError::MissingTable { + name: name.to_string(), + } + .into()); + }; + if *subtype != expected { + return Err(ApiError::WrongSubtype { + name: name.to_string(), + expected: subtype_label(expected), + actual: subtype_label(*subtype), + } + .into()); + } + Ok(schema) +} + fn lookup_action(registry: &ActionRegistry, name: &str) -> Result { registry.lookup_table(name).cloned().ok_or_else(|| { ApiError::MissingTable { diff --git a/src/lib.rs b/src/lib.rs index 9d26e17ad..62abcae34 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -310,8 +310,8 @@ pub struct EGraph { proof_check_program: Vec, /// Declared column sorts by table name. Shared (via `Arc>`) /// with the state wrappers, which cannot reach [`TypeInfo`] from a - /// primitive body, and read through [`Read::table_schema`]. `push`/`pop` - /// snapshot and restore its contents, so a popped table stops resolving. + /// primitive body. `push`/`pop` snapshot and restore its contents, so a + /// popped table stops resolving. pub(crate) function_schemas: Arc>, } From 765e83281a1922e22dbbc1d35d2cadc9bbfd1463 Mon Sep 17 00:00:00 2001 From: Oliver Flatt Date: Thu, 6 Aug 2026 21:34:50 +0000 Subject: [PATCH 3/8] Store a function's signature once, instead of adding a third copy The first cut of this branch added a FunctionSchemas registry so a primitive body could see column sorts. That was a third copy of data egglog already kept twice: TypeInfo::func_types holds a FuncType {name, subtype, input, output}, and Function holds the same content again as a ResolvedSchema plus decl.subtype (and a third time, unresolved, in decl.schema). What was actually missing was not the data but a way to reach it from a state wrapper. So there is now one store. TypeInfo::func_types becomes the shared cell - Arc>>> - and the state wrappers hold a handle to it, which is what backs constructor_schema / function_schema / table_subtype. FunctionSchemas is gone. Function points at the same Arc rather than storing its own copy, so Function::schema() -> &ResolvedSchema becomes Function::func_type() -> &FuncType and ResolvedSchema is removed, its get_by_pos moving to FuncType. declare_function reuses the signature typechecking already resolved instead of resolving the sorts a second time; it still resolves and records the functions desugaring generates (global bindings, proof tables), which never go through typechecking. Two things this had to get right. TypeInfo::clone deep-copies the map: a clone is an independent e-graph - a pushed copy, or the parallel typechecking the proof checker keeps - and declaring a function in one must not make it resolve in the other. And pop restores the pushed contents into the live cell the registered primitives already hold, rather than swapping in a cell they have no handle to. Also drops the second RwLock acquisition per primitive invocation the first cut introduced: the wrapper holds the unlocked handle and the schema accessors lock only when called, so a primitive that never asks pays nothing. Co-Authored-By: Claude Opus 5 --- CHANGELOG.md | 3 +- src/core.rs | 4 +- src/exec_state.rs | 92 +++++++++-------------- src/extract.rs | 32 ++++---- src/lib.rs | 129 ++++++++++++++------------------ src/proofs/proof_extraction.rs | 4 +- src/proofs/proof_extractor.rs | 6 +- src/serialize.rs | 10 +-- src/typechecking.rs | 132 +++++++++++++++++++++++++++------ 9 files changed, 233 insertions(+), 179 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ff5baf226..69a068b59 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,8 @@ ## [Unreleased] - ReleaseDate -- **Enough e-graph introspection for a primitive to walk and rebuild a sub-e-graph.** Three additions, together sufficient to implement substitution as an out-of-tree primitive (see `unstable-subst` in `egglog-experimental`): `Read::enodes_for_eclass(name, eclass, f)` walks a constructor's rows by output e-class through the backend's column index instead of scanning the table (with `ExecutionState::for_each_matching_col` and `TableAction::for_each_output_value` behind it); `Read::constructor_schema(name)`, `Read::function_schema(name)`, and `Read::table_subtype(name)` expose a table's declared column sorts and subtype to a primitive body, which cannot reach `TypeInfo` (`table_subtype` also replaces probing a table's subtype by starting a scan); and `Core::rebuild_container(type_id, value, remap)` remaps a container value's contents and interns the result, which out-of-tree code cannot do through `register_container` because it cannot name an arbitrary container's Rust type. +- **Enough e-graph introspection for a primitive to walk and rebuild a sub-e-graph**, and one home for a function's signature. `Read::enodes_for_eclass(name, eclass, f)` walks a constructor's rows by output e-class through the backend's column index instead of scanning the table (with `ExecutionState::for_each_matching_col` and `TableAction::for_each_output_value` behind it). `Read::constructor_schema(name)`, `Read::function_schema(name)`, and `Read::table_subtype(name)` expose a table's declared signature to a primitive body, which cannot reach `TypeInfo` (`table_subtype` also replaces probing a table's subtype by starting a scan). `Core::rebuild_container(type_id, value, remap)` remaps a container value's contents and interns the result, which out-of-tree code cannot do through `register_container` because it cannot name an arbitrary container's Rust type. Together these are what an out-of-tree substitution primitive needs; see `unstable-subst` in `egglog-experimental`. + **Breaking:** a signature is now stored once. `TypeInfo::func_types` holds every declared function's `FuncType` and is shared rather than copied, so `TypeInfo::get_func_type` returns `Option>` instead of `Option<&FuncType>`; `Function` points at the same `FuncType` and `Function::schema() -> &ResolvedSchema` becomes `Function::func_type() -> &FuncType`; `ResolvedSchema` is removed, with its `get_by_pos` moving to `FuncType`. `declare_function` no longer re-resolves the sorts typechecking already resolved. - 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. 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 64fe29205..3dcbb304a 100644 --- a/src/exec_state.rs +++ b/src/exec_state.rs @@ -32,6 +32,7 @@ use std::any::TypeId; use std::ops::Deref; +use std::sync::Arc; use crate::Error; use crate::api::{ApiError, IntoValue, IntoValues}; @@ -40,12 +41,10 @@ use crate::core_relations::{ Value, }; use crate::{ - ResolvedSchema, ast::{FunctionSubtype, Literal, ResolvedExpr}, core::ResolvedCall, sort::{F, S}, - typechecking::FuncType, - util::HashMap, + typechecking::{FuncType, SharedFuncTypes}, }; use egglog_bridge::{ActionRegistry, TableAction, TableKind}; use smallvec::SmallVec; @@ -87,24 +86,6 @@ impl Context { pub const ALL: [Context; 4] = [Context::Pure, Context::Write, Context::Read, Context::Full]; } -/// The declared column sorts of every table, by name. Read through -/// [`Read::constructor_schema`], [`Read::function_schema`], and -/// [`Read::table_subtype`]. -#[derive(Clone, Default)] -pub struct FunctionSchemas { - tables: HashMap, -} - -impl FunctionSchemas { - pub(crate) fn declare(&mut self, name: &str, subtype: FunctionSubtype, schema: ResolvedSchema) { - self.tables.insert(name.to_owned(), (subtype, schema)); - } - - fn get(&self, name: &str) -> Option<&(FunctionSubtype, ResolvedSchema)> { - self.tables.get(name) - } -} - // ===================================================================== // Sealed traits. // @@ -163,9 +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; - /// Borrowed for `'a`, not for `&self`, so a caller can hold a column sort - /// across a write through the same state. - fn schemas(&self) -> &'a FunctionSchemas; + /// 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; } // ===================================================================== @@ -398,22 +379,21 @@ pub trait Read<'a, 'db: 'a>: Core<'a, 'db> + RegistrySealed<'a, 'db> { Ok(action.lookup(self.es(), &key_values).is_some()) } - /// The declared sorts of a constructor's columns: its inputs followed by - /// the eclass column. Pair with [`Read::table_sizes`] to walk every table. + /// 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<&'a ResolvedSchema, Error> { - schema_of(self.schemas(), name, FunctionSubtype::Constructor) + fn constructor_schema(&self, name: &str) -> Result, Error> { + func_type_of(self.func_types(), name, FunctionSubtype::Constructor) } - /// The declared sorts of a function's columns: its inputs followed by the - /// output column. + /// 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<&'a ResolvedSchema, Error> { - schema_of(self.schemas(), name, FunctionSubtype::Custom) + 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` @@ -421,7 +401,8 @@ pub trait Read<'a, 'db: 'a>: Core<'a, 'db> + RegistrySealed<'a, 'db> { /// 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 { - self.schemas().get(name).map(|(subtype, _)| *subtype) + 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 @@ -651,26 +632,27 @@ fn subtype_label(subtype: FunctionSubtype) -> &'static str { } } -fn schema_of<'a>( - schemas: &'a FunctionSchemas, +fn func_type_of( + func_types: &SharedFuncTypes, name: &str, expected: FunctionSubtype, -) -> Result<&'a ResolvedSchema, Error> { - let Some((subtype, schema)) = schemas.get(name) else { +) -> 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 *subtype != expected { + if func_type.subtype != expected { return Err(ApiError::WrongSubtype { name: name.to_string(), expected: subtype_label(expected), - actual: subtype_label(*subtype), + actual: subtype_label(func_type.subtype), } .into()); } - Ok(schema) + Ok(func_type) } fn lookup_action(registry: &ActionRegistry, name: &str) -> Result { @@ -758,7 +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) schemas: &'a FunctionSchemas, + pub(crate) func_types: &'a SharedFuncTypes, pub(crate) ctx: Context, } @@ -773,7 +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) schemas: &'a FunctionSchemas, + pub(crate) func_types: &'a SharedFuncTypes, pub(crate) ctx: Context, } @@ -788,7 +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) schemas: &'a FunctionSchemas, + pub(crate) func_types: &'a SharedFuncTypes, pub(crate) ctx: Context, } @@ -805,13 +787,13 @@ impl<'a, 'db: 'a> ReadState<'a, 'db> { pub(crate) fn wrap( es: &'a mut ExecutionState<'db>, registry: &'a ActionRegistry, - schemas: &'a FunctionSchemas, + func_types: &'a SharedFuncTypes, ctx: Context, ) -> Self { Self { inner: es, registry, - schemas, + func_types, ctx, } } @@ -824,13 +806,13 @@ impl<'a, 'db: 'a> WriteState<'a, 'db> { pub(crate) fn wrap( es: &'a mut ExecutionState<'db>, registry: &'a ActionRegistry, - schemas: &'a FunctionSchemas, + func_types: &'a SharedFuncTypes, ctx: Context, ) -> Self { Self { inner: es, registry, - schemas, + func_types, ctx, } } @@ -843,13 +825,13 @@ impl<'a, 'db: 'a> FullState<'a, 'db> { pub(crate) fn wrap( es: &'a mut ExecutionState<'db>, registry: &'a ActionRegistry, - schemas: &'a FunctionSchemas, + func_types: &'a SharedFuncTypes, ctx: Context, ) -> Self { Self { inner: es, registry, - schemas, + func_types, ctx, } } @@ -894,8 +876,8 @@ impl<'a, 'db: 'a> RegistrySealed<'a, 'db> for ReadState<'a, 'db> { fn registry(&self) -> &ActionRegistry { self.registry } - fn schemas(&self) -> &'a FunctionSchemas { - self.schemas + fn func_types(&self) -> &'a SharedFuncTypes { + self.func_types } } impl<'a, 'db: 'a> Core<'a, 'db> for ReadState<'a, 'db> {} @@ -919,8 +901,8 @@ impl<'a, 'db: 'a> RegistrySealed<'a, 'db> for WriteState<'a, 'db> { fn registry(&self) -> &ActionRegistry { self.registry } - fn schemas(&self) -> &'a FunctionSchemas { - self.schemas + fn func_types(&self) -> &'a SharedFuncTypes { + self.func_types } } impl<'a, 'db: 'a> Core<'a, 'db> for WriteState<'a, 'db> {} @@ -944,8 +926,8 @@ impl<'a, 'db: 'a> RegistrySealed<'a, 'db> for FullState<'a, 'db> { fn registry(&self) -> &ActionRegistry { self.registry } - fn schemas(&self) -> &'a FunctionSchemas { - self.schemas + fn func_types(&self) -> &'a SharedFuncTypes { + self.func_types } } impl<'a, 'db: 'a> Core<'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 62abcae34..5475a2185 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -44,8 +44,7 @@ use egglog_core_relations as core_relations; use egglog_numeric_id as numeric_id; use egglog_reports::{ReportLevel, RunReport}; pub use exec_state::{ - Context, Core, Enode, FullState, FunctionEntry, FunctionSchemas, PureState, Read, ReadState, - Write, WriteState, + Context, Core, Enode, FullState, FunctionEntry, PureState, Read, ReadState, Write, WriteState, }; use extract::{DefaultCost, Extractor, TreeAdditiveCostModel}; use indexmap::map::Entry; @@ -69,9 +68,10 @@ use std::io::{Read as _, Write as _}; use std::iter::once; use std::ops::Deref; use std::path::PathBuf; -use std::sync::{Arc, RwLock}; +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; @@ -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.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 { @@ -308,11 +308,6 @@ pub struct EGraph { proof_state: EncodingState, /// In proof mode, this is the program before proof instrumentation and the version we use for proof checking. proof_check_program: Vec, - /// Declared column sorts by table name. Shared (via `Arc>`) - /// with the state wrappers, which cannot reach [`TypeInfo`] from a - /// primitive body. `push`/`pop` snapshot and restore its contents, so a - /// popped table stops resolving. - pub(crate) function_schemas: Arc>, } /// A user-defined command allows users to inject custom command that can be called @@ -332,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, } @@ -343,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. @@ -372,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() } } @@ -422,7 +401,6 @@ impl Default for EGraph { command_macros: Default::default(), proof_state, proof_check_program: vec![], - function_schemas: Default::default(), }; add_base_sort(&mut eg, UnitSort, span!()).unwrap(); add_base_sort(&mut eg, StringSort, span!()).unwrap(); @@ -721,11 +699,6 @@ impl EGraph { let prev_prev: Option> = self.pushed_egraph.take(); let mut prev = self.clone(); prev.pushed_egraph = prev_prev; - // The live schemas are shared with the registered primitives, so the - // saved copy gets its own cell rather than a handle to the one that - // keeps being appended to. - prev.function_schemas = - Arc::new(RwLock::new(self.function_schemas.read().unwrap().clone())); self.pushed_egraph = Some(Box::new(prev)); } @@ -741,13 +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 schemas into the cell the registered + // 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_schemas = self.function_schemas.clone(); - *live_schemas.write().unwrap() = e.function_schemas.read().unwrap().clone(); + let live_func_types = self.type_info.func_types_handle().clone(); + let restored = e.type_info.func_types_handle().read().unwrap().clone(); *self = *e; - self.function_schemas = live_schemas; + *live_func_types.write().unwrap() = restored; + self.type_info.set_func_types_handle(live_func_types); Ok(()) } None => Err(Error::Pop(span!())), @@ -817,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 @@ -841,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 { @@ -861,15 +850,9 @@ impl EGraph { can_subsume, }); - let schema = ResolvedSchema { input, output }; - self.function_schemas - .write() - .unwrap() - .declare(&decl.name, decl.subtype, schema.clone()); - let function = Function { decl: decl.clone(), - schema, + func_type, can_subsume, backend_id, }; @@ -1255,10 +1238,10 @@ 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 schemas = self.function_schemas.read().unwrap(); + let func_types = self.type_info.func_types_handle(); self.backend .with_execution_state_tracked(|es| { - f(ReadState::wrap(es, &guard, &schemas, Context::Read)) + f(ReadState::wrap(es, &guard, func_types, Context::Read)) }) .0 } @@ -1901,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())), @@ -1909,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())), } @@ -1925,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:?}"); @@ -2406,12 +2389,10 @@ impl EGraph { ) -> Result { let registry = self.backend.action_registry().clone(); let guard = registry.read().unwrap(); - let schemas = self.function_schemas.clone(); - let schemas = schemas.read().unwrap(); + 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, &schemas, Context::Full)) + f(FullState::wrap(es, &guard, &func_types, Context::Full)) }); - drop(schemas); 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 45c1925a5..3c479afe9 100644 --- a/src/typechecking.rs +++ b/src/typechecking.rs @@ -1,7 +1,6 @@ use std::hash::Hasher; use crate::Context; -use crate::FunctionSchemas; use crate::proofs::proof_container_rebuild::register_container_rebuild_from_spec; use crate::{ core::{CoreActionContext, CoreRule, GenericActionsExt, ResolvedCall}, @@ -45,7 +44,7 @@ impl ExternalFunction for PurePrimWrapper { struct RegistryPrimWrapper { prim: T, registry: Arc>, - schemas: Arc>, + func_types: SharedFuncTypes, /// Stamped onto the state wrapper. ctx: Context, _wrap: std::marker::PhantomData S>, @@ -58,7 +57,7 @@ trait RegistryWrap: Clone + Send + Sync { ctx: Context, args: &[Value], registry: &ActionRegistry, - schemas: &FunctionSchemas, + func_types: &SharedFuncTypes, ) -> Option; } @@ -72,9 +71,9 @@ impl RegistryWrap for WrapRead { ctx: Context, args: &[Value], registry: &ActionRegistry, - schemas: &FunctionSchemas, + func_types: &SharedFuncTypes, ) -> Option { - prim.apply(ReadState::wrap(exec_state, registry, schemas, ctx), args) + prim.apply(ReadState::wrap(exec_state, registry, func_types, ctx), args) } } #[derive(Clone)] @@ -87,9 +86,12 @@ impl RegistryWrap for WrapWrite { ctx: Context, args: &[Value], registry: &ActionRegistry, - schemas: &FunctionSchemas, + func_types: &SharedFuncTypes, ) -> Option { - prim.apply(WriteState::wrap(exec_state, registry, schemas, ctx), args) + prim.apply( + WriteState::wrap(exec_state, registry, func_types, ctx), + args, + ) } } #[derive(Clone)] @@ -102,9 +104,9 @@ impl RegistryWrap for WrapFull { ctx: Context, args: &[Value], registry: &ActionRegistry, - schemas: &FunctionSchemas, + func_types: &SharedFuncTypes, ) -> Option { - prim.apply(FullState::wrap(exec_state, registry, schemas, ctx), args) + prim.apply(FullState::wrap(exec_state, registry, func_types, ctx), args) } } @@ -113,8 +115,14 @@ impl + 'static> ExternalFun { fn invoke(&self, exec_state: &mut ExecutionState, args: &[Value]) -> Option { let registry = self.registry.read().unwrap(); - let schemas = self.schemas.read().unwrap(); - S::invoke(&self.prim, exec_state, self.ctx, args, ®istry, &schemas) + S::invoke( + &self.prim, + exec_state, + self.ctx, + args, + ®istry, + &self.func_types, + ) } } @@ -126,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 @@ -159,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. @@ -211,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 { @@ -244,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)); } @@ -335,12 +388,12 @@ impl EGraph { S: RegistryWrap + 'static, { let registry = self.backend.action_registry().clone(); - let schemas = self.function_schemas.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(), - schemas: schemas.clone(), + func_types: func_types.clone(), ctx, _wrap: std::marker::PhantomData, }) @@ -581,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 @@ -793,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(), @@ -814,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 @@ -1147,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) } From e424551c2feb442fd5620a3084bc533f7774906f Mon Sep 17 00:00:00 2001 From: Oliver Flatt Date: Thu, 6 Aug 2026 21:42:52 +0000 Subject: [PATCH 4/8] Shorten the changelog entry and fold it into the existing API entries The additions belong with the name-indexed e-graph access entry they extend, not as a standalone block longer than anything else on the list. The signature consolidation keeps a one-line breaking bullet next to the other breaking ones. Co-Authored-By: Claude Opus 5 --- CHANGELOG.md | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 69a068b59..2d5249da0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,12 +2,10 @@ ## [Unreleased] - ReleaseDate -- **Enough e-graph introspection for a primitive to walk and rebuild a sub-e-graph**, and one home for a function's signature. `Read::enodes_for_eclass(name, eclass, f)` walks a constructor's rows by output e-class through the backend's column index instead of scanning the table (with `ExecutionState::for_each_matching_col` and `TableAction::for_each_output_value` behind it). `Read::constructor_schema(name)`, `Read::function_schema(name)`, and `Read::table_subtype(name)` expose a table's declared signature to a primitive body, which cannot reach `TypeInfo` (`table_subtype` also replaces probing a table's subtype by starting a scan). `Core::rebuild_container(type_id, value, remap)` remaps a container value's contents and interns the result, which out-of-tree code cannot do through `register_container` because it cannot name an arbitrary container's Rust type. Together these are what an out-of-tree substitution primitive needs; see `unstable-subst` in `egglog-experimental`. - **Breaking:** a signature is now stored once. `TypeInfo::func_types` holds every declared function's `FuncType` and is shared rather than copied, so `TypeInfo::get_func_type` returns `Option>` instead of `Option<&FuncType>`; `Function` points at the same `FuncType` and `Function::schema() -> &ResolvedSchema` becomes `Function::func_type() -> &FuncType`; `ResolvedSchema` is removed, with its `get_by_pos` moving to `FuncType`. `declare_function` no longer re-resolves the sorts typechecking already resolved. - - 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`). @@ -30,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 From f910d36243a2937724db4c10349dff0d68fbaa84 Mon Sep 17 00:00:00 2001 From: Oliver Flatt Date: Fri, 7 Aug 2026 21:06:18 +0000 Subject: [PATCH 5/8] Let a primitive borrow the e-graph's signatures, instead of locking a handle The schema accessors read TypeInfo through an Arc> the primitive captured when it was registered. That lock was never guarding a race: declaring anything needs &mut EGraph, run_rules already holds that borrow for the whole execution, and nothing on Core/Read/Write can declare. It was there only because external functions are stored as Box, so a wrapper cannot hold a borrow of the e-graph. So change the route rather than the synchronization. ExecutionState now carries an ExternalContext - a borrowed value the caller of an operation makes visible to the external functions it reaches - and egglog passes &TypeInfo into run_rules and with_execution_state*. The state wrappers read it back with downcast_ref, so `constructor_schema` hands out a &FuncType borrowed from the e-graph itself. With that, the invariant is enforced by the borrow checker instead of stood in for by a lock: the &TypeInfo is live for exactly the operation that supplied it, so the e-graph cannot be mutably borrowed to declare a function while a rule is running. TypeInfo::func_types goes back to a plain owned map, so get_func_type keeps its original Option<&FuncType> signature and costs a map lookup again; TypeInfo::clone goes back to derive; and push/pop restore the map like every other field, with no cell to reattach. Rebuild paths pass None: they run generated rebuild rules whose actions evaluate in the Write context, which has no schema accessors. The one lock left on the invocation path is the pre-existing ActionRegistry read in RegistryPrimWrapper::invoke, which this does not touch. Co-Authored-By: Claude Opus 5 --- CHANGELOG.md | 5 +- core-relations/benches/writes.rs | 2 +- core-relations/src/action/mod.rs | 20 ++++ core-relations/src/containers/tests.rs | 6 +- core-relations/src/free_join/execute.rs | 10 +- core-relations/src/free_join/mod.rs | 26 ++++- .../src/hash_index/bench_support.rs | 2 +- core-relations/src/lib.rs | 2 +- core-relations/src/tests.rs | 35 ++++--- egglog-bridge/examples/ac.rs | 12 ++- egglog-bridge/examples/math.rs | 2 +- egglog-bridge/src/lib.rs | 53 +++++++--- egglog-bridge/src/tests.rs | 95 +++++++++++------- src/core.rs | 4 +- src/exec_state.rs | 57 +++++------ src/lib.rs | 68 +++++++------ src/scheduler.rs | 36 +++---- src/typechecking.rs | 96 ++++--------------- 18 files changed, 285 insertions(+), 246 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2d5249da0..e1d261174 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,8 @@ - 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:** a function's signature is stored once, shared between `TypeInfo` and `Function` instead of copied into both. `Function::schema() -> &ResolvedSchema` becomes `Function::func_type() -> &FuncType`, and `ResolvedSchema` is removed (its `get_by_pos` moves to `FuncType`). +- **Breaking:** an operation that runs external functions now takes an `ExternalContext` — a borrowed value it makes visible to them, which `ExecutionState::external_context` reads back. `Database::run_rule_set` / `with_execution_state` / `with_execution_state_tracked` and the bridge's `EGraph::run_rules` / `with_execution_state*` take one; pass `None` if there is nothing to share. egglog passes its `&TypeInfo`, which is how a primitive resolves a signature without the e-graph handing out a shared, lock-guarded copy: the borrow is live for exactly the operation that supplied it, so nothing can be declared while a rule is running. - **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`). @@ -28,7 +29,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`. 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`. +- **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 could not previously see at all), 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/benches/writes.rs b/core-relations/benches/writes.rs index e08689f82..3fbe16a0d 100644 --- a/core-relations/benches/writes.rs +++ b/core-relations/benches/writes.rs @@ -140,7 +140,7 @@ fn bench_workload( } } }); - db.with_execution_state(|es| table.merge(es)); + db.with_execution_state(None, |es| table.merge(es)); } }) }) diff --git a/core-relations/src/action/mod.rs b/core-relations/src/action/mod.rs index f8a92a918..d01089a64 100644 --- a/core-relations/src/action/mod.rs +++ b/core-relations/src/action/mod.rs @@ -3,6 +3,7 @@ //! This allows us to execute the "right-hand-side" of a rule. The //! implementation here is optimized to execute on a batch of rows at a time. use std::{ + any::Any, ops::Deref, sync::{ Arc, @@ -326,6 +327,7 @@ impl PredictedVals { #[derive(Copy, Clone)] pub(crate) struct DbView<'a> { + pub(crate) external_context: ExternalContext<'a>, pub(crate) table_info: &'a DenseIdMap, pub(crate) counters: &'a Counters, pub(crate) external_funcs: &'a ExternalFunctions, @@ -334,6 +336,18 @@ pub(crate) struct DbView<'a> { pub(crate) notification_list: &'a NotificationList, } +/// A borrowed value an embedder can make visible to every [`ExecutionState`] +/// created for one operation, for its [`ExternalFunction`]s to read back with +/// [`ExecutionState::external_context`]. +/// +/// Passing it per operation, rather than storing it in the database, is what +/// lets an external function see borrowed embedder state: the borrow is live +/// for exactly the call that supplied it, so the embedder cannot mutate that +/// state while a rule is running. +/// +/// [`ExternalFunction`]: crate::ExternalFunction +pub type ExternalContext<'a> = Option<&'a (dyn Any + Send + Sync)>; + /// A handle on a database that may be in the process of running a rule. /// /// An ExecutionState grants immutable access to the (much of) the database, and also provides a @@ -466,6 +480,12 @@ impl<'a> ExecutionState<'a> { self.db.external_funcs[func].invoke(self, args) } + /// The value the caller of this operation supplied as its + /// [`ExternalContext`], if any. + pub fn external_context(&self) -> ExternalContext<'a> { + self.db.external_context + } + pub fn inc_counter(&self, ctr: CounterId) -> usize { self.db.counters.inc(ctr) } diff --git a/core-relations/src/containers/tests.rs b/core-relations/src/containers/tests.rs index 54a2f953f..a0e4d453e 100644 --- a/core-relations/src/containers/tests.rs +++ b/core-relations/src/containers/tests.rs @@ -111,7 +111,7 @@ fn racing_inserts() { let env = env.clone(); let db = db.clone(); std::thread::spawn(move || { - db.with_execution_state(|es| { + db.with_execution_state(None, |es| { start.wait(); env.get_or_insert(&cont([1, 2, 3]), es) }) @@ -148,7 +148,7 @@ fn incremental_reinsert_canonicalizes_displaced_outer_id() { ); let container = cont([1, 2, 3]); - db.with_execution_state(|es| { + db.with_execution_state(None, |es| { let old_id = env.get_or_insert(&container, es); let new_id = Value::from_usize(old_id.index() + 1000); let hc = hash_container(&container); @@ -183,7 +183,7 @@ fn nonincremental_dirty_ids_only_include_stable_ids() { }), counter, ); - db.with_execution_state(|es| { + db.with_execution_state(None, |es| { let old_id = env.get_or_insert(&VecContainer(vec![old_inner]), es); let new_id = if outer_id_changes { Value::from_usize(old_id.index() + 1000) diff --git a/core-relations/src/free_join/execute.rs b/core-relations/src/free_join/execute.rs index 39d4b0f2b..ae1143559 100644 --- a/core-relations/src/free_join/execute.rs +++ b/core-relations/src/free_join/execute.rs @@ -10,6 +10,7 @@ use std::{ }; use crate::{ + action::ExternalContext, common::{HashMap, HashSet, IndexMap}, free_join::plan::{JoinStages, MatId, MatScanMode, MatSpec}, numeric_id::{DenseIdMap, IdVec, NumericId}, @@ -430,7 +431,12 @@ impl Prober { } impl Database { - pub fn run_rule_set(&mut self, rule_set: &RuleSet, report_level: ReportLevel) -> RuleSetReport { + pub fn run_rule_set( + &mut self, + rule_set: &RuleSet, + report_level: ReportLevel, + context: ExternalContext<'_>, + ) -> RuleSetReport { if rule_set.plans.is_empty() { return RuleSetReport::default(); } @@ -461,7 +467,7 @@ impl Database { let search_and_apply_timer = Instant::now(); // let mut rule_reports: HashMap>; let mut rule_reports: HashMap, Vec>; - let exec_state = ExecutionState::new(self.read_only_view(), Default::default()); + let exec_state = ExecutionState::new(self.read_only_view_with(context), Default::default()); if parallelize_db_level_op(self.total_size_estimate) { let dash_rule_reports: Arc, Vec>> = Arc::new(DashMap::default()); diff --git a/core-relations/src/free_join/mod.rs b/core-relations/src/free_join/mod.rs index d1f5bf3b0..7ae6e9b04 100644 --- a/core-relations/src/free_join/mod.rs +++ b/core-relations/src/free_join/mod.rs @@ -34,7 +34,7 @@ use crate::{ }; use self::plan::Plan; -use crate::action::ExecutionState; +use crate::action::{ExecutionState, ExternalContext}; pub(crate) mod execute; pub(crate) mod frame_update; @@ -357,7 +357,8 @@ impl Database { pub fn rebuild_containers(&mut self, table_id: TableId) -> ContainerRebuildSummary { let mut containers = mem::take(&mut self.container_values); let table = &self.tables[table_id].table; - let res = self.with_execution_state(|state| containers.rebuild_all(table_id, table, state)); + let res = + self.with_execution_state(None, |state| containers.rebuild_all(table_id, table, state)); self.container_values = containers; res } @@ -444,8 +445,15 @@ impl Database { } /// Run `f` with access to an `ExecutionState` mapped to this database. - pub fn with_execution_state(&self, f: impl FnOnce(&mut ExecutionState) -> R) -> R { - let mut state = ExecutionState::new(self.read_only_view(), Default::default()); + /// + /// `context` is visible to any external function the closure reaches; pass + /// `None` if there is nothing to share. + pub fn with_execution_state( + &self, + context: ExternalContext<'_>, + f: impl FnOnce(&mut ExecutionState) -> R, + ) -> R { + let mut state = ExecutionState::new(self.read_only_view_with(context), Default::default()); f(&mut state) } @@ -454,15 +462,23 @@ impl Database { /// flag to skip a subsequent `merge_all` when the closure was read-only. pub fn with_execution_state_tracked( &self, + context: ExternalContext<'_>, f: impl FnOnce(&mut ExecutionState) -> R, ) -> (R, bool) { - let mut state = ExecutionState::new(self.read_only_view(), Default::default()); + let mut state = ExecutionState::new(self.read_only_view_with(context), Default::default()); let result = f(&mut state); (result, state.changed) } pub(crate) fn read_only_view(&self) -> DbView<'_> { + self.read_only_view_with(None) + } + + /// Like [`Database::read_only_view`], but with an [`ExternalContext`] that + /// every [`ExecutionState`] built from the view will expose. + pub(crate) fn read_only_view_with<'a>(&'a self, context: ExternalContext<'a>) -> DbView<'a> { DbView { + external_context: context, table_info: &self.tables, counters: &self.counters, external_funcs: &self.external_functions, diff --git a/core-relations/src/hash_index/bench_support.rs b/core-relations/src/hash_index/bench_support.rs index a13495fbc..2a12125d0 100644 --- a/core-relations/src/hash_index/bench_support.rs +++ b/core-relations/src/hash_index/bench_support.rs @@ -113,7 +113,7 @@ impl IndexInput { } } let db = Database::default(); - db.with_execution_state(|es| { + db.with_execution_state(None, |es| { table.merge(es); }); IndexInput { diff --git a/core-relations/src/lib.rs b/core-relations/src/lib.rs index 1adf0c094..9b56d6128 100644 --- a/core-relations/src/lib.rs +++ b/core-relations/src/lib.rs @@ -23,7 +23,7 @@ pub(crate) mod uf; #[cfg(test)] mod tests; -pub use action::{ExecutionState, MergeVal, QueryEntry, WriteVal}; +pub use action::{ExecutionState, ExternalContext, MergeVal, QueryEntry, WriteVal}; pub use base_values::{BaseValue, BaseValueId, BaseValuePrinter, BaseValues, Boxed}; pub use common::Value; pub use containers::{ContainerRebuildSummary, ContainerValue, ContainerValueId, ContainerValues}; diff --git a/core-relations/src/tests.rs b/core-relations/src/tests.rs index ec6a14742..6d6e2282d 100644 --- a/core-relations/src/tests.rs +++ b/core-relations/src/tests.rs @@ -122,7 +122,7 @@ fn basic_query_inner() { rules.build_with_description("add"); let rule_set = rsb.build(); - let report = db.run_rule_set(&rule_set, ReportLevel::TimeOnly); + let report = db.run_rule_set(&rule_set, ReportLevel::TimeOnly, None); assert!(report.changed, "{report:?}"); assert_eq!(report.num_matches("add"), 5, "{report:?}"); @@ -192,7 +192,10 @@ fn line_graph_1_test(strat: PlanStrategy) { rule.build(); let rule_set = rsb.build(); - assert!(db.run_rule_set(&rule_set, ReportLevel::TimeOnly).changed); + assert!( + db.run_rule_set(&rule_set, ReportLevel::TimeOnly, None) + .changed + ); let mut expected = Vec::from_iter( nodes @@ -273,7 +276,10 @@ fn line_graph_2_test(strat: PlanStrategy) { rule.build(); let rule_set = rsb.build(); - assert!(db.run_rule_set(&rule_set, ReportLevel::TimeOnly).changed); + assert!( + db.run_rule_set(&rule_set, ReportLevel::TimeOnly, None) + .changed + ); let mut expected = Vec::from_iter( nodes.windows(2).map(|x| vec![x[0], x[1]]).chain( @@ -358,7 +364,10 @@ fn intersection_test(strat: PlanStrategy) { rule.build(); let rule_set = rsb.build(); - assert!(db.run_rule_set(&rule_set, ReportLevel::TimeOnly).changed); + assert!( + db.run_rule_set(&rule_set, ReportLevel::TimeOnly, None) + .changed + ); let expected = Vec::from_iter((6..10).map(|x| vec![Value::new(x)])); @@ -479,7 +488,7 @@ fn minimal_ac_inner() { rules.build(); let rule_set = rsb.build(); - db.run_rule_set(&rule_set, ReportLevel::TimeOnly); + db.run_rule_set(&rule_set, ReportLevel::TimeOnly, None); let add_table = db.get_table(add); let all_add = add_table.all(); let items = add_table.scan(all_add.as_ref()); @@ -674,7 +683,7 @@ fn ac_test_inner(strat: PlanStrategy) { .unwrap(); rules.build(); let rule_set = rsb.build(); - db.run_rule_set(&rule_set, ReportLevel::TimeOnly) + db.run_rule_set(&rule_set, ReportLevel::TimeOnly, None) }; let rebuild = |db: &mut Database, cur_ts: Value| -> (Value, bool) { @@ -788,7 +797,7 @@ fn ac_test_inner(strat: PlanStrategy) { .unwrap(); rules.build(); let rs = rsb.build(); - changed |= db.run_rule_set(&rs, ReportLevel::TimeOnly).changed; + changed |= db.run_rule_set(&rs, ReportLevel::TimeOnly, None).changed; let mut rsb = db.new_rule_set(); num_rebuild(&mut rsb, cur_ts, next_ts); let mut add_rebuild_l = rsb.new_rule(); @@ -829,7 +838,7 @@ fn ac_test_inner(strat: PlanStrategy) { rules.build(); let rs = rsb.build(); - changed |= db.run_rule_set(&rs, ReportLevel::TimeOnly).changed; + changed |= db.run_rule_set(&rs, ReportLevel::TimeOnly, None).changed; let mut rsb = db.new_rule_set(); num_rebuild(&mut rsb, cur_ts, next_ts); let mut add_rebuild_r = rsb.new_rule(); @@ -869,7 +878,7 @@ fn ac_test_inner(strat: PlanStrategy) { .unwrap(); rules.build(); let rs = rsb.build(); - changed |= db.run_rule_set(&rs, ReportLevel::TimeOnly).changed; + changed |= db.run_rule_set(&rs, ReportLevel::TimeOnly, None).changed; } else { // nonincremental. Just run one rule and recanonicalize everything. // add(x, y, id, t1) => @@ -918,7 +927,7 @@ fn ac_test_inner(strat: PlanStrategy) { .unwrap(); rules.build(); let rs = rsb.build(); - changed |= db.run_rule_set(&rs, ReportLevel::TimeOnly).changed; + changed |= db.run_rule_set(&rs, ReportLevel::TimeOnly, None).changed; } (next_ts, changed) }; @@ -1095,7 +1104,7 @@ fn lookup_with_fallback_partial_success_inner() { rb.insert(h, &[res.into(), y.into()]).unwrap(); rb.build(); let rs = rsb.build(); - assert!(db.run_rule_set(&rs, ReportLevel::TimeOnly).changed); + assert!(db.run_rule_set(&rs, ReportLevel::TimeOnly, None).changed); let h = db.get_table(h); let all = h.all(); @@ -1187,7 +1196,7 @@ fn call_external_with_fallback_inner() { rb.insert(h, &[res.into(), y.into()]).unwrap(); rb.build(); let rs = rsb.build(); - assert!(db.run_rule_set(&rs, ReportLevel::TimeOnly).changed); + assert!(db.run_rule_set(&rs, ReportLevel::TimeOnly, None).changed); let h = db.get_table(h); let all = h.all(); @@ -1253,7 +1262,7 @@ fn early_stop_inner() { rb.build_with_description("early_stop_test"); let rs = rsb.build(); - let report = db.run_rule_set(&rs, ReportLevel::TimeOnly); + let report = db.run_rule_set(&rs, ReportLevel::TimeOnly, None); let matches = report.num_matches("early_stop_test"); diff --git a/egglog-bridge/examples/ac.rs b/egglog-bridge/examples/ac.rs index e6a83b16e..894cebf6d 100644 --- a/egglog-bridge/examples/ac.rs +++ b/egglog-bridge/examples/ac.rs @@ -40,7 +40,12 @@ fn main() { }; // Running these rules on an empty database should change nothing. - assert!(!egraph.run_rules(&[add_comm, add_assoc]).unwrap().changed()); + assert!( + !egraph + .run_rules(&[add_comm, add_assoc], None) + .unwrap() + .changed() + ); // Fill the database. let mut ids = Vec::new(); @@ -81,7 +86,10 @@ fn main() { // Saturate for i in 0.. { let iter_start = web_time::Instant::now(); - let keep_going = egraph.run_rules(&[add_comm, add_assoc]).unwrap().changed(); + let keep_going = egraph + .run_rules(&[add_comm, add_assoc], None) + .unwrap() + .changed(); println!("Finished iteration {i} after {:?}", iter_start.elapsed()); if !keep_going { break; diff --git a/egglog-bridge/examples/math.rs b/egglog-bridge/examples/math.rs index 8c0eaa13b..010d8c3b1 100644 --- a/egglog-bridge/examples/math.rs +++ b/egglog-bridge/examples/math.rs @@ -231,7 +231,7 @@ fn main() { for i in 0..N { let start = Instant::now(); - let changed = egraph.run_rules(&rules).unwrap().changed(); + let changed = egraph.run_rules(&rules, None).unwrap().changed(); println!( "Iteration {i} finished, duration={:?} (saturated={})", start.elapsed(), diff --git a/egglog-bridge/src/lib.rs b/egglog-bridge/src/lib.rs index e527fc3a1..dceb68303 100644 --- a/egglog-bridge/src/lib.rs +++ b/egglog-bridge/src/lib.rs @@ -18,9 +18,9 @@ use std::{ use crate::core_relations::{ BaseValue, BaseValueId, BaseValues, ColumnId, Constraint, ContainerValue, ContainerValues, - CounterId, Database, DisplacedTable, ExecutionState, ExternalFunction, ExternalFunctionId, - MergeVal, Offset, PlanStrategy, SortedWritesTable, TableId, TaggedRowBuffer, Value, - WrappedTable, + CounterId, Database, DisplacedTable, ExecutionState, ExternalContext, ExternalFunction, + ExternalFunctionId, MergeVal, Offset, PlanStrategy, SortedWritesTable, TableId, + TaggedRowBuffer, Value, WrappedTable, }; use crate::numeric_id::{DenseIdMap, DenseIdMapWithReuse, NumericId, define_id}; use egglog_concurrency::ThreadPool; @@ -301,8 +301,9 @@ impl EGraph { /// Intern the given container value into the EGraph. pub fn get_container_value(&mut self, val: C) -> Value { self.register_container_ty::(); - self.db - .with_execution_state(|state| state.clone().container_values().register_val(val, state)) + self.db.with_execution_state(None, |state| { + state.clone().container_values().register_val(val, state) + }) } /// Register the given [`ContainerValue`] type with this EGraph. @@ -656,17 +657,31 @@ impl EGraph { /// Run the given rules, returning whether the database changed. /// /// If the given rules are malformed, this method can return an error. - pub fn run_rules(&mut self, rules: &[RuleId]) -> Result { + pub fn run_rules( + &mut self, + rules: &[RuleId], + context: ExternalContext<'_>, + ) -> Result { let thread_pool = self.thread_pool(); - install_thread_pool(thread_pool, || self.run_rules_inner(rules)) + install_thread_pool(thread_pool, || self.run_rules_inner(rules, context)) } - fn run_rules_inner(&mut self, rules: &[RuleId]) -> Result { + fn run_rules_inner( + &mut self, + rules: &[RuleId], + context: ExternalContext<'_>, + ) -> Result { let ts = self.next_ts(); let uf_size_before = self.db.get_table(self.uf_table).len(); - let rule_set_report = - run_rules_impl(&mut self.db, &mut self.rules, rules, ts, self.report_level)?; + let rule_set_report = run_rules_impl( + &mut self.db, + &mut self.rules, + rules, + ts, + self.report_level, + context, + )?; if let Some(message) = self.panic_message.lock().unwrap().take() { return Err(PanicError(message).into()); } @@ -789,6 +804,7 @@ impl EGraph { &[*rule], ts, ReportLevel::TimeOnly, + None, )? .changed; } @@ -804,6 +820,7 @@ impl EGraph { &[info.nonincremental_rebuild_rule], ts, ReportLevel::TimeOnly, + None, )? .changed; for rule in &info.incremental_rebuild_rules { @@ -877,6 +894,7 @@ impl EGraph { &scratch, ts, ReportLevel::TimeOnly, + None, )? .changed; scratch.clear(); @@ -893,6 +911,7 @@ impl EGraph { &scratch, ts, ReportLevel::TimeOnly, + None, )? .changed; scratch.clear(); @@ -1013,9 +1032,13 @@ impl EGraph { /// / global-action context: appropriate for one-shot database /// manipulation from outside any rule, not for use inside /// primitive implementations. - pub fn with_execution_state(&self, f: impl FnOnce(&mut ExecutionState<'_>) -> R) -> R { + pub fn with_execution_state( + &self, + context: ExternalContext<'_>, + f: impl FnOnce(&mut ExecutionState<'_>) -> R, + ) -> R { let thread_pool = self.thread_pool(); - install_thread_pool(thread_pool, || self.db.with_execution_state(f)) + install_thread_pool(thread_pool, || self.db.with_execution_state(context, f)) } /// Like [`EGraph::with_execution_state`], but also reports whether `f` @@ -1024,9 +1047,10 @@ impl EGraph { /// no-op merge plus a spurious timestamp bump. pub fn with_execution_state_tracked( &self, + context: ExternalContext<'_>, f: impl FnOnce(&mut ExecutionState<'_>) -> R, ) -> (R, bool) { - self.db.with_execution_state_tracked(f) + self.db.with_execution_state_tracked(context, f) } /// Flush the pending update buffers to the EGraph. @@ -1582,6 +1606,7 @@ fn run_rules_impl( rules: &[RuleId], next_ts: Timestamp, report_level: ReportLevel, + context: ExternalContext<'_>, ) -> Result { for rule in rules { let info = &mut rule_info[*rule]; @@ -1598,7 +1623,7 @@ fn run_rules_impl( info.last_run_at = next_ts; } let ruleset = rsb.build(); - Ok(db.run_rule_set(&ruleset, report_level)) + Ok(db.run_rule_set(&ruleset, report_level, context)) } // These markers are just used to make it easy to distinguish time spent in diff --git a/egglog-bridge/src/tests.rs b/egglog-bridge/src/tests.rs index 54df39f03..7deda5c41 100644 --- a/egglog-bridge/src/tests.rs +++ b/egglog-bridge/src/tests.rs @@ -58,7 +58,12 @@ fn ac_test(can_subsume: bool) { }; // Running these rules on an empty database should change nothing. - assert!(!egraph.run_rules(&[add_comm, add_assoc]).unwrap().changed()); + assert!( + !egraph + .run_rules(&[add_comm, add_assoc], None) + .unwrap() + .changed() + ); // Fill the database. let mut ids = Vec::new(); @@ -87,7 +92,11 @@ fn ac_test(can_subsume: bool) { (left_root, right_root) }; // Saturate - while egraph.run_rules(&[add_comm, add_assoc]).unwrap().changed() {} + while egraph + .run_rules(&[add_comm, add_assoc], None) + .unwrap() + .changed() + {} let canon_left = egraph.get_canon_in_uf(left_root); let canon_right = egraph.get_canon_in_uf(right_root); assert_eq!(canon_left, canon_right, "failed to reassociate!"); @@ -136,7 +145,12 @@ fn ac_fail() { }; // Running these rules on an empty database should change nothing. - assert!(!egraph.run_rules(&[add_comm, add_assoc]).unwrap().changed()); + assert!( + !egraph + .run_rules(&[add_comm, add_assoc], None) + .unwrap() + .changed() + ); // Fill the database. let mut ids = Vec::new(); @@ -174,7 +188,11 @@ fn ac_fail() { (left_root, right_root) }; // Saturate - while egraph.run_rules(&[add_comm, add_assoc]).unwrap().changed() {} + while egraph + .run_rules(&[add_comm, add_assoc], None) + .unwrap() + .changed() + {} let canon_left = egraph.get_canon_in_uf(left_root); let canon_right = egraph.get_canon_in_uf(right_root); assert_ne!(canon_left, canon_right); @@ -409,7 +427,7 @@ fn math_test(mut egraph: EGraph, can_subsume: bool) { } for _ in 0..N { - if !egraph.run_rules(&rules).unwrap().changed() { + if !egraph.run_rules(&rules, None).unwrap().changed() { break; } } @@ -675,20 +693,20 @@ fn container_test() { vec![vec![], vec![egraph.get_canon_in_uf(ids[1])]], ); - assert!(egraph.run_rules(&[vec_expand]).unwrap().changed()); + assert!(egraph.run_rules(&[vec_expand], None).unwrap().changed()); assert_eq!(dump_vecs(&egraph).len(), 4); // We have 2 new vectors with a last element. Each of those should spawn two more, adding 4. - assert!(egraph.run_rules(&[vec_expand]).unwrap().changed()); + assert!(egraph.run_rules(&[vec_expand], None).unwrap().changed()); assert_eq!(dump_vecs(&egraph).len(), 8); // We have 4 new vectors with a last element. Each of those should spawn two more, adding 8. - assert!(egraph.run_rules(&[vec_expand]).unwrap().changed()); + assert!(egraph.run_rules(&[vec_expand], None).unwrap().changed()); assert_eq!(dump_vecs(&egraph).len(), 16); // Now we want to saturate `eval_add`. This should collapse a bunch of new vectors. let mut saturated = false; for _ in 0..20 { - saturated = !egraph.run_rules(&[eval_add]).unwrap().changed(); + saturated = !egraph.run_rules(&[eval_add], None).unwrap().changed(); if saturated { break; } @@ -794,7 +812,10 @@ fn run_query_prim_container_match_case(seminaive: bool, seed_canonical: bool) -> let mut saturated = false; for _ in 0..8 { - saturated = !egraph.run_rules(&[w_rewrite, l_rewrite]).unwrap().changed(); + saturated = !egraph + .run_rules(&[w_rewrite, l_rewrite], None) + .unwrap() + .changed(); if saturated { break; } @@ -839,7 +860,7 @@ fn rhs_only_rule() { let mut contents = Vec::new(); assert!(contents.is_empty()); - assert!(egraph.run_rules(&[add_data]).unwrap().changed()); + assert!(egraph.run_rules(&[add_data], None).unwrap().changed()); egraph.for_each(num_table, |func_row| { assert!(!func_row.subsumed); contents.push(func_row.vals.to_vec()); @@ -868,9 +889,19 @@ fn rhs_only_rule_only_runs_once() { rb.build() }; - assert!(!egraph.run_rules(&[inc_counter_rule]).unwrap().changed()); + assert!( + !egraph + .run_rules(&[inc_counter_rule], None) + .unwrap() + .changed() + ); assert_eq!(counter.load(Ordering::SeqCst), 1); - assert!(!egraph.run_rules(&[inc_counter_rule]).unwrap().changed()); + assert!( + !egraph + .run_rules(&[inc_counter_rule], None) + .unwrap() + .changed() + ); assert_eq!(counter.load(Ordering::SeqCst), 1); } @@ -939,7 +970,7 @@ fn mergefn_arithmetic() { }; // Run the first rule and check state - assert!(egraph.run_rules(&[rule1]).unwrap().changed()); + assert!(egraph.run_rules(&[rule1], None).unwrap().changed()); let mut contents = Vec::new(); egraph.for_each(f_table, |func_row| { assert!(!func_row.subsumed); @@ -962,7 +993,7 @@ fn mergefn_arithmetic() { // Run the second rule and check state // Expected: (f 1 1) because 1 + (0 * 5) = 1 // Expected: (f 2 7) because 1 + (1 * 6) = 7 - assert!(egraph.run_rules(&[rule2]).unwrap().changed()); + assert!(egraph.run_rules(&[rule2], None).unwrap().changed()); contents.clear(); egraph.for_each(f_table, |func_row| { assert!(!func_row.subsumed); @@ -985,7 +1016,7 @@ fn mergefn_arithmetic() { // Run the third rule and check state // Expected: (f 1 4) because 1 + (1 * 3) = 4 // Expected: (f 2 29) because 1 + (7 * 4) = 29 - assert!(egraph.run_rules(&[rule3]).unwrap().changed()); + assert!(egraph.run_rules(&[rule3], None).unwrap().changed()); contents.clear(); egraph.for_each(f_table, |func_row| { assert!(!func_row.subsumed); @@ -1067,7 +1098,7 @@ fn mergefn_nested_function() { }; // First run of the rule - assert!(egraph.run_rules(&[write_rule]).unwrap().changed()); + assert!(egraph.run_rules(&[write_rule], None).unwrap().changed()); let f_entries_1 = get_f_entries(&egraph); let g_entries_1 = get_g_entries(&egraph); assert_eq!(f_entries_1.len(), 2); @@ -1092,7 +1123,7 @@ fn mergefn_nested_function() { }; // Second run of the rule - should trigger merging with previous values - assert!(egraph.run_rules(&[set_rule]).unwrap().changed()); + assert!(egraph.run_rules(&[set_rule], None).unwrap().changed()); let f_entries_2 = get_f_entries(&egraph); let g_entries_2 = get_g_entries(&egraph); assert_eq!(f_entries_2.len(), 2); @@ -1197,13 +1228,13 @@ fn constrain_prims_simple() { assert!(get_entries(&egraph, f_table).is_empty()); assert!(get_entries(&egraph, g_table).is_empty()); - egraph.run_rules(&[write_f]).unwrap(); + egraph.run_rules(&[write_f], None).unwrap(); let f = get_entries(&egraph, f_table); assert_eq!(f.len(), 3); - egraph.run_rules(&[copy_to_g]).unwrap(); + egraph.run_rules(&[copy_to_g], None).unwrap(); let invocations_after_first = query_prim_invocations.load(Ordering::Relaxed); assert!(invocations_after_first > 0); - assert!(!egraph.run_rules(&[copy_to_g]).unwrap().changed()); + assert!(!egraph.run_rules(&[copy_to_g], None).unwrap().changed()); assert_eq!( query_prim_invocations.load(Ordering::Relaxed), invocations_after_first @@ -1301,10 +1332,10 @@ fn constrain_prims_abstract() { assert!(get_entries(&egraph, f_table).is_empty()); assert!(get_entries(&egraph, g_table).is_empty()); - egraph.run_rules(&[write_f]).unwrap(); + egraph.run_rules(&[write_f], None).unwrap(); let f = get_entries(&egraph, f_table); assert_eq!(f.len(), 3); - egraph.run_rules(&[copy_to_g]).unwrap(); + egraph.run_rules(&[copy_to_g], None).unwrap(); let g = get_entries(&egraph, g_table); assert_eq!(g.len(), 2); assert_eq!(g, f[0..2]) @@ -1375,18 +1406,18 @@ fn basic_subsumption() { assert!(get_entries(&egraph, f_table).0.is_empty()); assert!(get_entries(&egraph, g_table).0.is_empty()); - egraph.run_rules(&[write_f]).unwrap(); + egraph.run_rules(&[write_f], None).unwrap(); let f = get_entries(&egraph, f_table); assert_eq!((f.0.len(), f.1), (2, 0)); assert_eq!(f.0.iter().map(|(x, _)| *x).collect::>(), vec![1, 2]); - egraph.run_rules(&[subsume_f]).unwrap(); + egraph.run_rules(&[subsume_f], None).unwrap(); let f = get_entries(&egraph, f_table); assert_eq!((f.0.len(), f.1), (3, 2)); assert_eq!( f.0.iter().map(|(x, _)| *x).collect::>(), vec![1, 2, 3] ); - egraph.run_rules(&[copy_to_g]).unwrap(); + egraph.run_rules(&[copy_to_g], None).unwrap(); let g = get_entries(&egraph, g_table); assert_eq!((g.0.len(), g.1), (1, 0)); assert_eq!(g.0[0], f.0[0]) @@ -1417,21 +1448,21 @@ fn lookup_failure_panics() { rb.set(f, &[value_2.clone(), value_2.clone()]); rb.build() }; - egraph.run_rules(&[write_f]).unwrap(); + egraph.run_rules(&[write_f], None).unwrap(); let lookup_success = { let mut rb = egraph.new_rule("lookup_success", true); rb.lookup(f, slice::from_ref(&value_1), String::new); rb.build() }; - egraph.run_rules(&[lookup_success]).unwrap(); + egraph.run_rules(&[lookup_success], None).unwrap(); let lookup_failure = { let mut rb = egraph.new_rule("lookup_fail", true); rb.lookup(f, slice::from_ref(&value_3), String::new); rb.build() }; - egraph.run_rules(&[lookup_failure]).err().unwrap(); + egraph.run_rules(&[lookup_failure], None).err().unwrap(); } #[test] @@ -1474,7 +1505,7 @@ fn primitive_failure_panics() { rb.build() }; - egraph.run_rules(&[assert_odd_rule]).err().unwrap(); + egraph.run_rules(&[assert_odd_rule], None).err().unwrap(); } #[test] @@ -1483,7 +1514,7 @@ fn panic_functions_trigger_early_stop() { let channel: crate::SideChannel = Default::default(); let panic_fn = super::Panic("panic".to_string(), channel.clone()); - let stopped = db.with_execution_state(|state| { + let stopped = db.with_execution_state(None, |state| { assert!(!state.should_stop()); let res = core_relations::ExternalFunction::invoke(&panic_fn, state, &[]); assert!(res.is_none()); @@ -1495,7 +1526,7 @@ fn panic_functions_trigger_early_stop() { let channel: crate::SideChannel = Default::default(); let lazy = Lazy::new(|| "lazy panic".to_string()); let panic_fn = super::LazyPanic(Arc::new(lazy), channel.clone()); - let stopped = db.with_execution_state(|state| { + let stopped = db.with_execution_state(None, |state| { assert!(!state.should_stop()); let res = core_relations::ExternalFunction::invoke(&panic_fn, state, &[]); assert!(res.is_none()); diff --git a/src/core.rs b/src/core.rs index eeeea5d0b..47a68abb8 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.as_ref().clone())); + return Some(ResolvedCall::Func(ty.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.as_ref().clone())); + return Ok(ResolvedCall::Func(ty.clone())); } } diff --git a/src/exec_state.rs b/src/exec_state.rs index 3dcbb304a..fd9291002 100644 --- a/src/exec_state.rs +++ b/src/exec_state.rs @@ -32,7 +32,6 @@ use std::any::TypeId; use std::ops::Deref; -use std::sync::Arc; use crate::Error; use crate::api::{ApiError, IntoValue, IntoValues}; @@ -41,10 +40,11 @@ use crate::core_relations::{ Value, }; use crate::{ + TypeInfo, ast::{FunctionSubtype, Literal, ResolvedExpr}, core::ResolvedCall, sort::{F, S}, - typechecking::{FuncType, SharedFuncTypes}, + typechecking::FuncType, }; use egglog_bridge::{ActionRegistry, TableAction, TableKind}; use smallvec::SmallVec; @@ -144,9 +144,18 @@ 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; + + /// The [`TypeInfo`] the e-graph made visible for this operation, or `None` + /// in an execution the e-graph did not supply one for (its internal + /// rebuild rules, whose actions never read a signature). + /// + /// This is a borrow rather than a shared handle, which is what keeps a + /// declaration from racing a running rule: the e-graph passes `&TypeInfo` + /// into the execution, so it cannot be `&mut` borrowed to declare anything + /// until that execution ends. + fn type_info(&self) -> Option<&'db TypeInfo> { + self.es().external_context()?.downcast_ref() + } } // ===================================================================== @@ -384,16 +393,16 @@ pub trait Read<'a, 'db: 'a>: Core<'a, 'db> + RegistrySealed<'a, 'db> { /// /// **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) + fn constructor_schema(&self, name: &str) -> Result<&'db FuncType, Error> { + func_type_of(self.type_info(), 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) + fn function_schema(&self, name: &str) -> Result<&'db FuncType, Error> { + func_type_of(self.type_info(), name, FunctionSubtype::Custom) } /// Whether the named table is a `constructor` or a `function`, or `None` @@ -401,8 +410,7 @@ pub trait Read<'a, 'db: 'a>: Core<'a, 'db> + RegistrySealed<'a, 'db> { /// 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) + Some(self.type_info()?.get_func_type(name)?.subtype) } /// Return the current row count for the named table, or `None` if no table @@ -632,13 +640,12 @@ fn subtype_label(subtype: FunctionSubtype) -> &'static str { } } -fn func_type_of( - func_types: &SharedFuncTypes, +fn func_type_of<'db>( + type_info: Option<&'db TypeInfo>, name: &str, expected: FunctionSubtype, -) -> Result, Error> { - let func_type = func_types.read().unwrap().get(name).cloned(); - let Some(func_type) = func_type else { +) -> Result<&'db FuncType, Error> { + let Some(func_type) = type_info.and_then(|type_info| type_info.get_func_type(name)) else { return Err(ApiError::MissingTable { name: name.to_string(), } @@ -740,7 +747,6 @@ 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, } @@ -755,7 +761,6 @@ 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, } @@ -770,7 +775,6 @@ 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, } @@ -787,13 +791,11 @@ 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, } } @@ -806,13 +808,11 @@ 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, } } @@ -825,13 +825,11 @@ 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, } } @@ -876,9 +874,6 @@ 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> {} @@ -901,9 +896,6 @@ 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> {} @@ -926,9 +918,6 @@ 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/lib.rs b/src/lib.rs index 5475a2185..2dfbb2073 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -714,14 +714,7 @@ 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!())), @@ -795,7 +788,7 @@ impl EGraph { // 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) { + let func_type = match self.type_info.func_type_arc(&decl.name) { Some(func_type) => func_type, None => { let get_sort = |name: &String| match self.type_info.get_sort_by_name(name) { @@ -1134,7 +1127,7 @@ impl EGraph { let iteration_report = self .backend - .run_rules(&rule_ids) + .run_rules(&rule_ids, Some(&self.type_info)) .map_err(|e| Error::BackendError(e.to_string()))?; Ok(RunReport::singleton(ruleset, iteration_report)) @@ -1211,7 +1204,7 @@ impl EGraph { ); translator.actions(&actions)?; let id = translator.build(); - let result = self.backend.run_rules(&[id]); + let result = self.backend.run_rules(&[id], Some(&self.type_info)); self.backend.free_rule(id); match result { @@ -1238,10 +1231,9 @@ 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, func_types, Context::Read)) + .with_execution_state_tracked(Some(&self.type_info), |es| { + f(ReadState::wrap(es, &guard, Context::Read)) }) .0 } @@ -1521,7 +1513,7 @@ impl EGraph { ); let id = translator.build(); - let rule_result = self.backend.run_rules(&[id]); + let rule_result = self.backend.run_rules(&[id], Some(&self.type_info)); self.backend.free_rule(id); self.backend.free_external_func(ext_id); let _ = rule_result.map_err(|e| { @@ -1588,7 +1580,7 @@ impl EGraph { "this function will never panic".to_string() }); let id = translator.build(); - let run_result = self.backend.run_rules(&[id]); + let run_result = self.backend.run_rules(&[id], Some(&self.type_info)); self.backend.free_rule(id); self.backend.free_external_func(ext_id); run_result.map_err(|e| Error::BackendError(e.to_string()))?; @@ -1969,21 +1961,23 @@ impl EGraph { let table_action = egglog_bridge::TableAction::new(&self.backend, func.backend_id); if function_type.subtype != FunctionSubtype::Constructor { - self.backend.with_execution_state(|es| { - for row in parsed_contents.iter() { - table_action.insert(es, row.iter().copied()); - } - Some(unit_val) - }); + self.backend + .with_execution_state(Some(&self.type_info), |es| { + for row in parsed_contents.iter() { + table_action.insert(es, row.iter().copied()); + } + Some(unit_val) + }); } else { - self.backend.with_execution_state(|es| { - for row in parsed_contents.iter() { - // Constructor semantics: mint a fresh eclass id for - // each missing key. - table_action.lookup_or_insert(es, row); - } - Some(unit_val) - }); + self.backend + .with_execution_state(Some(&self.type_info), |es| { + for row in parsed_contents.iter() { + // Constructor semantics: mint a fresh eclass id for + // each missing key. + table_action.lookup_or_insert(es, row); + } + Some(unit_val) + }); } self.backend.flush_updates(); @@ -2277,9 +2271,10 @@ impl EGraph { /// Convert from a Rust container type to an egglog value. pub fn container_to_value(&mut self, x: T) -> Value { - self.backend.with_execution_state(|state| { - self.backend.container_values().register_val::(x, state) - }) + self.backend + .with_execution_state(Some(&self.type_info), |state| { + self.backend.container_values().register_val::(x, state) + }) } /// Get the size of a function in the e-graph. @@ -2389,10 +2384,11 @@ impl EGraph { ) -> Result { let registry = self.backend.action_registry().clone(); let guard = registry.read().unwrap(); - 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)) - }); + let (result, changed) = self + .backend + .with_execution_state_tracked(Some(&self.type_info), |es| { + f(FullState::wrap(es, &guard, 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/scheduler.rs b/src/scheduler.rs index 2592f79c0..f34af869c 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -255,25 +255,27 @@ impl EGraph { let query_iter_report = self .backend - .run_rules(&query_rules) + .run_rules(&query_rules, Some(&self.type_info)) .map_err(|e| Error::BackendError(e.to_string()))?; // Step 3: let the scheduler decide which matches need to be kept - self.backend.with_execution_state(|state| { - for (rule_id, _rule) in rules.iter() { - let rule_info = record.rule_info.get_mut(rule_id).unwrap(); - - let matches: Vec = - std::mem::take(rule_info.matches.lock().unwrap().as_mut()); - let mut matches = Matches::new(matches, rule_info.free_vars.clone()); - rule_info.should_seek = - record - .scheduler - .filter_matches(rule_id, ruleset, &mut matches); - let table_action = TableAction::new(&self.backend, rule_info.decided); - *rule_info.matches.lock().unwrap() = matches.instantiate(state, &table_action); - } - }); + self.backend + .with_execution_state(Some(&self.type_info), |state| { + for (rule_id, _rule) in rules.iter() { + let rule_info = record.rule_info.get_mut(rule_id).unwrap(); + + let matches: Vec = + std::mem::take(rule_info.matches.lock().unwrap().as_mut()); + let mut matches = Matches::new(matches, rule_info.free_vars.clone()); + rule_info.should_seek = + record + .scheduler + .filter_matches(rule_id, ruleset, &mut matches); + let table_action = TableAction::new(&self.backend, rule_info.decided); + *rule_info.matches.lock().unwrap() = + matches.instantiate(state, &table_action); + } + }); self.backend.flush_updates(); // Step 4: run the action rules @@ -286,7 +288,7 @@ impl EGraph { .collect::>(); let action_iter_report = self .backend - .run_rules(&action_rules) + .run_rules(&action_rules, Some(&self.type_info)) .map_err(|e| Error::BackendError(e.to_string()))?; // Step 5: combine the reports diff --git a/src/typechecking.rs b/src/typechecking.rs index 3c479afe9..6ff01069c 100644 --- a/src/typechecking.rs +++ b/src/typechecking.rs @@ -44,7 +44,6 @@ 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>, @@ -57,7 +56,6 @@ trait RegistryWrap: Clone + Send + Sync { ctx: Context, args: &[Value], registry: &ActionRegistry, - func_types: &SharedFuncTypes, ) -> Option; } @@ -71,9 +69,8 @@ impl RegistryWrap for WrapRead { ctx: Context, args: &[Value], registry: &ActionRegistry, - func_types: &SharedFuncTypes, ) -> Option { - prim.apply(ReadState::wrap(exec_state, registry, func_types, ctx), args) + prim.apply(ReadState::wrap(exec_state, registry, ctx), args) } } #[derive(Clone)] @@ -86,12 +83,8 @@ impl RegistryWrap for WrapWrite { ctx: Context, args: &[Value], registry: &ActionRegistry, - func_types: &SharedFuncTypes, ) -> Option { - prim.apply( - WriteState::wrap(exec_state, registry, func_types, ctx), - args, - ) + prim.apply(WriteState::wrap(exec_state, registry, ctx), args) } } #[derive(Clone)] @@ -104,9 +97,8 @@ impl RegistryWrap for WrapFull { ctx: Context, args: &[Value], registry: &ActionRegistry, - func_types: &SharedFuncTypes, ) -> Option { - prim.apply(FullState::wrap(exec_state, registry, func_types, ctx), args) + prim.apply(FullState::wrap(exec_state, registry, ctx), args) } } @@ -115,14 +107,7 @@ 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, - &self.func_types, - ) + S::invoke(&self.prim, exec_state, self.ctx, args, ®istry) } } @@ -179,15 +164,6 @@ 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. @@ -240,37 +216,19 @@ impl Debug for PrimitiveWithId { } /// Stores resolved typechecking information. -#[derive(Default)] +#[derive(Clone, 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: SharedFuncTypes, + func_types: HashMap>, 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 { @@ -291,13 +249,7 @@ impl EGraph { span: Span, ) -> Result<(), TypeError> { let name = name.into(); - if self - .type_info - .func_types - .read() - .unwrap() - .contains_key(&name) - { + if self.type_info.func_types.contains_key(&name) { return Err(TypeError::FunctionAlreadyBound(name, span)); } @@ -388,12 +340,10 @@ 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, }) @@ -634,10 +584,7 @@ impl EGraph { span.clone(), )); } - ResolvedNCommand::ProveExists( - span.clone(), - ResolvedCall::Func(func_type.as_ref().clone()), - ) + ResolvedNCommand::ProveExists(span.clone(), ResolvedCall::Func(func_type.clone())) } NCommand::Output { span, file, exprs } => { let exprs = exprs @@ -851,8 +798,6 @@ impl TypeInfo { let ftype = self.function_to_functype(fdecl)?; if self .func_types - .write() - .unwrap() .insert(fdecl.name.clone(), Arc::new(ftype)) .is_some() { @@ -879,7 +824,6 @@ impl TypeInfo { resolved_schema: ResolvedCall::Func( self.get_func_type(&fdecl.name) .expect("just inserted") - .as_ref() .clone(), ), merge: match &fdecl.merge { @@ -1214,28 +1158,20 @@ 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> { - self.func_types.read().unwrap().get(sym).cloned() + pub fn get_func_type(&self, sym: &str) -> Option<&FuncType> { + self.func_types.get(sym).map(Arc::as_ref) + } + + /// The signature as a shared handle, for callers that keep it around. + pub(crate) fn func_type_arc(&self, sym: &str) -> Option> { + self.func_types.get(sym).cloned() } /// 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 - .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; + self.func_types.insert(func_type.name.clone(), func_type); } pub fn is_constructor(&self, sym: &str) -> bool { From f3e3a855d6c5907252de6fc1b8af90776beb49f4 Mon Sep 17 00:00:00 2001 From: Oliver Flatt Date: Fri, 7 Aug 2026 21:28:09 +0000 Subject: [PATCH 6/8] Share the signature a function decl carries, and say why it carries one GenericFunctionDecl held both an untyped `schema: Schema` and a `resolved_schema: Head`, which for a resolved decl was a whole owned FuncType inside a ResolvedCall - a fourth copy of the signature - and for an unresolved one was a placeholder `String::new()` standing in for "not resolved yet". It is now `Option>`: no placeholder, and a typechecked decl points at the same FuncType TypeInfo holds rather than cloning it. Worth recording why the field survives at all, since it looks redundant next to TypeInfo. Proof encoding runs over decls that desugaring generates *after* typechecking - the functions global bindings lower to - which no TypeInfo knows about until they are executed. Deleting the field and looking signatures up by name works for every ordinary declaration and then panics on `$I`. The doc comment now says so, so the next reader does not have to rediscover it. Also drops ResolvedCall::view_types, whose only caller was one of the two readers this touches. Co-Authored-By: Claude Opus 5 --- CHANGELOG.md | 2 +- src/ast/mod.rs | 16 +++++++++++----- src/ast/proof_global_remover.rs | 5 +++-- src/ast/remove_globals.rs | 5 +++-- src/core.rs | 11 ----------- src/proofs/proof_encoding.rs | 29 +++++++++++++++++++++++++---- src/typechecking.rs | 7 +++---- 7 files changed, 46 insertions(+), 29 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e1d261174..463ad2e37 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +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. `Function::schema() -> &ResolvedSchema` becomes `Function::func_type() -> &FuncType`, and `ResolvedSchema` is removed (its `get_by_pos` moves to `FuncType`). +- **Breaking:** a function's signature is stored once, shared between `TypeInfo`, `Function`, and `GenericFunctionDecl` instead of copied into each. `Function::schema() -> &ResolvedSchema` becomes `Function::func_type() -> &FuncType` and `ResolvedSchema` is removed (its `get_by_pos` moves to `FuncType`); `GenericFunctionDecl::resolved_schema` becomes `Option>`, so an unresolved decl holds `None` rather than a placeholder `String`. - **Breaking:** an operation that runs external functions now takes an `ExternalContext` — a borrowed value it makes visible to them, which `ExecutionState::external_context` reads back. `Database::run_rule_set` / `with_execution_state` / `with_execution_state_tracked` and the bridge's `EGraph::run_rules` / `with_execution_state*` take one; pass `None` if there is nothing to share. egglog passes its `&TypeInfo`, which is how a primitive resolves a signature without the e-graph handing out a shared, lock-guarded copy: the borrow is live for exactly the operation that supplied it, so nothing can be declared while a rule is running. - **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). diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 7970552cb..9edb0131e 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -1320,10 +1320,16 @@ where { pub name: String, pub subtype: FunctionSubtype, - /// Untyped schema + /// Untyped schema, as written in the source. pub schema: Schema, - /// Resolved schema after typechecking is stored here, otherwise "". - pub resolved_schema: Head, + /// The resolved signature, once typechecking has run. + /// + /// The same signature is in `TypeInfo`, keyed by `name`, and that is where + /// to read it from. This carries it for the one consumer that cannot: + /// proof encoding runs over decls that desugaring generated *after* + /// typechecking (global bindings), which no `TypeInfo` knows about until + /// they are executed. + pub resolved_schema: Option>, pub merge: Option>, pub cost: Option, pub unextractable: bool, @@ -1391,7 +1397,7 @@ impl FunctionDecl { name, subtype: FunctionSubtype::Custom, schema, - resolved_schema: String::new(), + resolved_schema: None, merge, cost: None, unextractable: true, @@ -1414,7 +1420,7 @@ impl FunctionDecl { Self { name, subtype: FunctionSubtype::Constructor, - resolved_schema: String::new(), + resolved_schema: None, schema, merge: None, cost, diff --git a/src/ast/proof_global_remover.rs b/src/ast/proof_global_remover.rs index 2cbbb4b77..e0a35f74a 100644 --- a/src/ast/proof_global_remover.rs +++ b/src/ast/proof_global_remover.rs @@ -75,12 +75,13 @@ fn remove_globals_cmd(cmd: ResolvedNCommand) -> Vec { panic!("Global variable {} has non-eq sort {}", name, ty.name()); } - let resolved_call = ResolvedCall::Func(FuncType { + let func_type = Arc::new(FuncType { name: name.name.clone(), subtype: FunctionSubtype::Constructor, input: vec![], output: ty.clone(), }); + let resolved_call = ResolvedCall::Func(func_type.as_ref().clone()); let func_decl = ResolvedFunctionDecl { name: name.name, subtype: FunctionSubtype::Constructor, @@ -88,7 +89,7 @@ fn remove_globals_cmd(cmd: ResolvedNCommand) -> Vec { input: vec![], output: ty.name().to_owned(), }, - resolved_schema: resolved_call.clone(), + resolved_schema: Some(func_type.clone()), merge: None, cost: None, unextractable: true, diff --git a/src/ast/remove_globals.rs b/src/ast/remove_globals.rs index c37f908af..d9d15b81e 100644 --- a/src/ast/remove_globals.rs +++ b/src/ast/remove_globals.rs @@ -93,12 +93,13 @@ impl GlobalRemover<'_> { GenericAction::Let(span, name, expr) => { let ty = expr.output_type(); - let resolved_call = ResolvedCall::Func(FuncType { + let func_type = Arc::new(FuncType { name: name.name.clone(), subtype: FunctionSubtype::Custom, input: vec![], output: ty.clone(), }); + let resolved_call = ResolvedCall::Func(func_type.as_ref().clone()); let func_decl = ResolvedFunctionDecl { name: name.name, subtype: FunctionSubtype::Custom, @@ -106,7 +107,7 @@ impl GlobalRemover<'_> { input: vec![], output: ty.name().to_owned(), }, - resolved_schema: resolved_call.clone(), + resolved_schema: Some(func_type.clone()), merge: None, cost: None, unextractable: true, diff --git a/src/core.rs b/src/core.rs index 47a68abb8..d2e0dde6b 100644 --- a/src/core.rs +++ b/src/core.rs @@ -134,17 +134,6 @@ impl ResolvedCall { /// Gives the types for a term's child with the given resolved call. /// For functions this includes the output sort, for constructors it's just the inputs. - pub(crate) fn view_types(&self) -> Vec { - match self { - ResolvedCall::Func(func) => { - let mut types = func.input.clone(); - types.push(func.output.clone()); - types - } - ResolvedCall::Primitive(prim) => prim.input().to_vec(), - } - } - // Different from `from_resolution`, this function only considers function types and not primitives. // As a result, it only requires input argument types, so types.len() == func.input.len(), // while for `from_resolution`, types.len() == func.input.len() + 1 to account for the output type diff --git a/src/proofs/proof_encoding.rs b/src/proofs/proof_encoding.rs index 9804be1cd..ac57cda56 100644 --- a/src/proofs/proof_encoding.rs +++ b/src/proofs/proof_encoding.rs @@ -60,6 +60,29 @@ pub(crate) struct ProofInstrumentor<'a> { pub(crate) egraph: &'a mut EGraph, } +/// The signature typechecking resolved for `fdecl`. +/// +/// Encoding also runs over decls desugaring generated after typechecking, which +/// no `TypeInfo` knows about yet, so it reads the signature off the decl rather +/// than looking it up by name. +fn resolved_signature(fdecl: &ResolvedFunctionDecl) -> &FuncType { + fdecl + .resolved_schema + .as_ref() + .unwrap_or_else(|| panic!("{} reached encoding unresolved", fdecl.name)) +} + +/// The declared sorts of `fdecl`'s columns, inputs followed by the output. +fn column_sorts(fdecl: &ResolvedFunctionDecl) -> Vec { + let func_type = resolved_signature(fdecl); + func_type + .input + .iter() + .chain([&func_type.output]) + .cloned() + .collect() +} + impl<'a> ProofInstrumentor<'a> { /// Make a term state and use it to instrument the code. pub(crate) fn add_term_encoding( @@ -557,7 +580,7 @@ impl<'a> ProofInstrumentor<'a> { /// Rules that update the views when children change. fn rebuilding_rules(&mut self, fdecl: &ResolvedFunctionDecl) -> Vec { - let types = fdecl.resolved_schema.view_types(); + let types = column_sorts(fdecl); let proofs_enabled = self.egraph.proof_state.proofs_enabled; // Check if there are any rebuildable columns at all; if not, no rule needed. @@ -734,9 +757,7 @@ impl<'a> ProofInstrumentor<'a> { /// 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 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") - }; + let input = &resolved_signature(fdecl).input; // Check if there are any eq-sort columns at all; if not, no rebuild rule needed. if !input.iter().any(|t| t.is_eq_sort()) { diff --git a/src/typechecking.rs b/src/typechecking.rs index 6ff01069c..8e74c8419 100644 --- a/src/typechecking.rs +++ b/src/typechecking.rs @@ -821,10 +821,9 @@ impl TypeInfo { name: fdecl.name.clone(), subtype: fdecl.subtype, schema: fdecl.schema.clone(), - resolved_schema: ResolvedCall::Func( - self.get_func_type(&fdecl.name) - .expect("just inserted") - .clone(), + resolved_schema: Some( + self.func_type_arc(&fdecl.name) + .expect("the signature was just recorded"), ), merge: match &fdecl.merge { // Merge expressions run as part of action-side table updates: From 491dc530fb4c4816e1899281ca95e82fc7141002 Mon Sep 17 00:00:00 2001 From: Oliver Flatt Date: Fri, 7 Aug 2026 21:44:22 +0000 Subject: [PATCH 7/8] Store a function's signature once, everywhere Two copies were left after the last pass. GenericFunctionDecl::resolved_schema is gone. I had concluded it was load bearing, because deleting it and looking signatures up by name panics on `$I`: proof encoding runs over decls desugaring generates after typechecking, which no TypeInfo knows about yet. That was the wrong conclusion. The encoder never needed a FuncType - it needed the column *sorts*, and sorts resolve by name from TypeInfo like any other. Resolving them from the decl's own untyped schema drops the field, its ResolvedCall::view_types reader, and the placeholder `String::new()` an unresolved decl carried in it. (Populating TypeInfo with the generated signatures instead does not work: the desugared program is typechecked again afterwards, so recording early trips the already-bound check.) ResolvedCall::Func now holds an Arc rather than an owned one, so a resolved call site shares the signature instead of cloning it - and there are as many of those as there are calls in every rule. The patterns that matched `ResolvedCall::Func(FuncType { subtype: Custom, .. })` cannot look through an Arc, so they match the variant and test the subtype in a guard via a new ResolvedCall::is_custom_func. Co-Authored-By: Claude Opus 5 --- CHANGELOG.md | 2 +- src/ast/mod.rs | 14 +----- src/ast/proof_global_remover.rs | 10 ++--- src/ast/remove_globals.rs | 10 ++--- src/core.rs | 20 ++++++--- src/lib.rs | 63 ++++++++++++++------------ src/proofs/proof_checker.rs | 16 ++----- src/proofs/proof_encoding.rs | 79 +++++++++++++++++++-------------- src/proofs/proof_format.rs | 15 ++----- src/proofs/proof_normal_form.rs | 29 ++++-------- src/typechecking.rs | 9 ++-- 11 files changed, 126 insertions(+), 141 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 463ad2e37..b01bb9c2d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +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`, `Function`, and `GenericFunctionDecl` instead of copied into each. `Function::schema() -> &ResolvedSchema` becomes `Function::func_type() -> &FuncType` and `ResolvedSchema` is removed (its `get_by_pos` moves to `FuncType`); `GenericFunctionDecl::resolved_schema` becomes `Option>`, so an unresolved decl holds `None` rather than a placeholder `String`. +- **Breaking:** a function's signature is stored once, in `TypeInfo`, instead of copied into `Function`, `GenericFunctionDecl`, and every resolved call site. `Function::schema() -> &ResolvedSchema` becomes `Function::func_type() -> &FuncType` and `ResolvedSchema` is removed (its `get_by_pos` moves to `FuncType`); `ResolvedCall::Func` holds an `Arc`; and `GenericFunctionDecl::resolved_schema` is gone, along with the placeholder `String` an unresolved decl used to carry in it. - **Breaking:** an operation that runs external functions now takes an `ExternalContext` — a borrowed value it makes visible to them, which `ExecutionState::external_context` reads back. `Database::run_rule_set` / `with_execution_state` / `with_execution_state_tracked` and the bridge's `EGraph::run_rules` / `with_execution_state*` take one; pass `None` if there is nothing to share. egglog passes its `&TypeInfo`, which is how a primitive resolves a signature without the e-graph handing out a shared, lock-guarded copy: the borrow is live for exactly the operation that supplied it, so nothing can be declared while a rule is running. - **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). diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 9edb0131e..5241f8ecc 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -1320,16 +1320,9 @@ where { pub name: String, pub subtype: FunctionSubtype, - /// Untyped schema, as written in the source. + /// The schema as written. The resolved signature lives in `TypeInfo`, + /// keyed by `name`. pub schema: Schema, - /// The resolved signature, once typechecking has run. - /// - /// The same signature is in `TypeInfo`, keyed by `name`, and that is where - /// to read it from. This carries it for the one consumer that cannot: - /// proof encoding runs over decls that desugaring generated *after* - /// typechecking (global bindings), which no `TypeInfo` knows about until - /// they are executed. - pub resolved_schema: Option>, pub merge: Option>, pub cost: Option, pub unextractable: bool, @@ -1397,7 +1390,6 @@ impl FunctionDecl { name, subtype: FunctionSubtype::Custom, schema, - resolved_schema: None, merge, cost: None, unextractable: true, @@ -1420,7 +1412,6 @@ impl FunctionDecl { Self { name, subtype: FunctionSubtype::Constructor, - resolved_schema: None, schema, merge: None, cost, @@ -1446,7 +1437,6 @@ where name: self.name, subtype: self.subtype, schema: self.schema, - resolved_schema: self.resolved_schema, merge: self.merge.map(|expr| expr.visit_exprs(f)), cost: self.cost, unextractable: self.unextractable, diff --git a/src/ast/proof_global_remover.rs b/src/ast/proof_global_remover.rs index e0a35f74a..0c315f1da 100644 --- a/src/ast/proof_global_remover.rs +++ b/src/ast/proof_global_remover.rs @@ -37,12 +37,12 @@ fn resolved_var_to_call(var: &ResolvedVar) -> ResolvedCall { var.is_global_ref, "resolved_var_to_call called on non-global var" ); - ResolvedCall::Func(FuncType { + ResolvedCall::Func(Arc::new(FuncType { name: var.name.clone(), subtype: FunctionSubtype::Constructor, input: vec![], output: var.sort.clone(), - }) + })) } /// TODO (yz) it would be better to implement replace_global_var @@ -75,13 +75,12 @@ fn remove_globals_cmd(cmd: ResolvedNCommand) -> Vec { panic!("Global variable {} has non-eq sort {}", name, ty.name()); } - let func_type = Arc::new(FuncType { + let resolved_call = ResolvedCall::Func(Arc::new(FuncType { name: name.name.clone(), subtype: FunctionSubtype::Constructor, input: vec![], output: ty.clone(), - }); - let resolved_call = ResolvedCall::Func(func_type.as_ref().clone()); + })); let func_decl = ResolvedFunctionDecl { name: name.name, subtype: FunctionSubtype::Constructor, @@ -89,7 +88,6 @@ fn remove_globals_cmd(cmd: ResolvedNCommand) -> Vec { input: vec![], output: ty.name().to_owned(), }, - resolved_schema: Some(func_type.clone()), merge: None, cost: None, unextractable: true, diff --git a/src/ast/remove_globals.rs b/src/ast/remove_globals.rs index d9d15b81e..b0e371408 100644 --- a/src/ast/remove_globals.rs +++ b/src/ast/remove_globals.rs @@ -57,12 +57,12 @@ fn resolved_var_to_call(var: &ResolvedVar) -> ResolvedCall { var.is_global_ref, "resolved_var_to_call called on non-global var" ); - ResolvedCall::Func(FuncType { + ResolvedCall::Func(Arc::new(FuncType { name: var.name.clone(), subtype: FunctionSubtype::Custom, input: vec![], output: var.sort.clone(), - }) + })) } /// TODO (yz) it would be better to implement replace_global_var @@ -93,13 +93,12 @@ impl GlobalRemover<'_> { GenericAction::Let(span, name, expr) => { let ty = expr.output_type(); - let func_type = Arc::new(FuncType { + let resolved_call = ResolvedCall::Func(Arc::new(FuncType { name: name.name.clone(), subtype: FunctionSubtype::Custom, input: vec![], output: ty.clone(), - }); - let resolved_call = ResolvedCall::Func(func_type.as_ref().clone()); + })); let func_decl = ResolvedFunctionDecl { name: name.name, subtype: FunctionSubtype::Custom, @@ -107,7 +106,6 @@ impl GlobalRemover<'_> { input: vec![], output: ty.name().to_owned(), }, - resolved_schema: Some(func_type.clone()), merge: None, cost: None, unextractable: true, diff --git a/src/core.rs b/src/core.rs index d2e0dde6b..29abcf308 100644 --- a/src/core.rs +++ b/src/core.rs @@ -113,7 +113,9 @@ impl Hash for SpecializedPrimitive { #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub enum ResolvedCall { - Func(FuncType), + /// Shares the signature `TypeInfo` resolved, rather than copying it into + /// every resolved call site. + Func(Arc), Primitive(SpecializedPrimitive), } @@ -142,12 +144,12 @@ impl ResolvedCall { types: &[ArcSort], typeinfo: &TypeInfo, ) -> Option { - if let Some(ty) = typeinfo.get_func_type(head) { + if let Some(ty) = typeinfo.func_type_arc(head) { // As long as input types match, a result is returned. 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)); } } None @@ -160,11 +162,11 @@ impl ResolvedCall { ctx: crate::Context, span: &Span, ) -> Result { - if let Some(ty) = typeinfo.get_func_type(head) { + if let Some(ty) = typeinfo.func_type_arc(head) { 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)); } } @@ -201,6 +203,14 @@ impl ResolvedCall { } } +impl ResolvedCall { + /// Whether this call is to a `function` table, as opposed to a constructor + /// or a primitive. + pub(crate) fn is_custom_func(&self) -> bool { + matches!(self, ResolvedCall::Func(func) if func.subtype == FunctionSubtype::Custom) + } +} + impl Display for ResolvedCall { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { match self { diff --git a/src/lib.rs b/src/lib.rs index 2dfbb2073..a5843ea17 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -783,36 +783,41 @@ impl EGraph { } } - fn declare_function(&mut self, decl: &ResolvedFunctionDecl) -> Result<(), Error> { - // 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.func_type_arc(&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 - } + /// The signature for `decl`, recording it if this is the first time the + /// e-graph has seen the declaration. + /// + /// Typechecking records what it resolves, so this returns that. Desugaring + /// also generates declarations (the functions global bindings lower to, + /// proof tables) without typechecking them; resolving those here is what + /// keeps every declared function's signature reachable by name. + fn record_signature(&mut self, decl: &ResolvedFunctionDecl) -> Result, Error> { + if let Some(func_type) = self.type_info.func_type_arc(&decl.name) { + return Ok(func_type); + } + 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()); + Ok(func_type) + } + + fn declare_function(&mut self, decl: &ResolvedFunctionDecl) -> Result<(), Error> { + let func_type = self.record_signature(decl)?; let can_subsume = match decl.subtype { FunctionSubtype::Constructor => true, diff --git a/src/proofs/proof_checker.rs b/src/proofs/proof_checker.rs index daef8ef8a..4acf21677 100644 --- a/src/proofs/proof_checker.rs +++ b/src/proofs/proof_checker.rs @@ -15,7 +15,6 @@ use crate::{ }, core::ResolvedCall, proofs::proof_format::{Justification, ProofId, ProofStore, Proposition}, - typechecking::FuncType, util::{HashMap, HashSet, SymbolGen}, }; use thiserror::Error; @@ -1077,17 +1076,10 @@ impl ProofStore { // In the term representation, custom functions store output as last arg: f(args..., v) ResolvedFact::Eq( _, - ResolvedExpr::Call( - _, - ResolvedCall::Func(FuncType { - subtype: FunctionSubtype::Custom, - name, - .. - }), - args, - ), + ResolvedExpr::Call(_, call @ ResolvedCall::Func(_), args), ResolvedExpr::Var(_, v), - ) => { + ) if call.is_custom_func() => { + let name = call.name(); // Get the output variable's term let var_term = subst_with_globals.get(&v.name).copied().ok_or_else(|| { ProofCheckErrorKind::UnboundVariable { @@ -1113,7 +1105,7 @@ impl ProofStore { // Add the output variable as the last argument arg_terms.push(var_term); - let expected_term_id = self.term_dag.app(name.clone(), arg_terms); + let expected_term_id = self.term_dag.app(name.to_owned(), arg_terms); // The proposition should be a reflexive equality for this term if lhs != expected_term_id || rhs != expected_term_id { diff --git a/src/proofs/proof_encoding.rs b/src/proofs/proof_encoding.rs index ac57cda56..fab90bcb5 100644 --- a/src/proofs/proof_encoding.rs +++ b/src/proofs/proof_encoding.rs @@ -60,29 +60,6 @@ pub(crate) struct ProofInstrumentor<'a> { pub(crate) egraph: &'a mut EGraph, } -/// The signature typechecking resolved for `fdecl`. -/// -/// Encoding also runs over decls desugaring generated after typechecking, which -/// no `TypeInfo` knows about yet, so it reads the signature off the decl rather -/// than looking it up by name. -fn resolved_signature(fdecl: &ResolvedFunctionDecl) -> &FuncType { - fdecl - .resolved_schema - .as_ref() - .unwrap_or_else(|| panic!("{} reached encoding unresolved", fdecl.name)) -} - -/// The declared sorts of `fdecl`'s columns, inputs followed by the output. -fn column_sorts(fdecl: &ResolvedFunctionDecl) -> Vec { - let func_type = resolved_signature(fdecl); - func_type - .input - .iter() - .chain([&func_type.output]) - .cloned() - .collect() -} - impl<'a> ProofInstrumentor<'a> { /// Make a term state and use it to instrument the code. pub(crate) fn add_term_encoding( @@ -578,9 +555,50 @@ impl<'a> ProofInstrumentor<'a> { "".to_string() } + /// The sort named `name`. + /// + /// Encoding runs over the program the pre-instrumentation typechecking + /// resolved, so that is where to look first — the same way a container + /// sort is resolved. + fn sort_by_name(&self, name: &str) -> ArcSort { + self.egraph + .proof_state + .original_typechecking + .as_ref() + .and_then(|typechecking| typechecking.get_sort_by_name(name)) + .or_else(|| self.egraph.get_sort_by_name(name)) + .cloned() + .unwrap_or_else(|| panic!("sort {name} not found while encoding")) + } + + /// The declared sorts of `fdecl`'s columns, inputs followed by the output. + /// + /// Resolved from the declaration's own schema rather than looked up by + /// name: encoding also runs over declarations desugaring generated, which + /// are not typechecked until after this pass. + fn column_sorts(&self, fdecl: &ResolvedFunctionDecl) -> Vec { + fdecl + .schema + .input + .iter() + .chain([&fdecl.schema.output]) + .map(|name| self.sort_by_name(name)) + .collect() + } + + /// The declared sorts of `fdecl`'s input columns. + fn input_sorts(&self, fdecl: &ResolvedFunctionDecl) -> Vec { + fdecl + .schema + .input + .iter() + .map(|name| self.sort_by_name(name)) + .collect() + } + /// Rules that update the views when children change. fn rebuilding_rules(&mut self, fdecl: &ResolvedFunctionDecl) -> Vec { - let types = column_sorts(fdecl); + let types = self.column_sorts(fdecl); let proofs_enabled = self.egraph.proof_state.proofs_enabled; // Check if there are any rebuildable columns at all; if not, no rule needed. @@ -757,7 +775,7 @@ impl<'a> ProofInstrumentor<'a> { /// 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 fn rebuilding_subsumed_rules(&mut self, fdecl: &ResolvedFunctionDecl) -> Vec { - let input = &resolved_signature(fdecl).input; + let input = &self.input_sorts(fdecl); // Check if there are any eq-sort columns at all; if not, no rebuild rule needed. if !input.iter().any(|t| t.is_eq_sort()) { @@ -844,17 +862,10 @@ impl<'a> ProofInstrumentor<'a> { // In proof normal form, this is the only way that function calls appear. ResolvedFact::Eq( _span, - ResolvedExpr::Call( - _span2, - head @ ResolvedCall::Func(FuncType { - subtype: FunctionSubtype::Custom, - .. - }), - args, - ), + ResolvedExpr::Call(_span2, head @ ResolvedCall::Func(_), args), // TODO this could actually be arbitrary pretty easily, it's just nested functions that are hard. ResolvedExpr::Var(_span3, v), - ) => { + ) if head.is_custom_func() => { let mut new_args = vec![]; let mut arg_proofs = vec![]; for arg in args { diff --git a/src/proofs/proof_format.rs b/src/proofs/proof_format.rs index 9861de45e..b8730a137 100644 --- a/src/proofs/proof_format.rs +++ b/src/proofs/proof_format.rs @@ -1,8 +1,8 @@ use crate::{ ResolvedCall, Term, TermDag, TermId, - ast::{FunctionSubtype, ResolvedExpr, ResolvedFact, ResolvedNCommand}, + ast::{ResolvedExpr, ResolvedFact, ResolvedNCommand}, proofs::{proof_checker::gather_globals, proof_encoding_helpers::EncodingNames}, - typechecking::{FuncType, PrimitiveValidator}, + typechecking::PrimitiveValidator, util::{HEntry, HashMap, IndexSet, SymbolGen}, }; use egglog_ast::generic_ast::Literal; @@ -598,16 +598,9 @@ impl ProofStore { // In proof normal form, this is the only way that function calls apppear. ResolvedFact::Eq( _span, - ResolvedExpr::Call( - _span2, - head @ ResolvedCall::Func(FuncType { - subtype: FunctionSubtype::Custom, - .. - }), - args, - ), + ResolvedExpr::Call(_span2, head @ ResolvedCall::Func(_), args), ResolvedExpr::Var(_span3, v), - ) => { + ) if head.is_custom_func() => { let term = proof.rhs(); let children = match self.term_dag.get(term) { Term::App(head_name, children) if head_name == head.name() => children.clone(), diff --git a/src/proofs/proof_normal_form.rs b/src/proofs/proof_normal_form.rs index 39a8f9c9b..0c982b2ac 100644 --- a/src/proofs/proof_normal_form.rs +++ b/src/proofs/proof_normal_form.rs @@ -1,6 +1,6 @@ -use crate::ast::{FunctionSubtype, ResolvedNCommand}; +use crate::ast::ResolvedNCommand; +use crate::core::ResolvedCall; use crate::*; -use crate::{core::ResolvedCall, typechecking::FuncType}; use egglog_ast::generic_ast::GenericExpr; /// Transforms queries into "proof normal form" by lifting subexpressions to the @@ -45,16 +45,9 @@ fn proof_form_fact( match fact { ResolvedFact::Eq( span, - ResolvedExpr::Call( - span2, - head @ ResolvedCall::Func(FuncType { - subtype: FunctionSubtype::Custom, - .. - }), - args, - ), + ResolvedExpr::Call(span2, head @ ResolvedCall::Func(_), args), ResolvedExpr::Var(span3, v), - ) => { + ) if head.is_custom_func() => { let mut new_args = vec![]; for arg in args { new_args.push(proof_form_expr(arg, res, fresh)); @@ -84,13 +77,9 @@ fn proof_form_expr( match fact { ref fact @ ResolvedExpr::Call( ref span, - ref head @ ResolvedCall::Func(FuncType { - subtype: FunctionSubtype::Custom, - ref output, - .. - }), + ref head @ ResolvedCall::Func(ref func_type), ref args, - ) => { + ) if head.is_custom_func() => { // bind this to a new variable let new_args = args .iter() @@ -100,7 +89,7 @@ fn proof_form_expr( span.clone(), ResolvedVar { name: fresh.fresh("n"), - sort: output.clone(), + sort: func_type.output.clone(), is_global_ref: false, }, ); @@ -131,7 +120,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(ref func_type), ref inner_args, ) => { // First recursively normalize the inner arguments @@ -145,7 +134,7 @@ fn proof_form_expr( arg_span.clone(), ResolvedVar { name: fresh.fresh("v"), - sort: output.clone(), + sort: func_type.output.clone(), is_global_ref: false, }, ); diff --git a/src/typechecking.rs b/src/typechecking.rs index 8e74c8419..8abd36b8e 100644 --- a/src/typechecking.rs +++ b/src/typechecking.rs @@ -584,7 +584,10 @@ impl EGraph { span.clone(), )); } - ResolvedNCommand::ProveExists(span.clone(), ResolvedCall::Func(func_type.clone())) + ResolvedNCommand::ProveExists( + span.clone(), + ResolvedCall::Func(Arc::new(func_type.clone())), + ) } NCommand::Output { span, file, exprs } => { let exprs = exprs @@ -821,10 +824,6 @@ impl TypeInfo { name: fdecl.name.clone(), subtype: fdecl.subtype, schema: fdecl.schema.clone(), - resolved_schema: Some( - self.func_type_arc(&fdecl.name) - .expect("the signature was just recorded"), - ), 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 From 5fc5050fa226dddd3847858d1276792caceaf2dc Mon Sep 17 00:00:00 2001 From: Oliver Flatt Date: Fri, 7 Aug 2026 22:37:52 +0000 Subject: [PATCH 8/8] Review pass: tidy the diff's docs, cover the new API, fix two defects Two reviews, one of them independent, converged on the same two defects: - Deleting `ResolvedCall::view_types` left its `///` block behind, so the public `from_resolution_func_types` was documented as a different function. - `func_type_of` reported `ApiError::MissingTable` for a table that exists whenever the execution carried no context, conflating "no declarations here" with "no such table". Now `ApiError::SchemasUnavailable`. Coverage: the codecov report on the PR put `src/exec_state.rs` at 0% for this patch - none of the new accessors were exercised in-tree, only from egglog-experimental. `tests/api_introspection.rs` covers them, pinning `enodes_for_eclass` against a filtered `constructor_enodes` scan (with a row-count assertion so the comparison cannot pass vacuously), the subtype split of the schema accessors, and `rebuild_container`'s three outcomes. `core-relations` covers the context mechanism itself: an external function reads back what the caller supplied, `None` when nothing was, and `None` rather than a misparse when the type does not match. Also from review: `record_signature` now debug-asserts a cached signature agrees with the declaration it is reused for, `enodes_for_eclass` documents that it matches the eclass column as stored, `is_custom_func` moved into the existing inherent impl, the `context` parameters are documented on the entry points that gained them, and `FuncType::get_by_pos` - carried over from the deleted `ResolvedSchema` with no callers in the workspace - is dropped. Doc tidy per the repo's tidy-diff-docs skill: five comments stated mechanism instead of contract (the lazy-column-index narration on `for_each_matching_col` and `for_each_output_value`, the rationale paragraph on `ExternalContext`, the borrow-checker explanation on `type_info`, and "shares rather than copies" on two `Arc` fields, which the type already says). Co-Authored-By: Claude Opus 5 --- core-relations/src/action/mod.rs | 13 +- core-relations/src/free_join/execute.rs | 4 + core-relations/src/free_join/mod.rs | 7 +- core-relations/src/tests.rs | 37 +++++ egglog-bridge/src/lib.rs | 12 +- src/api/mod.rs | 7 + src/core.rs | 6 - src/exec_state.rs | 21 +-- src/lib.rs | 13 +- src/typechecking.rs | 12 -- tests/api_introspection.rs | 180 ++++++++++++++++++++++++ 11 files changed, 266 insertions(+), 46 deletions(-) create mode 100644 tests/api_introspection.rs diff --git a/core-relations/src/action/mod.rs b/core-relations/src/action/mod.rs index d01089a64..96096ded6 100644 --- a/core-relations/src/action/mod.rs +++ b/core-relations/src/action/mod.rs @@ -340,10 +340,8 @@ pub(crate) struct DbView<'a> { /// created for one operation, for its [`ExternalFunction`]s to read back with /// [`ExecutionState::external_context`]. /// -/// Passing it per operation, rather than storing it in the database, is what -/// lets an external function see borrowed embedder state: the borrow is live -/// for exactly the call that supplied it, so the embedder cannot mutate that -/// state while a rule is running. +/// It is borrowed for exactly the operation that supplied it, so an embedder +/// cannot mutate the state it shared while that operation runs. /// /// [`ExternalFunction`]: crate::ExternalFunction pub type ExternalContext<'a> = Option<&'a (dyn Any + Send + Sync)>; @@ -505,11 +503,8 @@ 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. + /// Call `f` on each visible row in `table` whose `col` equals `value`, + /// with the whole row as a value slice. pub fn for_each_matching_col( &self, table: TableId, diff --git a/core-relations/src/free_join/execute.rs b/core-relations/src/free_join/execute.rs index ae1143559..bbef288f6 100644 --- a/core-relations/src/free_join/execute.rs +++ b/core-relations/src/free_join/execute.rs @@ -431,6 +431,10 @@ impl Prober { } impl Database { + /// Run `rule_set` to completion. + /// + /// `context` is visible to any external function the rules reach; pass + /// `None` if there is nothing to share. pub fn run_rule_set( &mut self, rule_set: &RuleSet, diff --git a/core-relations/src/free_join/mod.rs b/core-relations/src/free_join/mod.rs index 7ae6e9b04..11a2e0ea8 100644 --- a/core-relations/src/free_join/mod.rs +++ b/core-relations/src/free_join/mod.rs @@ -457,9 +457,10 @@ impl Database { f(&mut state) } - /// Like [`Database::with_execution_state`], but also reports whether `f` - /// staged any mutation through the execution state. Callers can use the - /// flag to skip a subsequent `merge_all` when the closure was read-only. + /// Like [`Database::with_execution_state`], including its `context`, but + /// also reports whether `f` staged any mutation through the execution + /// state. Callers can use the flag to skip a subsequent `merge_all` when + /// the closure was read-only. pub fn with_execution_state_tracked( &self, context: ExternalContext<'_>, diff --git a/core-relations/src/tests.rs b/core-relations/src/tests.rs index 6d6e2282d..b2587fe20 100644 --- a/core-relations/src/tests.rs +++ b/core-relations/src/tests.rs @@ -1292,3 +1292,40 @@ fn early_stop_inner() { "External function called {final_count} times, should be much less than 10k" ); } + +/// An external function sees the [`ExternalContext`] the caller of the +/// operation supplied, and `None` when the caller supplied none. +#[test] +fn external_context_reaches_external_functions() { + use std::sync::{Arc, Mutex}; + + let mut db = Database::new(); + // What the external function read back, per invocation. + let seen: Arc>>> = Default::default(); + let recorder = seen.clone(); + let read_context = db.add_external_function(Box::new(make_external_func( + move |exec_state: &mut crate::action::ExecutionState, _args: &[Value]| { + let context = exec_state + .external_context() + .and_then(|context| context.downcast_ref::()) + .copied(); + recorder.lock().unwrap().push(context); + Some(Value::new(0)) + }, + ))); + + let shared: i64 = 7; + db.with_execution_state(Some(&shared), |state| { + state.call_external_func(read_context, &[]); + }); + db.with_execution_state(None, |state| { + state.call_external_func(read_context, &[]); + }); + // A context of the wrong type reads back as absent rather than misparsed. + let wrong_type: u8 = 7; + db.with_execution_state(Some(&wrong_type), |state| { + state.call_external_func(read_context, &[]); + }); + + assert_eq!(*seen.lock().unwrap(), vec![Some(7), None, None]); +} diff --git a/egglog-bridge/src/lib.rs b/egglog-bridge/src/lib.rs index dceb68303..c25fe8594 100644 --- a/egglog-bridge/src/lib.rs +++ b/egglog-bridge/src/lib.rs @@ -656,6 +656,9 @@ impl EGraph { /// Run the given rules, returning whether the database changed. /// + /// `context` is visible to any external function the rules reach; pass + /// `None` if there is nothing to share. + /// /// If the given rules are malformed, this method can return an error. pub fn run_rules( &mut self, @@ -1023,6 +1026,9 @@ impl EGraph { /// The staged updates are not immediately reflected in the EGraph, so you may want to /// manually flush the updates using [`EGraph::flush_updates`]. /// + /// `context` is visible to any external function the closure reaches; pass + /// `None` if there is nothing to share. + /// /// # Seminaive-safety trust boundary /// /// This method hands out a raw `&mut ExecutionState`, which bypasses @@ -1478,11 +1484,7 @@ 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. + /// Call `f` on each row whose output column is `value`. pub fn for_each_output_value( &self, state: &ExecutionState, diff --git a/src/api/mod.rs b/src/api/mod.rs index 6af535d62..d19da2929 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -36,6 +36,13 @@ use thiserror::Error; /// those, see [`enum@crate::Error`]. #[derive(Debug, Error)] pub enum ApiError { + /// The execution was created without the e-graph's declarations, so a + /// table's signature cannot be resolved. Distinct from a missing table: + /// nothing was looked up. + #[error( + "declarations are not available in this execution, so the signature of `{name}` cannot be read" + )] + SchemasUnavailable { name: String }, #[error("no table named `{name}` is registered")] MissingTable { name: String }, diff --git a/src/core.rs b/src/core.rs index 29abcf308..ba7298854 100644 --- a/src/core.rs +++ b/src/core.rs @@ -113,8 +113,6 @@ impl Hash for SpecializedPrimitive { #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub enum ResolvedCall { - /// Shares the signature `TypeInfo` resolved, rather than copying it into - /// every resolved call site. Func(Arc), Primitive(SpecializedPrimitive), } @@ -134,8 +132,6 @@ impl ResolvedCall { } } - /// Gives the types for a term's child with the given resolved call. - /// For functions this includes the output sort, for constructors it's just the inputs. // Different from `from_resolution`, this function only considers function types and not primitives. // As a result, it only requires input argument types, so types.len() == func.input.len(), // while for `from_resolution`, types.len() == func.input.len() + 1 to account for the output type @@ -201,9 +197,7 @@ impl ResolvedCall { span: span.clone(), }) } -} -impl ResolvedCall { /// Whether this call is to a `function` table, as opposed to a constructor /// or a primitive. pub(crate) fn is_custom_func(&self) -> bool { diff --git a/src/exec_state.rs b/src/exec_state.rs index fd9291002..d55143cf3 100644 --- a/src/exec_state.rs +++ b/src/exec_state.rs @@ -149,10 +149,7 @@ pub(crate) trait RegistrySealed<'a, 'db: 'a>: Internal<'a, 'db> { /// in an execution the e-graph did not supply one for (its internal /// rebuild rules, whose actions never read a signature). /// - /// This is a borrow rather than a shared handle, which is what keeps a - /// declaration from racing a running rule: the e-graph passes `&TypeInfo` - /// into the execution, so it cannot be `&mut` borrowed to declare anything - /// until that execution ends. + /// Borrowed for the operation, so nothing can be declared while it runs. fn type_info(&self) -> Option<&'db TypeInfo> { self.es().external_context()?.downcast_ref() } @@ -460,11 +457,9 @@ pub trait Read<'a, 'db: 'a>: Core<'a, 'db> + RegistrySealed<'a, 'db> { } /// 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. + /// eclass column holds exactly `eclass`. Rows whose eclass has since been + /// merged into another are matched under the id they store, not their + /// canonical one. Errors with `WrongSubtype` if `name` is a function. fn enodes_for_eclass( &self, name: &str, @@ -645,7 +640,13 @@ fn func_type_of<'db>( name: &str, expected: FunctionSubtype, ) -> Result<&'db FuncType, Error> { - let Some(func_type) = type_info.and_then(|type_info| type_info.get_func_type(name)) else { + let Some(type_info) = type_info else { + return Err(ApiError::SchemasUnavailable { + name: name.to_string(), + } + .into()); + }; + let Some(func_type) = type_info.get_func_type(name) else { return Err(ApiError::MissingTable { name: name.to_string(), } diff --git a/src/lib.rs b/src/lib.rs index a5843ea17..0efcb79c3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -327,7 +327,6 @@ pub trait UserDefinedCommand: Send + Sync { #[derive(Clone)] pub struct Function { decl: ResolvedFunctionDecl, - /// Shared with [`TypeInfo`], which is where a signature is resolved. func_type: Arc, can_subsume: bool, backend_id: egglog_bridge::FunctionId, @@ -792,6 +791,18 @@ impl EGraph { /// keeps every declared function's signature reachable by name. fn record_signature(&mut self, decl: &ResolvedFunctionDecl) -> Result, Error> { if let Some(func_type) = self.type_info.func_type_arc(&decl.name) { + debug_assert!( + func_type.subtype == decl.subtype + && func_type.input.len() == decl.schema.input.len() + && func_type + .input + .iter() + .zip(&decl.schema.input) + .all(|(sort, name)| sort.name() == name) + && func_type.output.name() == decl.schema.output, + "recorded signature for {} disagrees with its declaration", + decl.name + ); return Ok(func_type); } let get_sort = |name: &String| match self.type_info.get_sort_by_name(name) { diff --git a/src/typechecking.rs b/src/typechecking.rs index 8abd36b8e..2a31ea750 100644 --- a/src/typechecking.rs +++ b/src/typechecking.rs @@ -119,18 +119,6 @@ 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 diff --git a/tests/api_introspection.rs b/tests/api_introspection.rs new file mode 100644 index 000000000..27b4f4169 --- /dev/null +++ b/tests/api_introspection.rs @@ -0,0 +1,180 @@ +//! Tests for the e-graph introspection a primitive body can reach: +//! `Read::enodes_for_eclass`, `Read::constructor_schema` / +//! `function_schema` / `table_subtype`, and `Core::rebuild_container`. + +use egglog::prelude::*; +use egglog::sort::{MapContainer, VecContainer}; +use egglog::{ApiError, Core, Error, Read, Value}; +use std::any::TypeId; + +const MATH: &str = " +(datatype Math + (Num i64) + (Var String) + (Add Math Math)) +(function cost (Math) i64 :no-merge) +"; + +/// `enodes_for_eclass` returns exactly the rows a full scan would, filtered on +/// the eclass column — that equivalence is the invariant the indexed path has +/// to preserve. +#[test] +fn enodes_for_eclass_agrees_with_a_filtered_scan() -> Result<(), Error> { + let mut egraph = EGraph::default(); + egraph.parse_and_run_program( + None, + &format!( + "{MATH} +(let $a (Add (Num 1) (Num 2))) +(let $b (Add (Num 2) (Num 1))) +(union $a $b) +(let $c (Add (Num 9) (Num 9))) +" + ), + )?; + + // Every eclass `Add` produces, so the comparison covers a class with two + // e-nodes and one with a single e-node. + let mut eclasses: Vec = Vec::new(); + egraph.constructor_enodes("Add", |enode| eclasses.push(enode.eclass))?; + eclasses.dedup(); + assert!(eclasses.len() >= 2, "expected several Add eclasses"); + + let mut total_indexed = 0; + for eclass in eclasses { + let mut scanned: Vec> = Vec::new(); + egraph.constructor_enodes("Add", |enode| { + if enode.eclass == eclass { + scanned.push(enode.children.to_vec()); + } + })?; + + let mut indexed: Vec> = Vec::new(); + egraph.read(|state| { + state.enodes_for_eclass("Add", eclass, |enode| { + indexed.push(enode.children.to_vec()); + }) + })?; + + scanned.sort(); + indexed.sort(); + assert_eq!(indexed, scanned, "mismatch for eclass {eclass:?}"); + total_indexed += indexed.len(); + } + + // Guards against both sides being vacuously empty: every row of the table + // has to have been reached through the indexed lookup. + let rows = egraph.update(|state| Ok(state.table_size("Add")))?.unwrap(); + assert_eq!(total_indexed, rows, "indexed lookup missed rows"); + Ok(()) +} + +#[test] +fn enodes_for_eclass_rejects_a_function_table() -> Result<(), Error> { + let mut egraph = EGraph::default(); + egraph.parse_and_run_program(None, MATH)?; + let err = egraph + .read(|state| state.enodes_for_eclass("cost", Value::new_const(0), |_| {})) + .unwrap_err(); + assert!( + matches!(err, Error::ApiError(ApiError::WrongSubtype { .. })), + "expected WrongSubtype, got {err}" + ); + Ok(()) +} + +/// The schema accessors report a table's declared signature, and split by +/// subtype the way the rest of `Read` does. +#[test] +fn schema_accessors_report_the_declaration() -> Result<(), Error> { + let mut egraph = EGraph::default(); + egraph.parse_and_run_program(None, MATH)?; + + egraph.read(|state| { + let add = state.constructor_schema("Add").unwrap(); + let input: Vec<&str> = add.input.iter().map(|sort| sort.name()).collect(); + assert_eq!(input, ["Math", "Math"]); + assert_eq!(add.output.name(), "Math"); + + let cost = state.function_schema("cost").unwrap(); + assert_eq!(cost.output.name(), "i64"); + + // Each rejects the other's subtype rather than answering. + assert!(matches!( + state.constructor_schema("cost"), + Err(Error::ApiError(ApiError::WrongSubtype { .. })) + )); + assert!(matches!( + state.function_schema("Add"), + Err(Error::ApiError(ApiError::WrongSubtype { .. })) + )); + assert!(matches!( + state.constructor_schema("nonesuch"), + Err(Error::ApiError(ApiError::MissingTable { .. })) + )); + + // `table_subtype` answers for either subtype without erroring. + assert_eq!( + state.table_subtype("Add"), + Some(egglog::ast::FunctionSubtype::Constructor) + ); + assert_eq!( + state.table_subtype("cost"), + Some(egglog::ast::FunctionSubtype::Custom) + ); + assert_eq!(state.table_subtype("nonesuch"), None); + }); + Ok(()) +} + +/// `rebuild_container` remaps a container's contents and interns the result, +/// leaving the original alone. +#[test] +fn rebuild_container_remaps_contents() -> Result<(), Error> { + let mut egraph = EGraph::default(); + egraph.parse_and_run_program( + None, + &format!( + "{MATH} +(sort MathVec (Vec Math)) +(let $one (Num 1)) +(let $two (Num 2)) +(let $vec (vec-of $one $two)) +(let $swapped (vec-of $two $one)) +" + ), + )?; + + let one = egraph.eval_expr(&exprs::var("$one"))?.1; + let two = egraph.eval_expr(&exprs::var("$two"))?.1; + let vec = egraph.eval_expr(&exprs::var("$vec"))?.1; + let swapped = egraph.eval_expr(&exprs::var("$swapped"))?.1; + + let rebuilt = egraph.update(|mut state| { + Ok( + state.rebuild_container(TypeId::of::(), vec, &|value| { + if value == one { + two + } else if value == two { + one + } else { + value + } + }), + ) + })?; + assert_eq!(rebuilt, swapped, "the swap should intern to the same vec"); + + // A remap that changes nothing hands back the original value. + let unchanged = egraph.update(|mut state| { + Ok(state.rebuild_container(TypeId::of::(), vec, &|value| value)) + })?; + assert_eq!(unchanged, vec); + + // A value that is not a container of that type is returned untouched. + let not_a_container = egraph.update(|mut state| { + Ok(state.rebuild_container(TypeId::of::(), vec, &|value| value)) + })?; + assert_eq!(not_a_container, vec); + Ok(()) +}