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

egglog/declarations.py

@@ -13,10 +13,11 @@ from typing import TYPE_CHECKING, Literal, Protocol, TypeAlias, Union, runtime_c
 from typing_extensions import Self, assert_never
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Iterable
+    from collections.abc import Callable, Iterable, Mapping
 
 
 __all__ = [
+    "replace_typed_expr",
     "Declarations",
     "DeclerationsLike",
     "DelayedDeclerations",
@@ -29,6 +30,7 @@ __all__ = [
     "MethodRef",
     "ClassMethodRef",
     "FunctionRef",
+    "UnnamedFunctionRef",
     "ConstantRef",
     "ClassVariableRef",
     "PropertyRef",
@@ -71,6 +73,9 @@ __all__ = [
     "CommandDecl",
     "SpecialFunctions",
     "FunctionSignature",
+    "DefaultRewriteDecl",
+    "InitRef",
+    "HasDeclerations",
 ]
 
 
@@ -80,7 +85,14 @@ class DelayedDeclerations:
 
     @property
     def __egg_decls__(self) -> Declarations:
-        return self.__egg_decls_thunk__()
+        thunk = self.__egg_decls_thunk__
+        try:
+            return thunk()
+        # Catch attribute error, so that it isn't bubbled up as a missing attribute and fallbacks on `__getattr__`
+        # instead raise explicitly
+        except AttributeError as err:
+            msg = f"Cannot resolve declerations for {self}"
+            raise RuntimeError(msg) from err
 
 
 @runtime_checkable
@@ -108,11 +120,18 @@ def upcast_declerations(declerations_like: Iterable[DeclerationsLike]) -> list[D
 
 @dataclass
 class Declarations:
+    _unnamed_functions: set[UnnamedFunctionRef] = field(default_factory=set)
     _functions: dict[str, FunctionDecl | RelationDecl] = field(default_factory=dict)
     _constants: dict[str, ConstantDecl] = field(default_factory=dict)
     _classes: dict[str, ClassDecl] = field(default_factory=dict)
     _rulesets: dict[str, RulesetDecl | CombinedRulesetDecl] = field(default_factory=lambda: {"": RulesetDecl([])})
 
+    @property
+    def default_ruleset(self) -> RulesetDecl:
+        ruleset = self._rulesets[""]
+        assert isinstance(ruleset, RulesetDecl)
+        return ruleset
+
     @classmethod
     def create(cls, *others: DeclerationsLike) -> Declarations:
         others = upcast_declerations(others)
@@ -127,7 +146,7 @@ class Declarations:
 
     def copy(self) -> Declarations:
         new = Declarations()
-        new |= self
+        self.update_other(new)
         return new
 
     def update(self, *others: DeclerationsLike) -> None:
@@ -154,9 +173,13 @@ class Declarations:
         other._functions |= self._functions
         other._classes |= self._classes
         other._constants |= self._constants
+        # Must combine rulesets bc the empty ruleset might be different, bc DefaultRewriteDecl
+        # is added to functions.
+        combined_default_rules: set[RewriteOrRuleDecl] = {*self.default_ruleset.rules, *other.default_ruleset.rules}
         other._rulesets |= self._rulesets
+        other._rulesets[""] = RulesetDecl(list(combined_default_rules))
 
-    def get_callable_decl(self, ref: CallableRef) -> CallableDecl:
+    def get_callable_decl(self, ref: CallableRef) -> CallableDecl:  # noqa: PLR0911
         match ref:
             case FunctionRef(name):
                 return self._functions[name]
@@ -170,8 +193,31 @@ class Declarations:
                 return self._classes[class_name].class_methods[name]
             case PropertyRef(class_name, property_name):
                 return self._classes[class_name].properties[property_name]
+            case InitRef(class_name):
+                init_fn = self._classes[class_name].init
+                assert init_fn
+                return init_fn
+            case UnnamedFunctionRef():
+                return ref.to_function_decl()
         assert_never(ref)
 
+    def set_function_decl(
+        self, ref: FunctionRef | MethodRef | ClassMethodRef | PropertyRef | InitRef, decl: FunctionDecl
+    ) -> None:
+        match ref:
+            case FunctionRef(name):
+                self._functions[name] = decl
+            case MethodRef(class_name, method_name):
+                self._classes[class_name].methods[method_name] = decl
+            case ClassMethodRef(class_name, name):
+                self._classes[class_name].class_methods[name] = decl
+            case PropertyRef(class_name, property_name):
+                self._classes[class_name].properties[property_name] = decl
+            case InitRef(class_name):
+                self._classes[class_name].init = decl
+            case _:
+                assert_never(ref)
+
     def has_method(self, class_name: str, method_name: str) -> bool | None:
         """
         Returns whether the given class has the given method, or None if we cant find the class.
@@ -183,12 +229,20 @@ class Declarations:
     def get_class_decl(self, name: str) -> ClassDecl:
         return self._classes[name]
 
+    def get_paramaterized_class(self, name: str) -> TypeRefWithVars:
+        """
+        Returns a class reference with type parameters, if the class is paramaterized.
+        """
+        type_vars = self._classes[name].type_vars
+        return TypeRefWithVars(name, tuple(map(ClassTypeVarRef, type_vars)))
+
 
 @dataclass
 class ClassDecl:
     egg_name: str | None = None
     type_vars: tuple[str, ...] = ()
     builtin: bool = False
+    init: FunctionDecl | None = None
     class_methods: dict[str, FunctionDecl] = field(default_factory=dict)
     # These have to be seperate from class_methods so that printing them can be done easily
     class_variables: dict[str, ConstantDecl] = field(default_factory=dict)
@@ -271,6 +325,37 @@ TypeOrVarRef: TypeAlias = ClassTypeVarRef | TypeRefWithVars
 ##
 
 
+@dataclass(frozen=True)
+class UnnamedFunctionRef:
+    """
+    A reference to a function that doesn't have a name, but does have a body.
+    """
+
+    # tuple of var arg names and their types
+    args: tuple[TypedExprDecl, ...]
+    res: TypedExprDecl
+
+    def to_function_decl(self) -> FunctionDecl:
+        arg_types = []
+        arg_names = []
+        for a in self.args:
+            arg_types.append(a.tp.to_var())
+            assert isinstance(a.expr, VarDecl)
+            arg_names.append(a.expr.name)
+        return FunctionDecl(
+            FunctionSignature(
+                arg_types=tuple(arg_types),
+                arg_names=tuple(arg_names),
+                arg_defaults=(None,) * len(self.args),
+                return_type=self.res.tp.to_var(),
+            ),
+        )
+
+    @property
+    def egg_name(self) -> None | str:
+        return None
+
+
 @dataclass(frozen=True)
 class FunctionRef:
     name: str
@@ -293,6 +378,11 @@ class ClassMethodRef:
     method_name: str
 
 
+@dataclass(frozen=True)
+class InitRef:
+    class_name: str
+
+
 @dataclass(frozen=True)
 class ClassVariableRef:
     class_name: str
@@ -305,7 +395,16 @@ class PropertyRef:
     property_name: str
 
 
-CallableRef: TypeAlias = FunctionRef | ConstantRef | MethodRef | ClassMethodRef | ClassVariableRef | PropertyRef
+CallableRef: TypeAlias = (
+    FunctionRef
+    | ConstantRef
+    | MethodRef
+    | ClassMethodRef
+    | InitRef
+    | ClassVariableRef
+    | PropertyRef
+    | UnnamedFunctionRef
+)
 
 
 ##
@@ -378,7 +477,6 @@ class FunctionSignature:
 @dataclass(frozen=True)
 class FunctionDecl:
     signature: FunctionSignature | SpecialFunctions = field(default_factory=FunctionSignature)
-
     # Egg params
     builtin: bool = False
     egg_name: str | None = None
@@ -402,6 +500,8 @@ CallableDecl: TypeAlias = RelationDecl | ConstantDecl | FunctionDecl
 @dataclass(frozen=True)
 class VarDecl:
     name: str
+    # Differentiate between let bound vars and vars created in rules so that they won't shadow in egglog, by adding a prefix
+    is_let: bool
 
 
 @dataclass(frozen=True)
@@ -458,7 +558,7 @@ class CallDecl:
     bound_tp_params: tuple[JustTypeRef, ...] | None = None
 
     def __post_init__(self) -> None:
-        if self.bound_tp_params and not isinstance(self.callable, ClassMethodRef):
+        if self.bound_tp_params and not isinstance(self.callable, ClassMethodRef | InitRef):
             msg = "Cannot bind type parameters to a non-class method callable."
             raise ValueError(msg)
 
@@ -508,6 +608,38 @@ class TypedExprDecl:
         return l
 
 
+def replace_typed_expr(typed_expr: TypedExprDecl, replacements: Mapping[TypedExprDecl, TypedExprDecl]) -> TypedExprDecl:
+    """
+    Replace all the typed expressions in the given typed expression with the replacements.
+    """
+    # keep track of the traversed expressions for memoization
+    traversed: dict[TypedExprDecl, TypedExprDecl] = {}
+
+    def _inner(typed_expr: TypedExprDecl) -> TypedExprDecl:
+        if typed_expr in traversed:
+            return traversed[typed_expr]
+        if typed_expr in replacements:
+            res = replacements[typed_expr]
+        else:
+            match typed_expr.expr:
+                case (
+                    CallDecl(callable, args, bound_tp_params)
+                    | PartialCallDecl(CallDecl(callable, args, bound_tp_params))
+                ):
+                    new_args = tuple(_inner(a) for a in args)
+                    call_decl = CallDecl(callable, new_args, bound_tp_params)
+                    res = TypedExprDecl(
+                        typed_expr.tp,
+                        call_decl if isinstance(typed_expr.expr, CallDecl) else PartialCallDecl(call_decl),
+                    )
+                case _:
+                    res = typed_expr
+        traversed[typed_expr] = res
+        return res
+
+    return _inner(typed_expr)
+
+
 ##
 # Schedules
 ##
@@ -629,7 +761,13 @@ class RuleDecl:
     name: str | None
 
 
-RewriteOrRuleDecl: TypeAlias = RewriteDecl | BiRewriteDecl | RuleDecl
+@dataclass(frozen=True)
+class DefaultRewriteDecl:
+    ref: CallableRef
+    expr: ExprDecl
+
+
+RewriteOrRuleDecl: TypeAlias = RewriteDecl | BiRewriteDecl | RuleDecl | DefaultRewriteDecl
 
 
 @dataclass(frozen=True)