egglog 7.0.0__cp312-none-win_amd64.whl → 7.2.0__cp312-none-win_amd64.whl

egglog/egraph.py

@@ -75,6 +75,7 @@ __all__ = [
     "seq",
     "Command",
     "simplify",
+    "unstable_combine_rulesets",
     "check",
     "GraphvizKwargs",
     "Ruleset",
@@ -88,6 +89,7 @@ __all__ = [
     "Fact",
     "Action",
     "Command",
+    "check_eq",
 ]
 
 T = TypeVar("T")
@@ -145,6 +147,23 @@ def simplify(x: EXPR, schedule: Schedule | None = None) -> EXPR:
     return EGraph().extract(x)
 
 
+def check_eq(x: EXPR, y: EXPR, schedule: Schedule | None = None) -> EGraph:
+    """
+    Verifies that two expressions are equal after running the schedule.
+    """
+    egraph = EGraph()
+    x_var = egraph.let("__check_eq_x", x)
+    y_var = egraph.let("__check_eq_y", y)
+    if schedule:
+        egraph.run(schedule)
+    fact = eq(x_var).to(y_var)
+    try:
+        egraph.check(fact)
+    except bindings.EggSmolError as err:
+        raise AssertionError(f"Failed {eq(x).to(y)}\n -> {ne(egraph.extract(x)).to(egraph.extract(y))})") from err
+    return egraph
+
+
 def check(x: FactLike, schedule: Schedule | None = None, *given: ActionLike) -> None:
     """
     Verifies that the fact is true given some assumptions and after running the schedule.
@@ -333,7 +352,7 @@ class _BaseModule:
         This is the same as defining a nullary function with a high cost.
         # TODO: Rename as declare to match eggglog?
         """
-        return constant(name, tp, egg_name)
+        return constant(name, tp, egg_name=egg_name)
 
     def register(
         self,
@@ -433,6 +452,7 @@ class _ExprMetaclass(type):
         namespace: dict[str, Any],
         egg_sort: str | None = None,
         builtin: bool = False,
+        ruleset: Ruleset | None = None,
     ) -> RuntimeClass | type:
         # If this is the Expr subclass, just return the class
         if not bases:
@@ -448,7 +468,14 @@ class _ExprMetaclass(type):
         # Otherwise, f_locals returns a copy
         # https://peps.python.org/pep-0667/
         decls_thunk = Thunk.fn(
-            _generate_class_decls, namespace, prev_frame, builtin, egg_sort, name, fallback=Declarations
+            _generate_class_decls,
+            namespace,
+            prev_frame,
+            builtin,
+            egg_sort,
+            name,
+            ruleset,
+            fallback=Declarations,
         )
         return RuntimeClass(decls_thunk, TypeRefWithVars(name))
 
@@ -457,7 +484,12 @@ class _ExprMetaclass(type):
 
 
 def _generate_class_decls(
-    namespace: dict[str, Any], frame: FrameType, builtin: bool, egg_sort: str | None, cls_name: str
+    namespace: dict[str, Any],
+    frame: FrameType,
+    builtin: bool,
+    egg_sort: str | None,
+    cls_name: str,
+    ruleset: Ruleset | None,
 ) -> Declarations:
     """
     Lazy constructor for class declerations to support classes with methods whose types are not yet defined.
@@ -479,9 +511,9 @@ def _generate_class_decls(
     for k, v in get_type_hints(_Dummytype, globalns=frame.f_globals, localns=frame.f_locals).items():
         if getattr(v, "__origin__", None) == ClassVar:
             (inner_tp,) = v.__args__
-            type_ref = resolve_type_annotation(decls, inner_tp).to_just()
-            cls_decl.class_variables[k] = ConstantDecl(type_ref)
-
+            type_ref = resolve_type_annotation(decls, inner_tp)
+            cls_decl.class_variables[k] = ConstantDecl(type_ref.to_just())
+            _add_default_rewrite(decls, ClassVariableRef(cls_name, k), type_ref, namespace.pop(k, None), ruleset)
         else:
             msg = f"On class {cls_name}, for attribute '{k}', expected a ClassVar, but got {v}"
             raise NotImplementedError(msg)
@@ -491,13 +523,15 @@ def _generate_class_decls(
     ##
 
     # The type ref of self is paramterized by the type vars
-    slf_type_ref = TypeRefWithVars(cls_name, tuple(map(ClassTypeVarRef, type_vars)))
+    TypeRefWithVars(cls_name, tuple(map(ClassTypeVarRef, type_vars)))
 
     # Get all the methods from the class
     filtered_namespace: list[tuple[str, Any]] = [
         (k, v) for k, v in namespace.items() if k not in IGNORED_ATTRIBUTES or isinstance(v, _WrappedMethod)
     ]
 
+    # all methods we should try adding default functions for
+    default_function_refs: dict[ClassMethodRef | MethodRef | PropertyRef, Callable] = {}
     # Then register each of its methods
     for method_name, method in filtered_namespace:
         is_init = method_name == "__init__"
@@ -517,33 +551,26 @@ def _generate_class_decls(
             continue
         locals = frame.f_locals
 
-        def create_decl(fn: object, first: Literal["cls"] | TypeRefWithVars) -> FunctionDecl:
-            return _fn_decl(
-                decls,
-                egg_fn,  # noqa: B023
-                fn,
-                locals,  # noqa: B023
-                default,  # noqa: B023
-                cost,  # noqa: B023
-                merge,  # noqa: B023
-                on_merge,  # noqa: B023
-                mutates,  # noqa: B023
-                builtin,
-                first,
-                is_init,  # noqa: B023
-                unextractable,  # noqa: B023
-            )
-
+        ref: ClassMethodRef | MethodRef | PropertyRef | InitRef
         match fn:
             case classmethod():
-                cls_decl.class_methods[method_name] = create_decl(fn.__func__, "cls")
+                ref = ClassMethodRef(cls_name, method_name)
+                fn = fn.__func__
             case property():
-                cls_decl.properties[method_name] = create_decl(fn.fget, slf_type_ref)
+                ref = PropertyRef(cls_name, method_name)
+                fn = fn.fget
             case _:
-                if is_init:
-                    cls_decl.class_methods[method_name] = create_decl(fn, slf_type_ref)
-                else:
-                    cls_decl.methods[method_name] = create_decl(fn, slf_type_ref)
+                ref = InitRef(cls_name) if is_init else MethodRef(cls_name, method_name)
+
+        _fn_decl(decls, egg_fn, ref, fn, locals, default, cost, merge, on_merge, mutates, builtin, unextractable)
+
+        if not builtin and not isinstance(ref, InitRef) and not mutates:
+            default_function_refs[ref] = fn
+
+    # Add all rewrite methods at the end so that all methods are registered first and can be accessed
+    # in the bodies
+    for ref, fn in default_function_refs.items():
+        _add_default_rewrite_function(decls, ref, fn, ruleset)
 
     return decls
 
@@ -562,6 +589,7 @@ def function(
     mutates_first_arg: bool = False,
     unextractable: bool = False,
     builtin: bool = False,
+    ruleset: Ruleset | None = None,
 ) -> Callable[[CALLABLE], CALLABLE]: ...
 
 
@@ -575,6 +603,7 @@ def function(
     on_merge: Callable[[EXPR, EXPR], Iterable[ActionLike]] | None = None,
     mutates_first_arg: bool = False,
     unextractable: bool = False,
+    ruleset: Ruleset | None = None,
 ) -> Callable[[Callable[P, EXPR]], Callable[P, EXPR]]: ...
 
 
@@ -605,15 +634,17 @@ class _FunctionConstructor:
     merge: Callable[[RuntimeExpr, RuntimeExpr], RuntimeExpr] | None = None
     on_merge: Callable[[RuntimeExpr, RuntimeExpr], Iterable[ActionLike]] | None = None
     unextractable: bool = False
+    ruleset: Ruleset | None = None
 
     def __call__(self, fn: Callable[..., RuntimeExpr]) -> RuntimeFunction:
         return RuntimeFunction(Thunk.fn(self.create_decls, fn), FunctionRef(fn.__name__))
 
     def create_decls(self, fn: Callable[..., RuntimeExpr]) -> Declarations:
         decls = Declarations()
-        decls._functions[fn.__name__] = _fn_decl(
+        _fn_decl(
             decls,
             self.egg_fn,
+            (ref := FunctionRef(fn.__name__)),
             fn,
             self.hint_locals,
             self.default,
@@ -624,12 +655,14 @@ class _FunctionConstructor:
             self.builtin,
             unextractable=self.unextractable,
         )
+        _add_default_rewrite_function(decls, ref, fn, self.ruleset)
         return decls
 
 
 def _fn_decl(
     decls: Declarations,
     egg_name: str | None,
+    ref: FunctionRef | MethodRef | PropertyRef | ClassMethodRef | InitRef,
     fn: object,
     # Pass in the locals, retrieved from the frame when wrapping,
     # so that we support classes and function defined inside of other functions (which won't show up in the globals)
@@ -640,22 +673,35 @@ def _fn_decl(
     on_merge: Callable[[RuntimeExpr, RuntimeExpr], Iterable[ActionLike]] | None,
     mutates_first_arg: bool,
     is_builtin: bool,
-    # The first arg is either cls, for a classmethod, a self type, or none for a function
-    first_arg: Literal["cls"] | TypeOrVarRef | None = None,
-    is_init: bool = False,
     unextractable: bool = False,
-) -> FunctionDecl:
+) -> None:
+    """
+    Sets the function decl for the function object.
+    """
     if not isinstance(fn, FunctionType):
         raise NotImplementedError(f"Can only generate function decls for functions not {fn}  {type(fn)}")
 
+    # partial function creation and calling are handled with a special case in the type checker, so don't
+    # use the normal logic
+    special_function_name: SpecialFunctions | None = (
+        "fn-partial" if egg_name == "unstable-fn" else "fn-app" if egg_name == "unstable-app" else None
+    )
+    if special_function_name:
+        decls.set_function_decl(ref, FunctionDecl(special_function_name, builtin=True, egg_name=egg_name))
+        return
+
     hint_globals = fn.__globals__.copy()
+    # Copy Callable into global if not present bc sometimes it gets automatically removed by ruff to type only block
+    # https://docs.astral.sh/ruff/rules/typing-only-standard-library-import/
+    if "Callable" not in hint_globals:
+        hint_globals["Callable"] = Callable
 
     hints = get_type_hints(fn, hint_globals, hint_locals)
 
     params = list(signature(fn).parameters.values())
 
-    # If this is an init function, or a classmethod, remove the first arg name
-    if is_init or first_arg == "cls":
+    # If this is an init function, or a classmethod, the first arg is not used
+    if isinstance(ref, ClassMethodRef | InitRef):
         params = params[1:]
 
     if _last_param_variable(params):
@@ -665,9 +711,8 @@ def _fn_decl(
     else:
         var_arg_type = None
     arg_types = tuple(
-        first_arg
-        # If the first arg is a self, and this not an __init__ fn, add this as a typeref
-        if i == 0 and isinstance(first_arg, ClassTypeVarRef | TypeRefWithVars) and not is_init
+        decls.get_paramaterized_class(ref.class_name)
+        if i == 0 and isinstance(ref, MethodRef | PropertyRef)
         else resolve_type_annotation(decls, hints[t.name])
         for i, t in enumerate(params)
     )
@@ -680,17 +725,15 @@ def _fn_decl(
 
     decls.update(*arg_defaults)
 
-    # If this is an init fn use the first arg as the return type
-    if is_init:
-        assert not mutates_first_arg
-        if not isinstance(first_arg, ClassTypeVarRef | TypeRefWithVars):
-            msg = "Init function must have a self type"
-            raise ValueError(msg)
-        return_type = first_arg
-    elif mutates_first_arg:
-        return_type = arg_types[0]
-    else:
-        return_type = resolve_type_annotation(decls, hints["return"])
+    return_type = (
+        decls.get_paramaterized_class(ref.class_name)
+        if isinstance(ref, InitRef)
+        else arg_types[0]
+        if mutates_first_arg
+        else resolve_type_annotation(decls, hints["return"])
+    )
+
+    arg_names = tuple(t.name for t in params)
 
     decls |= default
     merged = (
@@ -714,12 +757,14 @@ def _fn_decl(
         )
     )
     decls.update(*merge_action)
-    return FunctionDecl(
-        return_type=None if mutates_first_arg else return_type,
-        var_arg_type=var_arg_type,
-        arg_types=arg_types,
-        arg_names=tuple(t.name for t in params),
-        arg_defaults=tuple(a.__egg_typed_expr__.expr if a is not None else None for a in arg_defaults),
+    decl = FunctionDecl(
+        FunctionSignature(
+            return_type=None if mutates_first_arg else return_type,
+            var_arg_type=var_arg_type,
+            arg_types=arg_types,
+            arg_names=arg_names,
+            arg_defaults=tuple(a.__egg_typed_expr__.expr if a is not None else None for a in arg_defaults),
+        ),
         cost=cost,
         egg_name=egg_name,
         merge=merged.__egg_typed_expr__.expr if merged is not None else None,
@@ -728,6 +773,7 @@ def _fn_decl(
         default=None if default is None else default.__egg_typed_expr__.expr,
         on_merge=tuple(a.action for a in merge_action),
     )
+    decls.set_function_decl(ref, decl)
 
 
 # Overload to support aritys 0-4 until variadic generic support map, so we can map from type to value
@@ -769,19 +815,73 @@ def _relation_decls(name: str, tps: tuple[type, ...], egg_fn: str | None) -> Dec
     return decls
 
 
-def constant(name: str, tp: type[EXPR], egg_name: str | None = None) -> EXPR:
+def constant(
+    name: str,
+    tp: type[EXPR],
+    default_replacement: EXPR | None = None,
+    /,
+    *,
+    egg_name: str | None = None,
+    ruleset: Ruleset | None = None,
+) -> EXPR:
     """
     A "constant" is implemented as the instantiation of a value that takes no args.
     This creates a function with `name` and return type `tp` and returns a value of it being called.
     """
-    return cast(EXPR, RuntimeExpr(Thunk.fn(_constant_thunk, name, tp, egg_name)))
+    return cast(EXPR, RuntimeExpr(Thunk.fn(_constant_thunk, name, tp, egg_name, default_replacement, ruleset)))
 
 
-def _constant_thunk(name: str, tp: type, egg_name: str | None) -> tuple[Declarations, TypedExprDecl]:
+def _constant_thunk(
+    name: str, tp: type, egg_name: str | None, default_replacement: object, ruleset: Ruleset | None
+) -> tuple[Declarations, TypedExprDecl]:
     decls = Declarations()
-    type_ref = resolve_type_annotation(decls, tp).to_just()
-    decls._constants[name] = ConstantDecl(type_ref, egg_name)
-    return decls, TypedExprDecl(type_ref, CallDecl(ConstantRef(name)))
+    type_ref = resolve_type_annotation(decls, tp)
+    callable_ref = ConstantRef(name)
+    decls._constants[name] = ConstantDecl(type_ref.to_just(), egg_name)
+    _add_default_rewrite(decls, callable_ref, type_ref, default_replacement, ruleset)
+    return decls, TypedExprDecl(type_ref.to_just(), CallDecl(callable_ref))
+
+
+def _add_default_rewrite_function(decls: Declarations, ref: CallableRef, fn: Callable, ruleset: Ruleset | None) -> None:
+    """
+    Adds a default rewrite for a function, by calling the functions with vars and adding it if it is not None.
+    """
+    callable_decl = decls.get_callable_decl(ref)
+    assert isinstance(callable_decl, FunctionDecl)
+    signature = callable_decl.signature
+    assert isinstance(signature, FunctionSignature)
+
+    var_args: list[object] = [
+        RuntimeExpr.__from_value__(decls, TypedExprDecl(tp.to_just(), VarDecl(_rule_var_name(name))))
+        for name, tp in zip(signature.arg_names, signature.arg_types, strict=False)
+    ]
+    # If this is a classmethod, add the class as the first arg
+    if isinstance(ref, ClassMethodRef):
+        tp = decls.get_paramaterized_class(ref.class_name)
+        var_args.insert(0, RuntimeClass(Thunk.value(decls), tp))
+    _add_default_rewrite(decls, ref, signature.semantic_return_type, fn(*var_args), ruleset)
+
+
+def _add_default_rewrite(
+    decls: Declarations, ref: CallableRef, type_ref: TypeOrVarRef, default_rewrite: object, ruleset: Ruleset | None
+) -> None:
+    """
+    Adds a default rewrite for the callable, if the default rewrite is not None
+
+    Will add it to the ruleset if it is passed in, or add it to the default ruleset on the passed in decls if not.
+    """
+    if default_rewrite is None:
+        return
+    resolved_value = resolve_literal(type_ref, default_rewrite)
+    rewrite_decl = DefaultRewriteDecl(ref, resolved_value.__egg_typed_expr__.expr)
+    if ruleset:
+        ruleset_decls = ruleset._current_egg_decls
+        ruleset_decl = ruleset.__egg_ruleset__
+    else:
+        ruleset_decls = decls
+        ruleset_decl = decls.default_ruleset
+    ruleset_decl.rules.append(rewrite_decl)
+    ruleset_decls |= resolved_value
 
 
 def _last_param_variable(params: list[Parameter]) -> bool:
@@ -933,13 +1033,12 @@ class EGraph(_BaseModule):
         """
         Displays the e-graph in the notebook.
         """
-        graphviz = self.graphviz(**kwargs)
         if IN_IPYTHON:
             from IPython.display import SVG, display
 
             display(SVG(self.graphviz_svg(**kwargs)))
         else:
-            graphviz.render(view=True, format="svg", quiet=True)
+            self.graphviz(**kwargs).render(view=True, format="svg", quiet=True)
 
     def input(self, fn: Callable[..., String], path: str) -> None:
         """
@@ -957,6 +1056,7 @@ class EGraph(_BaseModule):
         action = let(name, expr)
         self.register(action)
         runtime_expr = to_runtime_expr(expr)
+        self._add_decls(runtime_expr)
         return cast(
             EXPR,
             RuntimeExpr.__from_value__(
@@ -982,7 +1082,8 @@ class EGraph(_BaseModule):
         self._add_decls(runtime_expr, schedule)
         egg_schedule = self._state.schedule_to_egg(schedule.schedule)
         typed_expr = runtime_expr.__egg_typed_expr__
-        egg_expr = self._state.expr_to_egg(typed_expr.expr)
+        # Must also register type
+        egg_expr = self._state.typed_expr_to_egg(typed_expr)
         self._egraph.run_program(bindings.Simplify(egg_expr, egg_schedule))
         extract_report = self._egraph.extract_report()
         if not isinstance(extract_report, bindings.Best):
@@ -1059,7 +1160,7 @@ class EGraph(_BaseModule):
         runtime_expr = to_runtime_expr(expr)
         self._add_decls(runtime_expr)
         typed_expr = runtime_expr.__egg_typed_expr__
-        extract_report = self._run_extract(typed_expr.expr, 0)
+        extract_report = self._run_extract(typed_expr, 0)
 
         if not isinstance(extract_report, bindings.Best):
             msg = "No extract report saved"
@@ -1079,15 +1180,15 @@ class EGraph(_BaseModule):
         self._add_decls(runtime_expr)
         typed_expr = runtime_expr.__egg_typed_expr__
 
-        extract_report = self._run_extract(typed_expr.expr, n)
+        extract_report = self._run_extract(typed_expr, n)
         if not isinstance(extract_report, bindings.Variants):
             msg = "Wrong extract report type"
             raise ValueError(msg)  # noqa: TRY004
         new_exprs = self._state.exprs_from_egg(extract_report.termdag, extract_report.terms, typed_expr.tp)
         return [cast(EXPR, RuntimeExpr.__from_value__(self.__egg_decls__, expr)) for expr in new_exprs]
 
-    def _run_extract(self, expr: ExprDecl, n: int) -> bindings._ExtractReport:
-        expr = self._state.expr_to_egg(expr)
+    def _run_extract(self, typed_expr: TypedExprDecl, n: int) -> bindings._ExtractReport:
+        expr = self._state.typed_expr_to_egg(typed_expr)
         self._egraph.run_program(bindings.ActionCommand(bindings.Extract(expr, bindings.Lit(bindings.Int(n)))))
         extract_report = self._egraph.extract_report()
         if not extract_report:
@@ -1146,7 +1247,7 @@ class EGraph(_BaseModule):
         runtime_expr = to_runtime_expr(expr)
         self._add_decls(runtime_expr)
         typed_expr = runtime_expr.__egg_typed_expr__
-        egg_expr = self._state.expr_to_egg(typed_expr.expr)
+        egg_expr = self._state.typed_expr_to_egg(typed_expr)
         match typed_expr.tp:
             case JustTypeRef("i64"):
                 return self._egraph.eval_i64(egg_expr)
@@ -1276,8 +1377,10 @@ def ruleset(
     """
     Creates a ruleset with the following rules.
 
-    If no name is provided, one is generated based on the current module
+    If no name is provided, try using the name of the funciton.
     """
+    if isinstance(rule_or_generator, FunctionType):
+        name = name or rule_or_generator.__name__
     r = Ruleset(name)
     if rule_or_generator is not None:
         r.register(rule_or_generator, *rules, _increase_frame=True)
@@ -1388,12 +1491,48 @@ class Ruleset(Schedule):
     def __repr__(self) -> str:
         return str(self)
 
+    def __or__(self, other: Ruleset | UnstableCombinedRuleset) -> UnstableCombinedRuleset:
+        return unstable_combine_rulesets(self, other)
+
     # Create a unique name if we didn't pass one from the user
     @property
     def __egg_name__(self) -> str:
         return self.name or f"ruleset_{id(self)}"
 
 
+@dataclass
+class UnstableCombinedRuleset(Schedule):
+    __egg_decls_thunk__: Callable[[], Declarations] = field(init=False)
+    schedule: RunDecl = field(init=False)
+    name: str | None
+    rulesets: InitVar[list[Ruleset | UnstableCombinedRuleset]]
+
+    def __post_init__(self, rulesets: list[Ruleset | UnstableCombinedRuleset]) -> None:
+        self.schedule = RunDecl(self.__egg_name__, ())
+        self.__egg_decls_thunk__ = Thunk.fn(self._create_egg_decls, *rulesets)
+
+    @property
+    def __egg_name__(self) -> str:
+        return self.name or f"combined_ruleset_{id(self)}"
+
+    def _create_egg_decls(self, *rulesets: Ruleset | UnstableCombinedRuleset) -> Declarations:
+        decls = Declarations.create(*rulesets)
+        decls._rulesets[self.__egg_name__] = CombinedRulesetDecl(tuple(r.__egg_name__ for r in rulesets))
+        return decls
+
+    def __or__(self, other: Ruleset | UnstableCombinedRuleset) -> UnstableCombinedRuleset:
+        return unstable_combine_rulesets(self, other)
+
+
+def unstable_combine_rulesets(
+    *rulesets: Ruleset | UnstableCombinedRuleset, name: str | None = None
+) -> UnstableCombinedRuleset:
+    """
+    Combine multiple rulesets into a single ruleset.
+    """
+    return UnstableCombinedRuleset(name, list(rulesets))
+
+
 @dataclass
 class RewriteOrRule:
     __egg_decls__: Declarations
@@ -1556,9 +1695,9 @@ def var(name: str, bound: type[EXPR]) -> EXPR:
 
 def _var(name: str, bound: object) -> RuntimeExpr:
     """Create a new variable with the given name and type."""
-    if not isinstance(bound, RuntimeClass):
-        raise TypeError(f"Unexpected type {type(bound)}")
-    return RuntimeExpr.__from_value__(bound.__egg_decls__, TypedExprDecl(bound.__egg_tp__.to_just(), VarDecl(name)))
+    decls = Declarations()
+    type_ref = resolve_type_annotation(decls, bound)
+    return RuntimeExpr.__from_value__(decls, TypedExprDecl(type_ref.to_just(), VarDecl(name)))
 
 
 def vars_(names: str, bound: type[EXPR]) -> Iterable[EXPR]:
@@ -1801,9 +1940,11 @@ def _rewrite_or_rule_generator(gen: RewriteOrRuleGenerator, frame: FrameType) ->
     """
     # Get the local scope from where the function is defined, so that we can get any type hints that are in the scope
     # but not in the globals
-
-    hints = get_type_hints(gen, gen.__globals__, frame.f_locals)
-    args = [_var(p.name, hints[p.name]) for p in signature(gen).parameters.values()]
+    globals = gen.__globals__.copy()
+    if "Callable" not in globals:
+        globals["Callable"] = Callable
+    hints = get_type_hints(gen, globals, frame.f_locals)
+    args = [_var(_rule_var_name(p.name), hints[p.name]) for p in signature(gen).parameters.values()]
     return list(gen(*args))  # type: ignore[misc]