Add custom cost model and ability to get costs

Author: saulshanabrook

docs/reference/python-integration.md

@@ -635,3 +635,50 @@ r = ruleset(
 )
 egraph.saturate(r)
 ```
+
+## Custom Cost Models
+
+Custom cost models are also supported by subclassing `CostModel[T]` or `DefaultCostModel` and passing in an instance as the `cost_model` kwargs to `EGraph.extract`. The `CostModel` is paramterized by a cost type `T`, which in the `DefaultCostModel` is `int`. Any cost must be able to be compared to choose the lowest cost.
+
+The `Expr`s passed to your cost model represent partial program trees. Any builtin values (containers or single primitives) will be fully evaluated, but any values that return user defined classes will be last as opaque "values",
+representing an e-class in the e-graph. The only thing you can do with values is to compare them to each other or
+use them in `EGraph.lookup_function_value` to lookup the resulting value of a call with values in it.
+
+For example, here is a cost model that uses boolean values to determine if a model is extractable or not:
+
+```{code-cell} python
+
+class MyExpr(Expr):
+    def __init__(self) -> None: ...
+
+
+class BooleanCostModel(CostModel[bool]):
+    cost_tp = bool
+
+    def primitive_cost(self, egraph: EGraph, value: Primitive) -> bool:
+        # Only allow extracting even integers
+        match value:
+            case i64(i) if i % 2 == 0:
+                return True
+        return False
+
+    def container_cost(self, egraph: EGraph, container: Expr, children_costs: list[bool]) -> bool:
+        # Only allow extracting Vecs of extractable values
+        match container:
+            case Vec():
+                return all(children_costs)
+        return False
+
+    def call_cost(self, egraph: EGraph, expr: Expr) -> bool:
+        # Only allow extracting calls to `my_f`
+        match expr:
+            case my_f():
+                return True
+        return False
+
+    def fold(self, callable: ExprCallable, children_costs: list[bool], head_cost: bool) -> bool:
+        # Only allow extracting calls where the head and all children are extractable
+        return head_cost and all(children_costs)
+
+assert EGraph().extract(i64(10), include_cost=True, cost_model=BooleanCostModel()) == (i64(10), True)
+```

pyproject.toml

@@ -55,6 +55,7 @@ test = [
     "pytest-codspeed",
     "pytest-benchmark",
     "pytest-xdist",
+    "pytest-mock",
 ]
 
 docs = [

python/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 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,33 @@ 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: ...
+    def __add__(self, other: _Cost) -> _Cost: ...
+
+_COST = TypeVar("_COST", bound=_Cost)
+
+@final
+class CostModel(Generic[_COST]):
+    def __init__(
+        self,
+        fold: Callable[[str, _COST, list[_COST]], _COST] | None,
+        enode_cost: Callable[[str, list[Value]], _COST] | None,
+        container_cost: Callable[[str, Value, list[_COST]], _COST] | None,
+        base_value_cost: Callable[[str, Value], _COST] | None,
+    ) -> None: ...
+
+@final
+class Extractor(Generic[_COST]):
+    def __init__(self, rootsorts: list[str] | None, egraph: EGraph, cost_model: CostModel[_COST]) -> 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]]: ...

python/egglog/builtins.py

@@ -33,10 +33,12 @@
     "BigRatLike",
     "Bool",
     "BoolLike",
+    "Container",
     "ExprValueError",
     "Map",
     "MapLike",
     "MultiSet",
+    "Primitive",
     "PyObject",
     "Rational",
     "Set",
@@ -57,6 +59,9 @@
     "py_exec",
 ]
 
+Container: TypeAlias = "Map | Set | MultiSet | Vec | UnstableFn"
+Primitive: TypeAlias = "String | Bool | i64 | f64 | Rational | BigInt | BigRat | PyObject | Unit"
+
 
 @dataclass
 class ExprValueError(AttributeError):

python/egglog/declarations.py

@@ -14,6 +14,8 @@
 
 from typing_extensions import Self, assert_never
 
+from .bindings import Value
+
 if TYPE_CHECKING:
     from collections.abc import Callable, Iterable, Mapping
 
@@ -49,6 +51,7 @@
     "FunctionDecl",
     "FunctionRef",
     "FunctionSignature",
+    "GetCostDecl",
     "HasDeclerations",
     "InitRef",
     "JustTypeRef",
@@ -82,6 +85,7 @@
     "UnboundVarDecl",
     "UnionDecl",
     "UnnamedFunctionRef",
+    "ValueDecl",
     "collect_unbound_vars",
     "replace_typed_expr",
     "upcast_declerations",
@@ -696,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)

python/egglog/deconstruct.py

@@ -11,7 +11,7 @@
 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,