egglog 7.1.0__cp311-none-win_amd64.whl → 8.0.0__cp311-none-win_amd64.whl

egglog/egraph_state.py

@@ -4,10 +4,10 @@ Implement conversion to/from egglog.
 
 from __future__ import annotations
 
+import re
 from collections import defaultdict
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, overload
-from weakref import WeakKeyDictionary
 
 from typing_extensions import assert_never
 
@@ -52,7 +52,7 @@ class EGraphState:
     type_ref_to_egg_sort: dict[JustTypeRef, str] = field(default_factory=dict)
 
     # Cache of egg expressions for converting to egg
-    expr_to_egg_cache: WeakKeyDictionary[ExprDecl, bindings._Expr] = field(default_factory=WeakKeyDictionary)
+    expr_to_egg_cache: dict[ExprDecl, bindings._Expr] = field(default_factory=dict)
 
     def copy(self) -> EGraphState:
         """
@@ -71,15 +71,15 @@ class EGraphState:
     def schedule_to_egg(self, schedule: ScheduleDecl) -> bindings._Schedule:
         match schedule:
             case SaturateDecl(schedule):
-                return bindings.Saturate(self.schedule_to_egg(schedule))
+                return bindings.Saturate(bindings.DUMMY_SPAN, self.schedule_to_egg(schedule))
             case RepeatDecl(schedule, times):
-                return bindings.Repeat(times, self.schedule_to_egg(schedule))
+                return bindings.Repeat(bindings.DUMMY_SPAN, times, self.schedule_to_egg(schedule))
             case SequenceDecl(schedules):
-                return bindings.Sequence([self.schedule_to_egg(s) for s in schedules])
+                return bindings.Sequence(bindings.DUMMY_SPAN, [self.schedule_to_egg(s) for s in schedules])
             case RunDecl(ruleset_name, until):
                 self.ruleset_to_egg(ruleset_name)
                 config = bindings.RunConfig(ruleset_name, None if not until else list(map(self.fact_to_egg, until)))
-                return bindings.Run(config)
+                return bindings.Run(bindings.DUMMY_SPAN, config)
             case _:
                 assert_never(schedule)
 
@@ -98,7 +98,8 @@ class EGraphState:
                 for rule in rules:
                     if rule in added_rules:
                         continue
-                    self.egraph.run_program(self.command_to_egg(rule, name))
+                    cmd = self.command_to_egg(rule, name)
+                    self.egraph.run_program(cmd)
                     added_rules.add(rule)
             case CombinedRulesetDecl(rulesets):
                 if name in self.rulesets:
@@ -115,8 +116,9 @@ class EGraphState:
             case RewriteDecl(tp, lhs, rhs, conditions) | BiRewriteDecl(tp, lhs, rhs, conditions):
                 self.type_ref_to_egg(tp)
                 rewrite = bindings.Rewrite(
-                    self.expr_to_egg(lhs),
-                    self.expr_to_egg(rhs),
+                    bindings.DUMMY_SPAN,
+                    self._expr_to_egg(lhs),
+                    self._expr_to_egg(rhs),
                     [self.fact_to_egg(c) for c in conditions],
                 )
                 return (
@@ -126,26 +128,44 @@ class EGraphState:
                 )
             case RuleDecl(head, body, name):
                 rule = bindings.Rule(
+                    bindings.DUMMY_SPAN,
                     [self.action_to_egg(a) for a in head],
                     [self.fact_to_egg(f) for f in body],
                 )
                 return bindings.RuleCommand(name or "", ruleset, rule)
+            # TODO: Replace with just constants value and looking at REF of function
+            case DefaultRewriteDecl(ref, expr):
+                decl = self.__egg_decls__.get_callable_decl(ref).to_function_decl()
+                sig = decl.signature
+                assert isinstance(sig, FunctionSignature)
+                # Replace args with rule_var_name mapping
+                arg_mapping = tuple(
+                    TypedExprDecl(tp.to_just(), VarDecl(name, False))
+                    for name, tp in zip(sig.arg_names, sig.arg_types, strict=False)
+                )
+                rewrite_decl = RewriteDecl(
+                    sig.semantic_return_type.to_just(), CallDecl(ref, arg_mapping), expr, (), False
+                )
+                return self.command_to_egg(rewrite_decl, ruleset)
             case _:
                 assert_never(cmd)
 
     def action_to_egg(self, action: ActionDecl) -> bindings._Action:
         match action:
             case LetDecl(name, typed_expr):
-                return bindings.Let(name, self.typed_expr_to_egg(typed_expr))
+                var_decl = VarDecl(name, True)
+                var_egg = self._expr_to_egg(var_decl)
+                self.expr_to_egg_cache[var_decl] = var_egg
+                return bindings.Let(bindings.DUMMY_SPAN, var_egg.name, self.typed_expr_to_egg(typed_expr))
             case SetDecl(tp, call, rhs):
                 self.type_ref_to_egg(tp)
-                call_ = self.expr_to_egg(call)
-                return bindings.Set(call_.name, call_.args, self.expr_to_egg(rhs))
+                call_ = self._expr_to_egg(call)
+                return bindings.Set(bindings.DUMMY_SPAN, call_.name, call_.args, self._expr_to_egg(rhs))
             case ExprActionDecl(typed_expr):
-                return bindings.Expr_(self.typed_expr_to_egg(typed_expr))
+                return bindings.Expr_(bindings.DUMMY_SPAN, self.typed_expr_to_egg(typed_expr))
             case ChangeDecl(tp, call, change):
                 self.type_ref_to_egg(tp)
-                call_ = self.expr_to_egg(call)
+                call_ = self._expr_to_egg(call)
                 egg_change: bindings._Change
                 match change:
                     case "delete":
@@ -154,12 +174,12 @@ class EGraphState:
                         egg_change = bindings.Subsume()
                     case _:
                         assert_never(change)
-                return bindings.Change(egg_change, call_.name, call_.args)
+                return bindings.Change(bindings.DUMMY_SPAN, egg_change, call_.name, call_.args)
             case UnionDecl(tp, lhs, rhs):
                 self.type_ref_to_egg(tp)
-                return bindings.Union(self.expr_to_egg(lhs), self.expr_to_egg(rhs))
+                return bindings.Union(bindings.DUMMY_SPAN, self._expr_to_egg(lhs), self._expr_to_egg(rhs))
             case PanicDecl(name):
-                return bindings.Panic(name)
+                return bindings.Panic(bindings.DUMMY_SPAN, name)
             case _:
                 assert_never(action)
 
@@ -167,9 +187,9 @@ class EGraphState:
         match fact:
             case EqDecl(tp, exprs):
                 self.type_ref_to_egg(tp)
-                return bindings.Eq([self.expr_to_egg(e) for e in exprs])
+                return bindings.Eq(bindings.DUMMY_SPAN, [self._expr_to_egg(e) for e in exprs])
             case ExprFactDecl(typed_expr):
-                return bindings.Fact(self.typed_expr_to_egg(typed_expr))
+                return bindings.Fact(self.typed_expr_to_egg(typed_expr, False))
             case _:
                 assert_never(fact)
 
@@ -180,7 +200,9 @@ class EGraphState:
         if ref in self.callable_ref_to_egg_fn:
             return self.callable_ref_to_egg_fn[ref]
         decl = self.__egg_decls__.get_callable_decl(ref)
-        self.callable_ref_to_egg_fn[ref] = egg_name = decl.egg_name or _generate_callable_egg_name(ref)
+        self.callable_ref_to_egg_fn[ref] = egg_name = decl.egg_name or _sanitize_egg_ident(
+            self._generate_callable_egg_name(ref)
+        )
         self.egg_fn_to_callable_refs[egg_name].add(ref)
         match decl:
             case RelationDecl(arg_types, _, _):
@@ -201,8 +223,8 @@ class EGraphState:
                             [self.type_ref_to_egg(a.to_just()) for a in signature.arg_types],
                             self.type_ref_to_egg(signature.semantic_return_type.to_just()),
                         ),
-                        self.expr_to_egg(decl.default) if decl.default else None,
-                        self.expr_to_egg(decl.merge) if decl.merge else None,
+                        self._expr_to_egg(decl.default) if decl.default else None,
+                        self._expr_to_egg(decl.merge) if decl.merge else None,
                         [self.action_to_egg(a) for a in decl.on_merge],
                         decl.cost,
                         decl.unextractable,
@@ -228,13 +250,14 @@ class EGraphState:
                     # UnstableFn is a special case, where the rest of args are collected into a call
                     type_args: list[bindings._Expr] = [
                         bindings.Call(
+                            bindings.DUMMY_SPAN,
                             self.type_ref_to_egg(ref.args[1]),
-                            [bindings.Var(self.type_ref_to_egg(a)) for a in ref.args[2:]],
+                            [bindings.Var(bindings.DUMMY_SPAN, self.type_ref_to_egg(a)) for a in ref.args[2:]],
                         ),
-                        bindings.Var(self.type_ref_to_egg(ref.args[0])),
+                        bindings.Var(bindings.DUMMY_SPAN, self.type_ref_to_egg(ref.args[0])),
                     ]
                 else:
-                    type_args = [bindings.Var(self.type_ref_to_egg(a)) for a in ref.args]
+                    type_args = [bindings.Var(bindings.DUMMY_SPAN, self.type_ref_to_egg(a)) for a in ref.args]
                 args = (self.type_ref_to_egg(JustTypeRef(ref.name)), type_args)
             else:
                 args = None
@@ -245,6 +268,8 @@ class EGraphState:
         if decl.builtin:
             for method in decl.class_methods:
                 self.callable_ref_to_egg(ClassMethodRef(ref.name, method))
+            if decl.init:
+                self.callable_ref_to_egg(InitRef(ref.name))
 
         return egg_name
 
@@ -259,31 +284,60 @@ class EGraphState:
             if len(v) == 1
         }
 
-    def typed_expr_to_egg(self, typed_expr_decl: TypedExprDecl) -> bindings._Expr:
+    def typed_expr_to_egg(self, typed_expr_decl: TypedExprDecl, transform_let: bool = True) -> bindings._Expr:
+        # transform all expressions with multiple parents into a let binding, so that less expressions
+        # are sent to egglog. Only for performance reasons.
+        if transform_let:
+            have_multiple_parents = _exprs_multiple_parents(typed_expr_decl)
+            for expr in reversed(have_multiple_parents):
+                self._transform_let(expr)
+
         self.type_ref_to_egg(typed_expr_decl.tp)
-        return self.expr_to_egg(typed_expr_decl.expr)
+        return self._expr_to_egg(typed_expr_decl.expr)
+
+    def _transform_let(self, typed_expr: TypedExprDecl) -> None:
+        """
+        Rewrites this expression as a let binding if it's not already a let binding.
+        """
+        var_decl = VarDecl(f"__expr_{hash(typed_expr)}", True)
+        if var_decl in self.expr_to_egg_cache:
+            return
+        var_egg = self._expr_to_egg(var_decl)
+        cmd = bindings.ActionCommand(
+            bindings.Let(bindings.DUMMY_SPAN, var_egg.name, self.typed_expr_to_egg(typed_expr))
+        )
+        try:
+            self.egraph.run_program(cmd)
+        # errors when creating let bindings for things like `(vec-empty)`
+        except bindings.EggSmolError:
+            return
+        self.expr_to_egg_cache[typed_expr.expr] = var_egg
+        self.expr_to_egg_cache[var_decl] = var_egg
 
     @overload
-    def expr_to_egg(self, expr_decl: CallDecl) -> bindings.Call: ...
+    def _expr_to_egg(self, expr_decl: CallDecl) -> bindings.Call: ...
 
     @overload
-    def expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr: ...
+    def _expr_to_egg(self, expr_decl: VarDecl) -> bindings.Var: ...
 
-    def expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr:
+    @overload
+    def _expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr: ...
+
+    def _expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr:
         """
         Convert an ExprDecl to an egg expression.
-
-        Cached using weakrefs to avoid memory leaks.
         """
         try:
             return self.expr_to_egg_cache[expr_decl]
         except KeyError:
             pass
-
         res: bindings._Expr
         match expr_decl:
-            case VarDecl(name):
-                res = bindings.Var(name)
+            case VarDecl(name, is_let):
+                # prefix let bindings with % to avoid name conflicts with rewrites
+                if is_let:
+                    name = f"%{name}"
+                res = bindings.Var(bindings.DUMMY_SPAN, name)
             case LitDecl(value):
                 l: bindings._Literal
                 match value:
@@ -299,19 +353,22 @@ class EGraphState:
                         l = bindings.String(s)
                     case _:
                         assert_never(value)
-                res = bindings.Lit(l)
+                res = bindings.Lit(bindings.DUMMY_SPAN, l)
             case CallDecl(ref, args, _):
                 egg_fn = self.callable_ref_to_egg(ref)
-                egg_args = [self.typed_expr_to_egg(a) for a in args]
-                res = bindings.Call(egg_fn, egg_args)
+                egg_args = [self.typed_expr_to_egg(a, False) for a in args]
+                res = bindings.Call(bindings.DUMMY_SPAN, egg_fn, egg_args)
             case PyObjectDecl(value):
                 res = GLOBAL_PY_OBJECT_SORT.store(value)
             case PartialCallDecl(call_decl):
-                egg_fn_call = self.expr_to_egg(call_decl)
-                res = bindings.Call("unstable-fn", [bindings.Lit(bindings.String(egg_fn_call.name)), *egg_fn_call.args])
+                egg_fn_call = self._expr_to_egg(call_decl)
+                res = bindings.Call(
+                    bindings.DUMMY_SPAN,
+                    "unstable-fn",
+                    [bindings.Lit(bindings.DUMMY_SPAN, bindings.String(egg_fn_call.name)), *egg_fn_call.args],
+                )
             case _:
                 assert_never(expr_decl.expr)
-
         self.expr_to_egg_cache[expr_decl] = res
         return res
 
@@ -330,6 +387,65 @@ class EGraphState:
         """
         return frozenset(tp for tp in self.type_ref_to_egg_sort if tp.name == cls_name)
 
+    def _generate_callable_egg_name(self, ref: CallableRef) -> str:
+        """
+        Generates a valid egg function name for a callable reference.
+        """
+        match ref:
+            case FunctionRef(name):
+                return name
+
+            case ConstantRef(name):
+                return name
+            case (
+                MethodRef(cls_name, name)
+                | ClassMethodRef(cls_name, name)
+                | ClassVariableRef(cls_name, name)
+                | PropertyRef(cls_name, name)
+            ):
+                return f"{cls_name}.{name}"
+            case InitRef(cls_name):
+                return f"{cls_name}.__init__"
+            case UnnamedFunctionRef(args, val):
+                parts = [str(self._expr_to_egg(a.expr)) + "-" + str(self.type_ref_to_egg(a.tp)) for a in args] + [
+                    str(self.typed_expr_to_egg(val, False))
+                ]
+                return "_".join(parts)
+            case _:
+                assert_never(ref)
+
+
+# https://chatgpt.com/share/9ab899b4-4e17-4426-a3f2-79d67a5ec456
+_EGGLOG_INVALID_IDENT = re.compile(r"[^\w\-+*/?!=<>&|^/%]")
+
+
+def _sanitize_egg_ident(input_string: str) -> str:
+    """
+    Replaces all invalid characters in an egg identifier with an underscore.
+    """
+    return _EGGLOG_INVALID_IDENT.sub("_", input_string)
+
+
+def _exprs_multiple_parents(typed_expr: TypedExprDecl) -> list[TypedExprDecl]:
+    """
+    Returns all expressions that have multiple parents (a list but semantically just an ordered set).
+    """
+    to_traverse = {typed_expr}
+    traversed = set[TypedExprDecl]()
+    traversed_twice = list[TypedExprDecl]()
+    while to_traverse:
+        typed_expr = to_traverse.pop()
+        if typed_expr in traversed:
+            traversed_twice.append(typed_expr)
+            continue
+        traversed.add(typed_expr)
+        expr = typed_expr.expr
+        if isinstance(expr, CallDecl):
+            to_traverse.update(expr.args)
+        elif isinstance(expr, PartialCallDecl):
+            to_traverse.update(expr.call.args)
+    return traversed_twice
+
 
 def _generate_type_egg_name(ref: JustTypeRef) -> str:
     """
@@ -341,24 +457,6 @@ def _generate_type_egg_name(ref: JustTypeRef) -> str:
     return f"{name}_{'_'.join(map(_generate_type_egg_name, ref.args))}"
 
 
-def _generate_callable_egg_name(ref: CallableRef) -> str:
-    """
-    Generates a valid egg function name for a callable reference.
-    """
-    match ref:
-        case FunctionRef(name) | ConstantRef(name):
-            return name
-        case (
-            MethodRef(cls_name, name)
-            | ClassMethodRef(cls_name, name)
-            | ClassVariableRef(cls_name, name)
-            | PropertyRef(cls_name, name)
-        ):
-            return f"{cls_name}_{name}"
-        case _:
-            assert_never(ref)
-
-
 @dataclass
 class FromEggState:
     """
@@ -380,7 +478,7 @@ class FromEggState:
         """
         expr_decl: ExprDecl
         if isinstance(term, bindings.TermVar):
-            expr_decl = VarDecl(term.name)
+            expr_decl = VarDecl(term.name, True)
         elif isinstance(term, bindings.TermLit):
             value = term.value
             expr_decl = LitDecl(None if isinstance(value, bindings.Unit) else value.value)
@@ -427,8 +525,11 @@ class FromEggState:
             possible_types: Iterable[JustTypeRef | None]
             signature = self.decls.get_callable_decl(callable_ref).to_function_decl().signature
             assert isinstance(signature, FunctionSignature)
-            if isinstance(callable_ref, ClassMethodRef):
-                possible_types = self.state._get_possible_types(callable_ref.class_name)
+            if isinstance(callable_ref, ClassMethodRef | InitRef | MethodRef):
+                # Need OR in case we have class method whose class whas never added as a sort, which would happen
+                # if the class method didn't return that type and no other function did. In this case, we don't need
+                # to care about the type vars and we we don't need to bind any possible type.
+                possible_types = self.state._get_possible_types(callable_ref.class_name) or [None]
                 cls_name = callable_ref.class_name
             else:
                 possible_types = [None]
@@ -437,7 +538,6 @@ class FromEggState:
                 tcs = TypeConstraintSolver(self.decls)
                 if possible_type and possible_type.args:
                     tcs.bind_class(possible_type)
-
                 try:
                     arg_types, bound_tp_params = tcs.infer_arg_types(
                         signature.arg_types, signature.semantic_return_type, signature.var_arg_type, tp, cls_name
@@ -445,7 +545,14 @@ class FromEggState:
                 except TypeConstraintError:
                     continue
                 args = tuple(self.resolve_term(a, tp) for a, tp in zip(term.args, arg_types, strict=False))
-                return CallDecl(callable_ref, args, bound_tp_params)
+
+                return CallDecl(
+                    callable_ref,
+                    args,
+                    # Don't include bound type params if this is just a method, we only needed them for type resolution
+                    # but dont need to store them
+                    bound_tp_params if isinstance(callable_ref, ClassMethodRef | InitRef) else None,
+                )
         raise ValueError(f"Could not find callable ref for call {term}")
 
     def resolve_term(self, term_id: int, tp: JustTypeRef) -> TypedExprDecl:

egglog/examples/higher_order_functions.py

@@ -0,0 +1,45 @@
+# mypy: disable-error-code="empty-body"
+"""
+Higher Order Functions
+======================
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from egglog import *
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+class Math(Expr):
+    def __init__(self, i: i64Like) -> None: ...
+    def __add__(self, other: Math) -> Math: ...
+
+
+class MathList(Expr):
+    def __init__(self) -> None: ...
+    def append(self, i: Math) -> MathList: ...
+    def map(self, f: Callable[[Math], Math]) -> MathList: ...
+
+
+@ruleset
+def math_ruleset(i: i64, j: i64, xs: MathList, x: Math, f: Callable[[Math], Math]):  # noqa: ANN201
+    yield rewrite(Math(i) + Math(j)).to(Math(i + j))
+    yield rewrite(xs.append(x).map(f)).to(xs.map(f).append(f(x)))
+    yield rewrite(MathList().map(f)).to(MathList())
+
+
+@function(ruleset=math_ruleset)
+def incr_list(xs: MathList) -> MathList:
+    return xs.map(lambda x: x + Math(1))
+
+
+egraph = EGraph()
+y = egraph.let("y", incr_list(MathList().append(Math(1)).append(Math(2))))
+egraph.run(math_ruleset.saturate())
+egraph.check(eq(y).to(MathList().append(Math(2)).append(Math(3))))
+
+egraph