From 4ed4acca313a14bdef3328f82fd05f714d79f20a Mon Sep 17 00:00:00 2001 From: Oliver Flatt Date: Tue, 7 Jul 2026 18:03:43 +0000 Subject: [PATCH] Register user primitives on the term-encoding typechecker A primitive registered through the Rust API (`add_pure_primitive` / `add_*_primitive` / the `add_primitive!` macros) after e-graph construction was only registered on the running e-graph, not on the separate `original_typechecking` e-graph that term encoding / proofs use to typecheck the encoded program. The encoder then reported the primitive as an unbound function, so callers had to manually register it on `proof_state.original_typechecking` as well (a proof test did exactly that). `register_per_context` now walks the `original_typechecking` chain and registers the primitive on each e-graph. A typechecker only typechecks and never evaluates, so the wrapper's runtime state is irrelevant there, and both e-graphs register the same built-ins at construction, so a primitive added to both gets the same `ExternalFunctionId`. A sort's primitives already reach the typechecker through its own `add_arcsort` when the typechecker typechecks the sort command, so `add_arcsort` detaches the typechecker while a sort registers to avoid double-registering (which would make primitive resolution ambiguous). Co-Authored-By: Claude Opus 4.8 --- CHANGELOG.md | 1 + src/prelude.rs | 17 +++++++++++++++ src/proofs/proof_tests.rs | 10 ++------- src/typechecking.rs | 44 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 64 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index dff0897c1..17c748058 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## [Unreleased] - ReleaseDate +- Fix user-defined primitives (registered through the Rust API after construction) being reported as unbound under term encoding / proofs: primitive registration now also reaches the term-encoding typechecker, so the encoder can typecheck the encoded program. Previously callers had to manually register the primitive on `proof_state.original_typechecking` as well. - Add `make nightly` and `scripts/nightly_bench.py`, a hyperfine-based benchmark harness that measures every `tests/**/*.egg` program at 1/2/4/8 threads and (where supported) in proof-testing mode, caps each run at a 2-minute timeout, skips sub-50ms programs, and emits an HTML dashboard (one row per benchmark, one column per configuration) for nightly.cs.washington.edu. The dashboard uses [eval-live](https://github.com/oflatt/eval-live) for interactive filtering and sorting. - Add typed `EGraph` extension state that clones with `EGraph` and is restored by `push`/`pop`. - Fix custom scheduler queries so subsumed rows are not offered as fresh matches. diff --git a/src/prelude.rs b/src/prelude.rs index 30194d0b3..65fbe48c2 100644 --- a/src/prelude.rs +++ b/src/prelude.rs @@ -1040,6 +1040,23 @@ mod tests { Ok(egraph) } + #[test] + fn add_primitive_under_term_encoding() -> Result<(), Error> { + // A primitive registered through the Rust API AFTER construction must + // reach the term-encoding typechecker (a separate e-graph), not just the + // running e-graph — otherwise the encoder reports it as unbound. + let mut egraph = EGraph::new_with_term_encoding(); + add_literal_prim!(&mut egraph, "tripleit" = |x: i64| -> i64 { x * 3 }); + egraph.parse_and_run_program( + None, + "(function f (i64) i64 :merge (min old new)) +(set (f 0) (tripleit 7)) +(run 1) +(check (= (f 0) 21))", + )?; + Ok(()) + } + #[test] fn rust_api_query() -> Result<(), Error> { let mut egraph = build_test_database()?; diff --git a/src/proofs/proof_tests.rs b/src/proofs/proof_tests.rs index 3103ffe12..4f44e79ec 100644 --- a/src/proofs/proof_tests.rs +++ b/src/proofs/proof_tests.rs @@ -200,6 +200,8 @@ mod tests { .get_sort_by_name("EqContainer") .expect("EqContainer sort") .clone(); + // Registering on the egraph also registers on its term-encoding + // typechecker (see `register_per_context`), so this one call is enough. let validator = |_: &mut TermDag, args: &[TermId]| -> Option { args.first().copied() }; add_primitive_with_validator!( @@ -207,14 +209,6 @@ mod tests { "proof-container-id" = |x: # (eq_container_sort)| -> # (eq_container_sort) { x }, validator ); - let validator = - |_: &mut TermDag, args: &[TermId]| -> Option { args.first().copied() }; - let original_typechecking = egraph.proof_state.original_typechecking.as_mut().unwrap(); - add_primitive_with_validator!( - &mut **original_typechecking, - "proof-container-id" = |x: # (eq_container_sort)| -> # (eq_container_sort) { x }, - validator - ); egraph .parse_and_run_program( diff --git a/src/typechecking.rs b/src/typechecking.rs index 39f3749c4..68269e71c 100644 --- a/src/typechecking.rs +++ b/src/typechecking.rs @@ -264,7 +264,15 @@ impl EGraph { HEntry::Occupied(_) => Err(TypeError::SortAlreadyBound(name.to_owned(), span)), HEntry::Vacant(e) => { e.insert(sort.clone()); + // A sort's primitives already reach the term-encoding typechecker + // through its OWN `add_arcsort` when it typechecks the sort + // command, so don't propagate them again from here (that would + // double-register and make primitive resolution ambiguous). + // Detach the typechecker while the sort registers, so only direct + // `add_*_primitive` calls propagate to it. + let saved = self.proof_state.original_typechecking.take(); sort.register_primitives(self); + self.proof_state.original_typechecking = saved; Ok(()) } } @@ -356,6 +364,42 @@ impl EGraph { ) where T: Primitive + Clone, F: FnMut(T, Context) -> Box, + { + // Register on this e-graph AND every term-encoding typechecker down the + // chain. Each typechecker is a separate e-graph that typechecks the + // encoded program (see `typecheck_program`); a primitive added after + // construction is otherwise unknown to it and reported as unbound. A + // typechecker only typechecks and never evaluates, so the wrapper's + // runtime state is irrelevant there, and — since both e-graphs register + // the same built-ins during construction — a primitive added to both + // gets the same `ExternalFunctionId`. + let mut eg: &mut EGraph = self; + loop { + eg.register_primitive_local( + x.clone(), + validator.clone(), + valid_ctxs, + &mut build_wrapper, + ); + match eg.proof_state.original_typechecking.as_deref_mut() { + Some(next) => eg = next, + None => break, + } + } + } + + /// Register one primitive on THIS e-graph only (its `type_info` plus a + /// backend `ExternalFunctionId` per valid context). Non-generic over the + /// wrapper builder so the chain walk in [`register_per_context`] doesn't + /// recurse through ever-growing `&mut F` types. + fn register_primitive_local( + &mut self, + x: T, + validator: Option, + valid_ctxs: &[Context], + build_wrapper: &mut dyn FnMut(T, Context) -> Box, + ) where + T: Primitive + Clone, { let primitive: Arc = Arc::new(x.clone()); let name = primitive.name().to_owned();