Add primitive body typecheck helpers

CHANGELOG.md

@@ -7,6 +7,7 @@
 - Add a BigRat to-i64 primitive for integral rationals.
 - Add f64 exp, log, and sqrt primitives.
 - Add `RunReport::can_stop` so scheduler progress can be reported separately from database updates.
+- Add `EGraph::typecheck_expr_with_bindings_and_output`, `Core::eval_resolved_expr`, and `Core::apply_primitive` for body-defined primitive support.
 - Desugar `relation`s to `constructor`s to simplify the language and implementation. Relations no longer return unit `()` values.
 - Refactored API to use [`TermId`] more consistently instead of `Term` where possible, simplifying egglog code.
 - **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)`.

src/exec_state.rs

@@ -30,12 +30,18 @@
 //! [`ReadPrim`]: crate::ReadPrim
 //! [`FullPrim`]: crate::FullPrim
 
-use std::ops::Deref;
+use std::{fmt, ops::Deref};
 
 use crate::core_relations::{
     BaseValue, BaseValues, ContainerValue, ContainerValues, ExecutionState, ExternalFunctionId,
     TableId, Value,
 };
+use crate::{
+    ast::{FunctionSubtype, Literal, ResolvedExpr},
+    core::ResolvedCall,
+    sort::{F, S},
+    typechecking::FuncType,
+};
 use egglog_bridge::{ActionRegistry, TableAction};
 
 /// The four contexts a primitive may run in, named after the
@@ -101,6 +107,26 @@ pub(crate) trait Internal<'a, 'db: 'a>: 'a {
     fn raw_exec_state(&mut self) -> &mut ExecutionState<'db> {
         self.es_mut()
     }
+    fn apply_resolved_function(&mut self, _func: &FuncType, _args: &[Value]) -> Option<Value> {
+        None
+    }
+
+    fn apply_table_function(
+        &mut self,
+        subtype: FunctionSubtype,
+        action: &TableAction,
+        args: &[Value],
+    ) -> Option<Value> {
+        match (subtype, self.ctx()) {
+            (FunctionSubtype::Constructor, Context::Write | Context::Full) => {
+                action.lookup_or_insert(self.es_mut(), args)
+            }
+            (FunctionSubtype::Custom, Context::Read | Context::Full) => {
+                action.lookup(self.es(), args)
+            }
+            _ => None,
+        }
+    }
 }
 
 /// Sealed accessor for the [`ActionRegistry`]. Implemented by every
@@ -197,6 +223,97 @@ pub trait Core<'a, 'db: 'a>: Internal<'a, 'db> {
         let mut pure = PureState::wrap(self.raw_exec_state(), ctx);
         fc.apply(&mut pure, args)
     }
+
+    /// Dispatch an already type-specialized primitive in the current
+    /// call-site context.
+    ///
+    /// This is a trusted evaluator hook, not an authorization boundary:
+    /// callers must only pass primitives that were resolved for this same
+    /// call-site context and whose surrounding expression has already been
+    /// checked to require no more capability than this state wrapper provides.
+    /// For example, a primitive body evaluator should typecheck the body under
+    /// the runtime context it will register for, infer the body's required
+    /// context, and register the primitive using the matching state wrapper.
+    fn apply_primitive(
+        &mut self,
+        primitive: &crate::core::SpecializedPrimitive,
+        args: &[Value],
+    ) -> Option<Value> {
+        let id = primitive.external_id(self.ctx());
+        self.es_mut().call_external_func(id, args)
+    }
+
+    /// Evaluate an already typechecked expression in this primitive call
+    /// context, using `bindings` for local variables.
+    ///
+    /// Primitive calls dispatch through the current call-site context.
+    /// Table-backed function calls follow the same capability split as
+    /// `unstable-app`: custom functions require a read-capable state,
+    /// and constructors require a write-capable state that can mint on miss.
+    ///
+    /// This method assumes `expr` came from a trusted preparation pipeline that
+    /// typechecked it for this exact runtime context and rejected expressions
+    /// whose required capabilities exceed the receiver's state wrapper. It
+    /// should not be used to evaluate arbitrary resolved expressions from a
+    /// wider context inside a less-capable wrapper.
+    fn eval_resolved_expr(
+        &mut self,
+        expr: &ResolvedExpr,
+        bindings: &[(&str, Value)],
+    ) -> Option<Value> {
+        eval_resolved_expr_result(self, expr, bindings).ok()
+    }
+}
+
+#[derive(Debug)]
+pub(crate) enum EvalError {
+    UnboundVariable(String),
+    FunctionFailed(String),
+    PrimitiveFailed(String),
+}
+
+impl fmt::Display for EvalError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            EvalError::UnboundVariable(name) => write!(f, "unbound variable {name}"),
+            EvalError::FunctionFailed(name) => write!(f, "lookup of function {name} failed"),
+            EvalError::PrimitiveFailed(name) => write!(f, "call of primitive {name} failed"),
+        }
+    }
+}
+
+pub(crate) fn eval_resolved_expr_result<'a, 'db: 'a>(
+    state: &mut (impl Core<'a, 'db> + ?Sized),
+    expr: &ResolvedExpr,
+    bindings: &[(&str, Value)],
+) -> Result<Value, EvalError> {
+    match expr {
+        ResolvedExpr::Lit(_, literal) => Ok(match literal {
+            Literal::Int(x) => state.base_to_value(*x),
+            Literal::Float(x) => state.base_to_value(F::from(*x)),
+            Literal::String(x) => state.base_to_value(S::new(x.clone())),
+            Literal::Bool(x) => state.base_to_value(*x),
+            Literal::Unit => state.base_to_value(()),
+        }),
+        ResolvedExpr::Var(_, resolved_var) => bindings
+            .iter()
+            .find_map(|(name, value)| (*name == resolved_var.name).then_some(*value))
+            .ok_or_else(|| EvalError::UnboundVariable(resolved_var.name.clone())),
+        ResolvedExpr::Call(_, resolved_call, children) => {
+            let mut values = Vec::with_capacity(children.len());
+            for child in children {
+                values.push(eval_resolved_expr_result(state, child, bindings)?);
+            }
+            match resolved_call {
+                ResolvedCall::Primitive(primitive) => state
+                    .apply_primitive(primitive, &values)
+                    .ok_or_else(|| EvalError::PrimitiveFailed(primitive.name().to_owned())),
+                ResolvedCall::Func(func) => state
+                    .apply_resolved_function(func, &values)
+                    .ok_or_else(|| EvalError::FunctionFailed(func.name.clone())),
+            }
+        }
+    }
 }
 
 /// Read-side methods — name-indexed table lookup. Implemented for
@@ -270,6 +387,15 @@ fn lookup_action<'r>(registry: &'r ActionRegistry, name: &str) -> &'r TableActio
         .unwrap_or_else(|| panic!("missing table action for table: {name}"))
 }
 
+fn apply_registered_function<'a, 'db: 'a>(
+    state: &mut impl RegistrySealed<'a, 'db>,
+    func: &FuncType,
+    args: &[Value],
+) -> Option<Value> {
+    let action = lookup_action(state.registry(), &func.name).clone();
+    state.apply_table_function(func.subtype, &action, args)
+}
+
 // =====================================================================
 // The four wrapper types — plain structs, methods come from traits.
 // =====================================================================
@@ -422,6 +548,9 @@ impl<'a, 'db: 'a> Internal<'a, 'db> for ReadState<'a, 'db> {
     fn ctx(&self) -> Context {
         self.ctx
     }
+    fn apply_resolved_function(&mut self, func: &FuncType, args: &[Value]) -> Option<Value> {
+        apply_registered_function(self, func, args)
+    }
 }
 impl<'a, 'db: 'a> RegistrySealed<'a, 'db> for ReadState<'a, 'db> {
     fn registry(&self) -> &ActionRegistry {
@@ -441,6 +570,9 @@ impl<'a, 'db: 'a> Internal<'a, 'db> for WriteState<'a, 'db> {
     fn ctx(&self) -> Context {
         self.ctx
     }
+    fn apply_resolved_function(&mut self, func: &FuncType, args: &[Value]) -> Option<Value> {
+        apply_registered_function(self, func, args)
+    }
 }
 impl<'a, 'db: 'a> RegistrySealed<'a, 'db> for WriteState<'a, 'db> {
     fn registry(&self) -> &ActionRegistry {
@@ -460,6 +592,9 @@ impl<'a, 'db: 'a> Internal<'a, 'db> for FullState<'a, 'db> {
     fn ctx(&self) -> Context {
         self.ctx
     }
+    fn apply_resolved_function(&mut self, func: &FuncType, args: &[Value]) -> Option<Value> {
+        apply_registered_function(self, func, args)
+    }
 }
 impl<'a, 'db: 'a> RegistrySealed<'a, 'db> for FullState<'a, 'db> {
     fn registry(&self) -> &ActionRegistry {