From 80b41058d560907d5a4dc56329eb005301fe6a09 Mon Sep 17 00:00:00 2001 From: Saul Shanabrook Date: Sun, 10 May 2026 15:57:28 -0700 Subject: [PATCH 1/4] Allow unstable-fn to target primitives --- CHANGELOG.md | 1 + src/lib.rs | 251 +++++++++++++++++++++++++++++++++++-------------- src/sort/fn.rs | 160 +++++++++++++++++++++---------- 3 files changed, 289 insertions(+), 123 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f081fae03..bc9bc4ff6 100644 --- a/CHANGELOG.md +++ b/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. +- Allow `unstable-fn` function containers to target primitive overloads. - 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)`. diff --git a/src/lib.rs b/src/lib.rs index 300a2ae3a..887211f77 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -712,10 +712,31 @@ impl EGraph { )) } GenericExpr::Call(_, ResolvedCall::Primitive(p), args) => { - let translated_args = args + let mut translated_args = args .iter() .map(|arg| self.translate_expr_to_mergefn(arg)) .collect::, _>>()?; + if p.name() == "unstable-fn" { + let Some(GenericExpr::Lit(_, Literal::String(name))) = args.first() else { + return Err(Error::BackendError( + "expected string literal after `unstable-fn`".into(), + )); + }; + let resolved = resolve_function_container_target_with_context( + &self.backend, + &self.functions, + &self.type_info, + name, + p, + self.backend + .action_registry() + .read() + .unwrap() + .default_panic_id(), + )?; + translated_args[0] = + egglog_bridge::MergeFn::Const(self.backend.base_values().get(resolved)); + } Ok(egglog_bridge::MergeFn::Primitive( p.external_id(crate::Context::Write), translated_args, @@ -1944,6 +1965,127 @@ impl EGraph { } } +fn resolve_function_container_target_with_context( + backend: &egglog_bridge::EGraph, + functions: &IndexMap, + type_info: &TypeInfo, + name: &str, + primitive: &core::SpecializedPrimitive, + panic_id: ExternalFunctionId, +) -> Result { + let Some(target_function) = type_info + .get_sorts::() + .into_iter() + .find(|function| function.name() == primitive.output().name()) + else { + return Err(Error::BackendError(format!( + "`unstable-fn` output sort `{}` is not a function sort", + primitive.output().name() + ))); + }; + + let partial_arcsorts: Vec<_> = primitive.input().iter().skip(1).cloned().collect(); + let remaining_inputs = target_function.inputs(); + let output = target_function.output(); + + let id = if let Some(func) = functions.get(name) { + let func_type = type_info + .get_func_type(name) + .ok_or_else(|| TypeError::UnboundFunction(name.to_string(), span!()))?; + let expected_arity = partial_arcsorts.len() + remaining_inputs.len(); + let actual_inputs: Vec<_> = func_type.input.iter().map(|sort| sort.name()).collect(); + let expected_inputs: Vec<_> = partial_arcsorts + .iter() + .chain(remaining_inputs) + .map(|sort| sort.name()) + .collect(); + if func_type.input.len() != expected_arity + || actual_inputs != expected_inputs + || func_type.output.name() != output.name() + { + return Err(Error::BackendError(format!( + "function container lookup for `{name}` expected ({}) -> {}, found ({}) -> {}", + expected_inputs.join(", "), + output.name(), + actual_inputs.join(", "), + func_type.output.name(), + ))); + } + + let action = egglog_bridge::TableAction::new(backend, func.backend_id); + match func_type.subtype { + ast::FunctionSubtype::Constructor => ResolvedFunctionId::Constructor(action), + ast::FunctionSubtype::Custom => ResolvedFunctionId::Function(action), + } + } else if let Some(primitives) = type_info.get_prims(name) { + let signature: Vec<_> = partial_arcsorts + .iter() + .chain(remaining_inputs) + .chain(once(&output)) + .cloned() + .collect(); + let candidates: Vec<_> = primitives + .iter() + .filter(|primitive| primitive.accept(&signature, type_info)) + .collect(); + + if candidates.is_empty() { + Err(Error::BackendError(format!( + "no primitive overload for `{name}` matches ({}) -> {}", + signature[..signature.len() - 1] + .iter() + .map(|sort| sort.name()) + .collect::>() + .join(", "), + output.name(), + )))?; + } + + let context_ids = enum_map::EnumMap::from_fn(|runtime_ctx| { + let mut ids = candidates + .iter() + .filter_map(|primitive| primitive.context_ids[runtime_ctx]); + match (ids.next(), ids.next()) { + (None, _) => None, + (Some(id), None) => Some(id), + (Some(_), Some(_)) => panic!( + "ambiguous primitive function container lookup for `{name}` \ + with signature ({}) -> {} in context {runtime_ctx:?}", + signature[..signature.len() - 1] + .iter() + .map(|sort| sort.name()) + .collect::>() + .join(", "), + output.name(), + ), + } + }); + + if !context_ids.iter().any(|(_, id)| id.is_some()) { + Err(Error::BackendError(format!( + "no primitive overload for `{name}` matches ({}) -> {}", + signature[..signature.len() - 1] + .iter() + .map(|sort| sort.name()) + .collect::>() + .join(", "), + output.name(), + )))?; + } + + ResolvedFunctionId::Primitive { context_ids } + } else { + Err(TypeError::UnboundFunction(name.to_string(), span!()))? + }; + + Ok(ResolvedFunction { + id, + partial_arcsorts, + name: name.to_owned(), + panic_id, + }) +} + struct BackendRule<'a> { rb: egglog_bridge::RuleBuilder<'a>, entries: HashMap, @@ -2038,82 +2180,21 @@ impl<'a> BackendRule<'a> { let core::ResolvedAtomTerm::Literal(_, Literal::String(ref name)) = args[0] else { panic!("expected string literal after `unstable-fn`") }; - let id = if let Some(f) = self.type_info.get_func_type(name) { - // Distinguish constructor-table lookups (write-on-miss - // via `lookup_or_insert`) from custom-function lookups - // (pure read via `lookup`). `FunctionContainer::apply` - // uses this distinction to allow `unstable-app` over a - // custom function in any DB-read-capable context (e.g. - // global checks) while refusing constructors outright - // anywhere they can't mint — a no-mint constructor - // would silently miss instead of producing the eclass - // the user asked for. - let action = egglog_bridge::TableAction::new(self.rb.egraph(), self.func(f)); - match f.subtype { - ast::FunctionSubtype::Constructor => ResolvedFunctionId::Constructor(action), - ast::FunctionSubtype::Custom => ResolvedFunctionId::Function(action), - } - } else { - let fn_sort = Arc::downcast::(prim.output().clone().as_arc_any()) - .unwrap_or_else(|_| panic!("expected `unstable-fn` to return a function sort")); - let types: Vec<_> = prim - .input() - .iter() - .skip(1) - .cloned() - .chain(fn_sort.inputs().iter().cloned()) - .chain(std::iter::once(fn_sort.output())) - .collect(); - // Bake every exact-signature registration. The - // build-site `ctx` here is intentionally unused: the - // *application*-time context (carried on `state.ctx` - // by the wrapping HOF body) selects which candidate - // dispatches at runtime. Filtering by build-site ctx - // would trap an `unstable-fn` value in the context it - // was constructed in (e.g. a `let`-bound function value - // built at top-level `Full` couldn't be applied later - // from a `Read` query). - let ps: Vec<_> = self - .type_info - .get_prims(name) - .into_iter() - .flatten() - .filter(|p| p.accept(&types, self.type_info)) - .collect(); - let context_ids = enum_map::EnumMap::from_fn(|runtime_ctx| { - let mut ids = ps.iter().filter_map(|p| p.context_ids[runtime_ctx]); - match (ids.next(), ids.next()) { - (None, _) => None, - (Some(id), None) => Some(id), - (Some(_), Some(_)) => panic!( - "Ambiguous primitive resolution for {name:?} in unstable-fn context {runtime_ctx:?}" - ), - } - }); - assert!( - context_ids.iter().any(|(_, id)| id.is_some()), - "no callable for {name}" - ); - ResolvedFunctionId::Primitive { context_ids } - }; - let partial_arcsorts = prim.input().iter().skip(1).cloned().collect(); - - // Pre-register a panic id used by `FunctionContainer::apply` - // when the wrapped function is applied in a context that - // doesn't admit it. Triggered at runtime via the egglog - // panic side channel so misuse surfaces as an `Err` from - // `run_rules` rather than a thread unwind. let panic_id = self.rb.new_panic(format!( "unstable-fn over `{name}` was applied in a context where its wrapped \ function is not valid for this call site, if in a rule, add :naive." )); - - qe_args[0] = self.rb.egraph().base_value_constant(ResolvedFunction { - id, - partial_arcsorts, - name: name.clone(), + let resolved = resolve_function_container_target_with_context( + self.rb.egraph(), + self.functions, + self.type_info, + name, + prim, panic_id, - }); + ) + .unwrap_or_else(|err| panic!("{err}")); + + qe_args[0] = self.rb.egraph().base_value_constant(resolved); } ( @@ -2376,6 +2457,32 @@ mod tests { .unwrap(); } + #[test] + fn test_unstable_function_container_primitives() { + let mut egraph = EGraph::default(); + egraph + .parse_and_run_program( + None, + r#" + (sort IntThunk (UnstableFn () i64)) + (sort F64Thunk (UnstableFn () f64)) + (sort F64ToF64 (UnstableFn (f64) f64)) + (sort I64ToI64 (UnstableFn (i64) i64)) + (sort UnitThunk (UnstableFn () Unit)) + (sort UnitToF64 (UnstableFn (Unit) f64)) + + (function merge-plus () i64 :merge (unstable-app (unstable-fn "+" old) new)) + (set (merge-plus) 1) + (set (merge-plus) 2) + (check (= (merge-plus) 3)) + + (check (= (unstable-app (unstable-fn "+" 1.0) 2.0) 3.0)) + (check (= (unstable-app (unstable-fn "+" 1) 2) 3)) + "#, + ) + .unwrap(); + } + // Test that an `EGraph` is `Send` & `Sync` #[test] fn test_egraph_send_sync() { diff --git a/src/sort/fn.rs b/src/sort/fn.rs index fc1b83386..b2839af4b 100644 --- a/src/sort/fn.rs +++ b/src/sort/fn.rs @@ -6,10 +6,10 @@ //! To create a function value, use the `(unstable-fn "name" [])` primitive and to apply it use the `(unstable-app function arg1 arg2 ...)` primitive. //! The number of args must match the number of arguments in the function sort. //! -//! //! The value is stored similar to the `vec` sort, as an index into a set, where each item in //! the set is a `(Symbol, Vec<(Sort, Value)>)` pairs. The Symbol is the function name, and the `Vec<(Sort, Value)>` is //! the list of partially applied arguments. +use std::any::TypeId; use std::sync::Mutex; use crate::exec_state::Internal; @@ -272,61 +272,119 @@ impl TypeConstraint for FunctionCTorTypeConstraint { ); // If first arg is a literal string and we know the name of the function and can use that to know what // types to expect - if let AtomTerm::Literal(_, Literal::String(ref name)) = arguments[0] - && let Some(func_type) = typeinfo.get_func_type(name) - { + if let AtomTerm::Literal(_, Literal::String(ref name)) = arguments[0] { // The arguments contains the return sort as well as the function name let n_partial_args = arguments.len() - 2; - // the number of partial args must match the number of inputs from the func type minus the number from - // this function sort - if self.function.inputs.len() + n_partial_args != func_type.input.len() { - return vec![constraint::impossible( - constraint::ImpossibleConstraint::ArityMismatch { - atom: core::Atom { - span: self.span.clone(), - head: self.name.clone(), - args: arguments.to_vec(), + if let Some(func_type) = typeinfo.get_func_type(name) { + // the number of partial args must match the number of inputs from the func type minus the number from + // this function sort + if self.function.inputs.len() + n_partial_args != func_type.input.len() { + return vec![constraint::impossible( + constraint::ImpossibleConstraint::ArityMismatch { + atom: core::Atom { + span: self.span.clone(), + head: self.name.clone(), + args: arguments.to_vec(), + }, + expected: self.function.inputs.len() + func_type.input.len() + 1, }, - expected: self.function.inputs.len() + func_type.input.len() + 1, - }, - )]; - } - // the output type and input types (starting after the partial args) must match between these functions - let expected_output = self.function.output.clone(); - let expected_input = self.function.inputs.clone(); - let actual_output = func_type.output.clone(); - let actual_input: Vec = func_type - .input - .iter() - .skip(n_partial_args) - .cloned() - .collect(); - if expected_output.name() != actual_output.name() - || expected_input + )]; + } + // the output type and input types (starting after the partial args) must match between these functions + let expected_output = self.function.output.clone(); + let expected_input = self.function.inputs.clone(); + let actual_output = func_type.output.clone(); + let actual_input: Vec = func_type + .input .iter() - .map(|s| s.name()) - .ne(actual_input.iter().map(|s| s.name())) - { - return vec![constraint::impossible( - constraint::ImpossibleConstraint::FunctionMismatch { - expected_output, - expected_input, - actual_output, - actual_input, - }, - )]; + .skip(n_partial_args) + .cloned() + .collect(); + if expected_output.name() != actual_output.name() + || expected_input + .iter() + .map(|s| s.name()) + .ne(actual_input.iter().map(|s| s.name())) + { + return vec![constraint::impossible( + constraint::ImpossibleConstraint::FunctionMismatch { + expected_output, + expected_input, + actual_output, + actual_input, + }, + )]; + } + // if they match, then just make sure the partial args match as well + return func_type + .input + .iter() + .take(n_partial_args) + .zip(arguments.iter().skip(1)) + .map(|(expected_sort, actual_term)| { + constraint::assign(actual_term.clone(), expected_sort.clone()) + }) + .chain(once(output_sort_constraint)) + .collect(); + } + + if let Some(primitives) = typeinfo.get_prims(name) { + let mut primitive_constraints = Vec::with_capacity(primitives.len()); + for primitive in primitives { + let mut primitive_args = arguments[1..arguments.len() - 1].to_vec(); + let mut constraints = Vec::new(); + for (index, sort) in self + .function + .inputs + .iter() + .chain(once(&self.function.output)) + .enumerate() + { + let term = AtomTerm::Var( + self.span.clone(), + format!( + "__unstable_fn_target_{}_{}_arg_{index}", + name, + self.function.name() + ), + ); + constraints.push(constraint::assign(term.clone(), sort.clone())); + primitive_args.push(term); + } + constraints.extend( + primitive + .primitive + .get_type_constraints(&self.span) + .get(&primitive_args, typeinfo), + ); + primitive_constraints.push(constraints); + } + + return match primitive_constraints.len() { + 0 => vec![constraint::impossible( + constraint::ImpossibleConstraint::ArityMismatch { + atom: core::Atom { + span: self.span.clone(), + head: self.name.clone(), + args: arguments.to_vec(), + }, + expected: n_partial_args + self.function.inputs.len() + 2, + }, + )], + 1 => once(output_sort_constraint) + .chain(primitive_constraints.pop().unwrap()) + .collect(), + _ => vec![ + output_sort_constraint, + constraint::xor( + primitive_constraints + .into_iter() + .map(constraint::and) + .collect(), + ), + ], + }; } - // if they match, then just make sure the partial args match as well - return func_type - .input - .iter() - .take(n_partial_args) - .zip(arguments.iter().skip(1)) - .map(|(expected_sort, actual_term)| { - constraint::assign(actual_term.clone(), expected_sort.clone()) - }) - .chain(once(output_sort_constraint)) - .collect(); } // Otherwise we just try assuming it's this function, we don't know if it is or not From 554343d3ae34692363baf605b89da3fc191b2abc Mon Sep 17 00:00:00 2001 From: Saul Shanabrook Date: Tue, 26 May 2026 12:34:47 -0700 Subject: [PATCH 2/4] Allow unstable-fn to target primitives Teach function-container construction to resolve primitive targets as well as ordinary egglog functions. A primitive target is matched against the requested UnstableFn sort plus any explicit partial arguments, and the resulting FunctionContainer stores the runtime primitive ids for each execution context. This keeps the existing application-time dispatch behavior: the context at the unstable-app call site chooses the runtime id, so a function value built in one context is not trapped there. This also handles merge expressions where unstable-fn appears inside a primitive merge function. The merge-function translator now replaces the target name with the same resolved FunctionContainer value used by rule/query lowering, so partially applied primitive targets like `(unstable-fn "+" old)` work as merge functions. Tests now cover primitive targets across nullary, unary, Unit-returning, overloaded, higher-order, container-specific, and merge contexts: - `vec-empty` and `map-empty` as nullary container primitives. - `not` and `<` as unary / Unit-returning primitives. - `+` for both i64 and f64 overloads. - `unstable-vec-map` wrapped through unstable-fn as a higher-order primitive. - a merge function using partially applied `+`. - a fail-typecheck case showing `(unstable-fn "+")` is not enough for an `UnstableFn (i64) i64` merge target without the missing `old` argument. One cleanup is intentionally deferred. Direct duplicate same-signature primitive registrations already panic in `ResolvedCall::from_resolution` after typechecking has concrete sorts but more than one indistinguishable primitive registration remains. This commit keeps primitive unstable-fn targets aligned with that behavior: duplicate exact primitive matches still panic rather than threading a new TypeError/Error path through rule lowering and scheduler rule construction. That deferred work is tracked in #898. The follow-up should make primitive resolution return typed errors for no-match and ambiguous exact-match cases in both direct calls and unstable-fn targets. It should also harden `step_rules_with_scheduler`: that method temporarily takes `self.rulesets` and `self.schedulers`, and the existing fallible `backend.run_rules(...)` paths can return before those fields are restored. A previous version of this patch exposed the same restoration problem through fallible scheduler rule construction, but fixing that scheduler robustness issue is broader than this primitive-target feature and is left for the follow-up. Validation: - cargo fmt --check - cargo test --test typed_primitive - cargo test --test files --features bin typed_primitive_unstable_app - cargo test --test files --features bin unstable_fn_primitive_missing_partial_arg - cargo test --lib scheduler::test - make nits - make test Reviews: - CodeRabbit reported a trivial panic-message issue, fixed here. - Independent sub-agent review reported no remaining blocker, major, or minor findings after the added map-empty / unstable-vec-map coverage. --- src/lib.rs | 129 +++++++----------- src/sort/fn.rs | 17 +++ ...table-fn-primitive-missing-partial-arg.egg | 7 + ...able_fn_primitive_missing_partial_arg.snap | 15 ++ ...snapshot_typed_primitive_unstable_app.snap | 3 +- tests/typed_primitive_unstable_app.egg | 29 +++- 6 files changed, 117 insertions(+), 83 deletions(-) create mode 100644 tests/fail-typecheck/unstable-fn-primitive-missing-partial-arg.egg create mode 100644 tests/snapshots/files__fail-typecheck__unstable_fn_primitive_missing_partial_arg.snap diff --git a/src/lib.rs b/src/lib.rs index 887211f77..074b798c0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -733,7 +733,7 @@ impl EGraph { .read() .unwrap() .default_panic_id(), - )?; + ); translated_args[0] = egglog_bridge::MergeFn::Const(self.backend.base_values().get(resolved)); } @@ -1972,16 +1972,16 @@ fn resolve_function_container_target_with_context( name: &str, primitive: &core::SpecializedPrimitive, panic_id: ExternalFunctionId, -) -> Result { +) -> ResolvedFunction { let Some(target_function) = type_info .get_sorts::() .into_iter() .find(|function| function.name() == primitive.output().name()) else { - return Err(Error::BackendError(format!( + panic!( "`unstable-fn` output sort `{}` is not a function sort", primitive.output().name() - ))); + ); }; let partial_arcsorts: Vec<_> = primitive.input().iter().skip(1).cloned().collect(); @@ -1991,25 +1991,36 @@ fn resolve_function_container_target_with_context( let id = if let Some(func) = functions.get(name) { let func_type = type_info .get_func_type(name) - .ok_or_else(|| TypeError::UnboundFunction(name.to_string(), span!()))?; - let expected_arity = partial_arcsorts.len() + remaining_inputs.len(); - let actual_inputs: Vec<_> = func_type.input.iter().map(|sort| sort.name()).collect(); - let expected_inputs: Vec<_> = partial_arcsorts + .unwrap_or_else(|| panic!("No resolution for {name:?}")); + let expected_inputs = partial_arcsorts .iter() .chain(remaining_inputs) - .map(|sort| sort.name()) - .collect(); - if func_type.input.len() != expected_arity - || actual_inputs != expected_inputs - || func_type.output.name() != output.name() - { - return Err(Error::BackendError(format!( + .collect::>(); + let inputs_match = func_type.input.len() == expected_inputs.len() + && func_type + .input + .iter() + .zip(&expected_inputs) + .all(|(actual, expected)| actual.name() == expected.name()); + if !inputs_match || func_type.output.name() != output.name() { + let expected_input_names = expected_inputs + .iter() + .map(|sort| sort.name()) + .collect::>() + .join(", "); + let actual_input_names = func_type + .input + .iter() + .map(|sort| sort.name()) + .collect::>() + .join(", "); + panic!( "function container lookup for `{name}` expected ({}) -> {}, found ({}) -> {}", - expected_inputs.join(", "), + expected_input_names, output.name(), - actual_inputs.join(", "), + actual_input_names, func_type.output.name(), - ))); + ); } let action = egglog_bridge::TableAction::new(backend, func.backend_id); @@ -2028,19 +2039,6 @@ fn resolve_function_container_target_with_context( .iter() .filter(|primitive| primitive.accept(&signature, type_info)) .collect(); - - if candidates.is_empty() { - Err(Error::BackendError(format!( - "no primitive overload for `{name}` matches ({}) -> {}", - signature[..signature.len() - 1] - .iter() - .map(|sort| sort.name()) - .collect::>() - .join(", "), - output.name(), - )))?; - } - let context_ids = enum_map::EnumMap::from_fn(|runtime_ctx| { let mut ids = candidates .iter() @@ -2049,41 +2047,37 @@ fn resolve_function_container_target_with_context( (None, _) => None, (Some(id), None) => Some(id), (Some(_), Some(_)) => panic!( - "ambiguous primitive function container lookup for `{name}` \ - with signature ({}) -> {} in context {runtime_ctx:?}", - signature[..signature.len() - 1] - .iter() - .map(|sort| sort.name()) - .collect::>() - .join(", "), - output.name(), + "Ambiguous primitive resolution for {name:?} in unstable-fn context {runtime_ctx:?}" ), } }); - if !context_ids.iter().any(|(_, id)| id.is_some()) { - Err(Error::BackendError(format!( - "no primitive overload for `{name}` matches ({}) -> {}", - signature[..signature.len() - 1] - .iter() - .map(|sort| sort.name()) - .collect::>() - .join(", "), - output.name(), - )))?; + let (output_sort, input_sorts) = signature + .split_last() + .expect("primitive signature should include an output sort"); + let input_names = input_sorts + .iter() + .map(|sort| sort.name()) + .collect::>() + .join(", "); + panic!( + "no primitive overload matched expected signature for {name:?}: ({}) -> {}; \ + context ids: {context_ids:?}", + input_names, + output_sort.name(), + ); } - ResolvedFunctionId::Primitive { context_ids } } else { - Err(TypeError::UnboundFunction(name.to_string(), span!()))? + panic!("No resolution for {name:?}"); }; - Ok(ResolvedFunction { + ResolvedFunction { id, partial_arcsorts, name: name.to_owned(), panic_id, - }) + } } struct BackendRule<'a> { @@ -2191,8 +2185,7 @@ impl<'a> BackendRule<'a> { name, prim, panic_id, - ) - .unwrap_or_else(|err| panic!("{err}")); + ); qe_args[0] = self.rb.egraph().base_value_constant(resolved); } @@ -2457,32 +2450,6 @@ mod tests { .unwrap(); } - #[test] - fn test_unstable_function_container_primitives() { - let mut egraph = EGraph::default(); - egraph - .parse_and_run_program( - None, - r#" - (sort IntThunk (UnstableFn () i64)) - (sort F64Thunk (UnstableFn () f64)) - (sort F64ToF64 (UnstableFn (f64) f64)) - (sort I64ToI64 (UnstableFn (i64) i64)) - (sort UnitThunk (UnstableFn () Unit)) - (sort UnitToF64 (UnstableFn (Unit) f64)) - - (function merge-plus () i64 :merge (unstable-app (unstable-fn "+" old) new)) - (set (merge-plus) 1) - (set (merge-plus) 2) - (check (= (merge-plus) 3)) - - (check (= (unstable-app (unstable-fn "+" 1.0) 2.0) 3.0)) - (check (= (unstable-app (unstable-fn "+" 1) 2) 3)) - "#, - ) - .unwrap(); - } - // Test that an `EGraph` is `Send` & `Sync` #[test] fn test_egraph_send_sync() { diff --git a/src/sort/fn.rs b/src/sort/fn.rs index b2839af4b..ab7583dce 100644 --- a/src/sort/fn.rs +++ b/src/sort/fn.rs @@ -329,6 +329,19 @@ impl TypeConstraint for FunctionCTorTypeConstraint { } if let Some(primitives) = typeinfo.get_prims(name) { + // Primitive targets are checked by asking each overload whether + // a full call would typecheck after stitching together: + // + // explicit partial args from `(unstable-fn "name" ...)` + // + synthetic future args from the requested UnstableFn sort + // + one synthetic output term + // + // For example, `(unstable-fn "+" old)` as `UnstableFn (i64) i64` + // checks each `+` overload as though it were called with + // `(old, future_arg) -> future_output`. The i64 overload matches; + // f64/string/etc. overloads become impossible constraints. If + // `old` is omitted, the same sort only provides one future arg, + // so no binary `+` overload has enough arguments to match. let mut primitive_constraints = Vec::with_capacity(primitives.len()); for primitive in primitives { let mut primitive_args = arguments[1..arguments.len() - 1].to_vec(); @@ -360,6 +373,10 @@ impl TypeConstraint for FunctionCTorTypeConstraint { primitive_constraints.push(constraints); } + // No alternatives is defensive, one alternative is ordinary + // non-overloaded primitive resolution, and multiple alternatives + // are overloaded primitives such as `+`; the xor lets the type + // solver pick exactly one viable overload. return match primitive_constraints.len() { 0 => vec![constraint::impossible( constraint::ImpossibleConstraint::ArityMismatch { diff --git a/tests/fail-typecheck/unstable-fn-primitive-missing-partial-arg.egg b/tests/fail-typecheck/unstable-fn-primitive-missing-partial-arg.egg new file mode 100644 index 000000000..a04f04ffc --- /dev/null +++ b/tests/fail-typecheck/unstable-fn-primitive-missing-partial-arg.egg @@ -0,0 +1,7 @@ +;; A primitive target must be partially applied enough to match the declared +;; function sort. Here `+` is treated as an i64 => i64 function, but no `old` +;; argument was supplied, so no `+` overload can resolve to that sort. + +(sort I64ToI64 (UnstableFn (i64) i64)) +(function merge-plus () i64 :merge (unstable-app (unstable-fn "+") new)) +(set (merge-plus) 1) diff --git a/tests/snapshots/files__fail-typecheck__unstable_fn_primitive_missing_partial_arg.snap b/tests/snapshots/files__fail-typecheck__unstable_fn_primitive_missing_partial_arg.snap new file mode 100644 index 000000000..c996070bd --- /dev/null +++ b/tests/snapshots/files__fail-typecheck__unstable_fn_primitive_missing_partial_arg.snap @@ -0,0 +1,15 @@ +--- +source: tests/files.rs +expression: err_msg +--- +All alternative definitions considered failed + In 6:50-66 of unstable-fn-primitive-missing-partial-arg.egg: (unstable-fn "+") + Expect expression __unstable_fn_target_+_I64ToI64_arg_1 to have type String, but get type i64 + In 6:50-66 of unstable-fn-primitive-missing-partial-arg.egg: (unstable-fn "+") +Arity mismatch, expected 2 args: (+ __unstable_fn_target_+_I64ToI64_arg_0) + In 6:50-66 of unstable-fn-primitive-missing-partial-arg.egg: (unstable-fn "+") +Arity mismatch, expected 2 args: (+ __unstable_fn_target_+_I64ToI64_arg_0) + In 6:50-66 of unstable-fn-primitive-missing-partial-arg.egg: (unstable-fn "+") +Arity mismatch, expected 2 args: (+ __unstable_fn_target_+_I64ToI64_arg_0) + In 6:50-66 of unstable-fn-primitive-missing-partial-arg.egg: (unstable-fn "+") +Arity mismatch, expected 2 args: (+ __unstable_fn_target_+_I64ToI64_arg_0) diff --git a/tests/snapshots/files__shared_snapshot_typed_primitive_unstable_app.snap b/tests/snapshots/files__shared_snapshot_typed_primitive_unstable_app.snap index 14dad6481..7bb2edc9f 100644 --- a/tests/snapshots/files__shared_snapshot_typed_primitive_unstable_app.snap +++ b/tests/snapshots/files__shared_snapshot_typed_primitive_unstable_app.snap @@ -5,4 +5,5 @@ expression: snapshot_content_across_treatments ((Num 1) (Trigger 1) (dbl 1) - (dbl-out 1)) + (dbl-out 1) + (merge-plus 1)) diff --git a/tests/typed_primitive_unstable_app.egg b/tests/typed_primitive_unstable_app.egg index 392235ff4..f8e0d1561 100644 --- a/tests/typed_primitive_unstable_app.egg +++ b/tests/typed_primitive_unstable_app.egg @@ -10,6 +10,34 @@ (let plus (unstable-fn "+")) (check (= (unstable-app plus 2 3) 5)) +; --- 1b. Primitive targets cover nullary, unary, Unit-returning, overloaded, +; higher-order, container-specific, and merge-function contexts. The merge +; case partially applies `+` to `old`, so the resulting unstable function +; consumes only `new`. +(sort IntVec (Vec i64)) +(sort IntVecThunk (UnstableFn () IntVec)) +(sort BoolVec (Vec bool)) +(sort MapI64String (Map i64 String)) +(sort MapI64StringThunk (UnstableFn () MapI64String)) +(sort BoolToBool (UnstableFn (bool) bool)) +(sort VecBoolMap (UnstableFn (BoolToBool BoolVec) BoolVec)) +(sort IntToUnit (UnstableFn (i64) Unit)) +(sort F64ToF64 (UnstableFn (f64) f64)) +(sort I2I (UnstableFn (i64) i64)) + +(function merge-plus () i64 :merge (unstable-app (unstable-fn "+" old) new)) +(set (merge-plus) 1) +(set (merge-plus) 2) +(check (= (merge-plus) 3)) + +(check (= (unstable-app (unstable-fn "vec-empty")) (vec-empty))) +(check (= (unstable-app (unstable-fn "map-empty")) (map-empty))) +(check (= (unstable-app (unstable-fn "not") true) false)) +(check (= (unstable-app (unstable-fn "unstable-vec-map") (unstable-fn "not") (vec-of true false)) (vec-of false true))) +(check (= (unstable-app (unstable-fn "<" 1) 2) ())) +(check (= (unstable-app (unstable-fn "+" 1.0) 2.0) 3.0)) +(check (= (unstable-app (unstable-fn "+" 1) 2) 3)) + ; --- 2. unstable-app over a constructor in an action mints the eclass. (datatype Math (Num i64)) (sort MathFn (UnstableFn (i64) Math)) @@ -43,7 +71,6 @@ ; `Context::Full` with `:naive`. (function dbl (i64) i64 :no-merge) (set (dbl 7) 14) -(sort I2I (UnstableFn (i64) i64)) (let twice (unstable-fn "dbl")) (function dbl-out (i64) i64 :no-merge) (rule ((Trigger)) ((set (dbl-out 7) (unstable-app twice 7))) :naive) From aef6b97e59eac7c727f10a142f6b64a3ef1d1009 Mon Sep 17 00:00:00 2001 From: Saul Shanabrook Date: Sun, 31 May 2026 06:18:19 -0700 Subject: [PATCH 3/4] Address unstable-fn primitive review --- src/sort/fn.rs | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/sort/fn.rs b/src/sort/fn.rs index ab7583dce..e7f3f8ed4 100644 --- a/src/sort/fn.rs +++ b/src/sort/fn.rs @@ -345,7 +345,8 @@ impl TypeConstraint for FunctionCTorTypeConstraint { let mut primitive_constraints = Vec::with_capacity(primitives.len()); for primitive in primitives { let mut primitive_args = arguments[1..arguments.len() - 1].to_vec(); - let mut constraints = Vec::new(); + primitive_constraints.push(Vec::new()); + let alternative_constraints = primitive_constraints.last_mut().unwrap(); for (index, sort) in self .function .inputs @@ -361,16 +362,16 @@ impl TypeConstraint for FunctionCTorTypeConstraint { self.function.name() ), ); - constraints.push(constraint::assign(term.clone(), sort.clone())); + alternative_constraints + .push(constraint::assign(term.clone(), sort.clone())); primitive_args.push(term); } - constraints.extend( + alternative_constraints.extend( primitive .primitive .get_type_constraints(&self.span) .get(&primitive_args, typeinfo), ); - primitive_constraints.push(constraints); } // No alternatives is defensive, one alternative is ordinary From 62121dfc6add323ba473c8919e028b15e41685f4 Mon Sep 17 00:00:00 2001 From: Saul Shanabrook Date: Sun, 31 May 2026 06:26:33 -0700 Subject: [PATCH 4/4] Refine unstable-fn primitive arity test --- ...table-fn-primitive-missing-partial-arg.egg | 21 ++++++++++++++----- ...able_fn_primitive_missing_partial_arg.snap | 19 ++++++++++++----- 2 files changed, 30 insertions(+), 10 deletions(-) diff --git a/tests/fail-typecheck/unstable-fn-primitive-missing-partial-arg.egg b/tests/fail-typecheck/unstable-fn-primitive-missing-partial-arg.egg index a04f04ffc..ebe090838 100644 --- a/tests/fail-typecheck/unstable-fn-primitive-missing-partial-arg.egg +++ b/tests/fail-typecheck/unstable-fn-primitive-missing-partial-arg.egg @@ -1,7 +1,18 @@ -;; A primitive target must be partially applied enough to match the declared -;; function sort. Here `+` is treated as an i64 => i64 function, but no `old` -;; argument was supplied, so no `+` overload can resolve to that sort. +;; A primitive target must be partially applied enough to match the inferred +;; function sort. The first two checks resolve binary `+` without partial args. +;; The next two partial applications turn binary `+` into unary functions for +;; i64 and f64. The final check still calls the inferred unary function with the +;; right number of arguments, but no partial argument was supplied to make +;; binary `+` unary. +(sort I64Add (UnstableFn (i64 i64) i64)) +(sort F64Add (UnstableFn (f64 f64) f64)) (sort I64ToI64 (UnstableFn (i64) i64)) -(function merge-plus () i64 :merge (unstable-app (unstable-fn "+") new)) -(set (merge-plus) 1) +(sort F64ToF64 (UnstableFn (f64) f64)) + +(check (= (unstable-app (unstable-fn "+") 1 2) 3)) +(check (= (unstable-app (unstable-fn "+") 1.0 2.0) 3.0)) +(check (= (unstable-app (unstable-fn "+" 1) 2) 3)) +(check (= (unstable-app (unstable-fn "+" 1.0) 2.0) 3.0)) + +(check (= (unstable-app (unstable-fn "+") 2) 3)) diff --git a/tests/snapshots/files__fail-typecheck__unstable_fn_primitive_missing_partial_arg.snap b/tests/snapshots/files__fail-typecheck__unstable_fn_primitive_missing_partial_arg.snap index c996070bd..b11319ab0 100644 --- a/tests/snapshots/files__fail-typecheck__unstable_fn_primitive_missing_partial_arg.snap +++ b/tests/snapshots/files__fail-typecheck__unstable_fn_primitive_missing_partial_arg.snap @@ -1,15 +1,24 @@ --- source: tests/files.rs +assertion_line: 106 expression: err_msg --- All alternative definitions considered failed - In 6:50-66 of unstable-fn-primitive-missing-partial-arg.egg: (unstable-fn "+") + In 18:25-41 of unstable-fn-primitive-missing-partial-arg.egg: (unstable-fn "+") + Expect expression @unstable-fn8 to have type I64Add, but get type I64ToI64 + In 18:25-41 of unstable-fn-primitive-missing-partial-arg.egg: (unstable-fn "+") + Expect expression @unstable-fn8 to have type F64Add, but get type I64ToI64 + All alternative definitions considered failed + In 18:25-41 of unstable-fn-primitive-missing-partial-arg.egg: (unstable-fn "+") Expect expression __unstable_fn_target_+_I64ToI64_arg_1 to have type String, but get type i64 - In 6:50-66 of unstable-fn-primitive-missing-partial-arg.egg: (unstable-fn "+") + In 18:25-41 of unstable-fn-primitive-missing-partial-arg.egg: (unstable-fn "+") Arity mismatch, expected 2 args: (+ __unstable_fn_target_+_I64ToI64_arg_0) - In 6:50-66 of unstable-fn-primitive-missing-partial-arg.egg: (unstable-fn "+") + In 18:25-41 of unstable-fn-primitive-missing-partial-arg.egg: (unstable-fn "+") Arity mismatch, expected 2 args: (+ __unstable_fn_target_+_I64ToI64_arg_0) - In 6:50-66 of unstable-fn-primitive-missing-partial-arg.egg: (unstable-fn "+") + In 18:25-41 of unstable-fn-primitive-missing-partial-arg.egg: (unstable-fn "+") Arity mismatch, expected 2 args: (+ __unstable_fn_target_+_I64ToI64_arg_0) - In 6:50-66 of unstable-fn-primitive-missing-partial-arg.egg: (unstable-fn "+") + In 18:25-41 of unstable-fn-primitive-missing-partial-arg.egg: (unstable-fn "+") Arity mismatch, expected 2 args: (+ __unstable_fn_target_+_I64ToI64_arg_0) + + In 18:25-41 of unstable-fn-primitive-missing-partial-arg.egg: (unstable-fn "+") + Expect expression @unstable-fn8 to have type F64ToF64, but get type I64ToI64