Upgrade support to new egglog version

Author: saulshanabrook

Cargo.lock

@@ -31,6 +31,9 @@ __all__ = [
     "Fact",
     "Fail",
     "Float",
+    "FrozenEGraph",
+    "FrozenFunction",
+    "FrozenRow",
     "Function",
     "FunctionCommand",
     "FusedIntersect",
@@ -55,6 +58,9 @@ __all__ = [
     "PrintFunctionSize",
     "PrintOverallStatistics",
     "PrintSize",
+    "Prove",
+    "ProveExists",
+    "ProveExistsOutput",
     "Push",
     "Relation",
     "Repeat",
@@ -269,6 +275,7 @@ class TermApp:
     def __new__(cls, name: str, args: list[int]) -> TermApp: ...
 
 _Term: TypeAlias = TermLit | TermVar | TermApp
+_TermId: TypeAlias = int
 
 ##
 # Facts
@@ -531,15 +538,20 @@ class PrintAllFunctionsSize:
 @final
 class ExtractVariants:
     termdag: TermDag
-    terms: list[_Term]
-    def __new__(cls, termdag: TermDag, terms: list[_Term]) -> ExtractVariants: ...
+    terms: list[_TermId]
+    def __new__(cls, termdag: TermDag, terms: list[_TermId]) -> ExtractVariants: ...
 
 @final
 class ExtractBest:
     termdag: TermDag
     cost: int
-    term: _Term
-    def __new__(cls, termdag: TermDag, cost: int, term: _Term) -> ExtractBest: ...
+    term: _TermId
+    def __new__(cls, termdag: TermDag, cost: int, term: _TermId) -> ExtractBest: ...
+
+@final
+class ProveExistsOutput:
+    proof: str
+    def __new__(cls, proof: str) -> ProveExistsOutput: ...
 
 @final
 class OverallStatistics:
@@ -555,10 +567,10 @@ class RunScheduleOutput:
 class PrintFunctionOutput:
     function: Function
     termdag: TermDag
-    terms: list[tuple[_Term, _Term]]
+    terms: list[tuple[_TermId, _TermId]]
     mode: _PrintFunctionMode
     def __new__(
-        cls, function: Function, termdag: TermDag, terms: list[tuple[_Term, _Term]], mode: _PrintFunctionMode
+        cls, function: Function, termdag: TermDag, terms: list[tuple[_TermId, _TermId]], mode: _PrintFunctionMode
     ) -> PrintFunctionOutput: ...
 
 @final
@@ -571,6 +583,7 @@ _CommandOutput: TypeAlias = (
     | PrintAllFunctionsSize
     | ExtractVariants
     | ExtractBest
+    | ProveExistsOutput
     | OverallStatistics
     | RunScheduleOutput
     | PrintFunctionOutput
@@ -718,6 +731,18 @@ class Check:
     facts: list[_Fact]
     def __new__(cls, span: _Span, facts: list[_Fact]) -> Check: ...
 
+@final
+class Prove:
+    span: _Span
+    facts: list[_Fact]
+    def __new__(cls, span: _Span, facts: list[_Fact]) -> Prove: ...
+
+@final
+class ProveExists:
+    span: _Span
+    expr: str
+    def __new__(cls, span: _Span, expr: str) -> ProveExists: ...
+
 @final
 class PrintFunction:
     span: _Span
@@ -823,6 +848,8 @@ _Command: TypeAlias = (
     | RunSchedule
     | Extract
     | Check
+    | Prove
+    | ProveExists
     | PrintFunction
     | PrintSize
     | Output
@@ -845,14 +872,14 @@ _Command: TypeAlias = (
 @final
 class TermDag:
     def size(self) -> int: ...
-    def lookup(self, node: _Term) -> int: ...
-    def get(self, id: int) -> _Term: ...
-    def app(self, sym: str, children: list[int]) -> _Term: ...
-    def lit(self, lit: _Literal) -> _Term: ...
-    def var(self, sym: str) -> _Term: ...
-    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: ...
+    def lookup(self, node: _Term) -> _TermId: ...
+    def get(self, id: _TermId) -> _Term: ...
+    def app(self, sym: str, children: list[_TermId]) -> _TermId: ...
+    def lit(self, lit: _Literal) -> _TermId: ...
+    def var(self, sym: str) -> _TermId: ...
+    def expr_to_term(self, expr: _Expr) -> _TermId: ...
+    def term_to_expr(self, term: _TermId, span: _Span) -> _Expr: ...
+    def to_string(self, term: _TermId) -> str: ...
 
 ##
 # Extraction
@@ -882,10 +909,10 @@ class Extractor(Generic[_COST]):
     def __new__(
         cls, rootsorts: list[str] | None, egraph: EGraph, cost_model: CostModel[_COST, Any]
     ) -> Extractor[_COST]: ...
-    def extract_best(self, egraph: EGraph, termdag: TermDag, value: Value, sort: str) -> tuple[_COST, _Term]: ...
+    def extract_best(self, egraph: EGraph, termdag: TermDag, value: Value, sort: str) -> tuple[_COST, _TermId]: ...
     def extract_variants(
         self, egraph: EGraph, termdag: TermDag, value: Value, nvariants: int, sort: str
-    ) -> list[tuple[_COST, _Term]]: ...
+    ) -> list[tuple[_COST, _TermId]]: ...
 
 ##
 # Frozen

python/egglog/bindings.pyi

@@ -656,7 +656,7 @@ def multiset_not_contains_swapped(x: T, xs: MultiSet[T]) -> Unit: ...
 def multiset_contains_swapped(x: T, xs: MultiSet[T]) -> Unit: ...
 
 
-@function(egg_fn="unstable-multiset-fold", builtin=True)
+@function(egg_fn="unstable-multiset-reduce", builtin=True)
 def multiset_fold(f: Callable[[T, T], T], initial: T, xs: MultiSet[T]) -> T: ...
 
 
@@ -875,7 +875,7 @@ def bool_le(self, other: BigIntLike) -> Bool: ...
     def bool_ge(self, other: BigIntLike) -> Bool: ...
 
 
-converter(i64, BigInt, lambda i: BigInt(i))
+converter(i64, BigInt, BigInt)
 
 BigIntLike: TypeAlias = BigInt | i64Like
 
@@ -1113,6 +1113,7 @@ def __call__(self, *args: *TS) -> T: ...
 
 # Method Type is for builtins like __getitem__
 converter(MethodType, UnstableFn, lambda m: UnstableFn[*get_type_args()](m.__func__, m.__self__))  # type: ignore[operator, misc]
+# Ignore PLW0108.
 converter(RuntimeFunction, UnstableFn, lambda rf: UnstableFn[*get_type_args()](rf))  # type: ignore[operator, misc]
 # converter(RuntimeClass, UnstableFn, lambda rc: UnstableFn[*get_type_args()](rc))  # type: ignore[operator, misc]
 converter(partial, UnstableFn, lambda p: UnstableFn[*get_type_args()](p.func, *p.args))  # type: ignore[operator, misc]

python/egglog/builtins.py

@@ -868,6 +868,7 @@ class EGraph:
     # For storing the global "current" egraph
     _token_stack: list[EGraph] = field(default_factory=list, repr=False)
 
+    # TODO: Add EGraph(...commands) constructor
     def __post_init__(self, seminaive: bool, save_egglog_string: bool) -> None:
         egraph = bindings.EGraph(seminaive=seminaive, record=save_egglog_string)
         self._state = EGraphState(egraph)
@@ -1037,7 +1038,7 @@ def extract(
             res = self._from_termdag(termdag, term, tp)
         return (res, cost) if include_cost else res
 
-    def _from_termdag(self, termdag: bindings.TermDag, term: bindings._Term, tp: JustTypeRef) -> Any:
+    def _from_termdag(self, termdag: bindings.TermDag, term: int, tp: JustTypeRef) -> Any:
         (new_typed_expr,) = self._state.exprs_from_egg(termdag, [term], tp)
         return RuntimeExpr.__from_values__(self.__egg_decls__, new_typed_expr)
 
@@ -1179,13 +1180,14 @@ def display(self, graphviz: bool = False, **kwargs: Unpack[GraphvizKwargs]) -> N
             serialized = self._serialize(**kwargs)
             VisualizerWidget(egraphs=[serialized.to_json()]).display_or_open()
 
-    def saturate(
+    def saturate(  # noqa: C901
         self,
         schedule: Schedule | None = None,
         *,
         expr: Expr | None = None,
         max: int = 1000,
         visualize: bool = True,
+        print_frozen: bool = False,
         **kwargs: Unpack[GraphvizKwargs],
     ) -> None:
         """
@@ -1210,9 +1212,14 @@ def to_json() -> str:
                 i += 1
                 if visualize:
                     egraphs.append(to_json())
+                if print_frozen:
+                    print(f"After iteration {i}:")
+                    print(self.freeze())
         except:
             if visualize:
                 egraphs.append(to_json())
+            if print_frozen:
+                print(self.freeze())
             raise
         finally:
             if visualize:

python/egglog/egraph.py

@@ -631,14 +631,12 @@ def translate_call(self, expr: CallDecl | GetCostDecl) -> tuple[str, list[TypedE
             case _:
                 assert_never(expr)
 
-    def exprs_from_egg(
-        self, termdag: bindings.TermDag, terms: list[bindings._Term], tp: JustTypeRef
-    ) -> Iterable[TypedExprDecl]:
+    def exprs_from_egg(self, termdag: bindings.TermDag, terms: list[int], tp: JustTypeRef) -> Iterable[TypedExprDecl]:
         """
         Create a function that can convert from an egg term to a typed expr.
         """
         state = FromEggState(self, termdag)
-        return [state.from_expr(tp, term) for term in terms]
+        return [state.resolve_term(term_id, tp) for term_id in terms]
 
     def _get_possible_types(self, cls_ident: Ident) -> frozenset[JustTypeRef]:
         """

python/egglog/egraph_state.py

@@ -63,8 +63,8 @@ def __bool__(self) -> bool:
             not _CURRENT_EGRAPH
             and (
                 args := get_callable_args(self, Int.__eq__)
-                or get_callable_args(self, Boolean.__eq__)
-                or get_callable_args(self, Value.__eq__)
+                or get_callable_args(self, Boolean.__eq__)  # type: ignore[arg-type]
+                or get_callable_args(self, Value.__eq__)  # type: ignore[arg-type]
             )
             is not None
         ):

python/egglog/exp/array_api.py

@@ -47,7 +47,7 @@ def bind_class(self, ref: JustTypeRef, decls: Declarations) -> None:
         try:
             cls_typevars = decls.get_class_decl(ref.ident).type_vars
         except KeyError:
-            cls_typevars = []
+            cls_typevars = ()
         for typevar, arg in zip(cls_typevars, ref.args, strict=True):
             self.infer_typevars(typevar, arg)