Add primitive body typecheck helpers

[saulshanabrook]

Context

This is the egglog-side support needed by egraphs-good/egglog-experimental#46 for body-defined (primitive ...) commands. It intentionally contains reusable core APIs only; the experimental command implementation stays in egglog-experimental.

Current shape

• Adds EGraph::typecheck_expr_with_bindings_and_output for expression typechecking with explicit local bindings, an expected output sort, and a runtime primitive context.
• Ensures the output-constrained helper rejects resolved expressions whose final output sort still differs from the requested output sort, including literal bodies.
• Rewrites global references returned by that helper into the same zero-argument function calls used by normal command execution, so direct runtime evaluation can handle globals consistently.
• Adds Core::eval_resolved_expr so primitive bodies can evaluate already-typechecked and prepared expressions through the current execution state.
• Adds Core::apply_primitive for evaluating an already-specialized primitive against runtime values.
• Exposes PrimitiveWithId::is_valid_in_context so downstream code can infer the minimum runtime context for a resolved primitive target without reaching into the internal context-id map.
• Adds EGraph::prepare_unstable_fn_targets_for_eval, a narrow bridge that rewrites literal (unstable-fn "...") targets into hidden bindings containing the same ResolvedFunction value that backend rule lowering injects.
• Keeps top-level EGraph::eval_resolved_expr on the existing backend-rule execution path.
• Adds focused tests and an unreleased changelog entry.

Tried and decided against

• Tried a general PreparedResolvedExpr / ResolvedCoreActions bridge for direct runtime expression evaluation. Dropped it because it introduced a broader public concept and action-interpreter surface just to support one hidden unstable-fn runtime value.
• Considered migrating top-level EGraph::eval_resolved_expr to direct ExecutionState evaluation. Kept the backend-rule path instead, since top-level eval already gets the full lowering behavior there and does not need to share the primitive-body direct-eval path.
• Explored representing the resolved function target directly in the lowered IR / atom term layer. Parked that as a larger semantic cleanup so this PR stays isolated and minimal for the experimental primitive-body API.
• Does not include the egglog-experimental (primitive ...) command implementation and does not include the unrelated core-relations ordering fix from PR Fix multi-column column index rebuild ordering #899.

Validation

• Latest after merging upstream/main:
  • cargo test --lib typecheck_expr_with --quiet
  • cargo test --test typed_primitive --quiet
  • cargo check --lib
  • git diff --check
• Earlier on this PR:
  • make nits
  • cargo check --locked --lib
  • cargo test --lib --quiet
  • cargo test --test naive_unstable_fn_via_global --quiet
  • cargo test --test files --features bin -- unstable_fn --quiet
  • cargo test --test files --features bin -- typed_primitive_unstable_app --quiet
  • Separate review agent over this PR plus the dependent experimental PR: no remaining findings after the dependency pin was updated.

[codspeed-hq]

Merging this PR will not alter performance

✅ 26 untouched benchmarks
⏩ 216 skipped benchmarks¹

---

_{Comparing saulshanabrook:codex/split-primitive-body-runtime-apis (a80b235) with main (57108da)}

Footnotes

1. 216 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩

[codecov-commenter]

Codecov Report

❌ Patch coverage is 54.59770% with 158 lines in your changes missing coverage. Please review.
✅ Project coverage is 86.66%. Comparing base (57108da) to head (a80b235).

Files with missing lines	Patch %	Lines
src/lib.rs	63.75%	87 Missing ⚠️
src/exec_state.rs	0.00%	66 Missing ⚠️
src/typechecking.rs	87.80%	5 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #881      +/-   ##
==========================================
- Coverage   87.17%   86.66%   -0.52%     
==========================================
  Files          88       88              
  Lines       26106    26439     +333     
==========================================
+ Hits        22759    22913     +154     
- Misses       3347     3526     +179     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[saulshanabrook]

@coderabbitai full review

[coderabbitai]

Actionable comments posted: 3

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/exec_state.rs`:
- Around line 286-289: The Var arm in ResolvedExpr (ResolvedExpr::Var, using
resolved_var and bindings) only looks up caller-supplied local bindings and thus
treats globals as unbound; update this branch to explicitly handle
resolved_var.is_global_ref: if is_global_ref is true, look up the name in the
runtime/global store (the global bindings source used elsewhere in exec_state)
and return that value, otherwise fall back to the existing local bindings lookup
and preserve the EvalError::UnboundVariable behavior; alternatively, if globals
should be rejected earlier, add a guard where
typecheck_expr_with_bindings_and_output produces ResolvedVar with is_global_ref
and return a clear error before exposing the resolved expression API.
- Around line 114-129: The current apply_table_function gates authorization by
the runtime context via self.ctx(), which allows less-capable wrappers running
under a wider stamped Context to escalate; instead, check the wrapper/state
capability itself before calling action.lookup_or_insert or action.lookup.
Update apply_table_function to require the wrapper's explicit read/write
capability (e.g., a method like allows_write()/allows_read() or a
wrapper_capabilities() bitset) when matching FunctionSubtype::Constructor and
FunctionSubtype::Custom, and use es_mut()/es() only after that capability check;
apply the same change pattern to the other similar block mentioned (lines
231-252) so authorization is based on wrapper capability, not self.ctx().

In `@src/lib.rs`:
- Around line 2241-2250: The code currently calls
resolved_function_for_unstable_fn(...).unwrap_or_else(|err| panic!("{err}")),
which converts a modeled Result failure into a hard panic; instead change
BackendRule::prim (and callers in actions) to return Result and propagate the
error: remove the unwrap_or_else panic branch, propagate the Err(err) from
resolved_function_for_unstable_fn back to the caller of prim (and update call
sites in actions to handle/propagate that Result), and ensure any
panic_id/new_panic usage is only created when you will continue successfully;
i.e., construct panic_id only after resolved_function_for_unstable_fn succeeds
or discard it when returning Err so compilation yields a normal Error rather
than aborting.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: 1a3cac49-7839-4b2c-8349-0a2c63acfe89

📥 Commits

Reviewing files that changed from the base of the PR and between 2e1796e and cf4fe23.

📒 Files selected for processing (4)

• CHANGELOG.md
• src/exec_state.rs
• src/lib.rs
• src/typechecking.rs

📜 Review details

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (7)

• GitHub Check: benchmark (ubuntu-latest, conv1d_128)
• GitHub Check: benchmark (ubuntu-latest, eggcc-2mm)
• GitHub Check: benchmark (ubuntu-latest, rectangle)
• GitHub Check: benchmark (ubuntu-latest, taylor51)
• GitHub Check: benchmark (ubuntu-latest, math-microbenchmark)
• GitHub Check: coverage
• GitHub Check: test

🧰 Additional context used

🪛 LanguageTool

CHANGELOG.md

[style] ~10-~10: Three successive sentences begin with the same word. Consider rewording the sentence or use a thesaurus to find a synonym.
Context: ...ted separately from database updates. - Add `EGraph::typecheck_expr_with_bindings_a...

(ENGLISH_WORD_REPEAT_BEGINNING_RULE)

🔇 Additional comments (1)

src/lib.rs (1)

2535-2662: LGTM!

[yihozhang]

Some small questions. Overall looks good

[yihozhang]

Why did this function get refactored? What changed?

[saulshanabrook]

I believe from the merge from main, with the context stuff and also with supporting primtives doing the more advanced typea analysis. LMK if you have issues with how it is now.

[yihozhang]

Had a small question. Otherwise looks good. Thanks!

[oflatt]

Hmm, can we chat about this PR before we merge it? I'd like some context

[oflatt]

Had some questions

[oflatt]

Could you explain why we need this? I'm wary of new unstable-fn related code in core

[saulshanabrook]

This is needed to to do the same translation as in our pass that tests if fn_name == "unstable-fn" or whatever, that replaces the first arg of the function name as a string with the real resolved function value. So it just exposes the current behavior we have in core so others can you use it when preparing a resolved expression. Like this is how its used in the primitive command in experimental:

            match egraph.typecheck_expr_with_bindings_and_output(
                &args[3],
                &bindings,
                output_sort.clone(),
                context,
            ) {
                Ok(resolved)
                    if matches!(
                        (required_context(egraph, &resolved), context),
                        (Context::Pure, _)
                            | (Context::Read, Context::Read | Context::Full)
                            | (Context::Write, Context::Write | Context::Full)
                            | (Context::Full, Context::Full)
                    ) =>
                {
                    typechecked_body = Some((resolved, context));
                    break;
                }
                Ok(_) => {}
                Err(err) => last_error = Some(err),
            }
        }
        let Some((body, context)) = typechecked_body else {
            return Err(last_error
                .expect("primitive body typechecking always tries at least one context")
                .into());
        };

        let (body, hidden_bindings) = egraph.prepare_unstable_fn_targets_for_eval(&body)?;
        let primitive = DefinedPrimitive {
            name,
            input_vars: input_var_names,
            input: input_sorts,
            output: output_sort,
            body,
            hidden_bindings,
        };

I think to make this DRY and do the same code path on both would require more refactoring and make unstable function more in core.

I did have a go at this, but threw out that branch when it added more code, also after talking to @yihozhang about it. The "best" place to put this for overall design is in ResolvedExpr to add a specific case for a Function, but yeah this adds it even more inside core, so decided against that for now...

[oflatt]

Doesn't this mean if users get the binding types wrong they can get ill-typed programs? Is that a concern?

I imagine this is for custom commands

[saulshanabrook]

If they put in the binding types wrong, it won't typecheck I believe? Is that what you mean?

[oflatt]

Did we really not have this? Is this duplicate code with existing typecheck expr?

[yihozhang]

It's mostly not. We can extract some common code for setting up type checking with an optional output type, but this isn't much code anyway.

[saulshanabrook]

Yeah I did try to see if we could extract out some commonality, but I'll have another go at it, to see if it decreases repetition or just adds more issues.

[saulshanabrook]

I was able to refactor it so there are a few shared helpers here, do you guys like this more? A little more LoC but less duplication:

diff --git a/src/typechecking.rs b/src/typechecking.rs
index dedb17f6..c3623d0f 100644
--- a/src/typechecking.rs
+++ b/src/typechecking.rs
@@ -6,7 +6,8 @@ use crate::{
     *,
 };
 use ast::{
-    MappedExprExt, ResolvedAction, ResolvedExpr, ResolvedFact, ResolvedRule, ResolvedVar, Rule,
+    MappedActions, MappedExprExt, ResolvedAction, ResolvedExpr, ResolvedFact, ResolvedRule,
+    ResolvedVar, Rule,
 };
 use core_relations::ExternalFunction;
 use egglog_ast::generic_ast::GenericAction;
@@ -960,6 +961,18 @@ impl TypeInfo {
         binding: &IndexMap<&str, (Span, ArcSort)>,
         context: Context,
     ) -> Result<ResolvedActions, TypeError> {
+        let (problem, mapped_action) =
+            self.build_standalone_action_problem(symbol_gen, actions, binding, context)?;
+        self.solve_standalone_action_problem(problem, mapped_action, context)
+    }
+
+    fn build_standalone_action_problem(
+        &self,
+        symbol_gen: &mut SymbolGen,
+        actions: &Actions,
+        binding: &IndexMap<&str, (Span, ArcSort)>,
+        context: Context,
+    ) -> Result<(Problem<AtomTerm, ArcSort>, MappedActions<String, String>), TypeError> {
         let mut binding_set: IndexSet<String> =
             binding.keys().copied().map(str::to_string).collect();
         // We lower to core actions with `union_to_set_optimization`
@@ -975,6 +988,15 @@ impl TypeInfo {
             problem.assign_local_var_type(var, span.clone(), sort.clone())?;
         }
 
+        Ok((problem, mapped_action))
+    }
+
+    fn solve_standalone_action_problem(
+        &self,
+        problem: Problem<AtomTerm, ArcSort>,
+        mapped_action: MappedActions<String, String>,
+        context: Context,
+    ) -> Result<ResolvedActions, TypeError> {
         let assignment = problem
             .solve(|sort: &ArcSort| sort.name())
             .map_err(|e| e.to_type_error())?;
@@ -1008,29 +1030,20 @@ impl TypeInfo {
         context: Context,
     ) -> Result<ResolvedExpr, TypeError> {
         let action = Action::Expr(expr.span(), expr.clone());
-        let mut binding_set: IndexSet<String> =
-            binding.keys().copied().map(str::to_string).collect();
-        let mut ctx = CoreActionContext::new(self, &mut binding_set, symbol_gen, false);
-        let (actions, mapped_action) = Actions::singleton(action).to_core_actions(&mut ctx)?;
-        let mut problem = Problem::default();
-
-        problem.add_actions(&actions, self, symbol_gen, context)?;
-
-        for (var, (span, sort)) in binding {
-            problem.assign_local_var_type(var, span.clone(), sort.clone())?;
-        }
-
+        let (mut problem, mapped_action) = self.build_standalone_action_problem(
+            symbol_gen,
+            &Actions::singleton(action),
+            binding,
+            context,
+        )?;
         let [GenericAction::Expr(_, mapped_expr)] = mapped_action.0.as_slice() else {
             unreachable!("typechecking an expression should produce one expression action")
         };
         let output_atom = mapped_expr.get_corresponding_var_or_lit(self);
         problem.add_binding(output_atom, output_sort.clone());
 
-        let assignment = problem
-            .solve(|sort: &ArcSort| sort.name())
-            .map_err(|e| e.to_type_error())?;
-
-        let annotated_actions = assignment.annotate_actions(&mapped_action, self, context)?;
+        let annotated_actions =
+            self.solve_standalone_action_problem(problem, mapped_action, context)?;
         match annotated_actions.0.into_iter().next().unwrap() {
             ResolvedAction::Expr(_, resolved_expr) => {
                 let actual = resolved_expr.output_type();

[oflatt]

Do we really need new primitive declarations in the body of a rule? Is that the main motivation for this PR?

[saulshanabrook]

@oflatt Thanks for the comments!

This is for egraphs-good/egglog-experimental#46 which I needed for the A/C stuff with maps, so you can do:

(primitive my-prim (i64 i64) i64 (* (+ _0 _1) 3))

Do we really need new primitive declarations in the body of a rule? Is that the main motivation for this PR?

No we dont have new primitive declarations in the body of a rule, only at the top level!

We just have basically like "eval a resolved expression" in the body of a rule, since this primitive top level command is turned into resolving a primitive with the resolved body saved.

[saulshanabrook]

I talked to Oliver and he seems to generally feel OK about most of this PR since its pretty general, except for prepare_unstable_fn_targets_for_eval, because it leaks out details of how we are compiling unstable functions currently.

I tried to see if we could remove that API and instead just do it in experimental, however, we would have to expose the existing resolve_function_container_target_with_context function so that experimental could use it and this would be hard because it uses a bunch of things that are private.

We could make a wrapper around that to expose, but I am not sure if that resolves the objection.

Details from the agent:

Making the existing function pub is not sufficient. Experimental could not call it without exposing more internals.

Current parameters require:

• &egglog_bridge::EGraph: EGraph.backend is private.
• &IndexMap<String, Function>: EGraph.functions is private, and Function::backend_id is private.
• ExternalFunctionId: there is no public EGraph API to register the required panic handler.
• Primitive runtime IDs: the resolver reads context-specific IDs that are pub(crate).
• ResolvedFunctionId contains backend TableAction values, coupling experimental directly to backend internals.

ResolvedFunction, SpecializedPrimitive, TypeInfo, and EGraph::base_to_value are accessible, but the pieces needed to construct the value are not.

You could instead expose a wrapper:

pub fn resolve_function_target_for_eval(
    &mut self,
    name: &str,
    primitive: &SpecializedPrimitive,
) -> Result<Value, Error>

Because it takes &mut EGraph, its implementation could access all private fields internally. Experimental could then own the tree traversal and hidden bindings.

However, that does not really address Oliver’s objection. It replaces one high-level unstable-function API with a lower-level, more explicitly unstable-function API. The existing preparation method is arguably
less leaky because it does not expose ResolvedFunction construction details.

My recommendation:

1. If unstable-fn support in body-defined primitives is optional, remove it from this PR.
2. Remove prepare_unstable_fn_targets_for_eval.
3. Reject unstable-fn in experimental primitive bodies with a focused fail-typecheck test.
4. Leave proper function-value support for a later IR/prepared-expression refactor.

If support must remain, the current preparation method is probably the least damaging public boundary. Rust has no friend visibility between egglog and egglog-experimental, even in the same workspace.

Instead to do this "right" would require refactoring the compiler toolchain to handle unstable function more consistently and as first class. This would be a larger change and would be happy to look into this, but not sure if its right for this PR.

[oflatt]

Thanks for exploring. I'm going to approve to unblock this, and we can discuss next steps in the next egglog meeting.