Upgrade to current Egglog, expose extraction controls, and retain Param-Eq demos - #414
Draft
saulshanabrook wants to merge 4 commits into
Draft
Upgrade to current Egglog, expose extraction controls, and retain Param-Eq demos#414saulshanabrook wants to merge 4 commits into
saulshanabrook wants to merge 4 commits into
Conversation
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
Merging this PR will improve performance by 43.18%
|
| 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)
saulshanabrook
changed the title
| 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: |
Comment on lines
+550
to
+554
| 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
commented
Aug 31, 2026
Comment on lines
+582
to
+607
| @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) |
Member
Author
There was a problem hiding this comment.
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
commented
Aug 31, 2026
Comment on lines
+632
to
+678
| 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, | ||
| ) | ||
|
|
||
|
|
Member
Author
There was a problem hiding this comment.
these also shouldnt be here if they are not primitives... they should just be normal expressions where they are used
saulshanabrook
commented
Aug 31, 2026
| 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) |
Member
Author
There was a problem hiding this comment.
Why is this different from expr_to_egg_cache? Are both required?
saulshanabrook
commented
Aug 31, 2026
Comment on lines
+666
to
+667
| # Use constructor declaration instead of constant b/c constants cannot be extracted | ||
| # https://github.com/egraphs-good/egglog/issues/334 |
Member
Author
There was a problem hiding this comment.
Constants dont exist anymore
saulshanabrook
commented
Aug 31, 2026
Comment on lines
+430
to
+446
| 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(): |
Member
Author
There was a problem hiding this comment.
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
commented
Aug 31, 2026
Comment on lines
+181
to
+182
| assert '(let $__expr_0 (LetConflictNum_var "explicit"))' in egglog_string | ||
| assert '(let $__expr_1 (LetConflictNum_var "synthetic"))' in egglog_string |
Member
Author
There was a problem hiding this comment.
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
commented
Aug 31, 2026
Comment on lines
+569
to
+573
| 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) |
Member
Author
There was a problem hiding this comment.
like this one too, make sure the tests are verifying a user seeable output, not low level implementaiton details
saulshanabrook
changed the title
| 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, ...]: |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
Map.rebuild(),Set.rebuild(), andVec.rebuild(); Egglog 3 rebuilds container values in the backend.BigInt, generic extraction, constructor/relation lookup, correctly scaled run durations, and parse-and-run source filenames.PairandMaybe, undefined-resultcatch, map folding, container lengths, numeric conversions, guardedf64math, andf64.is_finite().Maybe,catch, and map folding intentionally remain unsupported in proof mode.Rationalwith publicRationalLikeconversions, reflected arithmetic, powers,min/max, comparisons, and partial-operation behavior consistent with surrounding numeric APIs.reverse_argscallables and isolate higher-order probing from live rulesets.RunReport.can_stop; expose safenaiveand opt-inunsafe-seminaiverule evaluation.EGraph, and allowrule(..., no_decomp=True)per rule. The process-globalRAYON_NUM_THREADSpath is removed.Extraction and dynamic costs
extractor="tree" | "greedy-dag"to extraction APIs.extract_multiple, including empty per-root results, and destructivekeep_bestfor table-backed callables. Multi-root and keep-best paths use the upstream dynamic cost model.multi-extractas one structured aggregate with a sharedTermDagand ordered per-root term-ID groups. Low-level bindings expose the known output throughUserDefinedCommandOutput.as_multi_extract() -> MultiExtractOutput | None; the high-level sequence API converts it automatically and clones the shared DAG only once.set_costtables 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.TreeCostModel, retainingCostModelas an alias, and add frozen additiveDagCostModelvalues usable by tree or greedy-DAG extraction.GreedyDagCost,GreedyDagCostModel, andgreedy_dag_cost_model; useDagCostModel(..., extractor="greedy-dag").Param-Eq handoff
egglog.exp.param_eqretains 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_eqremains 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
ee58d8a537cec77e07595b1f0b577eee617f34b6, based on current Egglogmainat integration.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 --lockedn=2groupinguv run pytest --benchmark-disable -q— 920 passed, 1 skipped, 3 expected xfailsmake mypymake stubtestuv run ruff check .uv run ruff format --check .uv lock --checkcargo fmt --checkcargo check --lockedcargo test --locked --lib— 2 passedcargo clippy --locked --all-targetsmake docs— all 12 gallery examples executed successfully