egglog 7.1.0__cp311-none-win_amd64.whl → 8.0.0__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
 
 
@@ -352,7 +351,7 @@ class _BaseModule:
         This is the same as defining a nullary function with a high cost.
         # TODO: Rename as declare to match eggglog?
         """
-        return constant(name, tp, egg_name)
+        return constant(name, tp, egg_name=egg_name)
 
     def register(
         self,
@@ -452,6 +451,7 @@ class _ExprMetaclass(type):
         namespace: dict[str, Any],
         egg_sort: str | None = None,
         builtin: bool = False,
+        ruleset: Ruleset | None = None,
     ) -> RuntimeClass | type:
         # If this is the Expr subclass, just return the class
         if not bases:
@@ -463,20 +463,37 @@ 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(
-            _generate_class_decls, namespace, prev_frame, builtin, egg_sort, name, fallback=Declarations
+        runtime_cls.__egg_decls_thunk__ = Thunk.fn(
+            _generate_class_decls,
+            namespace,
+            prev_frame,
+            builtin,
+            egg_sort,
+            name,
+            ruleset,
+            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(  # noqa: C901
-    namespace: dict[str, Any], frame: FrameType, builtin: bool, egg_sort: str | None, cls_name: str
+    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.
@@ -489,6 +506,8 @@ def _generate_class_decls(  # noqa: C901
     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
@@ -498,9 +517,9 @@ def _generate_class_decls(  # noqa: C901
     for k, v in get_type_hints(_Dummytype, globalns=frame.f_globals, localns=frame.f_locals).items():
         if getattr(v, "__origin__", None) == ClassVar:
             (inner_tp,) = v.__args__
-            type_ref = resolve_type_annotation(decls, inner_tp).to_just()
-            cls_decl.class_variables[k] = ConstantDecl(type_ref)
-
+            type_ref = resolve_type_annotation(decls, inner_tp)
+            cls_decl.class_variables[k] = ConstantDecl(type_ref.to_just())
+            _add_default_rewrite(decls, ClassVariableRef(cls_name, k), type_ref, namespace.pop(k, None), ruleset)
         else:
             msg = f"On class {cls_name}, for attribute '{k}', expected a ClassVar, but got {v}"
             raise NotImplementedError(msg)
@@ -509,14 +528,13 @@ def _generate_class_decls(  # noqa: C901
     # Register methods, classmethods, preserved methods, and properties
     ##
 
-    # The type ref of self is paramterized by the type vars
-    slf_type_ref = 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
+    add_default_funcs: list[Callable[[], None]] = []
     # Then register each of its methods
     for method_name, method in filtered_namespace:
         is_init = method_name == "__init__"
@@ -535,44 +553,35 @@ def _generate_class_decls(  # noqa: C901
             cls_decl.preserved_methods[method_name] = fn
             continue
         locals = frame.f_locals
-
-        def create_decl(fn: object, first: Literal["cls"] | TypeRefWithVars) -> FunctionDecl:
-            special_function_name: SpecialFunctions | None = (
-                "fn-partial" if egg_fn == "unstable-fn" else "fn-app" if egg_fn == "unstable-app" else None  # noqa: B023
-            )
-            if special_function_name:
-                return FunctionDecl(
-                    special_function_name,
-                    builtin=True,
-                    egg_name=egg_fn,  # noqa: B023
-                )
-
-            return _fn_decl(
-                decls,
-                egg_fn,  # noqa: B023
-                fn,
-                locals,  # noqa: B023
-                default,  # noqa: B023
-                cost,  # noqa: B023
-                merge,  # noqa: B023
-                on_merge,  # noqa: B023
-                mutates,  # noqa: B023
-                builtin,
-                first,
-                is_init,  # noqa: B023
-                unextractable,  # noqa: B023
-            )
-
+        ref: ClassMethodRef | MethodRef | PropertyRef | InitRef
         match fn:
             case classmethod():
-                cls_decl.class_methods[method_name] = create_decl(fn.__func__, "cls")
+                ref = ClassMethodRef(cls_name, method_name)
+                fn = fn.__func__
             case property():
-                cls_decl.properties[method_name] = create_decl(fn.fget, slf_type_ref)
+                ref = PropertyRef(cls_name, method_name)
+                fn = fn.fget
             case _:
-                if is_init:
-                    cls_decl.class_methods[method_name] = create_decl(fn, slf_type_ref)
-                else:
-                    cls_decl.methods[method_name] = create_decl(fn, slf_type_ref)
+                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
+
+        _, 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:
+            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 add_rewrite in add_default_funcs:
+        add_rewrite()
 
     return decls
 
@@ -591,6 +600,8 @@ def function(
     mutates_first_arg: bool = False,
     unextractable: bool = False,
     builtin: bool = False,
+    ruleset: Ruleset | None = None,
+    use_body_as_name: bool = False,
 ) -> Callable[[CALLABLE], CALLABLE]: ...
 
 
@@ -604,6 +615,8 @@ def function(
     on_merge: Callable[[EXPR, EXPR], Iterable[ActionLike]] | None = None,
     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]]: ...
 
 
@@ -634,15 +647,19 @@ class _FunctionConstructor:
     merge: Callable[[RuntimeExpr, RuntimeExpr], RuntimeExpr] | None = None
     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()
-        decls._functions[fn.__name__] = _fn_decl(
+        ref = None if self.use_body_as_name else FunctionRef(fn.__name__)
+        ref, add_rewrite = _fn_decl(
             decls,
             self.egg_fn,
+            ref,
             fn,
             self.hint_locals,
             self.default,
@@ -651,14 +668,18 @@ class _FunctionConstructor:
             self.on_merge,
             self.mutates_first_arg,
             self.builtin,
+            self.ruleset,
             unextractable=self.unextractable,
         )
-        return decls
+        add_rewrite()
+        return decls, ref
 
 
 def _fn_decl(
     decls: Declarations,
     egg_name: str | None,
+    # 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)
@@ -669,11 +690,12 @@ def _fn_decl(
     on_merge: Callable[[RuntimeExpr, RuntimeExpr], Iterable[ActionLike]] | None,
     mutates_first_arg: bool,
     is_builtin: bool,
-    # The first arg is either cls, for a classmethod, a self type, or none for a function
-    first_arg: Literal["cls"] | TypeOrVarRef | None = None,
-    is_init: bool = False,
+    ruleset: Ruleset | None = None,
     unextractable: bool = False,
-) -> FunctionDecl:
+) -> tuple[CallableRef, Callable[[], None]]:
+    """
+    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)}")
 
@@ -687,8 +709,8 @@ def _fn_decl(
 
     params = list(signature(fn).parameters.values())
 
-    # If this is an init function, or a classmethod, remove the first arg name
-    if is_init or first_arg == "cls":
+    # If this is an init function, or a classmethod, the first arg is not used
+    if isinstance(ref, ClassMethodRef | InitRef):
         params = params[1:]
 
     if _last_param_variable(params):
@@ -698,40 +720,37 @@ def _fn_decl(
     else:
         var_arg_type = None
     arg_types = tuple(
-        first_arg
-        # If the first arg is a self, and this not an __init__ fn, add this as a typeref
-        if i == 0 and isinstance(first_arg, ClassTypeVarRef | TypeRefWithVars) and not is_init
+        decls.get_paramaterized_class(ref.class_name)
+        if i == 0 and isinstance(ref, MethodRef | PropertyRef)
         else resolve_type_annotation(decls, hints[t.name])
         for i, t in enumerate(params)
     )
 
     # 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)
     ]
 
     decls.update(*arg_defaults)
 
-    # If this is an init fn use the first arg as the return type
-    if is_init:
-        assert not mutates_first_arg
-        if not isinstance(first_arg, ClassTypeVarRef | TypeRefWithVars):
-            msg = "Init function must have a self type"
-            raise ValueError(msg)
-        return_type = first_arg
-    elif mutates_first_arg:
-        return_type = arg_types[0]
-    else:
-        return_type = resolve_type_annotation(decls, hints["return"])
+    return_type = (
+        decls.get_paramaterized_class(ref.class_name)
+        if isinstance(ref, InitRef)
+        else arg_types[0]
+        if mutates_first_arg
+        else resolve_type_annotation(decls, hints["return"])
+    )
+
+    arg_names = tuple(t.name for t in params)
 
     decls |= default
     merged = (
         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
@@ -741,28 +760,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)
-    return 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=tuple(t.name for t in params),
+            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),
-    )
+        )
+        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
@@ -793,7 +831,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:
@@ -804,19 +842,85 @@ def _relation_decls(name: str, tps: tuple[type, ...], egg_fn: str | None) -> Dec
     return decls
 
 
-def constant(name: str, tp: type[EXPR], egg_name: str | None = None) -> EXPR:
+def constant(
+    name: str,
+    tp: type[EXPR],
+    default_replacement: EXPR | None = None,
+    /,
+    *,
+    egg_name: str | None = None,
+    ruleset: Ruleset | None = None,
+) -> EXPR:
     """
     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)))
+    return cast(
+        EXPR, RuntimeExpr(*split_thunk(Thunk.fn(_constant_thunk, name, tp, egg_name, default_replacement, ruleset)))
+    )
 
 
-def _constant_thunk(name: str, tp: type, egg_name: str | None) -> tuple[Declarations, TypedExprDecl]:
+def _constant_thunk(
+    name: str, tp: type, egg_name: str | None, default_replacement: object, ruleset: Ruleset | None
+) -> tuple[Declarations, TypedExprDecl]:
     decls = Declarations()
-    type_ref = resolve_type_annotation(decls, tp).to_just()
-    decls._constants[name] = ConstantDecl(type_ref, egg_name)
-    return decls, TypedExprDecl(type_ref, CallDecl(ConstantRef(name)))
+    type_ref = resolve_type_annotation(decls, tp)
+    callable_ref = ConstantRef(name)
+    decls._constants[name] = ConstantDecl(type_ref.to_just(), egg_name)
+    _add_default_rewrite(decls, callable_ref, type_ref, default_replacement, ruleset)
+    return decls, TypedExprDecl(type_ref.to_just(), CallDecl(callable_ref))
+
+
+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)
+        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(
+    decls: Declarations, ref: CallableRef, type_ref: TypeOrVarRef, default_rewrite: object, ruleset: Ruleset | None
+) -> None:
+    """
+    Adds a default rewrite for the callable, if the default rewrite is not None
+
+    Will add it to the ruleset if it is passed in, or add it to the default ruleset on the passed in decls if not.
+    """
+    if default_rewrite is None:
+        return
+    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
+        ruleset_decl = ruleset.__egg_ruleset__
+    else:
+        ruleset_decls = decls
+        ruleset_decl = decls.default_ruleset
+    ruleset_decl.rules.append(rewrite_decl)
+    ruleset_decls |= resolved_value
 
 
 def _last_param_variable(params: list[Parameter]) -> bool:
@@ -866,6 +970,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
@@ -908,81 +1014,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(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:
         """
@@ -991,10 +1035,11 @@ class EGraph(_BaseModule):
         action = let(name, expr)
         self.register(action)
         runtime_expr = to_runtime_expr(expr)
+        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))
             ),
         )
 
@@ -1016,14 +1061,15 @@ class EGraph(_BaseModule):
         self._add_decls(runtime_expr, schedule)
         egg_schedule = self._state.schedule_to_egg(schedule.schedule)
         typed_expr = runtime_expr.__egg_typed_expr__
-        egg_expr = self._state.expr_to_egg(typed_expr.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))
         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:
         """
@@ -1078,7 +1124,7 @@ class EGraph(_BaseModule):
         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: ...
@@ -1100,7 +1146,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
@@ -1118,12 +1164,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:
-        self._state.type_ref_to_egg(typed_expr.tp)
-        expr = self._state.expr_to_egg(typed_expr.expr)
-        self._egraph.run_program(bindings.ActionCommand(bindings.Extract(expr, bindings.Lit(bindings.Int(n)))))
+        expr = self._state.typed_expr_to_egg(typed_expr)
+        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"
@@ -1181,7 +1230,7 @@ class EGraph(_BaseModule):
         runtime_expr = to_runtime_expr(expr)
         self._add_decls(runtime_expr)
         typed_expr = runtime_expr.__egg_typed_expr__
-        egg_expr = self._state.expr_to_egg(typed_expr.expr)
+        egg_expr = self._state.typed_expr_to_egg(typed_expr)
         match typed_expr.tp:
             case JustTypeRef("i64"):
                 return self._egraph.eval_i64(egg_expr)
@@ -1195,42 +1244,110 @@ 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_e_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)
 
-        # import panel
+    def display(self, graphviz: bool = False, **kwargs: Unpack[GraphvizKwargs]) -> None:
+        """
+        Displays the e-graph.
 
-        # panel.extension("ipywidgets")
+        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
 
-        widget = self.saturate(performance=performance, **kwargs)
-        # panel.panel(widget).save(file)
+        from .visualizer_widget import VisualizerWidget
 
-        from ipywidgets.embed import embed_minimal_html
+        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()
 
-        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
+    def saturate(self, schedule: Schedule | 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
+
+        def to_json() -> str:
+            return self._serialize(**kwargs).to_json()
+
+        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:
@@ -1382,7 +1499,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
@@ -1622,16 +1740,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]:
@@ -1724,7 +1842,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__))
@@ -1893,3 +2011,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)