Upgrade to current Egglog, expose extraction controls, and retain Param-Eq demos by saulshanabrook · Pull Request #414 · egraphs-good/egglog-python

saulshanabrook · 2026-04-06T20:19:20Z

Summary

Upgrade the Rust extension and high-level Python APIs to current Egglog 3 and the matching experimental APIs. This brings the branch onto the built-in greedy-DAG and dynamic-cost infrastructure, exposes the new runtime and extraction controls in Python, and leaves the paused Param-Eq research as a maintained experimental module with bounded CI stress demos.

User-facing changes

Breaking: remove Map.rebuild(), Set.rebuild(), and Vec.rebuild(); Egglog 3 rebuilds container values in the backend.
Preserve Egglog 3 declarations, values, metadata, reports, and source syntax through the low-level bindings, including arbitrary-size BigInt, generic extraction, constructor/relation lookup, correctly scaled run durations, and parse-and-run source filenames.
Add generic Pair and Maybe, undefined-result catch, map folding, container lengths, numeric conversions, guarded f64 math, and f64.is_finite(). Maybe, catch, and map folding intentionally remain unsupported in proof mode.
Expand experimental Rational with public RationalLike conversions, reflected arithmetic, powers, min/max, comparisons, and partial-operation behavior consistent with surrounding numeric APIs.
Lower bodyful Python functions, methods, and constants as eager primitives by default while keeping explicit-ruleset bodies rewrite-backed. Preserve Python argument order for reverse_args callables and isolate higher-order probing from live rulesets.
Add persistent backoff schedules and RunReport.can_stop; expose safe naive and opt-in unsafe-seminaive rule evaluation.
Configure worker threads and global rule decomposition per EGraph, and allow rule(..., no_decomp=True) per rule. The process-global RAYON_NUM_THREADS path is removed.

Extraction and dynamic costs

Add extractor="tree" | "greedy-dag" to extraction APIs.
Add ordered heterogeneous multi-root extract_multiple, including empty per-root results, and destructive keep_best for table-backed callables. Multi-root and keep-best paths use the upstream dynamic cost model.
Consume experimental multi-extract as one structured aggregate with a shared TermDag and ordered per-root term-ID groups. Low-level bindings expose the known output through UserDefinedCommandOutput.as_multi_extract() -> MultiExtractOutput | None; the high-level sequence API converts it automatically and clones the shared DAG only once.
Use canonical set_cost tables consistently. Compatible raw tables and callable aliases share the backend table and survive freeze/replay; incompatible schemas fail clearly, and literal negative costs are rejected before reaching the backend.
Rename the general callback protocol to TreeCostModel, retaining CostModel as an alias, and add frozen additive DagCostModel values usable by tree or greedy-DAG extraction.
Keep custom Python models single-root. Read-only callback table lookups are safe for pre-registered tables keyed by callback-supplied values; derived same-graph keys are rejected while extraction holds the graph read-only.
Breaking: remove GreedyDagCost, GreedyDagCostModel, and greedy_dag_cost_model; use DagCostModel(..., extractor="greedy-dag").

Param-Eq handoff

egglog.exp.param_eq retains the binary and container representations plus a CLI. Three project-authored end-to-end cases run in normal pytest/CI, parse the extracted result, and compare it with the source at finite sample points. Constant-folding and domain regressions are covered directly.

experiments/param_eq remains the optional private-corpus runner with resource guards, provenance checks, and aggregate-only output. Private expressions and external archives are not redistributed. Both retained variants use ordinary persistent upstream backoff; the binary repeated-monomial case remains an explicitly reported iteration-limit stress case.

Dependencies

egraphs-good/egglog#1008 supplies rule-name round-tripping and safe user-output downcasting at ee58d8a537cec77e07595b1f0b577eee617f34b6, based on current Egglog main at integration.
egraphs-good/egglog-experimental#67 supplies the minimal reusable Param-Eq primitives plus structured shared-DAG multi-extraction and keep-best fixes at 2b4627a5806f8476bc34814ecf817e3b77c16f87.

Cargo patches the canonical Egglog source so the direct bindings and experimental dependency use one Rust type identity.

Validation

uv sync --reinstall-package egglog --all-extras --locked
focused downcast and multi-root extraction tests, including both extractors, empty roots, heterogeneous roots, and exact n=2 grouping
uv run pytest --benchmark-disable -q — 920 passed, 1 skipped, 3 expected xfails
make mypy
make stubtest
uv run ruff check .
uv run ruff format --check .
uv lock --check
cargo fmt --check
cargo check --locked
cargo test --locked --lib — 2 passed
cargo clippy --locked --all-targets
clean make docs — all 12 gallery examples executed successfully

codspeed-hq · 2026-04-06T20:38:30Z

Merging this PR will improve performance by 43.18%

⚠️

Different runtime environments detected

Some benchmarks with significant performance changes were compared across different runtime environments,
which may affect the accuracy of the results.

Open the report in CodSpeed to investigate

⚡ 2 improved benchmarks
✅ 12 untouched benchmarks

Performance Changes

	Mode	Benchmark	`BASE`	`HEAD`	Efficiency
⚡	Simulation	`test_jit[lda]`	7.4 s	5 s	+47.59%
⚡	WallTime	`test_jit[lda]`	7.9 s	5.7 s	+38.9%

Tip

Curious why performance improved? Comment @codspeedbot explain why performance improved on this PR, or directly use the CodSpeed MCP with your agent.

_{Comparing sr (ddae949) with main (6b2016e)}

+    def test_init(self):
+        class B(Expr):
+            def __init__(self, value: i64Like) -> None:
+                return B.wrap(value)  # type: ignore[return-value]  # noqa: PLE0101 - symbolic constructor body


+CallableDecl: TypeAlias = RelationDecl | ConstantDecl | FunctionCallableDecl
+
+
+def is_callable_decl_constructor(decls: Declarations, decl: CallableDecl) -> bool:


+    def action_to_egg(  # noqa: C901, PLR0911, PLR0912
+        self,
+        action: ActionDecl,
+        expr_to_let: bool = False,
+    ) -> bindings._Action | None:


-        return f"cost_table_{self.callable_ref_to_egg(ref)[0]}"
-
-    def fact_to_egg(self, fact: FactDecl) -> bindings._Fact:
+    def fact_to_egg(self, fact: FactDecl, *, expr_to_let: bool = False) -> bindings._Fact:


+    def _allocate_callable_egg_name(self, ref: CallableRef) -> str:
+        return self._allocate_name(self._generate_callable_egg_name_candidates(ref), self._backend_symbol_is_occupied)
+
+    def _generate_callable_egg_name_candidates(self, ref: CallableRef) -> tuple[str, ...]:


Migrate the bindings to Egglog 3, consolidate the general API and correctness fixes with tests and changelog coverage, and preserve Param-Eq as a reusable module and CLI, bounded CI stress cases, and an optional aggregate-only external harness. Pin the Egglog v3 compatibility fix and experimental primitives to immutable revisions, patching the core workspace source so direct and transitive dependencies share one Rust type identity.

saulshanabrook · 2026-08-31T20:57:45Z

+    @method(preserve=True)
+    def pick_key(self) -> T:
+        runtime_self = to_runtime_expr(self)
+        key_type, _value_type = runtime_self.__egg_typed_expr__.tp.args
+        maybe_type = RuntimeClass(
+            Thunk.value(Declarations.create(runtime_self, cast("HasDeclarations", Maybe))),
+            TypeRefWithVars(Ident.builtin("Maybe"), (key_type.to_var(),)),
+            _egg_has_params=True,
+        )
+        initial = cast("Maybe[T]", maybe_type.none())
+        return map_fold_kv(
+            lambda picked, key, _value: picked.match(lambda _: picked, cast("Maybe[T]", maybe_type.some(key))),
+            initial,
+            self,
+        ).unwrap()
+
+    @method(preserve=True)
+    def keys(self) -> Set[T]:
+        runtime_self = to_runtime_expr(self)
+        key_type, _value_type = runtime_self.__egg_typed_expr__.tp.args
+        set_type = RuntimeClass(
+            Thunk.value(Declarations.create(runtime_self, cast("HasDeclarations", Set))),
+            TypeRefWithVars(Ident.builtin("Set"), (key_type.to_var(),)),
+            _egg_has_params=True,
+        )
+        return map_fold_kv(lambda keys, key, _value: keys.insert(key), cast("Set[T]", set_type.empty()), self)


These should be normal expressions, not sure why they are like this... if they aren't builtins then we shouldnt include them here, instead they should just be inlined where we need them... as normal expression, we shouldnt use theprivate runtime expression APIs in public places like this

saulshanabrook · 2026-08-31T20:58:20Z

+def map_filter_kv(f: Callable[[T, V], Unit], xs: Map[T, V]) -> Map[T, V]:
+    runtime_xs = to_runtime_expr(xs)
+    map_type = RuntimeClass(
+        Thunk.value(Declarations.create(runtime_xs, cast("HasDeclarations", Map))),
+        runtime_xs.__egg_typed_expr__.tp.to_var(),
+        _egg_has_params=True,
+    )
+    return map_fold_kv(
+        lambda result, key, value: catch(lambda: f(key, value)).match(lambda _: result.insert(key, value), result),
+        cast("Map[T, V]", map_type.empty()),
+        xs,
+    )
+
+
+def map_map_values(f: Callable[[T, V], V2], xs: Map[T, V]) -> Map[T, V2]:
+    runtime_xs = to_runtime_expr(xs)
+    key_type, value_type = runtime_xs.__egg_typed_expr__.tp.args
+    probe_decls = runtime_xs.__egg_decls__.copy()
+    dummy_key = RuntimeExpr.__from_values__(probe_decls, TypedExprDecl(key_type, DummyDecl()))
+    dummy_value = RuntimeExpr.__from_values__(probe_decls, TypedExprDecl(value_type, DummyDecl()))
+    with set_current_ruleset(None):
+        transformed = cast("Callable[[RuntimeExpr, RuntimeExpr], object]", f)(dummy_key, dummy_value)
+    if not isinstance(transformed, RuntimeExpr):
+        raise TypeError(f"Map value transform must return an egglog expression, got {type(transformed)}")
+    output_type = transformed.__egg_typed_expr__.tp
+    map_type = RuntimeClass(
+        Thunk.value(Declarations.create(runtime_xs, transformed, cast("HasDeclarations", Map))),
+        TypeRefWithVars(Ident.builtin("Map"), (key_type.to_var(), output_type.to_var())),
+        _egg_has_params=True,
+    )
+    return map_fold_kv(
+        lambda result, key, value: result.insert(key, f(key, value)),
+        cast("Map[T, V2]", map_type.empty()),
+        xs,
+    )
+
+
+def map_merge_with(f: Callable[[V, V], V], left: Map[T, V], right: Map[T, V]) -> Map[T, V]:
+    return map_fold_kv(
+        lambda result, key, value: catch(lambda: result[key]).match(
+            lambda old: result.insert(key, f(old, value)), result.insert(key, value)
+        ),
+        left,
+        right,
+    )
+
+


these also shouldnt be here if they are not primitives... they should just be normal expressions where they are used

saulshanabrook · 2026-08-31T21:02:35Z

-    cost_callables: set[CallableRef] = field(default_factory=set)
+    # Cache of top-level expressions lowered with any available synthetic let
+    # references. Rules and rewrites must never read from this cache.
+    expr_to_let_egg_cache: dict[ExprDecl, bindings._Expr] = field(default_factory=dict)


Why is this different from expr_to_egg_cache? Are both required?

saulshanabrook · 2026-08-31T21:04:26Z

+                    # Use constructor declaration instead of constant b/c constants cannot be extracted
+                    # https://github.com/egraphs-good/egglog/issues/334


Constants dont exist anymore

saulshanabrook · 2026-08-31T21:06:55Z

+                from .builtins import ExprValueError  # noqa: PLC0415 - avoid a module import cycle
+                from .runtime import RuntimeExpr  # noqa: PLC0415 - avoid a module import cycle
+
+                if tp.ident == Ident.builtin("Map"):
+                    runtime_expr = RuntimeExpr.__from_values__(self.decls, decl)
+                    try:
+                        as_dict = cast("Map[BaseExpr, BaseExpr]", runtime_expr).value
+                    except ExprValueError:
+                        return self(expr, unwrap_lit=unwrap_lit, ruleset_ident=ruleset_ident, parens=parens), "expr"
+                    if unwrap_lit:
+                        items = ", ".join(
+                            f"{self(cast('RuntimeExpr', k).__egg_typed_expr__, unwrap_lit=True)}: {self(cast('RuntimeExpr', v).__egg_typed_expr__, unwrap_lit=True)}"
+                            for k, v in as_dict.items()
+                        )
+                        return f"{{{items}}}", "Map"
+                    map_str = f"{tp}.empty()"
+                    for key, value in as_dict.items():


We should not use the runtime to rebuild things here... why is this neccessary? Doesnt have to be canoncial, just print expressions as they are without getting value

saulshanabrook · 2026-08-31T21:10:19Z

+    assert '(let $__expr_0 (LetConflictNum_var "explicit"))' in egglog_string
+    assert '(let $__expr_1 (LetConflictNum_var "synthetic"))' in egglog_string


In these tests dont test for the exact string, this seems birttle... also all tests in here when they can should avoid using low level _ APIs and instead just use high level APIs that actually test behavior and not the low level behavior. remove any tests that aren't needed if we just are testing high level behavior

saulshanabrook · 2026-08-31T21:10:52Z

+def test_anonymous_combined_rulesets_use_deterministic_generated_names() -> None:
+    first = ruleset(name="combined_name_probe_first")
+    second = ruleset(name="combined_name_probe_second")
+    combined = unstable_combine_rulesets(first, second)
+    egraph = EGraph(save_egglog_string=True)


like this one too, make sure the tests are verifying a user seeable output, not low level implementaiton details

+            def __eq__(self, other: object) -> bool:
+                return False
+
+            def __lt__(self, other: CompareErrorCost) -> bool:


+                msg = "comparison failed"
+                raise RuntimeError(msg)
+
+            def __le__(self, other: CompareErrorCost) -> bool:


+                msg = "comparison failed"
+                raise RuntimeError(msg)
+
+            def __gt__(self, other: CompareErrorCost) -> bool:


+                msg = "comparison failed"
+                raise RuntimeError(msg)
+
+            def __ge__(self, other: CompareErrorCost) -> bool:


+ExtractionMode: TypeAlias = Literal["tree", "greedy-dag"]
+
+
+def _extractor_options(extractor: str) -> tuple[bindings._Expr, ...]:


saulshanabrook · 2026-09-03T16:09:20Z

+## Numeric Predicates and Exact Rationals
+
+The `f64.is_finite()` method returns a `Unit` fact when its value is neither
+infinite nor NaN. This makes it suitable for guarding rules that evaluate
+partial floating-point operations.
+
+The experimental exact `Rational` sort accepts `fractions.Fraction` and
+`i64Like` values in arithmetic, reflected arithmetic, powers, `min`/`max`, and
+ordering predicates. `RationalLike` is the corresponding public type alias.


this seems too specific for this level of doc

saulshanabrook · 2026-09-03T16:09:48Z

-mutate_egraph.register(x)
-mutate_egraph.run(10)
-mutate_egraph.check(eq(x).to(Int(10) + Int(1)))
+incremented = mutate_egraph.let("incremented", x)


why is the run removed here

saulshanabrook · 2026-09-03T16:12:22Z

+Dynamic row costs live in a canonical table named
+`cost_table_<egg-function-name>`. If a compatible, bodyless raw function with
+the same input sorts and `i64` output already has that name when the cost table
+is created, it is reused. An incompatible callable already occupying the name,
+or an incompatible overload that would map to the same canonical table, raises
+an error instead of making the cost table use a generated suffix, because the
+backend only consults the canonical name. Ordinary generated-name collision
+handling still applies to callables registered after the cost table.


this is like mostly low level details, lets not get into this as much

saulshanabrook · 2026-09-03T16:12:50Z

+
+At the low-level bindings layer, the experimental `multi-extract` command
+returns one {class}`egglog.bindings.UserDefinedOutput`.
+{meth}`egglog.bindings.UserDefinedCommandOutput.as_multi_extract` returns a
+{class}`egglog.bindings.MultiExtractOutput` whose `termdag` stores the shared
+term DAG and whose `terms` groups the variant term IDs in root order. It
+returns `None` for a different user-defined output. The high-level method
+performs this conversion automatically.


this is also too low level, we shouldnt talk about bindings here

saulshanabrook · 2026-09-03T16:14:10Z

+Configure worker threads per e-graph with `num_threads`. The default of `1`
+keeps execution serial; `0` uses the machine's available parallelism. You can
+change the setting later with `set_num_threads` and inspect it with
+`num_threads`. The bindings no longer read `RAYON_NUM_THREADS`.


dont reference old env variable

saulshanabrook · 2026-09-03T16:15:02Z

+_CallableMode: TypeAlias = Literal["function", "constructor", "eager", "rewrite"]
+
+
+def _normalize_callable_mode(  # noqa: C901, PLR0911, PLR0912


can you see if this could be reduced in loc somehow? its kinda verbose and hard to follow

saulshanabrook · 2026-09-03T16:16:29Z

+        callback_context = _COST_MODEL_CALLBACK_VALUES.get()
+        in_cost_model_callback = callback_context is not None and callback_context[0] is self
+        if in_cost_model_callback:
+            ref = typed_expr.expr.callable
+            table_is_registered = (
+                ref in self._state.callable_ref_to_egg_fn
+                if isinstance(typed_expr.expr, CallDecl)
+                else ref in self._state.cost_table_names
+            )
+            if not table_is_registered:
+                msg = "Tables queried by cost-model callbacks must be registered before extraction starts"
+                raise ValueError(msg)


this global seems messy why is it needed

saulshanabrook · 2026-09-03T16:17:23Z

+    """
+
+    marginal_cost: Callable[[EGraph, BaseExpr], DAG_COST]
+    identity: DAG_COST


do we really need identity?

saulshanabrook force-pushed the sr branch from 5dbbbaa to 9920fd1 Compare April 6, 2026 20:19

github-code-quality Bot found potential problems Apr 6, 2026

View reviewed changes

saulshanabrook force-pushed the sr branch from 009177d to 3b0ef4f Compare April 24, 2026 00:57

github-code-quality Bot found potential problems Apr 24, 2026

View reviewed changes

Comment thread python/egglog/exp/param_eq/pipeline.py Fixed

Comment thread python/egglog/builtins.py Fixed

Comment thread python/egglog/builtins.py Fixed

Comment thread python/egglog/builtins.py Fixed

Comment thread python/egglog/builtins.py Fixed

Comment thread python/egglog/builtins.py Fixed

saulshanabrook force-pushed the sr branch from 169f45b to 784b9ab Compare May 18, 2026 20:19

github-code-quality Bot found potential problems May 18, 2026

View reviewed changes

Comment thread python/egglog/builtins.py Fixed

Comment thread python/egglog/builtins.py Fixed

Comment thread python/egglog/builtins.py Fixed

Comment thread python/egglog/builtins.py Fixed

Comment thread python/egglog/builtins.py Fixed

Comment thread python/egglog/builtins.py Fixed

saulshanabrook changed the title Sr Upgrade the Python bindings to Egglog 3 and retain Param-Eq stress demos Aug 31, 2026

saulshanabrook force-pushed the sr branch from 22d9066 to ef9dc6e Compare August 31, 2026 20:32

github-code-quality Bot found potential problems Aug 31, 2026

View reviewed changes

saulshanabrook force-pushed the sr branch from ef9dc6e to a822f3f Compare August 31, 2026 20:51

saulshanabrook commented Aug 31, 2026

View reviewed changes

Address Egglog 3 review feedback

2495210

github-code-quality Bot found potential problems Aug 31, 2026

View reviewed changes

Comment thread python/tests/param_eq/test_domain.py Fixed

Comment thread python/egglog/exp/param_eq/domain.py Fixed

Update Egglog dependencies and extraction APIs

6bdaad7

saulshanabrook changed the title ~~Upgrade the Python bindings to Egglog 3 and retain Param-Eq stress demos~~ Upgrade to current Egglog, expose extraction controls, and retain Param-Eq demos Sep 1, 2026

github-code-quality Bot found potential problems Sep 1, 2026

View reviewed changes

Expose structured multi-extract outputs

ddae949

saulshanabrook commented Sep 3, 2026

View reviewed changes

		CallableDecl: TypeAlias = RelationDecl \| ConstantDecl \| FunctionCallableDecl


		def is_callable_decl_constructor(decls: Declarations, decl: CallableDecl) -> bool:

		# Use constructor declaration instead of constant b/c constants cannot be extracted
		# https://github.com/egraphs-good/egglog/issues/334

		assert '(let $__expr_0 (LetConflictNum_var "explicit"))' in egglog_string
		assert '(let $__expr_1 (LetConflictNum_var "synthetic"))' in egglog_string

		ExtractionMode: TypeAlias = Literal["tree", "greedy-dag"]


		def _extractor_options(extractor: str) -> tuple[bindings._Expr, ...]:

		_CallableMode: TypeAlias = Literal["function", "constructor", "eager", "rewrite"]


		def _normalize_callable_mode( # noqa: C901, PLR0911, PLR0912

Uh oh!

Conversation

saulshanabrook commented Apr 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

User-facing changes

Extraction and dynamic costs

Param-Eq handoff

Dependencies

Validation

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

codspeed-hq Bot commented Apr 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Merging this PR will improve performance by 43.18%

Performance Changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

saulshanabrook commented Apr 6, 2026 •

edited

Loading

codspeed-hq Bot commented Apr 6, 2026 •

edited

Loading