Respond to review notes and try to fix codspeed crash

saulshanabrook · saulshanabrook · commit 4e0dfb097184 · 2026-03-16T16:30:14.000-07:00
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -67,7 +67,7 @@ jobs:
       - run: |
           export UV_PROJECT_ENVIRONMENT="${pythonLocation}"
           uv sync --extra test --locked
-      - uses: CodSpeedHQ/action@v4.3.1
+      - uses: CodSpeedHQ/action@v4.11.1
         with:
           token: ${{ secrets.CODSPEED_TOKEN }}
           # allow updating snapshots due to indeterministic benchmarks
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -36,7 +36,6 @@ ordered-float = "5"
 uuid = { version = "1.18", features = ["v4"] }
 rayon = "1.11"
 base64 = "0.22.1"
-tempfile = "3.24.0"
 
 # Use patched version of egglog in experimental
 [patch.'https://github.com/egraphs-good/egglog']
diff --git a/docs/explanation/2026_02_containers.md b/docs/explanation/2026_02_containers.md
@@ -435,7 +435,7 @@ GradientBending_distributed = gradient_egraph.extract(GradientBending)
 FunctionBending_distributed
 ```
 
-We now have an expression that is mainly a sum of products, a multivariate polyonimal.
+We now have an expression that is mainly a sum of products, a multivariate polynomial.
 
 For some sense of their size, the `FunctionBending` has initial cost of 401 and the `GradientBending` has 20,570.
 This cost is produced by the Egglog extractor, corresponding roughly to one node per op like `*` and one per variable as a tree.
diff --git a/python/egglog/builtins.py b/python/egglog/builtins.py
@@ -1207,7 +1207,7 @@ def dict(cls, *keys_and_values: object) -> PyObject: ...
 
 
 @function(builtin=True, egg_fn="py-eval")
-def py_eval(code: StringLike, globals: object = PyObject.dict(), locals: object = PyObject.dict()) -> PyObject: ...
+def py_eval(code: StringLike, globals_: object = PyObject.dict(), locals_: object = PyObject.dict()) -> PyObject: ...
 
 
 class PyObjectFunction(Protocol):
@@ -1226,7 +1226,7 @@ def py_eval_fn(fn: Callable) -> PyObjectFunction:
 
 
 @function(builtin=True, egg_fn="py-exec")
-def py_exec(code: StringLike, globals: object = PyObject.dict(), locals: object = PyObject.dict()) -> PyObject:
+def py_exec(code: StringLike, globals_: object = PyObject.dict(), locals_: object = PyObject.dict()) -> PyObject:
     """
     Copies the locals, execs the Python code, and returns the locals with any updates.
     """
diff --git a/python/egglog/conversion.py b/python/egglog/conversion.py
@@ -228,9 +228,9 @@ def resolve_literal(
     if isinstance(arg, RuntimeExpr) and tp.matches_just({}, arg.__egg_typed_expr__.tp):
         return arg
     tp_just = tp.to_just()
-    arg_type = resolve_type(arg)
     if arg is DUMMY_VALUE:
         return RuntimeExpr.__from_values__(decls(), TypedExprDecl(tp_just, DummyDecl()))
+    arg_type = resolve_type(arg)
     if (conversion := _lookup_conversion(arg_type, tp_just)) is not None:
         with with_type_args(tp_just.args, decls):
             return conversion[1](arg)
diff --git a/python/egglog/declarations.py b/python/egglog/declarations.py
@@ -377,11 +377,11 @@ class EGraphDecl:
     def __hash__(self) -> int:
         return hash((
             type(self),
-            tuple(self.let_bindings.items()),
-            tuple((value, tp, exprs) for value, (tp, exprs) in self.e_classes.items()),
-            tuple(self.sets.items()),
+            frozenset(self.let_bindings.items()),
+            frozenset((value, tp, exprs) for value, (tp, exprs) in self.e_classes.items()),
+            frozenset(self.sets.items()),
             self.expr_actions,
-            tuple(self.costs.items()),
+            frozenset(self.costs.items()),
             self.subsumed,
         ))
 
@@ -509,7 +509,7 @@ def to_grounded(expr: ExprDecl) -> ExprDecl:
         )
         actions.extend(ChangeDecl(tp, cast("CallDecl", to_grounded(call)), "subsume") for tp, call in self.subsumed)
 
-        # Now add any remaining call    s that weren't part of any other actions
+        # Now add any remaining calls that weren't part of any other actions
         actions.extend(
             ExprActionDecl(TypedExprDecl(tp, to_grounded(expr)))
             for (tp, expr) in single_e_class_calls
diff --git a/python/egglog/deconstruct.py b/python/egglog/deconstruct.py
@@ -189,7 +189,7 @@ def _deconstruct_call_decl(
         ), arg_exprs
     egg_bound = (
         JustTypeRef(call.callable.ident, call.bound_tp_params)
-        if isinstance(call.callable, (ClassMethodRef, MethodRef, InitRef)) and call.bound_tp_params
+        if isinstance(call.callable, (ClassMethodRef, MethodRef)) and call.bound_tp_params
         else None
     )
 
diff --git a/python/egglog/egraph_state.py b/python/egglog/egraph_state.py
@@ -83,7 +83,7 @@ class EGraphState:
 
     def copy(self) -> EGraphState:
         """
-        Returns a copy of the state. Th egraph reference is kept the same. Used for pushing/popping.
+        Returns a copy of the state. The egraph reference is kept the same. Used for pushing/popping.
         """
         return EGraphState(
             egraph=self.egraph,
@@ -802,7 +802,7 @@ def value_to_expr(self, tp: JustTypeRef, value: bindings.Value) -> ExprDecl:  #
     def _unstable_fn_value_to_expr(
         self, name: str, partial_args: list[bindings.Value], return_tp: JustTypeRef, _arg_types: list[JustTypeRef]
     ) -> PartialCallDecl:
-        # Similar to FromEggState::from_call but accepts partial list of args and returns in values
+        # Similar to FromEggState::from_call but reconstructs a partial application from serialized values.
         # Find first callable ref whose return type matches and fill in arg types.
         for callable_ref in self.egg_fn_to_callable_refs[name]:
             signature = self.__egg_decls__.get_callable_decl(callable_ref).signature
diff --git a/python/egglog/exp/array_api_jit.py b/python/egglog/exp/array_api_jit.py
@@ -37,7 +37,11 @@ def jit(
     try:
         return cast("X", egraph.extract(fn_program.as_py_object).value)
     except bindings.EggSmolError as e:
-        e.add_note(f"Failed to get py object from {egraph.extract(fn_program)}")
+        try:
+            debug_program = egraph.extract(fn_program)
+        except bindings.EggSmolError:
+            debug_program = fn_program
+        e.add_note(f"Failed to get py object from {debug_program}")
         raise
 
 
diff --git a/python/egglog/exp/array_api_program_gen.py b/python/egglog/exp/array_api_program_gen.py
@@ -572,7 +572,7 @@ def vec_recursive_value_program(x: Vec[RecursiveValue]) -> Program: ...
 
 
 @array_api_program_gen_ruleset.register
-def _vec_recursive_value_program(v: Value, vv: Vec[RecursiveValue]):
+def _vec_recursive_value_program(vv: Vec[RecursiveValue]):
     yield rewrite(vec_recursive_value_program(Vec[RecursiveValue].empty())).to(Program(""))
     yield rewrite(vec_recursive_value_program(vv)).to(
         recursive_value_program(vv[0]) + ", " + vec_recursive_value_program(vv.remove(0)),
diff --git a/python/egglog/runtime.py b/python/egglog/runtime.py
@@ -306,9 +306,11 @@ def __call__(self, *args: object, **kwargs: object) -> RuntimeExpr | None:
             return_tp = res_typed_expr.tp
             assert isinstance(call, CallDecl), "partial function must be a call"
             # Clip off the remaining arguments
-            n_args = len(partial_args)
-            value = PartialCallDecl(replace(call, args=call.args[:n_args]))
-            remaining_arg_types = [a.tp for a in call.args[n_args:]]
+            captured_args = len(partial_args) + int(
+                isinstance(fn_arg, RuntimeFunction) and isinstance(fn_arg.__egg_bound__, RuntimeExpr)
+            )
+            value = PartialCallDecl(replace(call, args=call.args[:captured_args]))
+            remaining_arg_types = [a.tp for a in call.args[captured_args:]]
             type_ref = JustTypeRef(Ident.builtin("UnstableFn"), (return_tp, *remaining_arg_types))
             return RuntimeExpr.__from_values__(Declarations.create(self, res), TypedExprDecl(type_ref, value))
 
diff --git a/python/egglog/visualizer_widget.py b/python/egglog/visualizer_widget.py
@@ -1,3 +1,5 @@
+import base64
+import json
 import pathlib
 import webbrowser
 
@@ -24,21 +26,15 @@ class VisualizerWidget(anywidget.AnyWidget):
 
     def display_or_open(self) -> None:
         """
-        Displays the widget if we are in a Jupyter environment, otherwise saves it to a file and opens it.
+        Displays the widget if we are in a Jupyter environment, otherwise opens a standalone HTML page.
         """
         if IN_IPYTHON:
             display(self)
             return
-        # 1. Create a temporary html file that will stay open after close
-        # 2. Write the widget to it with embed_minimal_html
-        # 3. Open the file using the open function from graphviz
-        file = pathlib.Path.cwd() / "tmp.html"
-        file.write_text(HTML.replace("MAGIC_STRING", str(self.egraphs)))
-        # https://github.com/manzt/anywidget/issues/339#issuecomment-1755654547
-        # embed_minimal_html(file, views=[self], drop_defaults=False)
-        # Instead of embedding widget, just save
-        print("Visualizer widget saved to", file)
-        webbrowser.open(file.as_uri())
+        payload = json.dumps(self.egraphs).replace("</", "<\\/")
+        html = HTML.replace("MAGIC_STRING", payload)
+        data_uri = "data:text/html;base64," + base64.b64encode(html.encode()).decode()
+        webbrowser.open(data_uri)
 
 
 HTML = """
diff --git a/python/tests/test_array_api.py b/python/tests/test_array_api.py
@@ -456,7 +456,7 @@ def test_run_lda(fn_thunk, benchmark):
 
 
 @pytest.mark.parametrize(
-    ("input", "expected"),
+    ("input_expr", "expected"),
     [
         pytest.param(x * x, x**2, id="exp"),
         pytest.param(x * y + x * z, x * (y + z), id="factor"),
@@ -466,9 +466,9 @@ def test_run_lda(fn_thunk, benchmark):
         ),
     ],
 )
-def test_polynomial_factoring(input: Value, expected: Value):
+def test_polynomial_factoring(input_expr: Value, expected: Value):
     egraph = EGraph()
-    x = egraph.let("x", input)
+    x = egraph.let("x", input_expr)
     egraph.run(polynomial_schedule)
     equiv_expr = egraph.extract(x)
     # Normalized them both so that we don't have to worry about term order.
diff --git a/python/tests/test_unstable_fn.py b/python/tests/test_unstable_fn.py
@@ -53,6 +53,15 @@ def __mul__(self, x: MathLike) -> MathList: ...
 converter(type(None), MathList, lambda _: MathList.NIL)
 
 
+class Pair(Expr):
+    def __init__(self, x: i64Like) -> None: ...
+
+    def add(self, y: i64Like) -> Pair: ...
+
+
+converter(i64, Pair, Pair)
+
+
 @function
 def square(x: MathLike) -> Math: ...
 
@@ -74,6 +83,12 @@ def test_string_fn_partial():
     assert str(UnstableFn(Math.__mul__, Math(2))) == "partial(Math.__mul__, Math(2))"
 
 
+def test_bound_runtime_function_partial():
+    pair = Pair(2)
+    assert expr_parts(UnstableFn(pair.add)) == expr_parts(UnstableFn(Pair.add, pair))
+    assert expr_parts(UnstableFn(pair.add, 3)) == expr_parts(UnstableFn(Pair.add, pair, 3))
+
+
 @ruleset
 def map_ruleset(f: MathFn, x: Math, xs: MathList):
     yield rewrite(MathList.NIL.map(f)).to(MathList.NIL)
diff --git a/src/egraph.rs b/src/egraph.rs
@@ -66,9 +66,6 @@ impl EGraph {
             let cmd_string = cmd.to_string();
             cmds_str = cmds_str + &cmd_string + "\n";
         }
-        if let Some(cmds) = &mut self.cmds {
-            cmds.push_str(&cmds_str);
-        }
         info!("Running commands:\n{}", cmds_str);
         let res = py.detach(|| self.egraph.run_program(commands));
         if let Some(err) = PyErr::take(py) {
@@ -77,6 +74,9 @@ impl EGraph {
         match res {
             Err(e) => Err(WrappedError::Egglog(e)),
             Ok(outputs) => {
+                if let Some(cmds) = &mut self.cmds {
+                    cmds.push_str(&cmds_str);
+                }
                 Ok(outputs.into_iter().map(|o| o.into()).collect())
             }
         }
diff --git a/src/freeze.rs b/src/freeze.rs
@@ -1,4 +1,4 @@
-// Freeze an egglog, turning it an immutable structure that can be printed, serialized, or added back to an e-graph.
+// Freeze an egglog, turning it into an immutable structure that can be printed, serialized, or added back to an e-graph.
 
 use egglog::EGraph;
 use indexmap::IndexMap;
@@ -30,7 +30,7 @@ pub struct FrozenEGraph {
 }
 
 impl FrozenEGraph {
-    /// Convert this frozen e-graph into a list of egglog commands that can reconstruct it
+    /// Convert a live `EGraph` into an immutable `FrozenEGraph` snapshot.
     pub fn from_egraph(egraph: &EGraph) -> FrozenEGraph {
         let mut functions = IndexMap::new();
         for fname in egraph.get_function_names() {
diff --git a/uv.lock b/uv.lock

Original file line number	Diff line number	Diff line change
`@@ -189,7 +189,7 @@ def _deconstruct_call_decl(`
`189`	`189`	`), arg_exprs`
`190`	`190`	`egg_bound = (`
`191`	`191`	`JustTypeRef(call.callable.ident, call.bound_tp_params)`
`192`		`- if isinstance(call.callable, (ClassMethodRef, MethodRef, InitRef)) and call.bound_tp_params`
	`192`	`+ if isinstance(call.callable, (ClassMethodRef, MethodRef)) and call.bound_tp_params`
`193`	`193`	`else None`
`194`	`194`	`)`
`195`	`195`