egglog 7.2.0__cp311-none-win_amd64.whl → 8.0.1__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
 
@@ -19,7 +19,7 @@ from .type_constraint_solver import TypeConstraintError, TypeConstraintSolver
 if TYPE_CHECKING:
     from collections.abc import Iterable
 
-__all__ = ["EGraphState", "GLOBAL_PY_OBJECT_SORT", "_rule_var_name"]
+__all__ = ["EGraphState", "GLOBAL_PY_OBJECT_SORT"]
 
 # Create a global sort for python objects, so we can store them without an e-graph instance
 # Needed when serializing commands to egg commands when creating modules
@@ -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)
 
@@ -116,6 +116,7 @@ class EGraphState:
             case RewriteDecl(tp, lhs, rhs, conditions) | BiRewriteDecl(tp, lhs, rhs, conditions):
                 self.type_ref_to_egg(tp)
                 rewrite = bindings.Rewrite(
+                    bindings.DUMMY_SPAN,
                     self._expr_to_egg(lhs),
                     self._expr_to_egg(rhs),
                     [self.fact_to_egg(c) for c in conditions],
@@ -127,19 +128,24 @@ 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)
-                args = tuple(
-                    TypedExprDecl(tp.to_just(), VarDecl(_rule_var_name(name)))
+                # 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, args), expr, (), 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)
@@ -147,13 +153,16 @@ class EGraphState:
     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))
+                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)
@@ -165,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)
 
@@ -178,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)
 
@@ -191,22 +200,31 @@ 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, _, _):
-                self.egraph.run_program(bindings.Relation(egg_name, [self.type_ref_to_egg(a) for a in arg_types]))
+                self.egraph.run_program(
+                    bindings.Relation(bindings.DUMMY_SPAN, egg_name, [self.type_ref_to_egg(a) for a in arg_types])
+                )
             case ConstantDecl(tp, _):
                 # Use function decleration instead of constant b/c constants cannot be extracted
                 # https://github.com/egraphs-good/egglog/issues/334
                 self.egraph.run_program(
-                    bindings.Function(bindings.FunctionDecl(egg_name, bindings.Schema([], self.type_ref_to_egg(tp))))
+                    bindings.Function(
+                        bindings.FunctionDecl(
+                            bindings.DUMMY_SPAN, egg_name, bindings.Schema([], self.type_ref_to_egg(tp))
+                        )
+                    )
                 )
             case FunctionDecl():
                 if not decl.builtin:
                     signature = decl.signature
                     assert isinstance(signature, FunctionSignature), "Cannot turn special function to egg"
                     egg_fn_decl = bindings.FunctionDecl(
+                        bindings.DUMMY_SPAN,
                         egg_name,
                         bindings.Schema(
                             [self.type_ref_to_egg(a.to_just()) for a in signature.arg_types],
@@ -239,17 +257,18 @@ 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
-            self.egraph.run_program(bindings.Sort(egg_name, args))
+            self.egraph.run_program(bindings.Sort(bindings.DUMMY_SPAN, egg_name, args))
         # For builtin classes, let's also make sure we have the mapping of all egg fn names for class methods, because
         # these can be created even without adding them to the e-graph, like `vec-empty` which can be extracted
         # even if you never use that function.
@@ -272,31 +291,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)
 
+    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: ...
 
+    @overload
+    def _expr_to_egg(self, expr_decl: VarDecl) -> bindings.Var: ...
+
     @overload
     def _expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr: ...
 
-    def _expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr:
+    def _expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr:  # noqa: PLR0912,C901
         """
         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:
@@ -312,19 +360,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])
+                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
 
@@ -343,6 +394,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:
     """
@@ -354,26 +464,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 InitRef(cls_name):
-            return f"{cls_name}___init__"
-        case _:
-            assert_never(ref)
-
-
 @dataclass
 class FromEggState:
     """
@@ -395,7 +485,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)
@@ -403,7 +493,7 @@ class FromEggState:
             if term.name == "py-object":
                 call = bindings.termdag_term_to_expr(self.termdag, term)
                 expr_decl = PyObjectDecl(self.state.egraph.eval_py_object(call))
-            if term.name == "unstable-fn":
+            elif term.name == "unstable-fn":
                 # Get function name
                 fn_term, *arg_terms = term.args
                 fn_value = self.resolve_term(fn_term, JustTypeRef("String"))
@@ -478,10 +568,3 @@ class FromEggState:
         except KeyError:
             res = self.cache[term_id] = self.from_expr(tp, self.termdag.nodes[term_id])
             return res
-
-
-def _rule_var_name(s: str) -> str:
-    """
-    Create a hidden variable name, for rewrites, so that let bindings or function won't conflict with it
-    """
-    return f"__var__{s}"

egglog/examples/higher_order_functions.py

@@ -16,35 +16,30 @@ if TYPE_CHECKING:
 
 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
+def math_ruleset(i: i64, j: i64, xs: MathList, x: Math, f: Callable[[Math], Math]):
     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 increment_by_one(x: Math) -> Math:
-    return x + Math(1)
+def incr_list(xs: MathList) -> MathList:
+    return xs.map(lambda x: x + Math(1))
 
 
 egraph = EGraph()
-x = egraph.let("x", MathList().append(Math(1)).append(Math(2)))
-y = egraph.let("y", x.map(increment_by_one))
+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