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..074b798c0 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,121 @@ 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, +) -> ResolvedFunction { + let Some(target_function) = type_info + .get_sorts::() + .into_iter() + .find(|function| function.name() == primitive.output().name()) + else { + panic!( + "`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) + .unwrap_or_else(|| panic!("No resolution for {name:?}")); + let expected_inputs = partial_arcsorts + .iter() + .chain(remaining_inputs) + .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_input_names, + output.name(), + actual_input_names, + 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(); + 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 resolution for {name:?} in unstable-fn context {runtime_ctx:?}" + ), + } + }); + if !context_ids.iter().any(|(_, id)| id.is_some()) { + 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 { + panic!("No resolution for {name:?}"); + }; + + ResolvedFunction { + id, + partial_arcsorts, + name: name.to_owned(), + panic_id, + } +} + struct BackendRule<'a> { rb: egglog_bridge::RuleBuilder<'a>, entries: HashMap, @@ -2038,82 +2174,20 @@ 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, - }); + ); + + qe_args[0] = self.rb.egraph().base_value_constant(resolved); } ( diff --git a/src/sort/fn.rs b/src/sort/fn.rs index fc1b83386..e7f3f8ed4 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,137 @@ 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) { + // 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(); + primitive_constraints.push(Vec::new()); + let alternative_constraints = primitive_constraints.last_mut().unwrap(); + 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() + ), + ); + alternative_constraints + .push(constraint::assign(term.clone(), sort.clone())); + primitive_args.push(term); + } + alternative_constraints.extend( + primitive + .primitive + .get_type_constraints(&self.span) + .get(&primitive_args, typeinfo), + ); + } + + // 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 { + 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 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..ebe090838 --- /dev/null +++ b/tests/fail-typecheck/unstable-fn-primitive-missing-partial-arg.egg @@ -0,0 +1,18 @@ +;; 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)) +(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 new file mode 100644 index 000000000..b11319ab0 --- /dev/null +++ b/tests/snapshots/files__fail-typecheck__unstable_fn_primitive_missing_partial_arg.snap @@ -0,0 +1,24 @@ +--- +source: tests/files.rs +assertion_line: 106 +expression: err_msg +--- +All alternative definitions considered failed + 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 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 "+") +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 "+") +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 "+") +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 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)