egglog 11.1.0__cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl → 11.2.0__cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

egglog/bindings.pyi

@@ -7,15 +7,16 @@ from typing_extensions import final
 __all__ = [
     "ActionCommand",
     "AddRuleset",
-    "Best",
     "BiRewriteCommand",
     "Bool",
+    "CSVPrintFunctionMode",
     "Call",
     "Change",
     "Check",
     "Constructor",
     "Datatype",
     "Datatypes",
+    "DefaultPrintFunctionMode",
     "Delete",
     "EGraph",
     "EggSmolError",
@@ -23,10 +24,13 @@ __all__ = [
     "Eq",
     "Expr_",
     "Extract",
+    "ExtractBest",
+    "ExtractVariants",
     "Fact",
     "Fail",
     "Float",
     "Function",
+    "FunctionCommand",
     "IdentSort",
     "Include",
     "Input",
@@ -35,10 +39,14 @@ __all__ = [
     "Lit",
     "NewSort",
     "Output",
+    "OverallStatistics",
     "Panic",
     "PanicSpan",
     "Pop",
+    "PrintAllFunctionsSize",
     "PrintFunction",
+    "PrintFunctionOutput",
+    "PrintFunctionSize",
     "PrintOverallStatistics",
     "PrintSize",
     "Push",
@@ -53,13 +61,13 @@ __all__ = [
     "RunConfig",
     "RunReport",
     "RunSchedule",
+    "RunScheduleOutput",
     "RustSpan",
     "Saturate",
     "Schema",
     "Sequence",
     "SerializedEGraph",
     "Set",
-    "SetOption",
     "Sort",
     "SrcFile",
     "String",
@@ -73,13 +81,18 @@ __all__ = [
     "Unit",
     "UnstableCombinedRuleset",
     "UserDefined",
+    "UserDefinedCommandOutput",
+    "UserDefinedOutput",
     "Var",
     "Variant",
-    "Variants",
 ]
 
 @final
 class SerializedEGraph:
+    @property
+    def truncated_functions(self) -> list[str]: ...
+    @property
+    def discarded_functions(self) -> list[str]: ...
     def inline_leaves(self) -> None: ...
     def saturate_inline_leaves(self) -> None: ...
     def to_dot(self) -> str: ...
@@ -106,9 +119,7 @@ class EGraph:
     ) -> None: ...
     def parse_program(self, __input: str, /, filename: str | None = None) -> list[_Command]: ...
     def commands(self) -> str | None: ...
-    def run_program(self, *commands: _Command) -> list[str]: ...
-    def extract_report(self) -> _ExtractReport | None: ...
-    def run_report(self) -> RunReport | None: ...
+    def run_program(self, *commands: _Command) -> list[_CommandOutput]: ...
     def serialize(
         self,
         root_eclasses: list[_Expr],
@@ -356,6 +367,13 @@ class IdentSort:
     sort: str
     def __init__(self, ident: str, sort: str) -> None: ...
 
+@final
+class UserDefinedCommandOutput: ...
+
+@final
+class Function:
+    name: str
+
 @final
 class RunReport:
     updated: bool
@@ -375,20 +393,80 @@ class RunReport:
         rebuild_time_per_ruleset: dict[str, timedelta],
     ) -> None: ...
 
+##
+# Command Outputs
+##
+
+@final
+class PrintFunctionSize:
+    size: int
+    def __init__(self, size: int) -> None: ...
+
+@final
+class PrintAllFunctionsSize:
+    sizes: list[tuple[str, int]]
+    def __init__(self, sizes: list[tuple[str, int]]) -> None: ...
+
 @final
-class Variants:
+class ExtractVariants:
     termdag: TermDag
     terms: list[_Term]
     def __init__(self, termdag: TermDag, terms: list[_Term]) -> None: ...
 
 @final
-class Best:
+class ExtractBest:
     termdag: TermDag
     cost: int
     term: _Term
     def __init__(self, termdag: TermDag, cost: int, term: _Term) -> None: ...
 
-_ExtractReport: TypeAlias = Variants | Best
+@final
+class OverallStatistics:
+    report: RunReport
+    def __init__(self, report: RunReport) -> None: ...
+
+@final
+class RunScheduleOutput:
+    report: RunReport
+    def __init__(self, report: RunReport) -> None: ...
+
+@final
+class PrintFunctionOutput:
+    function: Function
+    termdag: TermDag
+    terms: list[tuple[_Term, _Term]]
+    mode: _PrintFunctionMode
+    def __init__(
+        self, function: Function, termdag: TermDag, terms: list[tuple[_Term, _Term]], mode: _PrintFunctionMode
+    ) -> None: ...
+
+@final
+class UserDefinedOutput:
+    output: UserDefinedCommandOutput
+    def __init__(self, output: UserDefinedCommandOutput) -> None: ...
+
+_CommandOutput: TypeAlias = (
+    PrintFunctionSize
+    | PrintAllFunctionsSize
+    | ExtractVariants
+    | ExtractBest
+    | OverallStatistics
+    | RunScheduleOutput
+    | PrintFunctionOutput
+    | UserDefinedOutput
+)
+
+##
+# Print Function Modes
+##
+
+@final
+class DefaultPrintFunctionMode: ...
+
+@final
+class CSVPrintFunctionMode: ...
+
+_PrintFunctionMode: TypeAlias = DefaultPrintFunctionMode | CSVPrintFunctionMode
 
 ##
 # Schedules
@@ -442,12 +520,6 @@ _Subdatatypes: TypeAlias = SubVariants | NewSort
 # Commands
 ##
 
-@final
-class SetOption:
-    name: str
-    value: _Expr
-    def __init__(self, name: str, value: _Expr) -> None: ...
-
 @final
 class Datatype:
     span: _Span
@@ -469,7 +541,7 @@ class Sort:
     def __init__(self, span: _Span, name: str, presort_and_args: tuple[str, list[_Expr]] | None = None) -> None: ...
 
 @final
-class Function:
+class FunctionCommand:
     span: _Span
     name: str
     schema: Schema
@@ -531,8 +603,12 @@ class Check:
 class PrintFunction:
     span: _Span
     name: str
-    length: int
-    def __init__(self, span: _Span, name: str, length: int) -> None: ...
+    length: int | None
+    filename: str | None
+    mode: _PrintFunctionMode
+    def __init__(
+        self, span: _Span, name: str, length: int | None, filename: str | None, mode: _PrintFunctionMode
+    ) -> None: ...
 
 @final
 class PrintSize:
@@ -613,11 +689,10 @@ class UnstableCombinedRuleset:
     def __init__(self, span: _Span, name: str, rulesets: list[str]) -> None: ...
 
 _Command: TypeAlias = (
-    SetOption
-    | Datatype
+    Datatype
     | Datatypes
     | Sort
-    | Function
+    | FunctionCommand
     | AddRuleset
     | RuleCommand
     | RewriteCommand

egglog/declarations.py

@@ -69,6 +69,7 @@ __all__ = [
     "SaturateDecl",
     "ScheduleDecl",
     "SequenceDecl",
+    "SetCostDecl",
     "SetDecl",
     "SpecialFunctions",
     "TypeOrVarRef",
@@ -854,7 +855,14 @@ class PanicDecl:
     msg: str
 
 
-ActionDecl: TypeAlias = LetDecl | SetDecl | ExprActionDecl | ChangeDecl | UnionDecl | PanicDecl
+@dataclass(frozen=True)
+class SetCostDecl:
+    tp: JustTypeRef
+    expr: CallDecl
+    cost: ExprDecl
+
+
+ActionDecl: TypeAlias = LetDecl | SetDecl | ExprActionDecl | ChangeDecl | UnionDecl | PanicDecl | SetCostDecl
 
 
 ##

egglog/egraph.py

@@ -23,6 +23,7 @@ from typing import (
     get_type_hints,
     overload,
 )
+from warnings import warn
 
 import graphviz
 from typing_extensions import Never, ParamSpec, Self, Unpack, assert_never
@@ -39,7 +40,7 @@ from .thunk import *
 from .version_compat import *
 
 if TYPE_CHECKING:
-    from .builtins import String, Unit
+    from .builtins import String, Unit, i64Like
 
 
 __all__ = [
@@ -83,6 +84,7 @@ __all__ = [
     "run",
     "seq",
     "set_",
+    "set_cost",
     "subsume",
     "union",
     "unstable_combine_rulesets",
@@ -804,7 +806,7 @@ class GraphvizKwargs(TypedDict, total=False):
     max_calls_per_function: int | None
     n_inline_leaves: int
     split_primitive_outputs: bool
-    split_functions: list[object]
+    split_functions: list[ExprCallable]
     include_temporary_functions: bool
 
 
@@ -851,12 +853,12 @@ class EGraph:
         """
         Loads a CSV file and sets it as *input, output of the function.
         """
-        self._egraph.run_program(bindings.Input(span(1), self._callable_to_egg(fn), path))
+        self._egraph.run_program(bindings.Input(span(1), self._callable_to_egg(fn)[1], path))
 
-    def _callable_to_egg(self, fn: object) -> str:
+    def _callable_to_egg(self, fn: ExprCallable) -> tuple[CallableRef, str]:
         ref, decls = resolve_callable(fn)
         self._add_decls(decls)
-        return self._state.callable_ref_to_egg(ref)[0]
+        return ref, self._state.callable_ref_to_egg(ref)[0]
 
     # TODO: Change let to be action...
     def let(self, name: str, expr: BASE_EXPR) -> BASE_EXPR:
@@ -904,12 +906,17 @@ class EGraph:
     def _run_schedule(self, schedule: Schedule) -> bindings.RunReport:
         self._add_decls(schedule)
         egg_schedule = self._state.schedule_to_egg(schedule.schedule)
-        self._egraph.run_program(bindings.RunSchedule(egg_schedule))
-        run_report = self._egraph.run_report()
-        if not run_report:
-            msg = "No run report saved"
-            raise ValueError(msg)
-        return run_report
+        (command_output,) = self._egraph.run_program(bindings.RunSchedule(egg_schedule))
+        assert isinstance(command_output, bindings.RunScheduleOutput)
+        return command_output.report
+
+    def stats(self) -> bindings.RunReport:
+        """
+        Returns the overall run report for the egraph.
+        """
+        (output,) = self._egraph.run_program(bindings.PrintOverallStatistics())
+        assert isinstance(output, bindings.OverallStatistics)
+        return output.report
 
     def check_bool(self, *facts: FactLike) -> bool:
         """
@@ -954,45 +961,41 @@ class EGraph:
         """
         runtime_expr = to_runtime_expr(expr)
         extract_report = self._run_extract(runtime_expr, 0)
-
-        if not isinstance(extract_report, bindings.Best):
-            msg = "No extract report saved"
-            raise ValueError(msg)  # noqa: TRY004
-        (new_typed_expr,) = self._state.exprs_from_egg(
-            extract_report.termdag, [extract_report.term], runtime_expr.__egg_typed_expr__.tp
-        )
-
-        res = cast("BASE_EXPR", RuntimeExpr.__from_values__(self.__egg_decls__, new_typed_expr))
+        assert isinstance(extract_report, bindings.ExtractBest)
+        res = self._from_termdag(extract_report.termdag, extract_report.term, runtime_expr.__egg_typed_expr__.tp)
         if include_cost:
             return res, extract_report.cost
         return res
 
+    def _from_termdag(self, termdag: bindings.TermDag, term: bindings._Term, tp: JustTypeRef) -> Any:
+        (new_typed_expr,) = self._state.exprs_from_egg(termdag, [term], tp)
+        return RuntimeExpr.__from_values__(self.__egg_decls__, new_typed_expr)
+
     def extract_multiple(self, expr: BASE_EXPR, n: int) -> list[BASE_EXPR]:
         """
         Extract multiple expressions from the egraph.
         """
         runtime_expr = to_runtime_expr(expr)
         extract_report = self._run_extract(runtime_expr, n)
-        if not isinstance(extract_report, bindings.Variants):
-            msg = "Wrong extract report type"
-            raise ValueError(msg)  # noqa: TRY004
+        assert isinstance(extract_report, bindings.ExtractVariants)
         new_exprs = self._state.exprs_from_egg(
             extract_report.termdag, extract_report.terms, runtime_expr.__egg_typed_expr__.tp
         )
         return [cast("BASE_EXPR", RuntimeExpr.__from_values__(self.__egg_decls__, expr)) for expr in new_exprs]
 
-    def _run_extract(self, expr: RuntimeExpr, n: int) -> bindings._ExtractReport:
+    def _run_extract(self, expr: RuntimeExpr, n: int) -> bindings._CommandOutput:
         self._add_decls(expr)
         expr = self._state.typed_expr_to_egg(expr.__egg_typed_expr__)
+        # If we have defined any cost tables use the custom extraction
+        args = (expr, bindings.Lit(span(2), bindings.Int(n)))
+        if self._state.cost_callables:
+            cmd: bindings._Command = bindings.UserDefined(span(2), "extract", list(args))
+        else:
+            cmd = bindings.Extract(span(2), *args)
         try:
-            self._egraph.run_program(bindings.Extract(span(2), expr, bindings.Lit(span(2), bindings.Int(n))))
+            return self._egraph.run_program(cmd)[0]
         except BaseException as e:
             raise add_note("Extracting: " + str(expr), e)  # noqa: B904
-        extract_report = self._egraph.extract_report()
-        if not extract_report:
-            msg = "No extract report saved"
-            raise ValueError(msg)
-        return extract_report
 
     def push(self) -> None:
         """
@@ -1037,8 +1040,14 @@ class EGraph:
             max_calls_per_function=max_calls_per_function,
             include_temporary_functions=include_temporary_functions,
         )
+        if serialized.discarded_functions:
+            msg = ", ".join(set(self._state.possible_egglog_functions(serialized.discarded_functions)))
+            warn(f"Omitted: {msg}", stacklevel=3)
+        if serialized.truncated_functions:
+            msg = ", ".join(set(self._state.possible_egglog_functions(serialized.truncated_functions)))
+            warn(f"Truncated: {msg}", stacklevel=3)
         if split_primitive_outputs or split_functions:
-            additional_ops = set(map(self._callable_to_egg, split_functions))
+            additional_ops = {self._callable_to_egg(f)[1] for f in split_functions}
             serialized.split_classes(self._egraph, additional_ops)
         serialized.map_ops(self._state.op_mapping())
 
@@ -1185,6 +1194,58 @@ class EGraph:
                 assert_never(cmd)
         return self._state.command_to_egg(cmd_decl, ruleset_name)
 
+    def function_size(self, fn: ExprCallable) -> int:
+        """
+        Returns the number of rows in a certain function
+        """
+        egg_name = self._callable_to_egg(fn)[1]
+        (output,) = self._egraph.run_program(bindings.PrintSize(span(1), egg_name))
+        assert isinstance(output, bindings.PrintFunctionSize)
+        return output.size
+
+    def all_function_sizes(self) -> list[tuple[ExprCallable, int]]:
+        """
+        Returns a list of all functions and their sizes.
+        """
+        (output,) = self._egraph.run_program(bindings.PrintSize(span(1), None))
+        assert isinstance(output, bindings.PrintAllFunctionsSize)
+        return [
+            (
+                cast(
+                    "ExprCallable",
+                    create_callable(self._state.__egg_decls__, next(iter(refs))),
+                ),
+                size,
+            )
+            for (name, size) in output.sizes
+            if (refs := self._state.egg_fn_to_callable_refs[name])
+        ]
+
+    def function_values(
+        self, fn: Callable[..., BASE_EXPR] | BASE_EXPR, length: int | None = None
+    ) -> dict[BASE_EXPR, BASE_EXPR]:
+        """
+        Given a callable that is a "function", meaning it returns a primitive or has a merge set,
+        returns a mapping of the function applied with its arguments to its values
+
+        If length is specified, only the first `length` values will be returned.
+        """
+        ref, egg_name = self._callable_to_egg(fn)
+        cmd = bindings.PrintFunction(span(1), egg_name, length, None, bindings.DefaultPrintFunctionMode())
+        (output,) = self._egraph.run_program(cmd)
+        assert isinstance(output, bindings.PrintFunctionOutput)
+        signature = self.__egg_decls__.get_callable_decl(ref).signature
+        assert isinstance(signature, FunctionSignature)
+        tp = signature.semantic_return_type.to_just()
+        return {
+            self._from_termdag(output.termdag, call, tp): self._from_termdag(output.termdag, res, tp)
+            for (call, res) in output.terms
+        }
+
+
+# Either a constant or a function.
+ExprCallable: TypeAlias = Callable[..., BaseExpr] | BaseExpr
+
 
 @dataclass(frozen=True)
 class _WrappedMethod:
@@ -1406,10 +1467,13 @@ class Fact:
         """
         Returns True if the two sides of an equality are structurally equal.
         """
-        if not isinstance(self.fact, EqDecl):
-            msg = "Can only check equality facts"
-            raise TypeError(msg)
-        return self.fact.left == self.fact.right
+        match self.fact:
+            case EqDecl(_, left, right):
+                return left == right
+            case ExprFactDecl(TypedExprDecl(_, CallDecl(FunctionRef("!="), (left_tp, right_tp)))):
+                return left_tp != right_tp
+        msg = f"Can only check equality for == or != not {self}"
+        raise ValueError(msg)
 
 
 @dataclass
@@ -1457,6 +1521,18 @@ def panic(message: str) -> Action:
     return Action(Declarations(), PanicDecl(message))
 
 
+def set_cost(expr: BaseExpr, cost: i64Like) -> Action:
+    """Set the cost of the given expression."""
+    from .builtins import i64  # noqa: PLC0415
+
+    expr_runtime = to_runtime_expr(expr)
+    typed_expr_decl = expr_runtime.__egg_typed_expr__
+    expr_decl = typed_expr_decl.expr
+    assert isinstance(expr_decl, CallDecl), "Can only set cost of calls, not literals or vars"
+    cost_decl = to_runtime_expr(convert(cost, i64)).__egg_typed_expr__.expr
+    return Action(expr_runtime.__egg_decls__, SetCostDecl(typed_expr_decl.tp, expr_decl, cost_decl))
+
+
 def let(name: str, expr: BaseExpr) -> Action:
     """Create a let binding."""
     runtime_expr = to_runtime_expr(expr)

egglog/egraph_state.py

@@ -6,7 +6,7 @@ from __future__ import annotations
 
 import re
 from collections import defaultdict
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from typing import TYPE_CHECKING, Literal, overload
 
 from typing_extensions import assert_never
@@ -71,6 +71,9 @@ class EGraphState:
     # Cache of egg expressions for converting to egg
     expr_to_egg_cache: dict[ExprDecl, bindings._Expr] = field(default_factory=dict)
 
+    # Callables which have cost tables associated with them
+    cost_callables: set[CallableRef] = field(default_factory=set)
+
     def copy(self) -> EGraphState:
         """
         Returns a copy of the state. Th egraph reference is kept the same. Used for pushing/popping.
@@ -83,6 +86,7 @@ class EGraphState:
             callable_ref_to_egg_fn=self.callable_ref_to_egg_fn.copy(),
             type_ref_to_egg_sort=self.type_ref_to_egg_sort.copy(),
             expr_to_egg_cache=self.expr_to_egg_cache.copy(),
+            cost_callables=self.cost_callables.copy(),
         )
 
     def schedule_to_egg(self, schedule: ScheduleDecl) -> bindings._Schedule:
@@ -212,9 +216,32 @@ class EGraphState:
                 return bindings.Union(span(), self._expr_to_egg(lhs), self._expr_to_egg(rhs))
             case PanicDecl(name):
                 return bindings.Panic(span(), name)
+            case SetCostDecl(tp, expr, cost):
+                self.type_ref_to_egg(tp)
+                cost_table = self.create_cost_table(expr.callable)
+                args_egg = [self.typed_expr_to_egg(x, False) for x in expr.args]
+                return bindings.Set(span(), cost_table, args_egg, self._expr_to_egg(cost))
             case _:
                 assert_never(action)
 
+    def create_cost_table(self, ref: CallableRef) -> str:
+        """
+        Creates the egg cost table if needed and gets the name of the table.
+        """
+        name = self.cost_table_name(ref)
+        if ref not in self.cost_callables:
+            self.cost_callables.add(ref)
+            signature = self.__egg_decls__.get_callable_decl(ref).signature
+            assert isinstance(signature, FunctionSignature), "Can only add cost tables for functions"
+            signature = replace(signature, return_type=TypeRefWithVars("i64"))
+            self.egraph.run_program(
+                bindings.FunctionCommand(span(), name, self._signature_to_egg_schema(signature), None)
+            )
+        return name
+
+    def cost_table_name(self, ref: CallableRef) -> str:
+        return f"cost_table_{self.callable_ref_to_egg(ref)[0]}"
+
     def fact_to_egg(self, fact: FactDecl) -> bindings._Fact:
         match fact:
             case EqDecl(tp, left, right):
@@ -225,7 +252,7 @@ class EGraphState:
             case _:
                 assert_never(fact)
 
-    def callable_ref_to_egg(self, ref: CallableRef) -> tuple[str, bool]:
+    def callable_ref_to_egg(self, ref: CallableRef) -> tuple[str, bool]:  # noqa: C901, PLR0912
         """
         Returns the egg function name for a callable reference, registering it if it is not already registered.
 
@@ -245,9 +272,12 @@ class EGraphState:
             case ConstantDecl(tp, _):
                 # Use constructor decleration instead of constant b/c constants cannot be extracted
                 # https://github.com/egraphs-good/egglog/issues/334
-                self.egraph.run_program(
-                    bindings.Constructor(span(), egg_name, bindings.Schema([], self.type_ref_to_egg(tp)), None, False)
-                )
+                is_function = self.__egg_decls__._classes[tp.name].builtin
+                schema = bindings.Schema([], self.type_ref_to_egg(tp))
+                if is_function:
+                    self.egraph.run_program(bindings.FunctionCommand(span(), egg_name, schema, None))
+                else:
+                    self.egraph.run_program(bindings.Constructor(span(), egg_name, schema, None, False))
             case FunctionDecl(signature, builtin, _, merge):
                 if isinstance(signature, FunctionSignature):
                     reverse_args = signature.reverse_args
@@ -263,7 +293,7 @@ class EGraphState:
                         self.egraph.run_program(bindings.Relation(span(), egg_name, schema.input))
                     else:
                         self.egraph.run_program(
-                            bindings.Function(
+                            bindings.FunctionCommand(
                                 span(),
                                 egg_name,
                                 self._signature_to_egg_schema(signature),
@@ -347,13 +377,26 @@ class EGraphState:
         """
         Create a mapping of egglog function name to Python function name, for use in the serialized format
         for better visualization.
+
+        Includes cost tables
         """
         return {
             k: pretty_callable_ref(self.__egg_decls__, next(iter(v)))
             for k, v in self.egg_fn_to_callable_refs.items()
             if len(v) == 1
+        } | {
+            self.cost_table_name(ref): f"cost({pretty_callable_ref(self.__egg_decls__, ref, include_all_args=True)})"
+            for ref in self.cost_callables
         }
 
+    def possible_egglog_functions(self, names: list[str]) -> Iterable[str]:
+        """
+        Given a list of egglog functions, returns all the possible Python function strings
+        """
+        for name in names:
+            for c in self.egg_fn_to_callable_refs[name]:
+                yield pretty_callable_ref(self.__egg_decls__, c)
+
     def typed_expr_to_egg(self, typed_expr_decl: TypedExprDecl, transform_let: bool = True) -> bindings._Expr:
         # transform all expressions with multiple parents into a let binding, so that less expressions
         # are sent to egglog. Only for performance reasons.

egglog/examples/jointree.py

@@ -0,0 +1,67 @@
+# mypy: disable-error-code="empty-body"
+
+"""
+Join Tree (custom costs)
+========================
+
+Example of using custom cost functions for jointree.
+
+From https://egraphs.zulipchat.com/#narrow/stream/328972-general/topic/How.20can.20I.20find.20the.20tree.20associated.20with.20an.20extraction.3F
+"""
+
+from __future__ import annotations
+
+from egglog import *
+
+
+class JoinTree(Expr):
+    def __init__(self, name: StringLike) -> None: ...
+
+    def join(self, other: JoinTree) -> JoinTree: ...
+
+    @method(merge=lambda old, new: old.min(new))  # type:ignore[prop-decorator]
+    @property
+    def size(self) -> i64: ...
+
+
+ra = JoinTree("a")
+rb = JoinTree("b")
+rc = JoinTree("c")
+rd = JoinTree("d")
+re = JoinTree("e")
+rf = JoinTree("f")
+
+query = ra.join(rb).join(rc).join(rd).join(re).join(rf)
+
+egraph = EGraph()
+egraph.register(
+    set_(ra.size).to(50),
+    set_(rb.size).to(200),
+    set_(rc.size).to(10),
+    set_(rd.size).to(123),
+    set_(re.size).to(10000),
+    set_(rf.size).to(1),
+)
+
+
+@egraph.register
+def _rules(s: String, a: JoinTree, b: JoinTree, c: JoinTree, asize: i64, bsize: i64):
+    # cost of relation is its size minus 1, since the string arg will have a cost of 1 as well
+    yield rule(JoinTree(s).size == asize).then(set_cost(JoinTree(s), asize - 1))
+    # cost/size of join is product of sizes
+    yield rule(a.join(b), a.size == asize, b.size == bsize).then(
+        set_(a.join(b).size).to(asize * bsize), set_cost(a.join(b), asize * bsize)
+    )
+    # associativity
+    yield rewrite(a.join(b)).to(b.join(a))
+    # commutativity
+    yield rewrite(a.join(b).join(c)).to(a.join(b.join(c)))
+
+
+egraph.register(query)
+egraph.run(1000)
+print(egraph.extract(query))
+print(egraph.extract(query.size))
+
+
+egraph

egglog/exp/array_api.py

@@ -729,7 +729,7 @@ int64 = DType.int64
 _DTYPES = [float64, float32, int32, int64, DType.object]
 
 converter(type, DType, lambda x: convert(np.dtype(x), DType))
-converter(type(np.dtype), DType, lambda x: getattr(DType, x.name))  # type: ignore[call-overload]
+converter(type(np.dtype), DType, lambda x: getattr(DType, x.name))  # type:ignore[call-overload]
 
 
 @array_api_ruleset.register

egglog/pretty.py

@@ -98,6 +98,7 @@ def pretty_callable_ref(
     ref: CallableRef,
     first_arg: ExprDecl | None = None,
     bound_tp_params: tuple[JustTypeRef, ...] | None = None,
+    include_all_args: bool = False,
 ) -> str:
     """
     Pretty print a callable reference, using a dummy value for
@@ -115,6 +116,13 @@ def pretty_callable_ref(
     # Either returns a function or a function with args. If args are provided, they would just be called,
     # on the function, so return them, because they are dummies
     if isinstance(res, tuple):
+        # If we want to include all args as ARG_STR, then we need to figure out how many to use
+        # used for set_cost so that `cost(E(...))` will show up as a call
+        if include_all_args:
+            signature = decls.get_callable_decl(ref).signature
+            assert isinstance(signature, FunctionSignature)
+            correct_args: list[ExprDecl] = [UnboundVarDecl(ARG_STR)] * len(signature.arg_types)
+            return f"{res[0]}({', '.join(context(a, parens=False, unwrap_lit=True) for a in correct_args)})"
         return res[0]
     return res
 
@@ -190,6 +198,9 @@ class TraverseContext:
                 pass
             case DefaultRewriteDecl():
                 pass
+            case SetCostDecl(_, e, c):
+                self(e)
+                self(c)
             case _:
                 assert_never(decl)
 
@@ -285,6 +296,8 @@ class PrettyContext:
                 return f"{change}({self(expr)})", "action"
             case PanicDecl(s):
                 return f"panic({s!r})", "action"
+            case SetCostDecl(_, expr, cost):
+                return f"set_cost({self(expr)}, {self(cost, unwrap_lit=True)})", "action"
             case EqDecl(_, left, right):
                 return f"eq({self(left)}).to({self(right)})", "fact"
             case RulesetDecl(rules):

egglog/runtime.py

@@ -20,6 +20,8 @@ from inspect import Parameter, Signature
 from itertools import zip_longest
 from typing import TYPE_CHECKING, Any, TypeVar, Union, cast, get_args, get_origin
 
+from typing_extensions import assert_never
+
 from .declarations import *
 from .pretty import *
 from .thunk import Thunk
@@ -36,6 +38,7 @@ __all__ = [
     "RuntimeClass",
     "RuntimeExpr",
     "RuntimeFunction",
+    "create_callable",
     "define_expr_method",
     "resolve_callable",
     "resolve_type_annotation",
@@ -340,7 +343,7 @@ class RuntimeClass(DelayedDeclerations, metaclass=ClassFactory):
 
     # Make hashable so can go in Union
     def __hash__(self) -> int:
-        return hash((id(self.__egg_decls_thunk__), self.__egg_tp__))
+        return hash(self.__egg_tp__)
 
     def __eq__(self, other: object) -> bool:
         """
@@ -478,6 +481,9 @@ class RuntimeFunction(DelayedDeclerations):
                 bound_tp_params = args
         return pretty_callable_ref(self.__egg_decls__, self.__egg_ref__, first_arg, bound_tp_params)
 
+    def __repr__(self) -> str:
+        return str(self)
+
 
 def to_py_signature(sig: FunctionSignature, decls: Declarations, optional_args: bool) -> Signature:
     """
@@ -670,11 +676,12 @@ def resolve_callable(callable: object) -> tuple[CallableRef, Declarations]:
     """
     Resolves a runtime callable into a ref
     """
+    # TODO: Make runtime class work with __match_args__
+    if isinstance(callable, RuntimeClass):
+        return InitRef(callable.__egg_tp__.name), callable.__egg_decls__
     match callable:
         case RuntimeFunction(decls, ref, _):
             return ref(), decls()
-        case RuntimeClass(thunk, tp):
-            return InitRef(tp.name), thunk()
         case RuntimeExpr(decl_thunk, expr_thunk):
             if not isinstance((expr := expr_thunk().expr), CallDecl) or not isinstance(
                 expr.callable, ConstantRef | ClassVariableRef
@@ -683,3 +690,23 @@ def resolve_callable(callable: object) -> tuple[CallableRef, Declarations]:
             return expr.callable, decl_thunk()
         case _:
             raise NotImplementedError(f"Cannot turn {callable} of type {type(callable)} into a callable ref")
+
+
+def create_callable(decls: Declarations, ref: CallableRef) -> RuntimeClass | RuntimeFunction | RuntimeExpr:
+    """
+    Creates a callable object from a callable ref. This might not actually be callable, if the ref is a constant
+    or classvar then it is a value
+    """
+    match ref:
+        case InitRef(name):
+            return RuntimeClass(Thunk.value(decls), TypeRefWithVars(name))
+        case FunctionRef() | MethodRef() | ClassMethodRef() | PropertyRef() | UnnamedFunctionRef():
+            bound = JustTypeRef(ref.class_name) if isinstance(ref, ClassMethodRef) else None
+            return RuntimeFunction(Thunk.value(decls), Thunk.value(ref), bound)
+        case ConstantRef(name):
+            tp = decls._constants[name].type_ref
+        case ClassVariableRef(cls_name, var_name):
+            tp = decls._classes[cls_name].class_variables[var_name].type_ref
+        case _:
+            assert_never(ref)
+    return RuntimeExpr.__from_values__(decls, TypedExprDecl(tp, CallDecl(ref)))

{egglog-11.1.0.dist-info → egglog-11.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: egglog
-Version: 11.1.0
+Version: 11.2.0
 Classifier: Environment :: MacOS X
 Classifier: Environment :: Win32 (MS Windows)
 Classifier: Intended Audience :: Developers

{egglog-11.1.0.dist-info → egglog-11.2.0.dist-info}/RECORD

@@ -1,16 +1,16 @@
-egglog-11.1.0.dist-info/METADATA,sha256=u53IcPcs6w0Rry3b4tw2IDFN0qwIOIgbGmj_3Op0FPo,4009
-egglog-11.1.0.dist-info/WHEEL,sha256=Z7DyTkt9YGMi8QI2LcigLpgUbyKncsalZXZkSHrE6qM,131
-egglog-11.1.0.dist-info/licenses/LICENSE,sha256=w7VlVv5O_FPZRo8Z-4Zb_q7D5ac3YDs8JUkMZ4Gq9CE,1070
+egglog-11.2.0.dist-info/METADATA,sha256=10JQyEe0Um6WF78Wz6E6XMAi2zeAF-fT0oAlCVkBh9Q,4009
+egglog-11.2.0.dist-info/WHEEL,sha256=wvPMIrp4MafmZyLFFKa_Gu58t57iXPruL8GrrDqQ9eI,131
+egglog-11.2.0.dist-info/licenses/LICENSE,sha256=w7VlVv5O_FPZRo8Z-4Zb_q7D5ac3YDs8JUkMZ4Gq9CE,1070
 egglog/__init__.py,sha256=0r3MzQbU-9U0fSCeAoJ3deVhZ77tI-1tf8A_WFOhbJs,344
-egglog/bindings.cpython-313-aarch64-linux-gnu.so,sha256=YIS8KRRs95okhZAIpXDAmlReGiX64c0dlbpcjqKM1AQ,201850408
-egglog/bindings.pyi,sha256=QDgs0hKl2xl4Q6f7hXjfDANoMADRrGTYTdGYdri7Vqg,13642
+egglog/bindings.cpython-313-aarch64-linux-gnu.so,sha256=f1Kj2I-3CkE9byzab8SV8-lHpwIWtJdD0GmooIFulLo,164237384
+egglog/bindings.pyi,sha256=Y_YpdAKmVHZ0nIHTTPeg0sigBEPiS8z_U-Z161zKSK4,15330
 egglog/builtins.py,sha256=bF3GanSmsrIIutvuMUmr4SnSNF0zAEwjmDYWbPEEayo,30204
 egglog/config.py,sha256=yM3FIcVCKnhWZmHD0pxkzx7ah7t5PxZx3WUqKtA9tjU,168
 egglog/conversion.py,sha256=DO76lxRbbTqHs6hRo_Lckvtwu0c6LaKoX7k5_B2AfuY,11238
-egglog/declarations.py,sha256=CA8gWolJChVt6-xDL7XvV5ORjNj3O8PREobYCzVFxDI,25784
+egglog/declarations.py,sha256=vBCcK94EqLDfMPasTwCQA_8EU0nVDSLfJoudboi4MgI,25920
 egglog/deconstruct.py,sha256=CovORrpROMIwOLgERPUw8doqRUDUehj6LJEB5FMbpmI,5635
-egglog/egraph.py,sha256=niBx_EAbQhzjupJs_wZ9oIbjF7W_ttx9D0u1iB059HM,60601
-egglog/egraph_state.py,sha256=8-u0BLAHlWgoo_KiyW_aqC6nj_X0XgMFRlXWIyPjF64,28356
+egglog/egraph.py,sha256=HVpHLHzHfvv1ZX7uxmsblUiEF8ZfJKkNjM-h7Cf8Zb0,64227
+egglog/egraph_state.py,sha256=vIa2ACOxXMjrIS4vwA0UkY1I8hDjDUnS9Pq8awFCVfA,30515
 egglog/examples/README.rst,sha256=ztTvpofR0eotSqGoCy_C1fPLDPCncjvcqDanXtLHNNU,232
 egglog/examples/__init__.py,sha256=wm9evUbMPfbtylXIjbDdRTAVMLH4OjT4Z77PCBFyaPU,31
 egglog/examples/bignum.py,sha256=jfL57XXpQqIqizQQ3sSUCCjTrkdjtB71BmjrQIQorQk,535
@@ -18,6 +18,7 @@ egglog/examples/bool.py,sha256=e0z2YoYJsLlhpSADZK1yRYHzilyxSZWGiYAaM0DQ_Gw,695
 egglog/examples/eqsat_basic.py,sha256=2xtM81gG9Br72mr58N-2BUeksR7C_UXnZJ4MvzSPplc,869
 egglog/examples/fib.py,sha256=BOHxKWA7jGx4FURBmfmuZKfLo6xq9-uXAwAXjYid7LU,492
 egglog/examples/higher_order_functions.py,sha256=DNLIQfPJCX_DOLbHNiaYsfvcFIYCYOsRUqp99r9bpc8,1063
+egglog/examples/jointree.py,sha256=60lS8sCgg0V8O6QqBKpN1y7Tz5BRFsvN8L1jM-9mTSE,1723
 egglog/examples/lambda_.py,sha256=iQvwaXVhp2VNOMS7j1WwceZaiq3dqqilwUkMcW5GFBE,8194
 egglog/examples/matrix.py,sha256=7_mPcMcgE-t_GJDyf76-nv3xhPIeN2mvFkc_p_Gnr8g,4961
 egglog/examples/multiset.py,sha256=IBOsB80DkXQ07dYnk1odi27q77LH80Z8zysuLE-Q8F8,1445
@@ -25,7 +26,7 @@ egglog/examples/ndarrays.py,sha256=mfr410eletH8gfdg-P8L90vlF6TUifvYV_-ryOwvZZE,4
 egglog/examples/resolution.py,sha256=BJd5JClA3DBVGfiVRa-H0gbbFvIqeP3uYbhCXHblSQc,2119
 egglog/examples/schedule_demo.py,sha256=JbXdPII7_adxtgyKVAiqCyV2sj88VZ-DhomYrdn8vuc,618
 egglog/exp/__init__.py,sha256=nPtzrH1bz1LVZhZCuS0S9Qild8m5gEikjOVqWAFIa88,49
-egglog/exp/array_api.py,sha256=B4tr5DyC8jkKxdax7k98LK-iXOrsBceSaijwYTEnG-Q,65548
+egglog/exp/array_api.py,sha256=dKgEufUIyoT7J_RvnyGtOkg_DK25ZnxIgt7olVygaH8,65547
 egglog/exp/array_api_jit.py,sha256=Ak4QhmfYLKimjPf8ffUvPv62OhxOneJ9NEWQJuMxKJc,1680
 egglog/exp/array_api_loopnest.py,sha256=-kbyorlGxvlaNsLx1nmLfEZHQM7VMEBwSKtV0l-bs0g,2444
 egglog/exp/array_api_numba.py,sha256=X3H1TnCjPL92uVm6OvcWMJ11IeorAE58zWiOX6huPv4,2696
@@ -33,13 +34,13 @@ egglog/exp/array_api_program_gen.py,sha256=qnve8iqklRQVyGChllG8ZAjAffRpezmdxc3Id
 egglog/exp/program_gen.py,sha256=CavsD70x0ERS87V4OU9zkgMvLXswGEpb1ZZFK0WyN_g,13033
 egglog/exp/siu_examples.py,sha256=yZ-sgH2Y12iTdwBUumP7D2OtCGL83M6pPW7PMobVFXc,719
 egglog/ipython_magic.py,sha256=2hs3g2cSiyDmbCvE2t1OINmu17Bb8MWV--2DpEWwO7I,1189
-egglog/pretty.py,sha256=TCkwZIecdFT5eMIy0GGAZuKvPUt0fQeh5K1srgnCJLo,20136
+egglog/pretty.py,sha256=bX1rZnujgo_oUK0QRveS5ES1hECpR665__qksUxeEKQ,20921
 egglog/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-egglog/runtime.py,sha256=8y1YQ5vFuoHr9B8LGTeQBHZin7S8LYHlsR93LlNl2uo,27979
+egglog/runtime.py,sha256=NJZEE0PJGLfZMeBxhBphKFAuGPq01AaDoHOQxs3k79Y,29160
 egglog/thunk.py,sha256=MrAlPoGK36VQrUrq8PWSaJFu42sPL0yupwiH18lNips,2271
 egglog/type_constraint_solver.py,sha256=U2GjLgbebTLv5QY8_TU0As5wMKL5_NxkHLen9rpfMwI,4518
 egglog/version_compat.py,sha256=EaKRMIOPcatrx9XjCofxZD6Nr5WOooiWNdoapkKleww,3512
 egglog/visualizer.css,sha256=eL0POoThQRc0P4OYnDT-d808ln9O5Qy6DizH9Z5LgWc,259398
 egglog/visualizer.js,sha256=2qZZ-9W_INJx4gZMYjnVXl27IjT_JNuQyEeI2dbjWoU,3753315
 egglog/visualizer_widget.py,sha256=LtVfzOtv2WeKtNuILQQ_9SOHWvRr8YdBYQDKQSgry_s,1319
-egglog-11.1.0.dist-info/RECORD,,
+egglog-11.2.0.dist-info/RECORD,,

{egglog-11.1.0.dist-info → egglog-11.2.0.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: maturin (1.9.3)
+Generator: maturin (1.9.4)
 Root-Is-Purelib: false
 Tag: cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64