egglog/sort/
fn.rs

1//! Sort to represent functions as values.
2//!
3//! To declare the sort, you must specify the exact number of arguments and the sort of each, followed by the output sort:
4//! `(sort IntToString (UnstableFn (i64) String))`
5//!
6//! To create a function value, use the `(unstable-fn "name" [<partial args>])` primitive and to apply it use the `(unstable-app function arg1 arg2 ...)` primitive.
7//! The number of args must match the number of arguments in the function sort.
8//!
9//! The value is stored similar to the `vec` sort, as an index into a set, where each item in
10//! the set is a `(Symbol, Vec<(Sort, Value)>)` pairs. The Symbol is the function name, and the `Vec<(Sort, Value)>` is
11//! the list of partially applied arguments.
12use std::any::TypeId;
13use std::sync::Mutex;
14
15use crate::exec_state::Internal;
16use enum_map::EnumMap;
17
18use super::*;
19
20#[derive(Clone, Debug)]
21pub struct FunctionContainer(
22    pub ResolvedFunctionId,
23    pub Vec<(ArcSort, Value)>,
24    pub String,
25    /// Pre-registered panic id used by `FunctionContainer::apply`
26    /// on capability mismatch (see [`ResolvedFunction::panic_id`]).
27    /// Excluded from equality/hash — two function values that differ
28    /// only in their panic id are still the same function value.
29    pub ExternalFunctionId,
30);
31
32// implement hash and equality based on values only not arcsorts, since
33// arcsorts are not comparable and any two values that are equal must have the same sort
34
35impl PartialEq for FunctionContainer {
36    fn eq(&self, other: &Self) -> bool {
37        self.0 == other.0
38            && self.1.iter().map(|(_, v)| *v).collect::<Vec<_>>()
39                == other.1.iter().map(|(_, v)| *v).collect::<Vec<_>>()
40            && self.2 == other.2
41    }
42}
43
44impl Eq for FunctionContainer {}
45
46impl Hash for FunctionContainer {
47    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
48        self.0.hash(state);
49        for (_, v) in &self.1 {
50            v.hash(state);
51        }
52        self.2.hash(state);
53    }
54}
55
56impl ContainerValue for FunctionContainer {
57    fn rebuild_contents(&mut self, rebuilder: &dyn ValueRebuilder) -> bool {
58        let mut changed = false;
59        for (s, old) in &mut self.1 {
60            if s.is_eq_sort() || s.is_eq_container_sort() {
61                let new = rebuilder.rebuild_val(*old);
62                changed |= *old != new;
63                *old = new;
64            }
65        }
66        changed
67    }
68    fn iter(&self) -> impl Iterator<Item = Value> + '_ {
69        self.1.iter().map(|(_, v)| v).copied()
70    }
71}
72#[derive(Debug)]
73pub struct FunctionSort {
74    name: String,
75    inputs: Vec<ArcSort>,
76    output: ArcSort,
77    // store all the arcsorts for functions that were added as partial args to this function sort
78    // so that we can retrieve them during extraction
79    partial_arcsorts: Arc<Mutex<Vec<ArcSort>>>,
80}
81
82impl FunctionSort {
83    pub fn name(&self) -> &str {
84        &self.name
85    }
86
87    pub fn inputs(&self) -> &[ArcSort] {
88        &self.inputs
89    }
90
91    pub fn output(&self) -> ArcSort {
92        self.output.clone()
93    }
94}
95
96impl Presort for FunctionSort {
97    fn presort_name() -> &'static str {
98        "UnstableFn"
99    }
100
101    fn reserved_primitives() -> Vec<&'static str> {
102        vec!["unstable-fn", "unstable-app"]
103    }
104
105    fn make_sort(
106        typeinfo: &mut TypeInfo,
107        name: String,
108        args: &[Expr],
109    ) -> Result<ArcSort, TypeError> {
110        if let [inputs, Expr::Var(span, output)] = args {
111            let output_sort = typeinfo
112                .get_sort_by_name(output)
113                .ok_or(TypeError::UndefinedSort(output.clone(), span.clone()))?;
114
115            let input_sorts = match inputs {
116                Expr::Call(_, first, rest_args) => {
117                    let all_args = once(first).chain(rest_args.iter().map(|arg| {
118                        if let Expr::Var(_, arg) = arg {
119                            arg
120                        } else {
121                            panic!("function sort must be called with list of input sorts");
122                        }
123                    }));
124                    all_args
125                        .map(|arg| {
126                            typeinfo
127                                .get_sort_by_name(arg)
128                                .ok_or(TypeError::UndefinedSort(arg.clone(), span.clone()))
129                                .cloned()
130                        })
131                        .collect::<Result<Vec<_>, _>>()?
132                }
133                // an empty list of inputs args is parsed as a unit literal
134                Expr::Lit(_, Literal::Unit) => vec![],
135                _ => panic!("function sort must be called with list of input sorts"),
136            };
137
138            Ok(Arc::new(Self {
139                name,
140                inputs: input_sorts,
141                output: output_sort.clone(),
142                partial_arcsorts: Arc::new(Mutex::new(vec![])),
143            }))
144        } else {
145            panic!("function sort must be called with list of input args and output sort");
146        }
147    }
148}
149
150impl Sort for FunctionSort {
151    fn name(&self) -> &str {
152        &self.name
153    }
154
155    fn column_ty(&self, _backend: &egglog_bridge::EGraph) -> ColumnTy {
156        ColumnTy::Id
157    }
158
159    fn register_type(&self, backend: &mut egglog_bridge::EGraph) {
160        backend.register_container_ty::<FunctionContainer>();
161        backend
162            .base_values_mut()
163            .register_type::<ResolvedFunction>();
164    }
165
166    fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
167        self
168    }
169
170    fn is_container_sort(&self) -> bool {
171        true
172    }
173
174    fn is_eq_container_sort(&self) -> bool {
175        self.inputs
176            .iter()
177            .any(|s| s.is_eq_sort() || s.is_eq_container_sort())
178    }
179
180    fn serialized_name(&self, container_values: &ContainerValues, value: Value) -> String {
181        let val = container_values
182            .get_val::<FunctionContainer>(value)
183            .unwrap();
184        val.2.clone()
185    }
186
187    fn inner_sorts(&self) -> Vec<ArcSort> {
188        self.partial_arcsorts.lock().unwrap().clone()
189    }
190
191    fn inner_values(
192        &self,
193        container_values: &ContainerValues,
194        value: Value,
195    ) -> Vec<(ArcSort, Value)> {
196        let val = container_values
197            .get_val::<FunctionContainer>(value)
198            .unwrap();
199        val.1.clone()
200    }
201
202    fn register_primitives(self: Arc<Self>, eg: &mut EGraph) {
203        eg.add_pure_primitive(
204            Ctor {
205                name: "unstable-fn".into(),
206                function: self.clone(),
207            },
208            None,
209        );
210        eg.add_pure_primitive(
211            Apply {
212                name: "unstable-app".into(),
213                function: self.clone(),
214            },
215            None,
216        );
217
218        register_vec_primitives_for_function(eg, self.clone());
219        register_multiset_primitives_for_function(eg, self.clone());
220    }
221
222    fn value_type(&self) -> Option<TypeId> {
223        Some(TypeId::of::<FunctionContainer>())
224    }
225
226    fn reconstruct_termdag_container(
227        &self,
228        container_values: &ContainerValues,
229        value: Value,
230        termdag: &mut TermDag,
231        mut element_terms: Vec<TermId>,
232    ) -> TermId {
233        let name = &container_values
234            .get_val::<FunctionContainer>(value)
235            .unwrap()
236            .2;
237        let head = termdag.lit(Literal::String(name.clone()));
238        element_terms.insert(0, head);
239        termdag.app("unstable-fn".to_owned(), element_terms)
240    }
241}
242
243/// Takes a string and any number of partially applied args of any sort and returns a function
244struct FunctionCTorTypeConstraint {
245    name: String,
246    function: Arc<FunctionSort>,
247    span: Span,
248}
249
250impl TypeConstraint for FunctionCTorTypeConstraint {
251    fn get(
252        &self,
253        arguments: &[AtomTerm],
254        typeinfo: &TypeInfo,
255    ) -> Vec<Box<dyn Constraint<AtomTerm, ArcSort>>> {
256        // Must have at least one arg (plus the return value)
257        if arguments.len() < 2 {
258            return vec![constraint::impossible(
259                constraint::ImpossibleConstraint::ArityMismatch {
260                    atom: core::Atom {
261                        span: self.span.clone(),
262                        head: self.name.clone(),
263                        args: arguments.to_vec(),
264                    },
265                    expected: 2,
266                },
267            )];
268        }
269        let output_sort_constraint: Box<dyn Constraint<_, ArcSort>> = constraint::assign(
270            arguments[arguments.len() - 1].clone(),
271            self.function.clone(),
272        );
273        // If first arg is a literal string and we know the name of the function and can use that to know what
274        // types to expect
275        if let AtomTerm::Literal(_, Literal::String(ref name)) = arguments[0] {
276            // The arguments contains the return sort as well as the function name
277            let n_partial_args = arguments.len() - 2;
278            if let Some(func_type) = typeinfo.get_func_type(name) {
279                // the number of partial args must match the number of inputs from the func type minus the number from
280                // this function sort
281                if self.function.inputs.len() + n_partial_args != func_type.input.len() {
282                    return vec![constraint::impossible(
283                        constraint::ImpossibleConstraint::ArityMismatch {
284                            atom: core::Atom {
285                                span: self.span.clone(),
286                                head: self.name.clone(),
287                                args: arguments.to_vec(),
288                            },
289                            expected: self.function.inputs.len() + func_type.input.len() + 1,
290                        },
291                    )];
292                }
293                // the output type and input types (starting after the partial args) must match between these functions
294                let expected_output = self.function.output.clone();
295                let expected_input = self.function.inputs.clone();
296                let actual_output = func_type.output.clone();
297                let actual_input: Vec<ArcSort> = func_type
298                    .input
299                    .iter()
300                    .skip(n_partial_args)
301                    .cloned()
302                    .collect();
303                if expected_output.name() != actual_output.name()
304                    || expected_input
305                        .iter()
306                        .map(|s| s.name())
307                        .ne(actual_input.iter().map(|s| s.name()))
308                {
309                    return vec![constraint::impossible(
310                        constraint::ImpossibleConstraint::FunctionMismatch {
311                            expected_output,
312                            expected_input,
313                            actual_output,
314                            actual_input,
315                        },
316                    )];
317                }
318                // if they match, then just make sure the partial args match as well
319                return func_type
320                    .input
321                    .iter()
322                    .take(n_partial_args)
323                    .zip(arguments.iter().skip(1))
324                    .map(|(expected_sort, actual_term)| {
325                        constraint::assign(actual_term.clone(), expected_sort.clone())
326                    })
327                    .chain(once(output_sort_constraint))
328                    .collect();
329            }
330
331            if let Some(primitives) = typeinfo.get_prims(name) {
332                // Primitive targets are checked by asking each overload whether
333                // a full call would typecheck after stitching together:
334                //
335                //   explicit partial args from `(unstable-fn "name" ...)`
336                //   + synthetic future args from the requested UnstableFn sort
337                //   + one synthetic output term
338                //
339                // For example, `(unstable-fn "+" old)` as `UnstableFn (i64) i64`
340                // checks each `+` overload as though it were called with
341                // `(old, future_arg) -> future_output`. The i64 overload matches;
342                // f64/string/etc. overloads become impossible constraints. If
343                // `old` is omitted, the same sort only provides one future arg,
344                // so no binary `+` overload has enough arguments to match.
345                let mut primitive_constraints = Vec::with_capacity(primitives.len());
346                for primitive in primitives {
347                    let mut primitive_args = arguments[1..arguments.len() - 1].to_vec();
348                    primitive_constraints.push(Vec::new());
349                    let alternative_constraints = primitive_constraints.last_mut().unwrap();
350                    for (index, sort) in self
351                        .function
352                        .inputs
353                        .iter()
354                        .chain(once(&self.function.output))
355                        .enumerate()
356                    {
357                        let term = AtomTerm::Var(
358                            self.span.clone(),
359                            format!(
360                                "__unstable_fn_target_{}_{}_arg_{index}",
361                                name,
362                                self.function.name()
363                            ),
364                        );
365                        alternative_constraints
366                            .push(constraint::assign(term.clone(), sort.clone()));
367                        primitive_args.push(term);
368                    }
369                    alternative_constraints.extend(
370                        primitive
371                            .primitive
372                            .get_type_constraints(&self.span)
373                            .get(&primitive_args, typeinfo),
374                    );
375                }
376
377                // No alternatives is defensive, one alternative is ordinary
378                // non-overloaded primitive resolution, and multiple alternatives
379                // are overloaded primitives such as `+`; the xor lets the type
380                // solver pick exactly one viable overload.
381                return match primitive_constraints.len() {
382                    0 => vec![constraint::impossible(
383                        constraint::ImpossibleConstraint::ArityMismatch {
384                            atom: core::Atom {
385                                span: self.span.clone(),
386                                head: self.name.clone(),
387                                args: arguments.to_vec(),
388                            },
389                            expected: n_partial_args + self.function.inputs.len() + 2,
390                        },
391                    )],
392                    1 => once(output_sort_constraint)
393                        .chain(primitive_constraints.pop().unwrap())
394                        .collect(),
395                    _ => vec![
396                        output_sort_constraint,
397                        constraint::xor(
398                            primitive_constraints
399                                .into_iter()
400                                .map(constraint::and)
401                                .collect(),
402                        ),
403                    ],
404                };
405            }
406        }
407
408        // Otherwise we just try assuming it's this function, we don't know if it is or not
409        vec![
410            constraint::assign(arguments[0].clone(), StringSort.to_arcsort()),
411            output_sort_constraint,
412        ]
413    }
414}
415
416// (unstable-fn "name" [<arg1>, <arg2>, ...])
417#[derive(Clone)]
418struct Ctor {
419    name: String,
420    function: Arc<FunctionSort>,
421}
422
423// `Ctor` (`unstable-fn "name" [...]`) builds a `FunctionContainer` and
424// interns it via `register_container`. Container interning is idempotent,
425// so it's safe in every context; declaring `State = PureState`
426// permits this primitive inside rule queries, actions, and global
427// contexts alike.
428impl Primitive for Ctor {
429    fn name(&self) -> &str {
430        &self.name
431    }
432
433    fn get_type_constraints(&self, span: &Span) -> Box<dyn TypeConstraint> {
434        Box::new(FunctionCTorTypeConstraint {
435            name: self.name.clone(),
436            function: self.function.clone(),
437            span: span.clone(),
438        })
439    }
440}
441
442impl PurePrim for Ctor {
443    fn apply<'a, 'db>(
444        &self,
445        mut state: crate::PureState<'a, 'db>,
446        args: &[Value],
447    ) -> Option<Value> {
448        let (rf, args) = args.split_first().unwrap();
449        let ResolvedFunction {
450            id,
451            partial_arcsorts,
452            name,
453            panic_id,
454        } = state.base_values().unwrap(*rf);
455        self.function
456            .partial_arcsorts
457            .lock()
458            .unwrap()
459            .extend(partial_arcsorts.iter().cloned());
460        let args = partial_arcsorts
461            .iter()
462            .zip(args)
463            .map(|(b, x)| (b.clone(), *x))
464            .collect();
465        let y = FunctionContainer(id, args, name, panic_id);
466        Some(state.register_container(y))
467    }
468}
469
470#[derive(Clone, Debug)]
471pub struct ResolvedFunction {
472    pub id: ResolvedFunctionId,
473    pub partial_arcsorts: Vec<ArcSort>,
474    pub name: String,
475    /// Pre-registered runtime-panic id used by `FunctionContainer::apply`
476    /// when an `unstable-fn` value is applied in a context where its
477    /// wrapped function isn't valid (e.g. constructor minting in a
478    /// rule body without `:naive`). Calling this id writes a
479    /// descriptive message to the egraph's panic side channel and
480    /// triggers early stop, so `run_rules` returns an `Err` rather
481    /// than the calling thread unwinding.
482    pub panic_id: ExternalFunctionId,
483}
484// implement equality and hash based on id and  arcsort names, since arcsorts are not comparable
485
486impl PartialEq for ResolvedFunction {
487    fn eq(&self, other: &Self) -> bool {
488        self.id == other.id
489            && self
490                .partial_arcsorts
491                .iter()
492                .map(|s| s.name())
493                .collect::<Vec<_>>()
494                == other
495                    .partial_arcsorts
496                    .iter()
497                    .map(|s| s.name())
498                    .collect::<Vec<_>>()
499    }
500}
501
502impl Eq for ResolvedFunction {}
503
504impl Hash for ResolvedFunction {
505    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
506        self.id.hash(state);
507        for s in &self.partial_arcsorts {
508            s.name().hash(state);
509        }
510    }
511}
512
513impl BaseValue for ResolvedFunction {}
514
515#[derive(Clone, Debug, PartialEq, Eq, Hash)]
516pub enum ResolvedFunctionId {
517    /// Wraps a constructor-table lookup. Only admissible in
518    /// write-capable contexts (`Write`/`Full`), where
519    /// `FunctionContainer::apply` mints a fresh eclass via
520    /// `lookup_or_insert`. In any read-only context (`Read`/`Pure`)
521    /// it triggers the pre-registered runtime panic — a no-mint
522    /// constructor would silently miss instead of producing the
523    /// eclass the user asked for, so the call is rejected outright.
524    Constructor(egglog_bridge::TableAction),
525    /// Wraps a `(function …)` lookup — any non-constructor function,
526    /// regardless of its `:merge` strategy. `FunctionContainer::apply`
527    /// allows this only in DB-read-capable contexts (`Read`/`Full`);
528    /// `Pure` and `Write` would be untracked seminaive reads.
529    Function(egglog_bridge::TableAction),
530    /// Wraps a primitive. Carries the unique exact-signature runtime
531    /// id found for each context at build time. At dispatch time
532    /// `FunctionContainer::apply` picks the id for the application
533    /// context — so the runtime selection is independent of the
534    /// build-site context, and an `unstable-fn` value may flow freely
535    /// from one context to another.
536    Primitive {
537        context_ids: EnumMap<crate::Context, Option<ExternalFunctionId>>,
538    },
539}
540
541// (unstable-app <function> [<arg1>, <arg2>, ...])
542//
543// Registered as a `PurePrim`; `FunctionContainer::apply` reads the
544// runtime context to dispatch. Distinct `FunctionSort`s produce
545// different signature keys, so `unstable-app` for `MathFn` stays a
546// separate overload from `unstable-app` for `i64Fun`.
547
548#[derive(Clone)]
549struct Apply {
550    name: String,
551    function: Arc<FunctionSort>,
552}
553
554impl Primitive for Apply {
555    fn name(&self) -> &str {
556        &self.name
557    }
558
559    fn get_type_constraints(&self, span: &Span) -> Box<dyn TypeConstraint> {
560        let mut sorts: Vec<ArcSort> = vec![self.function.clone()];
561        sorts.extend(self.function.inputs.clone());
562        sorts.push(self.function.output.clone());
563        SimpleTypeConstraint::new(&self.name, sorts, span.clone()).into_box()
564    }
565}
566
567impl PurePrim for Apply {
568    fn apply<'a, 'db>(
569        &self,
570        mut state: crate::PureState<'a, 'db>,
571        args: &[Value],
572    ) -> Option<Value> {
573        let (fc_val, args) = args.split_first().unwrap();
574        let fc = state
575            .container_values()
576            .get_val::<FunctionContainer>(*fc_val)
577            .unwrap()
578            .clone();
579        state.apply_function(&fc, args)
580    }
581}
582
583impl FunctionContainer {
584    /// Apply the wrapped function. `state` is always a `PureState`
585    /// (the type every primitive's `apply` receives). The surrounding
586    /// context is stamped onto that state by the primitive wrapper, so
587    /// callers do not pass a second copy of the same context.
588    pub(crate) fn apply<'a, 'db>(
589        &self,
590        state: &mut crate::PureState<'a, 'db>,
591        args: &[Value],
592    ) -> Option<Value>
593    where
594        'db: 'a,
595    {
596        let ctx = state.ctx();
597        let args: Vec<_> = self.1.iter().map(|(_, x)| x).chain(args).copied().collect();
598        let can_mint = matches!(ctx, crate::Context::Write | crate::Context::Full);
599        let can_read = matches!(ctx, crate::Context::Read | crate::Context::Full);
600        let panic_id = self.3;
601        // On capability mismatch, trigger the egglog runtime panic
602        // pre-registered at the `unstable-fn` build site (see
603        // `BackendRule::prim`). The panic writes to the egraph's
604        // panic side channel and triggers early stop, so `run_rules`
605        // surfaces the misuse as an `Err`.
606        let mismatch = |state: &mut crate::PureState<'a, 'db>| -> Option<Value> {
607            state.call_external_func(panic_id, &[])
608        };
609        match &self.0 {
610            ResolvedFunctionId::Constructor(action) => {
611                if can_mint {
612                    action.lookup_or_insert(state.raw_exec_state(), &args)
613                } else {
614                    mismatch(state)
615                }
616            }
617            ResolvedFunctionId::Function(action) => {
618                if can_read {
619                    action.lookup(state.raw_exec_state(), &args)
620                } else {
621                    mismatch(state)
622                }
623            }
624            ResolvedFunctionId::Primitive { context_ids } => {
625                // Pick the runtime id whose context matches the
626                // application ctx.
627                match context_ids[ctx] {
628                    Some(id) => state.call_external_func(id, &args),
629                    None => mismatch(state),
630                }
631            }
632        }
633    }
634}
egglog/sort/fn.rs

egglog/sort/
fn.rs
