Fix quadratic let hoisting when converting deep terms

python/egglog/egraph_state.py

@@ -1305,16 +1305,25 @@ def typed_expr_to_egg(
         self,
         typed_expr_decl: TypedExprDecl,
         expr_to_let: bool = True,
+        *,
+        lets_hoisted: bool = False,
     ) -> bindings._Expr:
+        """
+        Convert a typed expression to an egg expression.
+
+        If `lets_hoisted` is set, an enclosing call has already hoisted this expression's shared
+        subexpressions into let bindings, so they are not searched for again. Searching once per
+        subterm would make converting a term quadratic in its depth.
+        """
         # transform all expressions with multiple parents into a let binding, so that less expressions
         # are sent to egglog. Only for performance reasons.
-        if expr_to_let:
+        if expr_to_let and not lets_hoisted:
             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, expr_to_let=expr_to_let)
+        return self._expr_to_egg(typed_expr_decl.expr, expr_to_let=expr_to_let, lets_hoisted=expr_to_let)
 
     def _transform_let(self, typed_expr: TypedExprDecl) -> TypedExprDecl | None:
         """
@@ -1342,17 +1351,27 @@ def _transform_let(self, typed_expr: TypedExprDecl) -> TypedExprDecl | None:
         return None
 
     @overload
-    def _expr_to_egg(self, expr_decl: CallDecl, *, expr_to_let: bool = ...) -> bindings.Call: ...
+    def _expr_to_egg(
+        self, expr_decl: CallDecl, *, expr_to_let: bool = ..., lets_hoisted: bool = ...
+    ) -> bindings.Call: ...
 
     @overload
-    def _expr_to_egg(self, expr_decl: UnboundVarDecl | LetRefDecl, *, expr_to_let: bool = ...) -> bindings.Var: ...
+    def _expr_to_egg(
+        self, expr_decl: UnboundVarDecl | LetRefDecl, *, expr_to_let: bool = ..., lets_hoisted: bool = ...
+    ) -> bindings.Var: ...
 
     @overload
-    def _expr_to_egg(self, expr_decl: ExprDecl, *, expr_to_let: bool = ...) -> bindings._Expr: ...
+    def _expr_to_egg(
+        self, expr_decl: ExprDecl, *, expr_to_let: bool = ..., lets_hoisted: bool = ...
+    ) -> bindings._Expr: ...
 
-    def _expr_to_egg(self, expr_decl: ExprDecl, *, expr_to_let: bool = False) -> bindings._Expr:  # noqa: PLR0912,C901
+    def _expr_to_egg(  # noqa: PLR0912,C901
+        self, expr_decl: ExprDecl, *, expr_to_let: bool = False, lets_hoisted: bool = False
+    ) -> bindings._Expr:
         """
         Convert an ExprDecl to an egg expression.
+
+        `lets_hoisted` is passed on to the arguments, see `typed_expr_to_egg`.
         """
         if expr_to_let:
             try:
@@ -1400,7 +1419,7 @@ def _expr_to_egg(self, expr_decl: ExprDecl, *, expr_to_let: bool = False) -> bin
                 res = bindings.Lit(span(), l)
             case CallDecl() | GetCostDecl():
                 egg_fn, typed_args = self.translate_call(expr_decl)
-                egg_args = [self.typed_expr_to_egg(a, expr_to_let) for a in typed_args]
+                egg_args = [self.typed_expr_to_egg(a, expr_to_let, lets_hoisted=lets_hoisted) for a in typed_args]
                 res = bindings.Call(span(), egg_fn, egg_args)
             case PyObjectDecl(value):
                 res = bindings.Call(
@@ -1415,7 +1434,7 @@ def _expr_to_egg(self, expr_decl: ExprDecl, *, expr_to_let: bool = False) -> bin
                     "unstable-fn",
                     [
                         bindings.Lit(span(), bindings.String(egg_fn)),
-                        *[self.typed_expr_to_egg(arg, expr_to_let) for arg in typed_args],
+                        *[self.typed_expr_to_egg(arg, expr_to_let, lets_hoisted=lets_hoisted) for arg in typed_args],
                     ],
                 )
             case ValueDecl():
@@ -1711,7 +1730,12 @@ def _sanitize_egg_ident(input_string: str) -> str:
 
 
 def _exprs_multiple_parents(typed_expr: TypedExprDecl) -> list[TypedExprDecl]:
-    """Return multiply-parented expressions in deterministic preorder for stable synthetic let names."""
+    """
+    Return multiply-parented expressions in deterministic preorder for stable synthetic let names.
+
+    Visits every child that `_expr_to_egg` converts with `lets_hoisted`, so the arguments of calls,
+    cost lookups and partial calls are not scanned again when they are converted.
+    """
     parent_counts: dict[TypedExprDecl, int] = {}
     traversal_order: list[TypedExprDecl] = []
     traversed: set[TypedExprDecl] = set()
@@ -1724,7 +1748,7 @@ def _exprs_multiple_parents(typed_expr: TypedExprDecl) -> list[TypedExprDecl]:
         if node is not typed_expr:
             traversal_order.append(node)
         match node.expr:
-            case CallDecl(args=args) | PartialCallDecl(CallDecl(args=args)):
+            case CallDecl(args=args) | GetCostDecl(args=args) | PartialCallDecl(CallDecl(args=args)):
                 for child in args:
                     parent_counts[child] = parent_counts.get(child, 0) + 1
                 stack.extend(reversed(args))
@@ -1746,7 +1770,7 @@ def _contains_unbound_var(typed_expr: TypedExprDecl) -> bool:
         match node.expr:
             case UnboundVarDecl():
                 return True
-            case CallDecl(args=args) | PartialCallDecl(CallDecl(args=args)):
+            case CallDecl(args=args) | GetCostDecl(args=args) | PartialCallDecl(CallDecl(args=args)):
                 stack.extend(args)
             case _:
                 pass

python/tests/test_high_level.py

@@ -19,6 +19,7 @@
 
 import egglog.bindings as egg_bindings
 import egglog.builtins as egg_builtins
+import egglog.egraph_state as egraph_state_module
 from egglog import *
 from egglog.declarations import (
     BiRewriteDecl,
@@ -1039,6 +1040,80 @@ def test_shared_subexpression_lowering_is_deterministic_across_processes() -> No
     assert transcripts[1:] == transcripts[:-1]
 
 
+def test_shared_subexpression_scan_runs_once_per_let(monkeypatch: pytest.MonkeyPatch) -> None:
+    """
+    Converting a deep term scans it for shared subexpressions once, plus once per synthetic let
+    body, instead of once per subterm (which made lowering quadratic in the depth of the term).
+    """
+
+    class ScanLeaf(Expr):
+        def __init__(self, name: StringLike) -> None: ...
+
+    class ScanList(Expr):
+        @classmethod
+        def nil(cls) -> ScanList: ...
+
+        @classmethod
+        def cons(cls, left: ScanLeaf, right: ScanLeaf, tail: ScanList) -> ScanList: ...
+
+    original_scan = egraph_state_module._exprs_multiple_parents
+    scans = 0
+
+    def counting_scan(typed_expr: TypedExprDecl) -> list[TypedExprDecl]:
+        nonlocal scans
+        scans += 1
+        return original_scan(typed_expr)
+
+    monkeypatch.setattr(egraph_state_module, "_exprs_multiple_parents", counting_scan)
+
+    def scans_to_let(n: int) -> int:
+        nonlocal scans
+        shared = ScanLeaf("shared")
+        term = ScanList.nil()
+        for i in range(n):
+            term = ScanList.cons(shared, ScanLeaf(f"x{i}"), term)
+        egraph = EGraph(save_egglog_string=True)
+        scans = 0
+        egraph.let("root", term)
+        # Only the shared leaf is hoisted, and every cons cell refers to it.
+        assert egraph.as_egglog_string.count("(let $__expr_") == 1
+        assert egraph.as_egglog_string.count("$__expr_0") == n + 1
+        return scans
+
+    assert scans_to_let(100) == scans_to_let(5)
+
+
+def test_shared_subexpression_inside_cost_lookup_is_hoisted() -> None:
+    """
+    The scan for shared subexpressions also looks inside a cost lookup's arguments, since they are
+    converted without scanning again (see `test_shared_subexpression_scan_runs_once_per_let`).
+    """
+
+    class CostLeaf(Expr):
+        def __init__(self, name: StringLike) -> None: ...
+
+    class CostPair(Expr):
+        def __init__(self, left: CostLeaf, right: CostLeaf) -> None: ...
+
+    class CostRoot(Expr):
+        def __init__(self, pair: CostPair) -> None: ...
+
+    class CostBox(Expr):
+        def __init__(self, cost: i64Like) -> None: ...
+
+    shared = CostLeaf("shared")
+    root = CostRoot(CostPair(shared, shared))
+    egraph = EGraph(save_egglog_string=True)
+    egraph.register(set_cost(root, i64(3)))
+    start = len(egraph.as_egglog_string)
+    egraph.register(CostBox(get_cost(root)))
+    commands = egraph.as_egglog_string[start:]
+    # The shared leaf is hoisted into one synthetic let, which the pair refers to twice.
+    assert commands.count("(let $__expr_") == 1
+    let_name = commands.split("(let ", 1)[1].split(" ", 1)[0]
+    assert commands.count(let_name) == 3
+
+
 def test_freeze_omits_synthetic_let_bindings() -> None:
     class FreezeLetNum(Expr):
         @classmethod