egglog 8.0.0__cp312-none-win_amd64.whl → 8.0.1__cp312-none-win_amd64.whl

egglog/bindings.pyi

@@ -5,8 +5,6 @@ from typing import TypeAlias
 
 from typing_extensions import final
 
-HIGH_COST: int
-
 @final
 class SerializedEGraph:
     def inline_leaves(self) -> None: ...
@@ -14,13 +12,14 @@ class SerializedEGraph:
     def to_dot(self) -> str: ...
     def to_json(self) -> str: ...
     def map_ops(self, map: dict[str, str]) -> None: ...
-    def split_e_classes(self, egraph: EGraph, ops: set[str]) -> None: ...
+    def split_classes(self, egraph: EGraph, ops: set[str]) -> None: ...
 
 @final
 class PyObjectSort:
     def __init__(self) -> None: ...
     def store(self, __o: object, /) -> _Expr: ...
 
+def parse_program(__input: str, /, filename: str | None = None) -> list[_Command]: ...
 @final
 class EGraph:
     def __init__(
@@ -29,11 +28,9 @@ class EGraph:
         *,
         fact_directory: str | Path | None = None,
         seminaive: bool = True,
-        terms_encoding: bool = False,
         record: bool = False,
     ) -> None: ...
     def commands(self) -> str | None: ...
-    def parse_program(self, __input: str, /, filename: str | None = None) -> list[_Command]: ...
     def run_program(self, *commands: _Command) -> list[str]: ...
     def extract_report(self) -> _ExtractReport | None: ...
     def run_report(self) -> RunReport | None: ...
@@ -250,6 +247,7 @@ _Action: TypeAlias = Let | Set | Change | Union | Panic | Expr_ | Extract
 
 @final
 class FunctionDecl:
+    span: Span
     name: str
     schema: Schema
     default: _Expr | None
@@ -261,6 +259,7 @@ class FunctionDecl:
 
     def __init__(
         self,
+        span: Span,
         name: str,
         schema: Schema,
         default: _Expr | None = None,
@@ -273,7 +272,8 @@ class FunctionDecl:
 
 @final
 class Variant:
-    def __init__(self, name: str, types: list[str], cost: int | None = None) -> None: ...
+    def __init__(self, span: Span, name: str, types: list[str], cost: int | None = None) -> None: ...
+    span: Span
     name: str
     types: list[str]
     cost: int | None
@@ -379,6 +379,23 @@ class Sequence:
 
 _Schedule: TypeAlias = Saturate | Repeat | Run | Sequence
 
+##
+# Subdatatypes
+##
+
+@final
+class SubVariants:
+    def __init__(self, variants: list[Variant]) -> None: ...
+    variants: list[Variant]
+
+@final
+class NewSort:
+    def __init__(self, name: str, args: list[_Expr]) -> None: ...
+    name: str
+    args: list[_Expr]
+
+_Subdatatypes: TypeAlias = SubVariants | NewSort
+
 ##
 # Commands
 ##
@@ -391,22 +408,23 @@ class SetOption:
 
 @final
 class Datatype:
+    span: Span
     name: str
     variants: list[Variant]
-    def __init__(self, name: str, variants: list[Variant]) -> None: ...
+    def __init__(self, span: Span, name: str, variants: list[Variant]) -> None: ...
 
 @final
-class Declare:
+class Datatypes:
     span: Span
-    name: str
-    sort: str
-    def __init__(self, span: Span, name: str, sort: str) -> None: ...
+    datatypes: list[tuple[Span, str, _Subdatatypes]]
+    def __init__(self, span: Span, datatypes: list[tuple[Span, str, _Subdatatypes]]) -> None: ...
 
 @final
 class Sort:
+    span: Span
     name: str
     presort_and_args: tuple[str, list[_Expr]] | None
-    def __init__(self, name: str, presort_and_args: tuple[str, list[_Expr]] | None = None) -> None: ...
+    def __init__(self, span: Span, name: str, presort_and_args: tuple[str, list[_Expr]] | None = None) -> None: ...
 
 @final
 class Function:
@@ -452,22 +470,17 @@ class RunSchedule:
 
 @final
 class Simplify:
+    span: Span
     expr: _Expr
     schedule: _Schedule
-    def __init__(self, expr: _Expr, schedule: _Schedule) -> None: ...
-
-@final
-class Calc:
-    span: Span
-    identifiers: list[IdentSort]
-    exprs: list[_Expr]
-    def __init__(self, span: Span, identifiers: list[IdentSort], exprs: list[_Expr]) -> None: ...
+    def __init__(self, span: Span, expr: _Expr, schedule: _Schedule) -> None: ...
 
 @final
 class QueryExtract:
+    span: Span
     variants: int
     expr: _Expr
-    def __init__(self, variants: int, expr: _Expr) -> None: ...
+    def __init__(self, span: Span, variants: int, expr: _Expr) -> None: ...
 
 @final
 class Check:
@@ -484,20 +497,23 @@ class PrintFunction:
 
 @final
 class PrintSize:
+    span: Span
     name: str | None
-    def __init__(self, name: str | None) -> None: ...
+    def __init__(self, span: Span, name: str | None) -> None: ...
 
 @final
 class Output:
+    span: Span
     file: str
     exprs: list[_Expr]
-    def __init__(self, file: str, exprs: list[_Expr]) -> None: ...
+    def __init__(self, span: Span, file: str, exprs: list[_Expr]) -> None: ...
 
 @final
 class Input:
+    span: Span
     name: str
     file: str
-    def __init__(self, name: str, file: str) -> None: ...
+    def __init__(self, span: Span, name: str, file: str) -> None: ...
 
 @final
 class Push:
@@ -506,29 +522,29 @@ class Push:
 
 @final
 class Pop:
+    span: Span
     length: int
-    def __init__(self, length: int) -> None: ...
+    def __init__(self, span: Span, length: int) -> None: ...
 
 @final
 class Fail:
+    span: Span
     command: _Command
-    def __init__(self, command: _Command) -> None: ...
+    def __init__(self, span: Span, command: _Command) -> None: ...
 
 @final
 class Include:
+    span: Span
     path: str
-    def __init__(self, path: str) -> None: ...
-
-@final
-class CheckProof:
-    def __init__(self) -> None: ...
+    def __init__(self, span: Span, path: str) -> None: ...
 
 @final
 class Relation:
+    span: Span
     constructor: str
     inputs: list[str]
 
-    def __init__(self, constructor: str, inputs: list[str]) -> None: ...
+    def __init__(self, span: Span, constructor: str, inputs: list[str]) -> None: ...
 
 @final
 class PrintOverallStatistics:
@@ -543,7 +559,7 @@ class UnstableCombinedRuleset:
 _Command: TypeAlias = (
     SetOption
     | Datatype
-    | Declare
+    | Datatypes
     | Sort
     | Function
     | AddRuleset
@@ -552,7 +568,6 @@ _Command: TypeAlias = (
     | BiRewriteCommand
     | ActionCommand
     | RunSchedule
-    | Calc
     | Simplify
     | QueryExtract
     | Check
@@ -564,7 +579,6 @@ _Command: TypeAlias = (
     | Pop
     | Fail
     | Include
-    | CheckProof
     | Relation
     | PrintOverallStatistics
     | UnstableCombinedRuleset

egglog/conversion.py

@@ -195,6 +195,6 @@ def _get_tp(x: object) -> TypeName | type:
         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:
+    if type(tp) is not type:
         return type(tp)
     return tp

egglog/declarations.py

@@ -622,9 +622,8 @@ def replace_typed_expr(typed_expr: TypedExprDecl, replacements: Mapping[TypedExp
             res = replacements[typed_expr]
         else:
             match typed_expr.expr:
-                case (
+                case CallDecl(callable, args, bound_tp_params) | PartialCallDecl(
                     CallDecl(callable, args, bound_tp_params)
-                    | PartialCallDecl(CallDecl(callable, args, bound_tp_params))
                 ):
                     new_args = tuple(_inner(a) for a in args)
                     call_decl = CallDecl(callable, new_args, bound_tp_params)

egglog/egraph.py

@@ -486,7 +486,7 @@ class _ExprMetaclass(type):
         return isinstance(instance, RuntimeExpr)
 
 
-def _generate_class_decls(  # noqa: C901
+def _generate_class_decls(  # noqa: C901,PLR0912
     namespace: dict[str, Any],
     frame: FrameType,
     builtin: bool,
@@ -704,8 +704,12 @@ def _fn_decl(
     # https://docs.astral.sh/ruff/rules/typing-only-standard-library-import/
     if "Callable" not in hint_globals:
         hint_globals["Callable"] = Callable
-
-    hints = get_type_hints(fn, hint_globals, hint_locals)
+    # Instead of passing both globals and locals, just pass the globals. Otherwise, for some reason forward references
+    # won't be resolved correctly
+    # We need this to be false so it returns "__forward_value__" https://github.com/python/cpython/blob/440ed18e08887b958ad50db1b823e692a747b671/Lib/typing.py#L919
+    # https://github.com/egraphs-good/egglog-python/issues/210
+    hint_globals.update(hint_locals)
+    hints = get_type_hints(fn, hint_globals)
 
     params = list(signature(fn).parameters.values())
 
@@ -1021,7 +1025,7 @@ class EGraph(_BaseModule):
         """
         Loads a CSV file and sets it as *input, output of the function.
         """
-        self._egraph.run_program(bindings.Input(self._callable_to_egg(fn), path))
+        self._egraph.run_program(bindings.Input(bindings.DUMMY_SPAN, self._callable_to_egg(fn), path))
 
     def _callable_to_egg(self, fn: object) -> str:
         ref, decls = resolve_callable(fn)
@@ -1063,7 +1067,7 @@ class EGraph(_BaseModule):
         typed_expr = runtime_expr.__egg_typed_expr__
         # Must also register type
         egg_expr = self._state.typed_expr_to_egg(typed_expr)
-        self._egraph.run_program(bindings.Simplify(egg_expr, egg_schedule))
+        self._egraph.run_program(bindings.Simplify(bindings.DUMMY_SPAN, egg_expr, egg_schedule))
         extract_report = self._egraph.extract_report()
         if not isinstance(extract_report, bindings.Best):
             msg = "No extract report saved"
@@ -1118,7 +1122,7 @@ class EGraph(_BaseModule):
         """
         Checks that one of the facts is not true
         """
-        self._egraph.run_program(bindings.Fail(self._facts_to_check(facts)))
+        self._egraph.run_program(bindings.Fail(bindings.DUMMY_SPAN, self._facts_to_check(facts)))
 
     def _facts_to_check(self, fact_likes: Iterable[FactLike]) -> bindings.Check:
         facts = _fact_likes(fact_likes)
@@ -1191,7 +1195,7 @@ class EGraph(_BaseModule):
         """
         Pop the current state of the egraph, reverting back to the previous state.
         """
-        self._egraph.run_program(bindings.Pop(1))
+        self._egraph.run_program(bindings.Pop(bindings.DUMMY_SPAN, 1))
         self._state = self._state_stack.pop()
 
     def __enter__(self) -> Self:
@@ -1204,7 +1208,7 @@ class EGraph(_BaseModule):
         self.push()
         return self
 
-    def __exit__(self, exc_type, exc, exc_tb) -> None:  # noqa: ANN001
+    def __exit__(self, exc_type, exc, exc_tb) -> None:
         CURRENT_EGRAPH.reset(self._token_stack.pop())
         self.pop()
 
@@ -1262,7 +1266,7 @@ class EGraph(_BaseModule):
         )
         if split_primitive_outputs or split_functions:
             additional_ops = set(map(self._callable_to_egg, split_functions))
-            serialized.split_e_classes(self._egraph, additional_ops)
+            serialized.split_classes(self._egraph, additional_ops)
         serialized.map_ops(self._state.op_mapping())
 
         for _ in range(n_inline_leaves):
@@ -1322,7 +1326,14 @@ class EGraph(_BaseModule):
             serialized = self._serialize(**kwargs)
             VisualizerWidget(egraphs=[serialized.to_json()]).display_or_open()
 
-    def saturate(self, schedule: Schedule | None = None, *, max: int = 1000, **kwargs: Unpack[GraphvizKwargs]) -> None:
+    def saturate(
+        self,
+        schedule: Schedule | None = None,
+        *,
+        expr: Expr | None = None,
+        max: int = 1000,
+        **kwargs: Unpack[GraphvizKwargs],
+    ) -> None:
         """
         Saturate the egraph, running the given schedule until the egraph is saturated.
         It serializes the egraph at each step and returns a widget to visualize the egraph.
@@ -1330,6 +1341,8 @@ class EGraph(_BaseModule):
         from .visualizer_widget import VisualizerWidget
 
         def to_json() -> str:
+            if expr:
+                print(self.extract(expr))
             return self._serialize(**kwargs).to_json()
 
         egraphs = [to_json()]

egglog/egraph_state.py

@@ -206,18 +206,25 @@ class EGraphState:
         self.egg_fn_to_callable_refs[egg_name].add(ref)
         match decl:
             case RelationDecl(arg_types, _, _):
-                self.egraph.run_program(bindings.Relation(egg_name, [self.type_ref_to_egg(a) for a in arg_types]))
+                self.egraph.run_program(
+                    bindings.Relation(bindings.DUMMY_SPAN, egg_name, [self.type_ref_to_egg(a) for a in arg_types])
+                )
             case ConstantDecl(tp, _):
                 # Use function decleration instead of constant b/c constants cannot be extracted
                 # https://github.com/egraphs-good/egglog/issues/334
                 self.egraph.run_program(
-                    bindings.Function(bindings.FunctionDecl(egg_name, bindings.Schema([], self.type_ref_to_egg(tp))))
+                    bindings.Function(
+                        bindings.FunctionDecl(
+                            bindings.DUMMY_SPAN, egg_name, bindings.Schema([], self.type_ref_to_egg(tp))
+                        )
+                    )
                 )
             case FunctionDecl():
                 if not decl.builtin:
                     signature = decl.signature
                     assert isinstance(signature, FunctionSignature), "Cannot turn special function to egg"
                     egg_fn_decl = bindings.FunctionDecl(
+                        bindings.DUMMY_SPAN,
                         egg_name,
                         bindings.Schema(
                             [self.type_ref_to_egg(a.to_just()) for a in signature.arg_types],
@@ -261,7 +268,7 @@ class EGraphState:
                 args = (self.type_ref_to_egg(JustTypeRef(ref.name)), type_args)
             else:
                 args = None
-            self.egraph.run_program(bindings.Sort(egg_name, args))
+            self.egraph.run_program(bindings.Sort(bindings.DUMMY_SPAN, egg_name, args))
         # For builtin classes, let's also make sure we have the mapping of all egg fn names for class methods, because
         # these can be created even without adding them to the e-graph, like `vec-empty` which can be extracted
         # even if you never use that function.
@@ -323,7 +330,7 @@ class EGraphState:
     @overload
     def _expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr: ...
 
-    def _expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr:
+    def _expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr:  # noqa: PLR0912,C901
         """
         Convert an ExprDecl to an egg expression.
         """
@@ -486,7 +493,7 @@ class FromEggState:
             if term.name == "py-object":
                 call = bindings.termdag_term_to_expr(self.termdag, term)
                 expr_decl = PyObjectDecl(self.state.egraph.eval_py_object(call))
-            if term.name == "unstable-fn":
+            elif term.name == "unstable-fn":
                 # Get function name
                 fn_term, *arg_terms = term.args
                 fn_value = self.resolve_term(fn_term, JustTypeRef("String"))

egglog/examples/higher_order_functions.py

@@ -26,7 +26,7 @@ class MathList(Expr):
 
 
 @ruleset
-def math_ruleset(i: i64, j: i64, xs: MathList, x: Math, f: Callable[[Math], Math]):  # noqa: ANN201
+def math_ruleset(i: i64, j: i64, xs: MathList, x: Math, f: Callable[[Math], Math]):
     yield rewrite(Math(i) + Math(j)).to(Math(i + j))
     yield rewrite(xs.append(x).map(f)).to(xs.map(f).append(f(x)))
     yield rewrite(MathList().map(f)).to(MathList())

egglog/exp/array_api.py

@@ -1538,20 +1538,23 @@ def try_evaling(expr: Expr, prim_expr: i64 | Bool) -> int | bool:
     egraph.register(expr)
     egraph.run(array_api_schedule)
     try:
-        return egraph.eval(prim_expr)
+        extracted = egraph.extract(prim_expr)
     except EggSmolError as exc:
-        egraph.display(n_inline_leaves=1, split_primitive_outputs=True)
+        # Try giving some context, by showing the smallest version of the larger expression
         try:
-            msg = f"Cannot simplify to primitive {egraph.extract(expr)}"
-        except EggSmolError:
-            msg = f"Cannot simplify to primitive or extract {expr}"
-
-        # string = (
-        #     egraph.as_egglog_string
-        #     + "\n"
-        #     + str(egraph._state.typed_expr_to_egg(cast(RuntimeExpr, prim_expr).__egg_typed_expr__))
-        # )
-        # # save to "tmp.egg"
-        # with open("tmp.egg", "w") as f:
-        #     f.write(string)
+            expr_extracted = egraph.extract(expr)
+        except EggSmolError as inner_exc:
+            raise ValueError(f"Cannot simplify {expr}") from inner_exc
+        egraph.display(n_inline_leaves=1, split_primitive_outputs=True)
+        msg = f"Cannot simplify to primitive {expr_extracted}"
         raise ValueError(msg) from exc
+    return egraph.eval(extracted)
+
+    # string = (
+    #     egraph.as_egglog_string
+    #     + "\n"
+    #     + str(egraph._state.typed_expr_to_egg(cast(RuntimeExpr, prim_expr).__egg_typed_expr__))
+    # )
+    # # save to "tmp.egg"
+    # with open("tmp.egg", "w") as f:
+    #     f.write(string)

egglog/exp/array_api_jit.py

@@ -17,18 +17,17 @@ def jit(fn: X) -> X:
     # 1. Create variables for each of the two args in the functions
     sig = inspect.signature(fn)
     arg1, arg2 = sig.parameters.keys()
-
-    with EGraph() as egraph:
+    egraph = EGraph()
+    with egraph:
         res = fn(NDArray.var(arg1), NDArray.var(arg2))
         egraph.register(res)
         egraph.run(array_api_numba_schedule)
         res_optimized = egraph.extract(res)
-        egraph.display(split_primitive_outputs=True, n_inline_leaves=3)
+        # egraph.display(split_primitive_outputs=True, n_inline_leaves=3)
 
-    egraph = EGraph()
     fn_program = ndarray_function_two(res_optimized, NDArray.var(arg1), NDArray.var(arg2))
     egraph.register(fn_program)
     egraph.run(array_api_program_gen_schedule)
-    fn = cast(X, egraph.eval(fn_program.py_object))
+    fn = cast(X, egraph.eval(egraph.extract(fn_program.py_object)))
     fn.expr = res_optimized  # type: ignore[attr-defined]
     return fn

egglog/exp/array_api_loopnest.py

@@ -28,7 +28,7 @@ class ShapeAPI(Expr):
 
 
 @array_api_ruleset.register
-def shape_api_ruleset(dims: TupleInt, axis: TupleInt):  # noqa: ANN201
+def shape_api_ruleset(dims: TupleInt, axis: TupleInt):
     s = ShapeAPI(dims)
     yield rewrite(s.deselect(axis)).to(
         ShapeAPI(TupleInt.range(dims.length()).filter(lambda i: ~axis.contains(i)).map(lambda i: dims[i]))
@@ -108,7 +108,7 @@ def _loopnest_api_ruleset(
 
 
 @function(ruleset=array_api_ruleset, unextractable=True)
-def linalg_norm(X: NDArray, axis: TupleIntLike) -> NDArray:  # noqa: N803
+def linalg_norm(X: NDArray, axis: TupleIntLike) -> NDArray:
     # peel off the outer shape for result array
     outshape = ShapeAPI(X.shape).deselect(axis).to_tuple()
     # get only the inner shape for reduction
@@ -126,20 +126,24 @@ def linalg_norm(X: NDArray, axis: TupleIntLike) -> NDArray:  # noqa: N803
 
 
 # %%
-egraph = EGraph(save_egglog_string=True)
+# egraph = EGraph(save_egglog_string=True)
+
+# egraph.register(val.shape)
+# egraph.run(array_api_ruleset.saturate())
+# egraph.extract_multiple(val.shape, 10)
+
+# %%
 
 X = NDArray.var("X")
 assume_shape(X, (3, 2, 3, 4))
 val = linalg_norm(X, (0, 1))
-egraph.register(val.shape)
-egraph.run(array_api_ruleset.saturate())
-egraph.extract_multiple(val.shape, 10)
-
-# %%
 egraph = EGraph()
-egraph.register(val.shape[2])
-egraph.run(array_api_ruleset.saturate())
-egraph.display(split_functions=[Int, TRUE, FALSE], n_inline_leaves=2)
+x = egraph.let("x", val.shape[2])
+# egraph.display(n_inline_leaves=0)
+# egraph.extract(x)
+# egraph.saturate(array_api_ruleset, expr=val.shape[2], split_functions=[Int, TRUE, FALSE], n_inline_leaves=2)
+# egraph.run(array_api_ruleset.saturate())
+# egraph.display()
 
 
 # %%

egglog/exp/array_api_program_gen.py

@@ -1,8 +1,6 @@
 # mypy: disable-error-code="empty-body"
 from __future__ import annotations
 
-import numpy as np
-
 from egglog import *
 
 from .array_api import *
@@ -13,9 +11,12 @@ from .program_gen import *
 # Depends on `np` as a global variable.
 ##
 
-array_api_program_gen_ruleset = ruleset()
+array_api_program_gen_ruleset = ruleset(name="array_api_program_gen_ruleset")
+array_api_program_gen_eval_ruleset = ruleset(name="array_api_program_gen_eval_ruleset")
 
-array_api_program_gen_schedule = array_api_program_gen_ruleset.saturate() + program_gen_ruleset.saturate()
+array_api_program_gen_schedule = (
+    array_api_program_gen_ruleset | program_gen_ruleset | array_api_program_gen_eval_ruleset | eval_program_rulseset
+).saturate()
 
 
 @function
@@ -98,17 +99,14 @@ def _tuple_int_program(i: Int, ti: TupleInt, k: i64, idx_fn: Callable[[Int], Int
 def ndarray_program(x: NDArray) -> Program: ...
 
 
-@function
-def ndarray_function_two(res: NDArray, l: NDArray, r: NDArray) -> Program: ...
+@function(ruleset=array_api_program_gen_ruleset)
+def ndarray_function_two_program(res: NDArray, l: NDArray, r: NDArray) -> Program:
+    return ndarray_program(res).function_two(ndarray_program(l), ndarray_program(r))
 
 
-@array_api_program_gen_ruleset.register
-def _ndarray_function_two(f: Program, res: NDArray, l: NDArray, r: NDArray, o: PyObject):
-    # When we have function, set the program and trigger it to be compiled
-    yield rule(eq(f).to(ndarray_function_two(res, l, r))).then(
-        union(f).with_(ndarray_program(res).function_two(ndarray_program(l), ndarray_program(r))),
-        f.eval_py_object({"np": np}),
-    )
+@function(ruleset=array_api_program_gen_eval_ruleset)
+def ndarray_function_two(res: NDArray, l: NDArray, r: NDArray) -> EvalProgram:
+    return EvalProgram(ndarray_function_two_program(res, l, r), {"np": np})
 
 
 @function

egglog/exp/program_gen.py

@@ -83,8 +83,18 @@ class Program(Expr):
         Only keeps the original parent, not any additional ones, so that each set of statements is only added once.
         """
 
-    @method(default=Unit())
-    def eval_py_object(self, globals: object) -> Unit:
+    @property
+    def is_identifer(self) -> Bool:
+        """
+        Returns whether the expression is an identifier. Used so that we don't re-assign any identifiers.
+        """
+
+
+converter(String, Program, Program)
+
+
+class EvalProgram(Expr):
+    def __init__(self, program: Program, globals: object) -> None:
         """
         Evaluates the program and saves as the py_object
         """
@@ -98,38 +108,34 @@ class Program(Expr):
         """
 
     @property
-    def is_identifer(self) -> Bool:
+    def statements(self) -> String:
         """
-        Returns whether the expression is an identifier. Used so that we don't re-assign any identifiers.
+        Returns the statements of the program, if it's been compiled
         """
 
 
-converter(String, Program, Program)
-
-program_gen_ruleset = ruleset()
-
-
-@program_gen_ruleset.register
-def _py_object(p: Program, expr: String, statements: String, g: PyObject):
+@ruleset
+def eval_program_rulseset(ep: EvalProgram, p: Program, expr: String, statements: String, g: PyObject):
     # When we evaluate a program, we first want to compile to a string
-    yield rule(p.eval_py_object(g)).then(p.compile())
+    yield rule(EvalProgram(p, g)).then(p.compile())
     # Then we want to evaluate the statements/expr
     yield rule(
-        p.eval_py_object(g),
+        eq(ep).to(EvalProgram(p, g)),
         eq(p.statements).to(statements),
         eq(p.expr).to(expr),
     ).then(
-        set_(p.py_object).to(
+        set_(ep.py_object).to(
             py_eval(
                 "l['___res']",
                 PyObject.dict(PyObject.from_string("l"), py_exec(join(statements, "\n", "___res = ", expr), g)),
             )
-        )
+        ),
+        set_(ep.statements).to(statements),
     )
 
 
-@program_gen_ruleset.register
-def _compile(
+@ruleset
+def program_gen_ruleset(
     s: String,
     s1: String,
     s2: String,

egglog/ipython_magic.py

@@ -14,7 +14,7 @@ if IN_IPYTHON:
 
     @needs_local_scope
     @register_cell_magic
-    def egglog(line, cell, local_ns):  # noqa: ANN001, ANN201
+    def egglog(line, cell, local_ns):
         """
         Run an egglog program.