egglog 11.2.0__cp310-cp310-manylinux_2_17_ppc64.manylinux2014_ppc64.whl → 11.4.0__cp310-cp310-manylinux_2_17_ppc64.manylinux2014_ppc64.whl

egglog/bindings.pyi

@@ -1,6 +1,8 @@
+from collections.abc import Callable
 from datetime import timedelta
+from fractions import Fraction
 from pathlib import Path
-from typing import TypeAlias
+from typing import Any, Generic, Protocol, TypeAlias, TypeVar
 
 from typing_extensions import final
 
@@ -14,6 +16,7 @@ __all__ = [
     "Change",
     "Check",
     "Constructor",
+    "CostModel",
     "Datatype",
     "Datatypes",
     "DefaultPrintFunctionMode",
@@ -26,6 +29,7 @@ __all__ = [
     "Extract",
     "ExtractBest",
     "ExtractVariants",
+    "Extractor",
     "Fact",
     "Fail",
     "Float",
@@ -83,6 +87,7 @@ __all__ = [
     "UserDefined",
     "UserDefinedCommandOutput",
     "UserDefinedOutput",
+    "Value",
     "Var",
     "Variant",
 ]
@@ -128,6 +133,31 @@ class EGraph:
         max_calls_per_function: int | None = None,
         include_temporary_functions: bool = False,
     ) -> SerializedEGraph: ...
+    def lookup_function(self, name: str, key: list[Value]) -> Value | None: ...
+    def eval_expr(self, expr: _Expr) -> tuple[str, Value]: ...
+    def value_to_i64(self, v: Value) -> int: ...
+    def value_to_f64(self, v: Value) -> float: ...
+    def value_to_string(self, v: Value) -> str: ...
+    def value_to_bool(self, v: Value) -> bool: ...
+    def value_to_rational(self, v: Value) -> Fraction: ...
+    def value_to_bigint(self, v: Value) -> int: ...
+    def value_to_bigrat(self, v: Value) -> Fraction: ...
+    def value_to_pyobject(self, py_object_sort: PyObjectSort, v: Value) -> object: ...
+    def value_to_map(self, v: Value) -> dict[Value, Value]: ...
+    def value_to_multiset(self, v: Value) -> list[Value]: ...
+    def value_to_vec(self, v: Value) -> list[Value]: ...
+    def value_to_function(self, v: Value) -> tuple[str, list[Value]]: ...
+    def value_to_set(self, v: Value) -> set[Value]: ...
+    # def dynamic_cost_model_enode_cost(self, func: str, args: list[Value]) -> int: ...
+
+@final
+class Value:
+    def __hash__(self) -> int: ...
+    def __eq__(self, value: object) -> bool: ...
+    def __lt__(self, other: object) -> bool: ...
+    def __le__(self, other: object) -> bool: ...
+    def __gt__(self, other: object) -> bool: ...
+    def __ge__(self, other: object) -> bool: ...
 
 @final
 class EggSmolError(Exception):
@@ -732,3 +762,34 @@ class TermDag:
     def expr_to_term(self, expr: _Expr) -> _Term: ...
     def term_to_expr(self, term: _Term, span: _Span) -> _Expr: ...
     def to_string(self, term: _Term) -> str: ...
+
+##
+# Extraction
+##
+class _Cost(Protocol):
+    def __lt__(self, other: _Cost) -> bool: ...
+    def __le__(self, other: _Cost) -> bool: ...
+    def __gt__(self, other: _Cost) -> bool: ...
+    def __ge__(self, other: _Cost) -> bool: ...
+
+_COST = TypeVar("_COST", bound=_Cost)
+
+_ENODE_COST = TypeVar("_ENODE_COST")
+
+@final
+class CostModel(Generic[_COST, _ENODE_COST]):
+    def __init__(
+        self,
+        fold: Callable[[str, _ENODE_COST, list[_COST]], _COST],
+        enode_cost: Callable[[str, list[Value]], _ENODE_COST],
+        container_cost: Callable[[str, Value, list[_COST]], _COST],
+        base_value_cost: Callable[[str, Value], _COST],
+    ) -> None: ...
+
+@final
+class Extractor(Generic[_COST]):
+    def __init__(self, rootsorts: list[str] | None, egraph: EGraph, cost_model: CostModel[_COST, Any]) -> None: ...
+    def extract_best(self, egraph: EGraph, termdag: TermDag, value: Value, sort: str) -> tuple[_COST, _Term]: ...
+    def extract_variants(
+        self, egraph: EGraph, termdag: TermDag, value: Value, nvariants: int, sort: str
+    ) -> list[tuple[_COST, _Term]]: ...

egglog/builtins.py

@@ -33,10 +33,12 @@ __all__ = [
     "BigRatLike",
     "Bool",
     "BoolLike",
+    "Container",
     "ExprValueError",
     "Map",
     "MapLike",
     "MultiSet",
+    "Primitive",
     "PyObject",
     "Rational",
     "Set",
@@ -103,6 +105,10 @@ class String(BuiltinExpr):
     @method(egg_fn="replace")
     def replace(self, old: StringLike, new: StringLike) -> String: ...
 
+    @method(preserve=True)
+    def __add__(self, other: StringLike) -> String:
+        return join(self, other)
+
 
 StringLike: TypeAlias = String | str
 
@@ -1131,3 +1137,7 @@ def _convert_function(fn: FunctionType) -> UnstableFn:
 
 
 converter(FunctionType, UnstableFn, _convert_function)
+
+
+Container: TypeAlias = Map | Set | MultiSet | Vec | UnstableFn
+Primitive: TypeAlias = String | Bool | i64 | f64 | Rational | BigInt | BigRat | PyObject | Unit

egglog/declarations.py

@@ -9,10 +9,13 @@ from __future__ import annotations
 from dataclasses import dataclass, field
 from functools import cached_property
 from typing import TYPE_CHECKING, ClassVar, Literal, Protocol, TypeAlias, TypeVar, Union, cast, runtime_checkable
+from uuid import UUID
 from weakref import WeakValueDictionary
 
 from typing_extensions import Self, assert_never
 
+from .bindings import Value
+
 if TYPE_CHECKING:
     from collections.abc import Callable, Iterable, Mapping
 
@@ -20,6 +23,7 @@ if TYPE_CHECKING:
 __all__ = [
     "ActionCommandDecl",
     "ActionDecl",
+    "BackOffDecl",
     "BiRewriteDecl",
     "CallDecl",
     "CallableDecl",
@@ -47,11 +51,13 @@ __all__ = [
     "FunctionDecl",
     "FunctionRef",
     "FunctionSignature",
+    "GetCostDecl",
     "HasDeclerations",
     "InitRef",
     "JustTypeRef",
     "LetDecl",
     "LetRefDecl",
+    "LetSchedulerDecl",
     "LitDecl",
     "LitType",
     "MethodRef",
@@ -79,6 +85,7 @@ __all__ = [
     "UnboundVarDecl",
     "UnionDecl",
     "UnnamedFunctionRef",
+    "ValueDecl",
     "collect_unbound_vars",
     "replace_typed_expr",
     "upcast_declerations",
@@ -636,7 +643,7 @@ class CallDecl:
     args: tuple[TypedExprDecl, ...] = ()
     # type parameters that were bound to the callable, if it is a classmethod
     # Used for pretty printing classmethod calls with type parameters
-    bound_tp_params: tuple[JustTypeRef, ...] | None = None
+    bound_tp_params: tuple[JustTypeRef, ...] = ()
 
     # pool objects for faster __eq__
     _args_to_value: ClassVar[WeakValueDictionary[tuple[object, ...], CallDecl]] = WeakValueDictionary({})
@@ -651,7 +658,7 @@ class CallDecl:
         # normalize the args/kwargs to a tuple so that they can be compared
         callable = args[0] if args else kwargs["callable"]
         args_ = args[1] if len(args) > 1 else kwargs.get("args", ())
-        bound_tp_params = args[2] if len(args) > 2 else kwargs.get("bound_tp_params")
+        bound_tp_params = args[2] if len(args) > 2 else kwargs.get("bound_tp_params", ())
 
         normalized_args = (callable, args_, bound_tp_params)
         try:
@@ -693,7 +700,20 @@ class PartialCallDecl:
     call: CallDecl
 
 
-ExprDecl: TypeAlias = UnboundVarDecl | LetRefDecl | LitDecl | CallDecl | PyObjectDecl | PartialCallDecl
+@dataclass(frozen=True)
+class GetCostDecl:
+    callable: CallableRef
+    args: tuple[TypedExprDecl, ...]
+
+
+@dataclass(frozen=True)
+class ValueDecl:
+    value: Value
+
+
+ExprDecl: TypeAlias = (
+    UnboundVarDecl | LetRefDecl | LitDecl | CallDecl | PyObjectDecl | PartialCallDecl | ValueDecl | GetCostDecl
+)
 
 
 @dataclass(frozen=True)
@@ -790,9 +810,24 @@ class SequenceDecl:
 class RunDecl:
     ruleset: str
     until: tuple[FactDecl, ...] | None
+    scheduler: BackOffDecl | None = None
+
+
+@dataclass(frozen=True)
+class LetSchedulerDecl:
+    scheduler: BackOffDecl
+    inner: ScheduleDecl
 
 
-ScheduleDecl: TypeAlias = SaturateDecl | RepeatDecl | SequenceDecl | RunDecl
+ScheduleDecl: TypeAlias = SaturateDecl | RepeatDecl | SequenceDecl | RunDecl | LetSchedulerDecl
+
+
+@dataclass(frozen=True)
+class BackOffDecl:
+    id: UUID
+    match_limit: int | None
+    ban_length: int | None
+
 
 ##
 # Facts

egglog/deconstruct.py

@@ -11,7 +11,7 @@ from typing import TYPE_CHECKING, TypeVar, overload
 from typing_extensions import TypeVarTuple, Unpack
 
 from .declarations import *
-from .egraph import BaseExpr
+from .egraph import BaseExpr, Expr
 from .runtime import *
 from .thunk import *
 
@@ -49,7 +49,11 @@ def get_literal_value(x: PyObject) -> object: ...
 def get_literal_value(x: UnstableFn[T, Unpack[TS]]) -> Callable[[Unpack[TS]], T] | None: ...
 
 
-def get_literal_value(x: String | Bool | i64 | f64 | PyObject | UnstableFn) -> object:
+@overload
+def get_literal_value(x: Expr) -> None: ...
+
+
+def get_literal_value(x: object) -> object:
     """
     Returns the literal value of an expression if it is a literal.
     If it is not a literal, returns None.
@@ -95,12 +99,9 @@ def get_var_name(x: BaseExpr) -> str | None:
     return None
 
 
-def get_callable_fn(x: T) -> Callable[..., T] | None:
+def get_callable_fn(x: T) -> Callable[..., T] | T | None:
     """
-    Gets the function of an expression if it is a call expression.
-    If it is not a call expression (a property, a primitive value, constants, classvars, a let value), return None.
-    For those values, you can check them by comparing them directly with equality or for primitives calling `.eval()`
-    to return the Python value.
+    Gets the function of an expression, or if it's a constant or classvar, return that.
     """
     if not isinstance(x, RuntimeExpr):
         raise TypeError(f"Expected Expression, got {type(x).__name__}")
@@ -159,6 +160,7 @@ def _deconstruct_call_decl(
     """
     args = call.args
     arg_exprs = tuple(RuntimeExpr(decls_thunk, Thunk.value(a)) for a in args)
+    # TODO: handle values? Like constants
     if isinstance(call.callable, InitRef):
         return RuntimeClass(
             decls_thunk,