diff --git a/Cargo.lock b/Cargo.lock index 2e62ff8..ccafc00 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -330,7 +330,7 @@ checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" [[package]] name = "egglog" version = "2.0.0" -source = "git+https://github.com/egraphs-good/egglog.git?rev=53b9721b9706741edff2b7c1e379fc37137184d9#53b9721b9706741edff2b7c1e379fc37137184d9" +source = "git+https://github.com/oflatt-claude/egglog.git?rev=5fc5050fa226dddd3847858d1276792caceaf2dc#5fc5050fa226dddd3847858d1276792caceaf2dc" dependencies = [ "chrono", "clap", @@ -362,7 +362,7 @@ dependencies = [ [[package]] name = "egglog-add-primitive" version = "2.0.0" -source = "git+https://github.com/egraphs-good/egglog.git?rev=53b9721b9706741edff2b7c1e379fc37137184d9#53b9721b9706741edff2b7c1e379fc37137184d9" +source = "git+https://github.com/oflatt-claude/egglog.git?rev=5fc5050fa226dddd3847858d1276792caceaf2dc#5fc5050fa226dddd3847858d1276792caceaf2dc" dependencies = [ "quote", "syn 2.0.108", @@ -371,7 +371,7 @@ dependencies = [ [[package]] name = "egglog-ast" version = "2.0.0" -source = "git+https://github.com/egraphs-good/egglog.git?rev=53b9721b9706741edff2b7c1e379fc37137184d9#53b9721b9706741edff2b7c1e379fc37137184d9" +source = "git+https://github.com/oflatt-claude/egglog.git?rev=5fc5050fa226dddd3847858d1276792caceaf2dc#5fc5050fa226dddd3847858d1276792caceaf2dc" dependencies = [ "ordered-float", ] @@ -379,7 +379,7 @@ dependencies = [ [[package]] name = "egglog-bridge" version = "2.0.0" -source = "git+https://github.com/egraphs-good/egglog.git?rev=53b9721b9706741edff2b7c1e379fc37137184d9#53b9721b9706741edff2b7c1e379fc37137184d9" +source = "git+https://github.com/oflatt-claude/egglog.git?rev=5fc5050fa226dddd3847858d1276792caceaf2dc#5fc5050fa226dddd3847858d1276792caceaf2dc" dependencies = [ "anyhow", "dyn-clone", @@ -402,7 +402,7 @@ dependencies = [ [[package]] name = "egglog-concurrency" version = "2.0.0" -source = "git+https://github.com/egraphs-good/egglog.git?rev=53b9721b9706741edff2b7c1e379fc37137184d9#53b9721b9706741edff2b7c1e379fc37137184d9" +source = "git+https://github.com/oflatt-claude/egglog.git?rev=5fc5050fa226dddd3847858d1276792caceaf2dc#5fc5050fa226dddd3847858d1276792caceaf2dc" dependencies = [ "arc-swap", "bumpalo", @@ -414,7 +414,7 @@ dependencies = [ [[package]] name = "egglog-core-relations" version = "2.0.0" -source = "git+https://github.com/egraphs-good/egglog.git?rev=53b9721b9706741edff2b7c1e379fc37137184d9#53b9721b9706741edff2b7c1e379fc37137184d9" +source = "git+https://github.com/oflatt-claude/egglog.git?rev=5fc5050fa226dddd3847858d1276792caceaf2dc#5fc5050fa226dddd3847858d1276792caceaf2dc" dependencies = [ "anyhow", "bumpalo", @@ -456,12 +456,12 @@ dependencies = [ [[package]] name = "egglog-numeric-id" version = "2.0.0" -source = "git+https://github.com/egraphs-good/egglog.git?rev=53b9721b9706741edff2b7c1e379fc37137184d9#53b9721b9706741edff2b7c1e379fc37137184d9" +source = "git+https://github.com/oflatt-claude/egglog.git?rev=5fc5050fa226dddd3847858d1276792caceaf2dc#5fc5050fa226dddd3847858d1276792caceaf2dc" [[package]] name = "egglog-reports" version = "2.0.0" -source = "git+https://github.com/egraphs-good/egglog.git?rev=53b9721b9706741edff2b7c1e379fc37137184d9#53b9721b9706741edff2b7c1e379fc37137184d9" +source = "git+https://github.com/oflatt-claude/egglog.git?rev=5fc5050fa226dddd3847858d1276792caceaf2dc#5fc5050fa226dddd3847858d1276792caceaf2dc" dependencies = [ "clap", "hashbrown 0.16.0", @@ -475,7 +475,7 @@ dependencies = [ [[package]] name = "egglog-union-find" version = "2.0.0" -source = "git+https://github.com/egraphs-good/egglog.git?rev=53b9721b9706741edff2b7c1e379fc37137184d9#53b9721b9706741edff2b7c1e379fc37137184d9" +source = "git+https://github.com/oflatt-claude/egglog.git?rev=5fc5050fa226dddd3847858d1276792caceaf2dc#5fc5050fa226dddd3847858d1276792caceaf2dc" dependencies = [ "crossbeam", "egglog-concurrency", diff --git a/Cargo.toml b/Cargo.toml index 2647184..92d7ccf 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,9 +17,13 @@ default = ["bin"] bin = ["egglog/bin"] [dependencies] -egglog = { git = "https://github.com/egraphs-good/egglog.git", rev = "53b9721b9706741edff2b7c1e379fc37137184d9", default-features = false } -egglog-ast = { git = "https://github.com/egraphs-good/egglog.git", rev = "53b9721b9706741edff2b7c1e379fc37137184d9", default-features = false } -egglog-reports = { git = "https://github.com/egraphs-good/egglog.git", rev = "53b9721b9706741edff2b7c1e379fc37137184d9", default-features = false } +# Points at egraphs-good/egglog#986, which adds the e-graph introspection +# `unstable-subst` is built on (Read::enodes_for_eclass, Read::table_schema / +# table_subtype, Core::rebuild_container). Move back to an egraphs-good rev once +# that PR merges. +egglog = { git = "https://github.com/oflatt-claude/egglog.git", rev = "5fc5050fa226dddd3847858d1276792caceaf2dc", default-features = false } +egglog-ast = { git = "https://github.com/oflatt-claude/egglog.git", rev = "5fc5050fa226dddd3847858d1276792caceaf2dc", default-features = false } +egglog-reports = { git = "https://github.com/oflatt-claude/egglog.git", rev = "5fc5050fa226dddd3847858d1276792caceaf2dc", default-features = false } num = "0.4.3" lazy_static = "1.4" diff --git a/SUBST_DESIGN.md b/SUBST_DESIGN.md new file mode 100644 index 0000000..1f0beb9 --- /dev/null +++ b/SUBST_DESIGN.md @@ -0,0 +1,177 @@ +# `unstable-subst`: a substitution primitive + +Status: exploration. Lives in `egglog-experimental`; the e-graph introspection +it needs lives in `egglog` (see "What egglog had to expose" below). + +## What it does + +``` +(unstable-subst root map) : (R, Map) -> R +``` + +`root` is an e-class of any eq-sort `R`; `map` is a `Map` whose key and value +sorts are the same eq-sort `K`. The primitive walks the sub-e-graph reachable +from `root`, copies the part of it that the substitution actually touches while +replacing every occurrence of a key e-class with its mapped value, and returns +the e-class of the copied root. + +Reachability follows **constructor** rows only (the term structure), never +`function` rows (those are analyses over the structure, not part of it). +Container-valued children (`Vec`, `Map`, `Set`, ...) are followed into and +rebuilt with substituted contents. + +## Semantics + +Let `σ` be the resulting map on values. + +* An e-class is **affected** if it is a key of `map`, or if one of its e-nodes + has an affected child (least fixpoint over the reachable subgraph). + Container values are affected if any of their contents are. +* `σ(v) = map[v]` for keys, `σ(v) = v` for unaffected `v`. +* For an affected non-key e-class `e`, every e-node `f(c1..cn) -> e` in the + snapshot is copied as `f(σ(c1)..σ(cn))`; all the copies are unioned together + and `σ(e)` is that class. Copying every e-node, not just the ones that change, + is what carries the region's equations over to the copy — see below. +* The return value is `σ(root)`. + +Consequences worth stating out loud: + +* **Nothing is copied when nothing changes.** An unaffected sub-e-graph is + shared with the original, so `(unstable-subst e (map-empty))` is exactly `e` + and allocates nothing. +* **The region's equations are substituted too, not just its terms.** An + e-class is a set of terms known equal, so copying it copies every one of its + e-nodes: if the e-graph knows `t1 = t2` and both are reachable, the copy + asserts `σ(t1) = σ(t2)`. An e-node with no substituted children copies to + itself, so `lookup_or_insert` finds the original row and the copy merges back + into the original class. + + That is what you want when the equations come from rewrite rules, which hold + for every value of the substituted classes. With + `(rewrite (Mul a (Num 0)) (Num 0))`, the class of `(Mul x (Num 0))` also holds + `(Num 0)`; substituting `x := 5` merges the copy back into that class and + returns it, and `5 * 0` really is `0`. + + It is wrong when the e-graph holds a **ground** equation pinning a substituted + class down. `(union (Add x (Num 1)) (Num 5))` asserts `x = 4`; substituting + `x := 9` copies the untouched `(Num 5)` unchanged, merges, and thereby asserts + `9 + 1 = 5`. So only substitute classes that behave like universally + quantified variables — being singleton is neither necessary nor sufficient. + (`tests/subst.rs` pins both directions.) + + Copying every e-node rather than only the ones that change is deliberate. + Copying only the changed e-nodes would never merge back into the original + class, but it would also drop the region's equations from the copy: the + `x * 0` result above would come back as a bare `(Mul (Num 5) (Num 0))` not + known to equal `0`. A running rule set would re-derive that; a one-shot + substitution would not. +* **Only committed rows are visible.** The walk reads tables, and an action's + writes are staged until the action finishes, so a term built in the same + action comes back unsubstituted. Terms from earlier commands and earlier rule + iterations are fine — which is what the `:naive` beta-reduction shape needs — + but `(unstable-subst (Mul x y) m)` in one action does not work. +* **No e-class id is ever invented.** Every copied e-node goes in through + `lookup_or_insert`, exactly as `(Add a b)` in an action does, so egglog names + the copy. The consequence is the cycle rule below. +* **Grounded cycles are copied; ungrounded ones error.** A cyclic e-class can + only be copied if one of its e-nodes has all its children outside the cycle: + that e-node's insert names the copy, and the cyclic e-node then unions into + it. `x = {Var "x", Add x (Num 0)}` qualifies and works. A cycle in which every + e-node points back into the cycle has no such starting point, and naming its + copy would mean inventing an e-class id — so it returns + `Error::SubstError` (`egglog_experimental::subst`) or panics with that reason (the + primitive) rather than producing a partial copy. +* **Subsumed e-nodes are skipped** — they are excluded from extraction, so + resurrecting them un-subsumed in a copy would be wrong. +* Registered by `new_experimental_egraph`, so a plain `EGraph::default` does not + have it. +* The snapshot is taken from live table contents, so `unstable-subst` is only + available where reads *and* writes are legal: top-level actions (`let`, + `eval`, action-mode `run-schedule`) and the head of a `:naive` rule. This is + the `Context::Full` capability, enforced by the typechecker. It is not + available in an ordinary seminaive rule head, because a rule that read live + state would not re-fire when the state it read grows. + +## Implementation + +Three passes over the reachable subgraph. The two that walk e-classes use an +explicit stack, so term depth is bounded by the heap rather than by Rust's +stack; only the descent into nested containers recurses, and that is bounded by +how deeply the program's sorts nest containers. + +1. **Collect** — DFS from `root` gathering, per reachable e-class, its e-nodes + (table name + child values), and per reachable container value, its + contents. The root's sort is not known at runtime (the primitive is shared + across all call sites), so the walk probes every eq-sorted constructor for + the root e-class and then uses the matched constructor's declared input + sorts for everything below it. E-class ids come from one global counter, so + probing the wrong sort's table simply finds nothing. +2. **Mark** — worklist propagation from the map keys up the parent relation + built in pass 1. +3. **Build** — a sweep over the affected e-classes in postorder, copying every + e-node whose children already have copies (`lookup_or_insert` per e-node; + the first copy of a class names it, later ones union into it). One sweep + finishes the acyclic case; the sweep repeats while it makes progress, which + is what lets a grounded cycle close. Anything still uncopied when progress + stops is an ungrounded cycle and is reported. Containers are rebuilt through + `ContainerValues::rebuild_val_with` with a remap table, and nothing is + interned until every value inside the container resolves, so a blocked + container leaves no half-substituted copy behind. + +### What egglog had to expose + +Nothing substitution-specific: three general pieces of e-graph introspection, +after which the whole primitive is ordinary out-of-tree code. + +* `Read::enodes_for_eclass` — indexed lookup of the constructor rows whose + output column is a given e-class, instead of scanning the table. Cherry-picked + from (still open), together + with the `core-relations` `ExecutionState::for_each_matching_col` and + `egglog-bridge` `TableAction::for_each_output_value` it rests on. +* `Read::table_schema` / `Read::table_subtype` — a table's declared column sorts + and subtype, from inside a primitive body. `EGraph::functions_iter` already + exposes this from `&EGraph`, but a primitive only sees the state wrapper, and + those carried no sort information at all. Backed by a `FunctionSchemas` map + the e-graph shares with the wrappers exactly as it already shares + `ActionRegistry`, snapshot/restored across `push`/`pop` so a popped table + stops resolving. +* `Core::rebuild_container` — remap a container value's contents and intern the + result, over the existing `ContainerValues::rebuild_val_with`. Out-of-tree + code cannot go through `Core::register_container`, which needs to name the + container's Rust type. + +Two things it did **not** need to expose, worth recording because they were the +expected blockers: `TypeInfo::get_arcsorts_by` is already public, so the type +constraint can enumerate the declared `Map` and eq-sorts itself; and a sort's +kind is recoverable from the public `Sort::value_type` and `Sort::inner_sorts`, +so a `Map` sort can be identified without downcasting to `MapSort` (whose +`ContainerSort` impl sits behind a private wrapper type). + +### Known limitations + +* `R` and `K` must be eq-sorts, and the map's key and value sorts must be + identical: a substitution that replaced a `K`-sorted child with a value of a + different sort would produce an ill-typed row. +* A cycle in the substituted region with no grounded e-node is rejected rather + than copied, as described above. +* No bound on snapshot size — a root that reaches the whole e-graph copies as + much of it as the substitution affects. +* Values are assumed canonical, which holds at the top level (egglog rebuilds + after every command) and in a `:naive` rule head. Term-encoding mode, where + canonicalization goes through a per-sort union-find table rather than the + backend's, is untested. +* Proof mode is unsupported: the copied rows carry no justification. + `egglog_experimental::subst` errors out (`ProofsIncompatibleApi`, from + `EGraph::update`), and the primitive is + registered without a proof validator, so `program_supports_proofs` already + excludes any program using it from the proof-checking test runs. A program + that enables proofs *and* calls the primitive is not caught. +* Tests live in `tests/subst.rs` rather than a `tests/*.egg` file, so they skip + the `files` harness's desugar / term-encoding / multi-thread variants. +* A failing substitution reaches an egglog program as the generic + "primitive panicked", with the reason in the log. Registering a custom panic + message needs `egglog_bridge::EGraph::new_panic`, and egglog exposes no + accessor for its backend; a `Write::panic_with(message)` would fix it. +* The constructor list is rebuilt from the table schemas on every call, which is + O(tables) per substitution. Fine at the scale of a typical program, and it is + what keeps the primitive correct as an e-graph gains constructors. diff --git a/src/keep_best.rs b/src/keep_best.rs index c209316..0e9496b 100644 --- a/src/keep_best.rs +++ b/src/keep_best.rs @@ -96,10 +96,10 @@ fn collect_and_extract( .ok_or_else(|| TypeError::UnboundFunction(table_name.clone(), span!()))?; let all_sorts: Vec = func - .schema() + .func_type() .input .iter() - .chain(std::iter::once(&func.schema().output)) + .chain(std::iter::once(&func.func_type().output)) .cloned() .collect(); diff --git a/src/lib.rs b/src/lib.rs index f316df3..b6dbb63 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -19,6 +19,8 @@ //! - [`(get-size!)` primitive](https://github.com/egraphs-good/egglog-experimental/blob/main/tests/web-demo/node-limit.egg) //! for inspecting total tuple counts or counts for specific tables //! - [Multi-extraction](https://github.com/egraphs-good/egglog-experimental/blob/main/tests/web-demo/multi-extract.egg) +//! - [`unstable-subst`](subst): substitution over the sub-e-graph reachable +//! from a root, driven by a `Map` from an eq-sort to itself //! - Body-defined primitives with `(primitive name (InputSort*) OutputSort body)`. //! Body variables are positional (`_0`, `_1`, ...), and a partial primitive //! body result propagates as primitive failure. The registered primitive uses @@ -61,6 +63,9 @@ pub use sugar::*; mod keep_best; pub use keep_best::KeepBestCommand; +mod subst; +pub use subst::{SUBST, Subst, subst, substitute}; + pub fn new_experimental_egraph() -> EGraph { let mut egraph = EGraph::default(); @@ -105,6 +110,9 @@ pub fn new_experimental_egraph() -> EGraph { egraph .add_command("primitive".into(), Arc::new(primitive::RegisterPrimitive)) .unwrap(); + + // Substitution over a reachable sub-e-graph. + egraph.add_full_primitive(Subst, None); egraph } diff --git a/src/subst.rs b/src/subst.rs new file mode 100644 index 0000000..5a166b2 --- /dev/null +++ b/src/subst.rs @@ -0,0 +1,602 @@ +//! Substitution over a reachable sub-e-graph: the `unstable-subst` primitive. +//! +//! `(unstable-subst root map)` takes an e-class `root` of any eq-sort and a +//! `Map` whose key and value sorts are the same eq-sort. It walks the +//! constructor rows reachable from `root`, copies the part of that sub-e-graph +//! the substitution actually touches while replacing each key e-class with its +//! mapped value, and returns the e-class of the copied root. +//! +//! The walk follows constructor rows only — `function` rows are analyses over +//! the term structure, not part of it — and reaches through container-valued +//! children, rebuilding them with substituted contents. +//! +//! Unaffected e-classes are shared with the original rather than copied, so +//! substituting an empty map returns `root` itself and writes nothing. +//! +//! Three properties are worth knowing before relying on this: +//! +//! - The region's equations are substituted along with its terms. Copying an +//! e-class copies every one of its e-nodes, so `t1 = t2` in the original +//! becomes `σ(t1) = σ(t2)` in the copy — and an e-node with no substituted +//! children copies to itself, merging the copy back into the original class. +//! That is correct for equations that hold for every value of the substituted +//! classes (anything a rewrite rule derived) and wrong for a ground `union` +//! pinning one of them down, so only substitute classes that behave like +//! universally quantified variables. +//! - The walk reads committed table contents, so a term the enclosing action +//! has only just staged is not visible to it and comes back unsubstituted. +//! Terms from earlier commands, and from earlier rule iterations, are fine. +//! - The snapshot comes from live table contents, so this is a `Context::Full` +//! primitive: top-level actions and `:naive` rule heads only. +//! +//! Copies are named by `lookup_or_insert`, the same way `(Add a b)` in an +//! action is, so no e-class id is ever invented here. A cyclic e-class can +//! therefore only be copied if it has an e-node whose children all lie outside +//! the cycle to name it first — `x = {Var "x", Add x (Num 0)}` does, and works. +//! A cycle with no such e-node is an error rather than a silent partial copy. + +use std::any::TypeId; +use std::collections::{BTreeMap, HashMap, HashSet, VecDeque}; + +use egglog::api::RawValues; +use egglog::ast::Span; +use egglog::constraint::{self, Constraint, ImpossibleConstraint, TypeConstraint}; +use egglog::sort::MapContainer; +use egglog::{ + ArcSort, Atom, AtomTerm, Core, EGraph, Error, FullPrim, FullState, FuncType, Primitive, Read, + TypeInfo, Value, Write, +}; + +/// The name of the primitive, as written in an egglog program. +pub const SUBST: &str = "unstable-subst"; + +/// A constructor the walk can follow. Resolved from the e-graph's signatures at +/// the start of each call, since an e-graph gains constructors over time. +type Constructor<'a> = &'a FuncType; + +/// What a column of a given sort holds, as far as the walk cares. +enum Kind { + /// An e-class: walk into it, and copy it if it is affected. + Eclass, + /// A container with e-classes somewhere inside: walk into its contents and + /// rebuild it if any of them change. Carries the Rust [`TypeId`] its values + /// are interned under. + Container(TypeId), + /// Nothing an e-class can hide in: copied through untouched. + Opaque, +} + +fn kind_of(sort: &ArcSort) -> Kind { + if sort.is_eq_sort() { + Kind::Eclass + } else if sort.is_eq_container_sort() { + Kind::Container( + sort.value_type() + .expect("a container sort has a value type"), + ) + } else { + Kind::Opaque + } +} + +/// One reachable constructor row, keyed by the constructor it came from so the +/// build pass can re-apply it. +struct ENode { + ctor: usize, + children: Vec, +} + +/// The reachable sub-e-graph. +#[derive(Default)] +struct Snapshot { + /// E-nodes of each reachable e-class. Subsumed rows are left out: they are + /// excluded from extraction, so a copy must not resurrect them. + nodes: HashMap>, + /// Contents of each reachable container value, with the [`TypeId`] to + /// rebuild it under. + containers: HashMap)>, + /// The e-classes each e-class references, flattened through containers. + /// Drives both the affected fixpoint and the build order. + deps: HashMap>, +} + +struct Walk<'a> { + ctors: Vec>, + /// Indices into `ctors`, by output sort name. + by_output: HashMap>, + map: &'a BTreeMap, + snapshot: Snapshot, + /// E-classes the substitution changes: a key, or a reference to an + /// affected e-class. + affected: HashSet, + /// The copy of each affected e-class that has been named so far. + images: HashMap, + container_images: HashMap, + /// The e-class leaves of a container value, flattened through nesting. + container_leaves: HashMap>, +} + +/// The constructors with an eq-sort output, which are the rows that make up the +/// term structure. Function tables are analyses over that structure, and +/// globals lower to function tables, so both stay out of the walk. +fn constructors<'db>(state: &FullState<'_, 'db>) -> Vec> { + let names: Vec = state + .table_sizes() + .into_iter() + .map(|(name, _)| name.to_owned()) + .collect(); + names + .into_iter() + .filter_map(|name| { + // `constructor_schema` rejects the function tables, which is also + // what keeps globals out: they lower to function tables. + let func_type = state.constructor_schema(&name).ok()?; + func_type.output.is_eq_sort().then_some(func_type) + }) + .collect() +} + +/// Substitute `map` through the sub-e-graph reachable from `root`, returning +/// the root of the copy. See the module docs for the semantics. +pub fn substitute<'db>( + state: &mut FullState<'_, 'db>, + root: Value, + map: &BTreeMap, +) -> Result { + if let Some(target) = map.get(&root) { + return Ok(*target); + } + if map.is_empty() { + return Ok(root); + } + + let ctors = constructors(state); + let mut by_output: HashMap> = HashMap::new(); + for (index, ctor) in ctors.iter().enumerate() { + by_output + .entry(ctor.output.name().to_owned()) + .or_default() + .push(index); + } + + let mut walk = Walk { + ctors, + by_output, + map, + snapshot: Snapshot::default(), + affected: HashSet::new(), + images: HashMap::new(), + container_images: HashMap::new(), + container_leaves: HashMap::new(), + }; + walk.collect(state, root)?; + walk.mark(); + if !walk.affected.contains(&root) { + return Ok(root); + } + walk.build(state, root) +} + +/// A pending e-class visit in the collect pass. The sort is unknown for the +/// root only. +struct Visit(Value, Option); + +impl Walk<'_> { + /// Gather the reachable e-nodes, container contents, and e-class + /// dependency edges, starting from `root`. + fn collect(&mut self, state: &FullState<'_, '_>, root: Value) -> Result<(), Error> { + let mut stack = vec![Visit(root, None)]; + while let Some(Visit(value, sort)) = stack.pop() { + if self.snapshot.nodes.contains_key(&value) { + continue; + } + // E-class ids all come from one counter, so probing a constructor + // of the wrong sort just finds nothing. That is what makes the + // root's unknown sort affordable: it costs one probe per eq-sorted + // constructor, once. + let candidates: Vec = match &sort { + Some(sort) => self.by_output.get(sort.name()).cloned().unwrap_or_default(), + None => (0..self.ctors.len()).collect(), + }; + let mut nodes = Vec::new(); + for index in candidates { + let mut rows = Vec::new(); + state.enodes_for_eclass(&self.ctors[index].name, value, |enode| { + if !enode.subsumed { + rows.push(enode.children.to_vec()); + } + })?; + nodes.extend(rows.into_iter().map(|children| ENode { + ctor: index, + children, + })); + } + + let mut deps = Vec::new(); + for node in &nodes { + let ctor = self.ctors[node.ctor]; + for (child, child_sort) in node.children.iter().zip(&ctor.input) { + match kind_of(child_sort) { + Kind::Eclass => { + deps.push(*child); + stack.push(Visit(*child, Some(child_sort.clone()))); + } + Kind::Container(type_id) => { + for (leaf_sort, leaf) in + self.container_leaves(state, *child, type_id, child_sort) + { + deps.push(leaf); + stack.push(Visit(leaf, Some(leaf_sort))); + } + } + Kind::Opaque => {} + } + } + } + // Recorded after the children are read, but before they are + // visited, so a cycle back to `value` terminates. + self.snapshot.nodes.insert(value, nodes); + self.snapshot.deps.insert(value, deps); + } + Ok(()) + } + + /// The e-class leaves of a container value, flattened through nested + /// containers, recording its contents on the way. Recursion is bounded by + /// how deeply the program's sorts nest containers. + fn container_leaves( + &mut self, + state: &FullState<'_, '_>, + value: Value, + type_id: TypeId, + sort: &ArcSort, + ) -> Vec<(ArcSort, Value)> { + if let Some(leaves) = self.container_leaves.get(&value) { + return leaves.clone(); + } + let contents = sort.inner_values(state.container_values(), value); + let mut leaves = Vec::new(); + for (inner_sort, inner) in &contents { + match kind_of(inner_sort) { + Kind::Eclass => leaves.push((inner_sort.clone(), *inner)), + Kind::Container(inner_type_id) => { + leaves.extend(self.container_leaves(state, *inner, inner_type_id, inner_sort)) + } + Kind::Opaque => {} + } + } + self.snapshot.containers.insert(value, (type_id, contents)); + self.container_leaves.insert(value, leaves.clone()); + leaves + } + + /// Least fixpoint of "references something the substitution changes", + /// seeded with the map keys the walk actually reached. + fn mark(&mut self) { + let mut users: HashMap> = HashMap::new(); + for (owner, deps) in &self.snapshot.deps { + for dep in deps { + users.entry(*dep).or_default().push(*owner); + } + } + + let mut queue: VecDeque = self + .map + .keys() + .copied() + .filter(|key| self.snapshot.nodes.contains_key(key)) + .collect(); + self.affected.extend(queue.iter().copied()); + while let Some(value) = queue.pop_front() { + for user in users.get(&value).into_iter().flatten() { + if self.affected.insert(*user) { + queue.push_back(*user); + } + } + } + } + + /// Copy the affected e-classes and return the root's image. + /// + /// Every copied e-node goes in through `lookup_or_insert`, so egglog names + /// the copy's e-class — nothing here invents an id. That is why this is a + /// sweep rather than a single postorder pass: an e-node can only be copied + /// once its children have copies, and a cycle in the copied region needs + /// one e-node whose children all lie outside it to get started. A cycle + /// with no such e-node is reported instead. + fn build(&mut self, state: &mut FullState<'_, '_>, root: Value) -> Result { + // Children before parents, so the acyclic case finishes in one sweep. + let order = self.postorder(root); + let mut pending: Vec<(Value, Vec)> = order + .into_iter() + .map(|eclass| { + let nodes = self.snapshot.nodes.remove(&eclass).unwrap_or_default(); + (eclass, nodes) + }) + .collect(); + + loop { + let mut progress = false; + for (eclass, nodes) in pending.iter_mut() { + let mut blocked = Vec::new(); + for node in std::mem::take(nodes) { + let Some(args) = self.copied_args(state, &node) else { + blocked.push(node); + continue; + }; + let copy = state.add(&self.ctors[node.ctor].name, RawValues(args))?; + match self.images.get(eclass) { + // The first e-node copied names the class; the rest are + // further ways to say the same one. + None => { + self.images.insert(*eclass, copy); + } + Some(image) if *image != copy => state.union(copy, *image)?, + Some(_) => {} + } + progress = true; + } + *nodes = blocked; + } + if !progress { + break; + } + } + + if let Some((eclass, blocked)) = pending.iter().find(|(_, nodes)| !nodes.is_empty()) { + let mut ctors: Vec<&str> = blocked + .iter() + .map(|node| self.ctors[node.ctor].name.as_str()) + .collect(); + ctors.dedup(); + return Err(error(format!( + "no order copies e-class {eclass:?}: its remaining e-nodes ({}) all refer to \ + copies that do not exist yet. Every cycle in the substituted region needs at \ + least one e-node whose children all lie outside it.", + ctors.join(", "), + ))); + } + + self.images + .get(&root) + .copied() + .ok_or_else(|| error(format!("the root e-class {root:?} was not copied"))) + } + + /// The affected e-classes that need copying, children before parents. + fn postorder(&self, root: Value) -> Vec { + enum Frame { + Enter(Value), + Exit(Value), + } + + let mut order = Vec::new(); + let mut seen = HashSet::new(); + let mut stack = vec![Frame::Enter(root)]; + while let Some(frame) = stack.pop() { + match frame { + Frame::Enter(eclass) => { + if !seen.insert(eclass) { + continue; + } + stack.push(Frame::Exit(eclass)); + for dep in self.snapshot.deps.get(&eclass).into_iter().flatten() { + if self.needs_copy(*dep) && !seen.contains(dep) { + stack.push(Frame::Enter(*dep)); + } + } + } + Frame::Exit(eclass) => order.push(eclass), + } + } + order + } + + /// Whether this e-class gets a copy: affected, and not a key (a key is + /// replaced outright rather than copied). + fn needs_copy(&self, eclass: Value) -> bool { + self.affected.contains(&eclass) && !self.map.contains_key(&eclass) + } + + /// The substituted children of `node`, or `None` if some child's copy does + /// not exist yet. + fn copied_args(&mut self, state: &mut FullState<'_, '_>, node: &ENode) -> Option> { + let ctor = self.ctors[node.ctor]; + let mut args = Vec::with_capacity(node.children.len()); + for (child, child_sort) in node.children.iter().zip(&ctor.input) { + let image = match kind_of(child_sort) { + Kind::Eclass => self.eclass_image(*child)?, + Kind::Container(type_id) => self.container_image(state, *child, type_id)?, + Kind::Opaque => *child, + }; + args.push(image); + } + Some(args) + } + + fn eclass_image(&self, eclass: Value) -> Option { + if let Some(target) = self.map.get(&eclass) { + return Some(*target); + } + if !self.affected.contains(&eclass) { + return Some(eclass); + } + self.images.get(&eclass).copied() + } + + /// The interned value of `container` with its contents substituted, or + /// `container` itself if nothing inside it changed. `None` while any + /// e-class inside it is still waiting for its copy. + fn container_image( + &mut self, + state: &mut FullState<'_, '_>, + container: Value, + type_id: TypeId, + ) -> Option { + if let Some(image) = self.container_images.get(&container) { + return Some(*image); + } + let Some((_, contents)) = self.snapshot.containers.get(&container).cloned() else { + return Some(container); + }; + + // Nothing is interned until every value inside resolves, so a blocked + // container leaves no half-substituted copy behind. + let mut remap: HashMap = HashMap::new(); + for (inner_sort, inner) in &contents { + let image = match kind_of(inner_sort) { + Kind::Eclass => self.eclass_image(*inner)?, + Kind::Container(inner_type_id) => { + self.container_image(state, *inner, inner_type_id)? + } + Kind::Opaque => continue, + }; + if image != *inner { + remap.insert(*inner, image); + } + } + + let image = if remap.is_empty() { + container + } else { + state.rebuild_container(type_id, container, &|value| { + remap.get(&value).copied().unwrap_or(value) + }) + }; + self.container_images.insert(container, image); + Some(image) + } +} + +fn error(message: String) -> Error { + Error::BackendError(format!("{SUBST}: {message}")) +} + +/// Substitute through the sub-e-graph reachable from `root`, returning the +/// e-class of the copy. The top-level form of the [`SUBST`] primitive. +/// +/// `map` must be a `Map` container value whose key and value sorts are the same +/// eq-sort; `root` may be of any eq-sort. Constructor rows reachable from +/// `root` are copied with each key e-class replaced by its mapped value; +/// e-classes the substitution does not affect are shared with the original +/// rather than copied. +/// +/// Errors if the substituted region contains a cycle in which every e-node +/// refers back into the cycle, since naming that copy would require an e-class +/// id no row produces. +pub fn subst(egraph: &mut EGraph, root: Value, map: Value) -> Result { + egraph.update(|mut state| { + let entries = match state.value_to_container::(map) { + Some(entries) => entries.data.clone(), + None => return Err(error(format!("{map:?} is not a Map container value"))), + }; + substitute(&mut state, root, &entries) + }) +} + +/// The `unstable-subst` primitive. +#[derive(Clone)] +pub struct Subst; + +impl Primitive for Subst { + fn name(&self) -> &str { + SUBST + } + + fn get_type_constraints(&self, span: &Span) -> Box { + Box::new(SubstTypeConstraint { span: span.clone() }) + } +} + +impl FullPrim for Subst { + fn apply<'a, 'db>(&self, mut state: FullState<'a, 'db>, args: &[Value]) -> Option { + let [root, map] = args else { return None }; + // Cloned out so the container registry is not still borrowed when the + // walk starts interning new containers. + let entries = state.value_to_container::(*map)?.data.clone(); + match substitute(&mut state, *root, &entries) { + Ok(image) => Some(image), + Err(err) => { + // A primitive cannot return an `Error`, and registering a + // custom panic message needs the backend, which egglog does not + // expose. So the reason goes to the log and the program sees a + // generic primitive panic. + log::error!("{err}"); + state.panic(); + None + } + } + } +} + +/// `(unstable-subst root map) : (R, Map) -> R` for any eq-sort `R`. +/// +/// `R` is free rather than pinned to `K` because a substitution reaches through +/// every sort in the term structure, so a root of one sort can perfectly well +/// be rewritten by a map over another. `K` must be an eq-sort mapping to +/// itself: replacing a `K`-sorted child with a value of another sort would +/// produce an ill-typed row. +struct SubstTypeConstraint { + span: Span, +} + +impl TypeConstraint for SubstTypeConstraint { + fn get( + &self, + arguments: &[AtomTerm], + typeinfo: &TypeInfo, + ) -> Vec>> { + let [root, map, out] = arguments else { + return vec![constraint::impossible( + ImpossibleConstraint::ArityMismatch { + atom: Atom { + span: self.span.clone(), + head: SUBST.to_owned(), + args: arguments.to_vec(), + }, + expected: 3, + }, + )]; + }; + + let mut cs: Vec>> = + vec![constraint::eq(root.clone(), out.clone())]; + + // One instantiation per declared sort that could stand in each + // position; `xor` defers until the surrounding program pins it down. + // + // A `Map` sort is identified by the Rust type its values intern + // under, since the `ContainerSort` impl behind an `ArcSort` is + // wrapped in a private type that out-of-tree code cannot downcast to. + let mut map_sorts: Vec = typeinfo.get_arcsorts_by(|sort| { + sort.value_type() == Some(TypeId::of::()) + && match sort.inner_sorts().as_slice() { + [key, value] => key.is_eq_sort() && key.name() == value.name(), + _ => false, + } + }); + map_sorts.sort_by_key(|sort| sort.name().to_owned()); + cs.push(constraint::xor( + map_sorts + .into_iter() + .map(|sort| constraint::assign(map.clone(), sort)) + .collect(), + )); + + let mut eq_sorts = typeinfo.get_arcsorts_by(|sort| sort.is_eq_sort()); + eq_sorts.sort_by_key(|sort| sort.name().to_owned()); + cs.push(constraint::xor( + eq_sorts + .into_iter() + .map(|sort| { + constraint::and(vec![ + constraint::assign(root.clone(), sort.clone()), + constraint::assign(out.clone(), sort), + ]) + }) + .collect(), + )); + + cs + } +} diff --git a/src/table_rows.rs b/src/table_rows.rs index dde91af..98926a2 100644 --- a/src/table_rows.rs +++ b/src/table_rows.rs @@ -6,7 +6,8 @@ //! helper dispatches on the subtype and hands the callback the whole row — //! inputs followed by the output (or eclass) column. -use egglog::{ApiError, EGraph, Error, Value}; +use egglog::ast::FunctionSubtype; +use egglog::{EGraph, Error, Read, Value}; /// Call `f` once per row of `name`, with every column in schema order. pub(crate) fn for_each_row( @@ -33,15 +34,8 @@ pub(crate) fn for_each_row( } /// Whether `name` is a constructor (or relation) table rather than a -/// `function` table. -/// -/// egglog exposes no subtype accessor on [`egglog::Function`], so probe with a -/// constructor scan that stops before reading a row and read the answer off the -/// subtype check, which runs before any iteration. +/// `function` table. Unknown tables report `false`; the scan that follows +/// reports the missing table. pub(crate) fn is_constructor(egraph: &EGraph, name: &str) -> Result { - match egraph.constructor_enodes_while(name, |_| false) { - Ok(()) => Ok(true), - Err(Error::ApiError(ApiError::WrongSubtype { .. })) => Ok(false), - Err(err) => Err(err), - } + Ok(egraph.read(|state| state.table_subtype(name)) == Some(FunctionSubtype::Constructor)) } diff --git a/src/table_stats.rs b/src/table_stats.rs index fc2b1d5..8db08ac 100644 --- a/src/table_stats.rs +++ b/src/table_stats.rs @@ -215,7 +215,7 @@ fn compute_table_stats(egraph: &EGraph, func_name: &str) -> Result = schema.input.iter().map(|s| s.name().to_owned()).collect(); column_types.push(schema.output.name().to_owned()); let n_cols = column_types.len(); diff --git a/tests/subst.rs b/tests/subst.rs new file mode 100644 index 0000000..3d84da7 --- /dev/null +++ b/tests/subst.rs @@ -0,0 +1,425 @@ +//! Tests for the `unstable-subst` primitive. + +use egglog::prelude::exprs::var; +use egglog::prelude::*; +use egglog::{EGraph, Error, Value}; +use egglog_experimental::new_experimental_egraph; + +const MATH: &str = r#" +(datatype Math + (Num i64) + (Var String) + (Add Math Math) + (Mul Math Math)) +(sort MathToMath (Map Math Math)) +"#; + +fn egraph(program: &str) -> EGraph { + let mut eg = new_experimental_egraph(); + eg.parse_and_run_program(None, MATH).unwrap(); + eg.parse_and_run_program(None, program).unwrap(); + eg +} + +/// The value of a global, for the assertions that need to compare e-classes +/// rather than terms. +fn global(eg: &mut EGraph, name: &str) -> Value { + eg.eval_expr(&var(&format!("${name}"))).unwrap().1 +} + +#[test] +fn replaces_a_variable_everywhere_it_occurs() { + egraph( + r#" +(let $x (Var "x")) +(let $y (Var "y")) +(let $e (Add $x (Mul $x (Num 2)))) +(let $copy (unstable-subst $e (map-insert (map-empty) $x $y))) +(check (= $copy (Add $y (Mul $y (Num 2))))) +(fail (check (= $copy $e))) +"#, + ); +} + +#[test] +fn substitutes_several_keys_at_once() { + egraph( + r#" +(let $x (Var "x")) +(let $y (Var "y")) +(let $e (Add $x $y)) +(let $copy (unstable-subst $e (map-insert (map-insert (map-empty) $x $y) $y $x))) +(check (= $copy (Add $y $x))) +"#, + ); +} + +#[test] +fn substituting_a_missing_key_returns_the_root_itself() { + let mut eg = egraph( + r#" +(let $x (Var "x")) +(let $e (Add $x (Num 1))) +(let $copy (unstable-subst $e (map-insert (map-empty) (Var "absent") (Num 0)))) +"#, + ); + assert_eq!(global(&mut eg, "e"), global(&mut eg, "copy")); +} + +#[test] +fn an_empty_map_returns_the_root_itself() { + let mut eg = egraph( + r#" +(let $e (Add (Var "x") (Num 1))) +(let $copy (unstable-subst $e (map-empty))) +"#, + ); + assert_eq!(global(&mut eg, "e"), global(&mut eg, "copy")); +} + +#[test] +fn substituting_the_root_returns_its_replacement() { + let mut eg = egraph( + r#" +(let $x (Var "x")) +(let $y (Var "y")) +(let $copy (unstable-subst $x (map-insert (map-empty) $x $y))) +"#, + ); + assert_eq!(global(&mut eg, "y"), global(&mut eg, "copy")); +} + +/// Only the spine above a substituted class is copied; e-classes the +/// substitution cannot reach stay shared with the original. +#[test] +fn unaffected_subterms_are_shared_not_copied() { + let mut eg = egraph( + r#" +(let $x (Var "x")) +(let $untouched (Mul (Num 3) (Num 4))) +(let $e (Add $x $untouched)) +(let $copy (unstable-subst $e (map-insert (map-empty) $x (Num 7)))) +(check (= $copy (Add (Num 7) $untouched))) +"#, + ); + // `Add` gained exactly one row: the copied root. + let adds = eg.update(|fs| Ok(fs.table_size("Add"))).unwrap(); + assert_eq!(adds, Some(2)); + let muls = eg.update(|fs| Ok(fs.table_size("Mul"))).unwrap(); + assert_eq!(muls, Some(1)); + assert_ne!(global(&mut eg, "e"), global(&mut eg, "copy")); +} + +/// Copying an e-class copies every one of its e-nodes, so an equation between +/// two reachable terms is carried over to the copy. For an equation a rewrite +/// rule derived — one that holds for every value of the substituted class — +/// that is exactly right: `x * 0 = 0` stays `5 * 0 = 0`. The untouched `(Num 0)` +/// e-node copies to itself, which is what merges the copy back into that class. +#[test] +fn an_equation_that_holds_for_every_value_survives_substitution() { + egraph( + r#" +(rewrite (Mul a (Num 0)) (Num 0)) +(let $x (Var "x")) +(let $e (Mul $x (Num 0))) +(run 1) +(check (= $e (Num 0))) +(let $copy (unstable-subst $e (map-insert (map-empty) $x (Num 5)))) +(check (= $copy (Mul (Num 5) (Num 0)))) +(check (= $copy (Num 0))) +"#, + ); +} + +/// The same mechanism is unsound on a *ground* equation, which holds only for +/// the class it pins down: `(union (Add x (Num 1)) (Num 5))` says `x = 4`, and +/// substituting `x := 9` into it asserts `9 + 1 = 5`. Only substitute classes +/// that behave like universally quantified variables. +#[test] +fn a_ground_equation_about_a_key_is_substituted_too() { + egraph( + r#" +(let $x (Var "x")) +(let $e (Add $x (Num 1))) +(union $e (Num 5)) +(let $copy (unstable-subst $e (map-insert (map-empty) $x (Num 9)))) +(check (= $copy (Add (Num 9) (Num 1)))) +(check (= $copy (Num 5))) +"#, + ); +} + +/// A row an action has only just staged is not in the tables the walk reads, +/// so a term built in the same action is invisible to it and comes back +/// unsubstituted. Terms from earlier commands, and from earlier rule +/// iterations, are fine. +#[test] +fn does_not_see_terms_built_in_the_same_action() { + let mut eg = egraph( + r#" +(let $x (Var "x")) +(let $copy (unstable-subst (Mul $x (Num 5)) (map-insert (map-empty) $x (Num 9)))) +(let $unsubstituted (Mul $x (Num 5))) +"#, + ); + assert_eq!(global(&mut eg, "copy"), global(&mut eg, "unsubstituted")); +} + +/// A cyclic e-class is copied as long as one of its e-nodes has all its +/// children outside the cycle: that e-node's `lookup_or_insert` names the copy, +/// and the cyclic e-node then unions into it. No e-class id is invented. +#[test] +fn a_grounded_cyclic_eclass_is_copied() { + egraph( + r#" +(let $x (Var "x")) +(let $loop (Add $x (Num 0))) +;; $loop now holds a grounded e-node and one referring to its own class. +(union $loop (Add $loop (Num 0))) +(let $copy (unstable-subst $loop (map-insert (map-empty) $x (Num 7)))) +(check (= $copy (Add (Num 7) (Num 0)))) +(check (= $copy (Add $copy (Num 0)))) +"#, + ); +} + +/// `$a = {Add (Num 0) x, Add $b x}` with `$b = Mul $a (Num 1)`: affected through +/// `x`, cyclic through `$b`. Subsuming the one e-node whose children lie outside +/// the cycle leaves no order that names the copy first. +const UNGROUNDED_CYCLE: &str = r#" +(let $x (Var "x")) +(let $a (Add (Num 0) $x)) +(let $b (Mul $a (Num 1))) +(union $a (Add $b $x)) +(subsume (Add (Num 0) $x)) +(let $map (map-insert (map-empty) $x (Num 9))) +"#; + +/// The walk reports an ungrounded cycle instead of inventing an e-class id. +#[test] +fn an_ungrounded_cycle_is_an_error() { + let mut eg = egraph(UNGROUNDED_CYCLE); + let root = global(&mut eg, "a"); + let map = global(&mut eg, "map"); + let err = egglog_experimental::subst(&mut eg, root, map).unwrap_err(); + let message = err.to_string(); + assert!( + message.contains("unstable-subst") && message.contains("no order copies e-class"), + "expected an ungrounded-cycle error, got {message}" + ); +} + +/// A primitive cannot return an `Error`, and registering a custom panic message +/// needs the backend, which egglog does not expose to out-of-tree code. So the +/// same failure reaches an egglog program as a generic primitive panic, with +/// the reason in the log. +#[test] +fn an_ungrounded_cycle_panics_from_egglog() { + let mut eg = egraph(UNGROUNDED_CYCLE); + let err = eg + .parse_and_run_program(None, "(let $copy (unstable-subst $a $map))") + .unwrap_err(); + let message = err.to_string(); + assert!( + message.contains("panicked"), + "expected a primitive panic, got {message}" + ); +} + +/// The root's sort is free: it need not be the map's key sort, because a +/// substitution reaches through every sort in the term structure. +#[test] +fn the_root_may_be_a_different_sort_than_the_keys() { + egraph( + r#" +(datatype MathList (Nil) (Cons Math MathList)) +(let $x (Var "x")) +(let $list (Cons $x (Cons (Num 1) (Nil)))) +(let $copy (unstable-subst $list (map-insert (map-empty) $x (Num 8)))) +(check (= $copy (Cons (Num 8) (Cons (Num 1) (Nil))))) +"#, + ); +} + +/// E-classes reached only through a container child are substituted too, and +/// the container is rebuilt around the replacements. +#[test] +fn substitutes_inside_container_children() { + egraph( + r#" +(sort MathVec (Vec Math)) +(constructor Sum (MathVec) Math) +(let $x (Var "x")) +(let $e (Sum (vec-of $x (Num 1)))) +(let $copy (unstable-subst $e (map-insert (map-empty) $x (Num 4)))) +(check (= $copy (Sum (vec-of (Num 4) (Num 1))))) +"#, + ); +} + +/// Container children with no e-classes inside are never walked into, so a +/// substitution leaves them alone. +#[test] +fn leaves_containers_without_eclasses_alone() { + egraph( + r#" +(sort Ints (Vec i64)) +(constructor Tagged (Math Ints) Math) +(let $x (Var "x")) +(let $e (Tagged $x (vec-of 1 2))) +(let $copy (unstable-subst $e (map-insert (map-empty) $x (Num 5)))) +(check (= $copy (Tagged (Num 5) (vec-of 1 2)))) +"#, + ); +} + +/// Subsumed e-nodes are excluded from extraction, so a copy must not bring +/// them back un-subsumed. +#[test] +fn skips_subsumed_enodes() { + let mut eg = egraph( + r#" +(let $x (Var "x")) +(let $keep (Add $x (Num 1))) +(let $drop (Mul $x (Num 1))) +(union $keep $drop) +(subsume (Mul $x (Num 1))) +(let $copy (unstable-subst $keep (map-insert (map-empty) $x (Num 6)))) +(check (= $copy (Add (Num 6) (Num 1)))) +"#, + ); + // The one `Mul` row is the subsumed original; no copy was made. + let muls = eg.update(|fs| Ok(fs.table_size("Mul"))).unwrap(); + assert_eq!(muls, Some(1)); +} + +#[test] +fn runs_in_a_naive_rule_head() { + egraph( + r#" +(constructor Beta (Math Math Math) Math) +(rule ((= $lhs (Beta body from to))) + ((union $lhs (unstable-subst body (map-insert (map-empty) from to)))) + :naive) +(let $x (Var "x")) +(let $b (Beta (Add $x $x) $x (Num 3))) +(run 1) +(check (= $b (Add (Num 3) (Num 3)))) +"#, + ); +} + +/// Reads of live tables in a seminaive rule head would not re-fire when the +/// tables they read grow, so the typechecker rejects them. +#[test] +fn rejected_in_a_seminaive_rule_head() { + let mut eg = new_experimental_egraph(); + eg.parse_and_run_program(None, MATH).unwrap(); + let err = eg + .parse_and_run_program( + None, + r#" +(constructor Beta (Math Math Math) Math) +(rule ((= $lhs (Beta body from to))) + ((union $lhs (unstable-subst body (map-insert (map-empty) from to))))) +"#, + ) + .unwrap_err(); + assert!( + matches!(err, Error::TypeError(_) | Error::TypeErrors(_)), + "expected a type error, got {err}" + ); +} + +/// A map from one sort to another cannot be a substitution: the replacement +/// would not typecheck in the column it is written to. +#[test] +fn rejects_a_map_between_different_sorts() { + let mut eg = new_experimental_egraph(); + eg.parse_and_run_program(None, MATH).unwrap(); + let err = eg + .parse_and_run_program( + None, + r#" +(sort MathToInt (Map Math i64)) +(let $x (Var "x")) +(let $copy (unstable-subst $x (map-insert (map-empty) $x 3))) +"#, + ) + .unwrap_err(); + assert!( + matches!(err, Error::TypeError(_) | Error::TypeErrors(_)), + "expected a type error, got {err}" + ); +} + +/// Substitution is simultaneous, and a key's replacement is not itself +/// substituted — but a *reachable* class that happens to be a replacement is +/// still copied where it occurs on its own. +#[test] +fn a_replacement_that_is_itself_affected_is_copied_where_it_occurs() { + egraph( + r#" +(let $x (Var "x")) +(let $y (Add $x (Num 1))) +(let $e (Mul $x $y)) +(let $copy (unstable-subst $e (map-insert (map-empty) $x $y))) +;; x |-> y, and y itself becomes (Add y 1) where it occurs as a child. +(check (= $copy (Mul $y (Add $y (Num 1))))) +"#, + ); +} + +/// The walk uses an explicit stack, so depth is bounded by the heap rather than +/// by Rust's stack. +#[test] +fn handles_a_deep_spine() { + const DEPTH: usize = 20_000; + let mut program = String::from("(let $x (Var \"x\"))\n(let $e0 $x)\n"); + for i in 1..=DEPTH { + program.push_str(&format!("(let $e{i} (Add $x $e{}))\n", i - 1)); + } + let mut eg = egraph(&program); + eg.parse_and_run_program( + None, + &format!("(let $copy (unstable-subst $e{DEPTH} (map-insert (map-empty) $x (Num 0))))"), + ) + .unwrap(); + // Every one of the spine's `Add`s mentions `x`, so all of them are copied. + let adds = eg.update(|fs| Ok(fs.table_size("Add"))).unwrap(); + assert_eq!(adds, Some(2 * DEPTH)); +} + +/// The Rust-level entry point, which takes the map as a container value. +#[test] +fn the_rust_api_substitutes() { + let mut eg = egraph( + r#" +(let $x (Var "x")) +(let $e (Add $x (Num 1))) +(let $map (map-insert (map-empty) $x (Num 2))) +(let $expected (Add (Num 2) (Num 1))) +"#, + ); + let root = global(&mut eg, "e"); + let map = global(&mut eg, "map"); + let copy = egglog_experimental::subst(&mut eg, root, map).unwrap(); + assert_eq!(copy, global(&mut eg, "expected")); +} + +#[test] +fn a_popped_constructor_is_not_walked() { + let mut eg = egraph( + r#" +(let $x (Var "x")) +(push) +(constructor Neg (Math) Math) +(let $inner (Neg $x)) +(pop) +(let $e (Add $x (Num 1))) +(let $copy (unstable-subst $e (map-insert (map-empty) $x (Num 2)))) +(check (= $copy (Add (Num 2) (Num 1)))) +"#, + ); + let _ = &mut eg; +}