Register user primitives on the term-encoding typechecker

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.

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()?;

src/proofs/proof_tests.rs

@@ -200,21 +200,15 @@ 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<TermId> { args.first().copied() };
         add_primitive_with_validator!(
             &mut egraph,
             "proof-container-id" = |x: # (eq_container_sort)| -> # (eq_container_sort) { x },
             validator
         );
-        let validator =
-            |_: &mut TermDag, args: &[TermId]| -> Option<TermId> { 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(

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<dyn ExternalFunction>,
+    {
+        // 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<T>(
+        &mut self,
+        x: T,
+        validator: Option<PrimitiveValidator>,
+        valid_ctxs: &[Context],
+        build_wrapper: &mut dyn FnMut(T, Context) -> Box<dyn ExternalFunction>,
+    ) where
+        T: Primitive + Clone,
     {
         let primitive: Arc<dyn Primitive> = Arc::new(x.clone());
         let name = primitive.name().to_owned();