egglog 7.2.0__cp311-none-win_amd64.whl → 8.0.1__cp311-none-win_amd64.whl

egglog/egraph.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
+import contextlib
 import inspect
 import pathlib
 import tempfile
 from abc import abstractmethod
-from collections.abc import Callable, Iterable
+from collections.abc import Callable, Generator, Iterable
 from contextvars import ContextVar, Token
 from dataclasses import InitVar, dataclass, field
 from inspect import Parameter, currentframe, signature
@@ -37,8 +38,6 @@ from .runtime import *
 from .thunk import *
 
 if TYPE_CHECKING:
-    import ipywidgets
-
     from .builtins import Bool, PyObject, String, f64, i64
 
 
@@ -464,10 +463,14 @@ class _ExprMetaclass(type):
         prev_frame = frame.f_back
         assert prev_frame
 
+        # Pass in an instance of the class so that when we are generating the decls
+        # we can update them eagerly so that we can access the methods in the class body
+        runtime_cls = RuntimeClass(None, TypeRefWithVars(name))  # type: ignore[arg-type]
+
         # Store frame so that we can get live access to updated locals/globals
         # Otherwise, f_locals returns a copy
         # https://peps.python.org/pep-0667/
-        decls_thunk = Thunk.fn(
+        runtime_cls.__egg_decls_thunk__ = Thunk.fn(
             _generate_class_decls,
             namespace,
             prev_frame,
@@ -475,21 +478,22 @@ class _ExprMetaclass(type):
             egg_sort,
             name,
             ruleset,
-            fallback=Declarations,
+            runtime_cls,
         )
-        return RuntimeClass(decls_thunk, TypeRefWithVars(name))
+        return runtime_cls
 
     def __instancecheck__(cls, instance: object) -> bool:
         return isinstance(instance, RuntimeExpr)
 
 
-def _generate_class_decls(
+def _generate_class_decls(  # noqa: C901,PLR0912
     namespace: dict[str, Any],
     frame: FrameType,
     builtin: bool,
     egg_sort: str | None,
     cls_name: str,
     ruleset: Ruleset | None,
+    runtime_cls: RuntimeClass,
 ) -> Declarations:
     """
     Lazy constructor for class declerations to support classes with methods whose types are not yet defined.
@@ -502,6 +506,8 @@ def _generate_class_decls(
     del parameters
     cls_decl = ClassDecl(egg_sort, type_vars, builtin)
     decls = Declarations(_classes={cls_name: cls_decl})
+    # Update class think eagerly when resolving so that lookups work in methods
+    runtime_cls.__egg_decls_thunk__ = Thunk.value(decls)
 
     ##
     # Register class variables
@@ -522,16 +528,13 @@ def _generate_class_decls(
     # Register methods, classmethods, preserved methods, and properties
     ##
 
-    # The type ref of self is paramterized by the type vars
-    TypeRefWithVars(cls_name, tuple(map(ClassTypeVarRef, type_vars)))
-
     # Get all the methods from the class
     filtered_namespace: list[tuple[str, Any]] = [
         (k, v) for k, v in namespace.items() if k not in IGNORED_ATTRIBUTES or isinstance(v, _WrappedMethod)
     ]
 
     # all methods we should try adding default functions for
-    default_function_refs: dict[ClassMethodRef | MethodRef | PropertyRef, Callable] = {}
+    add_default_funcs: list[Callable[[], None]] = []
     # Then register each of its methods
     for method_name, method in filtered_namespace:
         is_init = method_name == "__init__"
@@ -550,7 +553,6 @@ def _generate_class_decls(
             cls_decl.preserved_methods[method_name] = fn
             continue
         locals = frame.f_locals
-
         ref: ClassMethodRef | MethodRef | PropertyRef | InitRef
         match fn:
             case classmethod():
@@ -561,16 +563,25 @@ def _generate_class_decls(
                 fn = fn.fget
             case _:
                 ref = InitRef(cls_name) if is_init else MethodRef(cls_name, method_name)
+        special_function_name: SpecialFunctions | None = (
+            "fn-partial" if egg_fn == "unstable-fn" else "fn-app" if egg_fn == "unstable-app" else None
+        )
+        if special_function_name:
+            decl = FunctionDecl(special_function_name, builtin=True, egg_name=egg_fn)
+            decls.set_function_decl(ref, decl)
+            continue
 
-        _fn_decl(decls, egg_fn, ref, fn, locals, default, cost, merge, on_merge, mutates, builtin, unextractable)
+        _, add_rewrite = _fn_decl(
+            decls, egg_fn, ref, fn, locals, default, cost, merge, on_merge, mutates, builtin, ruleset, unextractable
+        )
 
         if not builtin and not isinstance(ref, InitRef) and not mutates:
-            default_function_refs[ref] = fn
+            add_default_funcs.append(add_rewrite)
 
     # Add all rewrite methods at the end so that all methods are registered first and can be accessed
     # in the bodies
-    for ref, fn in default_function_refs.items():
-        _add_default_rewrite_function(decls, ref, fn, ruleset)
+    for add_rewrite in add_default_funcs:
+        add_rewrite()
 
     return decls
 
@@ -590,6 +601,7 @@ def function(
     unextractable: bool = False,
     builtin: bool = False,
     ruleset: Ruleset | None = None,
+    use_body_as_name: bool = False,
 ) -> Callable[[CALLABLE], CALLABLE]: ...
 
 
@@ -604,6 +616,7 @@ def function(
     mutates_first_arg: bool = False,
     unextractable: bool = False,
     ruleset: Ruleset | None = None,
+    use_body_as_name: bool = False,
 ) -> Callable[[Callable[P, EXPR]], Callable[P, EXPR]]: ...
 
 
@@ -635,16 +648,18 @@ class _FunctionConstructor:
     on_merge: Callable[[RuntimeExpr, RuntimeExpr], Iterable[ActionLike]] | None = None
     unextractable: bool = False
     ruleset: Ruleset | None = None
+    use_body_as_name: bool = False
 
     def __call__(self, fn: Callable[..., RuntimeExpr]) -> RuntimeFunction:
-        return RuntimeFunction(Thunk.fn(self.create_decls, fn), FunctionRef(fn.__name__))
+        return RuntimeFunction(*split_thunk(Thunk.fn(self.create_decls, fn)))
 
-    def create_decls(self, fn: Callable[..., RuntimeExpr]) -> Declarations:
+    def create_decls(self, fn: Callable[..., RuntimeExpr]) -> tuple[Declarations, CallableRef]:
         decls = Declarations()
-        _fn_decl(
+        ref = None if self.use_body_as_name else FunctionRef(fn.__name__)
+        ref, add_rewrite = _fn_decl(
             decls,
             self.egg_fn,
-            (ref := FunctionRef(fn.__name__)),
+            ref,
             fn,
             self.hint_locals,
             self.default,
@@ -653,16 +668,18 @@ class _FunctionConstructor:
             self.on_merge,
             self.mutates_first_arg,
             self.builtin,
+            self.ruleset,
             unextractable=self.unextractable,
         )
-        _add_default_rewrite_function(decls, ref, fn, self.ruleset)
-        return decls
+        add_rewrite()
+        return decls, ref
 
 
 def _fn_decl(
     decls: Declarations,
     egg_name: str | None,
-    ref: FunctionRef | MethodRef | PropertyRef | ClassMethodRef | InitRef,
+    # If ref is Callable, then generate the ref from the function name
+    ref: FunctionRef | MethodRef | PropertyRef | ClassMethodRef | InitRef | None,
     fn: object,
     # Pass in the locals, retrieved from the frame when wrapping,
     # so that we support classes and function defined inside of other functions (which won't show up in the globals)
@@ -673,30 +690,26 @@ def _fn_decl(
     on_merge: Callable[[RuntimeExpr, RuntimeExpr], Iterable[ActionLike]] | None,
     mutates_first_arg: bool,
     is_builtin: bool,
+    ruleset: Ruleset | None = None,
     unextractable: bool = False,
-) -> None:
+) -> tuple[CallableRef, Callable[[], None]]:
     """
-    Sets the function decl for the function object.
+    Sets the function decl for the function object and returns the ref as well as a thunk that sets the default callable.
     """
     if not isinstance(fn, FunctionType):
         raise NotImplementedError(f"Can only generate function decls for functions not {fn}  {type(fn)}")
 
-    # partial function creation and calling are handled with a special case in the type checker, so don't
-    # use the normal logic
-    special_function_name: SpecialFunctions | None = (
-        "fn-partial" if egg_name == "unstable-fn" else "fn-app" if egg_name == "unstable-app" else None
-    )
-    if special_function_name:
-        decls.set_function_decl(ref, FunctionDecl(special_function_name, builtin=True, egg_name=egg_name))
-        return
-
     hint_globals = fn.__globals__.copy()
     # Copy Callable into global if not present bc sometimes it gets automatically removed by ruff to type only block
     # 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())
 
@@ -719,7 +732,7 @@ def _fn_decl(
 
     # Resolve all default values as arg types
     arg_defaults = [
-        resolve_literal(t, p.default) if p.default is not Parameter.empty else None
+        resolve_literal(t, p.default, Thunk.value(decls)) if p.default is not Parameter.empty else None
         for (t, p) in zip(arg_types, params, strict=True)
     ]
 
@@ -740,8 +753,8 @@ def _fn_decl(
         None
         if merge is None
         else merge(
-            RuntimeExpr.__from_value__(decls, TypedExprDecl(return_type.to_just(), VarDecl("old"))),
-            RuntimeExpr.__from_value__(decls, TypedExprDecl(return_type.to_just(), VarDecl("new"))),
+            RuntimeExpr.__from_values__(decls, TypedExprDecl(return_type.to_just(), VarDecl("old", False))),
+            RuntimeExpr.__from_values__(decls, TypedExprDecl(return_type.to_just(), VarDecl("new", False))),
         )
     )
     decls |= merged
@@ -751,29 +764,47 @@ def _fn_decl(
         if on_merge is None
         else _action_likes(
             on_merge(
-                RuntimeExpr.__from_value__(decls, TypedExprDecl(return_type.to_just(), VarDecl("old"))),
-                RuntimeExpr.__from_value__(decls, TypedExprDecl(return_type.to_just(), VarDecl("new"))),
+                RuntimeExpr.__from_values__(decls, TypedExprDecl(return_type.to_just(), VarDecl("old", False))),
+                RuntimeExpr.__from_values__(decls, TypedExprDecl(return_type.to_just(), VarDecl("new", False))),
             )
         )
     )
     decls.update(*merge_action)
-    decl = FunctionDecl(
-        FunctionSignature(
+    # defer this in generator so it doesnt resolve for builtins eagerly
+    args = (TypedExprDecl(tp.to_just(), VarDecl(name, False)) for name, tp in zip(arg_names, arg_types, strict=True))
+    res_ref: FunctionRef | MethodRef | ClassMethodRef | PropertyRef | InitRef | UnnamedFunctionRef
+    res_thunk: Callable[[], object]
+    # If we were not passed in a ref, this is an unnamed funciton, so eagerly compute the value and use that to refer to it
+    if not ref:
+        tuple_args = tuple(args)
+        res = _create_default_value(decls, ref, fn, tuple_args, ruleset)
+        assert isinstance(res, RuntimeExpr)
+        res_ref = UnnamedFunctionRef(tuple_args, res.__egg_typed_expr__)
+        decls._unnamed_functions.add(res_ref)
+        res_thunk = Thunk.value(res)
+
+    else:
+        signature_ = FunctionSignature(
             return_type=None if mutates_first_arg else return_type,
             var_arg_type=var_arg_type,
             arg_types=arg_types,
             arg_names=arg_names,
             arg_defaults=tuple(a.__egg_typed_expr__.expr if a is not None else None for a in arg_defaults),
-        ),
-        cost=cost,
-        egg_name=egg_name,
-        merge=merged.__egg_typed_expr__.expr if merged is not None else None,
-        unextractable=unextractable,
-        builtin=is_builtin,
-        default=None if default is None else default.__egg_typed_expr__.expr,
-        on_merge=tuple(a.action for a in merge_action),
-    )
-    decls.set_function_decl(ref, decl)
+        )
+        decl = FunctionDecl(
+            signature=signature_,
+            cost=cost,
+            egg_name=egg_name,
+            merge=merged.__egg_typed_expr__.expr if merged is not None else None,
+            unextractable=unextractable,
+            builtin=is_builtin,
+            default=None if default is None else default.__egg_typed_expr__.expr,
+            on_merge=tuple(a.action for a in merge_action),
+        )
+        res_ref = ref
+        decls.set_function_decl(ref, decl)
+        res_thunk = Thunk.fn(_create_default_value, decls, ref, fn, args, ruleset)
+    return res_ref, Thunk.fn(_add_default_rewrite_function, decls, res_ref, return_type, ruleset, res_thunk)
 
 
 # Overload to support aritys 0-4 until variadic generic support map, so we can map from type to value
@@ -804,7 +835,7 @@ def relation(name: str, /, *tps: type, egg_fn: str | None = None) -> Callable[..
     Creates a function whose return type is `Unit` and has a default value.
     """
     decls_thunk = Thunk.fn(_relation_decls, name, tps, egg_fn)
-    return cast(Callable[..., Unit], RuntimeFunction(decls_thunk, FunctionRef(name)))
+    return cast(Callable[..., Unit], RuntimeFunction(decls_thunk, Thunk.value(FunctionRef(name))))
 
 
 def _relation_decls(name: str, tps: tuple[type, ...], egg_fn: str | None) -> Declarations:
@@ -828,7 +859,9 @@ def constant(
     A "constant" is implemented as the instantiation of a value that takes no args.
     This creates a function with `name` and return type `tp` and returns a value of it being called.
     """
-    return cast(EXPR, RuntimeExpr(Thunk.fn(_constant_thunk, name, tp, egg_name, default_replacement, ruleset)))
+    return cast(
+        EXPR, RuntimeExpr(*split_thunk(Thunk.fn(_constant_thunk, name, tp, egg_name, default_replacement, ruleset)))
+    )
 
 
 def _constant_thunk(
@@ -842,24 +875,34 @@ def _constant_thunk(
     return decls, TypedExprDecl(type_ref.to_just(), CallDecl(callable_ref))
 
 
-def _add_default_rewrite_function(decls: Declarations, ref: CallableRef, fn: Callable, ruleset: Ruleset | None) -> None:
-    """
-    Adds a default rewrite for a function, by calling the functions with vars and adding it if it is not None.
-    """
-    callable_decl = decls.get_callable_decl(ref)
-    assert isinstance(callable_decl, FunctionDecl)
-    signature = callable_decl.signature
-    assert isinstance(signature, FunctionSignature)
-
-    var_args: list[object] = [
-        RuntimeExpr.__from_value__(decls, TypedExprDecl(tp.to_just(), VarDecl(_rule_var_name(name))))
-        for name, tp in zip(signature.arg_names, signature.arg_types, strict=False)
-    ]
+def _create_default_value(
+    decls: Declarations,
+    ref: CallableRef | None,
+    fn: Callable,
+    args: Iterable[TypedExprDecl],
+    ruleset: Ruleset | None,
+) -> object:
+    args: list[object] = [RuntimeExpr.__from_values__(decls, a) for a in args]
+
     # If this is a classmethod, add the class as the first arg
     if isinstance(ref, ClassMethodRef):
         tp = decls.get_paramaterized_class(ref.class_name)
-        var_args.insert(0, RuntimeClass(Thunk.value(decls), tp))
-    _add_default_rewrite(decls, ref, signature.semantic_return_type, fn(*var_args), ruleset)
+        args.insert(0, RuntimeClass(Thunk.value(decls), tp))
+    with set_current_ruleset(ruleset):
+        return fn(*args)
+
+
+def _add_default_rewrite_function(
+    decls: Declarations,
+    ref: CallableRef,
+    res_type: TypeOrVarRef,
+    ruleset: Ruleset | None,
+    value_thunk: Callable[[], object],
+) -> None:
+    """
+    Helper functions that resolves a value thunk to create the default value.
+    """
+    _add_default_rewrite(decls, ref, res_type, value_thunk(), ruleset)
 
 
 def _add_default_rewrite(
@@ -872,7 +915,7 @@ def _add_default_rewrite(
     """
     if default_rewrite is None:
         return
-    resolved_value = resolve_literal(type_ref, default_rewrite)
+    resolved_value = resolve_literal(type_ref, default_rewrite, Thunk.value(decls))
     rewrite_decl = DefaultRewriteDecl(ref, resolved_value.__egg_typed_expr__.expr)
     if ruleset:
         ruleset_decls = ruleset._current_egg_decls
@@ -931,6 +974,8 @@ class GraphvizKwargs(TypedDict, total=False):
     max_calls_per_function: int | None
     n_inline_leaves: int
     split_primitive_outputs: bool
+    split_functions: list[object]
+    include_temporary_functions: bool
 
 
 @dataclass
@@ -973,81 +1018,19 @@ class EGraph(_BaseModule):
             raise ValueError(msg)
         return cmds
 
-    def _repr_mimebundle_(self, *args, **kwargs):
-        """
-        Returns the graphviz representation of the e-graph.
-        """
-        return {"image/svg+xml": self.graphviz().pipe(format="svg", quiet=True, encoding="utf-8")}
-
-    def graphviz(self, **kwargs: Unpack[GraphvizKwargs]) -> graphviz.Source:
-        # By default we want to split primitive outputs
-        kwargs.setdefault("split_primitive_outputs", True)
-        n_inline = kwargs.pop("n_inline_leaves", 0)
-        serialized = self._egraph.serialize([], **kwargs)  # type: ignore[misc]
-        serialized.map_ops(self._state.op_mapping())
-        for _ in range(n_inline):
-            serialized.inline_leaves()
-        original = serialized.to_dot()
-        # Add link to stylesheet to the graph, so that edges light up on hover
-        # https://gist.github.com/sverweij/93e324f67310f66a8f5da5c2abe94682
-        styles = """/* the lines within the edges */
-      .edge:active path,
-      .edge:hover path {
-        stroke: fuchsia;
-        stroke-width: 3;
-        stroke-opacity: 1;
-      }
-      /* arrows are typically drawn with a polygon */
-      .edge:active polygon,
-      .edge:hover polygon {
-        stroke: fuchsia;
-        stroke-width: 3;
-        fill: fuchsia;
-        stroke-opacity: 1;
-        fill-opacity: 1;
-      }
-      /* If you happen to have text and want to color that as well... */
-      .edge:active text,
-      .edge:hover text {
-        fill: fuchsia;
-      }"""
-        p = pathlib.Path(tempfile.gettempdir()) / "graphviz-styles.css"
-        p.write_text(styles)
-        with_stylesheet = original.replace("{", f'{{stylesheet="{p!s}"', 1)
-        return graphviz.Source(with_stylesheet)
-
-    def graphviz_svg(self, **kwargs: Unpack[GraphvizKwargs]) -> str:
-        return self.graphviz(**kwargs).pipe(format="svg", quiet=True, encoding="utf-8")
-
-    def _repr_html_(self) -> str:
-        """
-        Add a _repr_html_ to be an SVG to work with sphinx gallery.
-
-        ala https://github.com/xflr6/graphviz/pull/121
-        until this PR is merged and released
-        https://github.com/sphinx-gallery/sphinx-gallery/pull/1138
-        """
-        return self.graphviz_svg()
-
-    def display(self, **kwargs: Unpack[GraphvizKwargs]) -> None:
-        """
-        Displays the e-graph in the notebook.
-        """
-        if IN_IPYTHON:
-            from IPython.display import SVG, display
-
-            display(SVG(self.graphviz_svg(**kwargs)))
-        else:
-            self.graphviz(**kwargs).render(view=True, format="svg", quiet=True)
+    def _ipython_display_(self) -> None:
+        self.display()
 
     def input(self, fn: Callable[..., String], path: str) -> None:
         """
         Loads a CSV file and sets it as *input, output of the function.
         """
+        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)
         self._add_decls(decls)
-        fn_name = self._state.callable_ref_to_egg(ref)
-        self._egraph.run_program(bindings.Input(fn_name, path))
+        return self._state.callable_ref_to_egg(ref)
 
     def let(self, name: str, expr: EXPR) -> EXPR:
         """
@@ -1059,8 +1042,8 @@ class EGraph(_BaseModule):
         self._add_decls(runtime_expr)
         return cast(
             EXPR,
-            RuntimeExpr.__from_value__(
-                self.__egg_decls__, TypedExprDecl(runtime_expr.__egg_typed_expr__.tp, VarDecl(name))
+            RuntimeExpr.__from_values__(
+                self.__egg_decls__, TypedExprDecl(runtime_expr.__egg_typed_expr__.tp, VarDecl(name, True))
             ),
         )
 
@@ -1084,13 +1067,13 @@ 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"
             raise ValueError(msg)  # noqa: TRY004
         (new_typed_expr,) = self._state.exprs_from_egg(extract_report.termdag, [extract_report.term], typed_expr.tp)
-        return cast(EXPR, RuntimeExpr.__from_value__(self.__egg_decls__, new_typed_expr))
+        return cast(EXPR, RuntimeExpr.__from_values__(self.__egg_decls__, new_typed_expr))
 
     def include(self, path: str) -> None:
         """
@@ -1139,13 +1122,13 @@ 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)
         self._add_decls(*facts)
         egg_facts = [self._state.fact_to_egg(f.fact) for f in _fact_likes(facts)]
-        return bindings.Check(egg_facts)
+        return bindings.Check(bindings.DUMMY_SPAN, egg_facts)
 
     @overload
     def extract(self, expr: EXPR, /, include_cost: Literal[False] = False) -> EXPR: ...
@@ -1167,7 +1150,7 @@ class EGraph(_BaseModule):
             raise ValueError(msg)  # noqa: TRY004
         (new_typed_expr,) = self._state.exprs_from_egg(extract_report.termdag, [extract_report.term], typed_expr.tp)
 
-        res = cast(EXPR, RuntimeExpr.__from_value__(self.__egg_decls__, new_typed_expr))
+        res = cast(EXPR, RuntimeExpr.__from_values__(self.__egg_decls__, new_typed_expr))
         if include_cost:
             return res, extract_report.cost
         return res
@@ -1185,11 +1168,15 @@ class EGraph(_BaseModule):
             msg = "Wrong extract report type"
             raise ValueError(msg)  # noqa: TRY004
         new_exprs = self._state.exprs_from_egg(extract_report.termdag, extract_report.terms, typed_expr.tp)
-        return [cast(EXPR, RuntimeExpr.__from_value__(self.__egg_decls__, expr)) for expr in new_exprs]
+        return [cast(EXPR, RuntimeExpr.__from_values__(self.__egg_decls__, expr)) for expr in new_exprs]
 
     def _run_extract(self, typed_expr: TypedExprDecl, n: int) -> bindings._ExtractReport:
         expr = self._state.typed_expr_to_egg(typed_expr)
-        self._egraph.run_program(bindings.ActionCommand(bindings.Extract(expr, bindings.Lit(bindings.Int(n)))))
+        self._egraph.run_program(
+            bindings.ActionCommand(
+                bindings.Extract(bindings.DUMMY_SPAN, expr, bindings.Lit(bindings.DUMMY_SPAN, bindings.Int(n)))
+            )
+        )
         extract_report = self._egraph.extract_report()
         if not extract_report:
             msg = "No extract report saved"
@@ -1208,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:
@@ -1221,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()
 
@@ -1261,42 +1248,119 @@ class EGraph(_BaseModule):
                 return self._egraph.eval_py_object(egg_expr)
         raise TypeError(f"Eval not implemented for {typed_expr.tp}")
 
-    def saturate(
-        self, *, max: int = 1000, performance: bool = False, **kwargs: Unpack[GraphvizKwargs]
-    ) -> ipywidgets.Widget:
-        from .graphviz_widget import graphviz_widget_with_slider
+    def _serialize(
+        self,
+        **kwargs: Unpack[GraphvizKwargs],
+    ) -> bindings.SerializedEGraph:
+        max_functions = kwargs.pop("max_functions", None)
+        max_calls_per_function = kwargs.pop("max_calls_per_function", None)
+        split_primitive_outputs = kwargs.pop("split_primitive_outputs", True)
+        split_functions = kwargs.pop("split_functions", [])
+        include_temporary_functions = kwargs.pop("include_temporary_functions", False)
+        n_inline_leaves = kwargs.pop("n_inline_leaves", 1)
+        serialized = self._egraph.serialize(
+            [],
+            max_functions=max_functions,
+            max_calls_per_function=max_calls_per_function,
+            include_temporary_functions=include_temporary_functions,
+        )
+        if split_primitive_outputs or split_functions:
+            additional_ops = set(map(self._callable_to_egg, split_functions))
+            serialized.split_classes(self._egraph, additional_ops)
+        serialized.map_ops(self._state.op_mapping())
 
-        dots = [str(self.graphviz(**kwargs))]
-        i = 0
-        while self.run(1).updated and i < max:
-            i += 1
-            dots.append(str(self.graphviz(**kwargs)))
-        return graphviz_widget_with_slider(dots, performance=performance)
+        for _ in range(n_inline_leaves):
+            serialized.inline_leaves()
 
-    def saturate_to_html(
-        self, file: str = "tmp.html", performance: bool = False, **kwargs: Unpack[GraphvizKwargs]
-    ) -> None:
-        # raise NotImplementedError("Upstream bugs prevent rendering to HTML")
+        return serialized
+
+    def _graphviz(self, **kwargs: Unpack[GraphvizKwargs]) -> graphviz.Source:
+        serialized = self._serialize(**kwargs)
+
+        original = serialized.to_dot()
+        # Add link to stylesheet to the graph, so that edges light up on hover
+        # https://gist.github.com/sverweij/93e324f67310f66a8f5da5c2abe94682
+        styles = """/* the lines within the edges */
+      .edge:active path,
+      .edge:hover path {
+        stroke: fuchsia;
+        stroke-width: 3;
+        stroke-opacity: 1;
+      }
+      /* arrows are typically drawn with a polygon */
+      .edge:active polygon,
+      .edge:hover polygon {
+        stroke: fuchsia;
+        stroke-width: 3;
+        fill: fuchsia;
+        stroke-opacity: 1;
+        fill-opacity: 1;
+      }
+      /* If you happen to have text and want to color that as well... */
+      .edge:active text,
+      .edge:hover text {
+        fill: fuchsia;
+      }"""
+        p = pathlib.Path(tempfile.gettempdir()) / "graphviz-styles.css"
+        p.write_text(styles)
+        with_stylesheet = original.replace("{", f'{{stylesheet="{p!s}"', 1)
+        return graphviz.Source(with_stylesheet)
+
+    def display(self, graphviz: bool = False, **kwargs: Unpack[GraphvizKwargs]) -> None:
+        """
+        Displays the e-graph.
+
+        If in IPython it will display it inline, otherwise it will write it to a file and open it.
+        """
+        from IPython.display import SVG, display
 
-        # import panel
+        from .visualizer_widget import VisualizerWidget
 
-        # panel.extension("ipywidgets")
+        if graphviz:
+            if IN_IPYTHON:
+                svg = self._graphviz(**kwargs).pipe(format="svg", quiet=True, encoding="utf-8")
+                display(SVG(svg))
+            else:
+                self._graphviz(**kwargs).render(view=True, format="svg", quiet=True)
+        else:
+            serialized = self._serialize(**kwargs)
+            VisualizerWidget(egraphs=[serialized.to_json()]).display_or_open()
 
-        widget = self.saturate(performance=performance, **kwargs)
-        # panel.panel(widget).save(file)
+    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.
+        """
+        from .visualizer_widget import VisualizerWidget
 
-        from ipywidgets.embed import embed_minimal_html
+        def to_json() -> str:
+            if expr:
+                print(self.extract(expr))
+            return self._serialize(**kwargs).to_json()
 
-        embed_minimal_html("tmp.html", views=[widget], drop_defaults=False)
-        # Use panel while this issue persists
-        # https://github.com/jupyter-widgets/ipywidgets/issues/3761#issuecomment-1755563436
+        egraphs = [to_json()]
+        i = 0
+        while self.run(schedule or 1).updated and i < max:
+            i += 1
+            egraphs.append(to_json())
+        VisualizerWidget(egraphs=egraphs).display_or_open()
 
     @classmethod
     def current(cls) -> EGraph:
         """
         Returns the current egraph, which is the one in the context.
         """
-        return CURRENT_EGRAPH.get()
+        try:
+            return CURRENT_EGRAPH.get()
+        except LookupError:
+            return cls(save_egglog_string=True)
 
     @property
     def _egraph(self) -> bindings.EGraph:
@@ -1448,7 +1512,8 @@ class Ruleset(Schedule):
         To return the egg decls, we go through our deferred rules and add any we haven't yet
         """
         while self.deferred_rule_gens:
-            rules = self.deferred_rule_gens.pop()()
+            with set_current_ruleset(self):
+                rules = self.deferred_rule_gens.pop()()
             self._current_egg_decls.update(*rules)
             self.__egg_ruleset__.rules.extend(r.decl for r in rules)
         return self._current_egg_decls
@@ -1688,16 +1753,16 @@ def action_command(action: Action) -> Action:
     return action
 
 
-def var(name: str, bound: type[EXPR]) -> EXPR:
+def var(name: str, bound: type[T]) -> T:
     """Create a new variable with the given name and type."""
-    return cast(EXPR, _var(name, bound))
+    return cast(T, _var(name, bound))
 
 
 def _var(name: str, bound: object) -> RuntimeExpr:
     """Create a new variable with the given name and type."""
     decls = Declarations()
     type_ref = resolve_type_annotation(decls, bound)
-    return RuntimeExpr.__from_value__(decls, TypedExprDecl(type_ref.to_just(), VarDecl(name)))
+    return RuntimeExpr.__from_values__(decls, TypedExprDecl(type_ref.to_just(), VarDecl(name, False)))
 
 
 def vars_(names: str, bound: type[EXPR]) -> Iterable[EXPR]:
@@ -1790,7 +1855,7 @@ class _NeBuilder(Generic[EXPR]):
         lhs = to_runtime_expr(self.lhs)
         rhs = convert_to_same_type(rhs, lhs)
         assert isinstance(Unit, RuntimeClass)
-        res = RuntimeExpr.__from_value__(
+        res = RuntimeExpr.__from_values__(
             Declarations.create(Unit, lhs, rhs),
             TypedExprDecl(
                 JustTypeRef("Unit"), CallDecl(FunctionRef("!="), (lhs.__egg_typed_expr__, rhs.__egg_typed_expr__))
@@ -1944,7 +2009,7 @@ def _rewrite_or_rule_generator(gen: RewriteOrRuleGenerator, frame: FrameType) ->
     if "Callable" not in globals:
         globals["Callable"] = Callable
     hints = get_type_hints(gen, globals, frame.f_locals)
-    args = [_var(_rule_var_name(p.name), hints[p.name]) for p in signature(gen).parameters.values()]
+    args = [_var(p.name, hints[p.name]) for p in signature(gen).parameters.values()]
     return list(gen(*args))  # type: ignore[misc]
 
 
@@ -1959,3 +2024,19 @@ def _fact_like(fact_like: FactLike) -> Fact:
     if isinstance(fact_like, Expr):
         return expr_fact(fact_like)
     return fact_like
+
+
+_CURRENT_RULESET = ContextVar[Ruleset | None]("CURRENT_RULESET", default=None)
+
+
+def get_current_ruleset() -> Ruleset | None:
+    return _CURRENT_RULESET.get()
+
+
+@contextlib.contextmanager
+def set_current_ruleset(r: Ruleset | None) -> Generator[None, None, None]:
+    token = _CURRENT_RULESET.set(r)
+    try:
+        yield
+    finally:
+        _CURRENT_RULESET.reset(token)