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

egglog/bindings.pyi

@@ -494,6 +494,12 @@ class Relation:
 class PrintOverallStatistics:
     def __init__(self) -> None: ...
 
+@final
+class UnstableCombinedRuleset:
+    name: str
+    rulesets: list[str]
+    def __init__(self, name: str, rulesets: list[str]) -> None: ...
+
 _Command: TypeAlias = (
     SetOption
     | Datatype
@@ -521,6 +527,7 @@ _Command: TypeAlias = (
     | CheckProof
     | Relation
     | PrintOverallStatistics
+    | UnstableCombinedRuleset
 )
 
 def termdag_term_to_expr(termdag: TermDag, term: _Term) -> _Expr: ...

egglog/builtins.py

@@ -5,10 +5,14 @@ Builtin sorts and function to egg.
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Generic, Protocol, TypeAlias, TypeVar, Union
+from functools import partial
+from typing import TYPE_CHECKING, Generic, Protocol, TypeAlias, TypeVar, Union, overload
+
+from typing_extensions import TypeVarTuple, Unpack
 
 from .conversion import converter
 from .egraph import Expr, Unit, function, method
+from .runtime import RuntimeFunction
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -31,6 +35,7 @@ __all__ = [
     "py_eval",
     "py_exec",
     "py_eval_fn",
+    "UnstableFn",
 ]
 
 
@@ -461,3 +466,38 @@ def py_exec(code: StringLike, globals: object = PyObject.dict(), locals: object
     """
     Copies the locals, execs the Python code, and returns the locals with any updates.
     """
+
+
+TS = TypeVarTuple("TS")
+
+T1 = TypeVar("T1")
+T2 = TypeVar("T2")
+T3 = TypeVar("T3")
+
+
+class UnstableFn(Expr, Generic[T, Unpack[TS]], builtin=True):
+    @overload
+    def __init__(self, f: Callable[[Unpack[TS]], T]) -> None: ...
+
+    @overload
+    def __init__(self, f: Callable[[T1, Unpack[TS]], T], _a: T1, /) -> None: ...
+
+    @overload
+    def __init__(self, f: Callable[[T1, T2, Unpack[TS]], T], _a: T1, _b: T2, /) -> None: ...
+
+    # Removing due to bug in MyPy
+    # https://github.com/python/mypy/issues/17212
+    # @overload
+    # def __init__(self, f: Callable[[T1, T2, T3, Unpack[TS]], T], _a: T1, _b: T2, _c: T3, /) -> None: ...
+
+    # etc, for partial application
+
+    @method(egg_fn="unstable-fn")
+    def __init__(self, f, *partial) -> None: ...
+
+    @method(egg_fn="unstable-app")
+    def __call__(self, *args: Unpack[TS]) -> T: ...
+
+
+converter(RuntimeFunction, UnstableFn, UnstableFn)
+converter(partial, UnstableFn, lambda p: UnstableFn(p.func, *p.args))

egglog/conversion.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, TypeVar, cast
+from typing import TYPE_CHECKING, NewType, TypeVar, cast
 
 from .declarations import *
 from .pretty import *
@@ -16,7 +16,8 @@ if TYPE_CHECKING:
 
 __all__ = ["convert", "converter", "resolve_literal", "convert_to_same_type"]
 # Mapping from (source type, target type) to and function which takes in the runtimes values of the source and return the target
-CONVERSIONS: dict[tuple[type | JustTypeRef, JustTypeRef], tuple[int, Callable]] = {}
+TypeName = NewType("TypeName", str)
+CONVERSIONS: dict[tuple[type | TypeName, TypeName], tuple[int, Callable]] = {}
 # Global declerations to store all convertable types so we can query if they have certain methods or not
 # Defer it as a thunk so we can register conversions without triggering type signature loading
 CONVERSIONS_DECLS: Callable[[], Declarations] = Thunk.value(Declarations())
@@ -34,12 +35,12 @@ def converter(from_type: type[T], to_type: type[V], fn: Callable[[T], V], cost:
     Register a converter from some type to an egglog type.
     """
     to_type_name = process_tp(to_type)
-    if not isinstance(to_type_name, JustTypeRef):
+    if not isinstance(to_type_name, str):
         raise TypeError(f"Expected return type to be a egglog type, got {to_type_name}")
     _register_converter(process_tp(from_type), to_type_name, fn, cost)
 
 
-def _register_converter(a: type | JustTypeRef, b: JustTypeRef, a_b: Callable, cost: int) -> None:
+def _register_converter(a: type | TypeName, b: TypeName, a_b: Callable, cost: int) -> None:
     """
     Registers a converter from some type to an egglog type, if not already registered.
 
@@ -94,14 +95,17 @@ def convert_to_same_type(source: object, target: RuntimeExpr) -> RuntimeExpr:
     return resolve_literal(tp.to_var(), source)
 
 
-def process_tp(tp: type | RuntimeClass) -> JustTypeRef | type:
+def process_tp(tp: type | RuntimeClass) -> TypeName | type:
     """
     Process a type before converting it, to add it to the global declerations and resolve to a ref.
     """
     global CONVERSIONS_DECLS
     if isinstance(tp, RuntimeClass):
         CONVERSIONS_DECLS = Thunk.fn(_combine_decls, CONVERSIONS_DECLS, tp)
-        return tp.__egg_tp__.to_just()
+        egg_tp = tp.__egg_tp__
+        if egg_tp.args:
+            raise TypeError(f"Cannot register a converter for a generic type, got {tp}")
+        return TypeName(egg_tp.name)
     return tp
 
 
@@ -109,7 +113,7 @@ def _combine_decls(d: Callable[[], Declarations], x: HasDeclerations) -> Declara
     return Declarations.create(d(), x)
 
 
-def min_convertable_tp(a: object, b: object, name: str) -> JustTypeRef:
+def min_convertable_tp(a: object, b: object, name: str) -> TypeName:
     """
     Returns the minimum convertable type between a and b, that has a method `name`, raising a ConvertError if no such type exists.
     """
@@ -117,14 +121,14 @@ def min_convertable_tp(a: object, b: object, name: str) -> JustTypeRef:
     a_tp = _get_tp(a)
     b_tp = _get_tp(b)
     a_converts_to = {
-        to: c for ((from_, to), (c, _)) in CONVERSIONS.items() if from_ == a_tp and decls.has_method(to.name, name)
+        to: c for ((from_, to), (c, _)) in CONVERSIONS.items() if from_ == a_tp and decls.has_method(to, name)
     }
     b_converts_to = {
-        to: c for ((from_, to), (c, _)) in CONVERSIONS.items() if from_ == b_tp and decls.has_method(to.name, name)
+        to: c for ((from_, to), (c, _)) in CONVERSIONS.items() if from_ == b_tp and decls.has_method(to, name)
     }
-    if isinstance(a_tp, JustTypeRef):
+    if isinstance(a_tp, str):
         a_converts_to[a_tp] = 0
-    if isinstance(b_tp, JustTypeRef):
+    if isinstance(b_tp, str):
         b_converts_to[b_tp] = 0
     common = set(a_converts_to) & set(b_converts_to)
     if not common:
@@ -143,28 +147,29 @@ def resolve_literal(tp: TypeOrVarRef, arg: object) -> RuntimeExpr:
     try:
         tp_just = tp.to_just()
     except NotImplementedError:
-        # If this is a var, it has to be a runtime exprssions
+        # If this is a var, it has to be a runtime expession
         assert isinstance(arg, RuntimeExpr)
         return arg
-    if arg_type == tp_just:
+    tp_name = TypeName(tp_just.name)
+    if arg_type == tp_name:
         # If the type is an egg type, it has to be a runtime expr
         assert isinstance(arg, RuntimeExpr)
         return arg
     # Try all parent types as well, if we are converting from a Python type
     for arg_type_instance in arg_type.__mro__ if isinstance(arg_type, type) else [arg_type]:
         try:
-            fn = CONVERSIONS[(cast(JustTypeRef | type, arg_type_instance), tp_just)][1]
+            fn = CONVERSIONS[(cast(TypeName | type, arg_type_instance), tp_name)][1]
         except KeyError:
             continue
         break
     else:
-        raise ConvertError(f"Cannot convert {arg_type} to {tp_just}")
+        raise ConvertError(f"Cannot convert {arg_type} to {tp_name}")
     return fn(arg)
 
 
-def _get_tp(x: object) -> JustTypeRef | type:
+def _get_tp(x: object) -> TypeName | type:
     if isinstance(x, RuntimeExpr):
-        return x.__egg_typed_expr__.tp
+        return TypeName(x.__egg_typed_expr__.tp.name)
     tp = type(x)
     # If this value has a custom metaclass, let's use that as our index instead of the type
     if type(tp) != type:

egglog/declarations.py

@@ -39,6 +39,7 @@ __all__ = [
     "CallableDecl",
     "VarDecl",
     "PyObjectDecl",
+    "PartialCallDecl",
     "LitType",
     "LitDecl",
     "CallDecl",
@@ -46,6 +47,7 @@ __all__ = [
     "TypedExprDecl",
     "ClassDecl",
     "RulesetDecl",
+    "CombinedRulesetDecl",
     "SaturateDecl",
     "RepeatDecl",
     "SequenceDecl",
@@ -67,6 +69,10 @@ __all__ = [
     "RewriteOrRuleDecl",
     "ActionCommandDecl",
     "CommandDecl",
+    "SpecialFunctions",
+    "FunctionSignature",
+    "DefaultRewriteDecl",
+    "InitRef",
 ]
 
 
@@ -76,7 +82,13 @@ class DelayedDeclerations:
 
     @property
     def __egg_decls__(self) -> Declarations:
-        return self.__egg_decls_thunk__()
+        try:
+            return self.__egg_decls_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 = "Error resolving declerations"
+            raise RuntimeError(msg) from err
 
 
 @runtime_checkable
@@ -88,9 +100,6 @@ class HasDeclerations(Protocol):
 DeclerationsLike: TypeAlias = Union[HasDeclerations, None, "Declarations"]
 
 
-# TODO: Make all ClassDecls take deferred type refs, which return new decls when resolving.
-
-
 def upcast_declerations(declerations_like: Iterable[DeclerationsLike]) -> list[Declarations]:
     d = []
     for l in declerations_like:
@@ -110,7 +119,13 @@ class Declarations:
     _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] = field(default_factory=lambda: {"": RulesetDecl([])})
+    _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:
@@ -126,7 +141,7 @@ class Declarations:
 
     def copy(self) -> Declarations:
         new = Declarations()
-        new |= self
+        self.update_other(new)
         return new
 
     def update(self, *others: DeclerationsLike) -> None:
@@ -153,9 +168,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]
@@ -169,8 +188,29 @@ 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
         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.
@@ -182,12 +222,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)
@@ -196,7 +244,7 @@ class ClassDecl:
     preserved_methods: dict[str, Callable] = field(default_factory=dict)
 
 
-@dataclass
+@dataclass(frozen=True)
 class RulesetDecl:
     rules: list[RewriteOrRuleDecl]
 
@@ -206,6 +254,11 @@ class RulesetDecl:
         return hash((type(self), tuple(self.rules)))
 
 
+@dataclass(frozen=True)
+class CombinedRulesetDecl:
+    rulesets: tuple[str, ...]
+
+
 # Have two different types of type refs, one that can include vars recursively and one that cannot.
 # We only use the one with vars for classmethods and methods, and the other one for egg references as
 # well as runtime values.
@@ -287,6 +340,11 @@ class ClassMethodRef:
     method_name: str
 
 
+@dataclass(frozen=True)
+class InitRef:
+    class_name: str
+
+
 @dataclass(frozen=True)
 class ClassVariableRef:
     class_name: str
@@ -299,7 +357,9 @@ class PropertyRef:
     property_name: str
 
 
-CallableRef: TypeAlias = FunctionRef | ConstantRef | MethodRef | ClassMethodRef | ClassVariableRef | PropertyRef
+CallableRef: TypeAlias = (
+    FunctionRef | ConstantRef | MethodRef | ClassMethodRef | InitRef | ClassVariableRef | PropertyRef
+)
 
 
 ##
@@ -316,10 +376,12 @@ class RelationDecl:
 
     def to_function_decl(self) -> FunctionDecl:
         return FunctionDecl(
-            arg_types=tuple(a.to_var() for a in self.arg_types),
-            arg_names=tuple(f"__{i}" for i in range(len(self.arg_types))),
-            arg_defaults=self.arg_defaults,
-            return_type=TypeRefWithVars("Unit"),
+            FunctionSignature(
+                arg_types=tuple(a.to_var() for a in self.arg_types),
+                arg_names=tuple(f"__{i}" for i in range(len(self.arg_types))),
+                arg_defaults=self.arg_defaults,
+                return_type=TypeRefWithVars("Unit"),
+            ),
             egg_name=self.egg_name,
             default=LitDecl(None),
         )
@@ -336,25 +398,40 @@ class ConstantDecl:
 
     def to_function_decl(self) -> FunctionDecl:
         return FunctionDecl(
-            arg_types=(),
-            arg_names=(),
-            arg_defaults=(),
-            return_type=self.type_ref.to_var(),
+            FunctionSignature(return_type=self.type_ref.to_var()),
             egg_name=self.egg_name,
         )
 
 
+# special cases for partial function creation and application, which cannot use the normal python rules
+SpecialFunctions: TypeAlias = Literal["fn-partial", "fn-app"]
+
+
 @dataclass(frozen=True)
-class FunctionDecl:
-    # All args are delayed except for relations converted to function decls
-    arg_types: tuple[TypeOrVarRef, ...]
-    arg_names: tuple[str, ...]
+class FunctionSignature:
+    arg_types: tuple[TypeOrVarRef, ...] = ()
+    arg_names: tuple[str, ...] = ()
     # List of defaults. None for any arg which doesn't have one.
-    arg_defaults: tuple[ExprDecl | None, ...]
+    arg_defaults: tuple[ExprDecl | None, ...] = ()
     # If None, then the first arg is mutated and returned
-    return_type: TypeOrVarRef | None
+    return_type: TypeOrVarRef | None = None
     var_arg_type: TypeOrVarRef | None = None
 
+    @property
+    def semantic_return_type(self) -> TypeOrVarRef:
+        """
+        The type that is returned by the function, which wil be in the first arg if it mutates it.
+        """
+        return self.return_type or self.arg_types[0]
+
+    @property
+    def mutates(self) -> bool:
+        return self.return_type is None
+
+
+@dataclass(frozen=True)
+class FunctionDecl:
+    signature: FunctionSignature | SpecialFunctions = field(default_factory=FunctionSignature)
     # Egg params
     builtin: bool = False
     egg_name: str | None = None
@@ -367,17 +444,6 @@ class FunctionDecl:
     def to_function_decl(self) -> FunctionDecl:
         return self
 
-    @property
-    def semantic_return_type(self) -> TypeOrVarRef:
-        """
-        The type that is returned by the function, which wil be in the first arg if it mutates it.
-        """
-        return self.return_type or self.arg_types[0]
-
-    @property
-    def mutates(self) -> bool:
-        return self.return_type is None
-
 
 CallableDecl: TypeAlias = RelationDecl | ConstantDecl | FunctionDecl
 
@@ -445,7 +511,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)
 
@@ -463,7 +529,20 @@ class CallDecl:
         return hash(self) == hash(other)
 
 
-ExprDecl: TypeAlias = VarDecl | LitDecl | CallDecl | PyObjectDecl
+@dataclass(frozen=True)
+class PartialCallDecl:
+    """
+    A partially applied function aka a function sort.
+
+    Note it does not need to have any args, in which case it's just a function pointer.
+
+    Seperated from the call decl so it's clear it is translated to a `unstable-fn` call.
+    """
+
+    call: CallDecl
+
+
+ExprDecl: TypeAlias = VarDecl | LitDecl | CallDecl | PyObjectDecl | PartialCallDecl
 
 
 @dataclass(frozen=True)
@@ -603,7 +682,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)