From a822f3fe4e7af9562e8b02a1e4e1aa24bcb85283 Mon Sep 17 00:00:00 2001 From: Saul Shanabrook Date: Sat, 22 Aug 2026 13:07:18 -0400 Subject: [PATCH 1/6] Prepare Egglog 3 migration and Param-Eq handoff 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. --- .github/workflows/CI.yml | 9 +- .readthedocs.yaml | 2 +- Cargo.lock | 71 +- Cargo.toml | 38 +- docs/changelog.md | 33 + docs/conf.py | 1 + docs/explanation/2026_02_containers.md | 34 +- .../parsing-and-running-program-strings.md | 26 + docs/reference/egglog-translation.md | 134 +- docs/reference/python-integration.md | 122 +- experiments/param_eq/Makefile | 25 + experiments/param_eq/NOTES.md | 116 ++ experiments/param_eq/README.md | 109 ++ experiments/param_eq/__init__.py | 1 + experiments/param_eq/aggregate.py | 476 ++++++ experiments/param_eq/corpus.py | 175 ++ experiments/param_eq/resource_guard.py | 144 ++ experiments/param_eq/results/manifest.toml | 6 + .../param_eq/results/paper-replication.csv | 1 + experiments/param_eq/results/raw/.gitignore | 2 + .../results/representation-comparison.csv | 1 + experiments/param_eq/run.py | 720 ++++++++ pyproject.toml | 12 +- python/egglog/bindings.pyi | 115 +- python/egglog/builtins.py | 258 ++- python/egglog/conversion.py | 5 +- python/egglog/declarations.py | 92 +- python/egglog/deconstruct.py | 58 +- python/egglog/egraph.py | 619 +++++-- python/egglog/egraph_state.py | 725 ++++++-- python/egglog/exp/array_api.py | 2 + python/egglog/exp/array_api_program_gen.py | 6 +- python/egglog/exp/param_eq/__init__.py | 48 + python/egglog/exp/param_eq/__main__.py | 31 + python/egglog/exp/param_eq/cases.py | 32 + python/egglog/exp/param_eq/domain.py | 635 +++++++ python/egglog/exp/param_eq/pipeline.py | 532 ++++++ python/egglog/ipython_magic.py | 2 +- python/egglog/pretty.py | 190 ++- python/egglog/run_report.py | 3 + python/egglog/runtime.py | 11 +- python/egglog/type_constraint_solver.py | 5 +- .../test_array_api/test_jit[lda][expr].py | 11 +- .../test_factor_multisets[code].py | 16 + python/tests/param_eq/__init__.py | 1 + python/tests/param_eq/conftest.py | 9 + python/tests/param_eq/evaluation.py | 53 + python/tests/param_eq/test_domain.py | 122 ++ python/tests/param_eq/test_pipeline.py | 129 ++ .../tests/param_eq/test_research_harness.py | 608 +++++++ python/tests/param_eq/test_resource_guard.py | 82 + python/tests/param_eq/test_semantics.py | 29 + python/tests/test_bindings.py | 348 +++- python/tests/test_high_level.py | 1481 ++++++++++++++++- python/tests/test_polynomials.py | 24 + python/tests/test_pretty.py | 60 + python/tests/test_run_report.py | 50 + python/tests/test_unstable_fn.py | 42 +- rust-toolchain.toml | 2 +- src/conversions.rs | 321 +++- src/egraph.rs | 153 +- src/extract.rs | 8 +- src/freeze.rs | 43 +- src/lib.rs | 1 + src/utils.rs | 4 +- test-data/unit/check-high-level.test | 29 + uv.lock | 191 ++- 67 files changed, 8634 insertions(+), 810 deletions(-) create mode 100644 experiments/param_eq/Makefile create mode 100644 experiments/param_eq/NOTES.md create mode 100644 experiments/param_eq/README.md create mode 100644 experiments/param_eq/__init__.py create mode 100644 experiments/param_eq/aggregate.py create mode 100644 experiments/param_eq/corpus.py create mode 100644 experiments/param_eq/resource_guard.py create mode 100644 experiments/param_eq/results/manifest.toml create mode 100644 experiments/param_eq/results/paper-replication.csv create mode 100644 experiments/param_eq/results/raw/.gitignore create mode 100644 experiments/param_eq/results/representation-comparison.csv create mode 100644 experiments/param_eq/run.py create mode 100644 python/egglog/exp/param_eq/__init__.py create mode 100644 python/egglog/exp/param_eq/__main__.py create mode 100644 python/egglog/exp/param_eq/cases.py create mode 100644 python/egglog/exp/param_eq/domain.py create mode 100644 python/egglog/exp/param_eq/pipeline.py create mode 100644 python/tests/__snapshots__/test_polynomials/test_factor_multisets[code].py create mode 100644 python/tests/param_eq/__init__.py create mode 100644 python/tests/param_eq/conftest.py create mode 100644 python/tests/param_eq/evaluation.py create mode 100644 python/tests/param_eq/test_domain.py create mode 100644 python/tests/param_eq/test_pipeline.py create mode 100644 python/tests/param_eq/test_research_harness.py create mode 100644 python/tests/param_eq/test_resource_guard.py create mode 100644 python/tests/param_eq/test_semantics.py create mode 100644 python/tests/test_polynomials.py diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index b1435c55..73156df3 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -29,7 +29,7 @@ jobs: with: enable-cache: true python-version: ${{ matrix.py }} - - uses: dtolnay/rust-toolchain@1.79.0 + - uses: dtolnay/rust-toolchain@1.91.0 - uses: Swatinem/rust-cache@v2 - run: uv sync --extra test --locked - run: uv run pytest --benchmark-disable -vvv --durations=10 @@ -41,11 +41,12 @@ jobs: - uses: astral-sh/setup-uv@v7 with: enable-cache: true - - uses: dtolnay/rust-toolchain@1.79.0 + - uses: dtolnay/rust-toolchain@1.91.0 - uses: Swatinem/rust-cache@v2 - run: uv sync --extra test --locked - run: make mypy - run: make stubtest + - run: cargo test --locked --lib benchmark: runs-on: ${{ matrix.runner }} @@ -64,7 +65,7 @@ jobs: with: enable-cache: true python-version: "3.12" - - uses: dtolnay/rust-toolchain@1.79.0 + - uses: dtolnay/rust-toolchain@1.91.0 - uses: Swatinem/rust-cache@v2 - run: | export UV_PROJECT_ENVIRONMENT="${pythonLocation}" @@ -83,7 +84,7 @@ jobs: - uses: astral-sh/setup-uv@v7 with: enable-cache: true - - uses: dtolnay/rust-toolchain@1.79.0 + - uses: dtolnay/rust-toolchain@1.91.0 - uses: Swatinem/rust-cache@v2 - name: Install graphviz run: | diff --git a/.readthedocs.yaml b/.readthedocs.yaml index 5a971056..df209af6 100644 --- a/.readthedocs.yaml +++ b/.readthedocs.yaml @@ -12,7 +12,7 @@ build: python: "3.12" # # You can also specify other tool versions: # # nodejs: "16" - rust: "1.78" + rust: "1.91" # golang: "1.17" apt_packages: - graphviz diff --git a/Cargo.lock b/Cargo.lock index ab79480e..19b90e3a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -362,8 +362,8 @@ checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" [[package]] name = "egglog" -version = "2.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=2e5657b#2e5657bbb2c1a90fba31002da61381815f891b6f" +version = "3.0.0" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=0476e56345f52e6fcba4b0c93f5d6d60a2b9e318#0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" dependencies = [ "csv", "dyn-clone", @@ -374,23 +374,24 @@ dependencies = [ "egglog-numeric-id", "egglog-reports", "egraph-serialize", + "enum-map", "hashbrown 0.16.1", "im-rc", "indexmap", "log", "num", "ordered-float", - "rayon", "rustc-hash", "serde_json", + "smallvec", "thiserror", "web-time", ] [[package]] name = "egglog-add-primitive" -version = "2.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=2e5657b#2e5657bbb2c1a90fba31002da61381815f891b6f" +version = "3.0.0" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=0476e56345f52e6fcba4b0c93f5d6d60a2b9e318#0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" dependencies = [ "quote", "syn 2.0.117", @@ -398,19 +399,20 @@ dependencies = [ [[package]] name = "egglog-ast" -version = "2.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=2e5657b#2e5657bbb2c1a90fba31002da61381815f891b6f" +version = "3.0.0" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=0476e56345f52e6fcba4b0c93f5d6d60a2b9e318#0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" dependencies = [ "ordered-float", ] [[package]] name = "egglog-bridge" -version = "2.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=2e5657b#2e5657bbb2c1a90fba31002da61381815f891b6f" +version = "3.0.0" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=0476e56345f52e6fcba4b0c93f5d6d60a2b9e318#0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" dependencies = [ "anyhow", "dyn-clone", + "egglog-concurrency", "egglog-core-relations", "egglog-numeric-id", "egglog-reports", @@ -421,7 +423,6 @@ dependencies = [ "num-rational", "once_cell", "ordered-float", - "rayon", "smallvec", "thiserror", "web-time", @@ -429,20 +430,20 @@ dependencies = [ [[package]] name = "egglog-concurrency" -version = "2.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=2e5657b#2e5657bbb2c1a90fba31002da61381815f891b6f" +version = "3.0.0" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=0476e56345f52e6fcba4b0c93f5d6d60a2b9e318#0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" dependencies = [ "arc-swap", "bumpalo", + "crossbeam", "egglog-numeric-id", - "rayon", "smallvec", ] [[package]] name = "egglog-core-relations" -version = "2.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=2e5657b#2e5657bbb2c1a90fba31002da61381815f891b6f" +version = "3.0.0" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=0476e56345f52e6fcba4b0c93f5d6d60a2b9e318#0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" dependencies = [ "anyhow", "bumpalo", @@ -461,7 +462,6 @@ dependencies = [ "num", "once_cell", "rand 0.9.4", - "rayon", "rustc-hash", "smallvec", "thiserror", @@ -470,8 +470,8 @@ dependencies = [ [[package]] name = "egglog-experimental" -version = "0.1.0" -source = "git+https://github.com/egraphs-good/egglog-experimental?branch=main#d443c29a58fd3affe5fcf92e6878168e89273752" +version = "3.0.0" +source = "git+https://github.com/egraphs-good/egglog-experimental.git?rev=78cdfa543de9e282c86f446f69943965119c2afe#78cdfa543de9e282c86f446f69943965119c2afe" dependencies = [ "egglog", "egglog-ast", @@ -483,16 +483,13 @@ dependencies = [ [[package]] name = "egglog-numeric-id" -version = "2.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=2e5657b#2e5657bbb2c1a90fba31002da61381815f891b6f" -dependencies = [ - "rayon", -] +version = "3.0.0" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=0476e56345f52e6fcba4b0c93f5d6d60a2b9e318#0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" [[package]] name = "egglog-reports" -version = "2.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=2e5657b#2e5657bbb2c1a90fba31002da61381815f891b6f" +version = "3.0.0" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=0476e56345f52e6fcba4b0c93f5d6d60a2b9e318#0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" dependencies = [ "clap", "hashbrown 0.16.1", @@ -505,8 +502,8 @@ dependencies = [ [[package]] name = "egglog-union-find" -version = "2.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=2e5657b#2e5657bbb2c1a90fba31002da61381815f891b6f" +version = "3.0.0" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=0476e56345f52e6fcba4b0c93f5d6d60a2b9e318#0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" dependencies = [ "crossbeam", "egglog-concurrency", @@ -565,6 +562,26 @@ version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" +[[package]] +name = "enum-map" +version = "2.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6866f3bfdf8207509a033af1a75a7b08abda06bbaaeae6669323fd5a097df2e9" +dependencies = [ + "enum-map-derive", +] + +[[package]] +name = "enum-map-derive" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f282cfdfe92516eb26c2af8589c274c7c17681f5ecc03c18255fe741c6aa64eb" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "equivalent" version = "1.0.2" diff --git a/Cargo.toml b/Cargo.toml index 1da061e7..82ed1f8b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,17 +18,12 @@ opentelemetry = "0.28" opentelemetry-otlp = { version = "0.28", features = ["http-proto", "reqwest-blocking-client", "trace"] } opentelemetry-stdout = { version = "0.28", features = ["trace"] } opentelemetry_sdk = "0.28" -# egglog = { path = "../egg-smol", default-features = false } -# egglog-bridge = { path = "../egg-smol/egglog-bridge" } -# egglog-core-relations = { path = "../egg-smol/core-relations" } -# egglog-ast = { path = "../egg-smol/egglog-ast" } -# egglog-reports = { path = "../egg-smol/egglog-reports" } -egglog = { git = "https://github.com/egraphs-good/egglog.git", rev = "2e5657b", default-features = false } -egglog-ast = { git = "https://github.com/egraphs-good/egglog.git", rev = "2e5657b" } -egglog-core-relations = { git = "https://github.com/egraphs-good/egglog.git", rev = "2e5657b" } -egglog-reports = { git = "https://github.com/egraphs-good/egglog.git", rev = "2e5657b" } -egglog-bridge = { git = "https://github.com/egraphs-good/egglog.git", rev = "2e5657b" } -egglog-experimental = { git = "https://github.com/egraphs-good/egglog-experimental", branch = "main", default-features = false } +egglog = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "0476e56345f52e6fcba4b0c93f5d6d60a2b9e318", default-features = false } +egglog-ast = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" } +egglog-core-relations = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" } +egglog-reports = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" } +egglog-bridge = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" } +egglog-experimental = { git = "https://github.com/egraphs-good/egglog-experimental.git", rev = "78cdfa543de9e282c86f446f69943965119c2afe", default-features = false } egraph-serialize = { version = "0.3", features = ["serde", "graphviz"] } serde_json = "1" pyo3-log = "*" @@ -42,20 +37,13 @@ uuid = { version = "1.18", features = ["v4"] } rayon = "1.11" base64 = "0.22.1" -# Use patched version of egglog in experimental -[patch.'https://github.com/egraphs-good/egglog'] -# egglog = { path = "../egg-smol" } -# egglog-core-relations = { path = "../egg-smol/core-relations" } -# egglog-ast = { path = "../egg-smol/egglog-ast" } -# egglog-reports = { path = "../egg-smol/egglog-reports" } -# egglog-bridge = { path = "../egg-smol/egglog-bridge" } - -egglog = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "2e5657b" } -egglog-ast = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "2e5657b" } -egglog-core-relations = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "2e5657b" } -egglog-bridge = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "2e5657b" } -egglog-reports = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "2e5657b" } - # enable debug symbols for easier profiling [profile.release] debug = true + +[patch."https://github.com/egraphs-good/egglog.git"] +egglog = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" } +egglog-ast = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" } +egglog-core-relations = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" } +egglog-reports = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" } +egglog-bridge = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" } diff --git a/docs/changelog.md b/docs/changelog.md index abf97873..21edb9aa 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -4,6 +4,39 @@ _This project uses semantic versioning_ ## UNRELEASED +- Upgrade the Python extension to Egglog 3 and the matching + `egglog-experimental` APIs [#414](https://github.com/egraphs-good/egglog-python/pull/414) + - BREAKING: remove `Map.rebuild()`, `Set.rebuild()`, and `Vec.rebuild()`; + rebuilding is now handled by the backend container sorts. + - Preserve the Egglog 3 sort, constructor, function, and proof metadata that + its source syntax can round-trip in the low-level AST bindings; support + arbitrary-size `BigInt` values, generic value extraction, + constructor/relation lookup, microsecond-resolution + run-report durations, program filenames, and all-or-nothing command + recording for parsed batches. + - Add generic `Pair` and `Maybe` values, undefined-result `catch`, map + folding and fold-derived map operations, map/set lengths, `f64` math + primitives and integer coercion, `i64`-to-`BigRat` coercion, and exact + `BigRat.to_i64()` conversion. + - Let Python function, method, and constant bodies lower as eager + primitives, preserving Python argument order for `reverse_args` callables, + while bodies attached to an explicit ruleset remain rewrite-backed; allow + `constant(..., merge=...)` for merged function-backed constants. + Higher-order callable probing is isolated from the active ruleset, and + `rule(..., eval_mode=...)` exposes safe `naive` and explicit + `unsafe-seminaive` rule evaluation when callbacks read mutable tables. + - Add persistent ordinary backoff schedules, expose `RunReport.can_stop`, + and keep high-level saturation running while a scheduler has deferred work. + - Improve source transcripts, diagnostics, generated-name collision + handling, shared-expression factoring without leaking synthetic bindings + into rules or checks, large-program AST parsing, custom-cost extraction of + literal roots, typed pretty/freeze round trips, map duplicate-key behavior, + and set iteration deduplication. + - Preserve the paused Param-Eq research as a small reusable experimental + module and CLI, three bounded public CI stress cases, and an optional + aggregate-only external-corpus harness with explicit iteration-limit, + timeout, and error accounting plus loaded-extension provenance. + ## 13.2.0 (2026-06-03) - Add Python-friendly `RunReport` wrapper that returns `CommandDecl` objects as rule keys instead of raw egglog s-expression strings, with pretty-printed Python syntax in `str()` output [#416](https://github.com/egraphs-good/egglog-python/pull/416) diff --git a/docs/conf.py b/docs/conf.py index 4d8f159b..f4e27cc9 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -34,6 +34,7 @@ "deflist", ] myst_fence_as_directive = ["mermaid"] +myst_heading_anchors = 3 ## # Built presentation in sphinx diff --git a/docs/explanation/2026_02_containers.md b/docs/explanation/2026_02_containers.md index e504eea1..eadcbc48 100644 --- a/docs/explanation/2026_02_containers.md +++ b/docs/explanation/2026_02_containers.md @@ -235,19 +235,22 @@ to do the constant folding: @ruleset def constant_fold_index(xs: MultiSet[Num], i: i64, k: i64): # For all sums, fill in the index function - yield rule(sum_(xs)).then(xs.fill_index(ms_num_index)) + yield rule(sum_(xs), eval_mode="naive").then(xs.fill_index(ms_num_index)) # Try replacing any sum with the folded version - yield rewrite(sum_(xs)).to( - # Replace the two numbers with their sum, by removing - # them and then inserting their sum back in - sum_(xs.remove(Num(i)).remove(Num(k)).insert(Num(i + k))), + yield rule( + sum_(xs), # These are conditions for the rewrite to match: # Look for a multiset that contains two numbers that # are not the same one ms_num_index(xs, Num(i)), ms_num_index(xs, Num(k)), i != k, + eval_mode="naive", + ).then( + # Replace the two numbers with their sum, by removing + # them and then inserting their sum back in + union(sum_(xs)).with_(sum_(xs.remove(Num(i)).remove(Num(k)).insert(Num(i + k)))) ) @@ -297,12 +300,15 @@ def constant_fold_sum(xs: MultiSet[Num]): remaining = xs - constants.map(UnstableFn(Num)) # Sum all the constants to fold them together folded = multiset_fold(i64.__add__, i64(0), constants) - yield rewrite(sum_(xs)).to( - # replace it with the non constants plus the folded - sum_(remaining.insert(Num(folded))), + yield rule( + sum_(xs), # Only run this rule if there are more than one # constant to fold together constants.length() > 1, + eval_mode="naive", + ).then( + # replace it with the non constants plus the folded + union(sum_(xs)).with_(sum_(remaining.insert(Num(folded)))) ) @@ -334,13 +340,11 @@ to the semantics of the your use case, compared to say a tree of binary operatio This work also highlights some of the current limitations of egglog. -One issue is that composing functions of primitives is currently very limited. The only tool we have is currying, but -it is not possible to reorder arguments or compose them in more complicated manners. This inevitably leads to -creating more bespoke functions. For example, I had to add a `multiset_contains_swapped` function that swaps the order -of the `contains` method, since I needed to partially apply it with the second argument. Further exploring this line of -work might lead to trying out different ways of enriching primitive functions, possibly by allowing a way at runtime -to create new ones by composing others, either through a DSL/JIT or a higher order composition approach like the -[compiling to categories](http://conal.net/papers/compiling-to-categories/) work. +The original version of this post identified primitive composition as a major +limitation: argument-order adapters required bespoke flipped backend +primitives. The current Python bindings can lower body-defined functions and +lambdas to primitives, so those adapters can now be expressed at the call site. +The retained examples therefore do not add new flipped container APIs. Implementing these higher order functional primitives on containers is also challenging, due to the lack of built-in generic type support in Egglog. Adding them currently is fiddly and requires careful thought over how to implement diff --git a/docs/how-to-guides/parsing-and-running-program-strings.md b/docs/how-to-guides/parsing-and-running-program-strings.md index a9d2c5d0..c2d78c6c 100644 --- a/docs/how-to-guides/parsing-and-running-program-strings.md +++ b/docs/how-to-guides/parsing-and-running-program-strings.md @@ -19,3 +19,29 @@ commands ```{code-cell} egraph.run_program(*commands) ``` + +When the program is already a string, use +{meth}`egglog.bindings.EGraph.parse_and_run_program` to parse and execute it in +one call. Supplying a filename adds that source name to parse and runtime error +locations: + +```{code-cell} +egraph.parse_and_run_program( + "(function double (i64) i64 :no-merge)\n(set (double 2) 4)", + filename="generated.egg", +) +``` + +The low-level binding can also record successfully executed command batches. +Recording is opt-in and preserves program order: + +```{code-cell} +recording_egraph = EGraph(record=True) +recording_egraph.parse_and_run_program("(let $answer 42)\n(check (= $answer 42))") +print(recording_egraph.commands()) +``` + +If a batch fails after executing an earlier command, Egglog keeps the mutation +made by that prefix, but the failed batch is not appended to `commands()`. Code +that needs transactional behavior should use a separate e-graph or explicit +`push`/`pop` boundaries. diff --git a/docs/reference/egglog-translation.md b/docs/reference/egglog-translation.md index 217360a2..b0c52853 100644 --- a/docs/reference/egglog-translation.md +++ b/docs/reference/egglog-translation.md @@ -6,11 +6,12 @@ file_format: mystnb The high level bindings available at the top module (`egglog`) expose most of the functionality of the `egglog` text format. This guide explains how to translate between the two. -Any EGraph can also be converted to egglog with the `egraph.as_egglog_string` property, as long as it was created with `Egraph(save_egglog_string=True)`. +Any EGraph can also be converted to egglog with the `egraph.as_egglog_string` property, as long as it was created with `EGraph(save_egglog_string=True)`. ## Builtin Types -The builtin types of Unit, String, Int, Map, and Rational are all exposed as Python classes. +Builtin sorts including `Unit`, `String`, `i64`, `f64`, `BigInt`, `BigRat`, +`Map`, `Set`, and `Vec` are exposed as Python classes. These can be imported from `egglog` can be instantiated using the class constructor from the equivalent Python type. Many of the functions on them are mapped to Python operators. For example, the `>>` operator is mapped to `__rshift__` so it can be used as `a >> b` in Python. @@ -41,6 +42,12 @@ i64(10) + 2 BigRat(1, 2) / BigRat(2, 1) ``` +The floating-point sort also exposes the backend's `exp()`, `log()`, and +`sqrt()` primitives. `BigRat.to_i64()` is partial: it is defined only when the +rational value is an integer that fits in `i64`. As with other partial +primitives, undefined use in a rule fact skips that match, while undefined use +in an action is an error. + ### `!=` Operator The `!=` function in egglog works on any two types with the same sort. In Python, this is mapped to the `ne` function: @@ -85,6 +92,38 @@ Since the generic types in the `Map` sort as specified with the `Generic` class, This doesn't require any custom type analysis on our part, only using Python's built in annotations with generic types. +### Generic container operations + +`Pair[L, R]` and `Maybe[T]` expose Egglog's generic product and optional-value +patterns in Python. `catch(lambda: expression)` converts an undefined partial +primitive call, such as a missing map lookup, into `Maybe.none()` instead of +failing the surrounding expression. + +```{code-cell} python +pair = Pair(i64(1), String("one")) +pair.left, pair.right + +present = Maybe[i64].some(1) +missing = catch(lambda: Map[i64, String].empty()[1]) +present, missing +``` + +Maps have a general `map_fold_kv` primitive. The higher-level +`map_filter_kv`, `map_map_values`, `map_merge_with`, `Map.keys()`, and +`Map.pick_key()` operations are composed from that fold and ordinary +container operations, so callbacks may be regular Egglog lambdas: + +```{code-cell} python +numbers = Map[i64, i64].empty().insert(1, 10).insert(2, 20) +map_map_values(lambda _key, value: value + 1, numbers) +``` + +Map folding uses opaque, e-graph-local `Value` order, not a semantic ordering +promised for arbitrary e-class keys. Prefer order-independent callbacks. An +undefined filter predicate skips that entry; +an undefined transform, merge callback, or fold callback makes the whole +operation undefined. + ## Declaring Functions In egglog, the most general way to declare a function is with the `(function ...)` command. In Python, we can use the `@function` decorator on a function with no body. The arg and return types are inferred from the function signature: @@ -98,15 +137,17 @@ def fib(n: i64Like) -> i64: Note that instead of using `i64` as the argument type, we used `i64Like` which is `i64 | int`. This allows us statically to declare that this function can take integers as well which will be upcasted to `i64` automatically. -The `function` decorator supports a number of options as well, which can be passed as keyword arguments, that correspond to the options in the egglog command: +The `function` decorator also accepts keyword arguments that map to backend features. Which ones are valid depends on how +the callable lowers, as described in [Functions vs Constructors](#functions-vs-constructors): - `egg_fn`: The name of the function in egglog. By default, this is the same as the Python function name. -- `cost`: The cost of the function. By default, this is 1. -- `merge`: A function to merge the results of the function. This must be a function that takes two arguments of the return type, the old and the new, and returns a single value of the return type. +- `merge`: A function to merge the results of function-style declarations. This must take the old and new return values and + return a single value of the same type. +- `cost`: The extraction cost for constructor-style declarations. ```{code-cell} python -# egg: (function foo () i64 :cost 10 :merge (max old new)) -@function(egg_fn="foo", cost=10, merge=lambda old, new: old.max(new)) +# egg: (function foo () i64 :merge (max old new)) +@function(egg_fn="foo", merge=lambda old, new: old.max(new)) def my_foo() -> i64: pass ``` @@ -115,13 +156,54 @@ The static types on the decorator preserve the type of the underlying function, ### Functions vs Constructors -Egglog has changed how it handles functions, seperating them into two seperate commands: +The Python bindings follow the backend split in egglog: + +- primitive-returning callables use function-style lowering +- eqsort-returning callables use constructor-style lowering + +That is not a Python-only policy choice. It comes from which backend features exist on each command: -- `function` which can include a `merge` expression. -- `constructor` which can include a cost and requires the result to be an "eqsort" aka a non builtin type. +- function-style declarations support `merge` +- constructor-style declarations support `cost` and `unextractable` +- `subsume` only applies to rewrite-backed bodies, so it only makes sense with an explicit `ruleset` -Since this was added after the Python API was first created, we added support to automatically choose between the two based on the return type of the function and whether a merge function is provided. If the return type is a builtin type, it will be a `function`, otherwise it will be a `constructor`, unless it has a merge function -provided then it will always be a `function`. +Python automatically infers which backend lowering to use from the callable shape. In practice, Python declarations can +lower to a `function`, a `constructor`, or an eager `primitive` depending on the return kind and whether a body/default +is present. + +For bodies and defaults, the canonical lowering mapping is: + +| Python shape | Lowering | +| --- | --- | +| primitive return, no body | lower to `function` | +| primitive return, body | lower to eager `primitive` | +| eqsort return, no body, no `merge` | lower to `constructor` | +| eqsort return, no body, with `merge` | lower to `function` | +| eqsort return, body, no `ruleset` | lower to eager `primitive` | +| eqsort return, body, explicit `ruleset` | lower to `constructor` plus rewrite-backed body | + +Constants and class-variable defaults are just zero-arg bodies/defaults, so they follow the same split based on their +declared return type: + +- no-default constants lower like zero-arg declarations, so primitive-returning constants lower as functions, while + eqsort-returning constants lower as constructors unless `merge` forces function-style lowering +- eqsort-returning defaults lower eagerly without a `ruleset`, and lower to rewrite-backed defaults with an explicit `ruleset` +- primitive-returning defaults lower eagerly, and cannot use an explicit `ruleset` + +Options follow that same backend split: + +- no-body function-style declarations may use `merge` +- builtin declarations are primitive/function-style only +- constructor-style declarations may use `cost` and `unextractable` +- `subsume` is only valid when an eqsort-returning body is lowered through an explicit `ruleset` +- `egg_fn` and mutating arguments are supported in every case +- direct top-level `@function(ruleset=...)` declarations require a body +- `constant(..., ruleset=...)` declarations require an eqsort-returning default +- `constant(..., merge=...)` declarations must not provide a default +- class-level `ruleset=` is still valid shorthand for attaching rewrite-backed eqsort method and class-variable defaults + +For the Python ergonomics of attaching rewrite-backed bodies/defaults to an explicit `ruleset`, see +[Python Integration](python-integration.md#default-replacements). ### Datatype functions @@ -314,6 +396,12 @@ egraph.register( ) ``` +Rules use semi-naive evaluation by default. A rule whose higher-order callback +must read tables populated during the same run can opt into naive evaluation +with `rule(..., eval_mode="naive")`. The third mode, +`eval_mode="unsafe-seminaive"`, skips semi-naive validation and should only be +used when the rule is known to be valid under that evaluation strategy. + ### Variables Unlike in egglog, variables must be declared before being use and must be given a type. They need a type both so that they can be checked statically and also so that we know what types are used to understand what how the names of the egg functions correspond to the method names. @@ -561,6 +649,22 @@ scheduler outside `* 10` lets its `times_banned` counters accumulate across all ten runs. Placing `bo.scope(...)` inside `* 10` creates a fresh scheduler each time, so every iteration starts with the initial `match_limit` and `ban_length`. +The scheduler bindings above are local to one call to `EGraph.run`. To carry a +scheduler's ban state across separate calls on the same e-graph, mark it as +persistent: + +```{code-cell} python +bo = back_off(match_limit=10).persistent() +step_egraph.run(run(step_right, scheduler=bo)) +step_egraph.run(run(step_right, scheduler=bo)) +``` + +The scheduler is registered once on that e-graph and reused by both calls. A +persistent scheduler has its own identity, so deriving it from another +scheduler configuration does not alias that configuration's local state. +High-level `EGraph.saturate()` also waits for `RunReport.can_stop`, so a +no-change round does not discard work deferred by a persistent scheduler. + ## Check The `(check ...)` command to verify that some facts are true, can be translated to Python with the `egraph.check` function: @@ -633,7 +737,8 @@ egraph.check_fail(eq(Math(0)).to(Math(1))) ## Function Sizes The `(print-size ?)` command is translated into either `egraph.function_size(fn)` to get the number of -rows of one function or `egraph.all_function_sizes()` to get a list of all the function sizes: +rows in one table-backed callable or `egraph.all_function_sizes()` to list the sizes of all registered function tables. +Relations, constructors, and bodyless functions have tables; eager and builtin primitives do not: ```{code-cell} python # (function-size Math) @@ -656,7 +761,8 @@ egraph.stats() ## Function Values -The `print-function` command is translated into `egraph.function_values(fn, [length]?)` to get the values of a specific function. Note that the function provided must either return a primitive or be created with a merge function. +The `print-function` command is translated into `egraph.function_values(fn, [length]?)` to get the rows of a +table-backed callable. As with `function_size`, eager and builtin primitives cannot be inspected this way. ```{code-cell} python # (print-function fib 3) diff --git a/docs/reference/python-integration.md b/docs/reference/python-integration.md index 23f52d8f..9a467ac5 100644 --- a/docs/reference/python-integration.md +++ b/docs/reference/python-integration.md @@ -250,6 +250,8 @@ Registering a conversion from A to B will also register all transitively reachab Math(2) + 30 + "x" ``` +When defining converters for a custom `Expr` sort, prefer registering conversions from egglog primitive sorts such as `i64`, `f64`, and `String` rather than directly from Python builtins like `int`, `float`, and `str`. The builtin promotions already handle those Python values transitively, so keeping the custom converters at the egglog-sort layer makes the promotion path clearer and usually leads to cleaner `...Like` aliases such as `Math | i64Like | f64Like | StringLike`. + If you want to have this work with the static type checker, you can define your own `Union` type, which MUST include the `Expr` class as the first item in the union. For example, in this case you could then define: @@ -456,7 +458,7 @@ assert str(-1.0 + Int.var("x")) == "Float(-1.0) + Float.from_int(Int.var(\"x\")) ### Mutating arguments -In order to support Python functions and methods which mutate their arguments, use the `mutates_first_arg` keyword argument on `@function` and the `mutates_self` keyword argument on `@method`. The runtime treats the mutated receiver as the return value of the egglog call, so the default rewrite points the call expression at the updated argument. +In order to support Python functions and methods which mutate their arguments, use the `mutates_first_arg` keyword argument on `@function` and the `mutates_self` keyword argument on `@method`. The runtime treats the mutated receiver as the semantic return value of the egglog call. A body lowers eagerly unless it is attached to an explicit `ruleset`, in which case the updated expression becomes that callable's default rewrite. Inside the Python implementation you can call `__replace_expr__` on an `Expr` instance to swap out its underlying egglog expression in-place. This keeps any existing Python references in sync while still allowing the e-graph to reason about the mutated value. The same helper works for methods that run immediately with `@method(preserve=True)`. @@ -489,19 +491,23 @@ mutate_egraph.register(rewrite(incr_i).to(i + Int(1)), x) mutate_egraph.run(10) mutate_egraph.check(eq(x).to(Int(10) + Int(1))) -# incr with the rewrite could also be written like this: +# The update can instead be defined as an eager body: @function(mutates_first_arg=True) def incr_other(x: Int) -> None: x.__replace_expr__(x + Int(1)) x = Int(10) incr_other(x) mutate_egraph = EGraph() -mutate_egraph.register(x) -mutate_egraph.run(10) -mutate_egraph.check(eq(x).to(Int(10) + Int(1))) +incremented = mutate_egraph.let("incremented", x) +mutate_egraph.check(eq(incremented).to(Int(10) + Int(1))) mutate_egraph ``` +The bodyful form lowers to an eager primitive. Because this example constructs +an e-class value, bind its result through an action before using that value in +a read-only check. Use an explicit `ruleset=` when the body should remain a +rewrite instead. + Note that dunder methods such as `__setitem__` will automatically be marked as mutating their first argument. ## Functions as Values @@ -549,11 +555,14 @@ We also support using normal python functions, either named or anonymous, as val ```{code-cell} python x = MathList.EMPTY.append(Math(1)) added_two = x.map(lambda x: x + Math(2)) -check_eq(added_two, MathList.EMPTY.append(Math(1) + Math(2)), (math_list_ruleset + run()) * 10) +check_eq(added_two, MathList.EMPTY.append(Math(1) + Math(2)), math_list_ruleset.saturate()) ``` -Their definition will be added to the default rulset, unless they are defined in the body of a function themselves or -in a rule function: +Converting the callback to `Callable`/`UnstableFn` materializes its body as an +eager anonymous primitive. It does not add a rewrite to the default ruleset or +require a separate `run()` step. An enclosing declared function still follows +the ordinary lowering rules; for example, an explicit `ruleset` makes this +outer body rewrite-backed: ```{code-cell} python @function(ruleset=math_list_ruleset) @@ -563,54 +572,25 @@ def map_add_two(x: MathList) -> MathList: check_eq(map_add_two(MathList.EMPTY.append(Math(1))), MathList.EMPTY.append(Math(1) + Math(2)), math_list_ruleset.saturate()) ``` -Their name will just be the body of the function, so that two anonymous functions with the same body will be considered equal. - -```{code-cell} python -added_two -``` +Generated primitive names are internal implementation details. ## Default Replacements -When defining a function or a constant, you can also provide a default replacement value. This is useful when -you might want both the original value and the replaced value in the e-graph, so that later rules could reference either. - -```{code-cell} python -@function -def math_float(f: f64Like) -> Math: - ... +The full lowering matrix for functions, constructors, primitives, and defaults is documented in +[Translation to/from egglog](egglog-translation.md#functions-vs-constructors). This section focuses on the +Python-only ergonomics for the explicit-`ruleset` case: when you pass a `ruleset`, the body/default is added as +a rewrite into that ruleset instead of being lowered eagerly. +This is useful when you want the declared name and the replacement body to both remain available to later rewrite rules. -# Can add a default replacement value for a constants -pi = constant("pi", Math, math_float(3.14)) +You can specify a ruleset for a default replacement by passing the `ruleset` keyword argument: +```{code-cell} python +math_ruleset = ruleset() -# or for a function by providing a body @function -def square(x: Math) -> Math: - return x * x - -# thse rewrites will be added to the e-graph under the default ruleset -egraph = EGraph() -egraph.register(pi) -egraph.register(square(Math.var('x'))) -egraph.run(1) -egraph.check(eq(pi).to(math_float(3.14))) -egraph.check(eq(square(Math.var('x'))).to(Math.var('x') * Math.var('x'))) -egraph -``` - -This is equivalent to adding the rewrite rules to the e-graph directly, like this, but just more succinct: +def math_float(value: f64Like) -> Math: ... -```python -x = var("x", Math) -egraph.register(rewrite(pi).to(math_float(3.14))) -egraph.register(rewrite(square(x)).to(x * x)) -``` - -You can also specify a ruleset to add the rewrites to, by passing in the `ruleset` keyword argument: - -```{code-cell} python -math_ruleset = ruleset() e_constant = constant("e", Math, math_float(2.71), ruleset=math_ruleset) @@ -626,9 +606,26 @@ egraph.check(eq(e_constant).to(math_float(2.71))) egraph.check(eq(cube(Math.var('x'))).to(Math.var('x') * Math.var('x') * Math.var('x'))) ``` +This rewrite-backed path is only available for eqsort-returning bodies and defaults. Primitive-returning defaults lower +eagerly and cannot use an explicit `ruleset`. + +When `subsume=True` is allowed for that callable shape, it applies on this rewrite-backed path as well. + +Constants without defaults can use merge functions because they lower as zero-argument function-style declarations: + +```{code-cell} python +best_score = constant("best_score", i64, merge=lambda old, new: old.max(new)) + +egraph = EGraph() +egraph.register(set_(best_score).to(i64(1)), set_(best_score).to(i64(2))) +egraph.check(eq(best_score).to(i64(2))) +``` + +Constants with eager or rewrite-backed defaults cannot also use `merge`. + ### Default Replacement for Classes -In classes, you can also provide a default replacement value for constants and methods, and an optional ruleset on the class constructor: +In classes, a `ruleset=` on the class means default method and class-variable bodies are also added to that ruleset as rewrites: ```{code-cell} python other_math_ruleset = ruleset() @@ -652,6 +649,25 @@ egraph.check(eq(x).to(WrappedMath(math_float(3.14)) + WrappedMath(math_float(3.1 egraph ``` +## Param-Eq Stress Demo + +The experimental `egglog.exp.param_eq` module preserves a bounded +parameter-reducing symbolic-regression pipeline. Its CLI runs either retained +representation and emits a JSON report: + +```{code-block} console +$ python -m egglog.exp.param_eq --expr '2.3 * (3.7*x0 + 5.1*x1) / 7.9' --variant container +``` + +Expressions use finite numeric literals, variables, Python arithmetic with +literal exponents, and `abs`, `exp`, `log`, `sqrt`, `plog`, `square`, or +`cube`. A `saturated` status means every inner schedule could stop; +`iteration_limit` means the retained 30-round boundary was reached. The rules +target real inputs where every relevant subexpression is defined, and the +included finite sample checks are regression tests rather than a proof of +universal equivalence. The container variant rejects inputs whose coefficient +normalization produces a non-finite `f64` value. + ## Debugging and Inspection When a rule does not fire or an equality appears unexpectedly, the most useful @@ -696,6 +712,10 @@ report = egraph.run(debug_rules) report.num_matches_per_rule ``` +`report.updated` records whether the run changed the database. The separate +`report.can_stop` flag is true only when the run observed no changes and its +scheduler has no deferred work that requires another iteration. + ### `stats` Use {meth}`egglog.egraph.EGraph.stats` when you want cumulative counters for the @@ -709,7 +729,8 @@ stats.num_matches_per_rule ### `function_values` Use {meth}`egglog.egraph.EGraph.function_values` to inspect the current rows in a -function table: +function table. This accepts relations, constructors, and bodyless functions; +eager and builtin primitives do not have tables to inspect: ```{code-cell} python egraph.function_values(score) @@ -738,7 +759,8 @@ egraph.display() ### `saturate` Use {meth}`egglog.egraph.EGraph.saturate` to keep running until the schedule -stops changing the graph while printing the extracted form after each step: +reports no graph changes or deferred scheduler work, while printing the +extracted form after each step: ```{code-cell} python egraph = EGraph() diff --git a/experiments/param_eq/Makefile b/experiments/param_eq/Makefile new file mode 100644 index 00000000..fe9e3ca9 --- /dev/null +++ b/experiments/param_eq/Makefile @@ -0,0 +1,25 @@ +.PHONY: smoke binary container haskell paired aggregate + +ROOT := $(abspath $(dir $(lastword $(MAKEFILE_LIST)))/../..) +RESULTS := $(ROOT)/experiments/param_eq/results +RAW := $(RESULTS)/raw +PAIRED := $(RAW)/paired.csv +EXECUTION_MODE ?= release + +smoke: + cd $(ROOT) && uv run pytest --benchmark-disable -q -m param_eq_smoke python/tests/param_eq + +binary: + cd $(ROOT) && uv run python -m experiments.param_eq.run --variant binary --execution-mode $(EXECUTION_MODE) --output $(RAW)/binary.csv + +container: + cd $(ROOT) && uv run python -m experiments.param_eq.run --variant container --execution-mode $(EXECUTION_MODE) --output $(RAW)/container.csv + +haskell: + cd $(ROOT) && uv run python -m experiments.param_eq.run --implementation haskell --variant binary --execution-mode $(EXECUTION_MODE) --output $(RAW)/haskell.csv + +paired: + cd $(ROOT) && uv run python -m experiments.param_eq.run --variant both --execution-mode $(EXECUTION_MODE) --output $(PAIRED) + +aggregate: paired + cd $(ROOT) && uv run python -m experiments.param_eq.aggregate --binary $(PAIRED) --container $(PAIRED) --output-dir $(RESULTS) diff --git a/experiments/param_eq/NOTES.md b/experiments/param_eq/NOTES.md new file mode 100644 index 00000000..cb91afc1 --- /dev/null +++ b/experiments/param_eq/NOTES.md @@ -0,0 +1,116 @@ +# Param-Eq research handoff + +## What is retained + +The binary pipeline is the fidelity baseline for de França and Kronberger's +parameter-reducing rewrite method. It uses a lexicographic extraction cost: +noninteger floating constants first, then expression nodes. Each outer pass +builds equivalences, extracts the best expression, and repeats once if the +representative changed. + +The container variant represents monomials and polynomials as nested +`Map`/`BigRat` values. It was introduced to make associative/commutative +polynomial structure canonical and to avoid enumerating every binary tree. +Only the general higher-order container operations needed by the retained +pipelines and public demonstrations should remain in `egglog-experimental`. +Argument-order adaptations use lambda-created primitives; flipped duplicate +builtins are not part of the retained design. + +The final public-case ablation retained small-polynomial coefficient factoring, +Horner factoring, and exact nested-polynomial flattening from the container +basic rules, together with constant/index collection and polynomial +normalization from the analysis rules. The specialized integer-residual, +subset-scale, and corpus-tail rules were unnecessary for all three public cases, +so the installed pipeline no longer references their primitives. +Literal float exponents are converted on the Python side with +`float.as_integer_ratio()`, constant-map collisions use a general old-value +merge, and rational float powers are composed from ordinary exponentiation and +`BigRat.to_f64()`. + +Container lowering distributes a literal scalar or one monomial into a single +polynomial and combines coefficient collisions. Polynomial-by-polynomial and +monomial-over-polynomial forms stay nested. This makes the public Equation 4 +case canonical without broadening the nested-flatten rewrite or recreating its +factor/flatten cycle. + +Both variants now use ordinary persistent upstream backoff. This is close to, +but not identical to, the prototype scheduler and is a documented fidelity +boundary rather than an unverified claim of exact replication. + +The binary repeated-monomial public stress case currently reaches the retained +30-round inner limit. CI still verifies every sample point of its extracted +expression, but the corpus runner records it as `iteration_limit` and excludes +it from successful aggregates. Revisit scheduler parity before resuming corpus +measurements. + +## Semantic limits + +The retained rules target real-valued expressions on inputs where the source +and every introduced subexpression are defined. Guards cover the literal and +structural domain boundaries needed by the retained cases, but the pipeline has +no general sign or interval analysis. CI compares every configured finite +sample point; those checks are regression evidence, not a universal proof of +equivalence. The container variant fails explicitly if coefficient +normalization becomes non-finite. Treat external-corpus measurements the same +way. + +## Performance evidence worth preserving + +The last pre-pause corpus artifacts repeatedly showed a smaller final container +e-graph alongside a slower wall time. Profiling indicated several reasons worth +testing later: + +- final `egraph_total_size` is not peak or cumulative work; +- constructing rules/declarations and materializing higher-order operations can + dominate small final graphs; +- embedded `Map` payload and custom extraction costs are not represented by a + simple row count; +- some exact flatten/scale rules create large intermediate polynomial spaces or + normalization cycles. + +These are hypotheses supported by local probes, not causal or portable +performance conclusions. The stale row-level artifacts and chronological debug +transcript were deliberately removed. A final dependency-compatible rerun is +required before publishing numerical corpus results. + +## Rejected or parked directions + +- Exact flattening and representative-scale rules improved isolated shapes but + created cycles or bad tail cases. +- Increasing the container match budget restored some reachability but made the + long tail worse. +- Final e-graph size alone was not a useful stop or performance proxy. +- Rank-miss-specific rules were not retained as a second undocumented rule set. +- The common-float-scale and integer-residual campaigns are recorded only as + rejected directions here; their commented implementations and specialized + helper APIs were deliberately removed from the installed package. + +## Restart checklist + +1. Record the Python, `egglog`, and `egglog-experimental` commits and rebuild the + extension in release mode. The runner records the clean source checkouts and + hashes the loaded native extension; the hash detects a changed executable + but does not itself prove which checkout produced it. +2. Run the three public cases in both variants and keep their independent + numeric checks green. Confirm the documented status boundary: binary + `repeated_monomial` reaches `iteration_limit`, while the other reports are + `saturated`. Resolve that limit before publishing new corpus measurements. +3. Set `EGGLOG_PARAM_EQ_DATA_DIR` to the private archive and + `EGGLOG_PARAM_EQ_EXPECTED_ARCHIVE_SHA256` to the stable value recovered from + private research records. The runner refuses an absent or mismatched hash. +4. Run binary and container rows with the same time/memory limits and inspect + every iteration-limit/timeout/error count. + Optionally run `make -C experiments/param_eq haskell`; this live baseline + compiles its temporary runner once, forces both result counts inside the + timed region, and requires Stack plus the external Haskell checkout. Its raw + output is diagnostic only; no tracked Haskell aggregate path is maintained. +5. Generate aggregate-only outputs and verify the manifest and paired hashes. + Confirm its Python and Rust versions, platform, CPU, execution mode, + requested/effective workers, ordering seed, loaded-extension hash, and + per-variant status counts; + aggregation rejects dirty or unidentified dependency worktrees and a dirty + egglog-python checkout. +6. Investigate peak/cumulative work and extraction payload costs before adding + more rewrite rules. +7. Obtain corrected license/redistribution permission before publishing any + source or row-level material. diff --git a/experiments/param_eq/README.md b/experiments/param_eq/README.md new file mode 100644 index 00000000..8a705f0f --- /dev/null +++ b/experiments/param_eq/README.md @@ -0,0 +1,109 @@ +# Param-Eq paused-research harness + +This directory preserves the optional corpus experiment without installing it +as part of `egglog`. The reusable expression domain and simplifier live in +`python/egglog/exp/param_eq`; this directory owns external-data loading, +isolated row execution, resource limits, aggregation, and the research handoff. + +The work is paused. The bounded public demonstrations remain maintained in CI, +while the private 714-row corpus is not run automatically. + +## Provenance and redistribution boundary + +This work reimplements the method published by Fabrício Olivetti de França and +Gabriel Kronberger in [Reducing Overparameterization of Symbolic Regression +Models with Equality Saturation](https://doi.org/10.1145/3583131.3590346). +Fabrício provided the original Haskell experiment repository and a separate +`pandoc-symreg` archive in personal correspondence. Those private files were +used for behavioral validation but are not redistributed here. The Python +implementation reproduces the published method without copying the prototype's +source text; checked-in result artifacts contain only aggregate measurements +and no source expressions. + +The supplied archive has placeholder copyright/author metadata, so attribution +alone is not sufficient redistribution permission. Keep it outside this +repository and configure it explicitly: + +```bash +export EGGLOG_PARAM_EQ_DATA_DIR=/absolute/path/to/param-eq-haskell +export EGGLOG_PARAM_EQ_EXPECTED_ARCHIVE_SHA256= +``` + +The expected hash is intentionally not checked into this public tree because +it fingerprints private material. Recover it from private research records +before resuming. If no prior value can be recovered, compute the current +candidate with `external_archive_hash()` only after independently verifying the +archive, store that value privately, and use it for every subsequent run. + +The loader applies the retained-paper policy by name: FEAT is omitted, SRjl and +GOMEA are displayed as PySR and GP-GOMEA, known unusable raw rows are omitted, +and rows without rank data are excluded. + +## Installed stress demo + +Run either retained representation on one public expression: + +```bash +python -m egglog.exp.param_eq --expr '2.3 * (3.7*x0 + 5.1*x1) / 7.9' --variant container +``` + +The restricted Python-like syntax accepts finite numeric literals, variables, +`+`, `-`, `*`, `/`, literal exponents, and the unary functions `abs`, `exp`, +`log`, `sqrt`, `plog`, `square`, and `cube`. The JSON report has status +`saturated` when every inner schedule can stop, or `iteration_limit` when the +retained 30-round boundary is reached. In either case the extracted expression +is available for independent checking; only saturated corpus rows contribute +to aggregates. The container variant rejects inputs whose coefficient +normalization produces a non-finite `f64` value. See `NOTES.md` for fidelity +and semantic limits. + +## External corpus commands + +```bash +make -C experiments/param_eq smoke +make -C experiments/param_eq binary +make -C experiments/param_eq container +make -C experiments/param_eq haskell +make -C experiments/param_eq aggregate +``` + +`binary` and `container` are useful for focused work and write expression-free +row metrics only under the ignored `results/raw/` directory. `aggregate` runs +the paired mode, alternating variant order by stable row hash, validates +identities, configuration, and input hashes, then replaces the tracked +aggregate CSVs and manifest. Never add files from `results/raw/`, the external +archive, source expressions, extracted expressions, or private absolute paths. + +The raw `external_archive_sha256` column intentionally repeats one hash of the +corpus inputs, Haskell source modules, and Stack/Cabal lock/configuration files; +it is not a hash of an individual expression. Raw rows +also record the Python and Rust versions, platform, CPU, declared debug/release +mode, requested and memory-capped worker counts, stable ordering seed, clean +source-checkout commits, and the SHA-256 of the loaded native extension. +Aggregation refuses to publish when those fields disagree, a recorded +worktree is dirty, the current repository is not the recorded producer commit, +or an identity is missing. The native hash identifies the code that executed; +the source commits remain procedural provenance, so rebuild the extension from +those clean checkouts before running. Set `EXECUTION_MODE=debug` on the make +command only when the installed extension was actually built in debug mode; +the default is `release`. + +Full timing comparisons are single-machine exploratory measurements. The +runner isolates each row, records iteration limits, timeouts, and errors instead +of dropping them, and alternates binary/container order by stable row hash when +`--variant both` is used. Aggregate rows retain separate counts for iteration, +timeout, memory, and execution failures. Ratio summaries omit pairs whose +binary denominator is zero and expose the remaining sample as `n_ratio`. + +The optional `haskell` target compiles one temporary runner against the +author-supplied implementation, then executes it once per isolated row. The +one-time Stack/GHC build is outside the per-row timer and has its own 600-second +guard; override that setup boundary with `--haskell-build-timeout-sec` when +calling the runner directly. The generated program looks up rows by public +metadata rather than embedding expression text, forces the input counts before +starting the clock, and forces both result counts before stopping it. Stack is +not installed or invoked by CI; program generation and expression-free output +parsing are unit tested. Haskell has no container representation, so this route +supports only `--variant binary`. Its expression-free CSV remains a local raw +diagnostic: `aggregate` does not currently publish a Haskell aggregate. Add +and validate a separate aggregate path before presenting those measurements. diff --git a/experiments/param_eq/__init__.py b/experiments/param_eq/__init__.py new file mode 100644 index 00000000..82268986 --- /dev/null +++ b/experiments/param_eq/__init__.py @@ -0,0 +1 @@ +"""Optional external-data harness for the paused Param-Eq research.""" diff --git a/experiments/param_eq/aggregate.py b/experiments/param_eq/aggregate.py new file mode 100644 index 00000000..3623ff0b --- /dev/null +++ b/experiments/param_eq/aggregate.py @@ -0,0 +1,476 @@ +"""Validate local row results and publish expression-free aggregate artifacts.""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import json +import math +import subprocess +from collections import Counter, defaultdict +from collections.abc import Callable, Iterable, Sequence +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +from experiments.param_eq.corpus import ALGORITHM_RENAMES, DATASETS, RAW_ALGORITHMS +from experiments.param_eq.run import RAW_COLUMNS, VARIANT_ORDER_SEED + +PAPER_COLUMNS = ( + "implementation", + "dataset", + "algorithm", + "input_kind", + "metric", + "n_total", + "n_success", + "n_iteration_limit", + "n_timeout", + "n_memory_limit", + "n_error", + "value", + "min", + "q1", + "median", + "q3", + "max", +) +COMPARISON_COLUMNS = ( + "slice", + "metric", + "n_pairs", + "n_ratio", + "n_binary_missing", + "n_container_missing", + "container_better", + "same", + "container_worse", + "ratio_p10", + "ratio_p25", + "ratio_median", + "ratio_p75", + "ratio_p90", +) +FORBIDDEN_COLUMNS = {"source", "expression", "rendered", "orig_expr", "simpl_expr", "private_path"} +RAW_STATUSES = {"saturated", "iteration_limit", "timeout", "memory_limit", "error"} +PROVENANCE_COLUMNS = ( + "execution_mode", + "workers_requested", + "workers_effective", + "ordering_seed", + "egglog_python_commit", + "egglog_python_clean", + "egglog_core_commit", + "egglog_core_clean", + "egglog_experimental_commit", + "egglog_experimental_clean", + "egglog_bindings_sha256", + "python_version", + "platform", + "rust_version", + "cpu", +) +PUBLIC_ALGORITHMS = frozenset(ALGORITHM_RENAMES.get(name, name) for name in RAW_ALGORITHMS if name != "FEAT") +INPUT_KINDS = frozenset({"original", "sympy"}) + + +def _validate_public_identity(row: dict[str, str], filename: str) -> None: + """Reject labels that could carry private data into aggregate artifacts.""" + if row["dataset"] not in DATASETS: + raise ValueError(f"Raw row in {filename} has an invalid dataset label") + if row["algorithm"] not in PUBLIC_ALGORITHMS: + raise ValueError(f"Raw row in {filename} has an invalid algorithm label") + if row["input_kind"] not in INPUT_KINDS: + raise ValueError(f"Raw row in {filename} has an invalid input-kind label") + identity = row["row_id"].split("/") + if ( + len(identity) != 5 + or identity[0] != row["dataset"] + or identity[2] != row["algorithm"] + or identity[4] != row["input_kind"] + or not identity[1].isdecimal() + or not identity[3].isdecimal() + or int(identity[3]) < 1 + ): + raise ValueError(f"Raw row in {filename} has an invalid or inconsistent row identity") + + +def _number(row: dict[str, str], key: str) -> float: + """Return one required finite, nonnegative raw measurement.""" + value = row[key] + if value == "": + raise ValueError(f"Raw row {row['row_id']} has no {key}") + try: + number = float(value) + except ValueError: + raise ValueError(f"Raw row {row['row_id']} has an invalid {key}") from None + if not math.isfinite(number) or number < 0.0: + raise ValueError(f"Raw row {row['row_id']} has an invalid {key}") + return number + + +def _validate_result_measurements(row: dict[str, str], filename: str) -> None: + """Enforce the result fields allowed by a row's status and implementation.""" + result_fields = ( + "runtime_ms", + "passes", + "total_size", + "before_nodes", + "before_params", + "after_nodes", + "after_params", + ) + if row["status"] == "saturated": + required = ("runtime_ms", "before_nodes", "before_params", "after_nodes", "after_params") + if row["implementation"] == "egglog": + required += ("passes", "total_size") + elif row["passes"] or row["total_size"]: + raise ValueError(f"Haskell row in {filename} unexpectedly has Egglog-only metrics") + for key in required: + number = _number(row, key) + if key != "runtime_ms" and not number.is_integer(): + raise ValueError(f"Raw row in {filename} has a noninteger {key}") + if key == "passes" and number == 0.0: + raise ValueError(f"Raw row in {filename} has no completed passes") + elif any(row[key] for key in result_fields): + raise ValueError(f"Unsuccessful row in {filename} unexpectedly has result metrics") + + +def _validate_measurements(row: dict[str, str], filename: str) -> None: + """Enforce numeric and status-dependent invariants for one raw row.""" + source_rank = _number(row, "source_n_rank") + timeout = _number(row, "timeout_sec") + memory_limit = _number(row, "memory_limit_mb") + sample_interval = _number(row, "sample_interval_sec") + if timeout == 0.0 or sample_interval == 0.0: + raise ValueError(f"Raw row in {filename} has an invalid timeout/sample configuration") + if memory_limit == 0.0 or not memory_limit.is_integer(): + raise ValueError(f"Raw row in {filename} has an invalid memory limit") + if not source_rank.is_integer(): + raise ValueError(f"Raw row in {filename} has a noninteger source_n_rank") + _validate_result_measurements(row, filename) + if row["peak_rss_mb"]: + _number(row, "peak_rss_mb") + + +def load_raw(path: Path, *, expected_variant: str, expected_implementation: str = "egglog") -> list[dict[str, str]]: + """Load a local raw result and reject provenance/schema mismatches.""" + with path.open(newline="", encoding="utf-8") as handle: + reader = csv.DictReader(handle) + fields = tuple(reader.fieldnames or ()) + if fields != RAW_COLUMNS: + extras = FORBIDDEN_COLUMNS.intersection(fields) + detail = f" forbidden columns={sorted(extras)}" if extras else "" + raise ValueError(f"Unexpected raw-result schema in {path.name}.{detail}") + all_rows = list(reader) + variants = {row["variant"] for row in all_rows} + if not variants.issubset({"binary", "container"}): + raise ValueError(f"Unknown variants in {path.name}: {sorted(variants)}") + implementations = {row["implementation"] for row in all_rows} + if not implementations.issubset({"egglog", "haskell"}): + raise ValueError(f"Unknown implementations in {path.name}: {sorted(implementations)}") + statuses = {row["status"] for row in all_rows} + if not statuses.issubset(RAW_STATUSES): + raise ValueError(f"Unknown statuses in {path.name}: {sorted(statuses)}") + for row in all_rows: + _validate_public_identity(row, path.name) + _validate_measurements(row, path.name) + rows = [ + row + for row in all_rows + if row["variant"] == expected_variant and row["implementation"] == expected_implementation + ] + if not rows: + raise ValueError(f"No {expected_implementation}/{expected_variant} rows in {path.name}") + seen: set[str] = set() + for row in rows: + if row["row_id"] in seen: + raise ValueError(f"Duplicate row identity in {path.name}: {row['row_id']}") + seen.add(row["row_id"]) + return rows + + +def _validated_provenance(rows: list[dict[str, str]]) -> dict[str, str]: + """Require one complete, clean dependency/runtime configuration.""" + provenance: dict[str, str] = {} + for column in PROVENANCE_COLUMNS: + values = {row[column] for row in rows} + if len(values) != 1 or not next(iter(values)): + raise ValueError(f"Raw rows do not share one nonempty {column}") + provenance[column] = next(iter(values)) + for column, length in ( + ("egglog_python_commit", 40), + ("egglog_core_commit", 40), + ("egglog_experimental_commit", 40), + ("egglog_bindings_sha256", 64), + ): + identity = provenance[column] + if len(identity) != length or any(character not in "0123456789abcdef" for character in identity.lower()): + raise ValueError(f"Raw rows have an invalid {column}") + for column in ("egglog_python_clean", "egglog_core_clean", "egglog_experimental_clean"): + if provenance[column] != "true": + raise ValueError(f"Refusing aggregate publication with dirty repository state: {column}") + if provenance["execution_mode"] not in {"debug", "release"}: + msg = "Raw rows have an invalid execution_mode" + raise ValueError(msg) + requested = int(provenance["workers_requested"]) + effective = int(provenance["workers_effective"]) + if requested < 1 or effective < 1 or effective > requested: + msg = "Raw rows have invalid worker counts" + raise ValueError(msg) + if provenance["ordering_seed"] != VARIANT_ORDER_SEED: + msg = "Raw rows use an unknown ordering seed" + raise ValueError(msg) + return provenance + + +def _quantile(values: Iterable[float], probability: float) -> float | str: + ordered = sorted(values) + if not ordered: + return "" + position = (len(ordered) - 1) * probability + lower = int(position) + upper = min(lower + 1, len(ordered) - 1) + fraction = position - lower + return ordered[lower] * (1.0 - fraction) + ordered[upper] * fraction + + +def _slice_predicates() -> dict[str, Callable[[dict[str, str]], bool]]: + return { + "all": lambda _row: True, + "original": lambda row: row["input_kind"] == "original", + "sympy": lambda row: row["input_kind"] == "sympy", + "pagie": lambda row: row["dataset"] == "pagie", + "kotanchek": lambda row: row["dataset"] == "kotanchek", + } + + +def build_representation_comparison( # noqa: C901 + binary_rows: list[dict[str, str]], container_rows: list[dict[str, str]] +) -> list[dict[str, Any]]: + """Build paired lower-is-better outcome and ratio summaries.""" + binary = {row["row_id"]: row for row in binary_rows} + container = {row["row_id"]: row for row in container_rows} + if set(binary) != set(container): + missing_binary = sorted(set(container) - set(binary)) + missing_container = sorted(set(binary) - set(container)) + raise ValueError( + f"Unpaired raw inputs: missing_binary={len(missing_binary)} missing_container={len(missing_container)}" + ) + archive_hashes = {row["external_archive_sha256"] for row in [*binary_rows, *container_rows]} + if len(archive_hashes) != 1: + msg = "Binary and container rows do not share one external-input hash" + raise ValueError(msg) + for row_id, left in binary.items(): + right = container[row_id] + for key in ( + "dataset", + "algorithm", + "input_kind", + "source_n_rank", + "timeout_sec", + "memory_limit_mb", + "sample_interval_sec", + *PROVENANCE_COLUMNS, + ): + if left[key] != right[key]: + raise ValueError(f"Paired metadata mismatch for {row_id}: {key}") + + output: list[dict[str, Any]] = [] + metrics = ("after_params", "after_nodes", "runtime_ms", "total_size") + for slice_name, predicate in _slice_predicates().items(): + identities = [row_id for row_id, row in binary.items() if predicate(row)] + for metric in metrics: + better = same = worse = 0 + ratios: list[float] = [] + binary_missing = container_missing = 0 + for row_id in identities: + left, right = binary[row_id], container[row_id] + left_ok = left["status"] == "saturated" + right_ok = right["status"] == "saturated" + binary_missing += int(not left_ok) + container_missing += int(not right_ok) + if not (left_ok and right_ok): + continue + binary_value = _number(left, metric) + container_value = _number(right, metric) + better += int(container_value < binary_value) + same += int(container_value == binary_value) + worse += int(container_value > binary_value) + if binary_value != 0.0: + ratios.append(container_value / binary_value) + output.append({ + "slice": slice_name, + "metric": metric, + "n_pairs": len(identities), + "n_ratio": len(ratios), + "n_binary_missing": binary_missing, + "n_container_missing": container_missing, + "container_better": better, + "same": same, + "container_worse": worse, + "ratio_p10": _quantile(ratios, 0.10), + "ratio_p25": _quantile(ratios, 0.25), + "ratio_median": _quantile(ratios, 0.50), + "ratio_p75": _quantile(ratios, 0.75), + "ratio_p90": _quantile(ratios, 0.90), + }) + return output + + +def build_paper_replication(rows: list[dict[str, str]]) -> list[dict[str, Any]]: + """Build aggregate Egglog rows without source identities or expressions.""" + groups: defaultdict[tuple[str, str, str, str, str], list[dict[str, str]]] = defaultdict(list) + for row in rows: + groups[(row["implementation"], row["variant"], row["dataset"], row["algorithm"], row["input_kind"])].append(row) + output: list[dict[str, Any]] = [] + for (implementation, variant, dataset, algorithm, input_kind), group in sorted(groups.items()): + successes = [row for row in group if row["status"] == "saturated"] + status_counts: defaultdict[str, int] = defaultdict(int) + for row in group: + status_counts[row["status"]] += 1 + metrics = { + "final_params": [_number(row, "after_params") for row in successes], + "parameter_reduction": [_number(row, "before_params") - _number(row, "after_params") for row in successes], + "parameter_rank_gap": [_number(row, "after_params") - _number(row, "source_n_rank") for row in successes], + } + for metric, values in metrics.items(): + output.append({ + "implementation": f"egglog-{variant}" if implementation == "egglog" else implementation, + "dataset": dataset, + "algorithm": algorithm, + "input_kind": input_kind, + "metric": metric, + "n_total": len(group), + "n_success": len(successes), + "n_iteration_limit": status_counts["iteration_limit"], + "n_timeout": status_counts["timeout"], + "n_memory_limit": status_counts["memory_limit"], + "n_error": status_counts["error"], + "value": "", + "min": min(values) if values else "", + "q1": _quantile(values, 0.25), + "median": _quantile(values, 0.50), + "q3": _quantile(values, 0.75), + "max": max(values) if values else "", + }) + return output + + +def _write_csv(path: Path, columns: tuple[str, ...], rows: Iterable[dict[str, Any]]) -> None: + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=columns) + writer.writeheader() + writer.writerows(rows) + + +def _repo_state() -> tuple[str, bool]: + root = Path(__file__).resolve().parents[2] + commit_process = subprocess.run( + ["git", "rev-parse", "HEAD"], # noqa: S607 + cwd=root, + capture_output=True, + check=False, + text=True, + ) + status_process = subprocess.run( + ["git", "status", "--porcelain"], # noqa: S607 + cwd=root, + capture_output=True, + check=False, + text=True, + ) + if commit_process.returncode != 0 or status_process.returncode != 0: + msg = "Could not resolve the egglog-python Git state" + raise RuntimeError(msg) + return commit_process.stdout.strip(), not status_process.stdout.strip() + + +def generate_aggregates(binary_path: Path, container_path: Path, output_dir: Path) -> None: + """Validate paired rows and write the only publishable research artifacts.""" + binary_rows = load_raw(binary_path, expected_variant="binary") + container_rows = load_raw(container_path, expected_variant="container") + all_rows = [*binary_rows, *container_rows] + provenance = _validated_provenance(all_rows) + comparison = build_representation_comparison(binary_rows, container_rows) + paper = build_paper_replication(all_rows) + repository_commit, repository_clean = _repo_state() + if not repository_clean: + msg = "Refusing aggregate publication from a dirty egglog-python worktree" + raise ValueError(msg) + if len(repository_commit) != 40 or any( + character not in "0123456789abcdef" for character in repository_commit.lower() + ): + msg = "Could not identify the egglog-python commit" + raise ValueError(msg) + if repository_commit != provenance["egglog_python_commit"]: + msg = "Refusing aggregate publication from raw rows produced by a different egglog-python commit" + raise ValueError(msg) + archive_hashes = {row["external_archive_sha256"] for row in all_rows} + archive_hash = next(iter(archive_hashes)) if len(archive_hashes) == 1 else "mismatch" + if len(archive_hash) != 64 or any(character not in "0123456789abcdef" for character in archive_hash.lower()): + msg = "Raw rows have an invalid external_archive_sha256" + raise ValueError(msg) + configurations = {(row["timeout_sec"], row["memory_limit_mb"], row["sample_interval_sec"]) for row in all_rows} + if len(configurations) != 1: + msg = "Raw rows do not share one timeout/memory/sample configuration" + raise ValueError(msg) + timeout_sec, memory_limit_mb, sample_interval_sec = next(iter(configurations)) + + output_dir.mkdir(parents=True, exist_ok=True) + paper_path = output_dir / "paper-replication.csv" + comparison_path = output_dir / "representation-comparison.csv" + _write_csv(paper_path, PAPER_COLUMNS, paper) + _write_csv(comparison_path, COMPARISON_COLUMNS, comparison) + binary_statuses = Counter(row["status"] for row in binary_rows) + container_statuses = Counter(row["status"] for row in container_rows) + raw_layout = "paired-single-file" if binary_path.resolve() == container_path.resolve() else "split-files" + manifest = "\n".join([ + "schema_version = 5", + f"generated_utc = {json.dumps(datetime.now(UTC).isoformat())}", + f"repository_commit = {json.dumps(repository_commit)}", + f"repository_clean = {str(repository_clean).lower()}", + f"egglog_core_commit = {json.dumps(provenance['egglog_core_commit'])}", + f"egglog_experimental_commit = {json.dumps(provenance['egglog_experimental_commit'])}", + f"egglog_bindings_sha256 = {json.dumps(provenance['egglog_bindings_sha256'])}", + f"python = {json.dumps(provenance['python_version'])}", + f"rust = {json.dumps(provenance['rust_version'])}", + f"platform = {json.dumps(provenance['platform'])}", + f"cpu = {json.dumps(provenance['cpu'])}", + f"execution_mode = {json.dumps(provenance['execution_mode'])}", + f"workers_requested = {provenance['workers_requested']}", + f"workers_effective = {provenance['workers_effective']}", + f"ordering_seed = {json.dumps(provenance['ordering_seed'])}", + f"raw_layout = {json.dumps(raw_layout)}", + f"binary_raw_sha256 = {json.dumps(hashlib.sha256(binary_path.read_bytes()).hexdigest())}", + f"container_raw_sha256 = {json.dumps(hashlib.sha256(container_path.read_bytes()).hexdigest())}", + f"timeout_sec = {json.dumps(timeout_sec)}", + f"memory_limit_mb = {json.dumps(memory_limit_mb)}", + f"sample_interval_sec = {json.dumps(sample_interval_sec)}", + f"paper_replication_sha256 = {json.dumps(hashlib.sha256(paper_path.read_bytes()).hexdigest())}", + f"representation_comparison_sha256 = {json.dumps(hashlib.sha256(comparison_path.read_bytes()).hexdigest())}", + f"binary_rows = {len(binary_rows)}", + f"container_rows = {len(container_rows)}", + *[f"binary_status_{status} = {binary_statuses[status]}" for status in sorted(RAW_STATUSES)], + *[f"container_status_{status} = {container_statuses[status]}" for status in sorted(RAW_STATUSES)], + 'ratio_definition = "container / binary; lower is better; zero binary values omitted; n_ratio counts the remaining ratios"', + 'recommended_full_rerun = "make -C experiments/param_eq aggregate"', + "", + ]) + (output_dir / "manifest.toml").write_text(manifest, encoding="utf-8") + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--binary", type=Path, required=True) + parser.add_argument("--container", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, default=Path(__file__).resolve().parent / "results") + args = parser.parse_args(argv) + generate_aggregates(args.binary, args.container, args.output_dir) + print(f"wrote aggregate-only artifacts to {args.output_dir}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/experiments/param_eq/corpus.py b/experiments/param_eq/corpus.py new file mode 100644 index 00000000..d435bcc8 --- /dev/null +++ b/experiments/param_eq/corpus.py @@ -0,0 +1,175 @@ +"""Load the author-supplied Param-Eq corpus without copying it into this repository.""" + +from __future__ import annotations + +import csv +import hashlib +import os +from collections import defaultdict +from dataclasses import dataclass +from pathlib import Path + +DATASETS = ("pagie", "kotanchek") +RAW_ALGORITHMS = ("Bingo", "EPLEX", "FEAT", "GOMEA", "Operon", "SBP", "SRjl") +DROP_INDEXES: dict[str, set[int]] = {"pagie": {16, 162}, "kotanchek": {1}} +ALGORITHM_RENAMES = {"GOMEA": "GP-GOMEA", "SRjl": "PySR"} + + +class ArchiveLayoutError(ValueError): + """Raised when the external research archive is absent or incomplete.""" + + +@dataclass(frozen=True) +class CorpusRow: + """One in-memory source expression; source text is never written to tracked output.""" + + row_id: str + dataset: str + raw_index: int + algorithm_raw: str + algorithm: str + algorithm_row: int + input_kind: str + source: str + source_n_rank: float + + +def external_archive_root(root: Path | None = None) -> Path: + """Resolve and validate the external archive root.""" + if root is None: + configured = os.environ.get("EGGLOG_PARAM_EQ_DATA_DIR") + if not configured: + msg = "Set EGGLOG_PARAM_EQ_DATA_DIR to the private param-eq-haskell archive checkout." + raise ArchiveLayoutError(msg) + root = Path(configured) + resolved = root.expanduser().resolve() + required = [resolved / "results" / f"{dataset}_table_counts.csv" for dataset in DATASETS] + required.extend(resolved / "results" / f"{dataset}_results" for dataset in DATASETS) + required.extend(resolved / "results" / directory for directory in ("exprs", "exprs_simpl")) + missing = [path.relative_to(resolved).as_posix() for path in required if not path.exists()] + if missing: + msg = f"External Param-Eq archive is incomplete; missing: {', '.join(missing)}" + raise ArchiveLayoutError(msg) + return resolved + + +def should_keep_row(dataset: str, raw_index: int, algorithm: str, n_rank: str | None) -> bool: + """Apply the documented retained-paper cleaning policy.""" + return algorithm != "FEAT" and raw_index not in DROP_INDEXES[dataset] and bool((n_rank or "").strip()) + + +def _read_expression_lines(path: Path) -> list[str]: + return [line.strip() for line in path.read_text(encoding="utf-8").splitlines() if line.strip()] + + +def _read_sympy_by_algorithm(path: Path) -> dict[str, list[str]]: + grouped: defaultdict[str, list[str]] = defaultdict(list) + with path.open(newline="", encoding="utf-8") as handle: + for row in csv.DictReader(handle): + grouped[row["algorithm"].strip()].append(row["expr_sympy"].strip()) + return dict(grouped) + + +def _raw_index(row: dict[str, str | None]) -> int: + for key in ("raw_index", "", "Unnamed: 0"): + value = row.get(key) + if value is not None and value != "": + return int(value) + msg = "A table-count row has no raw index column." + raise ArchiveLayoutError(msg) + + +def load_corpus_rows( # noqa: C901 + root: Path | None = None, + *, + dataset: str | None = None, + algorithm: str | None = None, + input_kind: str | None = None, + limit: int | None = None, +) -> list[CorpusRow]: + """Load retained rows in stable order while keeping expressions in memory only.""" + archive = external_archive_root(root) + if dataset is not None and dataset not in DATASETS: + raise ValueError(f"Unknown dataset: {dataset}") + if input_kind is not None and input_kind not in {"original", "sympy"}: + raise ValueError(f"Unknown input kind: {input_kind}") + + rows: list[CorpusRow] = [] + for current_dataset in DATASETS: + if dataset is not None and current_dataset != dataset: + continue + results_root = archive / "results" + sympy_by_algorithm = _read_sympy_by_algorithm(results_root / f"{current_dataset}_results") + originals = { + raw_algorithm: _read_expression_lines(results_root / "exprs" / f"{raw_algorithm}_exprs_{current_dataset}") + for raw_algorithm in RAW_ALGORITHMS + } + algorithm_positions: defaultdict[str, int] = defaultdict(int) + with (results_root / f"{current_dataset}_table_counts.csv").open(newline="", encoding="utf-8") as handle: + for count_row in csv.DictReader(handle): + raw_index = _raw_index(count_row) + raw_algorithm = (count_row.get("algorithm") or "").strip() + n_rank = count_row.get("n_rank") + # The known dropped indexes have no corresponding expression + # line. Rows missing rank data do have a line, so advance the + # per-algorithm position before omitting them from the study. + if raw_algorithm == "FEAT" or raw_index in DROP_INDEXES[current_dataset]: + continue + algorithm_positions[raw_algorithm] += 1 + algorithm_row = algorithm_positions[raw_algorithm] + if not should_keep_row(current_dataset, raw_index, raw_algorithm, n_rank): + continue + public_algorithm = ALGORITHM_RENAMES.get(raw_algorithm, raw_algorithm) + if algorithm not in (None, raw_algorithm, public_algorithm): + continue + try: + source_by_kind = { + "original": originals[raw_algorithm][algorithm_row - 1], + "sympy": sympy_by_algorithm[raw_algorithm][algorithm_row - 1], + } + except (KeyError, IndexError) as exc: + msg = ( + f"External archive rows are misaligned for {current_dataset}/{raw_algorithm} " + f"at retained algorithm row {algorithm_row}." + ) + raise ArchiveLayoutError(msg) from exc + for current_kind in ("original", "sympy"): + if input_kind is not None and current_kind != input_kind: + continue + rows.append( + CorpusRow( + row_id=(f"{current_dataset}/{raw_index}/{public_algorithm}/{algorithm_row}/{current_kind}"), + dataset=current_dataset, + raw_index=raw_index, + algorithm_raw=raw_algorithm, + algorithm=public_algorithm, + algorithm_row=algorithm_row, + input_kind=current_kind, + source=source_by_kind[current_kind], + source_n_rank=float(n_rank or "nan"), + ) + ) + if limit is not None and len(rows) >= limit: + return rows + return rows + + +def external_archive_hash(root: Path | None = None) -> str: + """Hash corpus rows plus the optional live-Haskell source/toolchain inputs.""" + archive = external_archive_root(root) + files = {path for path in (archive / "results").rglob("*") if path.is_file()} + files.update(path for path in (archive / "src").rglob("*.hs") if path.is_file()) + files.update( + path + for name in ("rewrite.cabal", "cabal.project", "stack.yaml", "stack.yaml.lock") + if (path := archive / name).is_file() + ) + digest = hashlib.sha256() + for path in sorted(files): + relative = path.relative_to(archive).as_posix().encode() + payload = path.read_bytes() + digest.update(len(relative).to_bytes(8, "big")) + digest.update(relative) + digest.update(len(payload).to_bytes(8, "big")) + digest.update(payload) + return digest.hexdigest() diff --git a/experiments/param_eq/resource_guard.py b/experiments/param_eq/resource_guard.py new file mode 100644 index 00000000..109943a3 --- /dev/null +++ b/experiments/param_eq/resource_guard.py @@ -0,0 +1,144 @@ +"""Memory and timeout guards for isolated Param-Eq corpus workers.""" + +from __future__ import annotations + +import os +import signal +import subprocess +import time +from contextlib import suppress +from dataclasses import dataclass +from multiprocessing.process import BaseProcess + +SAFE_MEMORY_FRACTION = 0.75 +DEFAULT_MEMORY_LIMIT_MB = 2048 +DEFAULT_SAMPLE_INTERVAL_SEC = 0.2 + + +@dataclass(frozen=True) +class WatchResult: + status: str + peak_rss_mb: float | None + + +def total_system_memory_bytes() -> int: + if "SC_PAGE_SIZE" in os.sysconf_names and "SC_PHYS_PAGES" in os.sysconf_names: + return int(os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")) + return int(subprocess.check_output(["sysctl", "-n", "hw.memsize"], text=True).strip()) # noqa: S607 + + +def cap_workers_for_memory( + requested_workers: int, + *, + memory_limit_mb: int = DEFAULT_MEMORY_LIMIT_MB, + total_memory_bytes_value: int | None = None, + safe_memory_fraction: float = SAFE_MEMORY_FRACTION, +) -> int: + total = total_system_memory_bytes() if total_memory_bytes_value is None else total_memory_bytes_value + allowed_mb = total / (1024.0 * 1024.0) * safe_memory_fraction + return max(1, min(requested_workers, max(1, int(allowed_mb // memory_limit_mb)))) + + +def _rss_mb(pid: int) -> float | None: + completed = subprocess.run( + ["ps", "-o", "rss=", "-p", str(pid)], # noqa: S607 + capture_output=True, + check=False, + text=True, + ) + if completed.returncode != 0 or not completed.stdout.strip(): + return None + try: + return float(completed.stdout.strip()) / 1024.0 + except ValueError: + return None + + +def _process_tree_rss_mb(root_pid: int) -> float | None: + """Return aggregate RSS for a subprocess and all of its descendants.""" + completed = subprocess.run( + ["ps", "-axo", "pid=,ppid=,rss="], # noqa: S607 + capture_output=True, + check=False, + text=True, + ) + if completed.returncode != 0: + return None + parents: dict[int, int] = {} + rss_kb: dict[int, int] = {} + for line in completed.stdout.splitlines(): + try: + pid_text, parent_text, rss_text = line.split() + pid = int(pid_text) + parents[pid] = int(parent_text) + rss_kb[pid] = int(rss_text) + except ValueError: + continue + descendants = {root_pid} + changed = True + while changed: + changed = False + for pid, parent in parents.items(): + if parent in descendants and pid not in descendants: + descendants.add(pid) + changed = True + measured = [rss_kb[pid] for pid in descendants if pid in rss_kb] + return sum(measured) / 1024.0 if measured else None + + +def _kill_subprocess_group(process: subprocess.Popen[str]) -> None: + """Kill a subprocess started in its own session, including descendants.""" + with suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + process.wait(timeout=1.0) + + +def watch_process( + process: BaseProcess, + *, + timeout_sec: float, + memory_limit_mb: int, + sample_interval_sec: float = DEFAULT_SAMPLE_INTERVAL_SEC, +) -> WatchResult: + start = time.monotonic() + peak_rss_mb = None + while process.is_alive(): + if process.pid is not None: + rss_mb = _rss_mb(process.pid) + if rss_mb is not None: + peak_rss_mb = rss_mb if peak_rss_mb is None else max(peak_rss_mb, rss_mb) + if rss_mb > memory_limit_mb: + process.kill() + process.join(timeout=1.0) + return WatchResult("memory_limit", peak_rss_mb) + if time.monotonic() - start > timeout_sec: + process.kill() + process.join(timeout=1.0) + return WatchResult("timeout", peak_rss_mb) + time.sleep(sample_interval_sec) + process.join(timeout=1.0) + return WatchResult("completed", peak_rss_mb) + + +def watch_subprocess( + process: subprocess.Popen[str], + *, + timeout_sec: float, + memory_limit_mb: int, + sample_interval_sec: float = DEFAULT_SAMPLE_INTERVAL_SEC, +) -> WatchResult: + """Apply a timeout/RSS boundary to an external process tree.""" + start = time.monotonic() + peak_rss_mb = None + while process.poll() is None: + rss_mb = _process_tree_rss_mb(process.pid) + if rss_mb is not None: + peak_rss_mb = rss_mb if peak_rss_mb is None else max(peak_rss_mb, rss_mb) + if rss_mb > memory_limit_mb: + _kill_subprocess_group(process) + return WatchResult("memory_limit", peak_rss_mb) + if time.monotonic() - start > timeout_sec: + _kill_subprocess_group(process) + return WatchResult("timeout", peak_rss_mb) + time.sleep(sample_interval_sec) + return WatchResult("completed", peak_rss_mb) diff --git a/experiments/param_eq/results/manifest.toml b/experiments/param_eq/results/manifest.toml new file mode 100644 index 00000000..2cc06be6 --- /dev/null +++ b/experiments/param_eq/results/manifest.toml @@ -0,0 +1,6 @@ +schema_version = 5 +status = "paused-before-final-dependency-compatible-rerun" +paper_replication_rows = 0 +representation_comparison_rows = 0 +note = "Run make -C experiments/param_eq aggregate after rebuilding the final dependency stack." +archive_identity = "Canonical archive hash is kept in private research records and required by the runner." diff --git a/experiments/param_eq/results/paper-replication.csv b/experiments/param_eq/results/paper-replication.csv new file mode 100644 index 00000000..e4056960 --- /dev/null +++ b/experiments/param_eq/results/paper-replication.csv @@ -0,0 +1 @@ +implementation,dataset,algorithm,input_kind,metric,n_total,n_success,n_iteration_limit,n_timeout,n_memory_limit,n_error,value,min,q1,median,q3,max diff --git a/experiments/param_eq/results/raw/.gitignore b/experiments/param_eq/results/raw/.gitignore new file mode 100644 index 00000000..d6b7ef32 --- /dev/null +++ b/experiments/param_eq/results/raw/.gitignore @@ -0,0 +1,2 @@ +* +!.gitignore diff --git a/experiments/param_eq/results/representation-comparison.csv b/experiments/param_eq/results/representation-comparison.csv new file mode 100644 index 00000000..5048014f --- /dev/null +++ b/experiments/param_eq/results/representation-comparison.csv @@ -0,0 +1 @@ +slice,metric,n_pairs,n_ratio,n_binary_missing,n_container_missing,container_better,same,container_worse,ratio_p10,ratio_p25,ratio_median,ratio_p75,ratio_p90 diff --git a/experiments/param_eq/run.py b/experiments/param_eq/run.py new file mode 100644 index 00000000..a823adfe --- /dev/null +++ b/experiments/param_eq/run.py @@ -0,0 +1,720 @@ +"""Run external Param-Eq corpus rows in isolated local workers.""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import json +import math +import os +import platform +import shutil +import subprocess +import tempfile +from collections import Counter +from collections.abc import Iterable, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor, as_completed +from contextlib import ExitStack +from multiprocessing import get_context +from multiprocessing.connection import Connection +from pathlib import Path +from typing import Any + +from experiments.param_eq.corpus import CorpusRow, external_archive_hash, external_archive_root, load_corpus_rows +from experiments.param_eq.resource_guard import ( + DEFAULT_MEMORY_LIMIT_MB, + DEFAULT_SAMPLE_INTERVAL_SEC, + cap_workers_for_memory, + watch_process, + watch_subprocess, +) + +from egglog import bindings as egglog_bindings + +RESULTS_DIR = Path(__file__).resolve().parent / "results" +RAW_RESULTS_DIR = RESULTS_DIR / "raw" +REPO_ROOT = Path(__file__).resolve().parents[2] +VARIANT_ORDER_SEED = "param-eq-variant-order-v1" +DEFAULT_HASKELL_BUILD_TIMEOUT_SEC = 600.0 +RAW_COLUMNS = ( + "row_id", + "dataset", + "algorithm", + "input_kind", + "implementation", + "variant", + "external_archive_sha256", + "source_n_rank", + "timeout_sec", + "memory_limit_mb", + "sample_interval_sec", + "execution_mode", + "workers_requested", + "workers_effective", + "ordering_seed", + "egglog_python_commit", + "egglog_python_clean", + "egglog_core_commit", + "egglog_core_clean", + "egglog_experimental_commit", + "egglog_experimental_clean", + "egglog_bindings_sha256", + "python_version", + "platform", + "rust_version", + "cpu", + "status", + "runtime_ms", + "peak_rss_mb", + "passes", + "total_size", + "before_nodes", + "before_params", + "after_nodes", + "after_params", +) + + +def _validate_haskell_checkout(archive_root: Path) -> None: + """Require the source modules and Stack configuration used by live mode.""" + required = [ + archive_root / "stack.yaml", + *( + archive_root / "src" / name + for name in ( + "FixTree.hs", + "KotanchekSR.hs", + "KotanchekSympy.hs", + "PagieSR.hs", + "PagieSympy.hs", + "Reparam.hs", + ) + ), + ] + missing = [path.relative_to(archive_root).as_posix() for path in required if not path.is_file()] + if missing: + raise ValueError(f"External Haskell checkout is missing: {', '.join(missing)}") + + +def _build_haskell_program() -> str: + """Build a reusable forced live-Haskell runner without embedding source text.""" + return "\n".join([ # noqa: FLY002 + "import Control.Exception (evaluate)", + "import Data.List (intercalate)", + "import qualified Data.Map as M", + "import Data.SRTree", + "import FixTree", + "import KotanchekSR (kotanchekSR)", + "import KotanchekSympy (kotanchekSympy)", + "import PagieSR (pagieSR)", + "import PagieSympy (pagieSympy)", + "import Reparam (replaceConstsWithParams)", + "import Data.Time.Clock.POSIX (getPOSIXTime)", + "import System.Environment (getArgs)", + "", + "lookupExpr :: String -> String -> String -> Int -> SRTree Int Double", + "lookupExpr dataset inputKind algorithm rowIndex = case (dataset, inputKind) of", + ' ("pagie", "original") -> (pagieSR M.! algorithm) !! rowIndex', + ' ("pagie", "sympy") -> (pagieSympy M.! algorithm) !! rowIndex', + ' ("kotanchek", "original") -> (kotanchekSR M.! algorithm) !! rowIndex', + ' ("kotanchek", "sympy") -> (kotanchekSympy M.! algorithm) !! rowIndex', + ' _ -> error "unknown dataset/input kind"', + "", + "emitExpr :: SRTree Int Double -> IO ()", + "emitExpr expr = do", + " beforeNodes <- evaluate (countNodes expr)", + " beforeParams <- evaluate (recountParams (replaceConstsWithParams expr))", + " start <- getPOSIXTime", + " let simplified = simplifyE expr", + " afterNodes <- evaluate (countNodes simplified)", + " afterParams <- evaluate (recountParams (replaceConstsWithParams simplified))", + " end <- getPOSIXTime", + " let runtimeMs = (realToFrac (end - start) :: Double) * 1000.0", + " fields = map show [fromIntegral beforeNodes, fromIntegral beforeParams,", + " fromIntegral afterNodes, fromIntegral afterParams, runtimeMs]", + ' putStrLn (intercalate "\\t" fields)', + "", + "main :: IO ()", + "main = do", + " args <- getArgs", + " case args of", + " [dataset, inputKind, algorithm, rowIndex] ->", + " emitExpr (lookupExpr dataset inputKind algorithm (read rowIndex))", + ' _ -> error "expected dataset input-kind algorithm zero-based-row-index"', + "", + ]) + + +def _parse_haskell_output(stdout: str) -> dict[str, Any]: + """Parse the expression-free, fully forced Haskell result row.""" + lines = [line for line in stdout.splitlines() if line.strip()] + if len(lines) != 1: + raise ValueError(f"Expected one Haskell output row, got {len(lines)}") + fields = lines[0].split("\t") + if len(fields) != 5: + raise ValueError(f"Expected five Haskell output fields, got {len(fields)}") + before_nodes, before_params, after_nodes, after_params, runtime_ms = map(float, fields) + counts = (before_nodes, before_params, after_nodes, after_params) + if ( + any(not math.isfinite(value) or value < 0.0 or not value.is_integer() for value in counts) + or not math.isfinite(runtime_ms) + or runtime_ms < 0.0 + ): + msg = "Haskell output contains invalid counts or runtime" + raise ValueError(msg) + return { + "status": "saturated", + "runtime_ms": runtime_ms, + "passes": "", + "total_size": "", + "before_nodes": before_nodes, + "before_params": before_params, + "after_nodes": after_nodes, + "after_params": after_params, + } + + +def _compile_haskell_runner( + archive_root: Path, + build_root: Path, + *, + execution_mode: str, + timeout_sec: float, + memory_limit_mb: int, + sample_interval_sec: float, +) -> Path: + """Compile the generated runner once so per-row timings exclude startup.""" + source_path = build_root / "ParamEqRunner.hs" + object_dir = build_root / "objects" + executable = build_root / "param-eq-haskell-runner" + object_dir.mkdir() + source_path.write_text(_build_haskell_program(), encoding="utf-8") + optimization = "-O2" if execution_mode == "release" else "-O0" + stack = shutil.which("stack") + if stack is None: + msg = "The live-Haskell runner requires Stack on PATH" + raise RuntimeError(msg) + process = subprocess.Popen( + [ + stack, + "ghc", + "--", + optimization, + "-rtsopts", + "-isrc", + "-outputdir", + str(object_dir), + "-o", + str(executable), + str(source_path), + ], + cwd=archive_root, + text=True, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=True, + ) + watched = watch_subprocess( + process, + timeout_sec=timeout_sec, + memory_limit_mb=memory_limit_mb, + sample_interval_sec=sample_interval_sec, + ) + if watched.status != "completed": + raise RuntimeError(f"Haskell runner compilation reached the {watched.status} guard") + if process.returncode != 0 or not executable.is_file(): + msg = "Could not compile the live-Haskell runner; run `stack build` in the external checkout first" + raise RuntimeError(msg) + return executable + + +def _collect_run_provenance(*, execution_mode: str, workers_requested: int, workers_effective: int) -> dict[str, str]: + """Resolve the local source checkouts, loaded extension, and machine state.""" + metadata_process = subprocess.run( + ["cargo", "metadata", "--format-version", "1", "--locked"], # noqa: S607 + cwd=REPO_ROOT, + capture_output=True, + check=False, + text=True, + ) + if metadata_process.returncode != 0: + raise RuntimeError(f"Could not resolve Cargo dependencies: {metadata_process.stderr.strip()}") + packages = json.loads(metadata_process.stdout)["packages"] + manifests = {package["name"]: Path(package["manifest_path"]) for package in packages} + repositories = {"egglog_python": REPO_ROOT} + for package_name, prefix in (("egglog", "egglog_core"), ("egglog-experimental", "egglog_experimental")): + try: + repositories[prefix] = manifests[package_name].parent + except KeyError as error: + raise RuntimeError(f"Cargo metadata has no {package_name!r} dependency") from error + dependency_values: dict[str, str] = {} + for prefix, worktree in repositories.items(): + commit_process = subprocess.run( + ["git", "rev-parse", "HEAD"], # noqa: S607 + cwd=worktree, + capture_output=True, + check=False, + text=True, + ) + status_process = subprocess.run( + ["git", "status", "--porcelain"], # noqa: S607 + cwd=worktree, + capture_output=True, + check=False, + text=True, + ) + if commit_process.returncode != 0 or status_process.returncode != 0: + raise RuntimeError(f"Could not resolve Git state for {prefix!r}") + dependency_values[f"{prefix}_commit"] = commit_process.stdout.strip() + dependency_values[f"{prefix}_clean"] = str(not status_process.stdout.strip()).lower() + + rust_process = subprocess.run( + ["rustc", "--version"], # noqa: S607 + capture_output=True, + check=False, + text=True, + ) + if rust_process.returncode != 0 or not rust_process.stdout.strip(): + msg = "Could not resolve the Rust compiler version" + raise RuntimeError(msg) + bindings_file = getattr(egglog_bindings, "__file__", None) + if bindings_file is None or not Path(bindings_file).is_file(): + msg = "Could not identify the loaded egglog native extension" + raise RuntimeError(msg) + with Path(bindings_file).open("rb") as artifact: + bindings_sha256 = hashlib.file_digest(artifact, "sha256").hexdigest() + cpu = platform.processor() + if platform.system() == "Darwin": + cpu_process = subprocess.run( + ["sysctl", "-n", "machdep.cpu.brand_string"], # noqa: S607 + capture_output=True, + check=False, + text=True, + ) + if cpu_process.returncode == 0 and cpu_process.stdout.strip(): + cpu = cpu_process.stdout.strip() + return { + "execution_mode": execution_mode, + "workers_requested": str(workers_requested), + "workers_effective": str(workers_effective), + "ordering_seed": VARIANT_ORDER_SEED, + **dependency_values, + "egglog_bindings_sha256": bindings_sha256, + "python_version": platform.python_version(), + "platform": platform.platform(), + "rust_version": rust_process.stdout.strip(), + "cpu": cpu or platform.machine(), + } + + +def _worker(connection: Connection, source: str, variant: str) -> None: + try: + from egglog.exp.param_eq import ( # noqa: PLC0415 + binary_to_containers, + parse_expression, + run_paper_pipeline, + run_paper_pipeline_container, + ) + + report = ( + run_paper_pipeline(parse_expression(source)) + if variant == "binary" + else run_paper_pipeline_container(binary_to_containers(parse_expression(source))) + ) + payload: dict[str, Any] = {"status": report.status} + if report.status == "saturated": + payload.update({ + "runtime_ms": report.total_sec * 1000.0, + "passes": report.passes, + "total_size": report.total_size, + "before_nodes": report.before_nodes, + "before_params": report.before_params, + "after_nodes": report.extracted_nodes, + "after_params": report.extracted_params, + }) + connection.send(payload) + except BaseException: # worker errors are accounted for without publishing private input text + connection.send({"status": "error"}) + finally: + connection.close() + + +def _raw_result( + row: CorpusRow, + implementation: str, + variant: str, + *, + external_archive_sha256: str, + timeout_sec: float, + memory_limit_mb: int, + sample_interval_sec: float, + provenance: Mapping[str, str], + payload: Mapping[str, Any], + peak_rss_mb: float | None, +) -> dict[str, Any]: + """Build the expression-free raw schema shared by both implementations.""" + return { + "row_id": row.row_id, + "dataset": row.dataset, + "algorithm": row.algorithm, + "input_kind": row.input_kind, + "implementation": implementation, + "variant": variant, + "external_archive_sha256": external_archive_sha256, + "source_n_rank": row.source_n_rank, + "timeout_sec": timeout_sec, + "memory_limit_mb": memory_limit_mb, + "sample_interval_sec": sample_interval_sec, + **provenance, + "status": payload["status"], + "runtime_ms": payload.get("runtime_ms", ""), + "peak_rss_mb": peak_rss_mb if peak_rss_mb is not None else "", + "passes": payload.get("passes", ""), + "total_size": payload.get("total_size", ""), + "before_nodes": payload.get("before_nodes", ""), + "before_params": payload.get("before_params", ""), + "after_nodes": payload.get("after_nodes", ""), + "after_params": payload.get("after_params", ""), + } + + +def _run_egglog_one( + row: CorpusRow, + variant: str, + *, + external_archive_sha256: str, + timeout_sec: float, + memory_limit_mb: int, + sample_interval_sec: float, + provenance: Mapping[str, str], +) -> dict[str, Any]: + context = get_context("spawn") + parent, child = context.Pipe(duplex=False) + process = context.Process(target=_worker, args=(child, row.source, variant)) + process.start() + child.close() + watched = watch_process( + process, + timeout_sec=timeout_sec, + memory_limit_mb=memory_limit_mb, + sample_interval_sec=sample_interval_sec, + ) + payload: dict[str, Any] + try: + if watched.status != "completed": + payload = {"status": watched.status} + elif parent.poll(): + try: + payload = parent.recv() + except EOFError: + payload = {"status": "error"} + else: + payload = {"status": "error"} + finally: + parent.close() + return _raw_result( + row, + "egglog", + variant, + external_archive_sha256=external_archive_sha256, + timeout_sec=timeout_sec, + memory_limit_mb=memory_limit_mb, + sample_interval_sec=sample_interval_sec, + provenance=provenance, + payload=payload, + peak_rss_mb=watched.peak_rss_mb, + ) + + +def _run_haskell_one( + row: CorpusRow, + *, + archive_root: Path, + executable: Path, + external_archive_sha256: str, + timeout_sec: float, + memory_limit_mb: int, + sample_interval_sec: float, + provenance: Mapping[str, str], +) -> dict[str, Any]: + peak_rss_mb = None + payload: dict[str, Any] = {"status": "error"} + try: + process = subprocess.Popen( + [ + str(executable), + row.dataset, + row.input_kind, + row.algorithm_raw, + str(row.algorithm_row - 1), + "+RTS", + "-K3G", + "-RTS", + ], + cwd=archive_root, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + start_new_session=True, + ) + watched = watch_subprocess( + process, + timeout_sec=timeout_sec, + memory_limit_mb=memory_limit_mb, + sample_interval_sec=sample_interval_sec, + ) + stdout, _stderr = process.communicate() + peak_rss_mb = watched.peak_rss_mb + if watched.status != "completed": + payload = {"status": watched.status} + elif process.returncode == 0: + payload = _parse_haskell_output(stdout) + except (OSError, ValueError): + pass + return _raw_result( + row, + "haskell", + "binary", + external_archive_sha256=external_archive_sha256, + timeout_sec=timeout_sec, + memory_limit_mb=memory_limit_mb, + sample_interval_sec=sample_interval_sec, + provenance=provenance, + payload=payload, + peak_rss_mb=peak_rss_mb, + ) + + +def _variant_order(row_id: str, requested: str) -> tuple[str, ...]: + if requested != "both": + return (requested,) + digest = hashlib.sha256(f"{VARIANT_ORDER_SEED}\0{row_id}".encode()).digest() + return ("binary", "container") if digest[0] % 2 == 0 else ("container", "binary") + + +def _run_row( + row: CorpusRow, + *, + implementation: str, + variant: str, + archive_root: Path, + haskell_executable: Path | None, + external_archive_sha256: str, + timeout_sec: float, + memory_limit_mb: int, + sample_interval_sec: float, + provenance: Mapping[str, str], +) -> list[dict[str, Any]]: + """Keep paired variants sequential while parallelizing independent rows.""" + if implementation == "haskell": + if haskell_executable is None: + msg = "The Haskell implementation requires a compiled runner" + raise ValueError(msg) + return [ + _run_haskell_one( + row, + archive_root=archive_root, + executable=haskell_executable, + external_archive_sha256=external_archive_sha256, + timeout_sec=timeout_sec, + memory_limit_mb=memory_limit_mb, + sample_interval_sec=sample_interval_sec, + provenance=provenance, + ) + ] + return [ + _run_egglog_one( + row, + current_variant, + external_archive_sha256=external_archive_sha256, + timeout_sec=timeout_sec, + memory_limit_mb=memory_limit_mb, + sample_interval_sec=sample_interval_sec, + provenance=provenance, + ) + for current_variant in _variant_order(row.row_id, variant) + ] + + +def run_rows( + rows: Iterable[CorpusRow], + *, + implementation: str, + variant: str, + archive_root: Path, + haskell_executable: Path | None, + external_archive_sha256: str, + workers: int, + timeout_sec: float, + memory_limit_mb: int, + sample_interval_sec: float = DEFAULT_SAMPLE_INTERVAL_SEC, + provenance: Mapping[str, str], +) -> list[dict[str, Any]]: + """Run rows with stable order balancing and explicit failure accounting.""" + results: list[dict[str, Any]] = [] + with ThreadPoolExecutor(max_workers=workers) as executor: + futures = { + executor.submit( + _run_row, + row, + implementation=implementation, + variant=variant, + archive_root=archive_root, + haskell_executable=haskell_executable, + external_archive_sha256=external_archive_sha256, + timeout_sec=timeout_sec, + memory_limit_mb=memory_limit_mb, + sample_interval_sec=sample_interval_sec, + provenance=provenance, + ): row.row_id + for row in rows + } + for future in as_completed(futures): + results.extend(future.result()) + return sorted(results, key=lambda item: (str(item["row_id"]), str(item["implementation"]), str(item["variant"]))) + + +def write_local_raw(rows: Iterable[dict[str, Any]], path: Path) -> None: + """Write untracked row metrics; expressions are never included.""" + resolved_path = path.expanduser().resolve() + resolved_parent = resolved_path.parent + if resolved_parent != RAW_RESULTS_DIR.resolve(): + msg = f"Row-level output must stay in the ignored directory {RAW_RESULTS_DIR}" + raise ValueError(msg) + relative_path = resolved_path.relative_to(REPO_ROOT.resolve()).as_posix() + git = shutil.which("git") + if git is None: + msg = "Git is required to verify that row-level output stays untracked" + raise RuntimeError(msg) + is_tracked = ( + subprocess.run( + [git, "ls-files", "--error-unmatch", "--", relative_path], + cwd=REPO_ROOT, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + check=False, + ).returncode + == 0 + ) + is_ignored = ( + subprocess.run( + [git, "check-ignore", "--quiet", "--", relative_path], + cwd=REPO_ROOT, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + check=False, + ).returncode + == 0 + ) + if is_tracked or not is_ignored: + msg = "Row-level output must be an untracked, Git-ignored file" + raise ValueError(msg) + resolved_parent.mkdir(parents=True, exist_ok=True) + with resolved_path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=RAW_COLUMNS) + writer.writeheader() + writer.writerows(rows) + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--implementation", choices=("egglog", "haskell"), default="egglog") + parser.add_argument("--variant", choices=("binary", "container", "both"), default="both") + parser.add_argument("--dataset", choices=("pagie", "kotanchek")) + parser.add_argument("--algorithm") + parser.add_argument("--input-kind", choices=("original", "sympy")) + parser.add_argument("--limit", type=int) + parser.add_argument("--workers", type=int, default=1) + parser.add_argument("--timeout-sec", type=float, default=60.0) + parser.add_argument("--memory-limit-mb", type=int, default=DEFAULT_MEMORY_LIMIT_MB) + parser.add_argument("--haskell-build-timeout-sec", type=float, default=DEFAULT_HASKELL_BUILD_TIMEOUT_SEC) + parser.add_argument("--execution-mode", choices=("debug", "release"), required=True) + parser.add_argument("--output", type=Path) + args = parser.parse_args(argv) + nonpositive = next( + ( + name + for name, value in { + "--workers": args.workers, + "--limit": args.limit, + "--timeout-sec": args.timeout_sec, + "--memory-limit-mb": args.memory_limit_mb, + "--haskell-build-timeout-sec": args.haskell_build_timeout_sec, + }.items() + if value is not None and value <= 0 + ), + None, + ) + if nonpositive is not None: + parser.error(f"{nonpositive} must be positive") + if args.implementation == "haskell" and args.variant != "binary": + parser.error("The live-Haskell implementation supports only --variant binary") + output = args.output or RAW_RESULTS_DIR / f"{args.implementation}-{args.variant}.csv" + archive_root = external_archive_root() + if args.implementation == "haskell": + try: + _validate_haskell_checkout(archive_root) + except ValueError as error: + parser.error(str(error)) + rows = load_corpus_rows( + archive_root, + dataset=args.dataset, + algorithm=args.algorithm, + input_kind=args.input_kind, + limit=args.limit, + ) + archive_sha256 = external_archive_hash(archive_root) + expected_archive_sha256 = os.environ.get("EGGLOG_PARAM_EQ_EXPECTED_ARCHIVE_SHA256", "").lower() + if len(expected_archive_sha256) != 64 or any( + character not in "0123456789abcdef" for character in expected_archive_sha256 + ): + parser.error("Set EGGLOG_PARAM_EQ_EXPECTED_ARCHIVE_SHA256 to the privately recorded archive hash") + if archive_sha256.lower() != expected_archive_sha256: + parser.error("The external Param-Eq archive does not match EGGLOG_PARAM_EQ_EXPECTED_ARCHIVE_SHA256") + if not rows: + parser.error("The selected corpus filters matched no rows") + effective_workers = cap_workers_for_memory(args.workers, memory_limit_mb=args.memory_limit_mb) + provenance = _collect_run_provenance( + execution_mode=args.execution_mode, + workers_requested=args.workers, + workers_effective=effective_workers, + ) + with ExitStack() as stack: + haskell_executable = None + if args.implementation == "haskell": + build_root = Path(stack.enter_context(tempfile.TemporaryDirectory(prefix="param-eq-haskell-"))) + try: + haskell_executable = _compile_haskell_runner( + archive_root, + build_root, + execution_mode=args.execution_mode, + timeout_sec=args.haskell_build_timeout_sec, + memory_limit_mb=args.memory_limit_mb, + sample_interval_sec=DEFAULT_SAMPLE_INTERVAL_SEC, + ) + except (OSError, RuntimeError) as error: + parser.error(str(error)) + results = run_rows( + rows, + implementation=args.implementation, + variant=args.variant, + archive_root=archive_root, + haskell_executable=haskell_executable, + external_archive_sha256=archive_sha256, + workers=effective_workers, + timeout_sec=args.timeout_sec, + memory_limit_mb=args.memory_limit_mb, + provenance=provenance, + ) + write_local_raw(results, output) + counts = Counter(str(row["status"]) for row in results) + print( + f"wrote {len(results)} rows to {output}; workers={effective_workers}/{args.workers}; " + f"statuses={dict(sorted(counts.items()))}" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/pyproject.toml b/pyproject.toml index b4a3724d..a48f7fb8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ classifiers = [ "Topic :: Software Development :: Interpreters", "Typing :: Typed", ] -dependencies = ["typing-extensions", "black", "graphviz", "anywidget", "cloudpickle>=3", "opentelemetry-api"] +dependencies = ["typing-extensions>=4.13", "black", "graphviz", "anywidget", "cloudpickle>=3", "opentelemetry-api"] [project.optional-dependencies] @@ -240,16 +240,19 @@ preview = true [tool.mypy] ignore_missing_imports = true -allow_redefinition = true +allow_redefinition_new = true exclude = ["__snapshots__", "_build", "^conftest.py$"] warn_unused_configs = true disallow_subclassing_any = true check_untyped_defs = true warn_redundant_casts = true warn_unused_ignores = true +warn_unreachable = true strict_equality = true extra_checks = true strict_equality_for_none = true +enable_incomplete_feature = ["TypeForm"] +local_partial_types = true [tool.maturin] python-source = "python" @@ -260,7 +263,10 @@ features = ["pyo3/extension-module"] addopts = ["--import-mode=importlib", "--doctest-modules"] testpaths = ["python"] python_files = ["test_*.py", "test.py"] -markers = ["slow: marks tests as slow (deselect with '-m \"not slow\"')"] +markers = [ + "slow: marks tests as slow (deselect with '-m \"not slow\"')", + "param_eq_smoke: bounded Param-Eq end-to-end binary/container cases", +] norecursedirs = ["__snapshots__"] filterwarnings = [ "ignore::numba.core.errors.NumbaPerformanceWarning", diff --git a/python/egglog/bindings.pyi b/python/egglog/bindings.pyi index 7199a065..a69ec9b9 100644 --- a/python/egglog/bindings.pyi +++ b/python/egglog/bindings.pyi @@ -14,6 +14,7 @@ __all__ = [ "Change", "Check", "Constructor", + "ContainerRebuildSpec", "CostModel", "Datatype", "Datatypes", @@ -45,6 +46,7 @@ __all__ = [ "IterationReport", "Let", "Lit", + "Naive", "NewSort", "Output", "OverallStatistics", @@ -58,6 +60,7 @@ __all__ = [ "PrintFunctionSize", "PrintOverallStatistics", "PrintSize", + "ProofConstructorNames", "Prove", "ProveExists", "ProveExistsOutput", @@ -79,6 +82,7 @@ __all__ = [ "Saturate", "Scan", "Schema", + "Seminaive", "Sequence", "SerializedEGraph", "Set", @@ -97,6 +101,7 @@ __all__ = [ "TimeOnly", "Union", "Unit", + "UnsafeSeminaive", "UnstableCombinedRuleset", "UserDefined", "UserDefinedCommandOutput", @@ -131,10 +136,18 @@ class EGraph: cls, *, fact_directory: str | Path | None = None, seminaive: bool = True, record: bool = False ) -> EGraph: ... def parse_program(self, __input: str, /, filename: str | None = None) -> list[_Command]: ... - def commands(self) -> str | None: ... + def parse_and_run_program( + self, + __input: str, + /, + filename: str | None = None, + traceparent: str | None = None, + tracestate: str | None = None, + ) -> list[_CommandOutput]: ... def run_program( self, *commands: _Command, traceparent: str | None = None, tracestate: str | None = None ) -> list[_CommandOutput]: ... + def commands(self) -> str | None: ... def serialize( self, root_eclasses: list[_Expr], @@ -147,6 +160,8 @@ class EGraph: ) -> SerializedEGraph: ... def set_report_level(self, level: _ReportLevel) -> None: ... def lookup_function(self, name: str, key: list[Value]) -> Value | None: ... + # `sort` must match the runtime sort returned with `value` by `eval_expr`. + def extract_value(self, value: Value, sort: str) -> tuple[TermDag, int, int]: ... def eval_expr( self, expr: _Expr, *, traceparent: str | None = None, tracestate: str | None = None ) -> tuple[str, Value]: ... @@ -169,11 +184,11 @@ class EGraph: @final class Value: def __hash__(self) -> int: ... - def __eq__(self, value: object) -> bool: ... - def __lt__(self, other: object) -> bool: ... - def __le__(self, other: object) -> bool: ... - def __gt__(self, other: object) -> bool: ... - def __ge__(self, other: object) -> bool: ... + def __eq__(self, value: object, /) -> bool: ... + def __lt__(self, other: object, /) -> bool: ... + def __le__(self, other: object, /) -> bool: ... + def __gt__(self, other: object, /) -> bool: ... + def __ge__(self, other: object, /) -> bool: ... @final class EggSmolError(Exception): @@ -395,7 +410,20 @@ class Rule: body: list[_Fact] name: str ruleset: str - def __new__(cls, span: _Span, head: list[_Action], body: list[_Fact], name: str, ruleset: str) -> Rule: ... + eval_mode: _RuleEvalMode + no_decomp: bool + include_subsumed: bool + def __new__( + cls, + span: _Span, + head: list[_Action], + body: list[_Fact], + name: str, + ruleset: str, + eval_mode: _RuleEvalMode = ..., + no_decomp: bool = ..., + include_subsumed: bool = ..., + ) -> Rule: ... @final class Rewrite: @@ -518,6 +546,7 @@ class Function: class RunReport: iterations: list[IterationReport] updated: bool + can_stop: bool search_and_apply_time_per_rule: dict[str, timedelta] num_matches_per_rule: dict[str, int] search_and_apply_time_per_ruleset: dict[str, timedelta] @@ -528,6 +557,7 @@ class RunReport: cls, iterations: list[IterationReport], updated: bool, + can_stop: bool, search_and_apply_time_per_rule: dict[str, timedelta], num_matches_per_rule: dict[str, int], search_and_apply_time_per_ruleset: dict[str, timedelta], @@ -616,6 +646,17 @@ class CSVPrintFunctionMode: ... _PrintFunctionMode: TypeAlias = DefaultPrintFunctionMode | CSVPrintFunctionMode +@final +class Seminaive: ... + +@final +class Naive: ... + +@final +class UnsafeSeminaive: ... + +_RuleEvalMode: TypeAlias = Seminaive | Naive | UnsafeSeminaive + ## # Schedules ## @@ -681,12 +722,41 @@ class Datatypes: datatypes: list[tuple[_Span, str, _Subdatatypes]] def __new__(cls, span: _Span, datatypes: list[tuple[_Span, str, _Subdatatypes]]) -> Datatypes: ... +@final +class ContainerRebuildSpec: + internal_rebuild_prim: str + internal_rebuild_proof_prim: str | None + def __new__( + cls, internal_rebuild_prim: str, internal_rebuild_proof_prim: str | None = ... + ) -> ContainerRebuildSpec: ... + +@final +class ProofConstructorNames: + congr: str + trans: str + sym: str + normalize: str + def __new__(cls, congr: str, trans: str, sym: str, normalize: str) -> ProofConstructorNames: ... + @final class Sort: span: _Span name: str presort_and_args: tuple[str, list[_Expr]] | None - def __new__(cls, span: _Span, name: str, presort_and_args: tuple[str, list[_Expr]] | None) -> Sort: ... + uf: tuple[str, str | None] | None + proof_func: str | None + container_rebuild: ContainerRebuildSpec | None + proof_constructors: ProofConstructorNames | None + def __new__( + cls, + span: _Span, + name: str, + presort_and_args: tuple[str, list[_Expr]] | None, + uf: tuple[str, str | None] | None = ..., + proof_func: str | None = ..., + container_rebuild: ContainerRebuildSpec | None = ..., + proof_constructors: ProofConstructorNames | None = ..., + ) -> Sort: ... @final class FunctionCommand: @@ -694,7 +764,21 @@ class FunctionCommand: name: str schema: Schema merge: _Expr | None - def __new__(cls, span: _Span, name: str, schema: Schema, merge: _Expr | None) -> FunctionCommand: ... + term_constructor: str | None + unextractable: bool + hidden: bool + let_binding: bool + def __new__( + cls, + span: _Span, + name: str, + schema: Schema, + merge: _Expr | None, + term_constructor: str | None = ..., + unextractable: bool = ..., + hidden: bool = ..., + let_binding: bool = ..., + ) -> FunctionCommand: ... @final class AddRuleset: @@ -826,7 +910,18 @@ class Constructor: schema: Schema cost: int | None unextractable: bool - def __new__(cls, span: _Span, name: str, schema: Schema, cost: int | None, unextractable: bool) -> Constructor: ... + hidden: bool + let_binding: bool + def __new__( + cls, + span: _Span, + name: str, + schema: Schema, + cost: int | None, + unextractable: bool, + hidden: bool = ..., + let_binding: bool = ..., + ) -> Constructor: ... @final class PrintOverallStatistics: diff --git a/python/egglog/builtins.py b/python/egglog/builtins.py index 5f11934b..126be588 100644 --- a/python/egglog/builtins.py +++ b/python/egglog/builtins.py @@ -19,8 +19,16 @@ from .conversion import convert, converter, get_type_args, resolve_literal from .declarations import * from .deconstruct import get_callable_args, get_literal_value -from .egraph import BaseExpr, BuiltinExpr, _add_default_rewrite_inner, expr_fact, function, get_current_ruleset, method -from .runtime import RuntimeExpr, RuntimeFunction, resolve_type_annotation_mutate +from .egraph import ( + BaseExpr, + BuiltinExpr, + expr_fact, + function, + method, + set_current_ruleset, + to_runtime_expr, +) +from .runtime import RuntimeClass, RuntimeExpr, RuntimeFunction, resolve_type_annotation_mutate from .thunk import Thunk if TYPE_CHECKING: @@ -38,8 +46,10 @@ "ExprValueError", "Map", "MapLike", + "Maybe", "MultiSet", "MultiSetLike", + "Pair", "Primitive", "PyObject", "Rational", @@ -51,11 +61,16 @@ "UnstableFn", "Vec", "VecLike", + "catch", "f64", "f64Like", "i64", "i64Like", "join", + "map_filter_kv", + "map_fold_kv", + "map_map_values", + "map_merge_with", "multiset_contains_swapped", "multiset_flat_map", "multiset_fold", @@ -366,6 +381,15 @@ def __rmod__(self, other: f64Like) -> f64: ... @method(egg_fn="abs") def __abs__(self) -> f64: ... + @method(egg_fn="exp") + def exp(self) -> f64: ... + + @method(egg_fn="log") + def log(self) -> f64: ... + + @method(egg_fn="sqrt") + def sqrt(self) -> f64: ... + @method(egg_fn="<") def __lt__(self, other: f64Like) -> Unit: # type: ignore[has-type] ... @@ -397,9 +421,9 @@ def from_i64(cls, i: i64) -> f64: ... def to_string(self) -> String: ... -f64Like: TypeAlias = f64 | float # noqa: N816, PYI042 - +f64Like: TypeAlias = f64 | float | int # noqa: N816, PYI042 +converter(int, f64, lambda i: f64(float(i))) converter(float, f64, f64) @@ -407,6 +431,101 @@ def to_string(self) -> String: ... V = TypeVar("V", bound=BaseExpr) +class Maybe(BuiltinExpr, Generic[T], egg_sort="Maybe"): + @method(preserve=True) + @deprecated("use .value") + def eval(self) -> T | None: + return self.value + + @method(preserve=True) # type: ignore[prop-decorator] + @property + def value(self) -> T | None: + if get_callable_args(self, Maybe.none) is not None: + return None + match get_callable_args(self, Maybe.some): + case (value,): + return value # type: ignore[unreachable] + raise ExprValueError(self, "Maybe.none() or Maybe.some(value)") + + __match_args__ = ("value",) + + @method(egg_fn="maybe-none") + @classmethod + def none(cls) -> Maybe[T]: ... + + @method(egg_fn="maybe-some") + @classmethod + def some(cls, value: T) -> Maybe[T]: ... + + @method(egg_fn="maybe-unwrap") + def unwrap(self) -> T: ... + + @method(egg_fn="maybe-unwrap-or") + def unwrap_or(self, default: T) -> T: ... + + @method(egg_fn="unstable-maybe-match") + def match(self, f: Callable[[T], V], n: V) -> V: ... + + +converter(type(None), Maybe, lambda _: Maybe[get_type_args()[0]].none()) # type: ignore[misc] + + +L = TypeVar("L", bound=BaseExpr) +R = TypeVar("R", bound=BaseExpr) +L2 = TypeVar("L2", bound=BaseExpr) +R2 = TypeVar("R2", bound=BaseExpr) + + +class Pair(BuiltinExpr, Generic[L, R], egg_sort="Pair"): + @method(preserve=True) # type: ignore[prop-decorator] + @property + def value(self) -> tuple[L, R]: + match get_callable_args(self, Pair[L, R]): + case (left, right): + return (left, right) + raise ExprValueError(self, "Pair(left, right)") + + __match_args__ = ("value",) + + @method(egg_fn="pair") + def __init__(self, left: L, right: R) -> None: ... + + @method(egg_fn="pair-first") # type: ignore[prop-decorator] + @property + def left(self) -> L: ... + + @method(egg_fn="pair-second") # type: ignore[prop-decorator] + @property + def right(self) -> R: ... + + @method(preserve=True) + def match(self, f: Callable[[L, R], V]) -> V: + return f(self.left, self.right) + + @method(preserve=True) + def map_left(self, f: Callable[[L], L2]) -> Pair[L2, R]: + return Pair(f(self.left), self.right) + + @method(preserve=True) + def map_right(self, f: Callable[[R], R2]) -> Pair[L, R2]: + return Pair(self.left, f(self.right)) + + +def _tuple_to_pair(value: tuple[object, ...]) -> Pair: + if len(value) != 2: + raise ValueError(f"Expected a tuple of length 2 for Pair conversion, got length {len(value)}") + left, right = value + left_type, right_type = get_type_args() + return Pair(convert(left, left_type), convert(right, right_type)) + + +converter(tuple, Pair, _tuple_to_pair) + + +@function(egg_fn="unstable-catch", builtin=True) +def catch(f: Callable[[], T]) -> Maybe[T]: ... + + class Map(BuiltinExpr, Generic[T, V], egg_sort="Map"): @method(preserve=True) @deprecated("use .value") @@ -416,13 +535,13 @@ def eval(self) -> dict[T, V]: @method(preserve=True) # type: ignore[prop-decorator] @property def value(self) -> dict[T, V]: - d = {} + items = [] while args := get_callable_args(self, Map.insert): # type: ignore[var-annotated] self, k, v = args # noqa: PLW0642 - d[k] = v + items.append((k, v)) if get_callable_args(self, Map.empty) is None: raise ExprValueError(self, "Map.empty or Map.insert") - return d + return dict(reversed(items)) __match_args__ = ("value",) @@ -457,12 +576,41 @@ def contains(self, key: T) -> Unit: ... @method(egg_fn="map-remove") def remove(self, key: T) -> Map[T, V]: ... - @method(egg_fn="rebuild") - def rebuild(self) -> Map[T, V]: ... + @method(egg_fn="map-length") + def length(self) -> i64: ... + + @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) TO = TypeVar("TO") VO = TypeVar("VO") +A = TypeVar("A", bound=BaseExpr) +V2 = TypeVar("V2", bound=BaseExpr) converter( dict, @@ -477,6 +625,57 @@ def rebuild(self) -> Map[T, V]: ... MapLike: TypeAlias = Map[T, V] | dict[TO, VO] +@function(egg_fn="map-fold-kv", builtin=True) +def map_fold_kv(f: Callable[[A, T, V], A], initial: A, xs: Map[T, V]) -> A: ... + + +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, + ) + + class Set(BuiltinExpr, Generic[T], egg_sort="Set"): @method(preserve=True) @deprecated("use .value") @@ -494,6 +693,8 @@ def value(self) -> set[T]: @method(preserve=True) def __iter__(self) -> Iterator[T]: + if (args := get_callable_args(self, Set[T])) is not None: + return iter(dict.fromkeys(args)) return iter(self.value) @method(preserve=True) @@ -532,8 +733,8 @@ def __sub__(self, other: Set[T]) -> Set[T]: ... @method(egg_fn="set-intersect") def __and__(self, other: Set[T]) -> Set[T]: ... - @method(egg_fn="rebuild") - def rebuild(self) -> Set[T]: ... + @method(egg_fn="set-length") + def length(self) -> i64: ... converter( @@ -910,6 +1111,9 @@ def __init__(self, num: BigIntLike, den: BigIntLike) -> None: ... @method(egg_fn="to-f64") def to_f64(self) -> f64: ... + @method(egg_fn="to-i64") + def to_i64(self) -> i64: ... + @method(egg_fn="+") def __add__(self, other: BigRatLike) -> BigRat: ... @@ -976,8 +1180,9 @@ def __ge__(self, other: BigRatLike) -> Unit: ... # type: ignore[has-type] def __le__(self, other: BigRatLike) -> Unit: ... +converter(i64, BigRat, lambda i: BigRat(BigInt(i), BigInt(1))) converter(Fraction, BigRat, lambda f: BigRat(f.numerator, f.denominator)) -BigRatLike: TypeAlias = BigRat | Fraction +BigRatLike: TypeAlias = BigRat | Fraction | i64Like class Vec(BuiltinExpr, Generic[T], egg_sort="Vec"): @@ -1037,9 +1242,6 @@ def length(self) -> i64: ... @method(egg_fn="vec-get") def __getitem__(self, index: i64Like) -> T: ... - @method(egg_fn="rebuild") - def rebuild(self) -> Vec[T]: ... - @method(egg_fn="vec-remove") def remove(self, index: i64Like) -> Vec[T]: ... @@ -1126,7 +1328,7 @@ def _convert_function(fn: FunctionType) -> UnstableFn: Would just be UnstableFn(function(a)) but we have to account for unbound vars within the body. This means that we have to turn all of those unbound vars into args to the function, and then - partially apply them, alongside creating a default rewrite for the function. + partially apply them, alongside storing the eager primitive body for the function. """ decls = Declarations() return_type, *arg_types = [resolve_type_annotation_mutate(decls, tp) for tp in get_type_args()] @@ -1134,19 +1336,25 @@ def _convert_function(fn: FunctionType) -> UnstableFn: arg_decls = [ TypedExprDecl(tp.to_just(), UnboundVarDecl(name)) for name, tp in zip(arg_names, arg_types, strict=True) ] - res = resolve_literal( - return_type, fn(*(RuntimeExpr.__from_values__(decls, a) for a in arg_decls)), Thunk.value(decls) - ) + with set_current_ruleset(None): + res = resolve_literal( + return_type, fn(*(RuntimeExpr.__from_values__(decls, a) for a in arg_decls)), Thunk.value(decls) + ) res_expr = res.__egg_typed_expr__ decls |= res # these are all the args that appear in the body that are not bound by the args of the function unbound_vars = list(collect_unbound_vars(res_expr) - set(arg_decls)) # prefix the args with them - fn_ref = UnnamedFunctionRef(tuple(unbound_vars + arg_decls), res_expr) - rewrite_decl = DefaultRewriteDecl(fn_ref, res_expr.expr, subsume=True) - ruleset_decls = _add_default_rewrite_inner(decls, rewrite_decl, get_current_ruleset()) - ruleset_decls |= res - + all_args = tuple(unbound_vars + arg_decls) + normalized_args = tuple( + TypedExprDecl( + typed_arg.tp, + UnboundVarDecl(cast("UnboundVarDecl", typed_arg.expr).name, f"_{i}"), + ) + for i, typed_arg in enumerate(all_args) + ) + res_expr = replace_typed_expr(res_expr, dict(zip(all_args, normalized_args, strict=True))) + fn_ref = UnnamedFunctionRef(normalized_args, res_expr) fn = RuntimeFunction(Thunk.value(decls), Thunk.value(fn_ref)) return UnstableFn(fn, *(RuntimeExpr.__from_values__(decls, v) for v in unbound_vars)) @@ -1232,5 +1440,5 @@ def py_exec(code: StringLike, globals_: object = PyObject.dict(), locals_: objec """ -Container: TypeAlias = Map | Set | MultiSet | Vec | UnstableFn +Container: TypeAlias = Map | Maybe | Pair | Set | MultiSet | Vec | UnstableFn Primitive: TypeAlias = String | Bool | i64 | f64 | Rational | BigInt | BigRat | PyObject | Unit diff --git a/python/egglog/conversion.py b/python/egglog/conversion.py index 13baac55..803ebabb 100644 --- a/python/egglog/conversion.py +++ b/python/egglog/conversion.py @@ -118,8 +118,9 @@ def convert(source: object, target: type[V]) -> V: """ Convert a source object to a target type. """ - assert isinstance(target, RuntimeClass) - return cast("V", resolve_literal(target.__egg_tp__, source, target.__egg_decls_thunk__)) + runtime_target: object = target + assert isinstance(runtime_target, RuntimeClass) + return cast("V", resolve_literal(runtime_target.__egg_tp__, source, runtime_target.__egg_decls_thunk__)) def convert_to_same_type(source: object, target: RuntimeExpr) -> RuntimeExpr: diff --git a/python/egglog/declarations.py b/python/egglog/declarations.py index 9a5d08e2..db79acd9 100644 --- a/python/egglog/declarations.py +++ b/python/egglog/declarations.py @@ -32,6 +32,8 @@ __all__ = [ + "BUILTIN_EGG_FN_NAMES", + "BUILTIN_EGG_SORT_NAMES", "ActionCommandDecl", "ActionDecl", "BackOffDecl", @@ -83,6 +85,7 @@ "RewriteDecl", "RewriteOrRuleDecl", "RuleDecl", + "RuleEvalMode", "RulesetDecl", "RunDecl", "SaturateDecl", @@ -106,6 +109,10 @@ ] +BUILTIN_EGG_FN_NAMES: set[str] = {"!="} +BUILTIN_EGG_SORT_NAMES: set[str] = set() + + @dataclass(match_args=False) class DelayedDeclarations: __egg_decls_thunk__: Callable[[], Declarations] = field(repr=False) @@ -160,24 +167,12 @@ def builtin(cls, name: str) -> Ident: return cls(name, "egglog.builtins") -default_ruleset_identifier = Ident("") - - @dataclass class Declarations: - _unnamed_functions: set[UnnamedFunctionRef] = field(default_factory=set) _functions: dict[Ident, FunctionDecl | RelationDecl | ConstructorDecl] = field(default_factory=dict) _constants: dict[Ident, ConstantDecl] = field(default_factory=dict) _classes: dict[Ident, ClassDecl] = field(default_factory=dict) - _rulesets: dict[Ident, RulesetDecl | CombinedRulesetDecl] = field( - default_factory=lambda: {default_ruleset_identifier: RulesetDecl([])} - ) - - @property - def default_ruleset(self) -> RulesetDecl: - ruleset = self._rulesets[default_ruleset_identifier] - assert isinstance(ruleset, RulesetDecl) - return ruleset + _rulesets: dict[Ident, RulesetDecl | CombinedRulesetDecl] = field(default_factory=dict) @classmethod def create(cls, *others: DeclarationsLike) -> Declarations: @@ -220,11 +215,7 @@ def update_other(self, other: Declarations) -> None: other._functions |= self._functions other._classes |= self._classes other._constants |= self._constants - # Must combine rulesets bc the empty ruleset might be different, bc DefaultRewriteDecl - # is added to functions. - combined_default_rules: set[RewriteOrRuleDecl] = {*self.default_ruleset.rules, *other.default_ruleset.rules} other._rulesets |= self._rulesets - other._rulesets[default_ruleset_identifier] = RulesetDecl(list(combined_default_rules)) def get_callable_decl(self, ref: CallableRef) -> CallableDecl: # noqa: PLR0911 match ref: @@ -245,14 +236,14 @@ def get_callable_decl(self, ref: CallableRef) -> CallableDecl: # noqa: PLR0911 assert init_fn, f"Class {class_name} does not have an init function." return init_fn case UnnamedFunctionRef(): - return ConstructorDecl(ref.signature) + return FunctionDecl(signature=ref.signature, body=None, builtin=False) assert_never(ref) def set_function_decl( self, ref: FunctionRef | MethodRef | ClassMethodRef | PropertyRef | InitRef, - decl: FunctionDecl | ConstructorDecl, + decl: FunctionCallableDecl, ) -> None: match ref: case FunctionRef(name): @@ -324,12 +315,12 @@ class ClassDecl: egg_name: str | None = None type_vars: tuple[TypeVarRef, ...] = () builtin: bool = False - init: ConstructorDecl | FunctionDecl | None = None - class_methods: dict[str, FunctionDecl | ConstructorDecl] = field(default_factory=dict) + init: FunctionCallableDecl | None = None + class_methods: dict[str, FunctionCallableDecl] = field(default_factory=dict) # These have to be separate from class_methods so that printing them can be done easily class_variables: dict[str, ConstantDecl] = field(default_factory=dict) - methods: dict[str, FunctionDecl | ConstructorDecl] = field(default_factory=dict) - properties: dict[str, FunctionDecl | ConstructorDecl] = field(default_factory=dict) + methods: dict[str, FunctionCallableDecl] = field(default_factory=dict) + properties: dict[str, FunctionCallableDecl] = field(default_factory=dict) preserved_methods: dict[str, Callable] = field(default_factory=dict) match_args: tuple[str, ...] = field(default=()) doc: str | None = field(default=None) @@ -739,11 +730,16 @@ def signature(self) -> FunctionSignature: @dataclass(frozen=True) class ConstantDecl: """ - Same as `(declare)` in egglog + Same as `(declare)` in egglog. + + `body is not None` means the constant lowers eagerly as a zero-arg primitive. + `merge is not None` means the constant lowers as a zero-arg function. """ type_ref: JustTypeRef egg_name: str | None = None + body: TypedExprDecl | None = None + merge: ExprDecl | None = None @property def signature(self) -> FunctionSignature: @@ -799,15 +795,31 @@ def all_args(self) -> Iterable[TypeOrVarRef]: @dataclass(frozen=True) class FunctionDecl: + """ + Function-style callable declaration. + + `builtin=True` is only used for builtin or special no-body callables. + `body is not None` means the callable lowers eagerly. + `merge is not None` is only used for no-body function-style declarations. + """ + signature: FunctionSignature | SpecialFunctions = field(default_factory=FunctionSignature) - builtin: bool = False egg_name: str | None = None + builtin: bool = False + body: TypedExprDecl | None = None merge: ExprDecl | None = None doc: str | None = None @dataclass(frozen=True) class ConstructorDecl: + """ + Constructor-style callable declaration for eqsort-returning callables. + + `cost` and `unextractable` only live on constructor declarations. + Rewrite-backed bodies are represented separately via explicit-ruleset rewrites. + """ + signature: FunctionSignature = field(default_factory=FunctionSignature) egg_name: str | None = None cost: int | None = None @@ -815,7 +827,28 @@ class ConstructorDecl: doc: str | None = None -CallableDecl: TypeAlias = RelationDecl | ConstantDecl | FunctionDecl | ConstructorDecl +FunctionCallableDecl: TypeAlias = FunctionDecl | ConstructorDecl +CallableDecl: TypeAlias = RelationDecl | ConstantDecl | FunctionCallableDecl + + +def is_callable_decl_constructor(decls: Declarations, decl: CallableDecl) -> bool: + """ + Check if a callable declaration will be compiled to a constructor in egglog, as opposed to a function or relation or primitive. + """ + match decl: + case ConstructorDecl(): + return True + case FunctionDecl() | RelationDecl(): + return False + case ConstantDecl() as const_decl: + return ( + not decls.get_class_decl(const_decl.type_ref.ident).builtin + and const_decl.body is None + and const_decl.merge is None + ) + case _: + assert_never(decl) + ## # Expressions @@ -1065,6 +1098,7 @@ class BackOffDecl: id: UUID match_limit: int | None ban_length: int | None + persistent: bool = False ## @@ -1160,16 +1194,20 @@ class BiRewriteDecl: conditions: tuple[FactDecl, ...] +RuleEvalMode: TypeAlias = Literal["seminaive", "naive", "unsafe-seminaive"] + + @dataclass(frozen=True) class RuleDecl: head: tuple[ActionDecl, ...] body: tuple[FactDecl, ...] name: str | None + eval_mode: RuleEvalMode = "seminaive" @dataclass(frozen=True) class DefaultRewriteDecl: - ref: CallableRef + ref: FunctionRef | ConstantRef | MethodRef | ClassMethodRef | InitRef | ClassVariableRef | PropertyRef expr: ExprDecl subsume: bool diff --git a/python/egglog/deconstruct.py b/python/egglog/deconstruct.py index 1a953b06..6ffa51a0 100644 --- a/python/egglog/deconstruct.py +++ b/python/egglog/deconstruct.py @@ -6,7 +6,7 @@ from collections.abc import Callable from functools import partial -from typing import TYPE_CHECKING, TypeVar, overload +from typing import TYPE_CHECKING, TypeVar, cast, overload import cloudpickle from typing_extensions import TypeVarTuple, Unpack @@ -86,9 +86,9 @@ def get_constant_name(x: BaseExpr) -> Ident | None: Check if the expression is a constant and return its name. If it is not a constant, return None. """ - if not isinstance(x, RuntimeExpr): + if not isinstance(cast("object", x), RuntimeExpr): raise TypeError(f"Expected Expression, got {type(x).__name__}") - match x.__egg_typed_expr__.expr: + match cast("RuntimeExpr", x).__egg_typed_expr__.expr: case CallDecl(ConstantRef(ident)): return ident return None @@ -99,9 +99,9 @@ def get_let_name(x: BaseExpr) -> str | None: Check if the expression is a `let` expression and return the name of the variable. If it is not a `let` expression, return None. """ - if not isinstance(x, RuntimeExpr): + if not isinstance(cast("object", x), RuntimeExpr): raise TypeError(f"Expected Expression, got {type(x).__name__}") - match x.__egg_typed_expr__.expr: + match cast("RuntimeExpr", x).__egg_typed_expr__.expr: case LetRefDecl(name): return name return None @@ -112,9 +112,9 @@ def get_var_name(x: BaseExpr) -> str | None: Check if the expression is a variable and return its name. If it is not a variable, return None. """ - if not isinstance(x, RuntimeExpr): + if not isinstance(cast("object", x), RuntimeExpr): raise TypeError(f"Expected Expression, got {type(x).__name__}") - match x.__egg_typed_expr__.expr: + match cast("RuntimeExpr", x).__egg_typed_expr__.expr: case UnboundVarDecl(name, _egg_name): return name return None @@ -124,17 +124,18 @@ def get_callable_fn(x: T) -> Callable[..., T] | T | None: """ Gets the function of an expression, or if it's a constant or classvar, return that. """ - if not isinstance(x, RuntimeExpr): + if not isinstance(cast("object", x), RuntimeExpr): raise TypeError(f"Expected Expression, got {type(x).__name__}") - match x.__egg_typed_expr__.expr: + runtime_x = cast("RuntimeExpr", x) + match runtime_x.__egg_typed_expr__.expr: case CallDecl() as call: - fn, _ = _deconstruct_call_decl(x.__egg_decls_thunk__, call) + fn, _ = _deconstruct_call_decl(runtime_x.__egg_decls_thunk__, call) return fn return None @overload -def get_callable_args(x: T, fn: None = ...) -> tuple[BaseExpr, ...]: ... +def get_callable_args(x: T, fn: None = ...) -> tuple[BaseExpr, ...] | None: ... @overload @@ -149,27 +150,30 @@ def get_callable_args(x: T, fn: Callable[[Unpack[TS]], T] | None = None) -> tupl Note that recursively calling the arguments is the safe way to walk the expression tree. """ - if not isinstance(x, RuntimeExpr): + if not isinstance(cast("object", x), RuntimeExpr): raise TypeError(f"Expected Expression, got {type(x).__name__}") - match x.__egg_typed_expr__.expr: + runtime_x = cast("RuntimeExpr", x) + match runtime_x.__egg_typed_expr__.expr: case CallDecl() as call: - actual_fn, args = _deconstruct_call_decl(x.__egg_decls_thunk__, call) + actual_fn, args = _deconstruct_call_decl(runtime_x.__egg_decls_thunk__, call) if fn is None: - return args + return cast("tuple[*TS]", args) # Compare functions and classes without considering bound type parameters, so that you can pass # in a binding like Vec[i64] and match Vec[i64](...) or Vec(...) calls. - if ( - isinstance(actual_fn, RuntimeFunction) - and isinstance(fn, RuntimeFunction) - and actual_fn.__egg_ref__ == fn.__egg_ref__ - ): - return args - if ( - isinstance(actual_fn, RuntimeClass) - and isinstance(fn, RuntimeClass) - and actual_fn.__egg_tp__.ident == fn.__egg_tp__.ident - ): - return args + if isinstance(cast("object", actual_fn), RuntimeFunction): + runtime_actual_fn = cast("RuntimeFunction", actual_fn) + if ( + isinstance(cast("object", fn), RuntimeFunction) + and runtime_actual_fn.__egg_ref__ == cast("RuntimeFunction", fn).__egg_ref__ + ): + return cast("tuple[*TS]", args) + if isinstance(cast("object", actual_fn), RuntimeClass): + runtime_actual_cls = cast("RuntimeClass", actual_fn) + if ( + isinstance(cast("object", fn), RuntimeClass) + and runtime_actual_cls.__egg_tp__.ident == cast("RuntimeClass", fn).__egg_tp__.ident + ): + return cast("tuple[*TS]", args) return None diff --git a/python/egglog/egraph.py b/python/egglog/egraph.py index 6fe782d5..4f679e19 100644 --- a/python/egglog/egraph.py +++ b/python/egglog/egraph.py @@ -3,10 +3,11 @@ import contextlib import inspect import pathlib +import sys import tempfile from collections.abc import Callable, Generator, Iterable from contextvars import ContextVar, Token -from dataclasses import InitVar, dataclass, field +from dataclasses import InitVar, dataclass, field, replace from functools import partial from inspect import Parameter, currentframe, getmodule, signature from types import FrameType, FunctionType @@ -32,13 +33,14 @@ import graphviz from opentelemetry import trace -from typing_extensions import ParamSpec, Unpack +from typing_extensions import ParamSpec, TypeForm, Unpack from . import bindings from ._tracing import call_with_current_trace from .conversion import * from .conversion import convert_to_same_type, resolve_literal from .declarations import * +from .declarations import is_callable_decl_constructor from .egraph_state import * from .ipython_magic import IN_IPYTHON from .pretty import pretty_decl @@ -70,6 +72,7 @@ "GraphvizKwargs", "GreedyDagCost", "RewriteOrRule", + "RuleEvalMode", "Ruleset", "RunReport", "Schedule", @@ -124,6 +127,7 @@ BE2 = TypeVar("BE2", bound="BaseExpr") BE3 = TypeVar("BE3", bound="BaseExpr") BE4 = TypeVar("BE4", bound="BaseExpr") +_MISSING = object() # Attributes which are sometimes added to classes by the interpreter or the dataclass decorator, or by ipython. # We ignore these when inspecting the class. @@ -188,11 +192,27 @@ def check(x: FactLike, schedule: Schedule | None = None, *given: ActionLike) -> # So that we can add the functions eagerly to the registry and wait on the methods till we process the class. +def _resolve_merge( + decls: Declarations, + return_type: TypeOrVarRef, + merge: Callable[[object, object], object] | None, +) -> ExprDecl | None: + if merge is None: + return None + old = RuntimeExpr.__from_values__(decls, TypedExprDecl(return_type.to_just(), UnboundVarDecl("old", "old"))) + new = RuntimeExpr.__from_values__(decls, TypedExprDecl(return_type.to_just(), UnboundVarDecl("new", "new"))) + resolved_merge = resolve_literal(return_type, merge(old, new), lambda: decls) + decls |= resolved_merge + return resolved_merge.__egg_typed_expr__.expr + + CALLABLE = TypeVar("CALLABLE", bound=Callable) CONSTRUCTOR_CALLABLE = TypeVar("CONSTRUCTOR_CALLABLE", bound=Callable[..., "Expr | None"]) EXPR_NONE = TypeVar("EXPR_NONE", bound="Expr | None") BASE_EXPR_NONE = TypeVar("BASE_EXPR_NONE", bound="BaseExpr | None") +USER_EXPR = TypeVar("USER_EXPR", bound="Expr") +BUILTIN_EXPR = TypeVar("BUILTIN_EXPR", bound="BuiltinExpr") @overload @@ -274,12 +294,21 @@ def function( *, egg_fn: str | None = ..., merge: Callable[[BASE_EXPR, BASE_EXPR], BASE_EXPR] | None = ..., - builtin: bool = ..., mutates_first_arg: bool = ..., ) -> Callable[[Callable[P, BASE_EXPR]], Callable[P, BASE_EXPR]]: ... # constructor +@overload +def function( + *, + egg_fn: str | None = ..., + cost: int | None = ..., + mutates_first_arg: bool = ..., + unextractable: bool = ..., +) -> Callable[[CONSTRUCTOR_CALLABLE], CONSTRUCTOR_CALLABLE]: ... + + @overload def function( *, @@ -296,10 +325,14 @@ def function(*args, **kwargs) -> Any: """ Decorate a function typing stub to create an egglog function for it. - If a body is included, it will be added to the `ruleset` passed in as a default rewrite. + Python automatically infers whether this lowers to a `function`, `constructor`, + or eager `primitive` from the return type and whether a body/default exists. - This will default to creating a "constructor" in egglog, unless a merge function is passed in or the return - type is a primtive, then it will be a "function". + Bodies and defaults lower eagerly unless an explicit `ruleset` is provided, + in which case they are added as rewrite-backed defaults in that ruleset. + Direct top-level functions can only use an explicit `ruleset` when they + provide a body to lower as a rewrite. `subsume` is only valid for eqsort- + returning bodies on that explicit-ruleset rewrite path. """ fn_locals = currentframe().f_back.f_locals # type: ignore[union-attr] @@ -330,6 +363,9 @@ def __new__( # type: ignore[misc] if not bases or bases == (BaseExpr,): return super().__new__(cls, name, bases, namespace) builtin = BuiltinExpr in bases + if builtin and egg_sort is not None: + BUILTIN_EGG_SORT_NAMES.add(egg_sort) + _register_builtin_class_egg_fns(namespace) frame = currentframe() assert frame @@ -430,10 +466,37 @@ def _generate_class_decls( # noqa: C901,PLR0912 if getattr(v, "__origin__", None) == ClassVar: (inner_tp,) = v.__args__ type_ref = resolve_type_annotation_mutate(decls, inner_tp) - cls_decl.class_variables[k] = ConstantDecl(type_ref.to_just()) - _add_default_rewrite( - decls, ClassVariableRef(cls_ident, k), type_ref, namespace.pop(k, None), ruleset, subsume=False + default_value = namespace.pop(k, _MISSING) + has_default = default_value is not _MISSING + return_type_is_eqsort = isinstance(type_ref, TypeRefWithVars) and not decls._classes[type_ref.ident].builtin + if has_default and ruleset is not None and not return_type_is_eqsort: + msg = "Primitive-returning defaults cannot use an explicit ruleset" + raise ValueError(msg) + default_mode = _normalize_callable_mode( + return_type_is_eqsort=return_type_is_eqsort, + returns_unit=type_ref == TypeRefWithVars(Ident.builtin("Unit")), + has_body=True if has_default else None, + has_ruleset=ruleset is not None, + require_body_for_ruleset=False, + has_merge=False, + builtin=False, + has_cost=False, + unextractable=False, + subsume=False, + ) + resolved_default = ( + resolve_literal(type_ref, default_value, Thunk.value(decls)) if default_mode == "eager" else None ) + if resolved_default is not None: + decls |= resolved_default + cls_decl.class_variables[k] = ConstantDecl( + type_ref.to_just(), + body=resolved_default.__egg_typed_expr__ if resolved_default is not None else None, + ) + if default_mode == "rewrite": + _add_default_rewrite( + decls, ClassVariableRef(cls_ident, k), type_ref, default_value, ruleset, subsume=False + ) else: msg = f"On class {cls_ident}, for attribute '{k}', expected a ClassVar, but got {v}" raise NotImplementedError(msg) @@ -470,7 +533,7 @@ def _generate_class_decls( # noqa: C901,PLR0912 ref: ClassMethodRef | MethodRef | PropertyRef | InitRef # TODO: Store deprecated message so we can get at runtime if (getattr(fn, "__deprecated__", None)) is not None: - fn = fn.__wrapped__ # type: ignore[attr-defined] + fn = fn.__wrapped__ # type: ignore[union-attr] match fn: case classmethod(): ref = ClassMethodRef(cls_ident, method_name) @@ -511,7 +574,7 @@ def _generate_class_decls( # noqa: C901,PLR0912 e.add_note(f"Error processing {cls_ident}.{method_name}") raise - if not builtin and not isinstance(ref, InitRef): + if not builtin: add_default_funcs.append(add_rewrite) # Add all rewrite methods at the end so that all methods are registered first and can be accessed @@ -521,6 +584,12 @@ def _generate_class_decls( # noqa: C901,PLR0912 return decls +def _register_builtin_class_egg_fns(namespace: dict[str, Any]) -> None: + for method in namespace.values(): + if isinstance(method, _WrappedMethod) and method.egg_fn is not None: + BUILTIN_EGG_FN_NAMES.add(method.egg_fn) + + @dataclass class _FunctionConstructor: hint_locals: dict[str, Any] @@ -533,6 +602,10 @@ class _FunctionConstructor: ruleset: Ruleset | None = None subsume: bool = False + def __post_init__(self) -> None: + if self.builtin and self.egg_fn is not None: + BUILTIN_EGG_FN_NAMES.add(self.egg_fn) + def __call__(self, fn: Callable) -> RuntimeFunction: return RuntimeFunction(*split_thunk(Thunk.fn(self.create_decls, fn))) @@ -582,12 +655,13 @@ def _fn_decl( if not isinstance(fn, FunctionType): raise NotImplementedError(f"Can only generate function decls for functions not {fn} {type(fn)}") - # Instead of passing both globals and locals, just pass the globals. Otherwise, for some reason forward references - # won't be resolved correctly + # Resolve annotations against the module, function globals, and captured definition locals. Passing the same merged + # namespace as globals and locals keeps forward references working across Python versions. # 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 = {**fn.__globals__, **hint_locals} - hints = get_type_hints(fn, hint_globals) + module_globals = vars(sys.modules[fn.__module__]) if fn.__module__ in sys.modules else {} + hint_namespace = {**module_globals, **fn.__globals__, **hint_locals} + hints = get_type_hints(fn, globalns=hint_namespace, localns=hint_namespace) params = list(signature(fn).parameters.values()) @@ -626,27 +700,23 @@ def _fn_decl( arg_names = tuple(t.name for t in params) - merged = ( - None - if merge is None - else resolve_literal( - return_type, - merge( - RuntimeExpr.__from_values__(decls, TypedExprDecl(return_type.to_just(), UnboundVarDecl("old", "old"))), - RuntimeExpr.__from_values__(decls, TypedExprDecl(return_type.to_just(), UnboundVarDecl("new", "new"))), - ), - lambda: decls, + merge_expr = _resolve_merge(decls, return_type, merge) + + # Keep these lazy so builtin declarations do not resolve them eagerly. + # Eager primitive bodies are bound in backend argument order. Rewrite-backed + # bodies keep Python-order variables because their call pattern is reversed + # separately during lowering. + reverse_body_args = reverse_args and ruleset is None + arg_count = len(arg_types) + args = ( + TypedExprDecl( + tp.to_just(), + UnboundVarDecl(name, f"_{arg_count - i - 1 if reverse_body_args else i}"), ) + for i, (name, tp) in enumerate(zip(arg_names, arg_types, strict=True)) ) - decls |= merged - # defer this in generator so it doesn't resolve for builtins eagerly - args = (TypedExprDecl(tp.to_just(), UnboundVarDecl(name)) for name, tp in zip(arg_names, arg_types, strict=True)) - - return_type_is_eqsort = ( - not decls._classes[return_type.ident].builtin if isinstance(return_type, TypeRefWithVars) else False - ) - is_constructor = not is_builtin and return_type_is_eqsort and merged is None + return_type_is_eqsort = isinstance(return_type, TypeRefWithVars) and not decls._classes[return_type.ident].builtin signature_ = FunctionSignature( return_type=None if mutates_first_arg else return_type, var_arg_type=var_arg_type, @@ -656,25 +726,35 @@ def _fn_decl( reverse_args=reverse_args, ) doc = fn.__doc__ + mode = _normalize_callable_mode( + return_type_is_eqsort=return_type_is_eqsort, + returns_unit=signature_.semantic_return_type == TypeRefWithVars(Ident.builtin("Unit")), + has_body=None, + has_ruleset=ruleset is not None, + require_body_for_ruleset=isinstance(ref, FunctionRef), + has_merge=merge_expr is not None, + builtin=is_builtin, + has_cost=cost is not None, + unextractable=unextractable, + subsume=subsume, + ) decl: ConstructorDecl | FunctionDecl - if is_constructor: + if mode == "constructor": decl = ConstructorDecl(signature_, egg_name, cost, unextractable, doc) else: - if cost is not None: - msg = "Cost can only be set for constructors" - raise ValueError(msg) - if unextractable: - msg = "Unextractable can only be set for constructors" - raise ValueError(msg) decl = FunctionDecl( signature=signature_, egg_name=egg_name, - merge=merged.__egg_typed_expr__.expr if merged is not None else None, + merge=merge_expr, builtin=is_builtin, doc=doc, ) decls.set_function_decl(ref, decl) - if is_builtin: + if is_builtin and ( + any(tp.vars for tp in arg_types) + or (var_arg_type is not None and bool(var_arg_type.vars)) + or bool(return_type.vars) + ): return lambda: None return Thunk.fn( _add_default_rewrite_function, @@ -746,14 +826,60 @@ def _relation_decls(ident: Ident, tps: tuple[type, ...], egg_fn: str | None) -> return decls +@overload +def constant( + name: str, + tp: type[USER_EXPR], + /, + *, + egg_name: str | None = ..., + merge: Callable[[USER_EXPR, USER_EXPR], USER_EXPR] | None = ..., +) -> USER_EXPR: ... + + +@overload +def constant( + name: str, + tp: type[BUILTIN_EXPR], + /, + *, + egg_name: str | None = ..., + merge: Callable[[BUILTIN_EXPR, BUILTIN_EXPR], BUILTIN_EXPR] | None = ..., +) -> BUILTIN_EXPR: ... + + +@overload +def constant( + name: str, + tp: type[USER_EXPR], + default_replacement: USER_EXPR, + /, + *, + egg_name: str | None = ..., + ruleset: Ruleset | None = ..., +) -> USER_EXPR: ... + + +@overload +def constant( + name: str, + tp: type[BUILTIN_EXPR], + default_replacement: BUILTIN_EXPR | None, + /, + *, + egg_name: str | None = ..., +) -> BUILTIN_EXPR: ... + + def constant( name: str, tp: type[BASE_EXPR], - default_replacement: BASE_EXPR | None = None, + default_replacement: BASE_EXPR | None | object = _MISSING, /, *, egg_name: str | None = None, ruleset: Ruleset | None = None, + merge: Callable[[BASE_EXPR, BASE_EXPR], BASE_EXPR] | None = None, ) -> BASE_EXPR: """ A "constant" is implemented as the instantiation of a value that takes no args. @@ -761,22 +887,136 @@ def constant( """ return cast( "BASE_EXPR", - RuntimeExpr(*split_thunk(Thunk.fn(_constant_thunk, name, tp, egg_name, default_replacement, ruleset))), + RuntimeExpr(*split_thunk(Thunk.fn(_constant_thunk, name, tp, egg_name, default_replacement, ruleset, merge))), ) def _constant_thunk( - name: str, tp: type, egg_name: str | None, default_replacement: object, ruleset: Ruleset | None + name: str, + tp: type, + egg_name: str | None, + default_replacement: object, + ruleset: Ruleset | None, + merge: Callable[[Any, Any], Any] | None = None, ) -> tuple[Declarations, TypedExprDecl]: decls = Declarations() type_ref = resolve_type_annotation_mutate(decls, tp) ident = Ident(name, _get_module()) callable_ref = ConstantRef(ident) - decls._constants[ident] = ConstantDecl(type_ref.to_just(), egg_name) - _add_default_rewrite(decls, callable_ref, type_ref, default_replacement, ruleset, subsume=False) + has_default = default_replacement is not _MISSING + return_type_is_eqsort = isinstance(type_ref, TypeRefWithVars) and not decls._classes[type_ref.ident].builtin + if ruleset is not None and not has_default: + msg = "Explicit rulesets require a default" + raise ValueError(msg) + if has_default and ruleset is not None and not return_type_is_eqsort: + msg = "Primitive-returning defaults cannot use an explicit ruleset" + raise ValueError(msg) + merge_expr = _resolve_merge(decls, type_ref, merge) + mode = _normalize_callable_mode( + return_type_is_eqsort=return_type_is_eqsort, + returns_unit=type_ref == TypeRefWithVars(Ident.builtin("Unit")), + has_body=True if has_default else None, + has_ruleset=ruleset is not None, + require_body_for_ruleset=False, + has_merge=merge_expr is not None, + builtin=False, + has_cost=False, + unextractable=False, + subsume=False, + ) + resolved_default = resolve_literal(type_ref, default_replacement, Thunk.value(decls)) if mode == "eager" else None + if resolved_default is not None: + decls |= resolved_default + decls._constants[ident] = ConstantDecl( + type_ref.to_just(), + egg_name, + resolved_default.__egg_typed_expr__ if resolved_default is not None else None, + merge_expr, + ) + if mode == "rewrite": + _add_default_rewrite(decls, callable_ref, type_ref, default_replacement, ruleset, subsume=False) return decls, TypedExprDecl(type_ref.to_just(), CallDecl(callable_ref)) +_CallableMode: TypeAlias = Literal["function", "constructor", "eager", "rewrite"] + + +def _normalize_callable_mode( # noqa: C901, PLR0911, PLR0912 + *, + return_type_is_eqsort: bool, + returns_unit: bool, + has_body: bool | None, + has_ruleset: bool, + require_body_for_ruleset: bool, + has_merge: bool, + builtin: bool, + has_cost: bool, + unextractable: bool, + subsume: bool, +) -> _CallableMode: + if builtin and has_merge: + msg = "Builtin callables cannot use merge" + raise ValueError(msg) + if require_body_for_ruleset and has_ruleset and has_body is False: + msg = "Explicit rulesets require a body" + raise ValueError(msg) + if subsume: + if not return_type_is_eqsort: + msg = "Primitive-returning callables cannot use subsume" + raise ValueError(msg) + if not has_ruleset: + msg = "subsume requires an explicit ruleset" + raise ValueError(msg) + if has_body is False: + msg = "subsume requires a body" + raise ValueError(msg) + + if return_type_is_eqsort: + if builtin: + msg = "Eqsort-returning callables cannot be builtin" + raise ValueError(msg) + if has_body is None: + return "function" if has_merge else "constructor" + if has_body: + if has_merge: + msg = "Eqsort-returning callables with bodies cannot use merge" + raise ValueError(msg) + if has_ruleset: + return "rewrite" + if has_cost: + msg = "Eqsort-returning eager bodies cannot use cost" + raise ValueError(msg) + if unextractable: + msg = "Eqsort-returning eager bodies cannot be unextractable" + raise ValueError(msg) + return "eager" + return "function" if has_merge else "constructor" + + if has_cost: + msg = "Primitive-returning callables cannot use cost" + raise ValueError(msg) + if unextractable: + msg = "Primitive-returning callables cannot be unextractable" + raise ValueError(msg) + if returns_unit and has_merge: + msg = "Functions that return Unit cannot use merge" + raise ValueError(msg) + if has_body is None: + return "function" + if has_body: + if has_ruleset: + msg = "Primitive-returning callables with bodies cannot use an explicit ruleset" + raise ValueError(msg) + if builtin: + msg = "Builtin callables cannot have a body" + raise ValueError(msg) + if has_merge: + msg = "Primitive-returning callables with bodies cannot use merge" + raise ValueError(msg) + return "eager" + return "function" + + def _add_default_rewrite_function( decls: Declarations, ref: FunctionRef | MethodRef | PropertyRef | ClassMethodRef | InitRef, @@ -793,29 +1033,64 @@ def _add_default_rewrite_function( if isinstance(ref, ClassMethodRef): tp = decls.get_parameterized_class(ref.ident) arg_exprs.insert(0, RuntimeClass(Thunk.value(decls), tp)) + elif isinstance(ref, InitRef): + tp = decls.get_parameterized_class(ref.ident) + arg_exprs.insert( + 0, + RuntimeExpr.__from_values__(decls, TypedExprDecl(tp.to_just(), CallDecl(ref, tuple(args)))), + ) with set_current_ruleset(ruleset): res = fn(*arg_exprs) - # If the function mutates the first arg and we have overwritten it, then use that as the result + # If this function mutates its first argument and that argument changed, use it as the result. if mutates_first_arg and arg_exprs[0].__egg_typed_expr__ != args[0]: res = arg_exprs[0] - _add_default_rewrite(decls, ref, res_type, res, ruleset, subsume) + decl = decls.get_callable_decl(ref) + assert isinstance(decl, ConstructorDecl | FunctionDecl) + mode = _normalize_callable_mode( + return_type_is_eqsort=isinstance(res_type, TypeRefWithVars) and not decls._classes[res_type.ident].builtin, + returns_unit=res_type == TypeRefWithVars(Ident.builtin("Unit")), + has_body=res is not None, + has_ruleset=ruleset is not None, + require_body_for_ruleset=isinstance(ref, FunctionRef), + has_merge=isinstance(decl, FunctionDecl) and decl.merge is not None, + builtin=isinstance(decl, FunctionDecl) and decl.builtin, + has_cost=isinstance(decl, ConstructorDecl) and decl.cost is not None, + unextractable=isinstance(decl, ConstructorDecl) and decl.unextractable, + subsume=subsume, + ) + if mode in ("function", "constructor"): + return + if mode == "rewrite": + _add_default_rewrite(decls, ref, res_type, res, ruleset, subsume) + return + + assert res is not None + resolved_value = resolve_literal(res_type, res, Thunk.value(decls)) + decls |= resolved_value + match decl: + case ConstructorDecl(signature, egg_name, _, _, doc): + decls.set_function_decl( + ref, + FunctionDecl(signature=signature, egg_name=egg_name, body=resolved_value.__egg_typed_expr__, doc=doc), + ) + case FunctionDecl(): + decls.set_function_decl(ref, replace(decl, body=resolved_value.__egg_typed_expr__)) def _add_default_rewrite( decls: Declarations, - ref: CallableRef, + ref: FunctionRef | ConstantRef | MethodRef | ClassMethodRef | InitRef | ClassVariableRef | PropertyRef, type_ref: TypeOrVarRef, default_rewrite: object, ruleset: Ruleset | None, subsume: bool, ) -> 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. + Adds a default rewrite for the callable when an explicit ruleset is provided. """ - if default_rewrite is None: - return + if ruleset is None: + msg = "Default rewrites require an explicit ruleset" + raise ValueError(msg) resolved_value = resolve_literal(type_ref, default_rewrite, Thunk.value(decls)) rewrite_decl = DefaultRewriteDecl(ref, resolved_value.__egg_typed_expr__.expr, subsume) ruleset_decls = _add_default_rewrite_inner(decls, rewrite_decl, ruleset) @@ -824,14 +1099,15 @@ def _add_default_rewrite( def _add_default_rewrite_inner( - decls: Declarations, rewrite_decl: DefaultRewriteDecl, ruleset: Ruleset | None + decls: Declarations, + rewrite_decl: DefaultRewriteDecl, + ruleset: Ruleset | None, ) -> Declarations: - if ruleset: - ruleset_decls = ruleset._current_egg_decls - ruleset_decl = ruleset.__egg_ruleset__ - else: - ruleset_decls = decls - ruleset_decl = decls.default_ruleset + if ruleset is None: + msg = "Default rewrites require an explicit ruleset" + raise ValueError(msg) + ruleset_decls = ruleset._current_egg_decls + ruleset_decl = ruleset.__egg_ruleset__ ruleset_decl.rules.append(rewrite_decl) return ruleset_decls @@ -886,7 +1162,7 @@ def __init__( ) -> None: with _TRACER.start_as_current_span("create"): with _TRACER.start_as_current_span("create_bindings"): - self._state = EGraphState(bindings.EGraph(seminaive=seminaive, record=save_egglog_string)) + self._state = EGraphState(bindings.EGraph(seminaive=seminaive), save_egglog_string=save_egglog_string) self._state_stack = [] self._token_stack = [] if actions: @@ -907,20 +1183,26 @@ def as_egglog_string(self) -> str: """ Returns the egglog string for this module. """ - cmds = self._egraph.commands() - if cmds is None: - msg = "Can't get egglog string unless EGraph created with save_egglog_string=True" - raise ValueError(msg) - return cmds + return self._state.egglog_string() + + def close(self) -> None: + """Close and remove the optional saved Egglog transcript.""" + self._state.close() 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. + Load a CSV file into a table-backed callable. + + Eager and builtin primitives do not have tables and cannot be input + targets. """ - self._run_program(bindings.Input(span(1), self._callable_to_egg(fn)[1], path)) + ref, decls = resolve_callable(fn) + self._add_decls(decls) + self._require_table_backed(ref) + self._state.run_program(bindings.Input(span(1), self._state.callable_ref_to_egg(ref)[0], path)) def _callable_to_egg(self, fn: ExprCallable) -> tuple[CallableRef, str]: ref, decls = resolve_callable(fn) @@ -955,13 +1237,21 @@ def output(self) -> None: raise NotImplementedError(msg) @overload - def run(self, limit: int, /, *until: Fact, ruleset: Ruleset | None = None) -> RunReport: ... + def run( + self, limit: int, /, *until: Fact, ruleset: Ruleset | UnstableCombinedRuleset | None = None + ) -> RunReport: ... @overload def run(self, schedule: Schedule, /) -> RunReport: ... @_TRACER.start_as_current_span("run") - def run(self, limit_or_schedule: int | Schedule, /, *until: Fact, ruleset: Ruleset | None = None) -> RunReport: + def run( + self, + limit_or_schedule: int | Schedule, + /, + *until: Fact, + ruleset: Ruleset | UnstableCombinedRuleset | None = None, + ) -> RunReport: """ Run the egraph until the given limit or until the given facts are true. """ @@ -972,7 +1262,7 @@ def run(self, limit_or_schedule: int | Schedule, /, *until: Fact, ruleset: Rules def _run_schedule(self, schedule: Schedule) -> RunReport: self._add_decls(schedule) cmd = self._state.run_schedule_to_egg(schedule.schedule) - (command_output,) = self._run_program(cmd) + (command_output,) = self._state.run_program(cmd) assert isinstance(command_output, bindings.RunScheduleOutput) return RunReport._from_bindings(command_output.report, self._state) @@ -980,7 +1270,7 @@ def stats(self) -> RunReport: """ Returns the overall run report for the egraph. """ - (output,) = self._run_program(bindings.PrintOverallStatistics(span(1), None)) + (output,) = self._state.run_program(bindings.PrintOverallStatistics(span(1), None)) assert isinstance(output, bindings.OverallStatistics) return RunReport._from_bindings(output.report, self._state) @@ -1002,13 +1292,13 @@ def check(self, *facts: FactLike) -> None: """ Check if a fact is true in the egraph. """ - self._run_program(self._facts_to_check(facts)) + self._state.run_program(self._facts_to_check(facts)) def check_fail(self, *facts: FactLike) -> None: """ Checks that one of the facts is not true """ - self._run_program(bindings.Fail(span(1), self._facts_to_check(facts))) + self._state.run_program(bindings.Fail(span(1), self._facts_to_check(facts))) def _facts_to_check(self, fact_likes: Iterable[FactLike]) -> bindings.Check: facts = _fact_likes(fact_likes) @@ -1047,8 +1337,8 @@ def extract( res = self._from_termdag(extract_report.termdag, extract_report.term, tp) cost = cast("COST", extract_report.cost) else: - # TODO: For some reason we need this or else it wont be registered. Not sure why - self.register(expr) + if isinstance(runtime_expr.__egg_typed_expr__.expr, CallDecl): + self._register_extract_root(runtime_expr) egg_cost_model = _CostModel(cost_model, self).to_bindings_cost_model() egg_sort = self._state.type_ref_to_egg(tp) extractor = call_with_current_trace(bindings.Extractor, [egg_sort], self._state.egraph, egg_cost_model) @@ -1079,12 +1369,12 @@ def _run_extract(self, expr: RuntimeExpr, n: int) -> bindings._CommandOutput: egg_expr = self._state.typed_expr_to_egg(expr.__egg_typed_expr__) # If we have defined any cost tables use the custom extraction args = (egg_expr, bindings.Lit(span(2), bindings.Int(n))) - if self._state.cost_callables: + if self._state.cost_table_names: cmd: bindings._Command = bindings.UserDefined(span(2), "extract", list(args)) else: cmd = bindings.Extract(span(2), *args) try: - return self._run_program(cmd)[0] + return self._state.run_program(cmd)[0] except BaseException as e: e.add_note("while extracting: " + str(expr)) raise @@ -1094,7 +1384,7 @@ def push(self) -> None: """ Push the current state of the egraph, so that it can be popped later and reverted back. """ - self._run_program(bindings.Push(1)) + self._state.run_program(bindings.Push(1)) self._state_stack.append(self._state) self._state = self._state.copy() @@ -1103,7 +1393,7 @@ def pop(self) -> None: """ Pop the current state of the egraph, reverting back to the previous state. """ - self._run_program(bindings.Pop(span(1), 1)) + self._state.run_program(bindings.Pop(span(1), 1)) self._state = self._state_stack.pop() def __enter__(self) -> Self: @@ -1231,7 +1521,10 @@ def to_json() -> str: i = 0 # Always visualize, even if we encounter an error try: - while (self.run(schedule or 1).updated) and i < max: + while i < max: + report = self.run(schedule or 1) + if report.can_stop: + break i += 1 if visualize: egraphs.append(to_json()) @@ -1253,9 +1546,6 @@ def to_json() -> str: def _egraph(self) -> bindings.EGraph: return self._state.egraph - def _run_program(self, *commands: bindings._Command) -> list[bindings._CommandOutput]: - return call_with_current_trace(self._egraph.run_program, *commands) - @property def __egg_decls__(self) -> Declarations: return self._state.__egg_decls__ @@ -1285,7 +1575,21 @@ def register( def _register_commands(self, cmds: list[Command]) -> None: self._add_decls(*cmds) egg_cmds = [egg_cmd for cmd in cmds if (egg_cmd := self._command_to_egg(cmd)) is not None] - self._run_program(*egg_cmds) + self._state.run_program(*egg_cmds) + + def _register_extract_root(self, runtime_expr: RuntimeExpr) -> None: + """ + Register the exact extraction root without synthetic let factoring. + + Synthetic lets are a command-size optimization for public registration, + but custom-cost extraction immediately evaluates the original root value. + Registering a let-factored presentation can leave the custom extractor + without a costed parent for that exact value in the direct command API. + """ + self._add_decls(runtime_expr) + action_egg = self._state.action_to_egg(ExprActionDecl(runtime_expr.__egg_typed_expr__), expr_to_let=False) + if action_egg is not None: + self._state.run_program(bindings.ActionCommand(action_egg)) def _command_to_egg(self, cmd: Command) -> bindings._Command | None: ruleset_ident = Ident("") @@ -1302,39 +1606,50 @@ def _command_to_egg(self, cmd: Command) -> bindings._Command | None: def function_size(self, fn: ExprCallable) -> int: """ - Returns the number of rows in a certain function + Return the number of rows in a table-backed callable. + + Relations, constructors, and bodyless functions have tables. Eager and + builtin primitives do not. """ - egg_name = self._callable_to_egg(fn)[1] - (output,) = self._run_program(bindings.PrintSize(span(1), egg_name)) + ref, decls = resolve_callable(fn) + self._add_decls(decls) + self._require_table_backed(ref) + egg_name = self._state.callable_ref_to_egg(ref)[0] + (output,) = self._state.run_program(bindings.PrintSize(span(1), egg_name)) assert isinstance(output, bindings.PrintFunctionSize) return output.size def all_function_sizes(self) -> list[tuple[ExprCallable, int]]: """ - Returns a list of all functions and their sizes. + Return the sizes of all registered function tables. """ - (output,) = self._run_program(bindings.PrintSize(span(1), None)) + (output,) = self._state.run_program(bindings.PrintSize(span(1), None)) assert isinstance(output, bindings.PrintAllFunctionsSize) return [(callables[0], size) for (name, size) in output.sizes if (callables := self._egg_fn_to_callables(name))] def _egg_fn_to_callables(self, egg_fn: str) -> list[ExprCallable]: return [ cast("ExprCallable", create_callable(self._state.__egg_decls__, ref)) - for ref in self._state.egg_fn_to_callable_refs[egg_fn] + for ref in self._state.egg_fn_to_callable_refs.get(egg_fn, ()) ] def function_values( self, fn: Callable[..., BASE_EXPR] | BASE_EXPR, length: int | None = None ) -> dict[BASE_EXPR, BASE_EXPR]: """ - Given a callable that is a "function", meaning it returns a primitive or has a merge set, - returns a mapping of the function applied with its arguments to its values + Return the rows of a table-backed callable as calls mapped to values. + + Relations, constructors, and bodyless functions have tables. Eager and + builtin primitives do not. If length is specified, only the first `length` values will be returned. """ - ref, egg_name = self._callable_to_egg(fn) + ref, decls = resolve_callable(fn) + self._add_decls(decls) + self._require_table_backed(ref) + egg_name = self._state.callable_ref_to_egg(ref)[0] cmd = bindings.PrintFunction(span(1), egg_name, length, None, bindings.DefaultPrintFunctionMode()) - (output,) = self._run_program(cmd) + (output,) = self._state.run_program(cmd) assert isinstance(output, bindings.PrintFunctionOutput) signature = self.__egg_decls__.get_callable_decl(ref).signature assert isinstance(signature, FunctionSignature) @@ -1346,11 +1661,17 @@ def function_values( def lookup_function_value(self, expr: BASE_EXPR) -> BASE_EXPR | None: """ - Given an expression that is a function call, looks up the value of the function call if it exists. + Look up the current value of a table-backed call, if that row exists. + + Cost lookups use their associated cost table. Eager and builtin + primitive calls cannot be inspected with this method. """ runtime_expr = to_runtime_expr(expr) + self._add_decls(runtime_expr) typed_expr = runtime_expr.__egg_typed_expr__ assert isinstance(typed_expr.expr, CallDecl | GetCostDecl) + if isinstance(typed_expr.expr, CallDecl): + self._require_table_backed(typed_expr.expr.callable) egg_fn, typed_args = self._state.translate_call(typed_expr.expr) values_args = [self._state.typed_expr_to_value(a) for a in typed_args] possible_value = self._egraph.lookup_function(egg_fn, values_args) @@ -1364,12 +1685,30 @@ def lookup_function_value(self, expr: BASE_EXPR) -> BASE_EXPR | None: ), ) + def _require_table_backed(self, ref: CallableRef) -> None: + """Reject callable shapes that Egglog lowers without a queryable table.""" + decl = self.__egg_decls__.get_callable_decl(ref) + match decl: + case RelationDecl() | ConstructorDecl() | ConstantDecl(body=None): + return + case FunctionDecl(body=None, builtin=False) if not isinstance(ref, UnnamedFunctionRef): + return + case ConstantDecl() | FunctionDecl(): + pass + case _: + assert_never(decl) + msg = ( + "This operation requires a table-backed relation, constructor, or bodyless function; " + "eager and builtin primitives do not have tables" + ) + raise ValueError(msg) + def has_custom_cost(self, fn: ExprCallable) -> bool: """ Checks if the any custom costs have been set for this expression callable. """ resolved, _ = resolve_callable(fn) - return resolved in self._state.cost_callables + return resolved in self._state.cost_table_names def freeze(self) -> FrozenEGraph: # noqa: C901,PLR0912 """ @@ -1407,9 +1746,10 @@ def append_e_class_row(output: bindings.Value, tp: JustTypeRef, call: CallDecl, if is_subsumed: subsumed.append((tp, call)) + synthetic_let_names = {var.name for var in self._state.expr_to_letref_cache.values()} for name, fn in frozen.functions.items(): if fn.is_let_binding: - if name.startswith("$__expr_"): + if name in synthetic_let_names: continue for row in fn.rows: output_tp = self._state.egg_sort_to_type_ref[fn.output_sort] @@ -1419,9 +1759,7 @@ def append_e_class_row(output: bindings.Value, tp: JustTypeRef, call: CallDecl, if name in self._state.egg_fn_to_callable_refs: (callable_ref,) = self._state.egg_fn_to_callable_refs[name] else: - (callable_ref,) = ( - ref for ref in self._state.cost_callables if name == self._state.cost_table_name(ref) - ) + (callable_ref,) = (ref for ref, cost_name in self._state.cost_table_names.items() if name == cost_name) is_cost = True callable_decl = self.__egg_decls__.get_callable_decl(callable_ref) signature = callable_decl.signature @@ -1455,12 +1793,12 @@ def append_e_class_row(output: bindings.Value, tp: JustTypeRef, call: CallDecl, case FunctionDecl(): set_tp = self._state.egg_sort_to_type_ref[fn.output_sort] sets[call] = TypedExprDecl(set_tp, self._state.value_to_expr(set_tp, row.output)) - case ConstantDecl(type_ref): - if type_ref.ident.module == Ident.builtin("").module: + case ConstantDecl(): + if is_callable_decl_constructor(self.__egg_decls__, callable_decl): + append_e_class_row(row.output, output_tp, call, row.subsumed) + else: set_tp = self._state.egg_sort_to_type_ref[fn.output_sort] sets[call] = TypedExprDecl(set_tp, self._state.value_to_expr(set_tp, row.output)) - continue - append_e_class_row(row.output, output_tp, call, row.subsumed) case RelationDecl(): if row.subsumed: raise TypeError(f"Cannot freeze subsumed relation row for {callable_ref}") @@ -1567,7 +1905,7 @@ class Schedule(DelayedDeclarations): repeat until the schedule stops changing the e-graph. """ - # Defer declerations so that we can have rule generators that used not yet defined yet + # Defer declarations so that we can have rule generators that used not yet defined yet schedule: ScheduleDecl def __str__(self) -> str: @@ -1677,12 +2015,18 @@ def __egg_ident__(self) -> Ident: @dataclass class UnstableCombinedRuleset(Schedule): + _next_generated_ident: ClassVar[int] = 0 + __egg_decls_thunk__: Callable[[], Declarations] = field(init=False) schedule: RunDecl = field(init=False) ident: Ident | None + _generated_ident: Ident | None = field(default=None, init=False, repr=False) rulesets: InitVar[list[Ruleset | UnstableCombinedRuleset]] def __post_init__(self, rulesets: list[Ruleset | UnstableCombinedRuleset]) -> None: + if self.ident is None: + self._generated_ident = Ident(f"_combined_ruleset_{UnstableCombinedRuleset._next_generated_ident}") + UnstableCombinedRuleset._next_generated_ident += 1 self.schedule = RunDecl(self.__egg_ident__, ()) # Don't use thunk so that this is re-evaluated each time its requsted, so that additions inside will # be added after its been evaluated once. @@ -1690,7 +2034,10 @@ def __post_init__(self, rulesets: list[Ruleset | UnstableCombinedRuleset]) -> No @property def __egg_ident__(self) -> Ident: - return self.ident or Ident(f"combined_ruleset_{id(self)}") + if self.ident: + return self.ident + assert self._generated_ident is not None + return self._generated_ident def _create_egg_decls(self, *rulesets: Ruleset | UnstableCombinedRuleset) -> Declarations: decls = Declarations.create(*rulesets) @@ -1857,12 +2204,17 @@ def set_(lhs: BASE_EXPR) -> _SetBuilder[BASE_EXPR]: return _SetBuilder(lhs=lhs) -def rule(*facts: FactLike, ruleset: None = None, name: str | None = None) -> _RuleBuilder: +def rule( + *facts: FactLike, + ruleset: None = None, + name: str | None = None, + eval_mode: RuleEvalMode = "seminaive", +) -> _RuleBuilder: """Create a rule with the given facts.""" - return _RuleBuilder(facts=_fact_likes(facts), name=name, ruleset=ruleset) + return _RuleBuilder(facts=_fact_likes(facts), name=name, ruleset=ruleset, eval_mode=eval_mode) -def var(name: str, bound: type[T], egg_name: str | None = None) -> T: +def var(name: str, bound: TypeForm[T], egg_name: str | None = None) -> T: """Create a new variable with the given name and type.""" return cast("T", _var(name, bound, egg_name=egg_name)) @@ -2031,12 +2383,18 @@ class _RuleBuilder: facts: tuple[Fact, ...] name: str | None ruleset: Ruleset | None + eval_mode: RuleEvalMode def then(self, *actions: ActionLike) -> RewriteOrRule: actions = _action_likes(actions) rule = RewriteOrRule( Declarations.create(self.ruleset, *actions, *self.facts), - RuleDecl(tuple(a.action for a in actions), tuple(f.fact for f in self.facts), self.name), + RuleDecl( + tuple(a.action for a in actions), + tuple(f.fact for f in self.facts), + self.name, + self.eval_mode, + ), ) if self.ruleset: self.ruleset.append(rule) @@ -2046,7 +2404,9 @@ def __str__(self) -> str: # TODO: Figure out how to stringify rulebuilder that preserves statements args = list(map(str, self.facts)) if self.name is not None: - args.append(f"name={self.name}") + args.append(f"name={self.name!r}") + if self.eval_mode != "seminaive": + args.append(f"eval_mode={self.eval_mode!r}") if self.ruleset is not None: args.append(f"ruleset={self.ruleset}") return f"rule({', '.join(args)})" @@ -2056,18 +2416,20 @@ def expr_parts(expr: BaseExpr) -> TypedExprDecl: """ Returns the underlying type and decleration of the expression. Useful for testing structural equality or debugging. """ - if not isinstance(expr, RuntimeExpr): + if not isinstance(cast("object", expr), RuntimeExpr): raise TypeError(f"Expected a RuntimeExpr not {expr}") - return expr.__egg_typed_expr__ + return cast("RuntimeExpr", expr).__egg_typed_expr__ def to_runtime_expr(expr: BaseExpr) -> RuntimeExpr: - if not isinstance(expr, RuntimeExpr): + if not isinstance(cast("object", expr), RuntimeExpr): raise TypeError(f"Expected a RuntimeExpr not {expr}") - return expr + return cast("RuntimeExpr", expr) -def run(ruleset: Ruleset | None = None, *until: FactLike, scheduler: BackOff | None = None) -> Schedule: +def run( + ruleset: Ruleset | UnstableCombinedRuleset | None = None, *until: FactLike, scheduler: BackOff | None = None +) -> Schedule: """ Create a one-step run schedule. @@ -2105,6 +2467,12 @@ def back_off(match_limit: None | int = None, ban_length: None | int = None) -> B class BackOff: scheduler: BackOffDecl + def persistent(self) -> BackOff: + """ + Reuse this scheduler across repeated runs on the same egraph. + """ + return BackOff(replace(self.scheduler, id=uuid4(), persistent=True)) + def scope(self, schedule: Schedule) -> Schedule: """ Defines the scheduler to be created directly before the inner schedule, instead of the default which is at the @@ -2214,14 +2582,17 @@ def get_cost(expr: BaseExpr) -> i64: """ Return a lookup of the cost of an expression. If not set, won't match. """ - assert isinstance(expr, RuntimeExpr) - expr_decl = expr.__egg_typed_expr__.expr + runtime_expr = to_runtime_expr(expr) + expr_decl = runtime_expr.__egg_typed_expr__.expr if not isinstance(expr_decl, CallDecl): msg = "Can only get cost of function calls, not literals or variables" raise TypeError(msg) - return RuntimeExpr.__from_values__( - expr.__egg_decls__, - TypedExprDecl(JustTypeRef("i64"), GetCostDecl(expr_decl.callable, expr_decl.args)), + return cast( + "i64", + RuntimeExpr.__from_values__( + runtime_expr.__egg_decls__, + TypedExprDecl(JustTypeRef(Ident.builtin("i64")), GetCostDecl(expr_decl.callable, expr_decl.args)), + ), ) diff --git a/python/egglog/egraph_state.py b/python/egglog/egraph_state.py index 876466e1..92683578 100644 --- a/python/egglog/egraph_state.py +++ b/python/egglog/egraph_state.py @@ -4,11 +4,14 @@ from __future__ import annotations +import contextlib import re +import tempfile +import weakref from base64 import standard_b64decode, standard_b64encode -from collections import defaultdict -from dataclasses import dataclass, field, replace -from typing import TYPE_CHECKING, Literal, assert_never, overload +from dataclasses import InitVar, dataclass, field, replace +from pathlib import Path +from typing import TYPE_CHECKING, Literal, TextIO, assert_never, cast, overload from uuid import UUID import cloudpickle @@ -17,12 +20,12 @@ from . import bindings from ._tracing import call_with_current_trace from .declarations import * -from .declarations import ConstructorDecl +from .declarations import ConstructorDecl, is_callable_decl_constructor from .pretty import * from .type_constraint_solver import * if TYPE_CHECKING: - from collections.abc import Iterable + from collections.abc import Callable, Iterable __all__ = ["EGraphState", "span"] @@ -30,6 +33,58 @@ _TRACER = trace.get_tracer(__name__) +@dataclass +class _SavedEgglogFile: + """ + Shared state for file-backed command execution. + + When `save_egglog_string=True`, the Python wrapper keeps one cumulative temp + `.egg` file per high-level egraph. Each command is reparsed-and-run from a + synthetic one-command program padded with blank lines so egglog error spans + point at the real file path and the command's true line numbers inside that + cumulative source log. + + We keep the append handle open for performance and for easy post-failure + inspection. Successful commands are saved normally; failed commands are + saved as expected failures with trailing error comments. + """ + + path: str + file: TextIO + line_count: int = 0 + _finalizer: weakref.finalize = field(init=False, repr=False) + + def __post_init__(self) -> None: + self._finalizer = weakref.finalize(self, _close_saved_egglog_file, self.file, self.path) + + def close(self) -> None: + self._finalizer() + + +def _close_saved_egglog_file(file: TextIO, path: str) -> None: + try: + file.close() + finally: + with contextlib.suppress(FileNotFoundError): + Path(path).unlink() + + +def _normalize_global_let_name(name: str) -> str: + return name if name.startswith("$") else f"${name}" + + +def _saved_egglog_failure_message(error: bindings.EggSmolError) -> str: + lines = [line.strip() for line in error.context.splitlines() if line.strip()] + for message in lines: + if "Failed" in message or "failed" in message: + return message + if len(lines) > 1 and lines[0].startswith("In "): + return lines[1] + if lines: + return lines[-1] + return str(error) + + def span(frame_index: int = 0) -> bindings.RustSpan: """ Returns a span for the current file and line. @@ -44,10 +99,6 @@ def span(frame_index: int = 0) -> bindings.RustSpan: return bindings.RustSpan("", 0, 0) -def _normalize_global_let_name(name: str) -> str: - return name if name.startswith("$") else f"${name}" - - @dataclass class EGraphState: """ @@ -57,16 +108,20 @@ class EGraphState: """ egraph: bindings.EGraph + save_egglog_string: InitVar[bool] = False + egglog_file_state: _SavedEgglogFile | None = field(default=None, repr=False) # The declarations we have added. __egg_decls__: Declarations = field(default_factory=Declarations) # Mapping of added rulesets to the added rules rulesets: dict[Ident, set[RewriteOrRuleDecl]] = field(default_factory=dict) + # Persistent schedulers live outside a single run-schedule command; only emit their let once per active scope. + registered_persistent_schedulers: set[UUID] = field(default_factory=set) # Bidirectional mapping between egg function names and python callable references. # Note that there are possibly multiple callable references for a single egg function name, like `+` # for both int and rational classes. egg_fn_to_callable_refs: dict[str, set[CallableRef]] = field( - default_factory=lambda: defaultdict(set, {"!=": {FunctionRef(Ident.builtin("!="))}}) + default_factory=lambda: {"!=": {FunctionRef(Ident.builtin("!="))}} ) callable_ref_to_egg_fn: dict[CallableRef, tuple[str, bool]] = field( default_factory=lambda: {FunctionRef(Ident.builtin("!=")): ("!=", False)} @@ -76,11 +131,19 @@ class EGraphState: type_ref_to_egg_sort: dict[JustTypeRef, str] = field(default_factory=dict) egg_sort_to_type_ref: dict[str, JustTypeRef] = field(default_factory=dict) - # Cache of egg expressions for converting to egg + # Cache of direct structural egg expressions for converting to egg. expr_to_egg_cache: dict[ExprDecl, bindings._Expr] = field(default_factory=dict) - - # Callables which have cost tables associated with them - 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) + # Cache of synthetic let references introduced for top-level command lowering. + # This stays separate from `expr_to_egg_cache` so nested rule/rewrite lowering + # can always rebuild structural surface syntax instead of leaking a previously + # synthesized `$__expr_n` binding across contexts. + expr_to_letref_cache: dict[ExprDecl, bindings.Var] = field(default_factory=dict) + + # Callables with custom-cost tables and the reserved backend names for those tables. + cost_table_names: dict[CallableRef, str] = field(default_factory=dict) # Counter for deterministic synthetic let bindings created while lowering expressions to egg. expr_to_let_counter: int = 0 # Counter for deterministic synthetic names assigned to unnamed functions. @@ -91,28 +154,94 @@ class EGraphState: # Mapping from numeric name (str) to command decl rule_name_to_command_decl: dict[str, RuleDecl | BiRewriteDecl | RewriteDecl] = field(default_factory=dict) + def __post_init__(self, save_egglog_string: bool) -> None: + if save_egglog_string and self.egglog_file_state is None: + # Keep one persistent temp `.egg` file per high-level egraph so parse errors + # can point at a stable filename the user can open after a failure. + egglog_file = tempfile.NamedTemporaryFile( # noqa: SIM115 - kept open for incremental appends + mode="w+", encoding="utf-8", suffix=".egg", delete=False + ) + self.egglog_file_state = _SavedEgglogFile(egglog_file.name, cast("TextIO", egglog_file)) + def copy(self) -> EGraphState: """ Returns a copy of the state. The egraph reference is kept the same. Used for pushing/popping. """ return EGraphState( egraph=self.egraph, + save_egglog_string=self.egglog_file_state is not None, + egglog_file_state=self.egglog_file_state, __egg_decls__=self.__egg_decls__.copy(), rulesets={k: v.copy() for k, v in self.rulesets.items()}, - egg_fn_to_callable_refs=defaultdict(set, {k: v.copy() for k, v in self.egg_fn_to_callable_refs.items()}), + registered_persistent_schedulers=self.registered_persistent_schedulers.copy(), + egg_fn_to_callable_refs={k: v.copy() for k, v in self.egg_fn_to_callable_refs.items()}, callable_ref_to_egg_fn=self.callable_ref_to_egg_fn.copy(), type_ref_to_egg_sort=self.type_ref_to_egg_sort.copy(), egg_sort_to_type_ref=self.egg_sort_to_type_ref.copy(), expr_to_egg_cache=self.expr_to_egg_cache.copy(), - cost_callables=self.cost_callables.copy(), + expr_to_let_egg_cache=self.expr_to_let_egg_cache.copy(), + expr_to_letref_cache=self.expr_to_letref_cache.copy(), + cost_table_names=self.cost_table_names.copy(), expr_to_let_counter=self.expr_to_let_counter, unnamed_function_counter=self.unnamed_function_counter, rule_name_counter=self.rule_name_counter, rule_name_to_command_decl=self.rule_name_to_command_decl.copy(), ) - def _run_program(self, *commands: bindings._Command) -> list[bindings._CommandOutput]: - return call_with_current_trace(self.egraph.run_program, *commands) + def egglog_string(self) -> str: + if self.egglog_file_state is None: + msg = "Can't get egglog string unless EGraph created with save_egglog_string=True" + raise ValueError(msg) + if self.egglog_file_state.file.closed: + msg = "Can't get egglog string after the saved transcript has been closed" + raise ValueError(msg) + # The append handle stays open for execution, so flush before reading the saved source. + self.egglog_file_state.file.flush() + with open(self.egglog_file_state.path, encoding="utf-8") as saved_file: + return saved_file.read() + + def close(self) -> None: + if self.egglog_file_state is not None: + self.egglog_file_state.close() + + def run_program(self, *commands: bindings._Command) -> list[bindings._CommandOutput]: + if not commands: + return [] + if self.egglog_file_state is None: + return call_with_current_trace(self.egraph.run_program, *commands) + + outputs: list[bindings._CommandOutput] = [] + for command in commands: + command_text = str(command).rstrip("\n") + "\n" + start_line = self.egglog_file_state.line_count + 1 + + # Parse and run just this command in Rust, but pad it with blank lines + # so the span lines match its location in the cumulative saved source file. + padded_command = ("\n" * (start_line - 1)) + command_text + try: + command_outputs = call_with_current_trace( + self.egraph.parse_and_run_program, padded_command, filename=self.egglog_file_state.path + ) + except bindings.EggSmolError as error: + fail_command_text = str(bindings.Fail(span(), command)).rstrip("\n") + saved_text = f"{fail_command_text} ; {_saved_egglog_failure_message(error)}\n" + self.egglog_file_state.file.write(saved_text) + self.egglog_file_state.file.flush() + self.egglog_file_state.line_count += saved_text.count("\n") + raise + self.egglog_file_state.file.write(command_text) + self.egglog_file_state.file.flush() + self.egglog_file_state.line_count += command_text.count("\n") + outputs.extend(command_outputs) + return outputs + + @staticmethod + def _persistent_scheduler_name(scheduler: BackOffDecl) -> str: + return f"_persistent_scheduler_{scheduler.id.hex}" + + @staticmethod + def _local_scheduler_name(index: int) -> str: + return f"_scheduler_{index}" @_TRACER.start_as_current_span("run_schedule_to_egg") def run_schedule_to_egg(self, schedule: ScheduleDecl) -> bindings._Command: @@ -122,9 +251,14 @@ def run_schedule_to_egg(self, schedule: ScheduleDecl) -> bindings._Command: If there exists any custom schedulers in the schedule, it will be turned into a custom extract command otherwise will be a normal run command. """ - processed_schedule = self._process_schedule(schedule) + processed_schedule, persistent_schedulers = self._process_schedule(schedule) if processed_schedule is None: return bindings.RunSchedule(self._schedule_to_egg(schedule)) + for scheduler in persistent_schedulers: + if scheduler.id in self.registered_persistent_schedulers: + continue + self.run_program(self._persistent_scheduler_to_egg(scheduler)) + self.registered_persistent_schedulers.add(scheduler.id) top_level_schedules = self._schedule_with_scheduler_to_egg(processed_schedule, []) if len(top_level_schedules) == 1: schedule_expr = top_level_schedules[0] @@ -132,7 +266,9 @@ def run_schedule_to_egg(self, schedule: ScheduleDecl) -> bindings._Command: schedule_expr = bindings.Call(span(), "seq", top_level_schedules) return bindings.UserDefined(span(), "run-schedule", [schedule_expr]) - def _process_schedule(self, schedule: ScheduleDecl) -> ScheduleDecl | None: + def _process_schedule( # noqa: C901 + self, schedule: ScheduleDecl + ) -> tuple[ScheduleDecl | None, tuple[BackOffDecl, ...]]: """ Processes a schedule to determine if it contains any custom schedulers. @@ -142,17 +278,27 @@ def _process_schedule(self, schedule: ScheduleDecl) -> ScheduleDecl | None: Also processes all rulesets in the schedule to make sure they are registered. """ bound_schedulers: list[UUID] = [] - unbound_schedulers: list[BackOffDecl] = [] + unbound_schedulers: dict[UUID, BackOffDecl] = {} + persistent_schedulers: dict[UUID, BackOffDecl] = {} + has_bound_scheduler = False def helper(s: ScheduleDecl) -> None: + nonlocal has_bound_scheduler match s: case LetSchedulerDecl(scheduler, inner): + has_bound_scheduler = True bound_schedulers.append(scheduler.id) - return helper(inner) + try: + return helper(inner) + finally: + bound_schedulers.pop() case RunDecl(ruleset_name, _, scheduler): self.ruleset_to_egg(ruleset_name) if scheduler and scheduler.id not in bound_schedulers: - unbound_schedulers.append(scheduler) + if scheduler.persistent: + persistent_schedulers[scheduler.id] = scheduler + else: + unbound_schedulers[scheduler.id] = scheduler case SaturateDecl(inner) | RepeatDecl(inner, _): return helper(inner) case SequenceDecl(schedules): @@ -163,11 +309,11 @@ def helper(s: ScheduleDecl) -> None: return None helper(schedule) - if not bound_schedulers and not unbound_schedulers: - return None - for scheduler in unbound_schedulers: + if not has_bound_scheduler and not unbound_schedulers and not persistent_schedulers: + return None, () + for scheduler in unbound_schedulers.values(): schedule = LetSchedulerDecl(scheduler, schedule) - return schedule + return schedule, tuple(persistent_schedulers.values()) def _schedule_to_egg(self, schedule: ScheduleDecl) -> bindings._Schedule: msg = "Should never reach this, let schedulers should be handled by custom scheduler" @@ -191,7 +337,7 @@ def _schedule_to_egg(self, schedule: ScheduleDecl) -> bindings._Schedule: assert_never(schedule) def _schedule_with_scheduler_to_egg( # noqa: C901, PLR0912 - self, schedule: ScheduleDecl, bound_schedulers: list[UUID] + self, schedule: ScheduleDecl, bound_schedulers: list[BackOffDecl] ) -> list[bindings._Expr]: """ Turns a scheduler into an egg expression, to be used with a custom extract command. @@ -199,9 +345,11 @@ def _schedule_with_scheduler_to_egg( # noqa: C901, PLR0912 The bound_schedulers is a list of all the schedulers that have been bound. We can lookup their name as `_scheduler_{index}`. """ match schedule: - case LetSchedulerDecl(BackOffDecl(id, match_limit, ban_length), inner): - name = f"_scheduler_{len(bound_schedulers)}" - bound_schedulers.append(id) + case LetSchedulerDecl(scheduler, inner): + match_limit = scheduler.match_limit + ban_length = scheduler.ban_length + name = self._local_scheduler_name(len(bound_schedulers)) + bound_schedulers.append(scheduler) args: list[bindings._Expr] = [] if match_limit is not None: args.append(bindings.Var(span(), ":match-limit")) @@ -211,12 +359,20 @@ def _schedule_with_scheduler_to_egg( # noqa: C901, PLR0912 args.append(bindings.Lit(span(), bindings.Int(ban_length))) back_off_decl = bindings.Call(span(), "back-off", args) let_decl = bindings.Call(span(), "let-scheduler", [bindings.Var(span(), name), back_off_decl]) - return [let_decl, *self._schedule_with_scheduler_to_egg(inner, bound_schedulers)] + try: + inner_exprs = self._schedule_with_scheduler_to_egg(inner, bound_schedulers) + finally: + bound_schedulers.pop() + return [bindings.Call(span(), "seq", [let_decl, *inner_exprs])] case RunDecl(ruleset_ident, until, scheduler): args = [bindings.Var(span(), str(ruleset_ident))] if scheduler: name = "run-with" - scheduler_name = f"_scheduler_{bound_schedulers.index(scheduler.id)}" + scheduler_name = self._persistent_scheduler_name(scheduler) + for i in range(len(bound_schedulers) - 1, -1, -1): + if bound_schedulers[i].id == scheduler.id: + scheduler_name = self._local_scheduler_name(i) + break args.insert(0, bindings.Var(span(), scheduler_name)) else: name = "run" @@ -254,15 +410,33 @@ def _schedule_with_scheduler_to_egg( # noqa: C901, PLR0912 case _: assert_never(schedule) - def ruleset_to_egg(self, ident: Ident) -> None: + def _persistent_scheduler_to_egg(self, scheduler: BackOffDecl) -> bindings._Command: + args: list[bindings._Expr] = [] + if scheduler.match_limit is not None: + args.append(bindings.Var(span(), ":match-limit")) + args.append(bindings.Lit(span(), bindings.Int(scheduler.match_limit))) + if scheduler.ban_length is not None: + args.append(bindings.Var(span(), ":ban-length")) + args.append(bindings.Lit(span(), bindings.Int(scheduler.ban_length))) + back_off_decl = bindings.Call(span(), "back-off", args) + return bindings.UserDefined( + span(), + "let-scheduler", + [bindings.Var(span(), self._persistent_scheduler_name(scheduler)), back_off_decl], + ) + + def ruleset_to_egg(self, ident: Ident) -> None: # noqa: C901 """ Registers a ruleset if it's not already registered. """ + if ident.name == "" and ident not in self.__egg_decls__._rulesets: + self.rulesets.setdefault(ident, set()) + return match self.__egg_decls__._rulesets[ident]: case RulesetDecl(rules): if ident not in self.rulesets: if str(ident): - self._run_program(bindings.AddRuleset(span(), str(ident))) + self.run_program(bindings.AddRuleset(span(), str(ident))) added_rules = self.rulesets[ident] = set() else: added_rules = self.rulesets[ident] @@ -271,7 +445,7 @@ def ruleset_to_egg(self, ident: Ident) -> None: continue cmd = self.command_to_egg(rule, ident) if cmd is not None: - self._run_program(cmd) + self.run_program(cmd) added_rules.add(rule) case CombinedRulesetDecl(rulesets): if ident in self.rulesets: @@ -279,7 +453,7 @@ def ruleset_to_egg(self, ident: Ident) -> None: self.rulesets[ident] = set() for ruleset in rulesets: self.ruleset_to_egg(ruleset) - self._run_program(bindings.UnstableCombinedRuleset(span(), str(ident), list(map(str, rulesets)))) + self.run_program(bindings.UnstableCombinedRuleset(span(), str(ident), list(map(str, rulesets)))) def command_to_egg(self, cmd: CommandDecl, ruleset: Ident) -> bindings._Command | None: match cmd: @@ -296,40 +470,62 @@ def command_to_egg(self, cmd: CommandDecl, ruleset: Ident) -> bindings._Command span(), self._expr_to_egg(lhs), self._expr_to_egg(rhs), - [self.fact_to_egg(c) for c in conditions], + [self.fact_to_egg(c, expr_to_let=False) for c in conditions], name, ) egg_cmd: bindings._Command if isinstance(cmd, RewriteDecl): self.rule_name_to_command_decl[name] = cmd egg_cmd = bindings.RewriteCommand(str(ruleset), rewrite, cmd.subsume) + # Saving a transcript executes the serialized command, whose syntax does not + # preserve the internal rewrite name. The engine then reports the serialized + # rewrite itself as its name, so retain that alias for RunReport translation. + serialized_name = str(egg_cmd) + self.rule_name_to_command_decl[serialized_name] = cmd + # Backend report names render every Symbol with single quotes, while + # command serialization accepts symbols as double-quoted literals. + reported_name = serialized_name.replace('"', "'") + self.rule_name_to_command_decl[reported_name] = cmd else: self.rule_name_to_command_decl[f"{name}=>"] = cmd self.rule_name_to_command_decl[f"{name}<="] = cmd egg_cmd = bindings.BiRewriteCommand(str(ruleset), rewrite) + serialized_name = str(egg_cmd) + reported_name = serialized_name.replace('"', "'") + for suffix in ("=>", "<="): + self.rule_name_to_command_decl[f"{serialized_name}{suffix}"] = cmd + self.rule_name_to_command_decl[f"{reported_name}{suffix}"] = cmd return egg_cmd - case RuleDecl(head, body, name): + case RuleDecl(head, body, name, eval_mode): if not name: name = str(self.rule_name_counter) self.rule_name_counter += 1 self.rule_name_to_command_decl[name] = cmd + binding_eval_mode = cast( + "bindings.Seminaive | bindings.Naive | bindings.UnsafeSeminaive", + { + "seminaive": bindings.Seminaive(), + "naive": bindings.Naive(), + "unsafe-seminaive": bindings.UnsafeSeminaive(), + }[eval_mode], + ) return bindings.RuleCommand( bindings.Rule( span(), [self.action_to_egg(a) for a in head], - [self.fact_to_egg(f) for f in body], - name, + [self.fact_to_egg(f, expr_to_let=False) for f in body], + name or "", str(ruleset), + binding_eval_mode, ) ) - # TODO: Replace with just constants value and looking at REF of function case DefaultRewriteDecl(ref, expr, subsume): sig = self.__egg_decls__.get_callable_decl(ref).signature assert isinstance(sig, FunctionSignature) # Replace args with rule_var_name mapping arg_mapping = tuple( - TypedExprDecl(tp.to_just(), UnboundVarDecl(name)) - for name, tp in zip(sig.arg_names, sig.arg_types, strict=False) + TypedExprDecl(tp.to_just(), UnboundVarDecl(name, f"_{i}")) + for i, (name, tp) in enumerate(zip(sig.arg_names, sig.arg_types, strict=True)) ) rewrite_decl = RewriteDecl( sig.semantic_return_type.to_just(), CallDecl(ref, arg_mapping), expr, (), subsume @@ -342,29 +538,48 @@ def command_to_egg(self, cmd: CommandDecl, ruleset: Ident) -> bindings._Command def action_to_egg(self, action: ActionDecl) -> bindings._Action: ... @overload - def action_to_egg(self, action: ActionDecl, expr_to_let: Literal[True] = ...) -> bindings._Action | None: ... + def action_to_egg( + self, + action: ActionDecl, + expr_to_let: Literal[True] = ..., + ) -> bindings._Action | None: ... + + @overload + def action_to_egg(self, action: ActionDecl, expr_to_let: bool) -> bindings._Action | None: ... - def action_to_egg(self, action: ActionDecl, expr_to_let: bool = False) -> bindings._Action | None: # noqa: C901, PLR0911, PLR0912 + def action_to_egg( # noqa: C901, PLR0911, PLR0912 + self, + action: ActionDecl, + expr_to_let: bool = False, + ) -> bindings._Action | None: match action: case LetDecl(name, typed_expr): var_decl = LetRefDecl(name) var_egg = self._expr_to_egg(var_decl) self.expr_to_egg_cache[var_decl] = var_egg - return bindings.Let(span(), var_egg.name, self.typed_expr_to_egg(typed_expr)) + return bindings.Let( + span(), + var_egg.name, + self.typed_expr_to_egg(typed_expr, expr_to_let=expr_to_let), + ) case SetDecl(tp, call, rhs): self.type_ref_to_egg(tp) egg_fn, typed_args = self.translate_call(call) return bindings.Set( - span(), egg_fn, [self.typed_expr_to_egg(arg, False) for arg in typed_args], self._expr_to_egg(rhs) + span(), + egg_fn, + [self.typed_expr_to_egg(arg, expr_to_let) for arg in typed_args], + self._expr_to_egg(rhs, expr_to_let=expr_to_let), ) case ExprActionDecl(typed_expr): - if expr_to_let: - maybe_typed_expr = self._transform_let(typed_expr) - if maybe_typed_expr: - typed_expr = maybe_typed_expr - else: - return None - return bindings.Expr_(span(), self.typed_expr_to_egg(typed_expr)) + if not isinstance(typed_expr.expr, CallDecl): + msg = "Top-level egglog expr commands must be calls" + raise ValueError(msg) # noqa: TRY004 - preserve the public validation error + egg_expr = self.typed_expr_to_egg(typed_expr, expr_to_let=expr_to_let) + if isinstance(egg_expr, bindings.Var): + return None + assert isinstance(egg_expr, bindings.Call) + return bindings.Expr_(span(), egg_expr) case ChangeDecl(tp, call, change): self.type_ref_to_egg(tp) egg_fn, typed_args = self.translate_call(call) @@ -377,18 +592,25 @@ def action_to_egg(self, action: ActionDecl, expr_to_let: bool = False) -> bindin case _: assert_never(change) return bindings.Change( - span(), egg_change, egg_fn, [self.typed_expr_to_egg(arg, False) for arg in typed_args] + span(), + egg_change, + egg_fn, + [self.typed_expr_to_egg(arg, expr_to_let) for arg in typed_args], ) case UnionDecl(tp, lhs, rhs): self.type_ref_to_egg(tp) - return bindings.Union(span(), self._expr_to_egg(lhs), self._expr_to_egg(rhs)) + return bindings.Union( + span(), + self._expr_to_egg(lhs, expr_to_let=expr_to_let), + self._expr_to_egg(rhs, expr_to_let=expr_to_let), + ) case PanicDecl(name): return bindings.Panic(span(), name) case SetCostDecl(tp, expr, cost): self.type_ref_to_egg(tp) - cost_table = self.create_cost_table(expr.callable) - args_egg = [self.typed_expr_to_egg(x, False) for x in expr.args] - return bindings.Set(span(), cost_table, args_egg, self._expr_to_egg(cost)) + cost_table, typed_args = self.translate_call(GetCostDecl(expr.callable, expr.args)) + args_egg = [self.typed_expr_to_egg(x, expr_to_let) for x in typed_args] + return bindings.Set(span(), cost_table, args_egg, self._expr_to_egg(cost, expr_to_let=expr_to_let)) case _: assert_never(action) @@ -396,25 +618,28 @@ def create_cost_table(self, ref: CallableRef) -> str: """ Creates the egg cost table if needed and gets the name of the table. """ - name = self.cost_table_name(ref) - if ref not in self.cost_callables: - self.cost_callables.add(ref) - signature = self.__egg_decls__.get_callable_decl(ref).signature - assert isinstance(signature, FunctionSignature), "Can only add cost tables for functions" - signature = replace(signature, return_type=TypeRefWithVars(Ident.builtin("i64"))) - self._run_program(bindings.FunctionCommand(span(), name, self._signature_to_egg_schema(signature), None)) + if ref in self.cost_table_names: + return self.cost_table_names[ref] + base_name = f"cost_table_{self.callable_ref_to_egg(ref)[0]}" + name = self._allocate_name((base_name,), self._backend_symbol_is_occupied) + signature = self.__egg_decls__.get_callable_decl(ref).signature + assert isinstance(signature, FunctionSignature), "Can only add cost tables for functions" + signature = replace(signature, return_type=TypeRefWithVars(Ident.builtin("i64"))) + self.run_program(bindings.FunctionCommand(span(), name, self._signature_to_egg_schema(signature), None)) + self.cost_table_names[ref] = name return name - def cost_table_name(self, ref: CallableRef) -> str: - 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: match fact: case EqDecl(tp, left, right): self.type_ref_to_egg(tp) - return bindings.Eq(span(), self._expr_to_egg(left), self._expr_to_egg(right)) + return bindings.Eq( + span(), + self._expr_to_egg(left, expr_to_let=expr_to_let), + self._expr_to_egg(right, expr_to_let=expr_to_let), + ) case ExprFactDecl(typed_expr): - return bindings.Fact(self.typed_expr_to_egg(typed_expr, False)) + return bindings.Fact(self.typed_expr_to_egg(typed_expr, expr_to_let=expr_to_let)) case _: assert_never(fact) @@ -427,45 +652,59 @@ def callable_ref_to_egg(self, ref: CallableRef) -> tuple[str, bool]: # noqa: C9 if ref in self.callable_ref_to_egg_fn: return self.callable_ref_to_egg_fn[ref] decl = self.__egg_decls__.get_callable_decl(ref) - egg_name = decl.egg_name or _sanitize_egg_ident(self._generate_callable_egg_name(ref)) - self.egg_fn_to_callable_refs[egg_name].add(ref) - reverse_args = False + egg_name = decl.egg_name or self._allocate_callable_egg_name(ref) + self.egg_fn_to_callable_refs.setdefault(egg_name, set()).add(ref) + callable_signature = decl.signature + reverse_args = callable_signature.reverse_args if isinstance(callable_signature, FunctionSignature) else False match decl: case RelationDecl(arg_types, _, _): - self._run_program(bindings.Relation(span(), egg_name, [self.type_ref_to_egg(a) for a in arg_types])) - case ConstantDecl(tp, _): - # Use constructor declaration instead of constant b/c constants cannot be extracted - # https://github.com/egraphs-good/egglog/issues/334 - is_function = self.__egg_decls__._classes[tp.ident].builtin - schema = bindings.Schema([], self.type_ref_to_egg(tp)) - if is_function: - self._run_program(bindings.FunctionCommand(span(), egg_name, schema, None)) + self.run_program(bindings.Relation(span(), egg_name, [self.type_ref_to_egg(a) for a in arg_types])) + case ConstantDecl(tp, _, body, merge): + if body is not None: + self.run_program(self._primitive_command_to_egg(egg_name, decl.signature, body)) else: - self._run_program(bindings.Constructor(span(), egg_name, schema, None, False)) - case FunctionDecl(signature, builtin, _, merge): - if isinstance(signature, FunctionSignature): - reverse_args = signature.reverse_args - if not builtin: - assert isinstance(signature, FunctionSignature), "Cannot turn special function to egg" - # Compile functions that return unit to relations, because these show up in methods where you - # cant use the relation helper - schema = self._signature_to_egg_schema(signature) - if signature.return_type == TypeRefWithVars(Ident.builtin("Unit")): - if merge: - msg = "Cannot specify a merge function for a function that returns unit" - raise ValueError(msg) - self._run_program(bindings.Relation(span(), egg_name, schema.input)) - else: - self._run_program( + # Use constructor declaration instead of constant b/c constants cannot be extracted + # https://github.com/egraphs-good/egglog/issues/334 + is_function = self.__egg_decls__._classes[tp.ident].builtin or merge is not None + schema = bindings.Schema([], self.type_ref_to_egg(tp)) + if is_function: + self.run_program( bindings.FunctionCommand( span(), egg_name, - self._signature_to_egg_schema(signature), + schema, self._expr_to_egg(merge) if merge else None, - ), + ) ) + else: + self.run_program(bindings.Constructor(span(), egg_name, schema, None, False)) + case FunctionDecl(signature=signature, builtin=builtin, body=body, merge=merge): + if not builtin: + assert isinstance(signature, FunctionSignature), "Cannot turn special function to egg" + if body is None and isinstance(ref, UnnamedFunctionRef): + body = ref.res + if body is not None: + self.run_program(self._primitive_command_to_egg(egg_name, signature, body)) + else: + # Compile functions that return unit to relations, because these show up in methods where you + # cant use the relation helper + schema = self._signature_to_egg_schema(signature) + if signature.return_type == TypeRefWithVars(Ident.builtin("Unit")): + if merge: + msg = "Cannot specify a merge function for a function that returns unit" + raise ValueError(msg) + self.run_program(bindings.Relation(span(), egg_name, schema.input)) + else: + self.run_program( + bindings.FunctionCommand( + span(), + egg_name, + schema, + self._expr_to_egg(merge) if merge else None, + ), + ) case ConstructorDecl(signature, _, cost, unextractable): - self._run_program( + self.run_program( bindings.Constructor( span(), egg_name, @@ -479,9 +718,40 @@ def callable_ref_to_egg(self, ref: CallableRef) -> tuple[str, bool]: # noqa: C9 self.callable_ref_to_egg_fn[ref] = egg_name, reverse_args return egg_name, reverse_args + def _primitive_command_to_egg( + self, + egg_name: str, + signature: FunctionSignature, + body: TypedExprDecl, + ) -> bindings.UserDefined: + backend_arg_types = signature.arg_types[::-1] if signature.reverse_args else signature.arg_types + input_sort_expr = self._primitive_input_sorts_to_egg([ + self.type_ref_to_egg(arg_type.to_just()) for arg_type in backend_arg_types + ]) + output_sort_expr = bindings.Var(span(), self.type_ref_to_egg(signature.semantic_return_type.to_just())) + return bindings.UserDefined( + span(), + "primitive", + [ + bindings.Var(span(), egg_name), + input_sort_expr, + output_sort_expr, + self.typed_expr_to_egg(body, expr_to_let=False), + ], + ) + + def _primitive_input_sorts_to_egg(self, sort_names: list[str]) -> bindings._Expr: + if not sort_names: + return bindings.Lit(span(), bindings.Unit()) + if len(sort_names) == 1: + return bindings.Var(span(), sort_names[0]) + first, *rest = sort_names + return bindings.Call(span(), first, [bindings.Var(span(), sort_name) for sort_name in rest]) + def _signature_to_egg_schema(self, signature: FunctionSignature) -> bindings.Schema: + backend_arg_types = signature.arg_types[::-1] if signature.reverse_args else signature.arg_types return bindings.Schema( - [self.type_ref_to_egg(a.to_just()) for a in signature.arg_types], + [self.type_ref_to_egg(a.to_just()) for a in backend_arg_types], self.type_ref_to_egg(signature.semantic_return_type.to_just()), ) @@ -495,7 +765,10 @@ def type_ref_to_egg(self, ref: JustTypeRef) -> str: except KeyError: pass decl = self.__egg_decls__._classes[ref.ident] - self.type_ref_to_egg_sort[ref] = egg_name = (not ref.args and decl.egg_name) or _generate_type_egg_name(ref) + arg_names = [self.type_ref_to_egg(arg) for arg in ref.args] + self.type_ref_to_egg_sort[ref] = egg_name = (not ref.args and decl.egg_name) or self._allocate_type_egg_name( + ref, decl, arg_names + ) self.egg_sort_to_type_ref[egg_name] = ref if decl.builtin: @@ -515,7 +788,7 @@ def type_ref_to_egg(self, ref: JustTypeRef) -> str: else: type_args = [bindings.Var(span(), self.type_ref_to_egg(a)) for a in ref.args] assert decl.egg_name - self._run_program(bindings.Sort(span(), egg_name, (decl.egg_name, type_args))) + self.run_program(bindings.Sort(span(), egg_name, (decl.egg_name, type_args))) # For builtin classes, let's also make sure we have the mapping of all egg fn names for class methods. # these can be created even without adding them to the e-graph, like `vec-empty` which can be extracted @@ -525,7 +798,7 @@ def type_ref_to_egg(self, ref: JustTypeRef) -> str: if decl.init: self.callable_ref_to_egg(InitRef(ref.ident)) else: - self._run_program(bindings.Sort(span(), egg_name, None)) + self.run_program(bindings.Sort(span(), egg_name, None)) return egg_name @@ -541,8 +814,8 @@ def op_mapping(self) -> dict[str, str]: for k, v in self.egg_fn_to_callable_refs.items() if len(v) == 1 } | { - self.cost_table_name(ref): f"cost({pretty_callable_ref(self.__egg_decls__, ref, include_all_args=True)})" - for ref in self.cost_callables + name: f"cost({pretty_callable_ref(self.__egg_decls__, ref, include_all_args=True)})" + for ref, name in self.cost_table_names.items() } def possible_egglog_functions(self, names: list[str]) -> Iterable[str]: @@ -550,54 +823,69 @@ def possible_egglog_functions(self, names: list[str]) -> Iterable[str]: Given a list of egglog functions, returns all the possible Python function strings """ for name in names: - for c in self.egg_fn_to_callable_refs[name]: + for c in self.egg_fn_to_callable_refs.get(name, ()): yield pretty_callable_ref(self.__egg_decls__, c) - def typed_expr_to_egg(self, typed_expr_decl: TypedExprDecl, transform_let: bool = True) -> bindings._Expr: + def typed_expr_to_egg( + self, + typed_expr_decl: TypedExprDecl, + expr_to_let: bool = True, + ) -> bindings._Expr: # transform all expressions with multiple parents into a let binding, so that less expressions # are sent to egglog. Only for performance reasons. - if transform_let: + if expr_to_let: have_multiple_parents = _exprs_multiple_parents(typed_expr_decl) for expr in reversed(have_multiple_parents): self._transform_let(expr) self.type_ref_to_egg(typed_expr_decl.tp) - return self._expr_to_egg(typed_expr_decl.expr) + return self._expr_to_egg(typed_expr_decl.expr, expr_to_let=expr_to_let) def _transform_let(self, typed_expr: TypedExprDecl) -> TypedExprDecl | None: """ Rewrites this expression as a let binding if it's not already a let binding. """ - if isinstance(self.expr_to_egg_cache.get(typed_expr.expr), bindings.Var): + if not isinstance(typed_expr.expr, CallDecl): + return typed_expr + if not is_callable_decl_constructor( + self.__egg_decls__, self.__egg_decls__.get_callable_decl(typed_expr.expr.callable) + ): + return typed_expr + if typed_expr.expr in self.expr_to_letref_cache: return None - var_decl = LetRefDecl(f"$__expr_{self.expr_to_let_counter}") - self.expr_to_let_counter += 1 + var_decl = LetRefDecl(self._allocate_synthetic_let_name()) var_egg = self._expr_to_egg(var_decl) - cmd = bindings.ActionCommand(bindings.Let(span(), var_egg.name, self.typed_expr_to_egg(typed_expr))) + cmd = bindings.ActionCommand(bindings.Let(span(), var_egg.name, self.typed_expr_to_egg(typed_expr, True))) try: - self._run_program(cmd) + self.run_program(cmd) # errors when creating let bindings for things like `(vec-empty)` except bindings.EggSmolError: return typed_expr - self.expr_to_egg_cache[typed_expr.expr] = var_egg + self.expr_to_letref_cache[typed_expr.expr] = var_egg self.expr_to_egg_cache[var_decl] = var_egg return None @overload - def _expr_to_egg(self, expr_decl: CallDecl) -> bindings.Call: ... + def _expr_to_egg(self, expr_decl: CallDecl, *, expr_to_let: bool = ...) -> bindings.Call: ... @overload - def _expr_to_egg(self, expr_decl: UnboundVarDecl | LetRefDecl) -> bindings.Var: ... + def _expr_to_egg(self, expr_decl: UnboundVarDecl | LetRefDecl, *, expr_to_let: bool = ...) -> bindings.Var: ... @overload - def _expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr: ... + def _expr_to_egg(self, expr_decl: ExprDecl, *, expr_to_let: bool = ...) -> bindings._Expr: ... - def _expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr: # noqa: PLR0912,C901 + def _expr_to_egg(self, expr_decl: ExprDecl, *, expr_to_let: bool = False) -> bindings._Expr: # noqa: PLR0912,C901 """ Convert an ExprDecl to an egg expression. """ + if expr_to_let: + try: + return self.expr_to_letref_cache[expr_decl] + except KeyError: + pass + cache = self.expr_to_let_egg_cache if expr_to_let else self.expr_to_egg_cache try: - return self.expr_to_egg_cache[expr_decl] + return cache[expr_decl] except KeyError: pass res: bindings._Expr @@ -624,7 +912,7 @@ def _expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr: # noqa: PLR0912, res = bindings.Lit(span(), l) case CallDecl() | GetCostDecl(): egg_fn, typed_args = self.translate_call(expr_decl) - egg_args = [self.typed_expr_to_egg(a, False) for a in typed_args] + egg_args = [self.typed_expr_to_egg(a, expr_to_let) for a in typed_args] res = bindings.Call(span(), egg_fn, egg_args) case PyObjectDecl(value): res = bindings.Call( @@ -633,11 +921,14 @@ def _expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr: # noqa: PLR0912, [bindings.Lit(span(), bindings.String(standard_b64encode(value).decode("utf-8")))], ) case PartialCallDecl(call_decl): - egg_fn_call = self._expr_to_egg(call_decl) + egg_fn, typed_args = self.translate_call(call_decl) res = bindings.Call( span(), "unstable-fn", - [bindings.Lit(span(), bindings.String(egg_fn_call.name)), *egg_fn_call.args], + [ + bindings.Lit(span(), bindings.String(egg_fn)), + *[self.typed_expr_to_egg(arg, expr_to_let) for arg in typed_args], + ], ) case ValueDecl(): msg = "Cannot turn a Value into an expression" @@ -647,7 +938,7 @@ def _expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr: # noqa: PLR0912, raise ValueError(msg) case _: assert_never(expr_decl.expr) - self.expr_to_egg_cache[expr_decl] = res + cache[expr_decl] = res return res def translate_call(self, expr: CallDecl | GetCostDecl) -> tuple[str, list[TypedExprDecl]]: @@ -657,15 +948,15 @@ def translate_call(self, expr: CallDecl | GetCostDecl) -> tuple[str, list[TypedE match expr: case CallDecl(ref, args, _): egg_fn, reverse_args = self.callable_ref_to_egg(ref) - args_list = list(args) - if reverse_args: - args_list.reverse() - return egg_fn, args_list case GetCostDecl(ref, args): - cost_table = self.create_cost_table(ref) - return cost_table, list(args) + egg_fn = self.create_cost_table(ref) + _, reverse_args = self.callable_ref_to_egg(ref) case _: assert_never(expr) + args_list = list(args) + if reverse_args: + args_list.reverse() + return egg_fn, args_list def exprs_from_egg(self, termdag: bindings.TermDag, terms: list[int], tp: JustTypeRef) -> Iterable[TypedExprDecl]: """ @@ -680,33 +971,84 @@ def _get_possible_types(self, cls_ident: Ident) -> frozenset[JustTypeRef]: """ return frozenset(tp for tp in self.type_ref_to_egg_sort if tp.ident == cls_ident) - def _generate_callable_egg_name(self, ref: CallableRef) -> str: + 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, ...]: """ - Generates a valid egg function name for a callable reference. + Generates short and fully-qualified egg function name candidates for a callable reference. """ match ref: case FunctionRef(ident): - return str(ident) - + return _name_candidates(ident.name, str(ident), sanitize=True) case ConstantRef(ident): # Prefix to avoid name collisions with local vars - return f"%{ident}" + return _name_candidates(f"%{ident.name}", f"%{ident}", sanitize=True) case ( MethodRef(cls_ident, name) | ClassMethodRef(cls_ident, name) | ClassVariableRef(cls_ident, name) | PropertyRef(cls_ident, name) ): - return f"{cls_ident}.{name}" + return _name_candidates(f"{cls_ident.name}.{name}", f"{cls_ident}.{name}", sanitize=True) case InitRef(cls_ident): - return f"{cls_ident}.__init__" + return _name_candidates(f"{cls_ident.name}.__init__", f"{cls_ident}.__init__", sanitize=True) case UnnamedFunctionRef(): name = f"_lambda_{self.unnamed_function_counter}" self.unnamed_function_counter += 1 - return name + return (name,) case _: assert_never(ref) + def _allocate_type_egg_name(self, ref: JustTypeRef, decl: ClassDecl, arg_names: list[str]) -> str: + return self._allocate_name( + self._generate_type_egg_name_candidates(ref, decl, arg_names), self._backend_symbol_is_occupied + ) + + def _backend_symbol_is_occupied(self, name: str) -> bool: + """Check Egglog's shared namespace for sorts, tables, and primitives.""" + return ( + bool(self.egg_fn_to_callable_refs.get(name)) + or name in self.egg_sort_to_type_ref + or name in self.cost_table_names.values() + or name in BUILTIN_EGG_FN_NAMES + or name in BUILTIN_EGG_SORT_NAMES + ) + + def _generate_type_egg_name_candidates( + self, ref: JustTypeRef, decl: ClassDecl, arg_names: list[str] + ) -> tuple[str, ...]: + base_short = decl.egg_name or ref.ident.name + base_full = decl.egg_name or str(ref.ident) + if not ref.args: + return _name_candidates(base_short, base_full, sanitize=False) + args = ",".join(arg_names) + return _name_candidates(f"{base_short}[{args}]", f"{base_full}[{args}]", sanitize=False) + + def _allocate_synthetic_let_name(self) -> str: + while True: + name = f"$__expr_{self.expr_to_let_counter}" + self.expr_to_let_counter += 1 + if name not in { + egg_expr.name + for decl, egg_expr in self.expr_to_egg_cache.items() + if isinstance(decl, LetRefDecl) and isinstance(egg_expr, bindings.Var) + }: + return name + + @staticmethod + def _allocate_name(candidates: Iterable[str], is_taken: Callable[[str], bool]) -> str: + candidate_list = tuple(dict.fromkeys(candidates)) + for candidate in candidate_list: + if not is_taken(candidate): + return candidate + + fallback = candidate_list[-1] + index = 1 + while is_taken(f"{fallback}_{index}"): + index += 1 + return f"{fallback}_{index}" + def typed_expr_to_value(self, typed_expr: TypedExprDecl) -> bindings.Value: if isinstance(typed_expr.expr, ValueDecl): return typed_expr.expr.value @@ -817,6 +1159,11 @@ def value_to_expr(self, tp: JustTypeRef, value: bindings.Value) -> ExprDecl: # _names, _args = self.egraph.value_to_function(value) return_tp, *arg_types = tp.args return self._unstable_fn_value_to_expr(_names, _args, return_tp, arg_types) + case "Pair" | "Maybe": + termdag, term, _cost = call_with_current_trace( + self.egraph.extract_value, value, self.type_ref_to_egg(tp) + ) + return FromEggState(self, termdag).resolve_term(term, tp).expr case _: # If this is not a builtin type, or we don't know how to convert it, just return as value return ValueDecl(value) @@ -826,7 +1173,7 @@ def _unstable_fn_value_to_expr( ) -> PartialCallDecl: # 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]: + for callable_ref in self.egg_fn_to_callable_refs.get(name, ()): signature = self.__egg_decls__.get_callable_decl(callable_ref).signature if not isinstance(signature, FunctionSignature): continue @@ -854,35 +1201,37 @@ def _sanitize_egg_ident(input_string: str) -> str: return _EGGLOG_INVALID_IDENT.sub("_", input_string) +def _name_candidates(short: str, full: str, *, sanitize: bool) -> tuple[str, ...]: + if sanitize: + short = _sanitize_egg_ident(short) + full = _sanitize_egg_ident(full) + return short, full + + def _exprs_multiple_parents(typed_expr: TypedExprDecl) -> list[TypedExprDecl]: """ Returns all expressions that have multiple parents (a list but semantically just an ordered set). """ - to_traverse = {typed_expr} - traversed = set[TypedExprDecl]() - traversed_twice = list[TypedExprDecl]() - while to_traverse: - typed_expr = to_traverse.pop() - if typed_expr in traversed: - traversed_twice.append(typed_expr) - continue - traversed.add(typed_expr) - expr = typed_expr.expr - if isinstance(expr, CallDecl): - to_traverse.update(expr.args) - elif isinstance(expr, PartialCallDecl): - to_traverse.update(expr.call.args) - return traversed_twice - - -def _generate_type_egg_name(ref: JustTypeRef) -> str: - """ - Generates an egg sort name for this type reference by linearizing the type. - """ - name = ref.ident - if not ref.args: - return str(name) - return f"{name}[{','.join(map(_generate_type_egg_name, ref.args))}]" + parent_counts: dict[TypedExprDecl, int] = {} + traversal_order: list[TypedExprDecl] = [] + traversed: set[TypedExprDecl] = set() + + def visit(node: TypedExprDecl) -> None: + if node in traversed: + return + traversed.add(node) + match node.expr: + case CallDecl(args=args) | PartialCallDecl(CallDecl(args=args)): + for child in args: + parent_counts[child] = parent_counts.get(child, 0) + 1 + if child not in traversed: + traversal_order.append(child) + visit(child) + case _: + pass + + visit(typed_expr) + return [node for node in traversal_order if parent_counts[node] > 1] @dataclass @@ -893,8 +1242,10 @@ class FromEggState: state: EGraphState termdag: bindings.TermDag - # Cache of termdag ID to TypedExprDecl - cache: dict[int, TypedExprDecl] = field(default_factory=dict) + # Cache of termdag ID and expected type to TypedExprDecl. Polymorphic + # zero-argument terms like map-empty can appear once in a termdag but be + # decoded at multiple concrete types. + cache: dict[tuple[int, JustTypeRef], TypedExprDecl] = field(default_factory=dict) @property def decls(self) -> Declarations: @@ -904,6 +1255,11 @@ def from_expr(self, tp: JustTypeRef, term: bindings._Term) -> TypedExprDecl: """ Convert an egg term to a typed expr. """ + # Extracted builtin values can use canonical constructors that were not + # present in the original Python expression, such as BigInt.from_string + # inside an extracted BigRat. Seed the expected type's callable mapping + # before resolving the term. + self.state.type_ref_to_egg(tp) expr_decl: ExprDecl if isinstance(term, bindings.TermVar): expr_decl = LetRefDecl(term.name) @@ -911,7 +1267,21 @@ def from_expr(self, tp: JustTypeRef, term: bindings._Term) -> TypedExprDecl: value = term.value expr_decl = LitDecl(None if isinstance(value, bindings.Unit) else value.value) elif isinstance(term, bindings.TermApp): - if term.name == "py-object": + if term.name == "map-of" and tp.ident == Ident.builtin("Map"): + if len(term.args) % 2: + raise ValueError(f"Expected alternating key/value terms in map-of, got {len(term.args)} terms") + key_tp, value_tp = tp.args + expr_decl = CallDecl(ClassMethodRef(Ident.builtin("Map"), "empty"), (), (key_tp, value_tp)) + for index in range(0, len(term.args), 2): + expr_decl = CallDecl( + MethodRef(Ident.builtin("Map"), "insert"), + ( + TypedExprDecl(tp, expr_decl), + self.resolve_term(term.args[index], key_tp), + self.resolve_term(term.args[index + 1], value_tp), + ), + ) + elif term.name == "py-object": (str_term,) = term.args call = self.termdag.get(str_term) assert isinstance(call, bindings.TermLit) @@ -946,12 +1316,14 @@ def from_call(self, tp: JustTypeRef, term: bindings.TermApp) -> CallDecl: we have values for. """ # Find the first callable ref that matches the call - for callable_ref in self.state.egg_fn_to_callable_refs[term.name]: + possible_callable_refs = self.state.egg_fn_to_callable_refs.get(term.name, ()) + for callable_ref in possible_callable_refs: # If this is a classmethod, we might need the type params that were bound for this type # This could be multiple types if the classmethod is ambiguous, like map create. possible_types: Iterable[JustTypeRef | None] signature = self.decls.get_callable_decl(callable_ref).signature assert isinstance(signature, FunctionSignature) + term_args = term.args[::-1] if signature.reverse_args else term.args if isinstance(callable_ref, ClassMethodRef | InitRef | MethodRef): # Need OR in case we have class method whose class was never added as a sort, which would happen # if the class method didn't return that type and no other function did. In this case, we don't need @@ -971,7 +1343,7 @@ def from_call(self, tp: JustTypeRef, term: bindings.TermApp) -> CallDecl: signature.arg_types, signature.semantic_return_type, signature.var_arg_type, tp ) # Include this in try because of iterable - a_tp = list(zip(term.args, arg_types, strict=False)) + a_tp = list(zip(term_args, arg_types, strict=False)) except TypeConstraintError: continue args = tuple(self.resolve_term(a, tp) for a, tp in a_tp) @@ -980,12 +1352,13 @@ def from_call(self, tp: JustTypeRef, term: bindings.TermApp) -> CallDecl: bound_tp_params = () if signature.semantic_return_type.vars.issubset(signature.arg_vars) else bound_args return CallDecl(callable_ref, args, bound_tp_params) raise ValueError( - f"Could not find callable ref for call {term}. None of these refs matched the types: {self.state.egg_fn_to_callable_refs[term.name]}" + f"Could not find callable ref for call {term}. None of these refs matched the types: {possible_callable_refs}" ) def resolve_term(self, term_id: int, tp: JustTypeRef) -> TypedExprDecl: + key = (term_id, tp) try: - return self.cache[term_id] + return self.cache[key] except KeyError: - res = self.cache[term_id] = self.from_expr(tp, self.termdag.get(term_id)) + res = self.cache[key] = self.from_expr(tp, self.termdag.get(term_id)) return res diff --git a/python/egglog/exp/array_api.py b/python/egglog/exp/array_api.py index 42ea19a0..2244a6bd 100644 --- a/python/egglog/exp/array_api.py +++ b/python/egglog/exp/array_api.py @@ -2867,6 +2867,7 @@ def to_polynomial_ruleset( mss1 == mss.map(partial(multiset_flat_map, get_monomial)), mss != mss1, # skip if this is a no-op name="unwrap monomial", + eval_mode="naive", ).then( union(n1).with_(polynomial(mss1)), delete(polynomial(mss)), @@ -2878,6 +2879,7 @@ def to_polynomial_ruleset( mss1 == multiset_flat_map(UnstableFn(get_sole_polynomial), mss), mss != mss1, name="unwrap polynomial", + eval_mode="naive", ).then( union(n1).with_(polynomial(mss1)), delete(polynomial(mss)), diff --git a/python/egglog/exp/array_api_program_gen.py b/python/egglog/exp/array_api_program_gen.py index bdafffd3..ff6908ea 100644 --- a/python/egglog/exp/array_api_program_gen.py +++ b/python/egglog/exp/array_api_program_gen.py @@ -71,11 +71,11 @@ def _int_program(i64_: i64, i: Int, j: Int, s: String, b: Boolean, ti: Callable[ def program_if(b: BooleanLike, t: Callable[[], Program], f: Callable[[], Program]) -> Program: ... -@function(ruleset=array_api_program_gen_ruleset) +@function def tuple_int_foldl_program(xs: TupleIntLike, f: Callable[[Program, Int], Program], init: ProgramLike) -> Program: ... -@function(ruleset=array_api_program_gen_ruleset) +@function def tuple_int_program(x: TupleIntLike) -> Program: ... @@ -186,7 +186,7 @@ def _value_program(i: Int, b: Boolean, f: Float, x: NDArray, v1: Value, v2: Valu yield rewrite(value_program(v1.conj())).to(Program("np.conj(") + value_program(v1) + ")") -@function(ruleset=array_api_program_gen_ruleset) +@function def tuple_value_foldl_program( xs: TupleValueLike, f: Callable[[Program, Value], Program], init: ProgramLike ) -> Program: ... diff --git a/python/egglog/exp/param_eq/__init__.py b/python/egglog/exp/param_eq/__init__.py new file mode 100644 index 00000000..001d3182 --- /dev/null +++ b/python/egglog/exp/param_eq/__init__.py @@ -0,0 +1,48 @@ +"""Experimental parameter-reducing simplifier for symbolic-regression expressions.""" + +from __future__ import annotations + +from .cases import DEMO_CASES, DemoCase +from .domain import ( + ContainerMonomial, + ContainerPolynomial, + Num, + ParamCost, + binary_to_containers, + container_cost_model, + containers_to_binary, + exp, + log, + param_cost_model, + parse_expression, + polynomial, + render_num, + sqrt, +) +from .pipeline import ( + PaperPipelineReport, + run_paper_pipeline, + run_paper_pipeline_container, +) + +__all__ = [ + "DEMO_CASES", + "ContainerMonomial", + "ContainerPolynomial", + "DemoCase", + "Num", + "PaperPipelineReport", + "ParamCost", + "binary_to_containers", + "container_cost_model", + "containers_to_binary", + "exp", + "log", + "param_cost_model", + "parse_expression", + "polynomial", + "render_num", + "run_paper_pipeline", + "run_paper_pipeline_container", + "sqrt", +] diff --git a/python/egglog/exp/param_eq/__main__.py b/python/egglog/exp/param_eq/__main__.py new file mode 100644 index 00000000..4bbb3167 --- /dev/null +++ b/python/egglog/exp/param_eq/__main__.py @@ -0,0 +1,31 @@ +"""Simplify one expression with the experimental Param-Eq pipeline.""" + +from __future__ import annotations + +import argparse +import json +from collections.abc import Sequence +from dataclasses import asdict + +from .domain import binary_to_containers, parse_expression +from .pipeline import run_paper_pipeline, run_paper_pipeline_container + + +def main(argv: Sequence[str] | None = None) -> int: + """Run the single-expression JSON CLI.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--expr", required=True, help="Python-like symbolic expression") + parser.add_argument("--variant", choices=("binary", "container"), default="binary") + args = parser.parse_args(argv) + + report = ( + run_paper_pipeline(parse_expression(args.expr)) + if args.variant == "binary" + else run_paper_pipeline_container(binary_to_containers(parse_expression(args.expr))) + ) + print(json.dumps({"variant": args.variant, **asdict(report)}, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/python/egglog/exp/param_eq/cases.py b/python/egglog/exp/param_eq/cases.py new file mode 100644 index 00000000..f262ef96 --- /dev/null +++ b/python/egglog/exp/param_eq/cases.py @@ -0,0 +1,32 @@ +"""Small public/project-authored cases used by bounded CI.""" + +from __future__ import annotations + +from dataclasses import dataclass + +_DEMO_SAMPLE_POINTS = ((-1.25, 0.5), (0.25, 1.5), (2.0, -0.75)) + + +@dataclass(frozen=True) +class DemoCase: + """A redistributable Param-Eq demonstration expression.""" + + name: str + source: str + sample_points: tuple[tuple[float, float], ...] = _DEMO_SAMPLE_POINTS + + +DEMO_CASES = ( + DemoCase( + name="paper_eq4_instantiation", + source="(2.3 * (3.7*x0 + 5.1*x1)) / 7.9", + ), + DemoCase( + name="repeated_monomial", + source="2.0*x0*x0 + 3.0*x0*x0 + 5.0*x1", + ), + DemoCase( + name="function_composition", + source="log(exp(2.0*x0 + 3.0))", + ), +) diff --git a/python/egglog/exp/param_eq/domain.py b/python/egglog/exp/param_eq/domain.py new file mode 100644 index 00000000..6c8283ed --- /dev/null +++ b/python/egglog/exp/param_eq/domain.py @@ -0,0 +1,635 @@ +""" +Expression domain for the retained Param-Eq stress cases. + +This module defines the restricted symbolic language, parser and renderer, +binary/container conversions, and lexicographic extraction costs. Container +polynomials use nested `Map` values to canonicalize repeated terms and factors; +lowering and decoding preserve supported expression meaning, not exact tree +shape. The detailed research method and restart notes live in +`experiments/param_eq/NOTES.md`. +""" + +# mypy: disable-error-code="empty-body" + +from __future__ import annotations + +import ast +import math +from collections.abc import Callable +from dataclasses import dataclass +from fractions import Fraction +from typing import TypeAlias, TypeVar, cast + +from typing_extensions import TypeIs + +from egglog import * + +from ...runtime import RuntimeClass, RuntimeExpr # noqa: TID252 + +_T_EXPR = TypeVar("_T_EXPR", bound=BaseExpr) + + +def _is_expr_instance(x: BaseExpr, cls: type[_T_EXPR]) -> TypeIs[_T_EXPR]: + """Check an expression's concrete Egglog type, including generic arguments.""" + if not isinstance(cast("object", x), RuntimeExpr): + raise TypeError(f"Expected Expression, got {type(x).__name__}") + if not isinstance(cast("object", cls), RuntimeClass): + raise TypeError(f"Expected expression class, got {type(cls).__name__}") + return cast("RuntimeExpr", x).__egg_typed_expr__.tp == cast("RuntimeClass", cls).__egg_tp__.to_just() + + +class Num(Expr): + """ + Paper EqSat language subset. + + This is deliberately closer to `FixTree`'s `SRTreeF` than to the broader + experimental translations that were removed during cleanup. The paper + corpus only needs constants, + variables, arithmetic, and a small unary-function set. + """ + + def __init__(self, value: f64Like) -> None: ... + + __match_args__ = ("value",) + + @method(preserve=True) # type: ignore[prop-decorator] + @property + def value(self) -> f64: + match get_callable_args(self, Num): + case (value,): + return cast("f64", value) + raise ExprValueError(self, "Num") + + @classmethod + def var(cls, name: StringLike) -> Num: ... + + def __add__(self, other: NumLike) -> Num: ... + + def __sub__(self, other: NumLike) -> Num: ... + + def __mul__(self, other: NumLike) -> Num: ... + + def __truediv__(self, other: NumLike) -> Num: ... + + def __pow__(self, other: NumLike) -> Num: ... + + def __abs__(self) -> Num: ... + + def __radd__(self, other: NumLike) -> Num: ... + + def __rsub__(self, other: NumLike) -> Num: ... + + def __rmul__(self, other: NumLike) -> Num: ... + + def __rtruediv__(self, other: NumLike) -> Num: ... + + def __rpow__(self, other: NumLike) -> Num: ... + + @method(preserve=True) + def __neg__(self) -> Num: + return Num(-1.0) * self + + +@function +def exp(num: NumLike) -> Num: ... + + +@function +def log(num: NumLike) -> Num: ... + + +@function +def sqrt(num: NumLike) -> Num: ... + + +@function +def polynomial(p: ContainerPolynomialLike) -> Num: ... + + +ContainerMonomial: TypeAlias = Map[Num, BigRat] +ContainerPolynomial: TypeAlias = Map[ContainerMonomial, f64] + +ContainerMonomialLike: TypeAlias = MapLike[Num, BigRat, "NumLike", BigRatLike] +ContainerPolynomialLike: TypeAlias = MapLike[ContainerMonomial, f64, ContainerMonomialLike, f64Like] +NumLike: TypeAlias = Num | StringLike | f64Like | i64Like | ContainerMonomialLike | ContainerPolynomialLike + +converter(f64, Num, Num) +converter(i64, Num, lambda value: Num(f64.from_i64(value))) +converter(String, Num, Num.var) +converter(ContainerPolynomial, Num, polynomial) +converter(ContainerMonomial, Num, lambda mono: polynomial(ContainerPolynomial.empty().insert(mono, f64(1.0)))) + + +def parse_expression(source: str) -> Num: + """ + Parse a string of the expression syntax into a `Num` expression. + """ + return convert(_from_ast(ast.parse(_normalize_expression_source(source), mode="eval")), Num) + + +def render_num(num: Num) -> str: + """Render a `Num` back into a Python-like surface syntax for reports.""" + # parse and unparse to remove redundant parentheses and spacing. + return ast.unparse(ast.parse(_render_num(num), mode="eval")) + + +def binary_to_containers(expr: Num) -> Num: + """ + Convert a binary expression to its container form. + """ + return convert(_binary_to_containers(expr), Num) + + +def containers_to_binary(num: Num) -> Num: + """ + Convert a container expression back to its binary form. + + Should be inverse of `binary_to_containers` (modulo ordering) + """ + match get_callable_args(num, polynomial): + case (poly,): + return _decode_container_polynomial(cast("ContainerPolynomial", poly)) + fn = get_callable_fn(num) + if fn in (Num, Num.var): + return num + args = get_callable_args(num) + if fn is None or args is None: + raise ValueError(f"Cannot decode container expression: {num}") + constructor = cast("Callable[..., Num]", fn) + return constructor(*(containers_to_binary(cast("Num", arg)) for arg in args)) + + +def _finite_num_literal(value: float) -> Num: + """Construct a numeric literal while enforcing the surface language's finite-value invariant.""" + finite_value = float(value) + if not math.isfinite(finite_value): + msg = "Numeric literals must be finite" + raise ValueError(msg) + return Num(finite_value) + + +def _from_ast(node: ast.AST) -> Num: # noqa: C901, PLR0911, PLR0912 + """ + Parse a subset of Python expressions into the `Num` DSL. + + Keep things as floats for as long as possible, so that when we convert to containers we know which terms are constants without running them through the e-graph. + """ + if isinstance(node, ast.Expression): + return _from_ast(node.body) + if isinstance(node, ast.Constant): + if isinstance(node.value, float | int) and not isinstance(node.value, bool): + return _finite_num_literal(node.value) + msg = f"Unsupported constant: {node.value!r}" + raise ValueError(msg) + if isinstance(node, ast.Name): + return Num.var(node.id) + if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub): + if ( + isinstance(node.operand, ast.Constant) + and isinstance(node.operand.value, int | float) + and not isinstance(node.operand.value, bool) + ): + return _finite_num_literal(-node.operand.value) + return -_from_ast(node.operand) + if isinstance(node, ast.BinOp): + lhs = _from_ast(node.left) + rhs = _from_ast(node.right) + if isinstance(node.op, ast.Add): + return lhs + rhs + if isinstance(node.op, ast.Sub): + return lhs - rhs + if isinstance(node.op, ast.Mult): + return lhs * rhs + if isinstance(node.op, ast.Div): + return lhs / rhs + if isinstance(node.op, ast.Pow): + if get_callable_fn(rhs) != Num: + msg = "Power exponent must be a numeric literal" + raise ValueError(msg) + return lhs**rhs + msg = f"Unsupported binary operator: {ast.dump(node.op)}" + raise TypeError(msg) + if isinstance(node, ast.Call): + if not isinstance(node.func, ast.Name): + msg = f"Unsupported call target: {ast.dump(node.func)}" + raise TypeError(msg) + name = node.func.id + if node.keywords: + msg = f"Function call {name!r} does not accept keyword arguments" + raise ValueError(msg) + if len(node.args) != 1: + msg = f"Function call {name!r} expects exactly one argument" + raise ValueError(msg) + (arg,) = [_from_ast(arg) for arg in node.args] + if name == "exp": + return exp(arg) + if name == "log": + return log(arg) + if name == "sqrt": + return sqrt(arg) + if name == "abs": + return arg.__abs__() + if name == "plog": + return log(arg.__abs__()) + if name == "square": + return arg**2 + if name == "cube": + return arg**3 + msg = f"Unsupported function call: {name}" + raise ValueError(msg) + msg = f"Unsupported AST node: {ast.dump(node)}" + raise TypeError(msg) + + +def _normalize_expression_source(source: str) -> str: + normalized = source.strip() + replacements = { + "Log(": "log(", + "Exp(": "exp(", + "Sqrt(": "sqrt(", + "Abs(": "abs(", + "^": "**", + } + for old, new in replacements.items(): + normalized = normalized.replace(old, new) + return normalized + + +def _float_to_bigrat(value: float) -> BigRat: + """Preserve a Python float's exact binary value as a rational exponent.""" + numerator, denominator = value.as_integer_ratio() + return BigRat(numerator, denominator) + + +def _binary_to_containers( # noqa: C901, PLR0911, PLR0912 + expr: Num, +) -> Num | ContainerPolynomial | ContainerMonomial: + """ + Turn all instances of *, +, etc. into container expressions. + """ + if get_callable_fn(expr) in {Num.var, Num}: + return expr + match get_callable_args(expr, Num.__add__): + case (lhs, rhs): + return map_merge_with( + lambda a, b: a + b, + _to_container_poly(_binary_to_containers(lhs)), + _to_container_poly(_binary_to_containers(cast("Num", rhs))), + ) + match get_callable_args(expr, Num.__sub__): + case (lhs, rhs): + return map_merge_with( + lambda a, b: a + b, + _to_container_poly(_binary_to_containers(lhs)), + map_map_values(lambda _, v: -v, _to_container_poly(_binary_to_containers(cast("Num", rhs)))), + ) + match get_callable_args(expr, Num.__mul__): + case (lhs, rhs): + lhs_mapped = _binary_to_containers(lhs) + rhs_mapped = _binary_to_containers(cast("Num", rhs)) + lhs_is_polynomial = _is_expr_instance(lhs_mapped, ContainerPolynomial) + rhs_is_polynomial = _is_expr_instance(rhs_mapped, ContainerPolynomial) + if lhs_is_polynomial: + lhs_poly = cast("ContainerPolynomial", lhs_mapped) + match get_callable_args(rhs_mapped, Num): + case (f64(scalar),): + return map_map_values(lambda _mono, coef: coef * scalar, lhs_poly) + if not rhs_is_polynomial: + return _multiply_container_polynomial_by_monomial(lhs_poly, _to_container_mono(rhs_mapped)) + if rhs_is_polynomial: + rhs_poly = cast("ContainerPolynomial", rhs_mapped) + match get_callable_args(lhs_mapped, Num): + case (f64(scalar),): + return map_map_values(lambda _mono, coef: coef * scalar, rhs_poly) + if not lhs_is_polynomial: + return _multiply_container_polynomial_by_monomial(rhs_poly, _to_container_mono(lhs_mapped)) + return map_merge_with(lambda a, b: a + b, _to_container_mono(lhs_mapped), _to_container_mono(rhs_mapped)) + match get_callable_args(expr, Num.__truediv__): + case (lhs, rhs): + lhs_mapped = _binary_to_containers(lhs) + rhs_mapped = _binary_to_containers(cast("Num", rhs)) + lhs_is_polynomial = _is_expr_instance(lhs_mapped, ContainerPolynomial) + rhs_is_polynomial = _is_expr_instance(rhs_mapped, ContainerPolynomial) + if lhs_is_polynomial: + lhs_poly = cast("ContainerPolynomial", lhs_mapped) + match get_callable_args(rhs_mapped, Num): + case (f64(scalar),): + return map_map_values(lambda _mono, coef: coef / scalar, lhs_poly) + denom = map_map_values(lambda _term, exponent: -exponent, _to_container_mono(rhs_mapped)) + if lhs_is_polynomial and not rhs_is_polynomial: + return _multiply_container_polynomial_by_monomial(cast("ContainerPolynomial", lhs_mapped), denom) + # If the numerator is just one, then dont add this as a term to the polynomial + if _is_expr_instance(lhs_mapped, Num) and lhs_mapped == Num(1.0): + return denom + num = _to_container_mono(lhs_mapped) + return map_merge_with(lambda a, b: a + b, num, denom) + match get_callable_args(expr, Num.__pow__): + case (n, Num(f64(f))): + n_mapped = _to_num(_binary_to_containers(n)) + if f == 1: + return n_mapped + return ContainerMonomial.empty().insert(n_mapped, _float_to_bigrat(f)) + match get_callable_args(expr, exp): + case (inner,): + return exp(_to_num(_binary_to_containers(cast("Num", inner)))) + match get_callable_args(expr, log): + case (inner,): + return log(_to_num(_binary_to_containers(cast("Num", inner)))) + match get_callable_args(expr, sqrt): + case (inner,): + return ContainerMonomial.empty().insert(_to_num(_binary_to_containers(cast("Num", inner))), BigRat(1, 2)) + match get_callable_args(expr, Num.__abs__): + case (inner,): + return abs(_to_num(_binary_to_containers(inner))) + raise ValueError(f"Cannot decode to container: {expr}") + + +def _multiply_container_polynomial_by_monomial( + poly: ContainerPolynomial, factor: ContainerMonomial +) -> ContainerPolynomial: + """Distribute one monomial into a polynomial and combine coefficient collisions.""" + return map_fold_kv( + lambda result, mono, coef: map_merge_with( + lambda old_coef, new_coef: old_coef + new_coef, + result, + ContainerPolynomial.empty().insert( + map_merge_with(lambda left_exp, right_exp: left_exp + right_exp, mono, factor), + coef, + ), + ), + ContainerPolynomial.empty(), + poly, + ) + + +def _to_container_poly(v: Num | ContainerPolynomial | ContainerMonomial | f64) -> ContainerPolynomial: + if _is_expr_instance(v, ContainerPolynomial): + return v + return ContainerPolynomial.empty().insert(_to_container_mono(v), f64(1.0)) + + +def _to_container_mono(v: Num | ContainerPolynomial | ContainerMonomial | f64) -> ContainerMonomial: + if _is_expr_instance(v, ContainerMonomial): + return v + return ContainerMonomial.empty().insert(_to_num(v), BigRat(1, 1)) + + +def _to_num(v: ContainerMonomial | ContainerPolynomial | Num | f64) -> Num: + if isinstance(v, Num): + return v + if _is_expr_instance(v, ContainerPolynomial): + return polynomial(v) + if _is_expr_instance(v, ContainerMonomial): + return polynomial(ContainerPolynomial.empty().insert(v, f64(1.0))) + return Num(v) + + +def _render_float(value: float) -> str: + if not math.isfinite(value): + msg = "Cannot render a non-finite Param-Eq result" + raise ValueError(msg) + if value == 0.0: + return "0.0" + if value.is_integer(): + return f"{value:.1f}" + return repr(value) + + +def _render_num(num: Num) -> str: # noqa: C901, PLR0911, PLR0912 + match get_callable_args(num, polynomial): + case (poly,) if isinstance(poly, Map): + return _render_num(containers_to_binary(num)) + match get_callable_args(num, Num): + case (f64(f),): + res = _render_float(f) + if f < 0.0: + res = f"({res})" + return res + match get_callable_args(num, Num.var): + case (String(s),): + return s + match get_callable_args(num, Num.__add__): + case (lhs, rhs): + return f"({_render_num(lhs)} + {_render_num(cast('Num', rhs))})" + match get_callable_args(num, Num.__sub__): + case (lhs, rhs): + return f"({_render_num(lhs)} - {_render_num(cast('Num', rhs))})" + match get_callable_args(num, Num.__mul__): + case (lhs, rhs): + return f"({_render_num(lhs)} * {_render_num(cast('Num', rhs))})" + match get_callable_args(num, Num.__truediv__): + case (lhs, rhs): + return f"({_render_num(lhs)} / {_render_num(cast('Num', rhs))})" + match get_callable_args(num, Num.__pow__): + case (lhs, rhs): + return f"({_render_num(lhs)} ** {_render_num(cast('Num', rhs))})" + match get_callable_args(num, exp): + case (inner,): + return f"exp({_render_num(cast('Num', inner))})" + match get_callable_args(num, log): + case (inner,): + return f"log({_render_num(cast('Num', inner))})" + match get_callable_args(num, sqrt): + case (inner,): + return f"sqrt({_render_num(cast('Num', inner))})" + match get_callable_args(num, Num.__abs__): + case (inner,): + return f"abs({_render_num(inner)})" + msg = f"Unsupported Num node for rendering: {num!r}" + raise TypeError(msg) + + +def _product(factors: list[Num], initial: float = 1.0) -> Num: + if not factors: + return Num(initial) + total = factors.pop(0) if initial == 1.0 else Num(initial) + for next_factor in factors: + total *= next_factor + return total + + +def _decode_container_mono_term(mono: dict[Num, BigRat], coef: float) -> Num: + if not mono: + return Num(coef) + numerator_factors: list[Num] = [] + denominator_factors: list[Num] = [] + for term, exp in mono.items(): + exp_value = exp.value + term_decoded = containers_to_binary(term) + abs_exp_value = abs(exp_value) + factor = ( + sqrt(term_decoded) + if abs_exp_value == Fraction(1, 2) + else term_decoded + if abs_exp_value == 1 + else term_decoded ** float(abs_exp_value) + ) + (denominator_factors if exp_value < 0 else numerator_factors).append(factor) + numerator = _product(numerator_factors, coef) + if not denominator_factors: + return numerator + return numerator / _product(denominator_factors) + + +def _decode_container_polynomial(poly: ContainerPolynomial) -> Num: + """ + Decode a polynomial into binary ops. Negative coefficients are turned into subtraction and 1.0 coefficients are elided. + """ + poly_items = [(mono.value, float(coef)) for (mono, coef) in poly.value.items()] + if not poly_items: + return Num(0.0) + mono, coef = poly_items.pop(0) + total = _decode_container_mono_term(mono, coef) + for mono, coef in poly_items: + if coef < 0.0: + total -= _decode_container_mono_term(mono, abs(coef)) + else: + total += _decode_container_mono_term(mono, coef) + return total + + +@dataclass(frozen=True, order=True) +class ParamCost: + """ + Custom cost type that prioritizes minimizing number of floats (which correspond to fitted parameters), and then on ties + minimizes the sum of ops and ints (which correspond to the complexity or cost of the operation). + """ + + # count of floats + floats: int = 0 + # count of +, *, /, **, exp, log, sqrt, abs operations plus any floats that can be parsed as ints (like 1, 2, etc) + ops_and_ints: int = 0 + + @property + def node_count(self) -> int: + return self.floats + self.ops_and_ints + + def __add__(self, other: ParamCost) -> ParamCost: + return ParamCost( + floats=self.floats + other.floats, + ops_and_ints=self.ops_and_ints + other.ops_and_ints, + ) + + def __str__(self) -> str: + return f"ParamCost({self.floats}, {self.ops_and_ints})" + + def __repr__(self) -> str: + return str(self) + + +def _float_cost(f: float) -> ParamCost: + """ + Dont count integer floats as floats, since they aren't parameters + """ + if f.is_integer(): + return ParamCost(ops_and_ints=1) + return ParamCost(floats=1) + + +def param_cost_model(egraph: EGraph, expr: BaseExpr, children_costs: list[ParamCost]) -> ParamCost: + if isinstance(expr, f64): + return _float_cost(float(expr)) + if isinstance(expr, String): + return ParamCost(ops_and_ints=1) + fn = get_callable_fn(expr) + if fn in (Num, Num.var): + cost = 0 + elif fn in (Num.__add__, Num.__sub__, Num.__mul__, Num.__truediv__, Num.__pow__, exp, log, sqrt, Num.__abs__): + cost = 1 + else: + raise ValueError(f"Unsupported expression in cost model: {expr}") + return sum(children_costs, start=ParamCost(ops_and_ints=cost)) + + +def _decoded_monomial_cost(mono: ContainerMonomial, children_costs: list[ParamCost]) -> ParamCost: + """ + Like _decode_container_mono_term. Assumes that if we have an empty numerator we include the 1.0 + """ + items = list(mono.value.items()) + if len(children_costs) != len(items) * 2: + msg = f"Expected {len(items) * 2} monomial child costs, got {len(children_costs)}" + raise ValueError(msg) + + numerator_factor_costs: list[ParamCost] = [] + denominator_factor_costs: list[ParamCost] = [] + for i, (_, exp) in enumerate(items): + exp_value = exp.value + abs_exp_value = abs(exp_value) + term_cost = children_costs[i * 2] + # factor cost + factor_cost = ( + # sqrt is one op plus the inside + term_cost + ParamCost(ops_and_ints=1) + if abs_exp_value == Fraction(1, 2) + else term_cost + if abs_exp_value == 1 + else term_cost + ParamCost(ops_and_ints=1) + _float_cost(float(abs_exp_value)) + ) + (denominator_factor_costs if exp_value < 0 else numerator_factor_costs).append(factor_cost) + + numerator_cost = ( + sum(numerator_factor_costs, ParamCost(ops_and_ints=len(numerator_factor_costs) - 1)) + if numerator_factor_costs + else ParamCost(ops_and_ints=1) + ) + if not denominator_factor_costs: + return numerator_cost + denominator_cost = sum(denominator_factor_costs, ParamCost(ops_and_ints=len(denominator_factor_costs))) + return numerator_cost + denominator_cost + + +def _decoded_polynomial_term_cost(mono: dict[Num, BigRat], coef: float, mono_cost: ParamCost) -> ParamCost: + """ + Gives the cost of one monomial and its coefficient based on the cost of the monomial. + + mirrors _decode_container_mono_term + """ + coef_cost = _float_cost(coef) + if not mono: + return coef_cost + has_empty_numerator = all(exp.value < 0 for exp in mono.values()) + if coef == 1.0: + return mono_cost + if has_empty_numerator: + # mono_cost includes the synthetic numerator `1`; decoding replaces + # that with the coefficient, so charge the coefficient instead. + return ParamCost( + floats=mono_cost.floats + coef_cost.floats, + ops_and_ints=mono_cost.ops_and_ints + coef_cost.ops_and_ints - 1, + ) + # if we are multiplying them, add their costs and the cost of the mul + return coef_cost + mono_cost + ParamCost(ops_and_ints=1) + + +def _decoded_polynomial_cost(poly: ContainerPolynomial, children_costs: list[ParamCost]) -> ParamCost: + """ + Should correspond to getting the cost from the return value of _decode_container_polynomial + """ + items = list(poly.value.items()) + if len(children_costs) != len(items) * 2: + msg = f"Expected {len(items) * 2} polynomial child costs, got {len(children_costs)}" + raise ValueError(msg) + if not items: + # empty is zero + return ParamCost(ops_and_ints=1) + + mono, coef = items[0] + total = _decoded_polynomial_term_cost(mono.value, float(coef), children_costs[0]) + for i, (mono, coef) in enumerate(items[1:], start=1): + term_cost = _decoded_polynomial_term_cost(mono.value, abs(float(coef)), children_costs[i * 2]) + # the cost is the cost of the monomial plus the cost of an add/sub + total += term_cost + ParamCost(ops_and_ints=1) + return total + + +# Container specific cost model that should give the same cost as the default cost model on the decoded expression. +def container_cost_model(egraph: EGraph, expr: BaseExpr, children_costs: list[ParamCost]) -> ParamCost: + if _is_expr_instance(expr, ContainerPolynomial): + return _decoded_polynomial_cost(expr, children_costs) + if _is_expr_instance(expr, ContainerMonomial): + return _decoded_monomial_cost(expr, children_costs) + if get_callable_fn(expr) == polynomial: + return children_costs[0] + if isinstance(expr, BigRat): + return ParamCost(ops_and_ints=1) + return param_cost_model(egraph, expr, children_costs) diff --git a/python/egglog/exp/param_eq/pipeline.py b/python/egglog/exp/param_eq/pipeline.py new file mode 100644 index 00000000..a24b5b38 --- /dev/null +++ b/python/egglog/exp/param_eq/pipeline.py @@ -0,0 +1,532 @@ +# mypy: disable-error-code="empty-body" + +"""Retained paper-era `param_eq` pipeline plus the experimental map variant.""" + +from __future__ import annotations + +import time +from collections.abc import Callable, Iterable +from dataclasses import dataclass +from functools import partial +from typing import Literal, TypeVar + +from egglog import * + +from .domain import * + +MAX_PASSES = 2 +HASKELL_INNER_ITERATION_LIMIT = 30 +BACKOFF_MATCH_LIMIT = 1000 +BACKOFF_BAN_LENGTH = 30 + +T = TypeVar("T", bound=BaseExpr) +V = TypeVar("V", bound=BaseExpr) + + +@function(builtin=True, egg_fn="f64-is-finite") +def _f64_is_finite(value: f64) -> Unit: ... + + +# Store discovered constants in a global map so semi-naive analysis can join +# them with polynomial terms. If two singleton updates collide, keep the first +# representative; tolerant float canonicalization is not part of this paused +# research slice. +CONSTS = constant( + "CONSTS", + Map[Num, f64], + merge=partial(map_merge_with, lambda old_value, _new_value: old_value), +) + +# Map a monomial of the form `{polynomial(P): 1}` to one representative `P`. +# This acts as an index from concrete monomial keys back to nested polynomial +# bodies. Directly matching `polynomial(P) == n` and then constructing +# `{n: 1}` is semantically equivalent, but it can make matching enumerate many +# unrelated polynomial e-classes before proving the singleton monomial exists. +# +# Concrete example: +# a*polynomial(P) + R -> a*P + R +# should start from monomial keys already present in the outer polynomial, not +# from every `polynomial(P) == n` relation in the e-graph. +POLYNOMIAL_MONOMIALS = constant( + "POLYNOMIAL_MONOMIALS", + Map[ContainerMonomial, ContainerPolynomial], + merge=partial(map_merge_with, lambda old_value, new_value: old_value), +) + + +def if_defined(cond: Unit, then: T, otherwise: T) -> T: + return catch(lambda: cond).match(lambda _: then, otherwise) + + +def try_match(expr: T, on_some: Callable[[T], V], default: V) -> V: + return catch(lambda: expr).match(on_some, default) + + +@ruleset +def shared_analysis_rules(a: f64) -> Iterable[RewriteOrRule]: + yield rewrite(exp(Num(a)), subsume=True).to(Num(a.exp()), _f64_is_finite(a.exp())) + yield rewrite(log(Num(a)), subsume=True).to(Num(a.log()), a > 0.0, _f64_is_finite(a.log())) + yield rule(log(Num(a)), a <= 0.0).then(panic("Log of non-positive number")) + yield rewrite(abs(Num(a)), subsume=True).to(Num(abs(a)), _f64_is_finite(abs(a))) + + +@ruleset +def binary_analysis_rules(x: Num, a: f64, b: f64) -> Iterable[RewriteOrRule]: + yield rewrite(Num(a) / Num(b), subsume=True).to(Num(a / b), b != f64(0.0), _f64_is_finite(a / b)) + yield rule(x / Num(0.0)).then(panic("Division by zero")) + yield rewrite(Num(a) - Num(b), subsume=True).to(Num(a - b), _f64_is_finite(a - b)) + yield rewrite(Num(a) * Num(b), subsume=True).to(Num(a * b), _f64_is_finite(a * b)) + yield rewrite(Num(a) + Num(b), subsume=True).to(Num(a + b), _f64_is_finite(a + b)) + + yield rewrite(Num(a) ** Num(b), subsume=True).to(Num(a**b), _f64_is_finite(a**b)) + yield rewrite(sqrt(Num(a)), subsume=True).to(Num(a.sqrt()), a >= 0.0, _f64_is_finite(a.sqrt())) + yield rule(sqrt(Num(a)), a < 0.0).then(panic("Sqrt of negative number")) + + # cancellations + yield rewrite(x - x, subsume=True).to(Num(0.0)) + yield rewrite(x / x, subsume=True).to(Num(1.0), x != Num(0.0)) + + # multiplicative of inverse + yield rewrite(x * (1 / x), subsume=True).to(Num(1.0), x != Num(0.0)) + + yield rewrite(0 * x, subsume=True).to(Num(0.0)) + yield rewrite(0 / x, subsume=True).to(Num(0.0), x != Num(0.0)) + + +@ruleset +def container_analysis_rules( + n: Num, + a: f64, + poly: ContainerPolynomial, + poly1: ContainerPolynomial, + consts: Map[Num, f64], +) -> Iterable[RewriteOrRule]: + yield rule(n == Num(a)).then(set_(CONSTS).to(Map[Num, f64].empty().insert(n, a))) + yield rule(polynomial(poly) == n).then( + set_(POLYNOMIAL_MONOMIALS).to( + Map[ContainerMonomial, ContainerPolynomial] + .empty() + .insert(ContainerMonomial.empty().insert(n, BigRat(1, 1)), poly) + ) + ) + # Constant fold polynomials so that in each monomial, + # all constants terms are pulled into a co-efficient + # and all empty terms are combined. Also drops terms with zero exponents. + # like: {{}: 3.14, {x: 2}: 2.71}} + + yield rewrite(polynomial(poly), subsume=True).to( + polynomial(poly1), + # pull in this so it gets joined in semi-naive + consts == CONSTS, + poly1 + == map_fold_kv( + lambda res_poly, mono, coef: ( + # split monomial into non constants and constants (which are combined into the coefficient): + map_fold_kv( + lambda res_mono_and_coef, term, exp: if_defined( + exp != BigRat(0, 1), + # if the exponent is not zero, process it + try_match( + consts[term], + # if it is a constant, multiply it into the coefficient and drop it from the monomial: + lambda v: if_defined( + exp != BigRat(-1, 1), + res_mono_and_coef.map_right(lambda prev_coef: prev_coef * (v ** exp.to_f64())), + if_defined( + v != f64(0.0), + res_mono_and_coef.map_right(lambda prev_coef: prev_coef / v), + res_mono_and_coef.map_left(lambda mono: mono.insert(term, exp)), + ), + ), + # if it is not a constant, keep it in the monomial + res_mono_and_coef.map_left(lambda mono: mono.insert(term, exp)), + ), + # if the exponent is zero, the term is just 1 and can be dropped from the monomial, so keep the monomial as is + res_mono_and_coef, + ), + Pair(ContainerMonomial.empty(), coef), + mono, + ).match( + lambda mono, coef: res_poly.insert(mono, coef + catch(lambda: res_poly[mono]).unwrap_or(f64(0.0))) + ) + ), + ContainerPolynomial.empty(), + poly, + ), + poly != poly1, + ) + + # Turn polynomials that are actually just constant factors into constants, so they can be used in more rewrites. + yield rewrite(polynomial(poly), subsume=True).to( + Num(poly[ContainerMonomial.empty()]), + poly.length() == i64(1), + # The only key is an empty monomial, so the polynomial is just a constant term: + ContainerMonomial.empty() == poly.pick_key(), + ) + + # remove monomials with zero coefficients + yield rewrite(polynomial(poly), subsume=True).to( + polynomial(poly1), + poly1 == map_filter_kv(lambda _key, value: value != f64(0.0), poly), + poly != poly1, + ) + + +@ruleset +def binary_basic_rules(x: Num, y: Num, z: Num, af: f64, bf: f64, cf: f64, df: f64) -> Iterable[RewriteOrRule]: + a = Num(af) + b = Num(bf) + c = Num(cf) + d = Num(df) + + # commutativity + yield rewrite(x + y).to(y + x) + yield rewrite(x * y).to(y * x) + + # associativity + yield rewrite(x + (y + z)).to((x + y) + z) # no-op + yield rewrite(x * (y * z)).to((x * y) * z) # no-op + yield rewrite(x * (y / z)).to((x * y) / z) # no-op + yield rewrite((x * y) / z).to(x * (y / z)) # no-op + yield rewrite((a * x) * (b * y)).to((a * b) * (x * y)) # no-op + yield rewrite(a * x + b).to(a * (x + b / a)) # no-op + yield rewrite(a * x - b).to(a * (x - b / a)) # no-op + yield rewrite(b - (a * x)).to(a * ((b / a) - x)) # no-op + yield rewrite(a * x + b * y).to( + a * (x + (b / a) * y) + ) # factoring out one constant from one term, and dividing the others who have constant terms to compensate + yield rewrite(a * x - b * y).to(a * (x - (b / a) * y)) # same as above + yield rewrite(a * x + b / y).to(a * (x + (b / a) / y)) # same as above + yield rewrite(a * x - b / y).to(a * (x - (b / a) / y)) # same as above + + yield rewrite(a / (b * x)).to((a / b) / x) # no-op + yield rewrite(x / (b * y)).to((1 / b) * x / y) # no-op + yield rewrite(x / a + b).to((x + b * a) / a) # same as above + yield rewrite(x / a - b).to((x - b * a) / a) # same as above + yield rewrite(b - x / a).to(((b * a) - x) / a) # same as above + yield rewrite(x / a + b * y).to((x + (b * a) * y) / a) # same as above + yield rewrite(x / a - b * y).to((x - (b * a) * y) / a) # same as above + yield rewrite((b + a * x) / (c + d * y)).to((a / d) * (b / a + x) / (c / d + y)) + yield rewrite((b + x) / (c + d * y)).to((1 / d) * (b + x) / (c / d + y)) + + # identities + yield rewrite(0 + x).to(x) + yield rewrite(x - 0).to(x) + yield rewrite(1 * x).to(x) + + # distributive and factorization + yield rewrite((x * y) + (x * z)).to(x * (y + z)) + yield rewrite(x - (y + z)).to((x - y) - z) + yield rewrite(x - (y - z)).to((x - y) + z) + yield rewrite(-(x + y)).to(-x - y) + yield rewrite(x - a).to(x + -a) + yield rewrite(x - (a * y)).to(x + -a * y) + yield rewrite((1 / x) * (1 / y)).to(1 / (x * y)) + + # negate + yield rewrite(x - -y).to(x + y) + yield rewrite(x + -y).to(x - y) + yield rewrite(0 - x).to(-x) + + +@ruleset +def container_basic_rules( + poly: ContainerPolynomial, + poly1: ContainerPolynomial, + poly2: ContainerPolynomial, + nonconst_poly: ContainerPolynomial, + coef: f64, + polynomial_monomials: Map[ContainerMonomial, ContainerPolynomial], + counts: MultiSet[Num], + n: Num, + poly_pair: Pair[ContainerPolynomial, ContainerPolynomial], + exp: BigRat, + mono: ContainerMonomial, +) -> Iterable[RewriteOrRule]: + # Factor one representative non-unit coefficient from a small polynomial: + # + # a*x + b -> a*(x + b/a) + yield rewrite(polynomial(poly)).to( + polynomial( + ContainerPolynomial.empty().insert(ContainerMonomial.empty().insert(polynomial(poly1), BigRat(1, 1)), coef) + ), + poly.length() > i64(1), + poly.length() <= i64(4), + nonconst_poly == map_filter_kv(lambda key, _value: key != ContainerMonomial.empty(), poly), + poly2 == map_filter_kv(lambda _key, value: value != f64(1.0), nonconst_poly), + coef == poly2[poly2.pick_key()], + poly2.length() == nonconst_poly.length(), + poly1 == map_map_values(lambda _key, value: value / coef, poly), + ) + + # Greedy multivariate Horner factorization for rational exponents. Choose + # the term present in the most monomials, then factor out its minimum + # exponent: + # + # x*y + x*z -> x*(y + z) + yield rewrite(polynomial(poly)).to( + polynomial( + poly_pair.right.insert( + ContainerMonomial.empty().insert(n, exp).insert(polynomial(poly_pair.left), BigRat(1, 1)), + f64(1.0), + ) + ), + counts + == map_fold_kv( + lambda counts, mono, _coef: map_fold_kv( + lambda updated_counts, term, _exp: updated_counts.insert(term), + counts, + mono, + ), + MultiSet[Num](), + poly, + ), + n == counts.pick_max(), + counts.count(n) > i64(1), + exp + == map_fold_kv( + lambda min_exp, mono, _coef: try_match(mono[n] < min_exp, lambda _: mono[n], min_exp), + BigRat(2**63 - 1, 1), + poly, + ), + poly_pair + == map_fold_kv( + lambda divided_and_remainder, mono, coef: try_match( + mono[n], + lambda current_exp: divided_and_remainder.map_left( + lambda divided: divided.insert(mono.insert(n, current_exp - exp), coef) + ), + divided_and_remainder.map_right(lambda remainder: remainder.insert(mono, coef)), + ), + Pair(ContainerPolynomial.empty(), ContainerPolynomial.empty()), + poly, + ), + ) + + # Flatten an exact nested polynomial term inside a larger polynomial: + # + # a*polynomial(P) + R -> a*P + R + # + # The whole monomial must be the nested polynomial at exponent one, which + # avoids distributing arbitrary products. + yield rewrite(polynomial(poly)).to( + polynomial( + map_merge_with( + lambda left, right: left + right, + poly.remove(mono), + map_map_values(lambda _nested_mono, nested_coef: nested_coef * poly[mono], poly1), + ) + ), + polynomial_monomials == POLYNOMIAL_MONOMIALS, + poly.length() > i64(1), + mono + == map_fold_kv( + lambda selected, candidate_mono, _candidate_coef: try_match( + polynomial_monomials[candidate_mono], + lambda _nested_poly: candidate_mono, + selected, + ), + ContainerMonomial.empty(), + poly, + ), + mono.length() > i64(0), + poly[mono] != f64(0.0), + poly1 == polynomial_monomials[mono], + poly1.length() > i64(1), + ) + + +@ruleset +def shared_fun_rules(x: Num) -> Iterable[RewriteOrRule]: + yield rewrite(log(exp(x))).to(x) + yield rewrite(log(abs(exp(x)))).to(x) + + +@ruleset +def binary_fun_rules(x: Num, y: Num, af: f64) -> Iterable[RewriteOrRule]: + a = Num(af) + + yield rewrite(log(a * y)).to(log(a) + log(y), af > 0.0, y != Num(0.0)) + yield rewrite(log(y * a)).to(log(y) + log(a), af > 0.0, y != Num(0.0)) + yield rewrite(log(a / y)).to(log(a) - log(y), af > 0.0, y != Num(0.0)) + yield rewrite(log(y / a)).to(log(y) - log(a), af > 0.0, y != Num(0.0)) + + yield rewrite(log(a**y)).to(y * log(a), af > 0.0) + yield rewrite(log(sqrt(x))).to(0.5 * log(x)) + yield rewrite(x**0.5).to(sqrt(x)) + + +@ruleset +def container_fun_rules(poly: ContainerPolynomial, m: ContainerMonomial, term: Num) -> Iterable[RewriteOrRule]: + # Preserve the one-factor identity needed by the public log(2*x0) case. + # Expanding multiple factors or an even power would be unsound when + # negative factors combine into a positive product. + yield rewrite(log(polynomial(poly))).to( + polynomial( + map_fold_kv( + lambda res_poly, term, exp: map_merge_with( + lambda old_coef, new_coef: old_coef + new_coef, + res_poly, + ContainerPolynomial.empty().insert( + ContainerMonomial.empty().insert(log(term), BigRat(1, 1)), + exp.to_f64(), + ), + ), + ContainerPolynomial.empty().insert(ContainerMonomial.empty(), poly[m].log()), + m, + ) + ), + poly.length() == i64(1), + m == poly.pick_key(), + poly[m] > f64(0.0), + m.length() == i64(1), + term == m.pick_key(), + m[term] == BigRat(1, 1), + ) + + +@dataclass(frozen=True) +class PaperPipelineReport: + """Bounded result; `saturated` describes inner schedules, not an outer fixed point.""" + + status: Literal["saturated", "iteration_limit"] + passes: int + total_sec: float + total_size: int + before_nodes: int + before_params: int + extracted: str + extracted_nodes: int + extracted_params: int + + +# Reporting and pipeline loop + + +binary_analysis_ruleset = shared_analysis_rules | binary_analysis_rules +container_analysis_ruleset = shared_analysis_rules | container_analysis_rules +binary_analysis_schedule = binary_analysis_ruleset.saturate() +containers_analysis_schedule = container_analysis_ruleset.saturate() +shared_rewrite_ruleset = shared_fun_rules +binary_rewrite_ruleset = shared_rewrite_ruleset | binary_basic_rules | binary_fun_rules +container_rewrite_ruleset = shared_rewrite_ruleset | container_basic_rules | container_fun_rules + + +def _graph_size(egraph: EGraph) -> int: + return sum(size for _, size in egraph.all_function_sizes()) + + +binary_schedule = run( + binary_rewrite_ruleset, + scheduler=back_off( + match_limit=BACKOFF_MATCH_LIMIT, + ban_length=BACKOFF_BAN_LENGTH, + ).persistent(), +) +container_schedule = run( + container_rewrite_ruleset, + scheduler=back_off( + # This lower container budget preserves corpus parameter parity while + # avoiding synthetic appended-expression slowdowns from over-searching + # equivalent container factorizations. + match_limit=5, + ban_length=10, + ).persistent(), +) + + +def _run_single_pass( + egraph: EGraph, + num: Num, + cost_model: CostModel[ParamCost], + analysis_schedule: Schedule, + schedule: Schedule, +) -> tuple[Num, ParamCost, int, bool]: + """ + Run one `rewriteTree`-like pass and return the populated e-graph. + + This mirrors Haskell at the control-flow level while using the ordinary + persistent backoff scheduler available upstream: + - up to 30 inner rewrite rounds + - one saturated analysis round after each rewrite round + - stop when both schedules report no changes or deferred work + """ + n = egraph.let("n", num) + current_size = _graph_size(egraph) + saturated = False + for _ in range(HASKELL_INNER_ITERATION_LIMIT): + analysis_report = egraph.run(analysis_schedule) + rewrite_report = egraph.run(schedule) + current_size = _graph_size(egraph) + if analysis_report.can_stop and rewrite_report.can_stop: + saturated = True + break + extracted, cost = egraph.extract(n, include_cost=True, cost_model=cost_model) + return extracted, cost, current_size, saturated + + +def run_paper_pipeline( + initial: Num, + decode: Callable[[Num], Num] = lambda x: x, + cost_model: CostModel[ParamCost] = param_cost_model, + schedule: Schedule = binary_schedule, + analysis_schedule: Schedule = binary_analysis_schedule, +) -> PaperPipelineReport: + current, before_cost = EGraph(save_egglog_string=False).extract(initial, include_cost=True, cost_model=cost_model) + # get schedule decls so that it's pre-cached + schedule.__egg_decls__ + analysis_schedule.__egg_decls__ + start = time.perf_counter() + # Add constants to the egraph so that they can be used in rules without needing to be registered each pass + egraph = EGraph(Num(0.0), save_egglog_string=False) + # pre-run rulesets so that we don't have to register them each pipeline pass + egraph.run(analysis_schedule) + last_cost = before_cost + max_size = 0 + passes = 0 + status: Literal["saturated", "iteration_limit"] = "saturated" + for pass_index in range(1, MAX_PASSES + 1): + with egraph: + extracted, last_cost, total_size, saturated = _run_single_pass( + egraph, + current, + cost_model=cost_model, + schedule=schedule, + analysis_schedule=analysis_schedule, + ) + max_size = max(max_size, total_size) + passes = pass_index + unchanged = extracted == current + current = extracted + if not saturated: + status = "iteration_limit" + break + if unchanged: + break + return PaperPipelineReport( + status=status, + passes=passes, + total_sec=time.perf_counter() - start, + total_size=max_size, + before_nodes=before_cost.node_count, + before_params=before_cost.floats, + extracted_nodes=last_cost.node_count, + extracted_params=last_cost.floats, + extracted=render_num(decode(current)), + ) + + +def run_paper_pipeline_container(initial: Num) -> PaperPipelineReport: + try: + return run_paper_pipeline( + initial, + decode=containers_to_binary, + cost_model=container_cost_model, + schedule=container_schedule, + analysis_schedule=containers_analysis_schedule, + ) + except ValueError as error: + if "non-finite" not in str(error): + raise + msg = "The Param-Eq container pipeline requires every coefficient normalization to remain finite" + raise ValueError(msg) from error diff --git a/python/egglog/ipython_magic.py b/python/egglog/ipython_magic.py index 2f2101e7..3c92b889 100644 --- a/python/egglog/ipython_magic.py +++ b/python/egglog/ipython_magic.py @@ -38,5 +38,5 @@ def egglog(line, cell, local_ns): if "output" in line: print("\n".join(res)) if "graph" in line: - return graphviz.Source(e.to_graphviz_string()) + return graphviz.Source(e.serialize([]).to_dot()) return None diff --git a/python/egglog/pretty.py b/python/egglog/pretty.py index 008fb1bd..a0755320 100644 --- a/python/egglog/pretty.py +++ b/python/egglog/pretty.py @@ -7,7 +7,7 @@ import ast from collections import Counter, defaultdict from dataclasses import dataclass, field -from typing import TYPE_CHECKING, TypeAlias, assert_never +from typing import TYPE_CHECKING, TypeAlias, assert_never, cast import black import cloudpickle @@ -17,6 +17,9 @@ if TYPE_CHECKING: from collections.abc import Mapping + from .builtins import BigRat, Map, Maybe, Pair + from .egraph import BaseExpr + __all__ = [ "BINARY_METHODS", @@ -86,6 +89,7 @@ | ScheduleDecl | BackOffDecl | EGraphDecl + | TypedExprDecl ) @@ -121,7 +125,7 @@ def pretty_decl( def pretty_callable_ref( decls: Declarations, ref: CallableRef, - first_arg: ExprDecl | None = None, + first_arg: TypedExprDecl | None = None, bound_tp_params: tuple[JustTypeRef, ...] | None = None, include_all_args: bool = False, ) -> str: @@ -131,11 +135,28 @@ def pretty_callable_ref( To be used in the visualization. """ + if ref == FunctionRef(Ident.builtin("!=")): + return "!=" + # Pass in three dummy args, which are the max used for any operation that # is not a generic function call - args: list[ExprDecl] = [UnboundVarDecl(ARG_STR)] * 3 - if first_arg: - args.insert(0, first_arg) + signature = decls.get_callable_decl(ref).signature + + def concrete_type(tp: TypeOrVarRef) -> JustTypeRef: + match tp: + case TypeVarRef(): + return JustTypeRef(Ident.builtin("Unit")) + case TypeRefWithVars(ident, type_args): + return JustTypeRef(ident, tuple(concrete_type(arg) for arg in type_args)) + + if isinstance(signature, FunctionSignature): + args = [TypedExprDecl(concrete_type(tp), UnboundVarDecl(ARG_STR)) for tp in signature.arg_types] + else: + args = [TypedExprDecl(JustTypeRef(Ident.builtin("Unit")), UnboundVarDecl(ARG_STR)) for _ in range(3)] + while len(args) < 3: + args.append(TypedExprDecl(JustTypeRef(Ident.builtin("Unit")), UnboundVarDecl(ARG_STR))) + if first_arg is not None: + args[0] = first_arg context = PrettyContext(decls, defaultdict(lambda: 0)) res = context._call_inner(ref, args, bound_tp_params=bound_tp_params, parens=False) # Either returns a function or a function with args. If args are provided, they would just be called, @@ -144,10 +165,9 @@ def pretty_callable_ref( # If we want to include all args as ARG_STR, then we need to figure out how many to use # used for set_cost so that `cost(E(...))` will show up as a call if include_all_args: - signature = decls.get_callable_decl(ref).signature assert isinstance(signature, FunctionSignature) - correct_args: list[ExprDecl] = [UnboundVarDecl(ARG_STR)] * len(signature.arg_types) - return f"{res[0]}({', '.join(context(a, parens=False, unwrap_lit=True) for a in correct_args)})" + correct_args = args[: len(signature.arg_types)] + return f"{res[0]}({', '.join(context(a, parens=False, unwrap_lit=res[2]) for a in correct_args)})" return res[0] return res @@ -187,7 +207,7 @@ def __call__(self, decl: AllDecls, toplevel: bool = False) -> None: # noqa: C90 self(rhs) for cond in conditions: self(cond) - case RuleDecl(head, body, _): + case RuleDecl(head, body, _, _): for action in head: self(action) for fact in body: @@ -196,7 +216,7 @@ def __call__(self, decl: AllDecls, toplevel: bool = False) -> None: # noqa: C90 self(lhs) self(rhs) case LetDecl(_, d) | ExprActionDecl(d) | ExprFactDecl(d): - self(d.expr) + self(d) case ChangeDecl(_, d, _) | SaturateDecl(d) | RepeatDecl(d, _) | ActionCommandDecl(d): self(d) case PanicDecl(_) | UnboundVarDecl(_) | LetRefDecl(_) | LitDecl(_) | PyObjectDecl(_): @@ -208,11 +228,11 @@ def __call__(self, decl: AllDecls, toplevel: bool = False) -> None: # noqa: C90 self(de) case CallDecl(ref, exprs, _) | GetCostDecl(ref, exprs): match ref: - case FunctionRef(UnnamedFunctionRef(_, res)): - self(res.expr) + case UnnamedFunctionRef(_, res): + self(res) case _: for e in exprs: - self(e.expr) + self(e) case RunDecl(_, until, scheduler): if until: for f in until: @@ -233,13 +253,13 @@ def __call__(self, decl: AllDecls, toplevel: bool = False) -> None: # noqa: C90 case LetSchedulerDecl(scheduler, schedule): self(scheduler) self(schedule) - case GetCostDecl(ref, args): - self(CallDecl(ref, args)) case DummyDecl(): pass case EGraphDecl() as eg: for a in eg.to_actions: self(a) + case TypedExprDecl(): + self(decl.expr) case _: assert_never(decl) @@ -312,7 +332,7 @@ def uncached( # noqa: C901, PLR0911, PLR0912 case CallDecl(_, _, _): return self._call(decl, parens) case PartialCallDecl(CallDecl(ref, typed_args, _)): - return self._pretty_partial(ref, [a.expr for a in typed_args], parens), "fn" + return self._pretty_partial(ref, list(typed_args), parens), "fn" case PyObjectDecl(pickled): value = cloudpickle.loads(pickled) value_str = repr(value) @@ -326,12 +346,14 @@ def uncached( # noqa: C901, PLR0911, PLR0912 args = ", ".join(map(self, (rhs, *conditions))) fn = "rewrite" if isinstance(decl, RewriteDecl) else "birewrite" return f"{fn}({self(lhs)}).to({args})", "rewrite" - case RuleDecl(head, body, name): - l = ", ".join(map(self, body)) + case RuleDecl(head, body, name, eval_mode): + args = list(map(self, body)) if name: - l += f", name={name}" + args.append(f"name={name!r}") + if eval_mode != "seminaive": + args.append(f"eval_mode={eval_mode!r}") r = ", ".join(map(self, head)) - return f"rule({l}).then({r})", "rule" + return f"rule({', '.join(args)}).then({r})", "rule" case SetDecl(_, lhs, rhs): return f"set_({self(lhs)}).to({self(rhs)})", "action" case UnionDecl(_, lhs, rhs): @@ -372,24 +394,30 @@ def uncached( # noqa: C901, PLR0911, PLR0912 case LetSchedulerDecl(scheduler, schedule): return f"{self(scheduler, parens=True)}.scope({self(schedule, parens=True)})", "schedule" case RunDecl(ruleset_ident, until, scheduler): - ruleset = self.decls._rulesets[ruleset_ident] - ruleset_str = self(ruleset, ruleset_ident=ruleset_ident) + if ruleset_ident.name == "" and ruleset_ident not in self.decls._rulesets: + ruleset_str = None + else: + ruleset = self.decls._rulesets[ruleset_ident] + ruleset_str = self(ruleset, ruleset_ident=ruleset_ident) if not until and not scheduler: - return ruleset_str, "schedule" - arg_lst = list(map(self, until or [])) + return "run()" if ruleset_str is None else ruleset_str, "schedule" + arg_lst = ["None" if ruleset_str is None else ruleset_str, *map(self, until or [])] if scheduler: arg_lst.append(f"scheduler={self(scheduler)}") - return f"run({ruleset_str}, {', '.join(arg_lst)})", "schedule" + return f"run({', '.join(arg_lst)})", "schedule" case DefaultRewriteDecl(): - msg = "default rewrites should not be pretty printed" + msg = "implicit rewrites should not be pretty printed" raise TypeError(msg) - case BackOffDecl(_, match_limit, ban_length): + case BackOffDecl(_, match_limit, ban_length, persistent): list_args = [] if match_limit is not None: list_args.append(f"match_limit={match_limit}") if ban_length is not None: list_args.append(f"ban_length={ban_length}") - return f"back_off({', '.join(list_args)})", "scheduler" + rendered = f"back_off({', '.join(list_args)})" + if persistent: + rendered += ".persistent()" + return rendered, "scheduler" case ValueDecl(value): return str(value), "value" case DummyDecl(): @@ -398,6 +426,57 @@ def uncached( # noqa: C901, PLR0911, PLR0912 return f"get_cost({self(CallDecl(ref, args))})", "get_cost" case EGraphDecl() as eg: return f"EGraph({', '.join(map(self, eg.to_actions))}).freeze()", "egraph" + case TypedExprDecl(tp, expr): + 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(): + key_str = self(cast("RuntimeExpr", key).__egg_typed_expr__) + value_str = self(cast("RuntimeExpr", value).__egg_typed_expr__) + map_str += f".insert({key_str}, {value_str})" + return map_str, "Map" + if tp.ident == Ident.builtin("BigRat"): + runtime_expr = RuntimeExpr.__from_values__(self.decls, decl) + try: + as_fraction = cast("BigRat", runtime_expr).value + except ExprValueError: + return self(expr, unwrap_lit=unwrap_lit, ruleset_ident=ruleset_ident, parens=parens), "expr" + return f"BigRat({as_fraction.numerator}, {as_fraction.denominator})", "BigRat" + if tp.ident == Ident.builtin("Pair"): + runtime_expr = RuntimeExpr.__from_values__(self.decls, decl) + try: + left, right = cast("Pair[BaseExpr, BaseExpr]", runtime_expr).value + except ExprValueError: + return self(expr, unwrap_lit=unwrap_lit, ruleset_ident=ruleset_ident, parens=parens), "expr" + left_str = self(cast("RuntimeExpr", left).__egg_typed_expr__) + right_str = self(cast("RuntimeExpr", right).__egg_typed_expr__) + return f"{tp}({left_str}, {right_str})", "Pair" + if tp.ident == Ident.builtin("Maybe"): + runtime_expr = RuntimeExpr.__from_values__(self.decls, decl) + try: + value = cast("Maybe[BaseExpr]", runtime_expr).value + except ExprValueError: + return self(expr, unwrap_lit=unwrap_lit, ruleset_ident=ruleset_ident, parens=parens), "expr" + if value is None: + return f"{tp}.none()", "Maybe" + return f"{tp}.some({self(cast('RuntimeExpr', value).__egg_typed_expr__)})", "Maybe" + return ( + self(expr, unwrap_lit=unwrap_lit, ruleset_ident=ruleset_ident, parens=parens), + tp.ident.name, + ) assert_never(decl) def _call( @@ -410,7 +489,7 @@ def _call( :param parens: If true, wrap the call in parens if it is a binary method call. """ - args = [a.expr for a in decl.args] + args = list(decl.args) ref = decl.callable # Special case != if decl.callable == FunctionRef(Ident.builtin("!=")): @@ -425,7 +504,7 @@ def _call( for arg, default in zip( reversed(args), reversed(signature.arg_defaults), strict=not signature.var_arg_type ): - if arg != default: + if arg.expr != default: break n_defaults += 1 if n_defaults: @@ -444,12 +523,12 @@ def _call( has_multiple_parents = self.parents[first_arg] > 1 self.names[decl] = expr_name = self._name_expr(tp_name, expr_str, copy_identifier=has_multiple_parents) # Set the first arg to be the name of the mutated arg and return the name - args[0] = LetRefDecl(expr_name) + args[0] = TypedExprDecl(args[0].tp, LetRefDecl(expr_name)) else: expr_name = None res = self._call_inner(ref, args, decl.bound_tp_params, parens) expr = ( - (f"{res[0]}({', '.join(self(a, parens=False, unwrap_lit=True) for a in res[1])})") + (f"{res[0]}({', '.join(self(a, parens=False, unwrap_lit=res[2]) for a in res[1])})") if isinstance(res, tuple) else res ) @@ -462,44 +541,49 @@ def _call( def _call_inner( # noqa: C901, PLR0911, PLR0912 self, ref: CallableRef, - args: list[ExprDecl], + args: list[TypedExprDecl], bound_tp_params: tuple[JustTypeRef, ...] | None, parens: bool, - ) -> tuple[str, list[ExprDecl]] | str: + ) -> tuple[str, list[TypedExprDecl], bool] | str: """ Pretty print the call, returning either the full function call or a tuple of the function and the args. """ match ref: case FunctionRef(Ident(name)): - return name, args + return name, args, True case ClassMethodRef(class_name, method_name): tp_ref = JustTypeRef(class_name, bound_tp_params or ()) - return f"{tp_ref}.{method_name}", args - case MethodRef(_class_name, method_name): + return f"{tp_ref}.{method_name}", args, True + case MethodRef(class_name, method_name): slf, *args = args non_str_slf = slf slf = self(slf, parens=True) + cls_has_generic_params = self.decls.get_class_decl(class_name).type_vars + # only unwrap literals if the class doesn't have generic params, otherwise we might lose type information that is needed for the method call + unwrap_lit = not cls_has_generic_params match method_name: case _ if method_name in UNARY_METHODS: expr = f"{UNARY_METHODS[method_name]}{slf}" return f"({expr})" if parens else expr case _ if method_name in BINARY_METHODS: - expr = f"{slf} {BINARY_METHODS[method_name]} {self(args[0], parens=True, unwrap_lit=True)}" + expr = ( + f"{slf} {BINARY_METHODS[method_name]} {self(args[0], parens=True, unwrap_lit=unwrap_lit)}" + ) return f"({expr})" if parens else expr case "__getitem__": - return f"{slf}[{self(args[0], unwrap_lit=True)}]" + return f"{slf}[{self(args[0], unwrap_lit=unwrap_lit)}]" case "__call__": - return slf, args + return slf, args, unwrap_lit case "__delitem__": - return f"del {slf}[{self(args[0], unwrap_lit=True)}]" + return f"del {slf}[{self(args[0], unwrap_lit=unwrap_lit)}]" case "__setitem__": - return f"{slf}[{self(args[0], unwrap_lit=True)}] = {self(args[1], unwrap_lit=True)}" + return f"{slf}[{self(args[0], unwrap_lit=unwrap_lit)}] = {self(args[1], unwrap_lit=unwrap_lit)}" case _ if method_name in NAMED_UNARY_METHODS: - return NAMED_UNARY_METHODS[method_name], [non_str_slf, *args] - case "__getattr__" if isinstance(args[0], LitDecl) and isinstance(args[0].value, str): - return f"{slf}.{args[0].value}" + return NAMED_UNARY_METHODS[method_name], [non_str_slf, *args], unwrap_lit + case "__getattr__" if isinstance(args[0].expr, LitDecl) and isinstance(args[0].expr.value, str): + return f"{slf}.{args[0].expr.value}" case _: - return f"{slf}.{method_name}", args + return f"{slf}.{method_name}", args, unwrap_lit case ConstantRef(Ident(name)): return name case ClassVariableRef(Ident(class_name), variable_name): @@ -508,10 +592,10 @@ def _call_inner( # noqa: C901, PLR0911, PLR0912 return f"{self(args[0], parens=True)}.{property_name}" case InitRef(class_name): tp_ref = JustTypeRef(class_name, bound_tp_params or ()) - return str(tp_ref), args + return str(tp_ref), args, True case UnnamedFunctionRef(): expr = self._pretty_function_body(ref, []) - return f"({expr})", args + return f"({expr})", args, True assert_never(ref) def _generate_name(self, typ: str) -> str: @@ -532,7 +616,7 @@ def _name_expr(self, tp_name: str, expr_str: str, copy_identifier: bool) -> str: self.statements.append(f"{name} = {expr_str}") return name - def _pretty_partial(self, ref: CallableRef, args: list[ExprDecl], parens: bool) -> str: + def _pretty_partial(self, ref: CallableRef, args: list[TypedExprDecl], parens: bool) -> str: """ Returns a partial function call as a string. """ @@ -566,19 +650,19 @@ def _pretty_partial(self, ref: CallableRef, args: list[ExprDecl], parens: bool) ) return f"partial({', '.join(arg_strs)})" - def _pretty_function_body(self, fn: UnnamedFunctionRef, args: list[ExprDecl]) -> str: + def _pretty_function_body(self, fn: UnnamedFunctionRef, args: list[TypedExprDecl]) -> str: """ Pretty print the body of a function, partially applying some arguments. """ var_args = fn.args - replacements = {var_arg: TypedExprDecl(var_arg.tp, arg) for var_arg, arg in zip(var_args, args, strict=False)} + replacements = dict(zip(var_args, args, strict=False)) var_args = var_args[len(args) :] res = replace_typed_expr(fn.res, replacements) arg_names = fn.args[len(args) :] prefix = "lambda" if arg_names: - prefix += f" {', '.join(self(a.expr) for a in arg_names)}" - return f"{prefix}: {self(res.expr)}" + prefix += f" {', '.join(self(a) for a in arg_names)}" + return f"{prefix}: {self(res)}" def is_valid_python_expr(s: str) -> bool: diff --git a/python/egglog/run_report.py b/python/egglog/run_report.py index b1772b18..fb9e79ff 100644 --- a/python/egglog/run_report.py +++ b/python/egglog/run_report.py @@ -86,6 +86,7 @@ class RunReport: _decls: Declarations = field(repr=False) iterations: list[IterationReport] = field(default_factory=list) updated: bool = False + can_stop: bool = False search_and_apply_time_per_rule: dict[RewriteOrRuleDecl, timedelta] = field(default_factory=dict) num_matches_per_rule: dict[RewriteOrRuleDecl, int] = field(default_factory=dict) search_and_apply_time_per_ruleset: dict[str, timedelta] = field(default_factory=dict) @@ -98,6 +99,7 @@ def __repr__(self) -> str: return ( f"RunReport(iterations={self.iterations}, " f"updated={self.updated}, " + f"can_stop={self.can_stop}, " f"search_and_apply_time_per_rule={time_per_rule}, " f"num_matches_per_rule={matches_per_rule}, " f"search_and_apply_time_per_ruleset={self.search_and_apply_time_per_ruleset}, " @@ -130,6 +132,7 @@ def _from_bindings(cls, report: bindings.RunReport, state: EGraphState) -> RunRe _decls=decls, iterations=[IterationReport._from_bindings(it, rule_map, decls) for it in report.iterations], updated=report.updated, + can_stop=report.can_stop, search_and_apply_time_per_rule=search_and_apply_time_per_rule, num_matches_per_rule=num_matches_per_rule, search_and_apply_time_per_ruleset=report.search_and_apply_time_per_ruleset, diff --git a/python/egglog/runtime.py b/python/egglog/runtime.py index 2ec5a453..70c0e02b 100644 --- a/python/egglog/runtime.py +++ b/python/egglog/runtime.py @@ -197,8 +197,10 @@ class BaseClassFactoryMeta(type): """ def __instancecheck__(cls, instance: object) -> bool: - assert isinstance(cls, RuntimeClass) - return isinstance(instance, RuntimeExpr) and cls.__egg_tp__.ident == instance.__egg_typed_expr__.tp.ident + runtime_cls = cast("RuntimeClass", cls) + return ( + isinstance(instance, RuntimeExpr) and runtime_cls.__egg_tp__.ident == instance.__egg_typed_expr__.tp.ident + ) class ClassFactory(type): @@ -549,6 +551,7 @@ def __call__( # noqa: C901,PLR0912 bound.apply_defaults() assert not bound.kwargs args = bound.args + decls.update(*(arg for arg in args if isinstance(arg, RuntimeExpr))) tcs = TypeConstraintSolver() if isinstance(self.__egg_bound__, JustTypeRef) and self.__egg_bound__.args: @@ -605,7 +608,7 @@ def __str__(self) -> str: first_arg, bound_tp_params = None, None match self.__egg_bound__: case RuntimeExpr(_): - first_arg = self.__egg_bound__.__egg_typed_expr__.expr + first_arg = self.__egg_bound__.__egg_typed_expr__ case JustTypeRef(_, args): bound_tp_params = args return pretty_callable_ref(self.__egg_decls__, self.__egg_ref__, first_arg, bound_tp_params) @@ -697,7 +700,7 @@ def __str__(self) -> str: return self.__egg_pretty__(None) def __egg_pretty__(self, wrapping_fn: str | None) -> str: - return pretty_decl(self.__egg_decls__, self.__egg_typed_expr__.expr, wrapping_fn=wrapping_fn) + return pretty_decl(self.__egg_decls__, self.__egg_typed_expr__, wrapping_fn=wrapping_fn) def _ipython_display_(self) -> None: from IPython.display import Code, display # noqa: PLC0415 diff --git a/python/egglog/type_constraint_solver.py b/python/egglog/type_constraint_solver.py index 1f2efc8c..532fe439 100644 --- a/python/egglog/type_constraint_solver.py +++ b/python/egglog/type_constraint_solver.py @@ -114,10 +114,7 @@ def substitute_typevars_try_function( # unnamed-function rewrites created while inferring types are discarded after the probe. probe_decls = decls().copy() dummy_args = [ - RuntimeExpr.__from_values__( - probe_decls, - TypedExprDecl(self.substitute_typevars(arg_tp), DummyDecl()), - ) + RuntimeExpr.__from_values__(probe_decls, TypedExprDecl(self.substitute_typevars(arg_tp), DummyDecl())) for arg_tp in tp.args[1:] ] try: diff --git a/python/tests/__snapshots__/test_array_api/test_jit[lda][expr].py b/python/tests/__snapshots__/test_array_api/test_jit[lda][expr].py index cda4a169..b880f65a 100644 --- a/python/tests/__snapshots__/test_array_api/test_jit[lda][expr].py +++ b/python/tests/__snapshots__/test_array_api/test_jit[lda][expr].py @@ -37,10 +37,7 @@ _TupleNDArray_1 = svd_( sqrt( asarray( - NDArray(RecursiveValue(Value.from_float(Float.rational(BigRat(BigInt.from_string("1"), BigInt.from_string("147")))))), - OptionalDType.some(DType.float64), - OptionalBool.none, - OptionalDevice.some(_NDArray_1.device), + NDArray(RecursiveValue(Value.from_float(Float.rational(BigRat(1, 147))))), OptionalDType.some(DType.float64), OptionalBool.none, OptionalDevice.some(_NDArray_1.device) ) ) * (_NDArray_8 / _NDArray_11), @@ -54,11 +51,7 @@ ).T / _TupleNDArray_1[Int(1)][IndexKey.slice(_Slice_1)] _TupleNDArray_2 = svd_( ( - sqrt( - NDArray(RecursiveValue(Value.from_int(Int(150)))) - * _NDArray_3 - * NDArray(RecursiveValue(Value.from_float(Float.rational(BigRat(BigInt.from_string("1"), BigInt.from_string("2")))))) - ) + sqrt(NDArray(RecursiveValue(Value.from_int(Int(150)))) * _NDArray_3 * NDArray(RecursiveValue(Value.from_float(Float.rational(BigRat(1, 2)))))) * (_NDArray_4 - _NDArray_3 @ _NDArray_4).T ).T @ _NDArray_12, diff --git a/python/tests/__snapshots__/test_polynomials/test_factor_multisets[code].py b/python/tests/__snapshots__/test_polynomials/test_factor_multisets[code].py new file mode 100644 index 00000000..20470786 --- /dev/null +++ b/python/tests/__snapshots__/test_polynomials/test_factor_multisets[code].py @@ -0,0 +1,16 @@ +_Value_1 = Value.var("bp3") * Value.var("q9") + Value.var("bp2") * Value.var("q6") + Value.var("bp1") * Value.var("q3") + Value.var("bp4") * Value.var("q12") +_Value_2 = Value.var("bpp2") * Value.var("q5") + Value.var("bpp1") * Value.var("q2") + Value.var("bpp3") * Value.var("q8") + Value.var("bpp4") * Value.var("q11") +_Value_3 = Value.var("bp3") * Value.var("q8") + Value.var("bp2") * Value.var("q5") + Value.var("bp1") * Value.var("q2") + Value.var("bp4") * Value.var("q11") +_Value_4 = Value.var("bpp2") * Value.var("q6") + Value.var("bpp1") * Value.var("q3") + Value.var("bpp3") * Value.var("q9") + Value.var("bpp4") * Value.var("q12") +_Value_5 = Value.var("q1") * Value.var("bp1") + Value.var("q4") * Value.var("bp2") + Value.var("q7") * Value.var("bp3") + Value.var("q10") * Value.var("bp4") +_Value_6 = Value.var("q4") * Value.var("bpp2") + Value.var("q1") * Value.var("bpp1") + Value.var("q7") * Value.var("bpp3") + Value.var("q10") * Value.var("bpp4") +NDArray( + RecursiveValue( + ( + (Value.from_int(Int(-1)) * (_Value_1 * _Value_2) + _Value_3 * _Value_4) ** Value.from_int(Int(2)) + + (Value.from_int(Int(-1)) * (_Value_5 * _Value_4) + _Value_1 * _Value_6) ** Value.from_int(Int(2)) + + (Value.from_int(Int(-1)) * (_Value_3 * _Value_6) + _Value_5 * _Value_2) ** Value.from_int(Int(2)) + ) + / (_Value_5 ** Value.from_int(Int(2)) + _Value_3 ** Value.from_int(Int(2)) + _Value_1 ** Value.from_int(Int(2))) ** Value.from_int(Int(3)) + ) +) \ No newline at end of file diff --git a/python/tests/param_eq/__init__.py b/python/tests/param_eq/__init__.py new file mode 100644 index 00000000..5bfd0745 --- /dev/null +++ b/python/tests/param_eq/__init__.py @@ -0,0 +1 @@ +"""Tests for the experimental Param-Eq package.""" diff --git a/python/tests/param_eq/conftest.py b/python/tests/param_eq/conftest.py new file mode 100644 index 00000000..0fa2a0e5 --- /dev/null +++ b/python/tests/param_eq/conftest.py @@ -0,0 +1,9 @@ +"""Make the non-installed research harness importable by its focused tests.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[3] +sys.path.insert(0, str(REPO_ROOT)) diff --git a/python/tests/param_eq/evaluation.py b/python/tests/param_eq/evaluation.py new file mode 100644 index 00000000..d6ca7e38 --- /dev/null +++ b/python/tests/param_eq/evaluation.py @@ -0,0 +1,53 @@ +"""Independent numeric evaluator for Param-Eq semantic tests.""" + +from __future__ import annotations + +import ast +import math +import operator +from collections.abc import Callable + +_BINARY: dict[type[ast.operator], Callable[[float, float], float]] = { + ast.Add: operator.add, + ast.Sub: operator.sub, + ast.Mult: operator.mul, + ast.Div: operator.truediv, + ast.Pow: operator.pow, +} +_UNARY: dict[type[ast.unaryop], Callable[[float], float]] = { + ast.UAdd: operator.pos, + ast.USub: operator.neg, +} +_FUNCTIONS: dict[str, Callable[[float], float]] = { + "abs": abs, + "exp": math.exp, + "log": math.log, + "sqrt": math.sqrt, +} + + +def evaluate(source: str, *, x0: float, x1: float) -> float: + """Evaluate the supported surface language without using Egglog.""" + + def visit(node: ast.AST) -> float: + if isinstance(node, ast.Expression): + return visit(node.body) + if isinstance(node, ast.Constant) and isinstance(node.value, int | float) and not isinstance(node.value, bool): + return float(node.value) + if isinstance(node, ast.Name) and node.id in {"x0", "x1"}: + return x0 if node.id == "x0" else x1 + if isinstance(node, ast.BinOp) and type(node.op) in _BINARY: + return _BINARY[type(node.op)](visit(node.left), visit(node.right)) + if isinstance(node, ast.UnaryOp) and type(node.op) in _UNARY: + return _UNARY[type(node.op)](visit(node.operand)) + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id in _FUNCTIONS + and len(node.args) == 1 + and not node.keywords + ): + return _FUNCTIONS[node.func.id](visit(node.args[0])) + raise ValueError(f"Unsupported numeric-test syntax: {ast.dump(node)}") + + return visit(ast.parse(source, mode="eval")) diff --git a/python/tests/param_eq/test_domain.py b/python/tests/param_eq/test_domain.py new file mode 100644 index 00000000..00e3eb7f --- /dev/null +++ b/python/tests/param_eq/test_domain.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +import math +from fractions import Fraction +from typing import cast + +import pytest + +from egglog import EGraph, get_callable_args, get_callable_fn +from egglog.exp.param_eq import ( + DEMO_CASES, + ContainerPolynomial, + DemoCase, + Num, + ParamCost, + binary_to_containers, + container_cost_model, + containers_to_binary, + param_cost_model, + parse_expression, + polynomial, + render_num, +) + +from .evaluation import evaluate + + +@pytest.mark.parametrize("source", [case.source for case in DEMO_CASES]) +def test_parse_render_parse_round_trip(source: str) -> None: + parsed = parse_expression(source) + rendered = render_num(parsed) + assert parse_expression(rendered) == parsed + + +@pytest.mark.parametrize("case", DEMO_CASES, ids=lambda case: case.name) +def test_container_round_trip_preserves_values(case: DemoCase) -> None: + container = EGraph().extract(binary_to_containers(parse_expression(case.source))) + decoded = render_num(containers_to_binary(container)) + for x0, x1 in case.sample_points: + assert math.isclose( + evaluate(case.source, x0=x0, x1=x1), + evaluate(decoded, x0=x0, x1=x1), + rel_tol=1e-9, + abs_tol=1e-9, + ) + + +@pytest.mark.parametrize("source", [*(case.source for case in DEMO_CASES), "3.5 * x0 ** 0.5 / x1 + 2.0"]) +def test_container_cost_matches_decoded_binary_cost(source: str) -> None: + container = binary_to_containers(parse_expression(source)) + extracted_container, container_cost = EGraph().extract( + container, include_cost=True, cost_model=container_cost_model + ) + decoded = containers_to_binary(extracted_container) + _, decoded_cost = EGraph().extract(decoded, include_cost=True, cost_model=param_cost_model) + + assert container_cost == decoded_cost + + +def test_lowering_distributes_scalar_products_into_a_polynomial() -> None: + lowered = binary_to_containers(parse_expression(DEMO_CASES[0].source)) + poly_args = get_callable_args(lowered, polynomial) + assert poly_args is not None + (poly_expr,) = poly_args + poly = EGraph().extract(cast("ContainerPolynomial", poly_expr)) + + assert len(poly.value) == 2 + assert sorted(coef.value for coef in poly.value.values()) == pytest.approx([2.3 / 7.9, 2.3 / 7.9]) + assert all(all(get_callable_fn(term) != polynomial for term in monomial.value) for monomial in poly.value) + + +def test_parser_rejects_unsupported_calls() -> None: + with pytest.raises(ValueError, match="Unsupported function call"): + parse_expression("sin(x0)") + + +def test_parser_rejects_nonliteral_power() -> None: + with pytest.raises(ValueError, match="Power exponent must be a numeric literal"): + parse_expression("x0 ** x1") + + +@pytest.mark.parametrize("source", ["log()", "log(x0, x1)"]) +def test_parser_rejects_wrong_function_arity(source: str) -> None: + with pytest.raises(ValueError, match="expects exactly one argument"): + parse_expression(source) + + +def test_parser_rejects_function_keywords() -> None: + with pytest.raises(ValueError, match="does not accept keyword arguments"): + parse_expression("log(num=x0)") + + +def test_parser_rejects_boolean_literals() -> None: + with pytest.raises(ValueError, match="Unsupported constant"): + parse_expression("True") + + +@pytest.mark.parametrize("source", ["1e309", "-1e309"]) +def test_parser_rejects_nonfinite_numeric_literals(source: str) -> None: + with pytest.raises(ValueError, match="Numeric literals must be finite"): + parse_expression(source) + + +def test_renderer_rejects_nonfinite_results() -> None: + with pytest.raises(ValueError, match="Cannot render a non-finite"): + render_num(Num(float("inf"))) + + +def test_float_exponents_keep_the_exact_python_ratio() -> None: + converted = binary_to_containers(parse_expression("x0 ** 0.1")) + poly_args = get_callable_args(converted, polynomial) + assert poly_args is not None + (poly_expr,) = poly_args + poly = cast("ContainerPolynomial", poly_expr) + (mono,) = poly.value + (exponent,) = mono.value.values() + assert EGraph().extract(exponent).value == Fraction(*(0.1).as_integer_ratio()) + + +def test_param_cost_is_lexicographic() -> None: + assert ParamCost(floats=1, ops_and_ints=100) < ParamCost(floats=2, ops_and_ints=0) + assert ParamCost(1, 2) + ParamCost(3, 4) == ParamCost(4, 6) diff --git a/python/tests/param_eq/test_pipeline.py b/python/tests/param_eq/test_pipeline.py new file mode 100644 index 00000000..2e254af1 --- /dev/null +++ b/python/tests/param_eq/test_pipeline.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +import json +import math + +import pytest + +from egglog import EGraph, back_off, eq, rewrite, ruleset, run, var +from egglog.exp.param_eq import ( + DEMO_CASES, + DemoCase, + Num, + binary_to_containers, + parse_expression, + pipeline, + run_paper_pipeline, + run_paper_pipeline_container, +) +from egglog.exp.param_eq.__main__ import main +from egglog.exp.param_eq.pipeline import container_schedule, containers_analysis_schedule + +from .evaluation import evaluate + + +@pytest.mark.param_eq_smoke +@pytest.mark.parametrize("variant", ["binary", "container"]) +@pytest.mark.parametrize("case", DEMO_CASES, ids=lambda case: case.name) +def test_public_end_to_end_cases(case: DemoCase, variant: str) -> None: + report = ( + run_paper_pipeline(parse_expression(case.source)) + if variant == "binary" + else run_paper_pipeline_container(binary_to_containers(parse_expression(case.source))) + ) + expected_status = "iteration_limit" if (case.name, variant) == ("repeated_monomial", "binary") else "saturated" + assert report.status == expected_status + assert 1 <= report.passes <= 2 + assert report.extracted_params <= report.before_params + assert report.extracted_params < report.before_params or report.extracted_nodes < report.before_nodes + parse_expression(report.extracted) + assert case.sample_points + for x0, x1 in case.sample_points: + expected = evaluate(case.source, x0=x0, x1=x1) + actual = evaluate(report.extracted, x0=x0, x1=x1) + assert math.isfinite(expected) + assert math.isfinite(actual) + assert math.isclose(actual, expected, rel_tol=1e-9, abs_tol=1e-9) + + +@pytest.mark.param_eq_smoke +def test_log_log_is_not_rewritten_to_its_argument() -> None: + report = run_paper_pipeline(parse_expression("log(log(x0))")) + assert math.isclose( + evaluate(report.extracted, x0=math.e**2, x1=0.0), + evaluate("log(log(x0))", x0=math.e**2, x1=0.0), + rel_tol=1e-9, + abs_tol=1e-9, + ) + + +@pytest.mark.param_eq_smoke +def test_container_log_product_rule_equates_expanded_log() -> None: + source = binary_to_containers(parse_expression("log(2.0 * x0)")) + expected = binary_to_containers(parse_expression("log(2.0) + log(x0)")) + egraph = EGraph(source, save_egglog_string=False) + + egraph.run(containers_analysis_schedule) + egraph.run(container_schedule) + + # Introduce and normalize the comparison only after the product rule has run, + # so it cannot pre-seed the constructors that the rule's RHS must create. + expected = egraph.let("expected", expected) + egraph.run(containers_analysis_schedule) + egraph.check(eq(source).to(expected)) + + +@pytest.mark.param_eq_smoke +def test_container_log_rule_preserves_even_power_domain() -> None: + source_text = "log(x0 ** 2.0)" + report = run_paper_pipeline_container(binary_to_containers(parse_expression(source_text))) + + assert math.isclose(evaluate(source_text, x0=-1.0, x1=0.0), 0.0, abs_tol=1e-9) + assert math.isclose(evaluate(report.extracted, x0=-1.0, x1=0.0), 0.0, abs_tol=1e-9) + + +@pytest.mark.param_eq_smoke +@pytest.mark.parametrize("source", ["exp(1000.5)", "(-1.5) ** 0.25"]) +def test_nonfinite_constant_results_are_not_folded(source: str) -> None: + report = run_paper_pipeline(parse_expression(source)) + + assert parse_expression(report.extracted) == parse_expression(source) + + +@pytest.mark.param_eq_smoke +@pytest.mark.parametrize("source", ["(-1.5) ** 0.25", "1e308 * 1e308", "(1e308 * x0) * 1e308"]) +def test_container_pipeline_rejects_nonfinite_coefficient_normalization(source: str) -> None: + with pytest.raises( + ValueError, match="container pipeline requires every coefficient normalization to remain finite" + ): + run_paper_pipeline_container(binary_to_containers(parse_expression(source))) + + +@pytest.mark.param_eq_smoke +def test_iteration_limited_pass_is_not_reported_as_saturated(monkeypatch: pytest.MonkeyPatch) -> None: + x = var("iteration_limit_x", Num) + limited_rules = ruleset(rewrite(x + Num(0.0)).to(x), name="iteration-limit-rules") + empty_analysis = ruleset(name="iteration-limit-analysis").saturate() + limited_schedule = run( + limited_rules, + scheduler=back_off(match_limit=0, ban_length=100).persistent(), + ) + monkeypatch.setattr(pipeline, "HASKELL_INNER_ITERATION_LIMIT", 1) + + report = run_paper_pipeline( + parse_expression("x0 + 0.0"), + schedule=limited_schedule, + analysis_schedule=empty_analysis, + ) + + assert report.status == "iteration_limit" + + +@pytest.mark.param_eq_smoke +def test_cli_emits_complete_json(capsys: pytest.CaptureFixture[str]) -> None: + assert main(["--variant", "binary", "--expr", DEMO_CASES[2].source]) == 0 + payload = json.loads(capsys.readouterr().out) + assert payload["variant"] == "binary" + assert payload["status"] == "saturated" + assert payload["extracted"] + assert payload["extracted_params"] <= payload["before_params"] diff --git a/python/tests/param_eq/test_research_harness.py b/python/tests/param_eq/test_research_harness.py new file mode 100644 index 00000000..8a233560 --- /dev/null +++ b/python/tests/param_eq/test_research_harness.py @@ -0,0 +1,608 @@ +from __future__ import annotations + +import csv +import hashlib +import json +import os +import subprocess +import tomllib +from pathlib import Path +from unittest.mock import Mock + +import pytest +from experiments.param_eq import run as param_eq_run +from experiments.param_eq.aggregate import ( + COMPARISON_COLUMNS, + PAPER_COLUMNS, + build_paper_replication, + build_representation_comparison, + generate_aggregates, + load_raw, +) +from experiments.param_eq.corpus import ( + RAW_ALGORITHMS, + ArchiveLayoutError, + CorpusRow, + external_archive_hash, + load_corpus_rows, +) +from experiments.param_eq.run import RAW_COLUMNS, _build_haskell_program, _parse_haskell_output + +from egglog.exp.param_eq import PaperPipelineReport + + +@pytest.fixture(autouse=True) +def _stable_repository_state(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("experiments.param_eq.aggregate._repo_state", lambda: ("d" * 40, True)) + + +def _archive(root: Path) -> Path: + results = root / "results" + (results / "exprs").mkdir(parents=True) + (results / "exprs_simpl").mkdir() + for dataset in ("pagie", "kotanchek"): + for algorithm in RAW_ALGORITHMS: + payload = "2.3*x0\n9.9*x0\n3.7*x0\n" if dataset == "pagie" and algorithm == "Bingo" else "" + (results / "exprs" / f"{algorithm}_exprs_{dataset}").write_text(payload, encoding="utf-8") + (results / "exprs_simpl" / f"{algorithm}_exprs_{dataset}").write_text(payload, encoding="utf-8") + (results / f"{dataset}_results").write_text( + "algorithm,expr,expr_sympy\n" + + ("Bingo,2.3*x0,2.3*x0\nBingo,9.9*x0,9.9*x0\nBingo,3.7*x0,3.7*x0\n" if dataset == "pagie" else ""), + encoding="utf-8", + ) + (results / f"{dataset}_table_counts.csv").write_text( + ( + ",orig_nodes,orig_params,simpl_nodes,simpl_params,orig_nodes_sympy,orig_params_sympy," + "simpl_nodes_sympy,simpl_params_sympy,algorithm,n_params,n_rank\n" + + ( + "0,3,1,3,1,3,1,3,1,Bingo,1,1\n2,3,1,3,1,3,1,3,1,Bingo,1,\n3,3,1,3,1,3,1,3,1,Bingo,1,1\n" + if dataset == "pagie" + else "" + ) + ), + encoding="utf-8", + ) + return root + + +def _raw_row(row_id: str, variant: str, *, status: str = "saturated") -> dict[str, str]: + values = { + "row_id": row_id, + "dataset": "pagie", + "algorithm": "Bingo", + "input_kind": "original", + "implementation": "egglog", + "variant": variant, + "external_archive_sha256": "a" * 64, + "source_n_rank": "1", + "timeout_sec": "60", + "memory_limit_mb": "2048", + "sample_interval_sec": "0.2", + "execution_mode": "release", + "workers_requested": "2", + "workers_effective": "2", + "ordering_seed": "param-eq-variant-order-v1", + "egglog_python_commit": "d" * 40, + "egglog_python_clean": "true", + "egglog_core_commit": "b" * 40, + "egglog_core_clean": "true", + "egglog_experimental_commit": "c" * 40, + "egglog_experimental_clean": "true", + "egglog_bindings_sha256": "e" * 64, + "python_version": "3.13.7", + "platform": "test-platform", + "rust_version": "rustc 1.91.0", + "cpu": "test-cpu", + "status": status, + "runtime_ms": "2.0" if variant == "binary" else "1.0", + "peak_rss_mb": "10.0", + "passes": "2", + "total_size": "20" if variant == "binary" else "10", + "before_nodes": "11", + "before_params": "4", + "after_nodes": "7", + "after_params": "2", + } + if status != "saturated": + for key in ( + "runtime_ms", + "passes", + "total_size", + "before_nodes", + "before_params", + "after_nodes", + "after_params", + ): + values[key] = "" + return values + + +def _write_raw(path: Path, rows: list[dict[str, str]]) -> None: + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=RAW_COLUMNS) + writer.writeheader() + writer.writerows(rows) + + +class _ExitOnUnpickle: + """Make a spawned worker hard-exit while restoring its process target.""" + + def __reduce__(self) -> tuple[object, tuple[int]]: + return os._exit, (17,) + + +def test_haskell_program_forces_results_and_parser_keeps_expression_text_out() -> None: + row = CorpusRow( + row_id="pagie/0/Bingo/1/original", + dataset="pagie", + raw_index=0, + algorithm_raw="Bingo", + algorithm="Bingo", + algorithm_row=1, + input_kind="original", + source="2.3*x0", + source_n_rank=1.0, + ) + program = _build_haskell_program() + assert "beforeNodes <- evaluate (countNodes expr)" in program + assert "beforeParams <- evaluate (recountParams (replaceConstsWithParams expr))" in program + assert "afterNodes <- evaluate (countNodes simplified)" in program + assert "afterParams <- evaluate (recountParams (replaceConstsWithParams simplified))" in program + assert row.source not in program + assert "[dataset, inputKind, algorithm, rowIndex]" in program + assert "lookupExpr dataset inputKind algorithm (read rowIndex)" in program + + parsed = _parse_haskell_output("11\t4\t7\t2\t1.25\n") + assert parsed == { + "status": "saturated", + "runtime_ms": 1.25, + "passes": "", + "total_size": "", + "before_nodes": 11.0, + "before_params": 4.0, + "after_nodes": 7.0, + "after_params": 2.0, + } + with pytest.raises(ValueError, match="five Haskell output fields"): + _parse_haskell_output("11\t4\n") + with pytest.raises(ValueError, match="invalid counts or runtime"): + _parse_haskell_output("11\t4\t7\t2\tnan\n") + + +def test_haskell_results_reuse_the_expression_free_raw_schema(tmp_path: Path) -> None: + path = tmp_path / "haskell.csv" + row = _raw_row("pagie/0/Bingo/1/original", "binary") + row["implementation"] = "haskell" + row["passes"] = row["total_size"] = "" + _write_raw(path, [row]) + + assert load_raw(path, expected_variant="binary", expected_implementation="haskell") == [row] + with pytest.raises(ValueError, match="No egglog/binary rows"): + load_raw(path, expected_variant="binary") + + +def test_iteration_limited_worker_result_has_no_publishable_metrics(monkeypatch: pytest.MonkeyPatch) -> None: + report = PaperPipelineReport( + status="iteration_limit", + passes=1, + total_sec=0.1, + total_size=10, + before_nodes=3, + before_params=1, + extracted="x0", + extracted_nodes=1, + extracted_params=0, + ) + monkeypatch.setattr("egglog.exp.param_eq.run_paper_pipeline", lambda _expr: report) + connection = Mock() + + param_eq_run._worker(connection, "x0", "binary") + + connection.send.assert_called_once_with({"status": "iteration_limit"}) + connection.close.assert_called_once_with() + + +def test_run_rows_records_error_when_worker_exits_without_sending( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + row = CorpusRow( + row_id="pagie/0/Bingo/1/original", + dataset="pagie", + raw_index=0, + algorithm_raw="Bingo", + algorithm="Bingo", + algorithm_row=1, + input_kind="original", + source="2.3*x0", + source_n_rank=1.0, + ) + expected = _raw_row(row.row_id, "binary", status="error") + provenance = { + column: expected[column] + for column in RAW_COLUMNS[RAW_COLUMNS.index("execution_mode") : RAW_COLUMNS.index("status")] + } + monkeypatch.setattr(param_eq_run, "_worker", _ExitOnUnpickle()) + + results = param_eq_run.run_rows( + [row], + implementation="egglog", + variant="binary", + archive_root=tmp_path, + haskell_executable=None, + external_archive_sha256="a" * 64, + workers=1, + timeout_sec=5.0, + memory_limit_mb=2048, + sample_interval_sec=0.01, + provenance=provenance, + ) + + assert len(results) == 1 + assert set(results[0]) == set(RAW_COLUMNS) + assert results[0]["status"] == "error" + assert all( + results[0][column] == "" + for column in ( + "runtime_ms", + "passes", + "total_size", + "before_nodes", + "before_params", + "after_nodes", + "after_params", + ) + ) + + +def test_run_provenance_hashes_the_loaded_native_extension(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + repository = tmp_path / "egglog-python" + core = tmp_path / "egglog" + experimental = tmp_path / "egglog-experimental" + artifact = tmp_path / "bindings.so" + artifact.write_bytes(b"loaded native extension") + metadata = { + "packages": [ + {"name": "egglog", "manifest_path": str(core / "Cargo.toml")}, + {"name": "egglog-experimental", "manifest_path": str(experimental / "Cargo.toml")}, + ] + } + + def run_metadata(command: list[str], **_kwargs: object) -> subprocess.CompletedProcess[str]: + if command[:2] == ["cargo", "metadata"]: + return subprocess.CompletedProcess(command, 0, json.dumps(metadata), "") + if command[:3] == ["git", "rev-parse", "HEAD"]: + return subprocess.CompletedProcess(command, 0, "d" * 40 + "\n", "") + if command[:3] == ["git", "status", "--porcelain"]: + return subprocess.CompletedProcess(command, 0, "", "") + if command[:2] == ["rustc", "--version"]: + return subprocess.CompletedProcess(command, 0, "rustc test\n", "") + if command[:2] == ["sysctl", "-n"]: + return subprocess.CompletedProcess(command, 0, "test cpu\n", "") + raise AssertionError(f"Unexpected command: {command}") + + monkeypatch.setattr(param_eq_run, "REPO_ROOT", repository) + monkeypatch.setattr(param_eq_run.egglog_bindings, "__file__", str(artifact)) + monkeypatch.setattr(param_eq_run.subprocess, "run", run_metadata) + monkeypatch.setattr(param_eq_run.platform, "processor", lambda: "test cpu") + monkeypatch.setattr(param_eq_run.platform, "system", lambda: "Test") + monkeypatch.setattr(param_eq_run.platform, "platform", lambda: "test platform") + + provenance = param_eq_run._collect_run_provenance( + execution_mode="release", workers_requested=2, workers_effective=1 + ) + + assert provenance["egglog_bindings_sha256"] == hashlib.sha256(artifact.read_bytes()).hexdigest() + + +def test_external_loader_is_explicit_and_expression_data_stays_in_memory(tmp_path: Path) -> None: + with pytest.raises(ArchiveLayoutError, match="incomplete"): + load_corpus_rows(tmp_path) + root = _archive(tmp_path) + rows = load_corpus_rows(root) + assert [(row.row_id, row.input_kind) for row in rows] == [ + ("pagie/0/Bingo/1/original", "original"), + ("pagie/0/Bingo/1/sympy", "sympy"), + ("pagie/3/Bingo/3/original", "original"), + ("pagie/3/Bingo/3/sympy", "sympy"), + ] + assert [row.source for row in rows] == ["2.3*x0", "2.3*x0", "3.7*x0", "3.7*x0"] + + +def test_external_hash_is_stable_and_content_sensitive(tmp_path: Path) -> None: + root = _archive(tmp_path) + first = external_archive_hash(root) + assert external_archive_hash(root) == first + target = root / "results" / "pagie_results" + target.write_text(target.read_text(encoding="utf-8") + "\n", encoding="utf-8") + assert external_archive_hash(root) != first + + +def test_aggregate_outputs_have_only_aggregate_schemas(tmp_path: Path) -> None: + binary = tmp_path / "binary.csv" + container = tmp_path / "container.csv" + output = tmp_path / "published" + _write_raw(binary, [_raw_row("pagie/0/Bingo/1/original", "binary")]) + _write_raw(container, [_raw_row("pagie/0/Bingo/1/original", "container")]) + generate_aggregates(binary, container, output) + with (output / "paper-replication.csv").open(newline="", encoding="utf-8") as handle: + assert tuple(csv.DictReader(handle).fieldnames or ()) == PAPER_COLUMNS + with (output / "representation-comparison.csv").open(newline="", encoding="utf-8") as handle: + reader = csv.DictReader(handle) + assert tuple(reader.fieldnames or ()) == COMPARISON_COLUMNS + all_runtime = next(row for row in reader if row["slice"] == "all" and row["metric"] == "runtime_ms") + assert all_runtime["container_better"] == "1" + combined = "\n".join(path.read_text(encoding="utf-8") for path in output.iterdir()) + assert str(tmp_path) not in combined + assert "expression" not in combined + assert "external_archive_sha256" not in combined + assert "a" * 64 not in combined + + +def test_paired_raw_file_drives_the_exact_aggregate_path(tmp_path: Path) -> None: + paired = tmp_path / "paired.csv" + output = tmp_path / "published" + _write_raw( + paired, + [ + _raw_row("pagie/0/Bingo/1/original", "binary"), + _raw_row("pagie/0/Bingo/1/original", "container"), + ], + ) + + generate_aggregates(paired, paired, output) + + manifest = tomllib.loads((output / "manifest.toml").read_text(encoding="utf-8")) + assert manifest["schema_version"] == 5 + assert manifest["egglog_bindings_sha256"] == "e" * 64 + assert manifest["python"] == "3.13.7" + assert manifest["platform"] == "test-platform" + assert manifest["raw_layout"] == "paired-single-file" + assert manifest["binary_raw_sha256"] == manifest["container_raw_sha256"] + assert manifest["binary_rows"] == manifest["container_rows"] == 1 + assert manifest["binary_status_saturated"] == manifest["container_status_saturated"] == 1 + assert manifest["recommended_full_rerun"] == "make -C experiments/param_eq aggregate" + assert "command" not in manifest + + +def test_paper_aggregate_keeps_failure_statuses_distinct() -> None: + statuses = ("saturated", "iteration_limit", "timeout", "memory_limit", "error") + rows = [ + _raw_row(f"pagie/{index}/Bingo/{index + 1}/original", "binary", status=status) + for index, status in enumerate(statuses) + ] + + aggregate = next(row for row in build_paper_replication(rows) if row["metric"] == "final_params") + + assert aggregate["n_total"] == 5 + assert aggregate["n_success"] == 1 + assert aggregate["n_iteration_limit"] == 1 + assert aggregate["n_timeout"] == 1 + assert aggregate["n_memory_limit"] == 1 + assert aggregate["n_error"] == 1 + + +def test_ratio_summary_discloses_zero_denominator_omissions() -> None: + binary = _raw_row("pagie/0/Bingo/1/original", "binary") + container = _raw_row("pagie/0/Bingo/1/original", "container") + binary["after_params"] = "0" + container["after_params"] = "1" + + comparison = build_representation_comparison([binary], [container]) + aggregate = next(row for row in comparison if row["slice"] == "all" and row["metric"] == "after_params") + + assert aggregate["n_pairs"] == 1 + assert aggregate["n_ratio"] == 0 + assert aggregate["container_worse"] == 1 + assert aggregate["ratio_median"] == "" + + +def test_aggregate_refuses_mismatched_producer_environment(tmp_path: Path) -> None: + binary_row = _raw_row("pagie/0/Bingo/1/original", "binary") + container_row = _raw_row("pagie/0/Bingo/1/original", "container") + container_row["platform"] = "other-platform" + binary = tmp_path / "binary.csv" + container = tmp_path / "container.csv" + _write_raw(binary, [binary_row]) + _write_raw(container, [container_row]) + + with pytest.raises(ValueError, match="nonempty platform"): + generate_aggregates(binary, container, tmp_path / "published") + + +@pytest.mark.parametrize( + ("column", "value", "message"), + [ + ("egglog_python_clean", "false", "dirty repository state"), + ("egglog_python_commit", "abcdef0", "invalid egglog_python_commit"), + ("egglog_core_clean", "false", "dirty repository state"), + ("egglog_core_commit", "unknown", "invalid egglog_core_commit"), + ("egglog_bindings_sha256", "not-a-hash", "invalid egglog_bindings_sha256"), + ], +) +def test_aggregate_refuses_unpublishable_dependency_provenance( + tmp_path: Path, column: str, value: str, message: str +) -> None: + binary_row = _raw_row("pagie/0/Bingo/1/original", "binary") + container_row = _raw_row("pagie/0/Bingo/1/original", "container") + binary_row[column] = value + container_row[column] = value + binary = tmp_path / "binary.csv" + container = tmp_path / "container.csv" + _write_raw(binary, [binary_row]) + _write_raw(container, [container_row]) + with pytest.raises(ValueError, match=message): + generate_aggregates(binary, container, tmp_path / "published") + + +def test_aggregate_refuses_a_dirty_egglog_python_worktree(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + binary = tmp_path / "binary.csv" + container = tmp_path / "container.csv" + output = tmp_path / "published" + _write_raw(binary, [_raw_row("pagie/0/Bingo/1/original", "binary")]) + _write_raw(container, [_raw_row("pagie/0/Bingo/1/original", "container")]) + monkeypatch.setattr("experiments.param_eq.aggregate._repo_state", lambda: ("d" * 40, False)) + + with pytest.raises(ValueError, match="dirty egglog-python worktree"): + generate_aggregates(binary, container, output) + assert not output.exists() + + +def test_aggregate_refuses_raw_from_a_different_egglog_python_commit(tmp_path: Path) -> None: + binary_row = _raw_row("pagie/0/Bingo/1/original", "binary") + container_row = _raw_row("pagie/0/Bingo/1/original", "container") + binary_row["egglog_python_commit"] = "e" * 40 + container_row["egglog_python_commit"] = "e" * 40 + binary = tmp_path / "binary.csv" + container = tmp_path / "container.csv" + output = tmp_path / "published" + _write_raw(binary, [binary_row]) + _write_raw(container, [container_row]) + + with pytest.raises(ValueError, match="produced by a different egglog-python commit"): + generate_aggregates(binary, container, output) + assert not output.exists() + + +@pytest.mark.parametrize( + ("status", "column", "value", "message"), + [ + ("saturated", "runtime_ms", "nan", "invalid runtime_ms"), + ("saturated", "after_nodes", "-1", "invalid after_nodes"), + ("saturated", "before_params", "1.5", "noninteger before_params"), + ("saturated", "passes", "0", "no completed passes"), + ("timeout", "runtime_ms", "1", "unexpectedly has result metrics"), + ], +) +def test_aggregate_rejects_invalid_or_status_inconsistent_metrics( + tmp_path: Path, status: str, column: str, value: str, message: str +) -> None: + binary_row = _raw_row("pagie/0/Bingo/1/original", "binary", status=status) + container_row = _raw_row("pagie/0/Bingo/1/original", "container", status=status) + binary_row[column] = value + container_row[column] = value + binary = tmp_path / "binary.csv" + container = tmp_path / "container.csv" + output = tmp_path / "published" + _write_raw(binary, [binary_row]) + _write_raw(container, [container_row]) + + with pytest.raises(ValueError, match=message): + generate_aggregates(binary, container, output) + assert not output.exists() + + +def test_aggregate_validates_all_publication_inputs_before_writing(tmp_path: Path) -> None: + binary_row = _raw_row("pagie/0/Bingo/1/original", "binary") + container_row = _raw_row("pagie/0/Bingo/1/original", "container") + binary_row["external_archive_sha256"] = "not-a-hash" + container_row["external_archive_sha256"] = "not-a-hash" + binary = tmp_path / "binary.csv" + container = tmp_path / "container.csv" + output = tmp_path / "published" + _write_raw(binary, [binary_row]) + _write_raw(container, [container_row]) + + with pytest.raises(ValueError, match="external_archive_sha256"): + generate_aggregates(binary, container, output) + assert not output.exists() + + +def test_aggregate_rejects_private_algorithm_label_without_leaking_or_writing(tmp_path: Path) -> None: + sentinel = "PRIVATE_EXPR_SENTINEL(x0)" + binary_row = _raw_row("pagie/0/Bingo/1/original", "binary") + container_row = _raw_row("pagie/0/Bingo/1/original", "container") + binary_row["algorithm"] = sentinel + container_row["algorithm"] = sentinel + binary = tmp_path / "binary.csv" + container = tmp_path / "container.csv" + output = tmp_path / "published" + _write_raw(binary, [binary_row]) + _write_raw(container, [container_row]) + + with pytest.raises(ValueError, match="invalid algorithm label") as error: + generate_aggregates(binary, container, output) + + assert sentinel not in str(error.value) + assert not output.exists() + + +def test_write_local_raw_requires_an_untracked_ignored_target(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + repo = tmp_path / "repo" + raw_results = repo / "results" / "raw" + raw_results.mkdir(parents=True) + ignore_file = raw_results / ".gitignore" + ignore_contents = "*\n!.gitignore\n" + ignore_file.write_text(ignore_contents, encoding="utf-8") + subprocess.run(["git", "init", "--quiet"], cwd=repo, check=True) # noqa: S607 + subprocess.run( + ["git", "add", "--", "results/raw/.gitignore"], # noqa: S607 + cwd=repo, + check=True, + ) + monkeypatch.setattr(param_eq_run, "REPO_ROOT", repo) + monkeypatch.setattr(param_eq_run, "RAW_RESULTS_DIR", raw_results) + row = _raw_row("pagie/0/Bingo/1/original", "binary") + + with pytest.raises(ValueError, match="untracked, Git-ignored"): + param_eq_run.write_local_raw([row], ignore_file) + assert ignore_file.read_text(encoding="utf-8") == ignore_contents + + output = raw_results / "binary.csv" + param_eq_run.write_local_raw([row], output) + assert load_raw(output, expected_variant="binary") == [row] + + +@pytest.mark.parametrize( + ("expected_hash", "message"), + [ + (None, "Set EGGLOG_PARAM_EQ_EXPECTED_ARCHIVE_SHA256"), + ("b" * 64, "does not match EGGLOG_PARAM_EQ_EXPECTED_ARCHIVE_SHA256"), + ], +) +def test_runner_requires_the_privately_recorded_archive_hash( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + expected_hash: str | None, + message: str, +) -> None: + monkeypatch.setattr(param_eq_run, "external_archive_root", lambda: tmp_path) + monkeypatch.setattr(param_eq_run, "load_corpus_rows", lambda *args, **kwargs: []) + monkeypatch.setattr(param_eq_run, "external_archive_hash", lambda root: "a" * 64) + if expected_hash is None: + monkeypatch.delenv("EGGLOG_PARAM_EQ_EXPECTED_ARCHIVE_SHA256", raising=False) + else: + monkeypatch.setenv("EGGLOG_PARAM_EQ_EXPECTED_ARCHIVE_SHA256", expected_hash) + + with pytest.raises(SystemExit, match="2"): + param_eq_run.main(["--execution-mode", "debug"]) + + stderr = capsys.readouterr().err + assert message in stderr + assert "a" * 64 not in stderr + + +def test_runner_rejects_filters_that_match_no_rows( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] +) -> None: + monkeypatch.setattr(param_eq_run, "external_archive_root", lambda: tmp_path) + monkeypatch.setattr(param_eq_run, "load_corpus_rows", lambda *args, **kwargs: []) + monkeypatch.setattr(param_eq_run, "external_archive_hash", lambda root: "a" * 64) + monkeypatch.setenv("EGGLOG_PARAM_EQ_EXPECTED_ARCHIVE_SHA256", "a" * 64) + + with pytest.raises(SystemExit, match="2"): + param_eq_run.main(["--execution-mode", "debug", "--algorithm", "typo"]) + + assert "selected corpus filters matched no rows" in capsys.readouterr().err + + +def test_raw_schema_rejects_expression_columns_and_unpaired_rows(tmp_path: Path) -> None: + bad = tmp_path / "bad.csv" + bad.write_text(",".join([*RAW_COLUMNS, "expression"]) + "\n", encoding="utf-8") + with pytest.raises(ValueError, match="forbidden columns"): + load_raw(bad, expected_variant="binary") + with pytest.raises(ValueError, match="Unpaired raw inputs"): + build_representation_comparison( + [_raw_row("pagie/0/Bingo/1/original", "binary")], + [_raw_row("pagie/3/Bingo/3/original", "container")], + ) diff --git a/python/tests/param_eq/test_resource_guard.py b/python/tests/param_eq/test_resource_guard.py new file mode 100644 index 00000000..fecbdd9f --- /dev/null +++ b/python/tests/param_eq/test_resource_guard.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +import subprocess +import sys +from unittest.mock import MagicMock + +import pytest +from experiments.param_eq import resource_guard + + +def test_cap_workers_for_memory_respects_safe_memory_budget() -> None: + gib = 1024**3 + + assert ( + resource_guard.cap_workers_for_memory( + 8, + memory_limit_mb=1024, + total_memory_bytes_value=8 * gib, + safe_memory_fraction=0.5, + ) + == 4 + ) + + +def test_process_tree_rss_includes_descendants_only(monkeypatch: pytest.MonkeyPatch) -> None: + completed = subprocess.CompletedProcess( + ["ps"], + 0, + stdout="100 1 1024\n101 100 2048\n102 101 3072\n200 1 4096\nmalformed\n", + ) + monkeypatch.setattr(resource_guard.subprocess, "run", lambda *args, **kwargs: completed) + + assert resource_guard._process_tree_rss_mb(100) == 6.0 + + +def test_watch_process_kills_multiprocessing_worker_at_memory_limit(monkeypatch: pytest.MonkeyPatch) -> None: + process = MagicMock() + process.pid = 1234 + process.is_alive.return_value = True + monkeypatch.setattr(resource_guard, "_rss_mb", lambda pid: 65.0) + + result = resource_guard.watch_process(process, timeout_sec=10.0, memory_limit_mb=64) + + assert result == resource_guard.WatchResult("memory_limit", 65.0) + process.kill.assert_called_once_with() + process.join.assert_called_once_with(timeout=1.0) + + +def test_watch_subprocess_kills_process_group_at_memory_limit(monkeypatch: pytest.MonkeyPatch) -> None: + process = MagicMock() + process.pid = 5678 + process.poll.return_value = None + kill_process_group = MagicMock() + monkeypatch.setattr(resource_guard, "_process_tree_rss_mb", lambda pid: 129.0) + monkeypatch.setattr(resource_guard, "_kill_subprocess_group", kill_process_group) + + result = resource_guard.watch_subprocess(process, timeout_sec=10.0, memory_limit_mb=128) + + assert result == resource_guard.WatchResult("memory_limit", 129.0) + kill_process_group.assert_called_once_with(process) + + +def test_watch_subprocess_timeout_terminates_real_process() -> None: + process = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(10)"], + text=True, + start_new_session=True, + ) + try: + result = resource_guard.watch_subprocess( + process, + timeout_sec=0.05, + memory_limit_mb=1_000_000_000, + sample_interval_sec=0.01, + ) + + assert result.status == "timeout" + assert process.poll() is not None + finally: + if process.poll() is None: + process.kill() + process.wait(timeout=1.0) diff --git a/python/tests/param_eq/test_semantics.py b/python/tests/param_eq/test_semantics.py new file mode 100644 index 00000000..43d77dea --- /dev/null +++ b/python/tests/param_eq/test_semantics.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +import math + +import pytest + +from .evaluation import evaluate + + +@pytest.mark.parametrize( + ("source", "expected"), + [ + ("2.0*x0 + 3.0*x1", 8.0), + ("log(exp(x0))", 1.0), + ("sqrt(x0**2)", 1.0), + ], +) +def test_independent_evaluator(source: str, expected: float) -> None: + assert math.isclose(evaluate(source, x0=1.0, x1=2.0), expected) + + +def test_independent_evaluator_rejects_arbitrary_python() -> None: + with pytest.raises(ValueError, match="Unsupported numeric-test syntax"): + evaluate("__import__('os')", x0=1.0, x1=2.0) + + +def test_independent_evaluator_rejects_boolean_literals() -> None: + with pytest.raises(ValueError, match="Unsupported numeric-test syntax"): + evaluate("True", x0=1.0, x1=2.0) diff --git a/python/tests/test_bindings.py b/python/tests/test_bindings.py index a5bc71ae..b26eeea4 100644 --- a/python/tests/test_bindings.py +++ b/python/tests/test_bindings.py @@ -4,6 +4,7 @@ import pathlib import subprocess from base64 import standard_b64encode +from datetime import timedelta from fractions import Fraction import black @@ -71,7 +72,7 @@ def test_example(example_file: pathlib.Path): def extract_best_term(program: str) -> str: - egraph = EGraph(record=True) + egraph = EGraph() outputs = egraph.run_program(*egraph.parse_program(program)) extract = next(output for output in outputs if isinstance(output, ExtractBest)) return extract.termdag.to_string(extract.term) @@ -99,6 +100,62 @@ def test_parse_and_run_program(self): assert egraph.run_program(*egraph.parse_program(program)) == [] + def test_parse_program_preserves_uf_extraction_behavior(self): + program = (EGG_SMOL_FOLDER / "tests" / "uf-extraction.egg").read_text() + + direct_egraph = EGraph() + direct_outputs = direct_egraph.parse_and_run_program(program) + direct_extract = next(output for output in direct_outputs if isinstance(output, ExtractBest)) + + converted_egraph = EGraph() + converted_outputs = converted_egraph.run_program(*converted_egraph.parse_program(program)) + converted_extract = next(output for output in converted_outputs if isinstance(output, ExtractBest)) + + assert direct_extract.termdag.to_string(direct_extract.term) == "(Foo)" + assert converted_extract.termdag.to_string(converted_extract.term) == "(Foo)" + + def test_parse_program_preserves_internal_hidden_behavior(self): + program = """ + (sort Expr) + (constructor Visible () Expr) + (constructor Hidden () Expr :internal-hidden) + (function visible-f () i64 :no-merge) + (function hidden-f () i64 :no-merge :internal-hidden) + (Visible) + (Hidden) + (set (visible-f) 1) + (set (hidden-f) 2) + (print-size) + """ + + direct_egraph = EGraph() + direct_outputs = direct_egraph.parse_and_run_program(program) + direct_sizes = next(output for output in direct_outputs if isinstance(output, PrintAllFunctionsSize)) + + converted_egraph = EGraph() + converted_outputs = converted_egraph.run_program(*converted_egraph.parse_program(program)) + converted_sizes = next(output for output in converted_outputs if isinstance(output, PrintAllFunctionsSize)) + + assert direct_sizes.sizes == [("Visible", 1), ("visible-f", 1)] + assert converted_sizes.sizes == direct_sizes.sizes + + @pytest.mark.parametrize("parse_and_run", [False, True]) + def test_command_recording(self, parse_and_run: bool): + program = """(function f (i64) i64 :no-merge) + (set (f 1) 2) + (check (= (f 1) 2))""" + assert EGraph().commands() is None + + egraph = EGraph(record=True) + assert egraph.commands() == "" + commands = egraph.parse_program(program) + expected = "".join(f"{command}\n" for command in commands) + if parse_and_run: + egraph.parse_and_run_program(program) + else: + egraph.run_program(*commands) + assert egraph.commands() == expected + def test_parse_and_run_program_exception(self): program = "(check (= 1 1.0))" egraph = EGraph() @@ -109,6 +166,21 @@ def test_parse_and_run_program_exception(self): ): egraph.run_program(*egraph.parse_program(program)) + def test_parse_and_run_program_error_keeps_recording_transactional(self): + program = """(function f (i64) i64 :no-merge) + (set (f 1) 2) + (check (= 1 1.0))""" + egraph = EGraph(record=True) + + with pytest.raises(EggSmolError, match=r"In 3:.*recording-error\.egg"): + egraph.parse_and_run_program(program, filename="recording-error.egg") + + assert egraph.commands() == "" + _, key = egraph.eval_expr(Lit(DUMMY_SPAN, Int(1))) + value = egraph.lookup_function("f", [key]) + assert value is not None + assert egraph.value_to_i64(value) == 2 + def test_run_rules(self): egraph = EGraph() res = egraph.run_program( @@ -128,6 +200,186 @@ def test_run_rules(self): assert len(res) == 1 assert isinstance(res[0], RunScheduleOutput) + @pytest.mark.parametrize( + ("label", "command"), + [ + ("ruleset", AddRuleset(DUMMY_SPAN, "rs")), + ("relation", Relation(DUMMY_SPAN, "rel", ["i64"])), + ( + "function", + FunctionCommand(DUMMY_SPAN, "f", Schema(["i64"], "i64"), None, "term-f", True), + ), + ("constructor", Constructor(DUMMY_SPAN, "C", Schema(["i64"], "Expr"), None, False)), + ("let-action", ActionCommand(Let(DUMMY_SPAN, "$x", Lit(DUMMY_SPAN, Int(1))))), + ("set-action", ActionCommand(Set(DUMMY_SPAN, "f", [Lit(DUMMY_SPAN, Int(1))], Lit(DUMMY_SPAN, Int(2))))), + ("union-action", ActionCommand(Union(DUMMY_SPAN, Lit(DUMMY_SPAN, Int(1)), Lit(DUMMY_SPAN, Int(2))))), + ( + "rewrite", + RewriteCommand( + "", + Rewrite( + DUMMY_SPAN, + Call(DUMMY_SPAN, "Add", [Var(DUMMY_SPAN, "a"), Var(DUMMY_SPAN, "b")]), + Call(DUMMY_SPAN, "Add", [Var(DUMMY_SPAN, "b"), Var(DUMMY_SPAN, "a")]), + ), + False, + ), + ), + ( + "rule", + RuleCommand( + Rule( + DUMMY_SPAN, + [ + Union( + DUMMY_SPAN, + Var(DUMMY_SPAN, "lhs"), + Call(DUMMY_SPAN, "Add", [Var(DUMMY_SPAN, "a"), Var(DUMMY_SPAN, "a")]), + ) + ], + [ + Eq( + DUMMY_SPAN, + Var(DUMMY_SPAN, "lhs"), + Call(DUMMY_SPAN, "Mul", [Var(DUMMY_SPAN, "a"), Lit(DUMMY_SPAN, Int(2))]), + ) + ], + "", + "", + ) + ), + ), + ("run-schedule", RunSchedule(Repeat(DUMMY_SPAN, 2, Run(DUMMY_SPAN, RunConfig(""))))), + ("check", Check(DUMMY_SPAN, [Eq(DUMMY_SPAN, Lit(DUMMY_SPAN, Int(1)), Lit(DUMMY_SPAN, Int(1)))])), + ( + "fail-check", + Fail(DUMMY_SPAN, Check(DUMMY_SPAN, [Eq(DUMMY_SPAN, Lit(DUMMY_SPAN, Int(1)), Lit(DUMMY_SPAN, Int(2)))])), + ), + ( + "fail-let-action", + Fail(DUMMY_SPAN, ActionCommand(Let(DUMMY_SPAN, "$x", Call(DUMMY_SPAN, "map-empty", [])))), + ), + ], + ) + def test_command_display_round_trip(self, label: str, command) -> None: + egraph = EGraph() + text = str(command) + (parsed,) = egraph.parse_program(text) + assert str(parsed) == text, label + + @pytest.mark.parametrize( + ("option", "mode_type"), + [("", Seminaive), (":naive", Naive), (":unsafe-seminaive", UnsafeSeminaive)], + ) + def test_rule_options_round_trip(self, option: str, mode_type: type) -> None: + text = f"(rule ((rel x)) ((rel x)) {option} :no-decomp :internal-include-subsumed)" + (parsed,) = EGraph().parse_program(text) + assert isinstance(parsed, RuleCommand) + assert isinstance(parsed.rule.eval_mode, mode_type) + assert parsed.rule.no_decomp + assert parsed.rule.include_subsumed + + (reparsed,) = EGraph().parse_program(str(parsed)) + assert str(reparsed) == str(parsed) + + def test_rewrite_name_conversion(self) -> None: + (parsed,) = EGraph().parse_program('(rewrite x x :name "binding-rewrite-name")') + assert isinstance(parsed, RewriteCommand) + assert parsed.rewrite.name == "binding-rewrite-name" + assert 'name: "binding-rewrite-name"' in str(parsed.rewrite) + + def test_function_command_defaults_and_metadata_round_trip(self) -> None: + default = FunctionCommand(DUMMY_SPAN, "default-f", Schema(["i64"], "i64"), None) + assert default.term_constructor is None + assert not default.unextractable + assert not default.hidden + assert not default.let_binding + + explicit = FunctionCommand(DUMMY_SPAN, "f", Schema(["i64"], "i64"), None, "term-f", True) + (parsed,) = EGraph().parse_program(str(explicit)) + assert isinstance(parsed, FunctionCommand) + assert parsed.term_constructor == "term-f" + assert parsed.unextractable + assert not parsed.hidden + assert not parsed.let_binding + + def test_internal_command_metadata_round_trip(self) -> None: + program = """ + (sort Expr :internal-uf UFExpr UFExprIndex :internal-proof-func ExprProof + :internal-proof-names Congr Trans Sym Normalize) + (sort ExprVec (Vec Expr) :internal-proof-func ExprVecProof + :internal-container-rebuild + (container-rebuild-spec rebuild-vec rebuild-vec-proof)) + (function view (Expr) Expr :no-merge :unextractable :internal-hidden + :internal-let :internal-term-constructor View) + (constructor Hidden () Expr :unextractable :internal-hidden :internal-let) + """ + sort, container_sort, function, constructor = EGraph().parse_program(program) + + assert isinstance(sort, Sort) + assert sort.uf == ("UFExpr", "UFExprIndex") + assert sort.proof_func == "ExprProof" + assert sort.proof_constructors == ProofConstructorNames("Congr", "Trans", "Sym", "Normalize") + assert sort.container_rebuild is None + + assert isinstance(container_sort, Sort) + assert container_sort.proof_func == "ExprVecProof" + assert container_sort.container_rebuild == ContainerRebuildSpec("rebuild-vec", "rebuild-vec-proof") + + assert isinstance(function, FunctionCommand) + assert function.hidden + assert function.let_binding + assert function.term_constructor == "View" + assert function.unextractable + + assert isinstance(constructor, Constructor) + assert constructor.hidden + assert constructor.let_binding + assert constructor.unextractable + + for command in (sort, container_sort, function, constructor): + (reparsed,) = EGraph().parse_program(str(command)) + assert str(reparsed) == str(command) + + default_sort = Sort(DUMMY_SPAN, "Default", None) + assert default_sort.uf is None + assert default_sort.proof_func is None + assert default_sort.container_rebuild is None + assert default_sort.proof_constructors is None + + @pytest.mark.parametrize( + "duration", + [timedelta(days=1, seconds=2, microseconds=345_678), timedelta.max], + ) + def test_report_duration_round_trip(self, duration: timedelta): + rule_report = RuleReport(None, duration, 7) + ruleset_report = RuleSetReport(True, {"rule": [rule_report]}, duration, duration) + iteration_report = IterationReport(ruleset_report, duration) + run_report = RunReport( + [iteration_report], + True, + True, + {"rule": duration}, + {"rule": 7}, + {"ruleset": duration}, + {"ruleset": duration}, + {"ruleset": duration}, + ) + + assert rule_report.search_and_apply_time == duration + assert ruleset_report.search_and_apply_time == duration + assert ruleset_report.merge_time == duration + assert iteration_report.rebuild_time == duration + assert run_report.search_and_apply_time_per_rule["rule"] == duration + assert run_report.search_and_apply_time_per_ruleset["ruleset"] == duration + assert run_report.merge_time_per_ruleset["ruleset"] == duration + assert run_report.rebuild_time_per_ruleset["ruleset"] == duration + + @pytest.mark.parametrize("duration", [timedelta(microseconds=-1), timedelta.min]) + def test_report_rejects_negative_duration(self, duration: timedelta): + with pytest.raises(ValueError, match="negative timedeltas"): + RuleReport(None, duration, 7) + def test_extract(self): # Example from extraction-cost egraph = EGraph() @@ -145,6 +397,26 @@ def test_extract(self): DUMMY_SPAN, "Num", [Lit(DUMMY_SPAN, Int(1))] ) + def test_extract_value(self): + egraph = EGraph() + egraph.parse_and_run_program("(sort Expr) (constructor Num (i64) Expr)") + sort, value = egraph.eval_expr(Call(DUMMY_SPAN, "Num", [Lit(DUMMY_SPAN, Int(42))])) + + termdag, term, cost = egraph.extract_value(value, sort) + assert termdag.to_string(term) == "(Num 42)" + assert cost > 0 + + with pytest.raises(EggSmolError, match="Undefined sort Missing"): + egraph.extract_value(value, "Missing") + + def test_extract_value_reports_extraction_failure(self): + egraph = EGraph() + egraph.parse_and_run_program("(sort Expr) (constructor Hidden (i64) Expr :unextractable)") + sort, value = egraph.eval_expr(Call(DUMMY_SPAN, "Hidden", [Lit(DUMMY_SPAN, Int(42))])) + + with pytest.raises(EggSmolError, match="Unable to find any valid extraction"): + egraph.extract_value(value, sort) + def test_sort_alias(self): # From map example egraph = EGraph() @@ -189,6 +461,37 @@ def test_sort_alias(self): Extract(DUMMY_SPAN, Var(DUMMY_SPAN, "my_map2"), Lit(DUMMY_SPAN, Int(0))), ) + def test_freeze_constructor_and_function_rows(self): + egraph = EGraph() + egraph.parse_and_run_program( + """ + (sort Expr) + (constructor Num (i64) Expr) + (function f (i64) i64 :no-merge) + (let $x (Num 1)) + (set (f 2) 3) + """ + ) + + functions = egraph.freeze().functions + num = functions["Num"] + function = functions["f"] + global_x = functions["$x"] + + assert num.input_sorts == ["i64"] + assert num.output_sort == "Expr" + assert len(num.rows) == 1 + assert egraph.value_to_i64(num.rows[0].inputs[0]) == 1 + + assert function.input_sorts == ["i64"] + assert function.output_sort == "i64" + assert len(function.rows) == 1 + assert egraph.value_to_i64(function.rows[0].inputs[0]) == 2 + assert egraph.value_to_i64(function.rows[0].output) == 3 + + assert global_x.is_let_binding + assert global_x.rows[0].output == num.rows[0].output + class TestVariant: def test_repr(self): @@ -259,6 +562,15 @@ def test_bigint(self): assert sort == "BigInt" assert egraph.value_to_bigint(value) == 100 + _, large_value = egraph.eval_expr( + Call( + DUMMY_SPAN, + "<<", + [Call(DUMMY_SPAN, "bigint", [Lit(DUMMY_SPAN, Int(1))]), Lit(DUMMY_SPAN, Int(200))], + ) + ) + assert egraph.value_to_bigint(large_value) == 1 << 200 + def test_bigrat(self): sort, value = egraph.eval_expr( Call( @@ -397,7 +709,33 @@ def test_fn(self): def test_lookup_function(): egraph = EGraph() egraph.run_program(*egraph.parse_program("(function hi (i64) i64 :no-merge)\n(set (hi 1) 2)")) - assert ( - egraph.lookup_function("hi", [egraph.eval_expr(Lit(DUMMY_SPAN, Int(1)))[1]]) - == egraph.eval_expr(Lit(DUMMY_SPAN, Int(2)))[1] - ) + _, one = egraph.eval_expr(Lit(DUMMY_SPAN, Int(1))) + _, two = egraph.eval_expr(Lit(DUMMY_SPAN, Int(2))) + _, absent = egraph.eval_expr(Lit(DUMMY_SPAN, Int(3))) + + assert egraph.lookup_function("hi", [one]) == two + assert egraph.lookup_function("hi", [absent]) is None + with pytest.raises(EggSmolError, match="no table named `missing`"): + egraph.lookup_function("missing", [one]) + + +def test_lookup_constructor(): + egraph = EGraph() + egraph.parse_and_run_program("(sort Expr) (constructor A (i64) Expr)") + _, one = egraph.eval_expr(Lit(DUMMY_SPAN, Int(1))) + _, two = egraph.eval_expr(Lit(DUMMY_SPAN, Int(2))) + _, a_one = egraph.eval_expr(Call(DUMMY_SPAN, "A", [Lit(DUMMY_SPAN, Int(1))])) + + assert egraph.lookup_function("A", [one]) == a_one + assert egraph.lookup_function("A", [two]) is None + + +def test_lookup_relation(): + egraph = EGraph() + egraph.parse_and_run_program("(relation R (i64)) (R 1)") + _, one = egraph.eval_expr(Lit(DUMMY_SPAN, Int(1))) + _, two = egraph.eval_expr(Lit(DUMMY_SPAN, Int(2))) + _, unit = egraph.eval_expr(Lit(DUMMY_SPAN, Unit())) + + assert egraph.lookup_function("R", [one]) == unit + assert egraph.lookup_function("R", [two]) is None diff --git a/python/tests/test_high_level.py b/python/tests/test_high_level.py index 0efd66f8..7a65b6a8 100644 --- a/python/tests/test_high_level.py +++ b/python/tests/test_high_level.py @@ -1,9 +1,11 @@ # mypy: disable-error-code="empty-body" from __future__ import annotations +import gc import importlib +import math import pathlib -from collections.abc import Iterator +from collections.abc import Callable, Iterator from copy import copy from fractions import Fraction from functools import partial @@ -12,8 +14,29 @@ import pytest +import egglog.builtins as egg_builtins from egglog import * -from egglog.declarations import CallDecl, FunctionRef, Ident, JustTypeRef, MethodRef, TypedExprDecl +from egglog import bindings +from egglog.declarations import ( + BUILTIN_EGG_FN_NAMES, + BUILTIN_EGG_SORT_NAMES, + CallableDecl, + CallDecl, + ClassDecl, + Declarations, + FunctionDecl, + FunctionRef, + FunctionSignature, + HasDeclarations, + Ident, + JustTypeRef, + LitDecl, + MethodRef, + TypedExprDecl, + TypeRefWithVars, + ValueDecl, +) +from egglog.egraph import get_current_ruleset from egglog.runtime import RuntimeExpr, RuntimeFunction @@ -26,6 +49,27 @@ def test_ne(self): assert str(ne(i64(1)).to(i64(2))) == "ne(i64(1)).to(i64(2))" +@pytest.mark.parametrize( + ("eval_mode", "binding_type"), + [ + ("seminaive", bindings.Seminaive), + ("naive", bindings.Naive), + ("unsafe-seminaive", bindings.UnsafeSeminaive), + ], +) +def test_rule_eval_mode_lowering(eval_mode: RuleEvalMode, binding_type: type) -> None: + rel = relation(f"eval_mode_{eval_mode}", i64) + x = var("x", i64) + high_level_rule = rule(rel(x), eval_mode=eval_mode).then(rel(x + 1)) + + egraph = EGraph() + egraph._add_decls(high_level_rule) + command = egraph._command_to_egg(high_level_rule) + + assert isinstance(command, bindings.RuleCommand) + assert isinstance(command.rule.eval_mode, binding_type) + + def test_eqsat_basic(): egraph = EGraph() @@ -60,6 +104,20 @@ def __mul__(self, other: Math) -> Math: ... egraph.check(eq(expr1).to(expr2)) +def test_lookup_function_value_constructor_row() -> None: + class A(Expr): + def __init__(self, value: i64Like) -> None: ... + + egraph = EGraph(A(1)) + + value = egraph.lookup_function_value(A(1)) + assert value is not None + a_typed_expr = cast("RuntimeExpr", A(1)).__egg_typed_expr__ + constructor_value = egraph._state.typed_expr_to_value(a_typed_expr) + assert cast("RuntimeExpr", value).__egg_typed_expr__ == TypedExprDecl(a_typed_expr.tp, ValueDecl(constructor_value)) + assert egraph.lookup_function_value(A(2)) is None + + def test_let_auto_prefixes_global_names(capfd: pytest.CaptureFixture[str]): egraph = EGraph(save_egglog_string=True) @@ -71,6 +129,723 @@ def test_let_auto_prefixes_global_names(capfd: pytest.CaptureFixture[str]): assert "(let $x " in egraph.as_egglog_string +def test_failed_check_does_not_materialize_shared_constructor_expressions() -> None: + class CheckEdge(Expr): + @classmethod + def leaf(cls, value: StringLike) -> CheckEdge: ... + + @classmethod + def pair(cls, left: CheckEdge, right: CheckEdge) -> CheckEdge: ... + + pair_rel = relation("check_pair_rel", CheckEdge, CheckEdge) + leaf = CheckEdge.leaf("shared") + pair = CheckEdge.pair(leaf, leaf) + egraph = EGraph() + + with pytest.raises(EggSmolError, match="Check failed"): + egraph.check(pair_rel(pair, pair)) + + assert egraph.function_size(CheckEdge.leaf) == 0 + assert egraph.function_size(CheckEdge.pair) == 0 + + +def test_synthetic_lets_use_reserved_expr_names() -> None: + class LetNum(Expr): + @classmethod + def var(cls, v: StringLike) -> LetNum: ... + + egraph = EGraph(save_egglog_string=True) + expr = LetNum.var("x") + runtime_expr = cast("RuntimeExpr", expr) + egraph._add_decls(runtime_expr) + + egraph._state._transform_let(runtime_expr.__egg_typed_expr__) + + assert '(let $__expr_0 (LetNum_var "x"))' in egraph.as_egglog_string + + +def test_synthetic_lets_skip_explicit_let_conflicts() -> None: + class LetConflictNum(Expr): + @classmethod + def var(cls, v: StringLike) -> LetConflictNum: ... + + egraph = EGraph(save_egglog_string=True) + egraph.let("__expr_0", LetConflictNum.var("explicit")) + expr = LetConflictNum.var("synthetic") + runtime_expr = cast("RuntimeExpr", expr) + egraph._add_decls(runtime_expr) + + egraph._state._transform_let(runtime_expr.__egg_typed_expr__) + + egglog_string = egraph.as_egglog_string + assert '(let $__expr_0 (LetConflictNum_var "explicit"))' in egglog_string + assert '(let $__expr_1 (LetConflictNum_var "synthetic"))' in egglog_string + + +def test_synthetic_let_names_do_not_shadow_default_rewrite_variables() -> None: + default_ruleset = ruleset(name="synthetic-let-shadow-default-rewrite") + + class LetShadowDefaultNum(Expr, ruleset=default_ruleset): + def __init__(self, value: i64Like) -> None: ... + + @classmethod + def make(cls, value: i64Like) -> LetShadowDefaultNum: + return LetShadowDefaultNum(value) + + egraph = EGraph(save_egglog_string=True) + expr = LetShadowDefaultNum(3) + runtime_expr = cast("RuntimeExpr", expr) + egraph._add_decls(runtime_expr) + egraph._state._transform_let(runtime_expr.__egg_typed_expr__) + + egraph.register(LetShadowDefaultNum.make(i64(1))) + egraph.run(run(default_ruleset)) + + egglog_string = egraph.as_egglog_string + assert "(let $__expr_0 (LetShadowDefaultNum___init__ 3))" in egglog_string + assert "(rewrite (LetShadowDefaultNum_make _0) (LetShadowDefaultNum___init__ _0)" in egglog_string + + +def test_save_egglog_string_defaults_to_disabled() -> None: + egraph = EGraph() + + with pytest.raises(ValueError, match="save_egglog_string=True"): + _ = egraph.as_egglog_string + assert egraph._state.egglog_file_state is None + + +def test_saved_egglog_transcript_close_is_idempotent() -> None: + egraph = EGraph(save_egglog_string=True) + assert egraph._state.egglog_file_state is not None + path = pathlib.Path(egraph._state.egglog_file_state.path) + assert path.exists() + + egraph.close() + egraph.close() + + assert not path.exists() + with pytest.raises(ValueError, match="has been closed"): + _ = egraph.as_egglog_string + + +def test_saved_egglog_transcript_is_removed_on_finalization() -> None: + egraph = EGraph(save_egglog_string=True) + assert egraph._state.egglog_file_state is not None + path = pathlib.Path(egraph._state.egglog_file_state.path) + assert path.exists() + + del egraph + gc.collect() + + assert not path.exists() + + +def test_saved_egglog_transcript_is_shared_across_push_and_pop() -> None: + egraph = EGraph(save_egglog_string=True) + transcript = egraph._state.egglog_file_state + assert transcript is not None + + egraph.push() + assert egraph._state.egglog_file_state is transcript + egraph.pop() + assert egraph._state.egglog_file_state is transcript + + egraph.close() + assert transcript.file.closed + + +def test_saved_egglog_string_uses_short_generated_sort_and_function_names() -> None: + class Num(Expr): + @classmethod + def var(cls, v: StringLike) -> Num: ... + + egraph = EGraph(save_egglog_string=True) + egraph.register(Num.var("x")) + egglog_string = egraph.as_egglog_string + + assert "(sort Num)" in egglog_string + assert "(constructor Num_var (String) Num)" in egglog_string + assert "test_high_level" not in egglog_string + + +def test_generated_names_fall_back_to_full_name_on_conflict() -> None: + state = EGraph(save_egglog_string=True)._state + ret1 = Ident("Ret", "pkg.one") + ret2 = Ident("Ret", "pkg.two") + fn1 = Ident("make", "pkg.one") + fn2 = Ident("make", "pkg.two") + state.__egg_decls__ |= Declarations( + _classes={ret1: ClassDecl(), ret2: ClassDecl()}, + _functions={ + fn1: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret1))), + fn2: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret2))), + }, + ) + + assert state.callable_ref_to_egg(FunctionRef(fn1))[0] == "make" + assert state.callable_ref_to_egg(FunctionRef(fn2))[0] == "pkg_two_make" + assert state.type_ref_to_egg(JustTypeRef(ret1)) == "Ret" + assert state.type_ref_to_egg(JustTypeRef(ret2)) == "pkg.two.Ret" + + +def test_missing_function_lookup_does_not_reserve_generated_name() -> None: + state = EGraph(save_egglog_string=True)._state + ret = Ident("LookupRet", "pkg.lookup") + fn = Ident("lookup_short_name", "pkg.lookup") + state.__egg_decls__ |= Declarations( + _classes={ret: ClassDecl()}, + _functions={fn: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret)))}, + ) + + assert list(state.possible_egglog_functions(["lookup_short_name"])) == [] + assert state.callable_ref_to_egg(FunctionRef(fn))[0] == "lookup_short_name" + + +def test_generated_names_fall_back_from_builtin_names() -> None: + state = EGraph(save_egglog_string=True)._state + ret = Ident("BuiltinConflictRet", "pkg.builtin_conflict") + fn = Ident("exp", "pkg.builtin_conflict") + sort = Ident("Map", "pkg.builtin_conflict") + state.__egg_decls__ |= Declarations( + _classes={ret: ClassDecl(), sort: ClassDecl()}, + _functions={fn: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret)))}, + ) + + assert state.callable_ref_to_egg(FunctionRef(fn))[0] == "pkg_builtin_conflict_exp" + assert state.type_ref_to_egg(JustTypeRef(sort)) == "pkg.builtin_conflict.Map" + + +def test_generated_callable_name_avoids_an_existing_cost_table() -> None: + state = EGraph(save_egglog_string=True)._state + state.__egg_decls__ |= cast("HasDeclarations", i64) + ret = Ident("CostRet", "pkg.cost") + fn = Ident("f", "pkg.cost") + conflict = Ident("cost_table_f", "pkg.cost") + state.__egg_decls__ |= Declarations( + _classes={ret: ClassDecl()}, + _functions={ + fn: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret))), + conflict: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret))), + }, + ) + + assert state.create_cost_table(FunctionRef(fn)) == "cost_table_f" + assert state.callable_ref_to_egg(FunctionRef(conflict))[0] == "pkg_cost_cost_table_f" + + +def test_generated_cost_table_name_avoids_an_existing_callable() -> None: + state = EGraph(save_egglog_string=True)._state + state.__egg_decls__ |= cast("HasDeclarations", i64) + ret = Ident("CostRet", "pkg.cost") + fn = Ident("f", "pkg.cost") + conflict = Ident("cost_table_f", "pkg.cost") + state.__egg_decls__ |= Declarations( + _classes={ret: ClassDecl()}, + _functions={ + fn: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret))), + conflict: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret))), + }, + ) + + assert state.callable_ref_to_egg(FunctionRef(conflict))[0] == "cost_table_f" + assert state.create_cost_table(FunctionRef(fn)) == "cost_table_f_1" + + +@pytest.mark.parametrize( + "body", + [ + pytest.param(None, id="bodyless-function"), + pytest.param( + TypedExprDecl(JustTypeRef(Ident.builtin("i64")), LitDecl(1)), + id="eager-primitive", + ), + ], +) +@pytest.mark.parametrize("sort_first", [True, False], ids=["sort-first", "callable-first"]) +def test_generated_names_share_the_backend_sort_and_callable_namespace( + body: TypedExprDecl | None, *, sort_first: bool +) -> None: + state = EGraph(save_egglog_string=True)._state + state.__egg_decls__ |= cast("HasDeclarations", i64) + sort_ident = Ident("Node", "pkg.sort") + fn_ident = Ident("Node", "pkg.fn") + sort_ref = JustTypeRef(sort_ident) + fn_ref = FunctionRef(fn_ident) + state.__egg_decls__ |= Declarations( + _classes={sort_ident: ClassDecl()}, + _functions={ + fn_ident: FunctionDecl( + signature=FunctionSignature(return_type=TypeRefWithVars(Ident.builtin("i64"))), + body=body, + ) + }, + ) + + if sort_first: + assert state.type_ref_to_egg(sort_ref) == "Node" + assert state.callable_ref_to_egg(fn_ref)[0] == "pkg_fn_Node" + else: + assert state.callable_ref_to_egg(fn_ref)[0] == "Node" + assert state.type_ref_to_egg(sort_ref) == "pkg.sort.Node" + + +def test_builtin_name_reservations_cover_builtins_module_declarations() -> None: + expected_fn_names = set[str]() + expected_sort_names = set[str]() + + def add_callable_name(decl: CallableDecl | None) -> None: + if decl is not None and decl.egg_name is not None: + expected_fn_names.add(decl.egg_name) + + for name in egg_builtins.__all__: + obj = getattr(egg_builtins, name, None) + if not isinstance(obj, HasDeclarations): + continue + decls = obj.__egg_decls__ + for decl in decls._functions.values(): + add_callable_name(decl) + for decl in decls._constants.values(): + add_callable_name(decl) + for decl in decls._classes.values(): + if decl.builtin and decl.egg_name is not None: + expected_sort_names.add(decl.egg_name) + add_callable_name(decl.init) + for callable_decl in ( + *decl.class_methods.values(), + *decl.class_variables.values(), + *decl.methods.values(), + *decl.properties.values(), + ): + add_callable_name(callable_decl) + + assert expected_fn_names <= BUILTIN_EGG_FN_NAMES + assert expected_sort_names <= BUILTIN_EGG_SORT_NAMES + + +def test_parameterized_sort_names_use_allocated_argument_names() -> None: + egraph = EGraph(save_egglog_string=True) + + egraph.register(Map[i64, BigRat].empty()) + + assert "(sort Map[i64,BigRat] (Map i64 BigRat))" in egraph.as_egglog_string + + +def test_non_constructor_map_empty_does_not_create_synthetic_let() -> None: + egraph = EGraph(save_egglog_string=True) + expr = Map[Map[String, i64], f64].empty() + runtime_expr = cast("RuntimeExpr", expr) + egraph._add_decls(runtime_expr) + + assert egraph._state._transform_let(runtime_expr.__egg_typed_expr__) is not None + + lines = egraph.as_egglog_string.splitlines() + assert "(let $__expr_0 (map-empty))" not in lines + assert not any(line.startswith("(fail ") for line in lines) + + +def test_non_constructor_maybe_none_does_not_create_synthetic_let() -> None: + egraph = EGraph(save_egglog_string=True) + expr = Maybe[Maybe[i64]].none() + runtime_expr = cast("RuntimeExpr", expr) + egraph._add_decls(runtime_expr) + + assert egraph._state._transform_let(runtime_expr.__egg_typed_expr__) is not None + + lines = egraph.as_egglog_string.splitlines() + assert "(let $__expr_0 (maybe-none))" not in lines + assert not any(line.startswith("(fail ") for line in lines) + + +def test_inferable_non_constructor_map_empty_does_not_create_synthetic_let() -> None: + egraph = EGraph(save_egglog_string=True) + expr = Map[String, i64].empty() + runtime_expr = cast("RuntimeExpr", expr) + egraph._add_decls(runtime_expr) + + assert egraph._state._transform_let(runtime_expr.__egg_typed_expr__) is not None + + lines = egraph.as_egglog_string.splitlines() + assert "(let $__expr_0 (map-empty))" not in lines + assert not any(line.startswith("(fail ") for line in lines) + + +def test_freeze_omits_synthetic_let_bindings() -> None: + class FreezeLetNum(Expr): + @classmethod + def var(cls, v: StringLike) -> FreezeLetNum: ... + + egraph = EGraph(save_egglog_string=True) + expr = FreezeLetNum.var("x") + runtime_expr = cast("RuntimeExpr", expr) + egraph._add_decls(runtime_expr) + egraph._state._transform_let(runtime_expr.__egg_typed_expr__) + + assert "(let $__expr_0 " in egraph.as_egglog_string + assert "$__expr_0" not in str(egraph.freeze()) + + +def test_popped_explicit_lets_do_not_block_synthetic_let_names() -> None: + class ScopedLetNum(Expr): + @classmethod + def var(cls, v: StringLike) -> ScopedLetNum: ... + + egraph = EGraph(save_egglog_string=True) + egraph.push() + egraph.let("__expr_0", ScopedLetNum.var("pushed")) + egraph.pop() + expr = ScopedLetNum.var("synthetic") + runtime_expr = cast("RuntimeExpr", expr) + egraph._add_decls(runtime_expr) + + egraph._state._transform_let(runtime_expr.__egg_typed_expr__) + + assert '(let $__expr_0 (ScopedLetNum_var "synthetic"))' in egraph.as_egglog_string + + +def test_registering_bare_variable_expression_raises() -> None: + egraph = EGraph() + + with pytest.raises(ValueError, match="must be calls"): + egraph.register(var("x", i64)) + + +def test_registering_let_reference_expression_raises() -> None: + egraph = EGraph() + x = egraph.let("x", i64(1)) + + with pytest.raises(ValueError, match="must be calls"): + egraph.register(x) + + +@pytest.mark.parametrize("save_egglog_string", [True, False]) +def test_nested_rule_lowering_does_not_reuse_top_level_synthetic_lets(save_egglog_string: bool) -> None: + class NestedRuleEdge(Expr): + @classmethod + def leaf(cls, value: StringLike) -> NestedRuleEdge: ... + + @classmethod + def pair(cls, left: NestedRuleEdge, right: NestedRuleEdge) -> NestedRuleEdge: ... + + egraph = EGraph(save_egglog_string=save_egglog_string) + done_rel = relation(f"done_rel_ctx_{int(save_egglog_string)}") + shared = NestedRuleEdge.leaf("shared") + pair = NestedRuleEdge.pair(shared, shared) + + egraph.register(pair) + egraph.register(rule(eq(var("x", NestedRuleEdge)).to(pair)).then(done_rel())) + egraph.register(subsume(shared)) + + egraph.run(1) + egraph.check_fail(done_rel()) + + if save_egglog_string: + assert '(= _x (NestedRuleEdge_pair (NestedRuleEdge_leaf "shared") (NestedRuleEdge_leaf "shared")))' in ( + egraph.as_egglog_string + ) + + +def test_top_level_action_factors_duplicate_sibling_edges() -> None: + class DuplicateEdge(Expr): + @classmethod + def leaf(cls, value: StringLike) -> DuplicateEdge: ... + + @classmethod + def pair(cls, left: DuplicateEdge, right: DuplicateEdge) -> DuplicateEdge: ... + + leaf = DuplicateEdge.leaf("shared") + first_pair = DuplicateEdge.pair(leaf, leaf) + second_pair = DuplicateEdge.pair(first_pair, first_pair) + egraph = EGraph(save_egglog_string=True) + + egraph.register(DuplicateEdge.pair(second_pair, second_pair)) + + transcript = egraph.as_egglog_string + assert transcript.count('(let $__expr_0 (DuplicateEdge_leaf "shared"))') == 1 + assert transcript.count("(let $__expr_1 (DuplicateEdge_pair $__expr_0 $__expr_0))") == 1 + assert transcript.count("(let $__expr_2 (DuplicateEdge_pair $__expr_1 $__expr_1))") == 1 + assert "(DuplicateEdge_pair $__expr_2 $__expr_2)" in transcript + + +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) + + egraph.run(combined) + + combined_line = next( + line for line in egraph.as_egglog_string.splitlines() if line.startswith("(unstable-combined-ruleset ") + ) + combined_name = combined_line.split()[1] + assert combined_name.startswith("_combined_ruleset_") + assert combined_name.removeprefix("_combined_ruleset_").isdigit() + assert f"(run-schedule (run {combined_name}))" in egraph.as_egglog_string + + +def test_integer_run_accepts_a_combined_ruleset() -> None: + seen = relation("combined_seen", i64) + copied = relation("combined_copied", i64) + x = var("x", i64) + copy_rules = ruleset(rule(seen(x)).then(copied(x))) + empty_rules = ruleset() + egraph = EGraph(seen(i64(1))) + + egraph.run(1, ruleset=copy_rules | empty_rules) + + egraph.check(copied(i64(1))) + + +def test_higher_order_builtin_callback_materializes_scalar_builtin_type_args() -> None: + check_eq(map_fold_kv(lambda acc, k, v: acc + v, f64(0.0), Map[i64, f64].empty()), f64(0.0)) + + +def test_higher_order_builtin_callback_materializes_parameterized_builtin_dummy_args() -> None: + input_map = Map[i64, Maybe[f64]].empty().insert(i64(1), Maybe[f64].some(f64(2.5))) + expected = Map[i64, f64].empty().insert(i64(1), f64(2.5)) + check_eq(map_map_values(lambda k, v: v.unwrap(), input_map), expected) + + +def test_higher_order_builtin_callback_materializes_rational_builtin_dummy_args() -> None: + input_map = Map[i64, Rational].empty().insert(i64(1), Rational(1, 2)) + expected = Map[i64, f64].empty().insert(i64(1), f64(1.5)) + check_eq(map_map_values(lambda k, v: v.to_f64() + 1.0, input_map), expected) + + +def test_map_map_values_generic_negation_callback_is_concretized() -> None: + input_map = Map[i64, f64].empty().insert(i64(1), f64(2.5)).insert(i64(2), f64(-4.0)) + expected = Map[i64, f64].empty().insert(i64(1), f64(-2.5)).insert(i64(2), f64(4.0)) + + check_eq(map_map_values(lambda _key, value: -value, input_map), expected) + + +def test_fold_derived_map_helpers() -> None: + left = Map[i64, f64].empty().insert(i64(1), f64(2.0)).insert(i64(2), f64(3.0)) + right = Map[i64, f64].empty().insert(i64(2), f64(4.0)).insert(i64(3), f64(5.0)) + + check_eq(map_filter_kv(lambda key, _value: key > 1, left), Map[i64, f64].empty().insert(i64(2), f64(3.0))) + check_eq( + map_merge_with(lambda old, new: old + new, left, right), + Map[i64, f64].empty().insert(i64(1), f64(2.0)).insert(i64(2), f64(7.0)).insert(i64(3), f64(5.0)), + ) + + check_eq(left.length(), i64(2)) + assert EGraph().extract(left.keys()).value == {i64(1), i64(2)} + check_eq(left.keys().length(), i64(2)) + EGraph().check(left.contains(left.pick_key())) + + +def test_higher_order_callable_inference_does_not_mutate_ambient_ruleset() -> None: + ambient = ruleset(name="hof-inference-ambient") + initial_rules = tuple(ambient.__egg_ruleset__.rules) + + with set_current_ruleset(ambient): + expr = map_map_values(lambda _key, value: -value, Map[i64, f64].empty().insert(i64(1), f64(2.0))) + _ = cast("RuntimeExpr", expr).__egg_decls__ + + assert tuple(ambient.__egg_ruleset__.rules) == initial_rules + + +def test_set_current_ruleset_restores_nested_contexts() -> None: + outer = ruleset(name="current-ruleset-outer") + inner = ruleset(name="current-ruleset-inner") + initial = get_current_ruleset() + + with set_current_ruleset(outer): + assert get_current_ruleset() is outer + with set_current_ruleset(inner): + assert get_current_ruleset() is inner + assert get_current_ruleset() is outer + + assert get_current_ruleset() is initial + + +def test_file_backed_errors_report_saved_file_line() -> None: + egraph = EGraph(save_egglog_string=True) + egraph.let("x", i64(1)) + egraph.let("y", i64(2)) + expected_line = len(egraph.as_egglog_string.splitlines()) + 1 + assert egraph._state.egglog_file_state is not None + path = egraph._state.egglog_file_state.path + + with pytest.raises(EggSmolError) as exc_info: + egraph.check(eq(i64(1)).to(i64(2))) + + error_text = exc_info.value.context + assert path in error_text + assert f"In {expected_line}:" in error_text + lines = egraph.as_egglog_string.splitlines() + assert "(fail (check (= 1 2))) ; Check failed:" in lines + assert "(check (= 1 2))" not in lines + + +def test_unnamed_lambda_returning_builtin_is_eager() -> None: + check_eq( + map_fold_kv( + lambda acc, k, v: acc + v, + f64(0.0), + Map[i64, f64].empty().insert(i64(1), f64(2.0)).insert(i64(2), f64(3.5)), + ), + f64(5.5), + ) + + +def test_unnamed_lambda_returning_eqsort_is_eager() -> None: + class Box(Expr): + def __init__(self, value: i64Like) -> None: ... + + expected = Map[i64, Box].empty().insert(i64(1), Box(i64(2))) + actual = cast("Map[i64, Box]", map_map_values(lambda k, v: Box(v), Map[i64, i64].empty().insert(i64(1), i64(2)))) + check_eq(cast("BaseExpr", actual), cast("BaseExpr", expected)) + + +def test_named_builtin_return_function_is_eager() -> None: + @function + def add1(x: i64Like) -> i64: + return cast("i64", x) + 1 + + check_eq(add1(i64(2)), i64(3)) + + +def test_named_container_return_function_is_eager() -> None: + @function + def singleton_map(x: i64Like) -> Map[i64, i64]: + runtime_x = cast("i64", x) + return Map[i64, i64].empty().insert(runtime_x, runtime_x + 1) + + expected = Map[i64, i64].empty().insert(i64(2), i64(3)) + check_eq(singleton_map(i64(2)), expected) + + +def test_reverse_args_eager_method_preserves_python_argument_order() -> None: + class ReverseBody(Expr): + def __init__(self, value: i64Like) -> None: ... + + @method(reverse_args=True) + def ordered(self, other: i64Like) -> Pair[ReverseBody, i64]: + return Pair(self, cast("i64", other)) + + extracted = EGraph().extract(ReverseBody(1).ordered(i64(2))) + + assert extracted.value == (ReverseBody(1), i64(2)) + + +def test_reverse_args_rewrite_method_preserves_python_argument_order() -> None: + reverse_ruleset = ruleset(name="reverse-args-rewrite") + + class ReverseRewrite(Expr, ruleset=reverse_ruleset): + def __init__(self, value: i64Like) -> None: ... + + @method(reverse_args=True) + def select_self(self, other: ReverseRewrite) -> ReverseRewrite: + return self + + expr = ReverseRewrite(1).select_self(ReverseRewrite(2)) + egraph = EGraph(expr) + + egraph.run(reverse_ruleset) + egraph.check(eq(expr).to(ReverseRewrite(1))) + + +def test_reverse_args_bodyless_method_uses_backend_argument_order() -> None: + class ReverseLookup(Expr): + def __init__(self, value: i64Like) -> None: ... + + @method(reverse_args=True) + def lookup(self, key: i64Like) -> String: + return None # type: ignore[return-value] # None denotes a bodyless symbolic function. + + value = ReverseLookup(1).lookup(i64(2)) + egraph = EGraph() + egraph.register(set_(value).to(String("found")), set_cost(value, 7)) + + egraph.check(eq(value).to(String("found")), eq(get_cost(value)).to(i64(7))) + assert egraph.function_values(ReverseLookup.lookup) == {value: String("found")} + + +def test_reverse_args_bodyless_constructor_round_trips_backend_argument_order() -> None: + class ReverseResult(Expr): ... + + class ReverseSource(Expr): + def __init__(self, value: i64Like) -> None: ... + + @method(reverse_args=True) + def make(self, key: i64Like) -> ReverseResult: ... + + expr = ReverseSource(1).make(i64(2)) + + assert expr_parts(EGraph(expr).extract(expr)) == expr_parts(expr) + + +def test_named_builtin_return_function_with_eqsort_input_is_eager() -> None: + class Box(Expr): + def __init__(self, value: i64Like) -> None: ... + + @function + def box_score(box: Box) -> i64: ... + + @function + def via_box(box: Box) -> i64: + return box_score(box) + 1 + + egraph = EGraph() + egraph.register(set_(box_score(Box(i64(4)))).to(i64(9))) + egraph.check(eq(via_box(Box(i64(4)))).to(i64(10))) + + +def test_named_builtin_return_function_with_none_body_stays_plain_function() -> None: + @function + def maybe_missing(x: i64Like) -> i64: + return None # type: ignore[return-value] # None denotes a bodyless symbolic function. + + egraph = EGraph() + egraph.check_fail(eq(maybe_missing(i64(1))).to(i64(1))) + egraph.register(set_(maybe_missing(i64(1))).to(i64(4))) + egraph.check(eq(maybe_missing(i64(1))).to(i64(4))) + + +def test_named_eqsort_function_body_is_eager() -> None: + class Box(Expr): + def __init__(self, value: i64Like) -> None: ... + + @function + def make_box(x: i64Like) -> Box: + return Box(x) + + check_eq(make_box(i64(4)), Box(i64(4))) + + +def test_mutating_eqsort_function_body_can_construct_a_bound_result() -> None: + class IntBox(Expr): + def __init__(self, value: i64Like) -> None: ... + + def __add__(self, other: IntBox) -> IntBox: ... + + @function(mutates_first_arg=True) + def increment(box: IntBox) -> None: + box.__replace_expr__(box + IntBox(1)) + + box = IntBox(10) + increment(box) + egraph = EGraph() + incremented = egraph.let("incremented", box) + egraph.check(eq(incremented).to(IntBox(10) + IntBox(1))) + + +def test_missing_function_row_inside_primitive_body_stays_undefined() -> None: + @function + def f_lookup(x: i64Like) -> i64: ... + + @function + def via_lookup(x: i64Like) -> i64: + return f_lookup(x) + 1 + + egraph = EGraph() + egraph.register(set_(f_lookup(i64(4))).to(i64(9))) + egraph.check(eq(via_lookup(i64(4))).to(i64(10))) + egraph.check_fail(eq(via_lookup(i64(5))).to(i64(0))) + + def test_fib(): egraph = EGraph() @@ -297,6 +1072,24 @@ def test_convert_int_float(): egraph = EGraph() egraph.check(eq(i64(1)).to(f64(1.0).to_i64())) egraph.check(eq(f64(1.0)).to(f64.from_i64(i64(1)))) + assert egraph.extract(f64(2.0) + 1).value == 3.0 + + +def test_f64_math_primitives() -> None: + egraph = EGraph() + assert egraph.extract(f64(1.0).exp()).value == pytest.approx(math.e) + assert egraph.extract(f64(math.e).log()).value == pytest.approx(1.0) + assert egraph.extract(f64(4.0).sqrt()).value == pytest.approx(2.0) + + +def test_bigrat_to_i64_is_exact_and_bounded() -> None: + assert EGraph().extract(BigRat(4, 2).to_i64()) == i64(2) + assert EGraph().extract(BigRat(1, 2) + i64(1)).value == Fraction(3, 2) + + for numerator, denominator in [(1, 2), (2**63, 1), (-(2**63) - 1, 1)]: + value = BigRat(BigInt.from_string(str(numerator)), BigInt(denominator)) + with pytest.raises(EggSmolError, match="primitive to-i64 failed"): + EGraph().extract(value.to_i64()) def test_f64_negation() -> None: @@ -495,6 +1288,11 @@ def test_map(self): assert String("a") in m assert String("c") not in m + def test_map_duplicate_key_uses_latest_value(self): + m = Map[String, i64].empty().insert(String("a"), i64(1)).insert(String("a"), i64(2)) + + assert EGraph().extract(m).value == {String("a"): i64(2)} + def test_set(self): assert EGraph().extract(Set[i64].empty()).value == set() s = Set(i64(1), i64(2)) @@ -504,6 +1302,7 @@ def test_set(self): assert len(s) == 2 assert i64(1) in s assert i64(3) not in s + assert list(Set(i64(1), i64(1))) == [i64(1)] def test_rational(self): assert Rational(1, 2).value == Fraction(1, 2) @@ -533,6 +1332,15 @@ def test_big_rat(self): assert float(br) == 1 / 2 assert br.value == Fraction(1, 2) + def test_extract_nested_maps_preserves_empty_map_type_params(self): + inner = Map[String, BigRat].empty().insert(String("x"), BigRat(2, 1)) + expr = Map[Map[String, BigRat], f64].empty().insert(inner, f64(1.0)) + + extracted = EGraph().extract(expr) + + assert "Map[String, BigRat].empty().insert" in str(extracted) + assert 'Map[Map[String, BigRat], f64].empty().insert(String("x")' not in str(extracted) + def test_multiset(self): assert list(MultiSet(i64(1), i64(1))) == [i64(1), i64(1)] @@ -680,12 +1488,71 @@ def __init__(self) -> None: ... class TestDefaultReplacements: + def test_builtin_function_without_body(self): + @function(builtin=True) + def f(x: i64Like) -> i64: ... + + assert expr_parts(f(1)) == expr_parts(f(i64(1))) + + def test_eqsort_merge_function_without_body(self): + @function(merge=lambda old, new: old) + def f() -> A: ... + + egraph = EGraph() + egraph.register(set_(f()).to(A())) + egraph.check(eq(f()).to(A())) + + def test_primitive_constant_with_merge(self): + best = constant("best", i64, merge=lambda old, new: old.max(new)) + + egraph = EGraph() + egraph.register(set_(best).to(i64(1)), set_(best).to(i64(2))) + + egraph.check(eq(best).to(i64(2))) + + def test_none_is_a_primitive_constant_default(self): + missing = constant("missing_default", Maybe[i64], None) + + check_eq(missing, Maybe[i64].none()) + + def test_none_is_a_primitive_class_variable_default(self): + class Defaults(Expr): + missing: ClassVar[Maybe[i64] | None] = None + + check_eq(cast("Maybe[i64]", Defaults.missing), Maybe[i64].none()) + + def test_bodyless_primitive_constant_is_not_a_synthetic_let(self): + value = constant("bodyless_primitive", i64) + egraph = EGraph(save_egglog_string=True) + egraph._add_decls(cast("RuntimeExpr", value)) + + assert egraph._state._transform_let(cast("RuntimeExpr", value).__egg_typed_expr__) is not None + assert "(let $__expr_0 bodyless_primitive)" not in egraph.as_egglog_string + + def test_bodyless_eqsort_constant_is_a_synthetic_let(self): + value = constant("bodyless_eqsort", A) + egraph = EGraph(save_egglog_string=True) + egraph._add_decls(cast("RuntimeExpr", value)) + + assert egraph._state._transform_let(cast("RuntimeExpr", value).__egg_typed_expr__) is None + assert "(let $__expr_0 (%bodyless_eqsort))" in egraph.as_egglog_string + + def test_eqsort_constant_with_merge(self): + merged = constant("merged", A, merge=lambda old, _new: old) + + egraph = EGraph() + egraph.register(set_(merged).to(A())) + + egraph.check(eq(merged).to(A())) + assert egraph.function_values(merged) == {merged: A()} + assert "set_(merged).to(A())" in str(egraph.freeze()) + def test_function(self): @function def f() -> A: return A() - check_eq(f(), A(), run()) + check_eq(f(), A()) def test_function_ruleset(self): r = ruleset() @@ -696,6 +1563,15 @@ def f() -> A: check_eq(f(), A(), r) + def test_function_ruleset_with_subsume(self): + r = ruleset() + + @function(ruleset=r, subsume=True) + def f() -> A: + return A() + + check_eq(f(), A(), r) + def test_function_ruleset_can_run_after_materialization_without_registration(self): r = ruleset() @@ -711,7 +1587,7 @@ def f() -> A: def test_constant(self): a = constant("a", A, A()) - check_eq(a, A(), run()) + check_eq(a, A()) def test_constant_ruleset(self): r = ruleset() @@ -725,7 +1601,7 @@ def __init__(self) -> None: ... def f(self) -> A: return A() - check_eq(B().f(), A(), run()) + check_eq(B().f(), A()) def test_method_ruleset(self): r = ruleset() @@ -743,7 +1619,27 @@ class B(Expr): def f(cls) -> A: return A() - check_eq(B.f(), A(), run()) + check_eq(B.f(), A()) + + def test_property(self): + class B(Expr): + def __init__(self, value: i64Like) -> None: ... + + @property + def a(self) -> A: + return A() + + check_eq(B(i64(1)).a, A()) + + 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 + + @classmethod + def wrap(cls, value: i64Like) -> B: ... + + check_eq(B(i64(1)), B.wrap(i64(1))) def test_classmethod_ruleset(self): r = ruleset() @@ -759,7 +1655,7 @@ def test_classvar(self): class B(Expr): a: ClassVar[A] = A() - check_eq(B.a, A(), run()) + check_eq(B.a, A()) def test_classvar_ruleset(self): r = ruleset() @@ -769,6 +1665,20 @@ class B(Expr, ruleset=r): check_eq(B.a, A(), r) + def test_constructor_unextractable(self): + class B(Expr): + def __init__(self, value: i64Like) -> None: ... + + @method(unextractable=True) + def opaque(self) -> B: ... + + def __add__(self, other: B) -> B: ... + + egraph = EGraph() + opaque = egraph.let("opaque", B(i64(1)).opaque()) + egraph.register(union(opaque).with_(B(i64(1)) + B(i64(1)))) + assert expr_parts(egraph.extract(opaque)) == expr_parts(B(i64(1)) + B(i64(1))) + def test_method_refer_to_later(self): """ Verify that an earlier method body can refer to values defined in later ones @@ -784,7 +1694,7 @@ def g(self) -> A: ... B() left = B().f() right = B().g() - check_eq(left, right, run()) + check_eq(left, right) def test_classmethod_own_class(self): class B(Expr): @@ -793,7 +1703,7 @@ def __init__(self) -> None: ... def f(cls) -> B: return B() - check_eq(B.f(), B(), run()) + check_eq(B.f(), B()) class TestIssue166: @@ -824,6 +1734,173 @@ def __init__(self) -> None: ... E() +class TestCallableValidation: + def test_primitive_function_ruleset_subsume_rejected(self): + r = ruleset() + + @function(ruleset=r, subsume=True) # type: ignore[type-var] # Deliberately invalid runtime API call. + def f() -> i64: + return i64(1) + + with pytest.raises(ValueError, match="Primitive-returning callables cannot use subsume"): + f() + + def test_no_body_function_cannot_use_explicit_ruleset(self): + r = ruleset() + + @function(ruleset=r) + def f() -> A: ... + + with pytest.raises(ValueError, match="Explicit rulesets require a body"): + f() + + def test_constant_without_default_cannot_use_explicit_ruleset(self): + r = ruleset() + + with pytest.raises(ValueError, match="Explicit rulesets require a default"): + EGraph().register(constant("no_default", A, ruleset=r)) # type: ignore[call-overload] + + def test_primitive_constant_default_cannot_use_explicit_ruleset(self): + r = ruleset() + + with pytest.raises(ValueError, match="Primitive-returning defaults cannot use an explicit ruleset"): + EGraph().register( + constant("primitive_default", i64, i64(1), ruleset=r) # type: ignore[call-overload] + ) + + def test_eqsort_constant_default_cannot_use_merge(self): + with pytest.raises(ValueError, match="Eqsort-returning callables with bodies cannot use merge"): + EGraph().register( + constant("default_merge", A, A(), merge=lambda old, _new: old) # type: ignore[call-overload] + ) + + def test_primitive_constant_default_cannot_use_merge(self): + with pytest.raises(ValueError, match="Primitive-returning callables with bodies cannot use merge"): + EGraph().register( + constant( # type: ignore[call-overload] + "primitive_default_merge", i64, i64(1), merge=lambda old, new: old.max(new) + ) + ) + + def test_unit_constant_cannot_use_merge(self): + with pytest.raises(ValueError, match="Functions that return Unit cannot use merge"): + EGraph().register(constant("unit_merge", Unit, merge=lambda old, _new: old)) + + def test_eqsort_eager_body_cannot_use_merge(self): + @function(merge=lambda old, new: old) + def f() -> A: + return A() + + with pytest.raises(ValueError, match="Eqsort-returning callables with bodies cannot use merge"): + f() + + def test_primitive_returning_functions_cannot_use_cost(self): + @function(cost=1) # type: ignore[type-var] # Deliberately invalid runtime API call. + def f() -> i64: ... + + with pytest.raises(ValueError, match="Primitive-returning callables cannot use cost"): + f() + + def test_primitive_returning_functions_cannot_be_unextractable(self): + @function(unextractable=True) # type: ignore[type-var] # Deliberately invalid runtime API call. + def f() -> i64: ... + + with pytest.raises(ValueError, match="Primitive-returning callables cannot be unextractable"): + f() + + def test_builtin_callables_cannot_use_merge(self): + @function(builtin=True, merge=lambda old, new: old) # type: ignore[call-overload] + def f() -> i64: ... + + with pytest.raises(ValueError, match="Builtin callables cannot use merge"): + f() + + def test_primitive_body_cannot_use_builtin(self): + @function(builtin=True) + def f() -> i64: + return i64(1) + + with pytest.raises(ValueError, match="Builtin callables cannot have a body"): + f() + + def test_primitive_body_cannot_use_merge(self): + @function(merge=lambda old, new: old) + def f() -> i64: + return i64(1) + + with pytest.raises(ValueError, match="Primitive-returning callables with bodies cannot use merge"): + f() + + def test_primitive_body_cannot_use_explicit_ruleset(self): + r = ruleset() + + @function(ruleset=r) # type: ignore[type-var] # Deliberately invalid runtime API call. + def f() -> i64: + return i64(1) + + with pytest.raises( + ValueError, match="Primitive-returning callables with bodies cannot use an explicit ruleset" + ): + f() + + def test_eqsort_body_cannot_use_merge(self): + r = ruleset() + + @function(ruleset=r, merge=lambda old, new: old) # type: ignore[call-overload] + def f() -> A: + return A() + + with pytest.raises(ValueError, match="Eqsort-returning callables with bodies cannot use merge"): + f() + + def test_eqsort_eager_body_cannot_use_cost(self): + @function(cost=1) + def f() -> A: + return A() + + with pytest.raises(ValueError, match="Eqsort-returning eager bodies cannot use cost"): + f() + + def test_eqsort_eager_body_cannot_be_unextractable(self): + @function(unextractable=True) + def f() -> A: + return A() + + with pytest.raises(ValueError, match="Eqsort-returning eager bodies cannot be unextractable"): + f() + + def test_no_body_function_cannot_use_subsume(self): + @function(subsume=True) + def f() -> A: ... + + with pytest.raises(ValueError, match="subsume requires an explicit ruleset"): + f() + + def test_primitive_method_subsume_rejected(self): + r = ruleset() + + class B(Expr, ruleset=r): + def __init__(self, value: i64Like) -> None: ... + + @method(subsume=True) # type: ignore[type-var] # Deliberately invalid runtime API call. + def f(self) -> i64: + return i64(1) + + with pytest.raises(ValueError, match="Primitive-returning callables cannot use subsume"): + B(i64(0)).f() + + def test_primitive_classvar_default_cannot_use_explicit_ruleset(self): + r = ruleset() + + class B(Expr, ruleset=r): + a: ClassVar[i64] = i64(1) + + def __init__(self) -> None: ... + + with pytest.raises(ValueError, match="Primitive-returning defaults cannot use an explicit ruleset"): + _ = B.a + + def test_vec_like_conversion(): """ Test that we can use a generic type alias for conversion @@ -852,6 +1929,79 @@ def my_fn(xs: MapLike[i64, String, i64Like, StringLike]) -> Unit: ... assert expr_parts(my_fn({})) == expr_parts(my_fn(Map[i64, String].empty())) +def test_maybe_builtin_surface(): + none_expr = EGraph().extract(Maybe[f64].none()) + assert none_expr.value is None + + some_expr = EGraph().extract(Maybe[f64].some(1.0)) # type: ignore[arg-type] # Runtime conversion. + assert some_expr.value is not None + assert some_expr.value.value == 1.0 + + assert EGraph().extract(Maybe[f64].some(1.0).unwrap()).value == 1.0 # type: ignore[arg-type] + assert EGraph().extract(Maybe[f64].none().unwrap_or(2.5)).value == 2.5 # type: ignore[arg-type] + + +def test_higher_order_maybe_pair_and_catch_builtins(): + assert EGraph().extract(Maybe[i64].some(2).match(lambda x: x + 3, i64(0))).value == 5 # type: ignore[arg-type] + assert EGraph().extract(Maybe[i64].none().match(lambda x: x + 3, i64(7))).value == 7 + + matched = EGraph().extract(Pair(i64(2), i64(3)).match(lambda left, right: left + right)) + assert matched.value == 5 + + mapped_left = EGraph().extract(Pair(i64(2), i64(3)).map_left(lambda left: left + 10)) + left, right = mapped_left.value + assert left.value == 12 + assert right.value == 3 + + mapped_right = EGraph().extract(Pair(i64(2), i64(3)).map_right(lambda right: right + 10)) + left, right = mapped_right.value + assert left.value == 2 + assert right.value == 13 + + caught = EGraph().extract(catch(lambda: Maybe[i64].some(4).unwrap())) # type: ignore[arg-type] + assert caught.value is not None + assert caught.value.value == 4 + + failed = EGraph().extract(catch(lambda: Maybe[i64].none().unwrap())) + assert failed.value is None + + +def test_nested_catch_match_with_different_inner_lambda_result_sort() -> None: + expr = catch(lambda: i64(1)).match( + lambda _: catch(lambda: f64(2.0)).match(lambda v: v, f64(0.0)), + f64(9.0), + ) + + assert EGraph().extract(expr).value == 2.0 + + +def test_maybe_conversion(): + @function + def maybe_identity(x: Maybe[i64]) -> Maybe[i64]: ... + + assert expr_parts(maybe_identity(None)) == expr_parts( # type: ignore[arg-type] # Runtime conversion. + maybe_identity(Maybe[i64].none()) + ) + + +def test_pair_conversion() -> None: + @function + def pair_identity(pair: Pair[i64, i64]) -> Pair[i64, i64]: ... + + assert expr_parts(pair_identity((1, 2))) == expr_parts( # type: ignore[arg-type] # Runtime conversion. + pair_identity(Pair(i64(1), i64(2))) + ) + + +@pytest.mark.parametrize("value", [(), (1,), (1, 2, 3)]) +def test_pair_conversion_requires_exactly_two_items(value: tuple[int, ...]) -> None: + @function + def pair_identity(pair: Pair[i64, i64]) -> Pair[i64, i64]: ... + + with pytest.raises(ValueError, match=rf"tuple of length 2.*length {len(value)}"): + pair_identity(value) # type: ignore[arg-type] # Deliberately malformed runtime conversion input. + + class TestEqNE: def test_eq(self): assert i64(3) == i64(3) @@ -917,7 +2067,7 @@ def __init__(self) -> None: ... case A(): pass case _: - msg = "Should have matched A" + msg = "Should have matched A" # type: ignore[unreachable] raise ValueError(msg) def test_literal(self): @@ -925,7 +2075,7 @@ def test_literal(self): case i64(i): assert i == 10 case _: - msg = "Should have matched i64(10)" + msg = "Should have matched i64(10)" # type: ignore[unreachable] raise ValueError(msg) def test_literal_fail(self): @@ -958,7 +2108,7 @@ def b(self) -> str: assert a == 1 assert b == "hi" case _: - msg = "Should have matched A" + msg = "Should have matched A" # type: ignore[unreachable] raise ValueError(msg) def test_custom_args_fail(self): @@ -1113,6 +2263,29 @@ def f(x: i64Like) -> i64: ... assert values == {f(i64(1)): i64(2)} +def test_table_inspection_rejects_eager_primitives() -> None: + @function + def eager_plus_one(x: i64Like) -> i64: + return cast("i64", x) + 1 + + egraph = EGraph() + assert egraph.extract(eager_plus_one(i64(1))) == i64(2) + + with pytest.raises(ValueError, match="table-backed"): + egraph.function_size(eager_plus_one) + with pytest.raises(ValueError, match="table-backed"): + egraph.function_values(eager_plus_one) + with pytest.raises(ValueError, match="table-backed"): + egraph.lookup_function_value(eager_plus_one(i64(1))) + + @function + def eager_text() -> String: + return String("value") + + with pytest.raises(ValueError, match="table-backed"): + egraph.input(eager_text, "unused.csv") + + def test_dynamic_cost(): """ https://github.com/egraphs-good/egglog-experimental/blob/6d07a34ac76deec751f86f70d9b9358cd3e236ca/tests/integration_test.rs#L5-L35 @@ -1200,7 +2373,9 @@ def test_sequence_repeat_saturate(self): def test_backoff_scheduler(self): """ Passing `scheduler=...` to run(...) hoists the scheduler to the - outer scope. This is equivalent to an explicit outer `bo.scope(...)`. + outer scope. This is equivalent to an explicit outer `bo.scope(...)` + around the whole repeated schedule. Scoping only one repetition creates + fresh scheduler state for each repeat. https://egraphs.zulipchat.com/#narrow/channel/375765-egg.2Fegglog/topic/.E2.9C.94.20Backoff.20Scheduler.20Example/with/538745863 """ @@ -1209,23 +2384,113 @@ def test_backoff_scheduler(self): grow = ruleset(rule(includes(x)).then(includes(x + 1))) shrink = ruleset(rule(includes(x)).then(includes(x - 1))) - e1 = EGraph() - e1.register(includes(i64(0))) - # default scheduler - with e1: - e1.run((grow + shrink) * 3) - e1.check(includes(i64(3)), includes(i64(-3))) - # back-off implicit outer hoisting bo = back_off(match_limit=1) - with e1: - e1.run((run(grow, scheduler=bo) + shrink) * 3) - e1.check(includes(i64(2)), includes(i64(-3))) - e1.check_fail(includes(i64(3))) - # back off inner hoisting - with e1: - e1.run(bo.scope(run(grow, scheduler=bo) + shrink) * 3) - e1.check(includes(i64(1)), includes(i64(-3))) - e1.check_fail(includes(i64(2))) + + def _run_and_collect(schedule: Schedule) -> set[int]: + egraph = EGraph() + egraph.register(includes(i64(0))) + with egraph: + egraph.run(schedule) + values = set() + for i in range(-3, 4): + try: + egraph.check(includes(i64(i))) + values.add(i) + except EggSmolError: + pass + return values + + default_values = _run_and_collect((grow + shrink) * 3) + assert default_values == {-3, -2, -1, 0, 1, 2, 3} + + implicit_values = _run_and_collect((run(grow, scheduler=bo) + shrink) * 3) + explicit_outer_values = _run_and_collect(bo.scope((run(grow, scheduler=bo) + shrink) * 3)) + explicit_inner_values = _run_and_collect(bo.scope(run(grow, scheduler=bo) + shrink) * 3) + + assert implicit_values == explicit_outer_values == {-3, -2, -1, 0, 1, 2} + assert explicit_inner_values == {-3, -2, -1, 0, 1} + + def test_persistent_scheduler_reuses_state_across_runs(self): + r = relation("R", i64) + s = relation("S", i64) + seed = relation("Seed") + x = var("x", i64) + + copy = ruleset(rule(r(x)).then(s(x)), name="copy") + grow = ruleset(rule(seed()).then(r(i64(3))), name="grow") + + def _run_and_collect() -> set[int]: + egraph = EGraph() + egraph.register(r(i64(0)), r(i64(1)), r(i64(2)), seed()) + scheduler = back_off(match_limit=2, ban_length=2).persistent() + egraph.run(run(copy, scheduler=scheduler)) + egraph.push() + egraph.pop() + egraph.run(run(grow)) + egraph.run(run(copy, scheduler=scheduler)) + values = set() + for i in range(4): + try: + egraph.check(s(i64(i))) + values.add(i) + except EggSmolError: + pass + return values + + assert _run_and_collect() == {0, 1, 2} + + def test_saturate_waits_for_deferred_persistent_scheduler_work(self): + source = relation("saturate_source", i64) + copied = relation("saturate_copied", i64) + x = var("x", i64) + copy = ruleset(rule(source(x)).then(copied(x)), name="saturate-copy") + + egraph = EGraph() + egraph.register(*(source(i64(i)) for i in range(3))) + scheduler = back_off(match_limit=1, ban_length=1).persistent() + + egraph.saturate(run(copy, scheduler=scheduler), max=4, visualize=False) + + egraph.check(*(copied(i64(i)) for i in range(3))) + + def test_persistent_scheduler_is_saved_once_across_runs(self): + r = relation("R_saved_scheduler", i64) + s = relation("S_saved_scheduler", i64) + x = var("x", i64) + copy = ruleset(rule(r(x)).then(s(x)), name="copy-saved-scheduler") + + egraph = EGraph(save_egglog_string=True) + egraph.register(r(i64(0)), r(i64(1))) + scheduler = back_off(match_limit=2, ban_length=2).persistent() + + egraph.run(run(copy, scheduler=scheduler)) + egraph.run(run(copy, scheduler=scheduler)) + + scheduler_lines = [line for line in egraph.as_egglog_string.splitlines() if line.startswith("(let-scheduler ")] + run_with_lines = [line for line in egraph.as_egglog_string.splitlines() if "(run-with " in line] + + assert len(scheduler_lines) == 1 + assert len(run_with_lines) == 2 + + def test_persistent_scheduler_gets_a_fresh_identity(self): + scheduler = back_off(match_limit=2, ban_length=2) + + persistent = scheduler.persistent() + + assert persistent.scheduler.id != scheduler.scheduler.id + + def test_scheduler_scope_does_not_leak_to_sequence_sibling(self): + r = ruleset(name="scheduler-lexical-scope") + scheduler = back_off(match_limit=2, ban_length=2) + egraph = EGraph(save_egglog_string=True) + + egraph.run(seq(scheduler.scope(run(r, scheduler=scheduler)), run(r, scheduler=scheduler))) + + run_schedule = next(line for line in egraph.as_egglog_string.splitlines() if line.startswith("(run-schedule ")) + assert run_schedule.count("(let-scheduler ") == 2 + ruleset_name = str(r.__egg_ident__) + assert f"(run-with _scheduler_0 {ruleset_name})" in run_schedule + assert f"(run-with _scheduler_1 {ruleset_name})" in run_schedule def test_custom_scheduler_invalid_until(self): """ @@ -1247,6 +2512,8 @@ def test_custom_scheduler_invalid_until(self): with pytest.raises(ValueError, match="Can only have one until fact with custom scheduler"): egraph.run(run(r, rel(i64(0)), rel(i64(1)), scheduler=bo)) + egraph.run(run(r, rel(i64(0)), scheduler=bo)) + @function def ff(x: i64Like, y: i64Like) -> E: ... @@ -1257,6 +2524,51 @@ def gg() -> E: ... class TestCustomExtract: + def test_literal_root(self) -> None: + def is_even_cost_model(egraph: EGraph, expr: BaseExpr, children_costs: list[int]) -> int: + del egraph, children_costs + return int(isinstance(expr, i64) and expr.value % 2 == 0) + + assert EGraph().extract(i64(10), include_cost=True, cost_model=is_even_cost_model) == (i64(10), 1) + assert EGraph().extract(i64(5), include_cost=True, cost_model=is_even_cost_model) == (i64(5), 0) + + @staticmethod + def _capture_container_children_costs( + root_expr: BaseExpr, + *, + leaf_cost: Callable[[BaseExpr], int], + should_capture: Callable[[BaseExpr, list[int]], bool], + ) -> tuple[BaseExpr, BaseExpr, list[int]]: + seen: dict[str, object] = {} + + def my_cost_model(egraph: EGraph, expr: BaseExpr, children_costs: list[int]) -> int: + match expr: + case i64() | String(): + return leaf_cost(expr) + case _: + if should_capture(expr, children_costs): + seen["expr"] = expr + seen["children_costs"] = children_costs.copy() + return default_cost_model(egraph, expr, children_costs) + + extracted, _cost = EGraph().extract(root_expr, include_cost=True, cost_model=my_cost_model) + + assert "expr" in seen + assert "children_costs" in seen + + return extracted, cast("BaseExpr", seen["expr"]), cast("list[int]", seen["children_costs"]) + + @staticmethod + def _small_leaf_cost(expr: BaseExpr) -> int: + match expr: + case i64(): + return {1: 11, 2: 22, 3: 33, 4: 44}[expr.value] + case String(): + return {"a": 101, "b": 202}[expr.value] + case _: + msg = f"Unexpected leaf {expr!r}" + raise AssertionError(msg) + @pytest.mark.parametrize( "expr", [ @@ -1272,6 +2584,9 @@ class TestCustomExtract: pytest.param(Set(i64(1), i64(2)), id="Set"), pytest.param(Map[i64, String].empty().insert(i64(1), String("hi")), id="Map"), pytest.param(MultiSet(i64(1), i64(1)), id="MultiSet"), + pytest.param(Pair(i64(1), String("hi")), id="Pair"), + pytest.param(Maybe[i64].some(i64(1)), id="Maybe some"), + pytest.param(Maybe[i64].none(), id="Maybe none"), pytest.param(Unit(), id="Unit"), pytest.param(UnstableFn[E, i64, i64](ff), id="fn"), pytest.param(UnstableFn[E, i64](ff, i64(1)), id="fn partial"), @@ -1336,6 +2651,103 @@ def my_cost_model(egraph: EGraph, expr: BaseExpr, children_costs: list[int]) -> my_cost_model.assert_any_call(egraph, x, []) my_cost_model.assert_any_call(egraph, called, [5]) + def test_map_container_children_costs_match_python_items_order(self): + expr = Map[i64, String].empty().insert(i64(2), String("b")).insert(i64(1), String("a")) + extracted, seen_expr, seen_children_costs = self._capture_container_children_costs( + expr, + leaf_cost=self._small_leaf_cost, + should_capture=lambda candidate, children_costs: isinstance(candidate, Map) and len(children_costs) == 4, + ) + + map_expr = cast("Map[i64, String]", seen_expr) + flattened_item_costs = [ + self._small_leaf_cost(item) for key, value in map_expr.value.items() for item in (key, value) + ] + + assert flattened_item_costs == seen_children_costs + assert flattened_item_costs == [11, 101, 22, 202] + assert list(cast("Map[i64, String]", extracted).value.items()) == list(map_expr.value.items()) + + def test_vec_container_children_costs_match_python_iteration_order(self): + expr = Vec(i64(2), i64(1), i64(3)) + extracted, _seen_expr, seen_children_costs = self._capture_container_children_costs( + expr, + leaf_cost=self._small_leaf_cost, + should_capture=lambda candidate, children_costs: isinstance(candidate, Vec) and len(children_costs) == 3, + ) + + extracted_vec = cast("Vec[i64]", extracted) + iter_costs = [self._small_leaf_cost(item) for item in extracted_vec] + + assert iter_costs == seen_children_costs + assert list(extracted_vec) == [i64(2), i64(1), i64(3)] + + def test_set_container_children_costs_match_python_iteration_order(self): + expr = Set(i64(2), i64(1), i64(2), i64(3)) + extracted, _seen_expr, seen_children_costs = self._capture_container_children_costs( + expr, + leaf_cost=self._small_leaf_cost, + should_capture=lambda candidate, children_costs: isinstance(candidate, Set) and len(children_costs) == 3, + ) + + extracted_set = cast("Set[i64]", extracted) + iter_costs = [self._small_leaf_cost(item) for item in extracted_set] + + assert iter_costs == seen_children_costs + assert list(extracted_set) == [i64(1), i64(2), i64(3)] + + def test_multiset_container_children_costs_match_python_iteration_order(self): + expr = MultiSet(i64(2), i64(1), i64(2), i64(3)) + extracted, _seen_expr, seen_children_costs = self._capture_container_children_costs( + expr, + leaf_cost=self._small_leaf_cost, + should_capture=lambda candidate, children_costs: isinstance(candidate, MultiSet) + and len(children_costs) == 4, + ) + + extracted_multiset = cast("MultiSet[i64]", extracted) + iter_costs = [self._small_leaf_cost(item) for item in extracted_multiset] + + assert iter_costs == seen_children_costs + assert list(extracted_multiset) == [i64(1), i64(2), i64(2), i64(3)] + + def test_pair_container_children_costs_match_python_value_order(self): + expr = Pair(i64(2), i64(1)) + extracted, _seen_expr, seen_children_costs = self._capture_container_children_costs( + expr, + leaf_cost=self._small_leaf_cost, + should_capture=lambda candidate, children_costs: isinstance(candidate, Pair) and len(children_costs) == 2, + ) + + pair_value = cast("Pair[i64, i64]", extracted).value + value_costs = [self._small_leaf_cost(item) for item in pair_value] + + assert value_costs == seen_children_costs + assert pair_value == (i64(2), i64(1)) + + def test_maybe_container_children_costs_match_python_value_order(self): + some_expr = Maybe[i64].some(i64(3)) + extracted_some, _seen_some, seen_some_children_costs = self._capture_container_children_costs( + some_expr, + leaf_cost=self._small_leaf_cost, + should_capture=lambda candidate, children_costs: isinstance(candidate, Maybe) and len(children_costs) == 1, + ) + + some_value = cast("Maybe[i64]", extracted_some).value + assert some_value is not None + assert [self._small_leaf_cost(some_value)] == seen_some_children_costs + assert some_value == i64(3) + + none_expr = Maybe[i64].none() + extracted_none, _seen_none, seen_none_children_costs = self._capture_container_children_costs( + none_expr, + leaf_cost=self._small_leaf_cost, + should_capture=lambda candidate, children_costs: isinstance(candidate, Maybe) and len(children_costs) == 0, + ) + + assert cast("Maybe[i64]", extracted_none).value is None + assert seen_none_children_costs == [] + @pytest.mark.xfail(reason="Errors dont bubble, just panic") def test_errors_bubble(self): def my_cost_model(egraph: EGraph, expr: BaseExpr, children_costs: list[int]) -> int: @@ -1549,11 +2961,11 @@ def __str__(self) -> str: return "hi" assert A.m(A()) == A().m() - assert isinstance(A.m, RuntimeFunction) + assert isinstance(cast("object", A.m), RuntimeFunction) assert eq(A.__eq__(A(), A())).to(A() == A()) - assert isinstance(A.__eq__, RuntimeFunction) + assert isinstance(cast("object", A.__eq__), RuntimeFunction) assert eq(A.__add__(A(), A())).to(A() + A()) - assert isinstance(A.__add__, RuntimeFunction) + assert isinstance(cast("object", A.__add__), RuntimeFunction) assert A.__str__(A()) == "hi" assert A.__str__.__doc__ == "Hi" @@ -1662,7 +3074,10 @@ def size(self) -> i64: ... def conv_cost(eg, expr, child_costs): if isinstance(expr, KAT): args = get_callable_args(expr) - return sum(int(eg.lookup_function_value(cast("KAT", a).size())) for a in args) + assert args is not None + values = [eg.lookup_function_value(cast("KAT", arg).size()) for arg in args] + assert all(value is not None for value in values) + return sum(int(value) for value in values if value is not None) return 2 diff --git a/python/tests/test_polynomials.py b/python/tests/test_polynomials.py new file mode 100644 index 00000000..2015f21f --- /dev/null +++ b/python/tests/test_polynomials.py @@ -0,0 +1,24 @@ +from egglog import * +from egglog.exp.array_api import factor_ruleset, from_polynomial_ruleset, to_polynomial_ruleset +from egglog.exp.polynomials import distribute, remove_subtraction, symbolic_bending_examples + + +def _factor_example(expr): + egraph = EGraph() + x = egraph.let("x", expr) + egraph.run(to_polynomial_ruleset.saturate() + factor_ruleset.saturate() + from_polynomial_ruleset.saturate()) + factored = egraph.extract(x) + egraph.check(eq(x).to(factored)) + return factored + + +def test_factor_multisets(snapshot_py): + function_bending, _gradient_bending = symbolic_bending_examples() + # remove subtraction and distribute first: + egraph = EGraph() + egraph.register(function_bending) + egraph.run(remove_subtraction.saturate()) + egraph.run(distribute.saturate()) + distributed = egraph.extract(function_bending) + factored = _factor_example(distributed) + assert str(factored) == snapshot_py(name="code") diff --git a/python/tests/test_pretty.py b/python/tests/test_pretty.py index d2a4dfc7..12cdf6f5 100644 --- a/python/tests/test_pretty.py +++ b/python/tests/test_pretty.py @@ -127,6 +127,7 @@ def my_very_long_function_name() -> A: ... r = ruleset(name="r") bo = back_off(ban_length=5) +bo_persistent = back_off(ban_length=5).persistent() class BadRepr: @@ -226,6 +227,12 @@ def __repr__(self) -> str: pytest.param(r + r, 'ruleset(name="r") + ruleset(name="r")', id="sequence"), pytest.param(seq(r, r, r), 'seq(ruleset(name="r"), ruleset(name="r"), ruleset(name="r"))', id="seq"), pytest.param(run(r, h()), 'run(ruleset(name="r"), h())', id="run"), + pytest.param(run(None, h()), "run(None, h())", id="default run with until"), + pytest.param( + run(None, h(), scheduler=bo), + "run(None, h(), scheduler=back_off(ban_length=5))", + id="default run with scheduler", + ), pytest.param( run(r, h(), scheduler=bo), 'run(ruleset(name="r"), h(), scheduler=back_off(ban_length=5))', @@ -236,6 +243,7 @@ def __repr__(self) -> str: '_scheduler_1 = back_off(ban_length=5)\n_scheduler_1.scope(run(ruleset(name="r"), scheduler=_scheduler_1))', id="scoped scheduler", ), + pytest.param(bo_persistent, "back_off(ban_length=5).persistent()", id="persistent scheduler"), # Functions pytest.param(f, "f", id="function"), pytest.param(A().method, "A().method", id="method"), @@ -252,6 +260,22 @@ def test_str(x: RuntimeExpr, s: str) -> None: assert str(x) == s +@pytest.mark.parametrize( + ("eval_mode", "option"), + [ + ("seminaive", ""), + ("naive", ', eval_mode="naive"'), + ("unsafe-seminaive", ', eval_mode="unsafe-seminaive"'), + ], +) +def test_rule_eval_mode_pretty_round_trip(eval_mode: RuleEvalMode, option: str) -> None: + original = rule(rel(g()), name="mode rule", eval_mode=eval_mode).then(rel(h())) + rendered = f'rule(rel(g()), name="mode rule"{option}).then(rel(h()))' + + assert str(original) == rendered + assert eval(rendered, globals()).decl == original.decl + + FREEZE_PARAMS = [ pytest.param((A(),), "EGraph(A()).freeze()", id="freeze add"), pytest.param((b,), "EGraph(b).freeze()", id="freeze constant"), @@ -295,3 +319,39 @@ def test_frozen_egraph_str_nested_vec_constructor() -> None: assert isinstance(frozen.decl, EGraphDecl) assert "Value(" not in str(frozen) assert str(frozen) == "EGraph(Wrapper(Box(Vec(A())))).freeze()" + + +@pytest.mark.parametrize( + "expr", + [ + pytest.param(BigRat(2, 1), id="integer BigRat"), + pytest.param(BigRat(1, 2), id="fractional BigRat"), + pytest.param(Map[String, BigRat].empty().insert(String("x"), BigRat(1, 2)), id="Map"), + pytest.param(Maybe[String].some(String("value")), id="Maybe some"), + pytest.param(Maybe[String].none(), id="Maybe none"), + pytest.param( + Pair[Map[String, BigRat], Maybe[String]]( + Map[String, BigRat].empty().insert(String("x"), BigRat(1, 2)), + Maybe[String].some(String("value")), + ), + id="Pair", + ), + ], +) +def test_extracted_typed_pretty_is_executable(expr: BaseExpr) -> None: + extracted = EGraph().extract(expr) + + rebuilt = eval(str(extracted), globals()) + + check_eq(extracted, rebuilt) + + +def test_frozen_typed_container_pretty_is_executable() -> None: + expr = Pair[Map[String, BigRat], Maybe[String]]( + Map[String, BigRat].empty().insert(String("x"), BigRat(1, 2)), Maybe[String].some(String("value")) + ) + frozen = EGraph(expr).freeze() + + rebuilt = eval(str(frozen), globals()) + + assert str(rebuilt) == str(frozen) diff --git a/python/tests/test_run_report.py b/python/tests/test_run_report.py index ec171001..0ae07fef 100644 --- a/python/tests/test_run_report.py +++ b/python/tests/test_run_report.py @@ -3,6 +3,8 @@ from datetime import timedelta +import pytest + from egglog import * from egglog.declarations import BiRewriteDecl, RewriteDecl, RuleDecl @@ -61,6 +63,13 @@ def test_updated_field(): assert report.updated is True +def test_can_stop_field(): + report = EGraph().run(1) + + assert report.can_stop is True + assert "can_stop=True" in repr(report) + + def test_num_matches(): egraph = _setup_simple_egraph() report = egraph.run(10) @@ -152,6 +161,17 @@ def __add__(self, other: Num) -> Num: ... assert "comm" in output, f"Expected rule name 'comm' in:\n{output}" +@pytest.mark.parametrize("name", ['a"b', r"a\b", 'a"b\\c']) +def test_saved_named_rule_round_trips_escaped_name(name: str) -> None: + egraph = EGraph(save_egglog_string=True) + seen = relation("seen", i64) + x = var("x", i64) + egraph.register(rule(seen(x), name=name).then(seen(x)), seen(i64(1))) + report = egraph.run(1) + + assert any(isinstance(decl, RuleDecl) and decl.name == name for decl in report.num_matches_per_rule) + + def test_unnamed_rule_decl(): egraph = EGraph() @@ -193,3 +213,33 @@ def __mul__(self, other: Num) -> Num: ... assert len(rule_keys) > 0 for key in rule_keys: assert isinstance(key, BiRewriteDecl) + + +def test_saved_transcript_report_translates_rewrite_with_string_literal() -> None: + class Word(Expr): + @classmethod + def named(cls, value: StringLike) -> Word: ... + + egraph = EGraph(save_egglog_string=True) + egraph.register(rewrite(Word.named("input")).to(Word.named("fixed"))) + egraph.register(Word.named("input")) + + report = egraph.run(1) + + assert report.search_and_apply_time_per_rule + assert all(isinstance(key, RewriteDecl) for key in report.search_and_apply_time_per_rule) + + +def test_saved_transcript_report_translates_birewrite_directions() -> None: + class Word(Expr): + @classmethod + def named(cls, value: StringLike) -> Word: ... + + egraph = EGraph(save_egglog_string=True) + egraph.register(birewrite(Word.named("input")).to(Word.named("fixed"))) + egraph.register(Word.named("input")) + + report = egraph.run(1) + + assert report.search_and_apply_time_per_rule + assert all(isinstance(key, BiRewriteDecl) for key in report.search_and_apply_time_per_rule) diff --git a/python/tests/test_unstable_fn.py b/python/tests/test_unstable_fn.py index 6d3b9b3b..2383886a 100644 --- a/python/tests/test_unstable_fn.py +++ b/python/tests/test_unstable_fn.py @@ -53,13 +53,13 @@ def __mul__(self, x: MathLike) -> MathList: ... converter(type(None), MathList, lambda _: MathList.NIL) -class Pair(Expr): +class Adder(Expr): def __init__(self, x: i64Like) -> None: ... - def add(self, y: i64Like) -> Pair: ... + def add(self, y: i64Like) -> Adder: ... -converter(i64, Pair, Pair) +converter(i64, Adder, Adder) @function @@ -84,9 +84,9 @@ def test_string_fn_partial(): 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)) + adder = Adder(2) + assert expr_parts(UnstableFn(adder.add)) == expr_parts(UnstableFn(Adder.add, adder)) + assert expr_parts(UnstableFn(adder.add, 3)) == expr_parts(UnstableFn(Adder.add, adder, 3)) @ruleset @@ -119,6 +119,26 @@ def test_partial_application(): ) +def test_partial_application_does_not_reuse_synthetic_let_for_its_call() -> None: + class PartialLet(Expr): + @classmethod + def value(cls, value: i64Like) -> PartialLet: ... + + @classmethod + def pair(cls, left: PartialLet, right: PartialLet) -> PartialLet: ... + + @function + def apply0(f: UnstableFn[PartialLet]) -> PartialLet: ... + + shared = PartialLet.value(1) + egraph = EGraph() + + egraph.register(PartialLet.pair(shared, shared)) + egraph.register(apply0(UnstableFn(PartialLet.value, i64(1)))) + + assert egraph.function_size(apply0) == 1 + + @function def composed_math(f: MathFn, g: MathFn, x: Math) -> Math: ... @@ -186,14 +206,16 @@ def test_callable_accepted_as_type(): @function def func(f: UnstableFn[C, A, B]) -> C: ... - assert isinstance(func, RuntimeFunction) - original = func.__egg_decls__, func.__egg_ref__ + original_func: object = func + assert isinstance(original_func, RuntimeFunction) + original = original_func.__egg_decls__, original_func.__egg_ref__ @function # type: ignore[no-redef] def func(f: Callable[[A, B], C]) -> C: ... - assert isinstance(func, RuntimeFunction) - converted = func.__egg_decls__, func.__egg_ref__ + converted_func: object = func + assert isinstance(converted_func, RuntimeFunction) + converted = converted_func.__egg_decls__, converted_func.__egg_ref__ assert converted == original diff --git a/rust-toolchain.toml b/rust-toolchain.toml index b67e7d53..d72668b0 100644 --- a/rust-toolchain.toml +++ b/rust-toolchain.toml @@ -1,2 +1,2 @@ [toolchain] -channel = "1.89.0" +channel = "1.91.0" diff --git a/src/conversions.rs b/src/conversions.rs index 039771fa..f1331b54 100644 --- a/src/conversions.rs +++ b/src/conversions.rs @@ -3,6 +3,7 @@ use crate::utils::*; use egglog::extract::DefaultCost; use ordered_float::OrderedFloat; +use pyo3::exceptions::{PyOverflowError, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyDelta, PyDeltaAccess}; use std::collections::HashMap; @@ -125,6 +126,17 @@ convert_enums!( _d -> egglog::ast::PrintFunctionMode::CSV, egglog::ast::PrintFunctionMode::CSV => CSVPrintFunctionMode {} }; + egglog::ast::RuleEvalMode: "{:?}" => RuleEvalMode { + Seminaive() + _s -> egglog::ast::RuleEvalMode::Seminaive, + egglog::ast::RuleEvalMode::Seminaive => Seminaive {}; + Naive() + _n -> egglog::ast::RuleEvalMode::Naive, + egglog::ast::RuleEvalMode::Naive => Naive {}; + UnsafeSeminaive() + _u -> egglog::ast::RuleEvalMode::UnsafeSeminaive, + egglog::ast::RuleEvalMode::UnsafeSeminaive => UnsafeSeminaive {} + }; egglog::ast::Command: "{}" => Command { Datatype(span: Span, name: String, variants: Vec) d -> egglog::ast::Command::Datatype { @@ -137,36 +149,81 @@ convert_enums!( name: name.to_string(), variants: variants.iter().map(|v| v.into()).collect() }; - Sort(span: Span, name: String, presort_and_args: Option<(String, Vec)>) + Sort( + span: Span, + name: String, + presort_and_args: Option<(String, Vec)>, + uf: Option<(String, Option)> = None, + proof_func: Option = None, + container_rebuild: Option = None, + proof_constructors: Option = None + ) s -> egglog::ast::Command::Sort { span: s.span.clone().into(), name: (&s.name).into(), presort_and_args: s.presort_and_args.as_ref().map(|(p, a)| (p.into(), a.iter().map(|e| e.into()).collect())), - uf: None, - proof_func: None, + uf: s.uf.clone(), + proof_func: s.proof_func.clone(), + container_rebuild: s.container_rebuild.as_ref().map(Into::into), + proof_constructors: s.proof_constructors.as_ref().map(Into::into), unionable: true }, - egglog::ast::Command::Sort { span, name, presort_and_args, .. } => Sort { + egglog::ast::Command::Sort { + span, + name, + presort_and_args, + uf, + proof_func, + container_rebuild, + proof_constructors, + unionable: _ + } => Sort { name: name.to_string(), presort_and_args: presort_and_args.as_ref().map(|(p, a)| (p.to_string(), a.iter().map(|e| e.into()).collect())), - span: span.into() + span: span.into(), + uf: uf.clone(), + proof_func: proof_func.clone(), + container_rebuild: container_rebuild.as_ref().map(Into::into), + proof_constructors: proof_constructors.as_ref().map(Into::into) }; - FunctionCommand(span: Span, name: String, schema: Schema, merge: Option) + FunctionCommand( + span: Span, + name: String, + schema: Schema, + merge: Option, + term_constructor: Option = None, + unextractable: bool = false, + hidden: bool = false, + let_binding: bool = false + ) f -> egglog::ast::Command::Function{ span: f.span.clone().into(), name: (&f.name).into(), schema: (&f.schema).into(), merge: f.merge.as_ref().map(|e| e.into()), - hidden: false, - let_binding: false, - term_constructor: None, - unextractable: false + hidden: f.hidden, + let_binding: f.let_binding, + term_constructor: f.term_constructor.clone(), + unextractable: f.unextractable }, - egglog::ast::Command::Function {span, name, schema, merge, .. } => FunctionCommand { + egglog::ast::Command::Function { + span, + name, + schema, + merge, + hidden, + let_binding, + term_constructor, + unextractable + } => FunctionCommand { span: span.into(), name: name.to_string(), schema: schema.into(), - merge: merge.as_ref().map(|e| e.into()) + merge: merge.as_ref().map(|e| e.into()), + term_constructor: term_constructor.clone(), + unextractable: *unextractable, + hidden: *hidden, + let_binding: *let_binding }; AddRuleset(span: Span, name: String) a -> egglog::ast::Command::AddRuleset( @@ -275,23 +332,42 @@ convert_enums!( Include(span: Span, path: String) i -> egglog::ast::Command::Include(i.span.clone().into(), (&i.path).into()), egglog::ast::Command::Include(span, p) => Include { span: span.into(), path: p.to_string() }; - Constructor(span: Span, name: String, schema: Schema, cost: Option, unextractable: bool) + Constructor( + span: Span, + name: String, + schema: Schema, + cost: Option, + unextractable: bool, + hidden: bool = false, + let_binding: bool = false + ) c -> egglog::ast::Command::Constructor { span: c.span.clone().into(), name: (&c.name).into(), schema: (&c.schema).into(), cost: c.cost, unextractable: c.unextractable, - hidden: false, - let_binding: false, + hidden: c.hidden, + let_binding: c.let_binding, term_constructor: None }, - egglog::ast::Command::Constructor {span, name, schema, cost, unextractable, .. } => Constructor { + egglog::ast::Command::Constructor { + span, + name, + schema, + cost, + unextractable, + hidden, + let_binding, + term_constructor: _ + } => Constructor { span: span.into(), name: name.to_string(), schema: schema.into(), cost: *cost, - unextractable: *unextractable + unextractable: *unextractable, + hidden: *hidden, + let_binding: *let_binding }; Relation(span: Span, name: String, inputs: Vec) r -> egglog::ast::Command::Relation { @@ -441,13 +517,13 @@ convert_enums!( EgglogSpan(file: SrcFile, i: usize, j: usize) e -> egglog_ast::span::Span::Egglog(Arc::new({ egglog_ast::span::EgglogSpan { - file: Arc::new(e.file.clone().into()), + file: e.file.0.clone(), i: e.i, j: e.j } })), egglog_ast::span::Span::Egglog(e) => EgglogSpan { - file: (*e.file.clone()).clone().into(), + file: SrcFile(e.file.clone()), i: e.i, j: e.j }; @@ -461,13 +537,90 @@ convert_enums!( } ); +impl Default for RuleEvalMode { + fn default() -> Self { + Self::Seminaive(Seminaive {}) + } +} + +#[pyclass(frozen)] +#[derive(Clone, PartialEq, Eq)] +pub struct SrcFile(Arc); + +#[pymethods] +impl SrcFile { + #[new] + fn new(name: Option, contents: String) -> Self { + Self(Arc::new(egglog_ast::span::SrcFile { name, contents })) + } + + #[getter] + fn name(&self) -> Option<&str> { + self.0.name.as_deref() + } + + #[getter] + fn contents(&self) -> &str { + &self.0.contents + } + + fn __repr__(slf: PyRef<'_, Self>, py: Python) -> PyResult { + data_repr(py, slf, vec!["name", "contents"]) + } + + fn __str__(&self) -> String { + format!("{:?}", self.0) + } + + fn __richcmp__( + &self, + other: &Self, + op: pyo3::basic::CompareOp, + py: Python<'_>, + ) -> PyResult> { + Ok(match op { + pyo3::basic::CompareOp::Eq => { + (self == other).into_pyobject(py)?.as_any().clone().unbind() + } + pyo3::basic::CompareOp::Ne => { + (self != other).into_pyobject(py)?.as_any().clone().unbind() + } + _ => py.NotImplemented(), + }) + } +} + convert_struct!( - egglog_ast::span::SrcFile: "{:?}" => SrcFile( - name: Option, - contents: String + egglog::ast::ContainerRebuildSpec: "{}" => ContainerRebuildSpec( + internal_rebuild_prim: String, + internal_rebuild_proof_prim: Option = None ) - s -> egglog_ast::span::SrcFile {name: s.name.clone(), contents: s.contents.clone()}, - s -> SrcFile {name: s.name.clone(), contents: s.contents.clone()}; + s -> egglog::ast::ContainerRebuildSpec { + internal_rebuild_prim: s.internal_rebuild_prim.clone(), + internal_rebuild_proof_prim: s.internal_rebuild_proof_prim.clone() + }, + s -> ContainerRebuildSpec { + internal_rebuild_prim: s.internal_rebuild_prim.clone(), + internal_rebuild_proof_prim: s.internal_rebuild_proof_prim.clone() + }; + egglog::ast::ProofConstructorNames: "{:?}" => ProofConstructorNames( + congr: String, + trans: String, + sym: String, + normalize: String + ) + n -> egglog::ast::ProofConstructorNames { + congr: n.congr.clone(), + trans: n.trans.clone(), + sym: n.sym.clone(), + normalize: n.normalize.clone() + }, + n -> ProofConstructorNames { + congr: n.congr.clone(), + trans: n.trans.clone(), + sym: n.sym.clone(), + normalize: n.normalize.clone() + }; egglog::ast::Variant: "{:?}" => Variant( span: Span, name: String, @@ -488,10 +641,31 @@ convert_struct!( head: Vec, body: Vec, name: String, - ruleset: String + ruleset: String, + eval_mode: RuleEvalMode = RuleEvalMode::default(), + no_decomp: bool = false, + include_subsumed: bool = false ) - r -> egglog::ast::GenericRule {span: r.span.clone().into(), head: egglog::ast::GenericActions(r.head.iter().map(|v| v.into()).collect()), body: r.body.iter().map(|v| v.into()).collect(), name: (&r.name).into(), ruleset: (&r.ruleset).into()}, - r -> Rule {span: r.span.clone().into(), head: r.head.0.iter().map(|v| v.into()).collect(), body: r.body.iter().map(|v| v.into()).collect(), name: r.name.to_string(), ruleset: r.ruleset.to_string()}; + r -> egglog::ast::GenericRule { + span: r.span.clone().into(), + head: egglog::ast::GenericActions(r.head.iter().map(|v| v.into()).collect()), + body: r.body.iter().map(|v| v.into()).collect(), + name: (&r.name).into(), + ruleset: (&r.ruleset).into(), + eval_mode: (&r.eval_mode).into(), + no_decomp: r.no_decomp, + include_subsumed: r.include_subsumed + }, + r -> Rule { + span: r.span.clone().into(), + head: r.head.0.iter().map(|v| v.into()).collect(), + body: r.body.iter().map(|v| v.into()).collect(), + name: r.name.to_string(), + ruleset: r.ruleset.to_string(), + eval_mode: (&r.eval_mode).into(), + no_decomp: r.no_decomp, + include_subsumed: r.include_subsumed + }; egglog::ast::GenericRewrite: "{:?}" => Rewrite( span: Span, lhs: Expr, @@ -644,6 +818,7 @@ convert_struct!( egglog_reports::RunReport: "{:?}" => RunReport( iterations: Vec, updated: bool, + can_stop: bool, search_and_apply_time_per_rule: HashMap, num_matches_per_rule: HashMap, search_and_apply_time_per_ruleset: HashMap, @@ -657,6 +832,7 @@ convert_struct!( .map(|i| Arc::new(i.clone().into())) .collect(), updated: r.updated, + can_stop: r.can_stop, search_and_apply_time_per_rule: r .search_and_apply_time_per_rule .iter() @@ -686,6 +862,7 @@ convert_struct!( r -> RunReport { iterations: r.iterations.iter().map(|i| i.as_ref().into()).collect(), updated: r.updated, + can_stop: r.can_stop, search_and_apply_time_per_rule: r .search_and_apply_time_per_rule .iter() @@ -791,30 +968,64 @@ impl<'py> FromPyObject<'_, 'py> for WrappedDuration { type Error = PyErr; fn extract(obj: Borrowed<'_, 'py, PyAny>) -> Result { let py_delta = obj.cast::()?; + let days = py_delta.get_days(); + let seconds = py_delta.get_seconds(); + let microseconds = py_delta.get_microseconds(); + if days < 0 { + return Err(PyValueError::new_err( + "negative timedeltas cannot be converted to Rust Duration", + )); + } + if seconds < 0 || microseconds < 0 { + return Err(PyValueError::new_err("invalid timedelta components")); + } + let seconds = (days as u64) + .checked_mul(SECONDS_PER_DAY) + .and_then(|day_seconds| day_seconds.checked_add(seconds as u64)) + .ok_or_else(|| PyOverflowError::new_err("timedelta is too large for Rust Duration"))?; + let nanoseconds = (microseconds as u32).checked_mul(1_000).ok_or_else(|| { + PyOverflowError::new_err("timedelta is too precise for Rust Duration") + })?; Ok(WrappedDuration(std::time::Duration::new( - py_delta.get_days() as u64 * 24 * 60 * 60 + py_delta.get_seconds() as u64, - py_delta.get_microseconds() as u32 * 1000, + seconds, + nanoseconds, ))) } } + +const SECONDS_PER_DAY: u64 = 24 * 60 * 60; + +fn duration_to_py_parts(duration: std::time::Duration) -> Option<(i32, i32, i32)> { + let total_seconds = duration.as_secs(); + Some(( + (total_seconds / SECONDS_PER_DAY).try_into().ok()?, + (total_seconds % SECONDS_PER_DAY).try_into().ok()?, + duration.subsec_micros().try_into().ok()?, + )) +} + impl<'py> IntoPyObject<'py> for WrappedDuration { type Target = PyDelta; // the Python type type Output = Bound<'py, Self::Target>; // in most cases this will be `Bound` type Error = pyo3::PyErr; fn into_pyobject(self, py: Python<'py>) -> Result { - let d = self.0; - Ok(pyo3::types::PyDelta::new( - py, - 0, - 0, - d.as_millis() - .try_into() - .expect("Failed to convert miliseconds to int32 when converting duration"), - true, - )? - .clone()) + let (days, seconds, microseconds) = duration_to_py_parts(self.0).ok_or_else(|| { + PyOverflowError::new_err("Rust Duration is too large for datetime.timedelta") + })?; + Ok(pyo3::types::PyDelta::new(py, days, seconds, microseconds, true)?.clone()) + } +} + +#[cfg(test)] +mod duration_tests { + use super::*; + + #[test] + fn rejects_duration_with_too_many_days_for_python() { + let seconds = (i32::MAX as u64 + 1) * SECONDS_PER_DAY; + assert!(duration_to_py_parts(std::time::Duration::from_secs(seconds)).is_none()); } } @@ -862,3 +1073,33 @@ impl PartialEq for Function { } impl std::cmp::Eq for Function {} + +#[cfg(test)] +mod span_tests { + use super::*; + + #[test] + fn egglog_span_conversion_shares_its_source_file() { + let file = Arc::new(egglog_ast::span::SrcFile { + name: Some("large.egg".to_owned()), + contents: "(relation R (i64))\n".repeat(1_000), + }); + let span = egglog_ast::span::Span::Egglog(Arc::new(egglog_ast::span::EgglogSpan { + file: file.clone(), + i: 0, + j: 18, + })); + + let Span::EgglogSpan(converted) = Span::from(&span) else { + panic!("expected an egglog span"); + }; + assert!(Arc::ptr_eq(&converted.file.0, &file)); + + let egglog_ast::span::Span::Egglog(round_tripped) = + egglog_ast::span::Span::from(&Span::EgglogSpan(converted)) + else { + panic!("expected an egglog span"); + }; + assert!(Arc::ptr_eq(&round_tripped.file, &file)); + } +} diff --git a/src/egraph.rs b/src/egraph.rs index 4cad3d30..fc3548b7 100644 --- a/src/egraph.rs +++ b/src/egraph.rs @@ -5,12 +5,12 @@ use crate::error::{EggResult, WrappedError}; use crate::freeze::FrozenEGraph; use crate::py_object_sort::{PyObjectSort, PyPickledValue, load}; use crate::serialize::SerializedEGraph; +use crate::termdag::TermDag; use crate::tracing_otel; -use egglog::prelude::{RustSpan, Span, add_base_sort}; -use egglog::{SerializeConfig, span}; +use egglog::prelude::add_base_sort; +use egglog::{RawValues, Read as _, SerializeConfig, span}; use log::info; -use num_bigint::BigInt; use num_rational::{BigRational, Rational64}; use pyo3::prelude::*; use std::collections::{BTreeMap, BTreeSet}; @@ -26,6 +26,51 @@ pub struct EGraph { cmds: Option, } +impl EGraph { + fn run_parsed_commands( + &mut self, + py: Python<'_>, + commands: Vec, + parsed_from_source: bool, + ) -> EggResult> { + let cmds_str = commands + .iter() + .map(|command| format!("{command}\n")) + .collect::(); + let res = if parsed_from_source { + let span = tracing::info_span!( + "bindings.parse_and_run_program", + command_count = commands.len(), + commands = tracing::field::display(cmds_str.trim_end()) + ); + let _entered = span.enter(); + info!("Running commands:\n{}", cmds_str); + py.detach(|| self.egraph.run_program(commands)) + } else { + let span = tracing::info_span!( + "bindings.run_program", + command_count = commands.len(), + commands = tracing::field::display(cmds_str.trim_end()) + ); + let _entered = span.enter(); + info!("Running commands:\n{}", cmds_str); + py.detach(|| self.egraph.run_program(commands)) + }; + if let Some(err) = PyErr::take(py) { + return Err(WrappedError::Py(err)); + } + 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()) + } + } + } +} + #[pymethods] impl EGraph { #[new] @@ -37,7 +82,7 @@ impl EGraph { add_base_sort(&mut egraph, PyObjectSort {}, span!()).unwrap(); Self { egraph, - cmds: if record { Some(String::new()) } else { None }, + cmds: record.then(String::new), } } @@ -51,6 +96,25 @@ impl EGraph { Ok(commands.into_iter().map(|x| x.into()).collect()) } + /// Parse a program and immediately run the parsed commands on the EGraph. + #[pyo3(signature = (input, /, filename=None, traceparent=None, tracestate=None))] + fn parse_and_run_program( + &mut self, + py: Python<'_>, + input: &str, + filename: Option, + traceparent: Option, + tracestate: Option, + ) -> EggResult> { + let _context_guard = + tracing_otel::attach_parent_context(traceparent.as_deref(), tracestate.as_deref()); + let commands = self + .egraph + .parser + .get_program_from_string(filename, input)?; + self.run_parsed_commands(py, commands, true) + } + /// Run a series of commands on the EGraph. /// Returns a list of strings representing the output. /// An EggSmolError is raised if there is problem parsing or executing. @@ -65,37 +129,10 @@ impl EGraph { let _context_guard = tracing_otel::attach_parent_context(traceparent.as_deref(), tracestate.as_deref()); let commands: Vec = commands.into_iter().map(|x| x.into()).collect(); - let mut cmds_str = String::new(); - - for cmd in &commands { - let cmd_string = cmd.to_string(); - cmds_str = cmds_str + &cmd_string + "\n"; - } - let span = tracing::info_span!( - "bindings.run_program", - command_count = commands.len(), - commands = tracing::field::display(cmds_str.trim_end()) - ); - let _entered = span.enter(); - info!("Running commands:\n{}", cmds_str); - let res = py.detach(|| self.egraph.run_program(commands)); - if let Some(err) = PyErr::take(py) { - return Err(WrappedError::Py(err)); - } - match res { - Err(e) => Err(WrappedError::Egglog(e)), - Ok(outputs) => { - if let Some(cmds) = &mut self.cmds { - cmds.push_str(&cmds_str); - } - let outputs = outputs.into_iter().map(|o| o.into()).collect(); - Ok(outputs) - } - } + self.run_parsed_commands(py, commands, false) } - /// Returns the text of the commands that have been run so far, if `record` was passed. - #[pyo3(signature = ())] + /// Returns the text of successfully run commands when recording is enabled. fn commands(&self) -> Option { self.cmds.clone() } @@ -146,13 +183,29 @@ impl EGraph { self.egraph.set_report_level(level.into()); } - fn lookup_function(&self, name: &str, key: Vec) -> Option { - self.egraph - .lookup_function( - name, - key.into_iter().map(|v| v.0).collect::>().as_slice(), - ) - .map(Value) + fn lookup_function(&self, name: &str, key: Vec) -> EggResult> { + let is_constructor = self.egraph.get_function(name).is_some_and(|function| { + function.func_type().subtype == egglog::ast::FunctionSubtype::Constructor + }); + let value = self.egraph.read(|state| { + let key = RawValues(key.into_iter().map(|value| value.0).collect()); + if is_constructor { + state.eclass_of(name, key) + } else { + state.lookup(name, key) + } + })?; + Ok(value.map(Value)) + } + + /// Extract `value` using its runtime sort. `sort` must match the sort returned with `value` + /// by `eval_expr`; passing a different existing sort is unsupported. + fn extract_value(&self, value: Value, sort: &str) -> EggResult<(TermDag, usize, u64)> { + let sort = self.egraph.get_sort_by_name(sort).ok_or_else(|| { + WrappedError::Egglog(egglog::TypeError::UndefinedSort(sort.to_owned(), span!()).into()) + })?; + let (termdag, term, cost) = self.egraph.extract_value(sort, value.0)?; + Ok((TermDag(termdag), term, cost)) } #[pyo3(signature = (expr, *, traceparent=None, tracestate=None))] @@ -184,9 +237,9 @@ impl EGraph { self.egraph.value_to_base(v.0) } - fn value_to_bigint(&self, v: Value) -> BigInt { + fn value_to_bigint<'py>(&self, py: Python<'py>, v: Value) -> PyResult> { let bi: egglog::sort::Z = self.egraph.value_to_base(v.0); - bi.0 + Ok(bi.0.into_pyobject(py)?.into_any()) } fn value_to_bigrat(&self, v: Value) -> BigRational { @@ -266,22 +319,6 @@ impl EGraph { fn freeze(&self) -> FrozenEGraph { FrozenEGraph::from_egraph(&self.egraph) } - - // fn dynamic_cost_model_enode_cost( - // &self, - // func: String, - // args: Vec, - // ) -> EggResult { - // let func = self.egraph.get_function(&func).ok_or_else(|| { - // WrappedError::Py(PyRuntimeError::new_err(format!("No such function: {func}"))) - // })?; - // let vals: Vec = args.into_iter().map(|v| v.0).collect(); - // let row = FunctionRow { - // vals: &vals, - // subsumed: false, - // }; - // Ok(egglog_experimental::DynamicCostModel {}.enode_cost(&self.egraph, &func, &row)) - // } } /// Wrapper around Egglog Value. Represents either a primitive base value or a reference to an e-class. diff --git a/src/extract.rs b/src/extract.rs index 5385610a..54306c8d 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -118,14 +118,10 @@ impl egglog::extract::CostModel for CostModel { &self, _egraph: &egglog::EGraph, func: &egglog::Function, - row: &egglog::FunctionRow<'_>, + enode: &egglog::Enode<'_>, ) -> Cost { Python::attach(|py| { - let mut values = row.vals.iter().map(|v| Value(*v)).collect::>(); - // Remove last element which is the output - // this is not needed because the only thing we can do with the output is look up an analysis - // which we can also do with the original function - values.pop().unwrap(); + let values = enode.children.iter().map(|v| Value(*v)).collect::>(); Cost(self.enode_cost.call1(py, (func.name(), values)).unwrap()) }) } diff --git a/src/freeze.rs b/src/freeze.rs index 28d1a1e8..c7d8969e 100644 --- a/src/freeze.rs +++ b/src/freeze.rs @@ -1,6 +1,6 @@ // 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 egglog::{EGraph, ast::FunctionSubtype}; use indexmap::IndexMap; use pyo3::prelude::*; @@ -33,29 +33,36 @@ impl FrozenEGraph { /// 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() { + for (fname, func) in egraph.functions_iter() { let mut rows = Vec::new(); - egraph.function_for_each(&fname, |row| { - let frozen_row = FrozenRow { - subsumed: row.subsumed, - inputs: row.vals[..row.vals.len() - 1] - .iter() - .cloned() - .map(Value) - .collect(), - output: Value(*row.vals.last().unwrap()), - }; - rows.push(frozen_row); - }).unwrap(); - let func = egraph.get_function(&fname).unwrap(); + match func.func_type().subtype { + FunctionSubtype::Constructor => egraph + .constructor_enodes(fname, |enode| { + rows.push(FrozenRow { + subsumed: enode.subsumed, + inputs: enode.children.iter().copied().map(Value).collect(), + output: Value(enode.eclass), + }); + }) + .unwrap(), + FunctionSubtype::Custom => egraph + .function_entries(fname, |entry| { + rows.push(FrozenRow { + subsumed: entry.subsumed, + inputs: entry.inputs.iter().copied().map(Value).collect(), + output: Value(entry.output), + }); + }) + .unwrap(), + } + let func_type = func.func_type(); let frozen_function = FrozenFunction { - input_sorts: func - .schema() + input_sorts: func_type .input .iter() .map(|s| s.name().to_string()) .collect(), - output_sort: func.schema().output.name().to_string(), + output_sort: func_type.output.name().to_string(), rows, is_let_binding: func.is_let_binding(), }; diff --git a/src/lib.rs b/src/lib.rs index 7590ddc7..36bbb772 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -54,6 +54,7 @@ fn bindings(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_function(wrap_pyfunction!(setup_tracing, m)?)?; m.add_function(wrap_pyfunction!(shutdown_tracing, m)?)?; crate::conversions::add_structs_to_module(m)?; diff --git a/src/utils.rs b/src/utils.rs index 34594407..63827b6b 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -29,7 +29,7 @@ macro_rules! convert_enums { ($( $from_type:ty: $str:literal $($trait_outer:ty)? => $to_type:ident { $( - $variant:ident$([name=$py_name:literal])?$([trait=$trait_inner:ty])?($($field:ident: $field_type:ty),*) + $variant:ident$([name=$py_name:literal])?$([trait=$trait_inner:ty])?($($field:ident: $field_type:ty$( = $default:expr)?),*) $from_ident:ident -> $from:expr, $to_pat:pat => $to:expr );* @@ -48,7 +48,7 @@ macro_rules! convert_enums { #[pymethods] impl $variant { #[new] - #[pyo3(signature=($($field),*))] + #[pyo3(signature=($($field $(= $default)?),*))] fn new($($field: $field_type),*) -> Self { Self { $($field),* diff --git a/test-data/unit/check-high-level.test b/test-data/unit/check-high-level.test index f1faabb6..020b2062 100644 --- a/test-data/unit/check-high-level.test +++ b/test-data/unit/check-high-level.test @@ -13,3 +13,32 @@ _ = eq(i64(0)).to(i64(0)) [case eqToNotAllowed] from egglog import * _ = eq(i64(0)).to(Unit()) # E: Argument 1 to "to" of "_EqBuilder" has incompatible type "Unit"; expected "i64" + +[case functionPrimitiveMergeAllowed] +from egglog import * +@function(merge=lambda old, new: old) +def f(x: i64Like) -> i64: ... + +[case functionEqsortCostAllowed] +from egglog import * +class A(Expr): ... +@function(cost=1) +def f() -> A: ... + +[case functionEqsortRulesetSubsumeAllowed] +from egglog import * +class A(Expr): ... +r = ruleset() +@function(ruleset=r, subsume=True) +def f() -> A: ... + +[case functionPrimitiveCostRejected] +from egglog import * +@function(cost=1) # E: Value of type variable "CONSTRUCTOR_CALLABLE" of function cannot be "Callable[[], i64]" +def f() -> i64: ... + +[case functionPrimitiveRulesetSubsumeRejected] +from egglog import * +r = ruleset() +@function(ruleset=r, subsume=True) # E: Value of type variable "CONSTRUCTOR_CALLABLE" of function cannot be "Callable[[], i64]" +def f() -> i64: ... diff --git a/uv.lock b/uv.lock index d97e6708..6c099f4e 100644 --- a/uv.lock +++ b/uv.lock @@ -2,7 +2,8 @@ version = 1 revision = 3 requires-python = ">=3.11" resolution-markers = [ - "python_full_version >= '3.14'", + "python_full_version >= '3.15'", + "python_full_version == '3.14.*'", "python_full_version >= '3.12' and python_full_version < '3.14'", "python_full_version < '3.12'", ] @@ -153,6 +154,44 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ed/c9/d7977eaacb9df673210491da99e6a247e93df98c715fc43fd136ce1d3d33/arrow-1.4.0-py3-none-any.whl", hash = "sha256:749f0769958ebdc79c173ff0b0670d59051a535fa26e8eba02953dc19eb43205", size = 68797, upload-time = "2025-10-18T17:46:45.663Z" }, ] +[[package]] +name = "ast-serialize" +version = "0.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a9/9d/912fefab0e30aee6a3af8a62bbea4a81b29afa4ba2c973d31170620a26de/ast_serialize-0.3.0.tar.gz", hash = "sha256:1bc3ca09a63a021376527c4e938deedd11d11d675ce850e6f9c7487f5889992b", size = 60689, upload-time = "2026-04-30T23:24:48.104Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6f/57/a54d4de491d6cdd7a4e4b0952cc3ca9f60dcefa7b5fb48d6d492debe1649/ast_serialize-0.3.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:3a867927df59f76a18dc1d874a0b2c079b42c58972dca637905576deb0912e14", size = 1182966, upload-time = "2026-04-30T23:23:57.376Z" }, + { url = "https://files.pythonhosted.org/packages/ee/9e/a5db014bb0f91b209236b57c429389e31290c0093532b8436d577699b2fa/ast_serialize-0.3.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:a6fb063bf040abf8321e7b8113a0554eda445ffc508aa51287f8808886a5ae22", size = 1171316, upload-time = "2026-04-30T23:23:59.63Z" }, + { url = "https://files.pythonhosted.org/packages/15/59/fd55133e478c4326f60a11df02573bf7ccb2ac685810b50f1803d0f68053/ast_serialize-0.3.0-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5075cd8482573d743586779e5f9b652a015e37d4e95132d7e5a9bc5c8f483d8f", size = 1232234, upload-time = "2026-04-30T23:24:01.168Z" }, + { url = "https://files.pythonhosted.org/packages/cc/79/0ca1d26357ecb4a697d74d00b73ef3137f24c140424125393a0de820eb09/ast_serialize-0.3.0-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:41560b27794f4553b0f77811e9fb325b77db4a2b39018d437e09932275306e66", size = 1233437, upload-time = "2026-04-30T23:24:03.151Z" }, + { url = "https://files.pythonhosted.org/packages/53/3e/7078ec94dd6e124b8e028ac77016a4f13c83fa1c145790f2e68f3816998b/ast_serialize-0.3.0-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b967c01ca74909c5d90e0fe4393401e2cc5da5ebd9a6262a19e45ffd3757dec8", size = 1440188, upload-time = "2026-04-30T23:24:04.717Z" }, + { url = "https://files.pythonhosted.org/packages/21/16/cca7195ef55a012f8013c3442afa91d287a0a36dcf88b480b262475135b3/ast_serialize-0.3.0-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:424ebb8f46cd993f7cec4009d119312d8433dd90e6b0df0499cd2c91bdcc5af9", size = 1254211, upload-time = "2026-04-30T23:24:06.18Z" }, + { url = "https://files.pythonhosted.org/packages/a0/0f/f3d4dfae67dee6580534361a6343367d34217e7d25cff858bd1d8f03b8ed/ast_serialize-0.3.0-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d14b1d566b56e2ee70b11fec1de7e0b94ec7cd83717ec7d189967841a361190e", size = 1255973, upload-time = "2026-04-30T23:24:07.772Z" }, + { url = "https://files.pythonhosted.org/packages/14/41/55fbfe02c42f40fbe3e74eda167d977d555ff720ce1abfa08515236efd88/ast_serialize-0.3.0-cp314-cp314t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:7ba30b18735f047ec11103d1ab92f4789cf1fea1e0dc89b04a2f5a0632fd79de", size = 1298629, upload-time = "2026-04-30T23:24:09.4Z" }, + { url = "https://files.pythonhosted.org/packages/28/36/7d2501cacc7989fb8504aa9da2a2022a174200a59d4e6639de4367a57fdd/ast_serialize-0.3.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e6ea0754cb7b0f682ebb005ffb0d18f8d17993490d9c289863cd69cacc4ab8df", size = 1408435, upload-time = "2026-04-30T23:24:11.013Z" }, + { url = "https://files.pythonhosted.org/packages/03/e7/54e3b469c3fa0bf9cd532fa643d1d33b73303f8d70beac3e366b68dd64b7/ast_serialize-0.3.0-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:a0c5aa1073a5ba7b2abaa4b54abe8b8d75c4d1e2d54a2ff70b0ca6222fea5728", size = 1508174, upload-time = "2026-04-30T23:24:12.635Z" }, + { url = "https://files.pythonhosted.org/packages/b5/2a/9b9621865b02c60539e26d9b114a312b4fa46aa703e33e79317174bfea21/ast_serialize-0.3.0-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:4e52650d834c1ea7791969a361de2c54c13b2fb4c519ec79445fa8b9021a147d", size = 1502354, upload-time = "2026-04-30T23:24:14.186Z" }, + { url = "https://files.pythonhosted.org/packages/34/dd/f138bc5c43b0c414fdd12eefe15677839323078b6e75301ad7f96cd26d45/ast_serialize-0.3.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:15bd6af3f136c61dae27805eb6b8f3269e85a545c4c27ffe9e530ead78d2b36d", size = 1450504, upload-time = "2026-04-30T23:24:16.076Z" }, + { url = "https://files.pythonhosted.org/packages/68/cf/97ef9e1c315601db74365955c8edd3292e3055500d6317602815dbdf08ae/ast_serialize-0.3.0-cp314-cp314t-win32.whl", hash = "sha256:d188bfe37b674b49708497683051d4b571366a668799c9b8e8a94513694969d9", size = 1058662, upload-time = "2026-04-30T23:24:17.535Z" }, + { url = "https://files.pythonhosted.org/packages/f8/d6/e2c3483c31580fdb623f92ad38d2f856cde4b9205a3e6bd84760f3de7d82/ast_serialize-0.3.0-cp314-cp314t-win_amd64.whl", hash = "sha256:5832c2fdf8f8a6cf682b4cfcf677f5eaf39b4ddbc490f5480cfccdd1e7ce8fa1", size = 1100349, upload-time = "2026-04-30T23:24:18.992Z" }, + { url = "https://files.pythonhosted.org/packages/ab/89/29abcb1fe18a429cda60c6e0bbd1d6e90499339842a2f548d7567542357e/ast_serialize-0.3.0-cp314-cp314t-win_arm64.whl", hash = "sha256:670f177188d128fb7f9f15b5ad0e1b553d22c34e3f584dcb83eb8077600437f0", size = 1072895, upload-time = "2026-04-30T23:24:20.706Z" }, + { url = "https://files.pythonhosted.org/packages/bc/93/72abad83966ed6235647c9f956417dc1e17e997696388521910e3d1fa3f4/ast_serialize-0.3.0-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:2ec2fafa5e4313cc8feed96e436ebe19ac7bc6fa41fbc2827e826c48b9e4c3a9", size = 1190024, upload-time = "2026-04-30T23:24:22.486Z" }, + { url = "https://files.pythonhosted.org/packages/85/4f/eb88584b2f0234e581762011208ca203252bf6c98e59b4769daa571f3576/ast_serialize-0.3.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:ef6d3c08b7b4cd29b48410338e134764a00e76d25841eb02c1084e868c888ecc", size = 1178633, upload-time = "2026-04-30T23:24:24.35Z" }, + { url = "https://files.pythonhosted.org/packages/56/51/cf1ec1ff3e616373d0dcbd5fad502e0029dc541f13ab642259762a7d127f/ast_serialize-0.3.0-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3d841424f41b886e98044abc80769c14a956e6e5ccd5fb5b0d9f5ead72be18a4", size = 1241351, upload-time = "2026-04-30T23:24:25.987Z" }, + { url = "https://files.pythonhosted.org/packages/0d/44/68fcf50478cf1093f2d423f034ae06453122c8b415d8e21a44668eca485d/ast_serialize-0.3.0-cp39-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:d21453734ad39367ede5d37efe4f59f830ce1c09f432fc72a90e368f77a4a3e7", size = 1239582, upload-time = "2026-04-30T23:24:27.808Z" }, + { url = "https://files.pythonhosted.org/packages/9d/c1/a6c9fa284eceb5fc6f21347e968445a051d7ca2c4d34e6a04314646dbcee/ast_serialize-0.3.0-cp39-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:f5e110cdce2a347e1dd987529c88ef54d26f67848dce3eba1b3b2cc2cf085c94", size = 1448853, upload-time = "2026-04-30T23:24:29.534Z" }, + { url = "https://files.pythonhosted.org/packages/23/5f/8ad3829a09e4e8c5328a53ce7d4711d660944e3e164c5f6abcc2c8f27167/ast_serialize-0.3.0-cp39-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:3b6e23a98e57560a055f5c4b68700a0fd5ce483d2814c23140b3638c7f5d1e61", size = 1262204, upload-time = "2026-04-30T23:24:31.482Z" }, + { url = "https://files.pythonhosted.org/packages/25/13/44aa28d97f10e25247e8576b5f6b2795d4fa1a80acc88acc942c508d06f7/ast_serialize-0.3.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c1c9e763d70293d65ce1e1ea8c943140c68d0953f0268c7ee0998f2e07f77dd0", size = 1266458, upload-time = "2026-04-30T23:24:33.088Z" }, + { url = "https://files.pythonhosted.org/packages/d8/58/b3a8be3777cd3744324fd5cec0d80d37cd96fc7cbb0fb010e03dff1e870f/ast_serialize-0.3.0-cp39-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:4388a1796c228f1ce5c391426f7d21a0003ad3b47f677dbeded9bd1a85c7209f", size = 1308700, upload-time = "2026-04-30T23:24:34.657Z" }, + { url = "https://files.pythonhosted.org/packages/13/03/f8312d6b57f5471a9dc7946f22b8798a1fc296d38c25766223aacadec42c/ast_serialize-0.3.0-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:5283cdcc0c64c3d8b9b688dc6aaa012d9c0cf1380a7f774a6bae6a1c01b3205a", size = 1416724, upload-time = "2026-04-30T23:24:36.562Z" }, + { url = "https://files.pythonhosted.org/packages/50/5d/13fc3789a7abac00559da2e2e9f386db4612aa1f84fc53d09bf714c37545/ast_serialize-0.3.0-cp39-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:f5ef88cc5842a5d7a6ac09dc0d5fc2c98f5d276c1f076f866d55047ce886785b", size = 1515441, upload-time = "2026-04-30T23:24:38.018Z" }, + { url = "https://files.pythonhosted.org/packages/eb/b9/7ab43fc7a23b1f970281093228f5f79bed6edeed7a3e672bde6d7a832a58/ast_serialize-0.3.0-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:cc14bf402bdc0978594ecce783793de2c7470cd4f5cd7eb286ca97ed8ff7cba9", size = 1510522, upload-time = "2026-04-30T23:24:39.798Z" }, + { url = "https://files.pythonhosted.org/packages/56/ec/d75fc2b788d319f1fad77c14156896f31afdfc68af85b505e5bdebcb9592/ast_serialize-0.3.0-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:11eae0cf1b7b3e0678133cc2daa974ea972caf02eb4b3aa062af6fa9acd52c57", size = 1460917, upload-time = "2026-04-30T23:24:41.305Z" }, + { url = "https://files.pythonhosted.org/packages/95/74/f99c81193a2725911e1911ae567ed27c2f2419332c7f3537366f9d238cac/ast_serialize-0.3.0-cp39-abi3-win32.whl", hash = "sha256:2db3dd99de5e6a5a11d7dda73de8750eb6e5baaf25245adf7bdcfe64b6108ae2", size = 1067804, upload-time = "2026-04-30T23:24:43.091Z" }, + { url = "https://files.pythonhosted.org/packages/16/81/76af00c47daa151e89f98ae21fbbcb2840aaa9f5766579c4da76a3c57188/ast_serialize-0.3.0-cp39-abi3-win_amd64.whl", hash = "sha256:a2cd125adccf7969470621905d302750cd25951f22ea430d9a25b7be031e5549", size = 1105561, upload-time = "2026-04-30T23:24:44.578Z" }, + { url = "https://files.pythonhosted.org/packages/bd/46/d3ec57ad500f598d1554bd14ce4df615960549ab2844961bc4e1f5fbd174/ast_serialize-0.3.0-cp39-abi3-win_arm64.whl", hash = "sha256:0dd00da29985f15f50dc35728b7e1e7c84507bccfea1d9914738530f1c72238a", size = 1077165, upload-time = "2026-04-30T23:24:46.377Z" }, +] + [[package]] name = "asttokens" version = "3.0.0" @@ -733,7 +772,7 @@ requires-dist = [ { name = "sphinx-gallery", marker = "extra == 'docs'" }, { name = "sphinxcontrib-mermaid", marker = "extra == 'docs'" }, { name = "syrupy", marker = "extra == 'test'", specifier = ">=5" }, - { name = "typing-extensions" }, + { name = "typing-extensions", specifier = ">=4.13" }, ] provides-extras = ["array", "dev", "docs", "test"] @@ -1482,6 +1521,79 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/82/3d/14ce75ef66813643812f3093ab17e46d3a206942ce7376d31ec2d36229e7/lark-1.3.1-py3-none-any.whl", hash = "sha256:c629b661023a014c37da873b4ff58a817398d12635d3bbb2c5a03be7fe5d1e12", size = 113151, upload-time = "2025-10-27T18:25:54.882Z" }, ] +[[package]] +name = "librt" +version = "0.11.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/40/08/9e7f6b5d2b5bed6ad055cdd5925f192bb403a51280f86b56554d9d0699a2/librt-0.11.0.tar.gz", hash = "sha256:075dc3ef4458a278e0195cbf6ac9d38808d9b906c5a6c7f7f79c3888276a3fb1", size = 200139, upload-time = "2026-05-10T18:17:25.138Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fe/87/2bf31fe17587b29e3f93ec31421e2b1e1c3e349b8bf6c7c313dbad1d5340/librt-0.11.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:93d95bd45b7d58343d8b90d904450a545144eec19a002511163426f8ab1fae29", size = 141092, upload-time = "2026-05-10T18:15:34.795Z" }, + { url = "https://files.pythonhosted.org/packages/cf/08/5c5bf772920b7ebac6e32bc91a643e0ab3870199c0b542356d3baa83970a/librt-0.11.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:4ee278c769a713638cdacd4c0436d72156e75df3ebc0166ab2b9dc43acc386c9", size = 142035, upload-time = "2026-05-10T18:15:36.242Z" }, + { url = "https://files.pythonhosted.org/packages/06/20/662a03d254e5b000d838e8b345d83303ddb768c080fd488e40634c0fa66b/librt-0.11.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f230cb1cbc9faaa616f9a678f530ebcf186e414b6bcbd88b960e4ba1b92428d5", size = 475022, upload-time = "2026-05-10T18:15:37.56Z" }, + { url = "https://files.pythonhosted.org/packages/de/f3/aa81523e45184c6ec23dc7f63263362ec55f80a09d424c012359ecbe7e35/librt-0.11.0-cp311-cp311-manylinux2014_i686.manylinux_2_17_i686.manylinux_2_28_i686.whl", hash = "sha256:5d63c855d86938d9de93e265c9bd8c705b51ec494de5738340ee93767a686e4b", size = 467273, upload-time = "2026-05-10T18:15:39.182Z" }, + { url = "https://files.pythonhosted.org/packages/6b/6f/59c74b560ca8853834d5501d589c8a2519f4184f273a085ffd0f37a1cc47/librt-0.11.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:993f028be9e96a08d31df3479ac80d99be374d17f3b78e4796b3fd3c913d4e89", size = 497083, upload-time = "2026-05-10T18:15:40.634Z" }, + { url = "https://files.pythonhosted.org/packages/fe/7b/5aa4d2c9600a719401160bf7055417df0b2a47439b9d88286ce45e56b65f/librt-0.11.0-cp311-cp311-manylinux_2_34_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:258d73a0aa66a055e65b2e4d1b8cdb23b9d132c5bb915d9547d804fcaed116cc", size = 489139, upload-time = "2026-05-10T18:15:41.934Z" }, + { url = "https://files.pythonhosted.org/packages/d6/31/9143803d7da6856a69153785768c4936864430eec0fd9461c3ea527d9922/librt-0.11.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:0827efe7854718f04aaddf6496e96960a956e676fe1d0f04eb41511fd8ad06d5", size = 508442, upload-time = "2026-05-10T18:15:43.206Z" }, + { url = "https://files.pythonhosted.org/packages/2f/5a/bce08184488426bda4ccc2c4964ac048c8f68ae89bd7120082eef4233cfd/librt-0.11.0-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:7753e57d6e12d019c0d8786f1c09c709f4c3fcc57c3887b24e36e6c06ec938b7", size = 514230, upload-time = "2026-05-10T18:15:44.761Z" }, + { url = "https://files.pythonhosted.org/packages/89/8c/bb5e213d254b7505a0e658da199d8ab719086632ce09eef311ab27976523/librt-0.11.0-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:11bd19822431cc21af9f27374e7ae2e58103c7d98bda823536a6c47f6bb2bb3d", size = 494231, upload-time = "2026-05-10T18:15:46.308Z" }, + { url = "https://files.pythonhosted.org/packages/9d/fb/541cdad5b1ab1300398c74c4c9a497b88e5074c21b1244c8f49731d3a284/librt-0.11.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:22bdf239b219d3993761a148ffa134b19e52e9989c84f845d5d7b71d70a17412", size = 537585, upload-time = "2026-05-10T18:15:47.629Z" }, + { url = "https://files.pythonhosted.org/packages/8f/f2/464bb69295c320cb06bddb4f14a4ec67934ee14b2bffb12b19fb7ab287ba/librt-0.11.0-cp311-cp311-win32.whl", hash = "sha256:46c60b61e308eb535fbd6fa622b1ee1bb2815691c1ad9c98bf7b84952ec3bc8d", size = 100509, upload-time = "2026-05-10T18:15:49.157Z" }, + { url = "https://files.pythonhosted.org/packages/6d/e7/a17ee1788f9e4fbf548c19f4afa07c92089b9e24fef6cb2410863781ef4c/librt-0.11.0-cp311-cp311-win_amd64.whl", hash = "sha256:902e546ff044f579ff1c953ff5fce97b636fe9e3943996b2177710c6ef076f73", size = 118628, upload-time = "2026-05-10T18:15:50.345Z" }, + { url = "https://files.pythonhosted.org/packages/cc/c7/6c766214f9f9903bcfcfbef97d807af8d8f5aa3502d247858ab17582d212/librt-0.11.0-cp311-cp311-win_arm64.whl", hash = "sha256:65ac3bc20f78aa0ee5ae84baa68917f89fef4af63e941084dd019a0d0e749f0c", size = 103122, upload-time = "2026-05-10T18:15:52.068Z" }, + { url = "https://files.pythonhosted.org/packages/8b/d0/07c77e067f0838949b43bd89232c29d72efebb9d2801a9750184eb706b71/librt-0.11.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b87504f1690a23b9a2cca841191a04f83895d4fc2dd04df91d82b1a04ca2ad46", size = 144147, upload-time = "2026-05-10T18:15:53.227Z" }, + { url = "https://files.pythonhosted.org/packages/7a/24/8493538fa4f62f982686398a5b8f68008138a75086abdea19ade64bf4255/librt-0.11.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40071fc5fe0ce8daa6de616702314a01e1250711682b0523d6ab8d4525910cb3", size = 143614, upload-time = "2026-05-10T18:15:54.657Z" }, + { url = "https://files.pythonhosted.org/packages/ff/1e/f8bad050810d9171f34a1648ed910e56814c2ba61639f2bd53c6377ae24b/librt-0.11.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:137e79445c896a0ea7b265f52d23954e05b64222ee1af69e2cb34219067cbb67", size = 485538, upload-time = "2026-05-10T18:15:56.117Z" }, + { url = "https://files.pythonhosted.org/packages/c0/fe/3594ebfbaf03084ba4b120c9ba5c3183fd938a48725e9bbe6ff0a5159ad8/librt-0.11.0-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.manylinux_2_28_i686.whl", hash = "sha256:cca6644054e78746d8d4ef238681f9c34ff8b584fe6b988ecebb8db3b15e622a", size = 479623, upload-time = "2026-05-10T18:15:57.544Z" }, + { url = "https://files.pythonhosted.org/packages/b0/da/5d1876984b3746c85dbd219dbfcb73c85f54ee263fd32e5b2a632ec14571/librt-0.11.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d5b0eea49f5562861ee8d757a32ef7d559c1d35be2aaaa1ec28941d74c9ffc8a", size = 513082, upload-time = "2026-05-10T18:15:58.805Z" }, + { url = "https://files.pythonhosted.org/packages/19/6e/55bdf5d5ca00c3e18430690bf2c953d8d3ffd3c337418173d33dec985dc9/librt-0.11.0-cp312-cp312-manylinux_2_34_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:0d1029d7e1ae1a7e647ed6fb5df8c4ce2dffefb7a9f5fd1376a4554d96dac09f", size = 508105, upload-time = "2026-05-10T18:16:00.2Z" }, + { url = "https://files.pythonhosted.org/packages/07/10/f1f23a7c595ee90ece4d35c851e5d104b1311a887ed1b4ac4c35bbd13da8/librt-0.11.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:bc3ce6b33c5828d9e80592011a5c584cb2ce86edbc4088405f70da47dc1d1b3b", size = 522268, upload-time = "2026-05-10T18:16:01.708Z" }, + { url = "https://files.pythonhosted.org/packages/b6/02/5720f5697a7f54b78b3aefbe20df3a48cedcff1276618c4aa481177942ed/librt-0.11.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:936c5995f3514a42111f20099397d8177c79b4d7e70961e396c6f5a0a3566766", size = 527348, upload-time = "2026-05-10T18:16:03.496Z" }, + { url = "https://files.pythonhosted.org/packages/50/db/b4a47c6f91db4ff76348a0b3dd0cc65e090a078b765a810a62ff9434c3d3/librt-0.11.0-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:9bc0ca6ad9381cbe8e4aa6e5726e4c80c78115a6e9723c599ed1d73e092bc49d", size = 516294, upload-time = "2026-05-10T18:16:05.173Z" }, + { url = "https://files.pythonhosted.org/packages/9e/58/9384b2f4eb1ed1d273d40948a7c5c4b2360213b402ef3be4641c06299f9c/librt-0.11.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:070aa8c26c0a74774317a72df8851facc7f0f012a5b406557ac56992d92e1ec8", size = 553608, upload-time = "2026-05-10T18:16:06.839Z" }, + { url = "https://files.pythonhosted.org/packages/21/7b/5aa8848a7c6a9278c79375146da1812e695754ceec5f005e6043461a7315/librt-0.11.0-cp312-cp312-win32.whl", hash = "sha256:6bf14feb84b05ae945277395451998c89c54d0def4070eb5c08de544930b245a", size = 101879, upload-time = "2026-05-10T18:16:08.103Z" }, + { url = "https://files.pythonhosted.org/packages/37/33/8a745436944947575b584231750a41417de1a38cf6a2e9251d1065651c09/librt-0.11.0-cp312-cp312-win_amd64.whl", hash = "sha256:75672f0bc524ede266287d532d7923dbce94c7514ad07627bac3d0c6d92cc4d9", size = 119831, upload-time = "2026-05-10T18:16:09.174Z" }, + { url = "https://files.pythonhosted.org/packages/59/67/a6739ac96e28b7855808bdb0370e250606104a859750d209e5a0716fe7ab/librt-0.11.0-cp312-cp312-win_arm64.whl", hash = "sha256:2f10cf143e4a9bb0f4f5af568a00df94a2d69ef41c2579584454bb0fe5cc642c", size = 103470, upload-time = "2026-05-10T18:16:10.369Z" }, + { url = "https://files.pythonhosted.org/packages/82/61/e59168d4d0bf2bf90f4f0caf7a001bfc60254c3af4586013b04dc3ef517b/librt-0.11.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:78dc31f7fdfe9c9d0eb0e8f42d139db230e826415bbcabd9f0e9faaaee909894", size = 144119, upload-time = "2026-05-10T18:16:11.771Z" }, + { url = "https://files.pythonhosted.org/packages/61/fd/caa1d60b12f7dd79ccea23054e06eeaebe266a5f52c40a6b651069200ce5/librt-0.11.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:fa475675db22290c3158e1d42326d0f5a65f04f44a0e68c3630a25b53560fb9c", size = 143565, upload-time = "2026-05-10T18:16:13.334Z" }, + { url = "https://files.pythonhosted.org/packages/b8/a9/dc744f5c2b4978d48db970be29f22716d3413d28b14ad99740817315cf2c/librt-0.11.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:621db29691044bdeda22e789e482e1b0f3a985d90e3426c9c6d17606416205ea", size = 485395, upload-time = "2026-05-10T18:16:14.729Z" }, + { url = "https://files.pythonhosted.org/packages/8f/21/7f8e97a1e4dae952a5a95948f6f8507a173bc1e669f54340bba6ca1ca31b/librt-0.11.0-cp313-cp313-manylinux2014_i686.manylinux_2_17_i686.manylinux_2_28_i686.whl", hash = "sha256:a9010e2ed5b3a9e158c5fd966b3ab7e834bb3d3aacc8f66c91dd4b57a3799230", size = 479383, upload-time = "2026-05-10T18:16:16.321Z" }, + { url = "https://files.pythonhosted.org/packages/a6/6d/d8ee9c114bebf2c50e29ec2aa940826fccb62a645c3e4c18760987d0e16d/librt-0.11.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7c39513d8b7477a2e1ed8c43fc21c524e8d5a0f8d4e8b7b074dbdbe7820a08e2", size = 513010, upload-time = "2026-05-10T18:16:17.647Z" }, + { url = "https://files.pythonhosted.org/packages/f0/43/0b5708af2bd30a46400e72ba6bdaa8f066f15fb9a688527e34220e8d6c06/librt-0.11.0-cp313-cp313-manylinux_2_34_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:7aef3cf1d5af86e770ab04bfd993dfc4ae8b8c17f66fb77dd4a7d50de7bbb1a3", size = 508433, upload-time = "2026-05-10T18:16:19.309Z" }, + { url = "https://files.pythonhosted.org/packages/4a/50/356187247d09013490481033183b3532b58acf8028bcb34b2b56a375c9b2/librt-0.11.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:557183ddc36babe46b27dd60facbd5adb4492181a5be887587d57cda6e092f21", size = 522595, upload-time = "2026-05-10T18:16:20.642Z" }, + { url = "https://files.pythonhosted.org/packages/40/e7/c6ac4240899c7f3248079d5a9900debe0dadb3fdeaf856684c987105ba47/librt-0.11.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:83d3e1f72bd42f6c5c0b7daec530c3f829bd02db42c70b8ddf0c2d90a2459930", size = 527255, upload-time = "2026-05-10T18:16:22.352Z" }, + { url = "https://files.pythonhosted.org/packages/eb/b5/a81322dbeedeeaf9c1ee6f001734d28a09d8383ac9e6779bc24bbd0743c6/librt-0.11.0-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:4ce1f21fbe589bc1afd7872dece84fb0e1144f794a288e58a10d2c54a55c43be", size = 516847, upload-time = "2026-05-10T18:16:23.627Z" }, + { url = "https://files.pythonhosted.org/packages/ae/66/6e6323787d592b55204a42595ff1102da5115601b53a7e9ddebc889a6da5/librt-0.11.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:970b09f7044ea2b64c9da42fd3d335666518cfd1c6e8a182c95da73d0214b41e", size = 553920, upload-time = "2026-05-10T18:16:25.025Z" }, + { url = "https://files.pythonhosted.org/packages/9c/21/623f8ca230857102066d9ca8c6c1734995908c4d0d1bee7bb2ef0021cb33/librt-0.11.0-cp313-cp313-win32.whl", hash = "sha256:78fddc31cd4d3caa897ad5d31f856b1faadc9474021ad6cb182b9018793e254e", size = 101898, upload-time = "2026-05-10T18:16:26.649Z" }, + { url = "https://files.pythonhosted.org/packages/b3/1d/b4ebd44dd723f768469007515cb92251e0ae286c94c140f374801140fa74/librt-0.11.0-cp313-cp313-win_amd64.whl", hash = "sha256:8ca8aa88751a775870b764e93bad5135385f563cb8dcee399abf034ea4d3cb47", size = 119812, upload-time = "2026-05-10T18:16:27.859Z" }, + { url = "https://files.pythonhosted.org/packages/3b/e4/b2f4ca7965ca373b491cdb4bc25cdb30c1649ca81a8782056a83850292a9/librt-0.11.0-cp313-cp313-win_arm64.whl", hash = "sha256:96f044bb325fd9cf1a723015638c219e9143f0dfbc0ca54c565df2b7fc748b44", size = 103448, upload-time = "2026-05-10T18:16:29.066Z" }, + { url = "https://files.pythonhosted.org/packages/29/eb/dbce197da4e227779e56b5735f2decc3eb36e55a1cdbf1bd65d6639d76c1/librt-0.11.0-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:4a017a95e5837dc15a8c5661d60e05daa96b90908b1aa6b7acdf443cd25c8ebd", size = 143345, upload-time = "2026-05-10T18:16:30.674Z" }, + { url = "https://files.pythonhosted.org/packages/76/a3/254bebd0c11c8ba684018efb8006ff22e466abce445215cca6c778e7d9de/librt-0.11.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:b1ecbd9819deccc39b7542bf4d2a740d8a620694d39989e58661d3763458f8d4", size = 143131, upload-time = "2026-05-10T18:16:32.037Z" }, + { url = "https://files.pythonhosted.org/packages/f1/3f/f77d6122d21ac7bf6ae8a7dfced1bd2a7ac545d3273ebdcaf8042f6d619f/librt-0.11.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7da327dacd7be8f8ec36547373550744a3cc0e536d54665cd83f8bcd961200e8", size = 477024, upload-time = "2026-05-10T18:16:33.493Z" }, + { url = "https://files.pythonhosted.org/packages/ac/0a/2c996dadebaa7d9bbbd43ef2d4f3e66b6da545f838a41694ef6172cebec8/librt-0.11.0-cp314-cp314-manylinux2014_i686.manylinux_2_17_i686.manylinux_2_28_i686.whl", hash = "sha256:0dc56b1f8d06e60db362cc3fdae206681817f86ce4725d34511473487f12a34b", size = 474221, upload-time = "2026-05-10T18:16:34.864Z" }, + { url = "https://files.pythonhosted.org/packages/0a/7e/f5d92af8486b8272c23b3e686b46ff72d89c8169585eb61eef01a2ac7147/librt-0.11.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:05fb8fb2ab90e21c8d12ea240d744ad514da9baf381ebfa70d91d20d21713175", size = 505174, upload-time = "2026-05-10T18:16:36.705Z" }, + { url = "https://files.pythonhosted.org/packages/af/1a/cb0734fe86398eb33193ab753b7326255c74cac5eb09e76b9b16536e7adb/librt-0.11.0-cp314-cp314-manylinux_2_34_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:cae74872be221df4374d10fec61f93ed1513b9546ea84f2c0bf73ab3e9bd0b03", size = 497216, upload-time = "2026-05-10T18:16:38.418Z" }, + { url = "https://files.pythonhosted.org/packages/18/06/094820f91558b66e29943c0ec41c9914f460f48dd51fc503c3101e10842d/librt-0.11.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:32bcc918c0148eb7e3d57385125bac7e5f9e4359d05f07448b09f6f778c2f31c", size = 513921, upload-time = "2026-05-10T18:16:39.848Z" }, + { url = "https://files.pythonhosted.org/packages/0b/c2/00de9018871a282f530cacb457d5ec0428f6ac7e6fedde9aff7468d9fb04/librt-0.11.0-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:f9743fc99135d5f78d2454435615f6dec0473ca507c26ce9d92b10b562a280d3", size = 520850, upload-time = "2026-05-10T18:16:41.471Z" }, + { url = "https://files.pythonhosted.org/packages/51/9d/64631832348fd1834fb3a61b996434edddaaf25a31d03b0a76273159d2cf/librt-0.11.0-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:5ba067f4aadae8fda802d91d2124c90c42195ff32d9161d3549e6d05cfe26f96", size = 504237, upload-time = "2026-05-10T18:16:43.15Z" }, + { url = "https://files.pythonhosted.org/packages/a5/ec/ae5525eb16edc827a044e7bb8777a455ff95d4bca9379e7e6bddd7383647/librt-0.11.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:de3bf945454d032f9e390b85c4072e0a0570bf825421c8be0e71209fa65e1abe", size = 546261, upload-time = "2026-05-10T18:16:44.408Z" }, + { url = "https://files.pythonhosted.org/packages/5a/09/adce371f27ca039411da9659f7430fcc2ba6cd0c7b3e4467a0f091be7fa9/librt-0.11.0-cp314-cp314-win32.whl", hash = "sha256:d2277a05f6dcb9fd13db9566aac4fabd68c3ea1ea46ee5567d4eef8efa495a2f", size = 96965, upload-time = "2026-05-10T18:16:46.039Z" }, + { url = "https://files.pythonhosted.org/packages/d6/ee/8ac720d98548f173c7ce2e632a7ca94673f74cacd5c8162a84af5b35958a/librt-0.11.0-cp314-cp314-win_amd64.whl", hash = "sha256:ab73e8db5e3f564d812c1f5c3a175930a5f9bc96ccb5e3b22a34d7858b401cf7", size = 115151, upload-time = "2026-05-10T18:16:47.133Z" }, + { url = "https://files.pythonhosted.org/packages/94/20/c900cf14efeb09b6bef2b2dff20779f73464b97fd58d1c6bccc379588ae3/librt-0.11.0-cp314-cp314-win_arm64.whl", hash = "sha256:aea3caa317752e3a466fa8af45d91ee0ea8c7fdd96e42b0a8dd9b76a7931eba1", size = 98850, upload-time = "2026-05-10T18:16:48.597Z" }, + { url = "https://files.pythonhosted.org/packages/0c/71/944bfe4b64e12abffcd3c15e1cce07f72f3d55655083786285f4dedeb532/librt-0.11.0-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:d1b36540d7aaf9b9101b3a6f376c8d8e9f7a9aec93ed05918f2c69d493ffef72", size = 151138, upload-time = "2026-05-10T18:16:49.839Z" }, + { url = "https://files.pythonhosted.org/packages/b6/10/99e64a5c86989357fda078c8143c533389585f6473b7439172dd8f3b3b2d/librt-0.11.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:efbb343ab2ce3540f4ecbe6315d677ed70f37cd9a72b1e58066c918ca83acbaa", size = 151976, upload-time = "2026-05-10T18:16:51.062Z" }, + { url = "https://files.pythonhosted.org/packages/21/31/5072ad880946d83e5ea4147d6d018c78eefce85b77819b19bdd0ee229435/librt-0.11.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:aa0dd688aab3f7914d3e6e5e3554978e0383312fb8e771d84be008a35b9ee548", size = 557927, upload-time = "2026-05-10T18:16:52.632Z" }, + { url = "https://files.pythonhosted.org/packages/5e/8d/70b5fb7cfbab60edbe7381614ab985da58e144fbf465c86d44c95f43cdca/librt-0.11.0-cp314-cp314t-manylinux2014_i686.manylinux_2_17_i686.manylinux_2_28_i686.whl", hash = "sha256:f5fb36b8c6c63fdcbb1d526d94c0d1331610d43f4118cc1beb4efef4f3faacb2", size = 539698, upload-time = "2026-05-10T18:16:53.934Z" }, + { url = "https://files.pythonhosted.org/packages/fa/a3/ba3495a0b3edbd24a4cae0d1d3c64f39a9fc45d06e812101289b50c1a619/librt-0.11.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4a9a237d13addb93715b6fee74023d5ee3469b53fce527626c0e088aa585805f", size = 577162, upload-time = "2026-05-10T18:16:55.589Z" }, + { url = "https://files.pythonhosted.org/packages/f7/db/36e25fb81f99937ff1b96612a1dc9fd66f039cb9cc3aee12c01fac31aab9/librt-0.11.0-cp314-cp314t-manylinux_2_34_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:5ddd17bd87b2c56ddd60e546a7984a2e64c4e8eab92fb4cf3830a48ad5469d51", size = 566494, upload-time = "2026-05-10T18:16:56.975Z" }, + { url = "https://files.pythonhosted.org/packages/33/0d/3f622b47f0b013eeb9cf4cc07ae9bfe378d832a4eec998b2b209fe84244d/librt-0.11.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:bd43992b4473d42f12ff9e68326079f0696d9d4e6000e8f39a0238d482ba6ee2", size = 596858, upload-time = "2026-05-10T18:16:58.374Z" }, + { url = "https://files.pythonhosted.org/packages/a9/02/71b90bc93039c46a2000651f6ad60122b114c8f54c4ad306e0e96f5b75ad/librt-0.11.0-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:f8e3e8056dd674e279741485e2e512d6e9a751c7455809d0114e6ebf8d781085", size = 590318, upload-time = "2026-05-10T18:16:59.676Z" }, + { url = "https://files.pythonhosted.org/packages/04/04/418cb3f75621e2b761fb1ab0f017f4d70a1a72a6e7c74ee4f7e8d198c2f3/librt-0.11.0-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:c1f708d8ae9c56cf38a903c44297243d2ec83fd82b396b977e0144a3e76217e3", size = 575115, upload-time = "2026-05-10T18:17:01.007Z" }, + { url = "https://files.pythonhosted.org/packages/cc/2c/5a2183ac58dd911f26b5d7e7d7d8f1d87fcecdddd99d6c12169a258ff62c/librt-0.11.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:0add982e0e7b9fc14cf4b33789d5f13f66581889b88c2f58099f6ce8f92617bd", size = 617918, upload-time = "2026-05-10T18:17:02.682Z" }, + { url = "https://files.pythonhosted.org/packages/15/1f/dc6771a52592a4451be6effa200cbfc9cec61e4393d3033d81a9d307961d/librt-0.11.0-cp314-cp314t-win32.whl", hash = "sha256:2b481d846ac894c4e8403c5fd0e87c5d11d6499e404b474602508a224ff531c8", size = 103562, upload-time = "2026-05-10T18:17:03.99Z" }, + { url = "https://files.pythonhosted.org/packages/62/4a/7d1415567027286a75ba1093ec4aca11f073e0f559c530cf3e0a757ad55c/librt-0.11.0-cp314-cp314t-win_amd64.whl", hash = "sha256:28edb433edde181112a908c78907af28f964eabc15f4dd16c9d66c834302677c", size = 124327, upload-time = "2026-05-10T18:17:05.465Z" }, + { url = "https://files.pythonhosted.org/packages/ce/62/b40b382fa0c66fee1478073eb8db352a4a6beda4a1adccf1df911d8c289c/librt-0.11.0-cp314-cp314t-win_arm64.whl", hash = "sha256:dee008f20b542e3cd162ba338a7f9ec0f6d23d395f66fe8aeeec3c9d067ea253", size = 102572, upload-time = "2026-05-10T18:17:06.809Z" }, +] + [[package]] name = "line-profiler" version = "5.0.0" @@ -1841,40 +1953,53 @@ wheels = [ [[package]] name = "mypy" -version = "1.18.2" +version = "2.1.0" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "ast-serialize" }, + { name = "librt", marker = "platform_python_implementation != 'PyPy'" }, { name = "mypy-extensions" }, { name = "pathspec" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/c0/77/8f0d0001ffad290cef2f7f216f96c814866248a0b92a722365ed54648e7e/mypy-1.18.2.tar.gz", hash = "sha256:06a398102a5f203d7477b2923dda3634c36727fa5c237d8f859ef90c42a9924b", size = 3448846, upload-time = "2025-09-19T00:11:10.519Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/88/87/cafd3ae563f88f94eec33f35ff722d043e09832ea8530ef149ec1efbaf08/mypy-1.18.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:807d9315ab9d464125aa9fcf6d84fde6e1dc67da0b6f80e7405506b8ac72bc7f", size = 12731198, upload-time = "2025-09-19T00:09:44.857Z" }, - { url = "https://files.pythonhosted.org/packages/0f/e0/1e96c3d4266a06d4b0197ace5356d67d937d8358e2ee3ffac71faa843724/mypy-1.18.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:776bb00de1778caf4db739c6e83919c1d85a448f71979b6a0edd774ea8399341", size = 11817879, upload-time = "2025-09-19T00:09:47.131Z" }, - { url = "https://files.pythonhosted.org/packages/72/ef/0c9ba89eb03453e76bdac5a78b08260a848c7bfc5d6603634774d9cd9525/mypy-1.18.2-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1379451880512ffce14505493bd9fe469e0697543717298242574882cf8cdb8d", size = 12427292, upload-time = "2025-09-19T00:10:22.472Z" }, - { url = "https://files.pythonhosted.org/packages/1a/52/ec4a061dd599eb8179d5411d99775bec2a20542505988f40fc2fee781068/mypy-1.18.2-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1331eb7fd110d60c24999893320967594ff84c38ac6d19e0a76c5fd809a84c86", size = 13163750, upload-time = "2025-09-19T00:09:51.472Z" }, - { url = "https://files.pythonhosted.org/packages/c4/5f/2cf2ceb3b36372d51568f2208c021870fe7834cf3186b653ac6446511839/mypy-1.18.2-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:3ca30b50a51e7ba93b00422e486cbb124f1c56a535e20eff7b2d6ab72b3b2e37", size = 13351827, upload-time = "2025-09-19T00:09:58.311Z" }, - { url = "https://files.pythonhosted.org/packages/c8/7d/2697b930179e7277529eaaec1513f8de622818696857f689e4a5432e5e27/mypy-1.18.2-cp311-cp311-win_amd64.whl", hash = "sha256:664dc726e67fa54e14536f6e1224bcfce1d9e5ac02426d2326e2bb4e081d1ce8", size = 9757983, upload-time = "2025-09-19T00:10:09.071Z" }, - { url = "https://files.pythonhosted.org/packages/07/06/dfdd2bc60c66611dd8335f463818514733bc763e4760dee289dcc33df709/mypy-1.18.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:33eca32dd124b29400c31d7cf784e795b050ace0e1f91b8dc035672725617e34", size = 12908273, upload-time = "2025-09-19T00:10:58.321Z" }, - { url = "https://files.pythonhosted.org/packages/81/14/6a9de6d13a122d5608e1a04130724caf9170333ac5a924e10f670687d3eb/mypy-1.18.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:a3c47adf30d65e89b2dcd2fa32f3aeb5e94ca970d2c15fcb25e297871c8e4764", size = 11920910, upload-time = "2025-09-19T00:10:20.043Z" }, - { url = "https://files.pythonhosted.org/packages/5f/a9/b29de53e42f18e8cc547e38daa9dfa132ffdc64f7250e353f5c8cdd44bee/mypy-1.18.2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5d6c838e831a062f5f29d11c9057c6009f60cb294fea33a98422688181fe2893", size = 12465585, upload-time = "2025-09-19T00:10:33.005Z" }, - { url = "https://files.pythonhosted.org/packages/77/ae/6c3d2c7c61ff21f2bee938c917616c92ebf852f015fb55917fd6e2811db2/mypy-1.18.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:01199871b6110a2ce984bde85acd481232d17413868c9807e95c1b0739a58914", size = 13348562, upload-time = "2025-09-19T00:10:11.51Z" }, - { url = "https://files.pythonhosted.org/packages/4d/31/aec68ab3b4aebdf8f36d191b0685d99faa899ab990753ca0fee60fb99511/mypy-1.18.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:a2afc0fa0b0e91b4599ddfe0f91e2c26c2b5a5ab263737e998d6817874c5f7c8", size = 13533296, upload-time = "2025-09-19T00:10:06.568Z" }, - { url = "https://files.pythonhosted.org/packages/9f/83/abcb3ad9478fca3ebeb6a5358bb0b22c95ea42b43b7789c7fb1297ca44f4/mypy-1.18.2-cp312-cp312-win_amd64.whl", hash = "sha256:d8068d0afe682c7c4897c0f7ce84ea77f6de953262b12d07038f4d296d547074", size = 9828828, upload-time = "2025-09-19T00:10:28.203Z" }, - { url = "https://files.pythonhosted.org/packages/5f/04/7f462e6fbba87a72bc8097b93f6842499c428a6ff0c81dd46948d175afe8/mypy-1.18.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:07b8b0f580ca6d289e69209ec9d3911b4a26e5abfde32228a288eb79df129fcc", size = 12898728, upload-time = "2025-09-19T00:10:01.33Z" }, - { url = "https://files.pythonhosted.org/packages/99/5b/61ed4efb64f1871b41fd0b82d29a64640f3516078f6c7905b68ab1ad8b13/mypy-1.18.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:ed4482847168439651d3feee5833ccedbf6657e964572706a2adb1f7fa4dfe2e", size = 11910758, upload-time = "2025-09-19T00:10:42.607Z" }, - { url = "https://files.pythonhosted.org/packages/3c/46/d297d4b683cc89a6e4108c4250a6a6b717f5fa96e1a30a7944a6da44da35/mypy-1.18.2-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c3ad2afadd1e9fea5cf99a45a822346971ede8685cc581ed9cd4d42eaf940986", size = 12475342, upload-time = "2025-09-19T00:11:00.371Z" }, - { url = "https://files.pythonhosted.org/packages/83/45/4798f4d00df13eae3bfdf726c9244bcb495ab5bd588c0eed93a2f2dd67f3/mypy-1.18.2-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a431a6f1ef14cf8c144c6b14793a23ec4eae3db28277c358136e79d7d062f62d", size = 13338709, upload-time = "2025-09-19T00:11:03.358Z" }, - { url = "https://files.pythonhosted.org/packages/d7/09/479f7358d9625172521a87a9271ddd2441e1dab16a09708f056e97007207/mypy-1.18.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:7ab28cc197f1dd77a67e1c6f35cd1f8e8b73ed2217e4fc005f9e6a504e46e7ba", size = 13529806, upload-time = "2025-09-19T00:10:26.073Z" }, - { url = "https://files.pythonhosted.org/packages/71/cf/ac0f2c7e9d0ea3c75cd99dff7aec1c9df4a1376537cb90e4c882267ee7e9/mypy-1.18.2-cp313-cp313-win_amd64.whl", hash = "sha256:0e2785a84b34a72ba55fb5daf079a1003a34c05b22238da94fcae2bbe46f3544", size = 9833262, upload-time = "2025-09-19T00:10:40.035Z" }, - { url = "https://files.pythonhosted.org/packages/5a/0c/7d5300883da16f0063ae53996358758b2a2df2a09c72a5061fa79a1f5006/mypy-1.18.2-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:62f0e1e988ad41c2a110edde6c398383a889d95b36b3e60bcf155f5164c4fdce", size = 12893775, upload-time = "2025-09-19T00:10:03.814Z" }, - { url = "https://files.pythonhosted.org/packages/50/df/2cffbf25737bdb236f60c973edf62e3e7b4ee1c25b6878629e88e2cde967/mypy-1.18.2-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:8795a039bab805ff0c1dfdb8cd3344642c2b99b8e439d057aba30850b8d3423d", size = 11936852, upload-time = "2025-09-19T00:10:51.631Z" }, - { url = "https://files.pythonhosted.org/packages/be/50/34059de13dd269227fb4a03be1faee6e2a4b04a2051c82ac0a0b5a773c9a/mypy-1.18.2-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6ca1e64b24a700ab5ce10133f7ccd956a04715463d30498e64ea8715236f9c9c", size = 12480242, upload-time = "2025-09-19T00:11:07.955Z" }, - { url = "https://files.pythonhosted.org/packages/5b/11/040983fad5132d85914c874a2836252bbc57832065548885b5bb5b0d4359/mypy-1.18.2-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d924eef3795cc89fecf6bedc6ed32b33ac13e8321344f6ddbf8ee89f706c05cb", size = 13326683, upload-time = "2025-09-19T00:09:55.572Z" }, - { url = "https://files.pythonhosted.org/packages/e9/ba/89b2901dd77414dd7a8c8729985832a5735053be15b744c18e4586e506ef/mypy-1.18.2-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:20c02215a080e3a2be3aa50506c67242df1c151eaba0dcbc1e4e557922a26075", size = 13514749, upload-time = "2025-09-19T00:10:44.827Z" }, - { url = "https://files.pythonhosted.org/packages/25/bc/cc98767cffd6b2928ba680f3e5bc969c4152bf7c2d83f92f5a504b92b0eb/mypy-1.18.2-cp314-cp314-win_amd64.whl", hash = "sha256:749b5f83198f1ca64345603118a6f01a4e99ad4bf9d103ddc5a3200cc4614adf", size = 9982959, upload-time = "2025-09-19T00:10:37.344Z" }, - { url = "https://files.pythonhosted.org/packages/87/e3/be76d87158ebafa0309946c4a73831974d4d6ab4f4ef40c3b53a385a66fd/mypy-1.18.2-py3-none-any.whl", hash = "sha256:22a1748707dd62b58d2ae53562ffc4d7f8bcc727e8ac7cbc69c053ddc874d47e", size = 2352367, upload-time = "2025-09-19T00:10:15.489Z" }, +sdist = { url = "https://files.pythonhosted.org/packages/82/15/cca9d88503549ed6fedeaa1d448cdddd542ee8a490232d732e278036fbf2/mypy-2.1.0.tar.gz", hash = "sha256:81e76ad12c2d804512e9b13240d1588316531bfba07558286078bfbce9613633", size = 3898359, upload-time = "2026-05-11T18:37:36.237Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0a/a1/639f3024794a2a15899cb90707fe02e044c4412794c39c5769fd3df2e2ef/mypy-2.1.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:a683016b16fe2f572dc04c72be7ee0504ac1605a265d0200f5cea695fb788f41", size = 14691685, upload-time = "2026-05-11T18:33:27.973Z" }, + { url = "https://files.pythonhosted.org/packages/3b/08/9a585dea4325f20d8b80dc78623fa50d1fd2173b710f6237afd6ba6ab39b/mypy-2.1.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:1a293c534adb55271fef24a26da04b855540a8c13cc07bc5917b9fd2c394f2ca", size = 13555165, upload-time = "2026-05-11T18:32:16.107Z" }, + { url = "https://files.pythonhosted.org/packages/81/dc/7c42cc9c6cb01e8eb09961f1f738741d3e9c7e9d5c5b30ec69222625cd5f/mypy-2.1.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7406f4d048e71e576f5356d317e5b0a9e666dfd966bd99f9d14ca06e1a341538", size = 13994376, upload-time = "2026-05-11T18:32:39.256Z" }, + { url = "https://files.pythonhosted.org/packages/d4/fa/285946c33bce716e082c11dfeee9ee196eaf1f5042efb3581a31f9f205e4/mypy-2.1.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e0210d626fc8b31ccc90233754c7bc90e1f43205e85d96387f7db1285b55c398", size = 14864618, upload-time = "2026-05-11T18:34:49.765Z" }, + { url = "https://files.pythonhosted.org/packages/2b/83/82397f48af6c27e295d57979ded8490c9829040152cf7571b2f026aeb9a0/mypy-2.1.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:3712c20deed54e814eaaa825603bada8ea1c390670a397c95b98405347acc563", size = 15102063, upload-time = "2026-05-11T18:34:05.855Z" }, + { url = "https://files.pythonhosted.org/packages/40/68/b02dec39057b88eb03dc0aa854732e26e8361f34f9d0e20c7614967d1eba/mypy-2.1.0-cp311-cp311-win_amd64.whl", hash = "sha256:fcaa0e479066e31f7cceb6a3bea39cb22b2ff51a6b2f24f193d19179ba17c389", size = 11060564, upload-time = "2026-05-11T18:35:36.494Z" }, + { url = "https://files.pythonhosted.org/packages/cf/a8/ea3dcbef31f99b634f2ee23bb0321cbc8c1b388b76a861eb849f13c347dc/mypy-2.1.0-cp311-cp311-win_arm64.whl", hash = "sha256:0b1a5260c95aa443083f9ed3592662941951bca3d4ca224a5dc517c38b7cf666", size = 9966983, upload-time = "2026-05-11T18:37:14.139Z" }, + { url = "https://files.pythonhosted.org/packages/95/b1/55861beb5c339b44f9a2ba92df9e2cb1eeb4ae1eee674cdf7772c797778b/mypy-2.1.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:244358bf1c0da7722230bce60683d52e8e9fd030554926f15b747a84efb5b3af", size = 14874381, upload-time = "2026-05-11T18:37:31.784Z" }, + { url = "https://files.pythonhosted.org/packages/0b/b3/b7f770114b7d0ac92d0f76e8d93c2780844a70488a90e91821927850da86/mypy-2.1.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:4ec7c57657493c7a75534df2751c8ae2cda383c16ecc55d2106c54476b1b16f6", size = 13665501, upload-time = "2026-05-11T18:34:23.063Z" }, + { url = "https://files.pythonhosted.org/packages/b6/f3/8ae2037967e2126689a0c11d99e2b707134a565191e92c60ca2572aec60a/mypy-2.1.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d8161b6ff4392410023224f0969d17db93e1e154bc3e4ba62598e720723ae211", size = 14045750, upload-time = "2026-05-11T18:31:48.151Z" }, + { url = "https://files.pythonhosted.org/packages/a0/32/615eb5911859e43d054941b0d0a7d06cfa2870eba86529cf385b052b111c/mypy-2.1.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bf03e12003084a67395184d3eb8cbd6a489dc3655b5664b28c210a9e2403ab0b", size = 15061630, upload-time = "2026-05-11T18:37:06.898Z" }, + { url = "https://files.pythonhosted.org/packages/d4/03/4eafbfff8bfab1b87082741eae6e6a624028c984e6708b73bce2a8570c9d/mypy-2.1.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:20509760fd791c51579d573153407d226385ec1f8bcce55d730b354f3336bc22", size = 15288831, upload-time = "2026-05-11T18:31:18.07Z" }, + { url = "https://files.pythonhosted.org/packages/99/ee/919661478e5891a3c96e549c036e467e64563ab85995b10c53c8358e16a3/mypy-2.1.0-cp312-cp312-win_amd64.whl", hash = "sha256:6753d0c1fdd6b1a23b9e4f283ce80b2153b724adcb2653b20b85a8a28ac6436b", size = 11135228, upload-time = "2026-05-11T18:34:31.23Z" }, + { url = "https://files.pythonhosted.org/packages/24/0a/6a12b9782ca0831a553192f351679f4548abc9d19a7cc93bb7feb02084c7/mypy-2.1.0-cp312-cp312-win_arm64.whl", hash = "sha256:98ebb6589bb3b6d0c6f0c459d53ca55b8091fbc13d277c4041c885392e8195e8", size = 10040684, upload-time = "2026-05-11T18:36:48.199Z" }, + { url = "https://files.pythonhosted.org/packages/6e/dd/c7191469c777f07689c032a8f7326e393ea34c92d6d76eb7ce5ba57ea66d/mypy-2.1.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:35aac3bb114e03888f535d5eb51b8bafbb3266586b599da1940f9b1be3ec5bd5", size = 14852174, upload-time = "2026-05-11T18:31:38.929Z" }, + { url = "https://files.pythonhosted.org/packages/55/8c/aed55408879043d72bb9135f4d0d19a02b886dd569631e113e3d2706cb8d/mypy-2.1.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:8de55a8c861f2a49331f807be98d90caeceeef520bde13d43a160207f8af613e", size = 13651542, upload-time = "2026-05-11T18:36:04.636Z" }, + { url = "https://files.pythonhosted.org/packages/3a/8e/f371a824b1f1fa8ea6e3dbb8703d232977d572be2329554a3bc4d960302f/mypy-2.1.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5fdf2941a07434af755837d9880f7d7d25f1dacb1af9dcd4b9b66f2220a3024e", size = 14033929, upload-time = "2026-05-11T18:35:55.742Z" }, + { url = "https://files.pythonhosted.org/packages/94/21/f54be870d6dd53a82c674407e0f8eed7174b05ec78d42e5abd7b42e84fd5/mypy-2.1.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e195b817c13f02352a9c124301f9f30f078405444679b6753c1b96b6eed37285", size = 15039200, upload-time = "2026-05-11T18:33:10.281Z" }, + { url = "https://files.pythonhosted.org/packages/17/99/bf21748626a40ce59fd29a39386ab46afec88b7bd2f0fa6c3a97c995523f/mypy-2.1.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:5431d42af987ebd92ba2f71d45c85ed41d8e6ca9f5fd209a69f68f707d2469e5", size = 15272690, upload-time = "2026-05-11T18:32:07.205Z" }, + { url = "https://files.pythonhosted.org/packages/d6/d7/9e90d2cf47100bea550ed2bc7b0d4de3a62181d84d5e37da0003e8462637/mypy-2.1.0-cp313-cp313-win_amd64.whl", hash = "sha256:767fe8c66dc3e01e19e1737d4c38ebefead16125e1b8e58ad421903b376f5c65", size = 11147435, upload-time = "2026-05-11T18:33:56.477Z" }, + { url = "https://files.pythonhosted.org/packages/ec/46/e5c449e858798e35ffc90946282a27c62a77be743fe17480e4977374eb91/mypy-2.1.0-cp313-cp313-win_arm64.whl", hash = "sha256:ecfe70d43775ab99562ab128ce49854a362044c9f894961f68f898c23cb7429d", size = 10035052, upload-time = "2026-05-11T18:32:30.049Z" }, + { url = "https://files.pythonhosted.org/packages/b0/ca/b279a672e874aedd5498ae25f722dacc8aa86bbffb939b3f97cbb1cf6686/mypy-2.1.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:7354c5a7f69d9345c3d6e69921d57088eea3ddeeb6b20d34c1b3855b02c36ec2", size = 14848422, upload-time = "2026-05-11T18:35:45.984Z" }, + { url = "https://files.pythonhosted.org/packages/27/e6/3efe56c631d959b9b4454e208b0ac4b7f4f58b404c89f8bec7b49efdfc21/mypy-2.1.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:49890d4f76ac9e06ec117f9e09f3174da70a620a0c300953d8595c926e80947f", size = 13677374, upload-time = "2026-05-11T18:36:57.188Z" }, + { url = "https://files.pythonhosted.org/packages/84/7f/8107ea87a44fd1f1b59882442f033c9c3488c127201b1d1d15f1cbd6022e/mypy-2.1.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:761be68e023ef5d94678772396a8af1220030f80837a3afd8d0aef3b419666f4", size = 14055743, upload-time = "2026-05-11T18:35:18.361Z" }, + { url = "https://files.pythonhosted.org/packages/51/4d/b6d34db183133b83761b9199a82d31557cdbb70a380d8c3b3438e11882a3/mypy-2.1.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c90345fc182dc363b891350457ec69c35140858538f38b4540845afcc32b1aef", size = 15020937, upload-time = "2026-05-11T18:34:59.618Z" }, + { url = "https://files.pythonhosted.org/packages/ff/d7/f08360c691d758acb02f45022c34d98b92892f4ea756644e1000d4b9f3d8/mypy-2.1.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:b84802e7b5a6daf1f5e15bc9fcd7ddae77be13981ffab037f1c67bb84d67d135", size = 15253371, upload-time = "2026-05-11T18:36:41.081Z" }, + { url = "https://files.pythonhosted.org/packages/67/1b/09460a13719530a19bce27bd3bc8449e83569dd2ba7faf51c9c3c30c0b61/mypy-2.1.0-cp314-cp314-win_amd64.whl", hash = "sha256:022c771234936ceac541ebaf836fe9e2abeb3f5e09aff21588fe543ff006fe21", size = 11326429, upload-time = "2026-05-11T18:34:13.526Z" }, + { url = "https://files.pythonhosted.org/packages/40/62/75dbf0f82f7b6680340efc614af29dd0b3c17b8a4f1cd09b8bd2fd6bc814/mypy-2.1.0-cp314-cp314-win_arm64.whl", hash = "sha256:498207db725cec88829a6a5c2fc771205fd043719ef98bc49aba8fb9fc4e6d57", size = 10218799, upload-time = "2026-05-11T18:32:23.491Z" }, + { url = "https://files.pythonhosted.org/packages/b2/66/caca04ed7d972fb6eb6dd1ccd6df1de5c38fae8c5b3dc1c4e8e0d85ee6b9/mypy-2.1.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:7d5e5cad0efeba72b93cd17490cc0d69c5ac9ca132994fe3fb0314808aeeb83e", size = 15923458, upload-time = "2026-05-11T18:35:28.64Z" }, + { url = "https://files.pythonhosted.org/packages/ed/52/2d90cbe49d014b13ed7ff337930c30bad35893fe38a1e4641e756bb62191/mypy-2.1.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:ff715050c127d724fd260a2e666e7747fdd83511c0c47d449d98238970aef780", size = 14757697, upload-time = "2026-05-11T18:36:14.208Z" }, + { url = "https://files.pythonhosted.org/packages/ac/37/d98f4a14e081b238992d0ed96b6d39c7cc0148c9699eb71eaa68629665ea/mypy-2.1.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:82208da9e09414d520e912d3e462d454854bed0810b71540bb016dcbca7308fd", size = 15405638, upload-time = "2026-05-11T18:33:48.249Z" }, + { url = "https://files.pythonhosted.org/packages/a3/c2/15c46613b24a84fad2aea1248bf9619b99c2767ae9071fe224c179a0b7d4/mypy-2.1.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e79ebc1b904b84f0310dff7469655a9c36c7a68bddb37bdd42b67a332df61d08", size = 16215852, upload-time = "2026-05-11T18:32:50.296Z" }, + { url = "https://files.pythonhosted.org/packages/5c/90/9c16a57f482c76d25f6379762b56bbf65c711d8158cf271fb2802cfb0640/mypy-2.1.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:e583edc957cfb0deb142079162ae826f58449b116c1d442f2d91c69d9fced081", size = 16452695, upload-time = "2026-05-11T18:33:38.182Z" }, + { url = "https://files.pythonhosted.org/packages/0f/4c/215a4eeb63cacc5f17f516691ea7285d11e249802b942476bff15922a314/mypy-2.1.0-cp314-cp314t-win_amd64.whl", hash = "sha256:b33b6cd332695bba180d55e717a79d3038e479a2c49cc5eb3d53603409b9a5d7", size = 12866622, upload-time = "2026-05-11T18:34:39.945Z" }, + { url = "https://files.pythonhosted.org/packages/4b/50/1043e1db5f455ffe4c9ab22747cd8ca2bc492b1e4f4e21b130a44ee2b217/mypy-2.1.0-cp314-cp314t-win_arm64.whl", hash = "sha256:4f910fe825376a7b66ef7ca8c98e5a149e8cd64c19ae71d84047a74ee060d4e6", size = 10610798, upload-time = "2026-05-11T18:36:31.444Z" }, + { url = "https://files.pythonhosted.org/packages/0d/2a/13ca1f292f6db1b98ff495ef3467736b331621c5917cad984b7043e7348d/mypy-2.1.0-py3-none-any.whl", hash = "sha256:a663814603a5c563fb87a4f96fb473eeb30d1f5a4885afcf44f9db000a366289", size = 2693302, upload-time = "2026-05-11T18:31:29.246Z" }, ] [[package]] @@ -2291,11 +2416,11 @@ wheels = [ [[package]] name = "pathspec" -version = "0.12.1" +version = "1.1.1" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/ca/bc/f35b8446f4531a7cb215605d100cd88b7ac6f44ab3fc94870c120ab3adbf/pathspec-0.12.1.tar.gz", hash = "sha256:a482d51503a1ab33b1c67a6c3813a26953dbdc71c31dacaef9a838c4e29f5712", size = 51043, upload-time = "2023-12-10T22:30:45Z" } +sdist = { url = "https://files.pythonhosted.org/packages/5a/82/42f767fc1c1143d6fd36efb827202a2d997a375e160a71eb2888a925aac1/pathspec-1.1.1.tar.gz", hash = "sha256:17db5ecd524104a120e173814c90367a96a98d07c45b2e10c2f3919fff91bf5a", size = 135180, upload-time = "2026-04-27T01:46:08.907Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/cc/20/ff623b09d963f88bfde16306a54e12ee5ea43e9b597108672ff3a408aad6/pathspec-0.12.1-py3-none-any.whl", hash = "sha256:a0d503e138a4c123b27490a4f7beda6a01c6f288df0e4a8b79c7eb0dc7b4cc08", size = 31191, upload-time = "2023-12-10T22:30:43.14Z" }, + { url = "https://files.pythonhosted.org/packages/f1/d9/7fb5aa316bc299258e68c73ba3bddbc499654a07f151cba08f6153988714/pathspec-1.1.1-py3-none-any.whl", hash = "sha256:a00ce642f577bf7f473932318056212bc4f8bfdf53128c78bbd5af0b9b20b189", size = 57328, upload-time = "2026-04-27T01:46:07.06Z" }, ] [[package]] From 24952104af3b23a81970920cc3be7a452606a58f Mon Sep 17 00:00:00 2001 From: Saul Shanabrook Date: Mon, 31 Aug 2026 15:11:43 -0700 Subject: [PATCH 2/6] Address Egglog 3 review feedback --- AGENTS.md | 3 + Cargo.toml | 2 + docs/changelog.md | 9 +- docs/reference/egglog-translation.md | 18 +- experiments/param_eq/README.md | 2 + experiments/param_eq/results/manifest.toml | 6 - python/egglog/bindings.pyi | 4 +- python/egglog/builtins.py | 95 +----- python/egglog/declarations.py | 6 +- python/egglog/egraph.py | 29 +- python/egglog/egraph_state.py | 4 +- python/egglog/exp/param_eq/domain.py | 76 ++++- python/egglog/exp/param_eq/pipeline.py | 191 +++++++---- python/egglog/pretty.py | 58 +--- python/egglog/run_report.py | 2 +- .../test_array_api/test_jit[lda][expr].py | 11 +- python/tests/test_bindings.py | 7 +- python/tests/test_high_level.py | 299 ++++++------------ python/tests/test_run_report.py | 8 +- src/conversions.rs | 4 +- 20 files changed, 358 insertions(+), 476 deletions(-) delete mode 100644 experiments/param_eq/results/manifest.toml diff --git a/AGENTS.md b/AGENTS.md index f8828ade..1f3775c8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -26,6 +26,8 @@ - Prefer relative imports inside `python/egglog`. - When changing public high-level APIs, update the public docs, stubs, and pretty/freeze round-trip expectations together. +- Keep `builtins.py` limited to operations implemented by Egglog primitives. Compose derived expressions at their domain call sites, without private runtime plumbing in public APIs. +- Pretty-print declaration structure directly; do not evaluate runtime expressions merely to canonicalize their output. - Higher-order callable type probing should stay isolated from the live ruleset: copy declarations and run with no current ruleset so inference does not register temporary unnamed functions or rewrites. ## Array API @@ -41,6 +43,7 @@ ## Verification - Prefer the minimal code change and the minimal diff that solves the task; only broaden the change if the smaller fix is not sufficient. +- High-level tests should assert observable behavior through public APIs. Avoid private `_` APIs and exact generated names unless serialized output is itself the public contract. - Run `make mypy` for typing changes. - Run targeted pytest for touched modules. - Run `make docs` for docs or public API changes. diff --git a/Cargo.toml b/Cargo.toml index 82ed1f8b..2424a7e8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,6 +18,8 @@ opentelemetry = "0.28" opentelemetry-otlp = { version = "0.28", features = ["http-proto", "reqwest-blocking-client", "trace"] } opentelemetry-stdout = { version = "0.28", features = ["trace"] } opentelemetry_sdk = "0.28" +# egglog-experimental 78cdfa5 targets Egglog 8879605. This pin is PR #1008's rule-name fix +# backported onto that compatible base, rather than the same patch on current Egglog main. egglog = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "0476e56345f52e6fcba4b0c93f5d6d60a2b9e318", default-features = false } egglog-ast = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" } egglog-core-relations = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" } diff --git a/docs/changelog.md b/docs/changelog.md index 21edb9aa..16b5d0ae 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -11,11 +11,12 @@ _This project uses semantic versioning_ - Preserve the Egglog 3 sort, constructor, function, and proof metadata that its source syntax can round-trip in the low-level AST bindings; support arbitrary-size `BigInt` values, generic value extraction, - constructor/relation lookup, microsecond-resolution - run-report durations, program filenames, and all-or-nothing command - recording for parsed batches. + constructor/relation lookup, correct run-report duration units with + microsecond resolution, direct parse-and-run execution with source + filenames, and batch-level command recording that omits failed parsed + batches. - Add generic `Pair` and `Maybe` values, undefined-result `catch`, map - folding and fold-derived map operations, map/set lengths, `f64` math + folding, map/set lengths, `f64` math primitives and integer coercion, `i64`-to-`BigRat` coercion, and exact `BigRat.to_i64()` conversion. - Let Python function, method, and constant bodies lower as eager diff --git a/docs/reference/egglog-translation.md b/docs/reference/egglog-translation.md index b0c52853..6603429c 100644 --- a/docs/reference/egglog-translation.md +++ b/docs/reference/egglog-translation.md @@ -108,21 +108,23 @@ missing = catch(lambda: Map[i64, String].empty()[1]) present, missing ``` -Maps have a general `map_fold_kv` primitive. The higher-level -`map_filter_kv`, `map_map_values`, `map_merge_with`, `Map.keys()`, and -`Map.pick_key()` operations are composed from that fold and ordinary -container operations, so callbacks may be regular Egglog lambdas: +Maps have a general `map_fold_kv` primitive. Derived operations can be written +as ordinary expressions at their use sites. Give the fold an explicitly typed +initial value when its result is another container: ```{code-cell} python numbers = Map[i64, i64].empty().insert(1, 10).insert(2, 20) -map_map_values(lambda _key, value: value + 1, numbers) +map_fold_kv( + lambda result, key, value: result.insert(key, value + 1), + Map[i64, i64].empty(), + numbers, +) ``` Map folding uses opaque, e-graph-local `Value` order, not a semantic ordering promised for arbitrary e-class keys. Prefer order-independent callbacks. An -undefined filter predicate skips that entry; -an undefined transform, merge callback, or fold callback makes the whole -operation undefined. +undefined callback makes the fold undefined; use `catch` in the callback when +undefined values should instead select a fallback expression. ## Declaring Functions diff --git a/experiments/param_eq/README.md b/experiments/param_eq/README.md index 8a705f0f..664a76ed 100644 --- a/experiments/param_eq/README.md +++ b/experiments/param_eq/README.md @@ -7,6 +7,8 @@ isolated row execution, resource limits, aggregation, and the research handoff. The work is paused. The bounded public demonstrations remain maintained in CI, while the private 714-row corpus is not run automatically. +The checked-in result CSVs contain headers only as schema examples; a manifest +is created only by a final dependency-compatible aggregate run. ## Provenance and redistribution boundary diff --git a/experiments/param_eq/results/manifest.toml b/experiments/param_eq/results/manifest.toml deleted file mode 100644 index 2cc06be6..00000000 --- a/experiments/param_eq/results/manifest.toml +++ /dev/null @@ -1,6 +0,0 @@ -schema_version = 5 -status = "paused-before-final-dependency-compatible-rerun" -paper_replication_rows = 0 -representation_comparison_rows = 0 -note = "Run make -C experiments/param_eq aggregate after rebuilding the final dependency stack." -archive_identity = "Canonical archive hash is kept in private research records and required by the runner." diff --git a/python/egglog/bindings.pyi b/python/egglog/bindings.pyi index a69ec9b9..fa51973b 100644 --- a/python/egglog/bindings.pyi +++ b/python/egglog/bindings.pyi @@ -546,23 +546,23 @@ class Function: class RunReport: iterations: list[IterationReport] updated: bool - can_stop: bool search_and_apply_time_per_rule: dict[str, timedelta] num_matches_per_rule: dict[str, int] search_and_apply_time_per_ruleset: dict[str, timedelta] merge_time_per_ruleset: dict[str, timedelta] rebuild_time_per_ruleset: dict[str, timedelta] + can_stop: bool def __new__( cls, iterations: list[IterationReport], updated: bool, - can_stop: bool, search_and_apply_time_per_rule: dict[str, timedelta], num_matches_per_rule: dict[str, int], search_and_apply_time_per_ruleset: dict[str, timedelta], merge_time_per_ruleset: dict[str, timedelta], rebuild_time_per_ruleset: dict[str, timedelta], + can_stop: bool = ..., ) -> RunReport: ... ## diff --git a/python/egglog/builtins.py b/python/egglog/builtins.py index 126be588..06e4d29e 100644 --- a/python/egglog/builtins.py +++ b/python/egglog/builtins.py @@ -26,9 +26,8 @@ function, method, set_current_ruleset, - to_runtime_expr, ) -from .runtime import RuntimeClass, RuntimeExpr, RuntimeFunction, resolve_type_annotation_mutate +from .runtime import RuntimeExpr, RuntimeFunction, resolve_type_annotation_mutate from .thunk import Thunk if TYPE_CHECKING: @@ -67,10 +66,7 @@ "i64", "i64Like", "join", - "map_filter_kv", "map_fold_kv", - "map_map_values", - "map_merge_with", "multiset_contains_swapped", "multiset_flat_map", "multiset_fold", @@ -472,8 +468,6 @@ def match(self, f: Callable[[T], V], n: V) -> V: ... L = TypeVar("L", bound=BaseExpr) R = TypeVar("R", bound=BaseExpr) -L2 = TypeVar("L2", bound=BaseExpr) -R2 = TypeVar("R2", bound=BaseExpr) class Pair(BuiltinExpr, Generic[L, R], egg_sort="Pair"): @@ -498,18 +492,6 @@ def left(self) -> L: ... @property def right(self) -> R: ... - @method(preserve=True) - def match(self, f: Callable[[L, R], V]) -> V: - return f(self.left, self.right) - - @method(preserve=True) - def map_left(self, f: Callable[[L], L2]) -> Pair[L2, R]: - return Pair(f(self.left), self.right) - - @method(preserve=True) - def map_right(self, f: Callable[[R], R2]) -> Pair[L, R2]: - return Pair(self.left, f(self.right)) - def _tuple_to_pair(value: tuple[object, ...]) -> Pair: if len(value) != 2: @@ -579,38 +561,10 @@ def remove(self, key: T) -> Map[T, V]: ... @method(egg_fn="map-length") def length(self) -> i64: ... - @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) - TO = TypeVar("TO") VO = TypeVar("VO") A = TypeVar("A", bound=BaseExpr) -V2 = TypeVar("V2", bound=BaseExpr) converter( dict, @@ -629,53 +583,6 @@ def keys(self) -> Set[T]: def map_fold_kv(f: Callable[[A, T, V], A], initial: A, xs: Map[T, V]) -> A: ... -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, - ) - - class Set(BuiltinExpr, Generic[T], egg_sort="Set"): @method(preserve=True) @deprecated("use .value") diff --git a/python/egglog/declarations.py b/python/egglog/declarations.py index db79acd9..f4b5c59a 100644 --- a/python/egglog/declarations.py +++ b/python/egglog/declarations.py @@ -730,10 +730,10 @@ def signature(self) -> FunctionSignature: @dataclass(frozen=True) class ConstantDecl: """ - Same as `(declare)` in egglog. + Python's zero-argument callable declaration. - `body is not None` means the constant lowers eagerly as a zero-arg primitive. - `merge is not None` means the constant lowers as a zero-arg function. + Depending on its return type and whether `body` or `merge` is present, it + lowers to a zero-argument constructor, function, or eager primitive. """ type_ref: JustTypeRef diff --git a/python/egglog/egraph.py b/python/egglog/egraph.py index 4f679e19..38ff4f86 100644 --- a/python/egglog/egraph.py +++ b/python/egglog/egraph.py @@ -1338,12 +1338,23 @@ def extract( cost = cast("COST", extract_report.cost) else: if isinstance(runtime_expr.__egg_typed_expr__.expr, CallDecl): - self._register_extract_root(runtime_expr) + # Register the root through the normal command path before computing costs, so shared subexpressions + # use synthetic lets and the extractor sees the already-materialized root value. + self.register(expr) egg_cost_model = _CostModel(cost_model, self).to_bindings_cost_model() egg_sort = self._state.type_ref_to_egg(tp) extractor = call_with_current_trace(bindings.Extractor, [egg_sort], self._state.egraph, egg_cost_model) termdag = bindings.TermDag() - value = self._state.typed_expr_to_value(runtime_expr.__egg_typed_expr__) + typed_expr = runtime_expr.__egg_typed_expr__ + if isinstance(typed_expr.expr, ValueDecl): + # Values returned by lookup_function_value already identify an e-graph value and cannot be lowered + # back into Egglog syntax. + value = self._state.typed_expr_to_value(typed_expr) + else: + # For call roots, evaluate the same let-factored presentation registered above. Keep + # typed_expr_to_value's direct lowering for non-registering callers such as lookup_function_value. + egg_expr = self._state.typed_expr_to_egg(typed_expr, expr_to_let=True) + value = call_with_current_trace(self._state.egraph.eval_expr, egg_expr)[1] cost, term = call_with_current_trace(extractor.extract_best, self._state.egraph, termdag, value, egg_sort) res = self._from_termdag(termdag, term, tp) return (res, cost) if include_cost else res @@ -1577,20 +1588,6 @@ def _register_commands(self, cmds: list[Command]) -> None: egg_cmds = [egg_cmd for cmd in cmds if (egg_cmd := self._command_to_egg(cmd)) is not None] self._state.run_program(*egg_cmds) - def _register_extract_root(self, runtime_expr: RuntimeExpr) -> None: - """ - Register the exact extraction root without synthetic let factoring. - - Synthetic lets are a command-size optimization for public registration, - but custom-cost extraction immediately evaluates the original root value. - Registering a let-factored presentation can leave the custom extractor - without a costed parent for that exact value in the direct command API. - """ - self._add_decls(runtime_expr) - action_egg = self._state.action_to_egg(ExprActionDecl(runtime_expr.__egg_typed_expr__), expr_to_let=False) - if action_egg is not None: - self._state.run_program(bindings.ActionCommand(action_egg)) - def _command_to_egg(self, cmd: Command) -> bindings._Command | None: ruleset_ident = Ident("") cmd_decl: CommandDecl diff --git a/python/egglog/egraph_state.py b/python/egglog/egraph_state.py index 92683578..985ae462 100644 --- a/python/egglog/egraph_state.py +++ b/python/egglog/egraph_state.py @@ -663,8 +663,8 @@ def callable_ref_to_egg(self, ref: CallableRef) -> tuple[str, bool]: # noqa: C9 if body is not None: self.run_program(self._primitive_command_to_egg(egg_name, decl.signature, body)) else: - # Use constructor declaration instead of constant b/c constants cannot be extracted - # https://github.com/egraphs-good/egglog/issues/334 + # Egglog v3 has no constant command, so lower Python constants + # and class variables as zero-argument functions or constructors. is_function = self.__egg_decls__._classes[tp.ident].builtin or merge is not None schema = bindings.Schema([], self.type_ref_to_egg(tp)) if is_function: diff --git a/python/egglog/exp/param_eq/domain.py b/python/egglog/exp/param_eq/domain.py index 6c8283ed..91e300d9 100644 --- a/python/egglog/exp/param_eq/domain.py +++ b/python/egglog/exp/param_eq/domain.py @@ -271,17 +271,25 @@ def _binary_to_containers( # noqa: C901, PLR0911, PLR0912 return expr match get_callable_args(expr, Num.__add__): case (lhs, rhs): - return map_merge_with( - lambda a, b: a + b, + return map_fold_kv( + lambda result, mono, coef: catch(lambda: result[mono]).match( + lambda old_coef: result.insert(mono, old_coef + coef), result.insert(mono, coef) + ), _to_container_poly(_binary_to_containers(lhs)), _to_container_poly(_binary_to_containers(cast("Num", rhs))), ) match get_callable_args(expr, Num.__sub__): case (lhs, rhs): - return map_merge_with( - lambda a, b: a + b, + return map_fold_kv( + lambda result, mono, coef: catch(lambda: result[mono]).match( + lambda old_coef: result.insert(mono, old_coef + coef), result.insert(mono, coef) + ), _to_container_poly(_binary_to_containers(lhs)), - map_map_values(lambda _, v: -v, _to_container_poly(_binary_to_containers(cast("Num", rhs)))), + map_fold_kv( + lambda result, mono, coef: result.insert(mono, -coef), + ContainerPolynomial.empty(), + _to_container_poly(_binary_to_containers(cast("Num", rhs))), + ), ) match get_callable_args(expr, Num.__mul__): case (lhs, rhs): @@ -293,17 +301,32 @@ def _binary_to_containers( # noqa: C901, PLR0911, PLR0912 lhs_poly = cast("ContainerPolynomial", lhs_mapped) match get_callable_args(rhs_mapped, Num): case (f64(scalar),): - return map_map_values(lambda _mono, coef: coef * scalar, lhs_poly) + return map_fold_kv( + lambda result, mono, coef: result.insert(mono, coef * scalar), + ContainerPolynomial.empty(), + lhs_poly, + ) if not rhs_is_polynomial: return _multiply_container_polynomial_by_monomial(lhs_poly, _to_container_mono(rhs_mapped)) if rhs_is_polynomial: rhs_poly = cast("ContainerPolynomial", rhs_mapped) match get_callable_args(lhs_mapped, Num): case (f64(scalar),): - return map_map_values(lambda _mono, coef: coef * scalar, rhs_poly) + return map_fold_kv( + lambda result, mono, coef: result.insert(mono, coef * scalar), + ContainerPolynomial.empty(), + rhs_poly, + ) if not lhs_is_polynomial: return _multiply_container_polynomial_by_monomial(rhs_poly, _to_container_mono(lhs_mapped)) - return map_merge_with(lambda a, b: a + b, _to_container_mono(lhs_mapped), _to_container_mono(rhs_mapped)) + return map_fold_kv( + lambda result, term, exponent: catch(lambda: result[term]).match( + lambda old_exponent: result.insert(term, old_exponent + exponent), + result.insert(term, exponent), + ), + _to_container_mono(lhs_mapped), + _to_container_mono(rhs_mapped), + ) match get_callable_args(expr, Num.__truediv__): case (lhs, rhs): lhs_mapped = _binary_to_containers(lhs) @@ -314,15 +337,30 @@ def _binary_to_containers( # noqa: C901, PLR0911, PLR0912 lhs_poly = cast("ContainerPolynomial", lhs_mapped) match get_callable_args(rhs_mapped, Num): case (f64(scalar),): - return map_map_values(lambda _mono, coef: coef / scalar, lhs_poly) - denom = map_map_values(lambda _term, exponent: -exponent, _to_container_mono(rhs_mapped)) + return map_fold_kv( + lambda result, mono, coef: result.insert(mono, coef / scalar), + ContainerPolynomial.empty(), + lhs_poly, + ) + denom = map_fold_kv( + lambda result, term, exponent: result.insert(term, -exponent), + ContainerMonomial.empty(), + _to_container_mono(rhs_mapped), + ) if lhs_is_polynomial and not rhs_is_polynomial: return _multiply_container_polynomial_by_monomial(cast("ContainerPolynomial", lhs_mapped), denom) # If the numerator is just one, then dont add this as a term to the polynomial if _is_expr_instance(lhs_mapped, Num) and lhs_mapped == Num(1.0): return denom num = _to_container_mono(lhs_mapped) - return map_merge_with(lambda a, b: a + b, num, denom) + return map_fold_kv( + lambda result, term, exponent: catch(lambda: result[term]).match( + lambda old_exponent: result.insert(term, old_exponent + exponent), + result.insert(term, exponent), + ), + num, + denom, + ) match get_callable_args(expr, Num.__pow__): case (n, Num(f64(f))): n_mapped = _to_num(_binary_to_containers(n)) @@ -349,11 +387,21 @@ def _multiply_container_polynomial_by_monomial( ) -> ContainerPolynomial: """Distribute one monomial into a polynomial and combine coefficient collisions.""" return map_fold_kv( - lambda result, mono, coef: map_merge_with( - lambda old_coef, new_coef: old_coef + new_coef, + lambda result, mono, coef: map_fold_kv( + lambda merged_poly, merged_mono, new_coef: catch(lambda: merged_poly[merged_mono]).match( + lambda old_coef: merged_poly.insert(merged_mono, old_coef + new_coef), + merged_poly.insert(merged_mono, new_coef), + ), result, ContainerPolynomial.empty().insert( - map_merge_with(lambda left_exp, right_exp: left_exp + right_exp, mono, factor), + map_fold_kv( + lambda merged_mono, term, exponent: catch(lambda: merged_mono[term]).match( + lambda old_exponent: merged_mono.insert(term, old_exponent + exponent), + merged_mono.insert(term, exponent), + ), + mono, + factor, + ), coef, ), ), diff --git a/python/egglog/exp/param_eq/pipeline.py b/python/egglog/exp/param_eq/pipeline.py index a24b5b38..89f4fce5 100644 --- a/python/egglog/exp/param_eq/pipeline.py +++ b/python/egglog/exp/param_eq/pipeline.py @@ -7,8 +7,7 @@ import time from collections.abc import Callable, Iterable from dataclasses import dataclass -from functools import partial -from typing import Literal, TypeVar +from typing import Literal from egglog import * @@ -19,14 +18,13 @@ BACKOFF_MATCH_LIMIT = 1000 BACKOFF_BAN_LENGTH = 30 -T = TypeVar("T", bound=BaseExpr) -V = TypeVar("V", bound=BaseExpr) - @function(builtin=True, egg_fn="f64-is-finite") def _f64_is_finite(value: f64) -> Unit: ... +# Keep derived map operations as explicitly typed folds in this research +# module; only map_fold_kv is a backend primitive and public builtin. # Store discovered constants in a global map so semi-naive analysis can join # them with polynomial terms. If two singleton updates collide, keep the first # representative; tolerant float canonicalization is not part of this paused @@ -34,7 +32,13 @@ def _f64_is_finite(value: f64) -> Unit: ... CONSTS = constant( "CONSTS", Map[Num, f64], - merge=partial(map_merge_with, lambda old_value, _new_value: old_value), + merge=lambda left, right: map_fold_kv( + lambda result, key, value: catch(lambda: result[key]).match( + lambda old_value: result.insert(key, old_value), result.insert(key, value) + ), + left, + right, + ), ) # Map a monomial of the form `{polynomial(P): 1}` to one representative `P`. @@ -50,18 +54,16 @@ def _f64_is_finite(value: f64) -> Unit: ... POLYNOMIAL_MONOMIALS = constant( "POLYNOMIAL_MONOMIALS", Map[ContainerMonomial, ContainerPolynomial], - merge=partial(map_merge_with, lambda old_value, new_value: old_value), + merge=lambda left, right: map_fold_kv( + lambda result, key, value: catch(lambda: result[key]).match( + lambda old_value: result.insert(key, old_value), result.insert(key, value) + ), + left, + right, + ), ) -def if_defined(cond: Unit, then: T, otherwise: T) -> T: - return catch(lambda: cond).match(lambda _: then, otherwise) - - -def try_match(expr: T, on_some: Callable[[T], V], default: V) -> V: - return catch(lambda: expr).match(on_some, default) - - @ruleset def shared_analysis_rules(a: f64) -> Iterable[RewriteOrRule]: yield rewrite(exp(Num(a)), subsume=True).to(Num(a.exp()), _f64_is_finite(a.exp())) @@ -120,35 +122,44 @@ def container_analysis_rules( consts == CONSTS, poly1 == map_fold_kv( - lambda res_poly, mono, coef: ( + lambda res_poly, mono, coef: res_poly.insert( # split monomial into non constants and constants (which are combined into the coefficient): - map_fold_kv( - lambda res_mono_and_coef, term, exp: if_defined( - exp != BigRat(0, 1), - # if the exponent is not zero, process it - try_match( - consts[term], - # if it is a constant, multiply it into the coefficient and drop it from the monomial: - lambda v: if_defined( - exp != BigRat(-1, 1), - res_mono_and_coef.map_right(lambda prev_coef: prev_coef * (v ** exp.to_f64())), - if_defined( - v != f64(0.0), - res_mono_and_coef.map_right(lambda prev_coef: prev_coef / v), - res_mono_and_coef.map_left(lambda mono: mono.insert(term, exp)), + ( + mono_and_coef := map_fold_kv( + lambda res_mono_and_coef, term, exp: catch(lambda: exp != BigRat(0, 1)).match( + # if the exponent is not zero, process it + lambda _: catch(lambda: consts[term]).match( + # if it is a constant, multiply it into the coefficient and drop it from the monomial: + lambda v: catch(lambda: exp != BigRat(-1, 1)).match( + lambda _: Pair( + res_mono_and_coef.left, + res_mono_and_coef.right * (v ** exp.to_f64()), + ), + catch(lambda: v != f64(0.0)).match( + lambda _: Pair( + res_mono_and_coef.left, + res_mono_and_coef.right / v, + ), + Pair( + res_mono_and_coef.left.insert(term, exp), + res_mono_and_coef.right, + ), + ), + ), + # if it is not a constant, keep it in the monomial + Pair( + res_mono_and_coef.left.insert(term, exp), + res_mono_and_coef.right, ), ), - # if it is not a constant, keep it in the monomial - res_mono_and_coef.map_left(lambda mono: mono.insert(term, exp)), + # if the exponent is zero, the term is just 1 and can be dropped from the monomial, so keep the monomial as is + res_mono_and_coef, ), - # if the exponent is zero, the term is just 1 and can be dropped from the monomial, so keep the monomial as is - res_mono_and_coef, - ), - Pair(ContainerMonomial.empty(), coef), - mono, - ).match( - lambda mono, coef: res_poly.insert(mono, coef + catch(lambda: res_poly[mono]).unwrap_or(f64(0.0))) - ) + Pair(ContainerMonomial.empty(), coef), + mono, + ) + ).left, + mono_and_coef.right + catch(lambda: res_poly[mono_and_coef.left]).unwrap_or(f64(0.0)), ), ContainerPolynomial.empty(), poly, @@ -161,13 +172,25 @@ def container_analysis_rules( Num(poly[ContainerMonomial.empty()]), poly.length() == i64(1), # The only key is an empty monomial, so the polynomial is just a constant term: - ContainerMonomial.empty() == poly.pick_key(), + ContainerMonomial.empty() + == map_fold_kv( + lambda picked, key, _value: picked.match(lambda _: picked, Maybe[ContainerMonomial].some(key)), + Maybe[ContainerMonomial].none(), + poly, + ).unwrap(), ) # remove monomials with zero coefficients yield rewrite(polynomial(poly), subsume=True).to( polynomial(poly1), - poly1 == map_filter_kv(lambda _key, value: value != f64(0.0), poly), + poly1 + == map_fold_kv( + lambda result, key, value: catch(lambda: value != f64(0.0)).match( + lambda _: result.insert(key, value), result + ), + ContainerPolynomial.empty(), + poly, + ), poly != poly1, ) @@ -252,11 +275,37 @@ def container_basic_rules( ), poly.length() > i64(1), poly.length() <= i64(4), - nonconst_poly == map_filter_kv(lambda key, _value: key != ContainerMonomial.empty(), poly), - poly2 == map_filter_kv(lambda _key, value: value != f64(1.0), nonconst_poly), - coef == poly2[poly2.pick_key()], + nonconst_poly + == map_fold_kv( + lambda result, key, value: catch(lambda: key != ContainerMonomial.empty()).match( + lambda _: result.insert(key, value), result + ), + ContainerPolynomial.empty(), + poly, + ), + poly2 + == map_fold_kv( + lambda result, key, value: catch(lambda: value != f64(1.0)).match( + lambda _: result.insert(key, value), result + ), + ContainerPolynomial.empty(), + nonconst_poly, + ), + coef + == poly2[ + map_fold_kv( + lambda picked, key, _value: picked.match(lambda _: picked, Maybe[ContainerMonomial].some(key)), + Maybe[ContainerMonomial].none(), + poly2, + ).unwrap() + ], poly2.length() == nonconst_poly.length(), - poly1 == map_map_values(lambda _key, value: value / coef, poly), + poly1 + == map_fold_kv( + lambda result, key, value: result.insert(key, value / coef), + ContainerPolynomial.empty(), + poly, + ), ) # Greedy multivariate Horner factorization for rational exponents. Choose @@ -285,18 +334,21 @@ def container_basic_rules( counts.count(n) > i64(1), exp == map_fold_kv( - lambda min_exp, mono, _coef: try_match(mono[n] < min_exp, lambda _: mono[n], min_exp), + lambda min_exp, mono, _coef: catch(lambda: mono[n] < min_exp).match(lambda _: mono[n], min_exp), BigRat(2**63 - 1, 1), poly, ), poly_pair == map_fold_kv( - lambda divided_and_remainder, mono, coef: try_match( - mono[n], - lambda current_exp: divided_and_remainder.map_left( - lambda divided: divided.insert(mono.insert(n, current_exp - exp), coef) + lambda divided_and_remainder, mono, coef: catch(lambda: mono[n]).match( + lambda current_exp: Pair( + divided_and_remainder.left.insert(mono.insert(n, current_exp - exp), coef), + divided_and_remainder.right, + ), + Pair( + divided_and_remainder.left, + divided_and_remainder.right.insert(mono, coef), ), - divided_and_remainder.map_right(lambda remainder: remainder.insert(mono, coef)), ), Pair(ContainerPolynomial.empty(), ContainerPolynomial.empty()), poly, @@ -311,18 +363,23 @@ def container_basic_rules( # avoids distributing arbitrary products. yield rewrite(polynomial(poly)).to( polynomial( - map_merge_with( - lambda left, right: left + right, + map_fold_kv( + lambda result, key, value: catch(lambda: result[key]).match( + lambda old_value: result.insert(key, old_value + value), result.insert(key, value) + ), poly.remove(mono), - map_map_values(lambda _nested_mono, nested_coef: nested_coef * poly[mono], poly1), + map_fold_kv( + lambda result, nested_mono, nested_coef: result.insert(nested_mono, nested_coef * poly[mono]), + ContainerPolynomial.empty(), + poly1, + ), ) ), polynomial_monomials == POLYNOMIAL_MONOMIALS, poly.length() > i64(1), mono == map_fold_kv( - lambda selected, candidate_mono, _candidate_coef: try_match( - polynomial_monomials[candidate_mono], + lambda selected, candidate_mono, _candidate_coef: catch(lambda: polynomial_monomials[candidate_mono]).match( lambda _nested_poly: candidate_mono, selected, ), @@ -364,8 +421,10 @@ def container_fun_rules(poly: ContainerPolynomial, m: ContainerMonomial, term: N yield rewrite(log(polynomial(poly))).to( polynomial( map_fold_kv( - lambda res_poly, term, exp: map_merge_with( - lambda old_coef, new_coef: old_coef + new_coef, + lambda res_poly, term, exp: map_fold_kv( + lambda result, mono, coef: catch(lambda: result[mono]).match( + lambda old_coef: result.insert(mono, old_coef + coef), result.insert(mono, coef) + ), res_poly, ContainerPolynomial.empty().insert( ContainerMonomial.empty().insert(log(term), BigRat(1, 1)), @@ -377,10 +436,20 @@ def container_fun_rules(poly: ContainerPolynomial, m: ContainerMonomial, term: N ) ), poly.length() == i64(1), - m == poly.pick_key(), + m + == map_fold_kv( + lambda picked, key, _value: picked.match(lambda _: picked, Maybe[ContainerMonomial].some(key)), + Maybe[ContainerMonomial].none(), + poly, + ).unwrap(), poly[m] > f64(0.0), m.length() == i64(1), - term == m.pick_key(), + term + == map_fold_kv( + lambda picked, key, _value: picked.match(lambda _: picked, Maybe[Num].some(key)), + Maybe[Num].none(), + m, + ).unwrap(), m[term] == BigRat(1, 1), ) diff --git a/python/egglog/pretty.py b/python/egglog/pretty.py index a0755320..c85d04a8 100644 --- a/python/egglog/pretty.py +++ b/python/egglog/pretty.py @@ -7,7 +7,7 @@ import ast from collections import Counter, defaultdict from dataclasses import dataclass, field -from typing import TYPE_CHECKING, TypeAlias, assert_never, cast +from typing import TYPE_CHECKING, TypeAlias, assert_never import black import cloudpickle @@ -17,9 +17,6 @@ if TYPE_CHECKING: from collections.abc import Mapping - from .builtins import BigRat, Map, Maybe, Pair - from .egraph import BaseExpr - __all__ = [ "BINARY_METHODS", @@ -427,52 +424,6 @@ def uncached( # noqa: C901, PLR0911, PLR0912 case EGraphDecl() as eg: return f"EGraph({', '.join(map(self, eg.to_actions))}).freeze()", "egraph" case TypedExprDecl(tp, expr): - 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(): - key_str = self(cast("RuntimeExpr", key).__egg_typed_expr__) - value_str = self(cast("RuntimeExpr", value).__egg_typed_expr__) - map_str += f".insert({key_str}, {value_str})" - return map_str, "Map" - if tp.ident == Ident.builtin("BigRat"): - runtime_expr = RuntimeExpr.__from_values__(self.decls, decl) - try: - as_fraction = cast("BigRat", runtime_expr).value - except ExprValueError: - return self(expr, unwrap_lit=unwrap_lit, ruleset_ident=ruleset_ident, parens=parens), "expr" - return f"BigRat({as_fraction.numerator}, {as_fraction.denominator})", "BigRat" - if tp.ident == Ident.builtin("Pair"): - runtime_expr = RuntimeExpr.__from_values__(self.decls, decl) - try: - left, right = cast("Pair[BaseExpr, BaseExpr]", runtime_expr).value - except ExprValueError: - return self(expr, unwrap_lit=unwrap_lit, ruleset_ident=ruleset_ident, parens=parens), "expr" - left_str = self(cast("RuntimeExpr", left).__egg_typed_expr__) - right_str = self(cast("RuntimeExpr", right).__egg_typed_expr__) - return f"{tp}({left_str}, {right_str})", "Pair" - if tp.ident == Ident.builtin("Maybe"): - runtime_expr = RuntimeExpr.__from_values__(self.decls, decl) - try: - value = cast("Maybe[BaseExpr]", runtime_expr).value - except ExprValueError: - return self(expr, unwrap_lit=unwrap_lit, ruleset_ident=ruleset_ident, parens=parens), "expr" - if value is None: - return f"{tp}.none()", "Maybe" - return f"{tp}.some({self(cast('RuntimeExpr', value).__egg_typed_expr__)})", "Maybe" return ( self(expr, unwrap_lit=unwrap_lit, ruleset_ident=ruleset_ident, parens=parens), tp.ident.name, @@ -553,7 +504,9 @@ def _call_inner( # noqa: C901, PLR0911, PLR0912 return name, args, True case ClassMethodRef(class_name, method_name): tp_ref = JustTypeRef(class_name, bound_tp_params or ()) - return f"{tp_ref}.{method_name}", args, True + # Generic calls retain wrappers so their class parameters remain inferable from the arguments. + unwrap_lit = not self.decls.get_class_decl(class_name).type_vars + return f"{tp_ref}.{method_name}", args, unwrap_lit case MethodRef(class_name, method_name): slf, *args = args non_str_slf = slf @@ -592,7 +545,8 @@ def _call_inner( # noqa: C901, PLR0911, PLR0912 return f"{self(args[0], parens=True)}.{property_name}" case InitRef(class_name): tp_ref = JustTypeRef(class_name, bound_tp_params or ()) - return str(tp_ref), args, True + unwrap_lit = not self.decls.get_class_decl(class_name).type_vars + return str(tp_ref), args, unwrap_lit case UnnamedFunctionRef(): expr = self._pretty_function_body(ref, []) return f"({expr})", args, True diff --git a/python/egglog/run_report.py b/python/egglog/run_report.py index fb9e79ff..26501d68 100644 --- a/python/egglog/run_report.py +++ b/python/egglog/run_report.py @@ -86,12 +86,12 @@ class RunReport: _decls: Declarations = field(repr=False) iterations: list[IterationReport] = field(default_factory=list) updated: bool = False - can_stop: bool = False search_and_apply_time_per_rule: dict[RewriteOrRuleDecl, timedelta] = field(default_factory=dict) num_matches_per_rule: dict[RewriteOrRuleDecl, int] = field(default_factory=dict) search_and_apply_time_per_ruleset: dict[str, timedelta] = field(default_factory=dict) merge_time_per_ruleset: dict[str, timedelta] = field(default_factory=dict) rebuild_time_per_ruleset: dict[str, timedelta] = field(default_factory=dict) + can_stop: bool = False def __repr__(self) -> str: time_per_rule = {pretty_decl(self._decls, k): v for k, v in self.search_and_apply_time_per_rule.items()} diff --git a/python/tests/__snapshots__/test_array_api/test_jit[lda][expr].py b/python/tests/__snapshots__/test_array_api/test_jit[lda][expr].py index b880f65a..cda4a169 100644 --- a/python/tests/__snapshots__/test_array_api/test_jit[lda][expr].py +++ b/python/tests/__snapshots__/test_array_api/test_jit[lda][expr].py @@ -37,7 +37,10 @@ _TupleNDArray_1 = svd_( sqrt( asarray( - NDArray(RecursiveValue(Value.from_float(Float.rational(BigRat(1, 147))))), OptionalDType.some(DType.float64), OptionalBool.none, OptionalDevice.some(_NDArray_1.device) + NDArray(RecursiveValue(Value.from_float(Float.rational(BigRat(BigInt.from_string("1"), BigInt.from_string("147")))))), + OptionalDType.some(DType.float64), + OptionalBool.none, + OptionalDevice.some(_NDArray_1.device), ) ) * (_NDArray_8 / _NDArray_11), @@ -51,7 +54,11 @@ ).T / _TupleNDArray_1[Int(1)][IndexKey.slice(_Slice_1)] _TupleNDArray_2 = svd_( ( - sqrt(NDArray(RecursiveValue(Value.from_int(Int(150)))) * _NDArray_3 * NDArray(RecursiveValue(Value.from_float(Float.rational(BigRat(1, 2)))))) + sqrt( + NDArray(RecursiveValue(Value.from_int(Int(150)))) + * _NDArray_3 + * NDArray(RecursiveValue(Value.from_float(Float.rational(BigRat(BigInt.from_string("1"), BigInt.from_string("2")))))) + ) * (_NDArray_4 - _NDArray_3 @ _NDArray_4).T ).T @ _NDArray_12, diff --git a/python/tests/test_bindings.py b/python/tests/test_bindings.py index b26eeea4..8054d848 100644 --- a/python/tests/test_bindings.py +++ b/python/tests/test_bindings.py @@ -358,12 +358,12 @@ def test_report_duration_round_trip(self, duration: timedelta): run_report = RunReport( [iteration_report], True, - True, {"rule": duration}, {"rule": 7}, {"ruleset": duration}, {"ruleset": duration}, {"ruleset": duration}, + can_stop=True, ) assert rule_report.search_and_apply_time == duration @@ -375,6 +375,11 @@ def test_report_duration_round_trip(self, duration: timedelta): assert run_report.merge_time_per_ruleset["ruleset"] == duration assert run_report.rebuild_time_per_ruleset["ruleset"] == duration + def test_run_report_can_stop_preserves_positional_constructor(self): + report = RunReport([], False, {}, {}, {}, {}, {}) + + assert report.can_stop is False + @pytest.mark.parametrize("duration", [timedelta(microseconds=-1), timedelta.min]) def test_report_rejects_negative_duration(self, duration: timedelta): with pytest.raises(ValueError, match="negative timedeltas"): diff --git a/python/tests/test_high_level.py b/python/tests/test_high_level.py index 7a65b6a8..9ece37dc 100644 --- a/python/tests/test_high_level.py +++ b/python/tests/test_high_level.py @@ -16,7 +16,6 @@ import egglog.builtins as egg_builtins from egglog import * -from egglog import bindings from egglog.declarations import ( BUILTIN_EGG_FN_NAMES, BUILTIN_EGG_SORT_NAMES, @@ -34,7 +33,6 @@ MethodRef, TypedExprDecl, TypeRefWithVars, - ValueDecl, ) from egglog.egraph import get_current_ruleset from egglog.runtime import RuntimeExpr, RuntimeFunction @@ -50,24 +48,17 @@ def test_ne(self): @pytest.mark.parametrize( - ("eval_mode", "binding_type"), - [ - ("seminaive", bindings.Seminaive), - ("naive", bindings.Naive), - ("unsafe-seminaive", bindings.UnsafeSeminaive), - ], + "eval_mode", + ["seminaive", "naive", "unsafe-seminaive"], ) -def test_rule_eval_mode_lowering(eval_mode: RuleEvalMode, binding_type: type) -> None: +def test_rule_eval_mode(eval_mode: RuleEvalMode) -> None: rel = relation(f"eval_mode_{eval_mode}", i64) x = var("x", i64) - high_level_rule = rule(rel(x), eval_mode=eval_mode).then(rel(x + 1)) - - egraph = EGraph() - egraph._add_decls(high_level_rule) - command = egraph._command_to_egg(high_level_rule) + egraph = EGraph(rel(i64(0))) - assert isinstance(command, bindings.RuleCommand) - assert isinstance(command.rule.eval_mode, binding_type) + egraph.register(rule(rel(x), eval_mode=eval_mode).then(rel(x + 1))) + egraph.run(1) + egraph.check(rel(i64(1))) def test_eqsat_basic(): @@ -108,13 +99,16 @@ def test_lookup_function_value_constructor_row() -> None: class A(Expr): def __init__(self, value: i64Like) -> None: ... + def score(self) -> i64: ... + egraph = EGraph(A(1)) + egraph.register(set_(A(1).score()).to(i64(7))) value = egraph.lookup_function_value(A(1)) assert value is not None - a_typed_expr = cast("RuntimeExpr", A(1)).__egg_typed_expr__ - constructor_value = egraph._state.typed_expr_to_value(a_typed_expr) - assert cast("RuntimeExpr", value).__egg_typed_expr__ == TypedExprDecl(a_typed_expr.tp, ValueDecl(constructor_value)) + score = egraph.lookup_function_value(value.score()) + assert score is not None + assert int(score) == 7 assert egraph.lookup_function_value(A(2)) is None @@ -149,37 +143,23 @@ def pair(cls, left: CheckEdge, right: CheckEdge) -> CheckEdge: ... assert egraph.function_size(CheckEdge.pair) == 0 -def test_synthetic_lets_use_reserved_expr_names() -> None: - class LetNum(Expr): - @classmethod - def var(cls, v: StringLike) -> LetNum: ... - - egraph = EGraph(save_egglog_string=True) - expr = LetNum.var("x") - runtime_expr = cast("RuntimeExpr", expr) - egraph._add_decls(runtime_expr) - - egraph._state._transform_let(runtime_expr.__egg_typed_expr__) - - assert '(let $__expr_0 (LetNum_var "x"))' in egraph.as_egglog_string - - def test_synthetic_lets_skip_explicit_let_conflicts() -> None: class LetConflictNum(Expr): @classmethod def var(cls, v: StringLike) -> LetConflictNum: ... - egraph = EGraph(save_egglog_string=True) - egraph.let("__expr_0", LetConflictNum.var("explicit")) - expr = LetConflictNum.var("synthetic") - runtime_expr = cast("RuntimeExpr", expr) - egraph._add_decls(runtime_expr) + @classmethod + def pair(cls, left: LetConflictNum, right: LetConflictNum) -> LetConflictNum: ... - egraph._state._transform_let(runtime_expr.__egg_typed_expr__) + egraph = EGraph() + explicit = egraph.let("__expr_0", LetConflictNum.var("explicit")) + shared = LetConflictNum.var("shared") + pair = LetConflictNum.pair(shared, shared) + egraph.register(pair) - egglog_string = egraph.as_egglog_string - assert '(let $__expr_0 (LetConflictNum_var "explicit"))' in egglog_string - assert '(let $__expr_1 (LetConflictNum_var "synthetic"))' in egglog_string + egraph.check(eq(explicit).to(LetConflictNum.var("explicit"))) + assert egraph.function_size(LetConflictNum.var) == 2 + assert egraph.function_size(LetConflictNum.pair) == 1 def test_synthetic_let_names_do_not_shadow_default_rewrite_variables() -> None: @@ -192,18 +172,15 @@ def __init__(self, value: i64Like) -> None: ... def make(cls, value: i64Like) -> LetShadowDefaultNum: return LetShadowDefaultNum(value) - egraph = EGraph(save_egglog_string=True) - expr = LetShadowDefaultNum(3) - runtime_expr = cast("RuntimeExpr", expr) - egraph._add_decls(runtime_expr) - egraph._state._transform_let(runtime_expr.__egg_typed_expr__) + @classmethod + def pair(cls, left: LetShadowDefaultNum, right: LetShadowDefaultNum) -> LetShadowDefaultNum: ... + egraph = EGraph() + shared = LetShadowDefaultNum(3) + egraph.register(LetShadowDefaultNum.pair(shared, shared)) egraph.register(LetShadowDefaultNum.make(i64(1))) egraph.run(run(default_ruleset)) - - egglog_string = egraph.as_egglog_string - assert "(let $__expr_0 (LetShadowDefaultNum___init__ 3))" in egglog_string - assert "(rewrite (LetShadowDefaultNum_make _0) (LetShadowDefaultNum___init__ _0)" in egglog_string + egraph.check(eq(LetShadowDefaultNum.make(i64(1))).to(LetShadowDefaultNum(i64(1)))) def test_save_egglog_string_defaults_to_disabled() -> None: @@ -211,14 +188,14 @@ def test_save_egglog_string_defaults_to_disabled() -> None: with pytest.raises(ValueError, match="save_egglog_string=True"): _ = egraph.as_egglog_string - assert egraph._state.egglog_file_state is None def test_saved_egglog_transcript_close_is_idempotent() -> None: egraph = EGraph(save_egglog_string=True) assert egraph._state.egglog_file_state is not None path = pathlib.Path(egraph._state.egglog_file_state.path) - assert path.exists() + egraph.let("x", i64(1)) + assert egraph.as_egglog_string egraph.close() egraph.close() @@ -242,16 +219,14 @@ def test_saved_egglog_transcript_is_removed_on_finalization() -> None: def test_saved_egglog_transcript_is_shared_across_push_and_pop() -> None: egraph = EGraph(save_egglog_string=True) - transcript = egraph._state.egglog_file_state - assert transcript is not None + egraph.let("outer", i64(1)) + outer_transcript = egraph.as_egglog_string egraph.push() - assert egraph._state.egglog_file_state is transcript + egraph.let("inner", i64(2)) egraph.pop() - assert egraph._state.egglog_file_state is transcript - - egraph.close() - assert transcript.file.closed + assert egraph.as_egglog_string.startswith(outer_transcript) + assert "(let $inner 2)" in egraph.as_egglog_string def test_saved_egglog_string_uses_short_generated_sort_and_function_names() -> None: @@ -430,78 +405,6 @@ def test_parameterized_sort_names_use_allocated_argument_names() -> None: assert "(sort Map[i64,BigRat] (Map i64 BigRat))" in egraph.as_egglog_string -def test_non_constructor_map_empty_does_not_create_synthetic_let() -> None: - egraph = EGraph(save_egglog_string=True) - expr = Map[Map[String, i64], f64].empty() - runtime_expr = cast("RuntimeExpr", expr) - egraph._add_decls(runtime_expr) - - assert egraph._state._transform_let(runtime_expr.__egg_typed_expr__) is not None - - lines = egraph.as_egglog_string.splitlines() - assert "(let $__expr_0 (map-empty))" not in lines - assert not any(line.startswith("(fail ") for line in lines) - - -def test_non_constructor_maybe_none_does_not_create_synthetic_let() -> None: - egraph = EGraph(save_egglog_string=True) - expr = Maybe[Maybe[i64]].none() - runtime_expr = cast("RuntimeExpr", expr) - egraph._add_decls(runtime_expr) - - assert egraph._state._transform_let(runtime_expr.__egg_typed_expr__) is not None - - lines = egraph.as_egglog_string.splitlines() - assert "(let $__expr_0 (maybe-none))" not in lines - assert not any(line.startswith("(fail ") for line in lines) - - -def test_inferable_non_constructor_map_empty_does_not_create_synthetic_let() -> None: - egraph = EGraph(save_egglog_string=True) - expr = Map[String, i64].empty() - runtime_expr = cast("RuntimeExpr", expr) - egraph._add_decls(runtime_expr) - - assert egraph._state._transform_let(runtime_expr.__egg_typed_expr__) is not None - - lines = egraph.as_egglog_string.splitlines() - assert "(let $__expr_0 (map-empty))" not in lines - assert not any(line.startswith("(fail ") for line in lines) - - -def test_freeze_omits_synthetic_let_bindings() -> None: - class FreezeLetNum(Expr): - @classmethod - def var(cls, v: StringLike) -> FreezeLetNum: ... - - egraph = EGraph(save_egglog_string=True) - expr = FreezeLetNum.var("x") - runtime_expr = cast("RuntimeExpr", expr) - egraph._add_decls(runtime_expr) - egraph._state._transform_let(runtime_expr.__egg_typed_expr__) - - assert "(let $__expr_0 " in egraph.as_egglog_string - assert "$__expr_0" not in str(egraph.freeze()) - - -def test_popped_explicit_lets_do_not_block_synthetic_let_names() -> None: - class ScopedLetNum(Expr): - @classmethod - def var(cls, v: StringLike) -> ScopedLetNum: ... - - egraph = EGraph(save_egglog_string=True) - egraph.push() - egraph.let("__expr_0", ScopedLetNum.var("pushed")) - egraph.pop() - expr = ScopedLetNum.var("synthetic") - runtime_expr = cast("RuntimeExpr", expr) - egraph._add_decls(runtime_expr) - - egraph._state._transform_let(runtime_expr.__egg_typed_expr__) - - assert '(let $__expr_0 (ScopedLetNum_var "synthetic"))' in egraph.as_egglog_string - - def test_registering_bare_variable_expression_raises() -> None: egraph = EGraph() @@ -517,8 +420,7 @@ def test_registering_let_reference_expression_raises() -> None: egraph.register(x) -@pytest.mark.parametrize("save_egglog_string", [True, False]) -def test_nested_rule_lowering_does_not_reuse_top_level_synthetic_lets(save_egglog_string: bool) -> None: +def test_nested_rule_lowering_does_not_reuse_top_level_synthetic_lets() -> None: class NestedRuleEdge(Expr): @classmethod def leaf(cls, value: StringLike) -> NestedRuleEdge: ... @@ -526,8 +428,8 @@ def leaf(cls, value: StringLike) -> NestedRuleEdge: ... @classmethod def pair(cls, left: NestedRuleEdge, right: NestedRuleEdge) -> NestedRuleEdge: ... - egraph = EGraph(save_egglog_string=save_egglog_string) - done_rel = relation(f"done_rel_ctx_{int(save_egglog_string)}") + egraph = EGraph() + done_rel = relation("done_rel_ctx") shared = NestedRuleEdge.leaf("shared") pair = NestedRuleEdge.pair(shared, shared) @@ -538,11 +440,6 @@ def pair(cls, left: NestedRuleEdge, right: NestedRuleEdge) -> NestedRuleEdge: .. egraph.run(1) egraph.check_fail(done_rel()) - if save_egglog_string: - assert '(= _x (NestedRuleEdge_pair (NestedRuleEdge_leaf "shared") (NestedRuleEdge_leaf "shared")))' in ( - egraph.as_egglog_string - ) - def test_top_level_action_factors_duplicate_sibling_edges() -> None: class DuplicateEdge(Expr): @@ -566,21 +463,22 @@ def pair(cls, left: DuplicateEdge, right: DuplicateEdge) -> DuplicateEdge: ... assert "(DuplicateEdge_pair $__expr_2 $__expr_2)" in transcript -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) +def test_freeze_omits_synthetic_let_bindings() -> None: + class FreezeLetNum(Expr): + @classmethod + def var(cls, value: StringLike) -> FreezeLetNum: ... - egraph.run(combined) + @classmethod + def pair(cls, left: FreezeLetNum, right: FreezeLetNum) -> FreezeLetNum: ... - combined_line = next( - line for line in egraph.as_egglog_string.splitlines() if line.startswith("(unstable-combined-ruleset ") - ) - combined_name = combined_line.split()[1] - assert combined_name.startswith("_combined_ruleset_") - assert combined_name.removeprefix("_combined_ruleset_").isdigit() - assert f"(run-schedule (run {combined_name}))" in egraph.as_egglog_string + shared = FreezeLetNum.var("shared") + egraph = EGraph(save_egglog_string=True) + egraph.register(FreezeLetNum.pair(shared, shared)) + + assert "$__expr_" in egraph.as_egglog_string + frozen = str(egraph.freeze()) + assert "FreezeLetNum.pair" in frozen + assert "$__expr_" not in frozen def test_integer_run_accepts_a_combined_ruleset() -> None: @@ -596,43 +494,34 @@ def test_integer_run_accepts_a_combined_ruleset() -> None: egraph.check(copied(i64(1))) -def test_higher_order_builtin_callback_materializes_scalar_builtin_type_args() -> None: - check_eq(map_fold_kv(lambda acc, k, v: acc + v, f64(0.0), Map[i64, f64].empty()), f64(0.0)) - - def test_higher_order_builtin_callback_materializes_parameterized_builtin_dummy_args() -> None: input_map = Map[i64, Maybe[f64]].empty().insert(i64(1), Maybe[f64].some(f64(2.5))) + initial: Map[i64, f64] = Map[i64, f64].empty() expected = Map[i64, f64].empty().insert(i64(1), f64(2.5)) - check_eq(map_map_values(lambda k, v: v.unwrap(), input_map), expected) + check_eq(map_fold_kv(lambda result, key, value: result.insert(key, value.unwrap()), initial, input_map), expected) def test_higher_order_builtin_callback_materializes_rational_builtin_dummy_args() -> None: input_map = Map[i64, Rational].empty().insert(i64(1), Rational(1, 2)) + initial: Map[i64, f64] = Map[i64, f64].empty() expected = Map[i64, f64].empty().insert(i64(1), f64(1.5)) - check_eq(map_map_values(lambda k, v: v.to_f64() + 1.0, input_map), expected) + check_eq( + map_fold_kv(lambda result, key, value: result.insert(key, value.to_f64() + 1.0), initial, input_map), + expected, + ) -def test_map_map_values_generic_negation_callback_is_concretized() -> None: +def test_map_fold_generic_negation_callback_is_concretized() -> None: input_map = Map[i64, f64].empty().insert(i64(1), f64(2.5)).insert(i64(2), f64(-4.0)) + initial: Map[i64, f64] = Map[i64, f64].empty() expected = Map[i64, f64].empty().insert(i64(1), f64(-2.5)).insert(i64(2), f64(4.0)) - check_eq(map_map_values(lambda _key, value: -value, input_map), expected) - + check_eq(map_fold_kv(lambda result, key, value: result.insert(key, -value), initial, input_map), expected) -def test_fold_derived_map_helpers() -> None: - left = Map[i64, f64].empty().insert(i64(1), f64(2.0)).insert(i64(2), f64(3.0)) - right = Map[i64, f64].empty().insert(i64(2), f64(4.0)).insert(i64(3), f64(5.0)) - check_eq(map_filter_kv(lambda key, _value: key > 1, left), Map[i64, f64].empty().insert(i64(2), f64(3.0))) - check_eq( - map_merge_with(lambda old, new: old + new, left, right), - Map[i64, f64].empty().insert(i64(1), f64(2.0)).insert(i64(2), f64(7.0)).insert(i64(3), f64(5.0)), - ) - - check_eq(left.length(), i64(2)) - assert EGraph().extract(left.keys()).value == {i64(1), i64(2)} - check_eq(left.keys().length(), i64(2)) - EGraph().check(left.contains(left.pick_key())) +def test_map_and_set_length_primitives() -> None: + check_eq(Map[i64, f64].empty().insert(i64(1), f64(2.0)).insert(i64(2), f64(3.0)).length(), i64(2)) + check_eq(Set(i64(1), i64(2)).length(), i64(2)) def test_higher_order_callable_inference_does_not_mutate_ambient_ruleset() -> None: @@ -640,7 +529,12 @@ def test_higher_order_callable_inference_does_not_mutate_ambient_ruleset() -> No initial_rules = tuple(ambient.__egg_ruleset__.rules) with set_current_ruleset(ambient): - expr = map_map_values(lambda _key, value: -value, Map[i64, f64].empty().insert(i64(1), f64(2.0))) + initial: Map[i64, f64] = Map[i64, f64].empty() + expr = map_fold_kv( + lambda result, key, value: result.insert(key, -value), + initial, + Map[i64, f64].empty().insert(i64(1), f64(2.0)), + ) _ = cast("RuntimeExpr", expr).__egg_decls__ assert tuple(ambient.__egg_ruleset__.rules) == initial_rules @@ -694,8 +588,13 @@ def test_unnamed_lambda_returning_eqsort_is_eager() -> None: class Box(Expr): def __init__(self, value: i64Like) -> None: ... + initial: Map[i64, Box] = Map[i64, Box].empty() expected = Map[i64, Box].empty().insert(i64(1), Box(i64(2))) - actual = cast("Map[i64, Box]", map_map_values(lambda k, v: Box(v), Map[i64, i64].empty().insert(i64(1), i64(2)))) + actual = map_fold_kv( + lambda result, key, value: result.insert(key, Box(value)), + initial, + Map[i64, i64].empty().insert(i64(1), i64(2)), + ) check_eq(cast("BaseExpr", actual), cast("BaseExpr", expected)) @@ -1521,22 +1420,6 @@ class Defaults(Expr): check_eq(cast("Maybe[i64]", Defaults.missing), Maybe[i64].none()) - def test_bodyless_primitive_constant_is_not_a_synthetic_let(self): - value = constant("bodyless_primitive", i64) - egraph = EGraph(save_egglog_string=True) - egraph._add_decls(cast("RuntimeExpr", value)) - - assert egraph._state._transform_let(cast("RuntimeExpr", value).__egg_typed_expr__) is not None - assert "(let $__expr_0 bodyless_primitive)" not in egraph.as_egglog_string - - def test_bodyless_eqsort_constant_is_a_synthetic_let(self): - value = constant("bodyless_eqsort", A) - egraph = EGraph(save_egglog_string=True) - egraph._add_decls(cast("RuntimeExpr", value)) - - assert egraph._state._transform_let(cast("RuntimeExpr", value).__egg_typed_expr__) is None - assert "(let $__expr_0 (%bodyless_eqsort))" in egraph.as_egglog_string - def test_eqsort_constant_with_merge(self): merged = constant("merged", A, merge=lambda old, _new: old) @@ -1945,18 +1828,10 @@ def test_higher_order_maybe_pair_and_catch_builtins(): assert EGraph().extract(Maybe[i64].some(2).match(lambda x: x + 3, i64(0))).value == 5 # type: ignore[arg-type] assert EGraph().extract(Maybe[i64].none().match(lambda x: x + 3, i64(7))).value == 7 - matched = EGraph().extract(Pair(i64(2), i64(3)).match(lambda left, right: left + right)) - assert matched.value == 5 - - mapped_left = EGraph().extract(Pair(i64(2), i64(3)).map_left(lambda left: left + 10)) - left, right = mapped_left.value - assert left.value == 12 - assert right.value == 3 - - mapped_right = EGraph().extract(Pair(i64(2), i64(3)).map_right(lambda right: right + 10)) - left, right = mapped_right.value - assert left.value == 2 - assert right.value == 13 + pair = Pair(i64(2), i64(3)) + assert EGraph().extract(pair.left).value == 2 + assert EGraph().extract(pair.right).value == 3 + assert EGraph().extract(pair).value == (i64(2), i64(3)) caught = EGraph().extract(catch(lambda: Maybe[i64].some(4).unwrap())) # type: ignore[arg-type] assert caught.value is not None @@ -2532,6 +2407,16 @@ def is_even_cost_model(egraph: EGraph, expr: BaseExpr, children_costs: list[int] assert EGraph().extract(i64(10), include_cost=True, cost_model=is_even_cost_model) == (i64(10), 1) assert EGraph().extract(i64(5), include_cost=True, cost_model=is_even_cost_model) == (i64(5), 0) + def test_lookup_value_root(self) -> None: + class LookupExpr(Expr): + def __init__(self, value: i64Like) -> None: ... + + egraph = EGraph(LookupExpr(1)) + value = egraph.lookup_function_value(LookupExpr(1)) + + assert value is not None + assert egraph.extract(value, include_cost=True, cost_model=default_cost_model) == (LookupExpr(1), 2) + @staticmethod def _capture_container_children_costs( root_expr: BaseExpr, diff --git a/python/tests/test_run_report.py b/python/tests/test_run_report.py index 0ae07fef..32620f68 100644 --- a/python/tests/test_run_report.py +++ b/python/tests/test_run_report.py @@ -6,7 +6,7 @@ import pytest from egglog import * -from egglog.declarations import BiRewriteDecl, RewriteDecl, RuleDecl +from egglog.declarations import BiRewriteDecl, Declarations, RewriteDecl, RuleDecl def _setup_simple_egraph(): @@ -70,6 +70,12 @@ def test_can_stop_field(): assert "can_stop=True" in repr(report) +def test_can_stop_preserves_positional_constructor(): + report = RunReport(Declarations(), [], False, {}, {}, {}, {}, {}) + + assert report.can_stop is False + + def test_num_matches(): egraph = _setup_simple_egraph() report = egraph.run(10) diff --git a/src/conversions.rs b/src/conversions.rs index f1331b54..73fb9d42 100644 --- a/src/conversions.rs +++ b/src/conversions.rs @@ -818,12 +818,12 @@ convert_struct!( egglog_reports::RunReport: "{:?}" => RunReport( iterations: Vec, updated: bool, - can_stop: bool, search_and_apply_time_per_rule: HashMap, num_matches_per_rule: HashMap, search_and_apply_time_per_ruleset: HashMap, merge_time_per_ruleset: HashMap, - rebuild_time_per_ruleset: HashMap + rebuild_time_per_ruleset: HashMap, + can_stop: bool = false ) r -> egglog_reports::RunReport { iterations: r From 6bdaad7d6a2958f8b210db748eacba8d3977c85e Mon Sep 17 00:00:00 2001 From: Saul Shanabrook Date: Tue, 1 Sep 2026 11:21:44 -0700 Subject: [PATCH 3/6] Update Egglog dependencies and extraction APIs --- Cargo.lock | 107 +++- Cargo.toml | 26 +- docs/changelog.md | 31 +- docs/reference/egglog-translation.md | 14 +- docs/reference/python-integration.md | 146 ++++- docs/reference/usage.md | 22 +- python/egglog/bindings.pyi | 49 +- python/egglog/builtins.py | 45 +- python/egglog/declarations.py | 1 + python/egglog/egraph.py | 556 ++++++++++++----- python/egglog/egraph_state.py | 58 +- python/egglog/exp/any_expr_example.ipynb | 2 +- python/egglog/exp/array_api_jit.py | 4 +- python/egglog/exp/param_eq/pipeline.py | 22 +- python/egglog/pretty.py | 6 +- .../test_array_api/test_jit[lda][code].py | 2 +- .../test_array_api/test_jit[lda][expr].py | 4 +- python/tests/param_eq/test_pipeline.py | 8 + python/tests/test_array_api.py | 4 +- python/tests/test_bindings.py | 88 +++ python/tests/test_high_level.py | 537 +++++++++++++++- python/tests/test_pretty.py | 8 + src/egraph.rs | 32 +- src/extract.rs | 584 +++++++++++++++--- src/lib.rs | 15 +- 25 files changed, 2004 insertions(+), 367 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 19b90e3a..1ad18c95 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -73,6 +73,15 @@ version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" +[[package]] +name = "ar_archive_writer" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73cd58deff2140a0a8eae87e417bd01db68a33e148aa93d1e8cd837e55e312b6" +dependencies = [ + "object", +] + [[package]] name = "arc-swap" version = "1.9.1" @@ -147,6 +156,16 @@ version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +[[package]] +name = "cc" +version = "1.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ad534f4357a5264cce5019c989cf66a4f0dc4e0d1b1d15f8aacec0ff7360273" +dependencies = [ + "find-msvc-tools", + "shlex", +] + [[package]] name = "cfg-if" version = "1.0.4" @@ -363,7 +382,7 @@ checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" [[package]] name = "egglog" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=0476e56345f52e6fcba4b0c93f5d6d60a2b9e318#0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=69177f4f37db9bc6c23899375988dc8b8ecedfc3#69177f4f37db9bc6c23899375988dc8b8ecedfc3" dependencies = [ "csv", "dyn-clone", @@ -391,7 +410,7 @@ dependencies = [ [[package]] name = "egglog-add-primitive" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=0476e56345f52e6fcba4b0c93f5d6d60a2b9e318#0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=69177f4f37db9bc6c23899375988dc8b8ecedfc3#69177f4f37db9bc6c23899375988dc8b8ecedfc3" dependencies = [ "quote", "syn 2.0.117", @@ -400,7 +419,7 @@ dependencies = [ [[package]] name = "egglog-ast" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=0476e56345f52e6fcba4b0c93f5d6d60a2b9e318#0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=69177f4f37db9bc6c23899375988dc8b8ecedfc3#69177f4f37db9bc6c23899375988dc8b8ecedfc3" dependencies = [ "ordered-float", ] @@ -408,7 +427,7 @@ dependencies = [ [[package]] name = "egglog-bridge" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=0476e56345f52e6fcba4b0c93f5d6d60a2b9e318#0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=69177f4f37db9bc6c23899375988dc8b8ecedfc3#69177f4f37db9bc6c23899375988dc8b8ecedfc3" dependencies = [ "anyhow", "dyn-clone", @@ -431,7 +450,7 @@ dependencies = [ [[package]] name = "egglog-concurrency" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=0476e56345f52e6fcba4b0c93f5d6d60a2b9e318#0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=69177f4f37db9bc6c23899375988dc8b8ecedfc3#69177f4f37db9bc6c23899375988dc8b8ecedfc3" dependencies = [ "arc-swap", "bumpalo", @@ -443,7 +462,7 @@ dependencies = [ [[package]] name = "egglog-core-relations" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=0476e56345f52e6fcba4b0c93f5d6d60a2b9e318#0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=69177f4f37db9bc6c23899375988dc8b8ecedfc3#69177f4f37db9bc6c23899375988dc8b8ecedfc3" dependencies = [ "anyhow", "bumpalo", @@ -471,25 +490,28 @@ dependencies = [ [[package]] name = "egglog-experimental" version = "3.0.0" -source = "git+https://github.com/egraphs-good/egglog-experimental.git?rev=78cdfa543de9e282c86f446f69943965119c2afe#78cdfa543de9e282c86f446f69943965119c2afe" +source = "git+https://github.com/egraphs-good/egglog-experimental.git?rev=6bfb589d1c4695232fc24e9bdf097c49af451823#6bfb589d1c4695232fc24e9bdf097c49af451823" dependencies = [ "egglog", "egglog-ast", "egglog-reports", + "fixedbitset", + "hashbrown 0.16.1", "lazy_static", "log", "num", + "stacker", ] [[package]] name = "egglog-numeric-id" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=0476e56345f52e6fcba4b0c93f5d6d60a2b9e318#0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=69177f4f37db9bc6c23899375988dc8b8ecedfc3#69177f4f37db9bc6c23899375988dc8b8ecedfc3" [[package]] name = "egglog-reports" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=0476e56345f52e6fcba4b0c93f5d6d60a2b9e318#0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=69177f4f37db9bc6c23899375988dc8b8ecedfc3#69177f4f37db9bc6c23899375988dc8b8ecedfc3" dependencies = [ "clap", "hashbrown 0.16.1", @@ -503,7 +525,7 @@ dependencies = [ [[package]] name = "egglog-union-find" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=0476e56345f52e6fcba4b0c93f5d6d60a2b9e318#0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=69177f4f37db9bc6c23899375988dc8b8ecedfc3#69177f4f37db9bc6c23899375988dc8b8ecedfc3" dependencies = [ "crossbeam", "egglog-concurrency", @@ -534,7 +556,6 @@ dependencies = [ "ordered-float", "pyo3", "pyo3-log", - "rayon", "serde_json", "tracing", "tracing-opentelemetry", @@ -604,6 +625,12 @@ version = "2.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" +[[package]] +name = "find-msvc-tools" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d45db016d36b838f563236e9193d0ee6ce38f3f68b6c94e914b4929c96bbb890" + [[package]] name = "fixedbitset" version = "0.5.7" @@ -1261,6 +1288,15 @@ dependencies = [ "autocfg", ] +[[package]] +name = "object" +version = "0.39.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e5a6c098c7a3b6547378093f5cc30bc54fd361ce711e05293a5cc589562739b" +dependencies = [ + "memchr", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -1531,6 +1567,16 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "psm" +version = "0.1.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4dcd034599e63b970727f70d79e02d62390a4a84f7c6b827c27c46d5ac3fa622" +dependencies = [ + "ar_archive_writer", + "cc", +] + [[package]] name = "pyo3" version = "0.27.2" @@ -1698,26 +1744,6 @@ dependencies = [ "rand_core 0.6.4", ] -[[package]] -name = "rayon" -version = "1.12.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" -dependencies = [ - "either", - "rayon-core", -] - -[[package]] -name = "rayon-core" -version = "1.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" -dependencies = [ - "crossbeam-deque", - "crossbeam-utils", -] - [[package]] name = "redox_syscall" version = "0.5.18" @@ -1897,6 +1923,12 @@ dependencies = [ "lazy_static", ] +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + [[package]] name = "sized-chunks" version = "0.6.5" @@ -1935,6 +1967,19 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" +[[package]] +name = "stacker" +version = "0.1.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "707f49d46706bacf8a2b00d51dace3f9de527c13eec3778f570c411f89e69967" +dependencies = [ + "cc", + "cfg-if", + "libc", + "psm", + "windows-sys", +] + [[package]] name = "strsim" version = "0.11.1" diff --git a/Cargo.toml b/Cargo.toml index 2424a7e8..80cb6386 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,14 +18,13 @@ opentelemetry = "0.28" opentelemetry-otlp = { version = "0.28", features = ["http-proto", "reqwest-blocking-client", "trace"] } opentelemetry-stdout = { version = "0.28", features = ["trace"] } opentelemetry_sdk = "0.28" -# egglog-experimental 78cdfa5 targets Egglog 8879605. This pin is PR #1008's rule-name fix -# backported onto that compatible base, rather than the same patch on current Egglog main. -egglog = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "0476e56345f52e6fcba4b0c93f5d6d60a2b9e318", default-features = false } -egglog-ast = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" } -egglog-core-relations = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" } -egglog-reports = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" } -egglog-bridge = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" } -egglog-experimental = { git = "https://github.com/egraphs-good/egglog-experimental.git", rev = "78cdfa543de9e282c86f446f69943965119c2afe", default-features = false } +# Egglog main at integration plus the rule-name round-trip fix used by this branch. +egglog = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "69177f4f37db9bc6c23899375988dc8b8ecedfc3", default-features = false } +egglog-ast = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "69177f4f37db9bc6c23899375988dc8b8ecedfc3" } +egglog-core-relations = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "69177f4f37db9bc6c23899375988dc8b8ecedfc3" } +egglog-reports = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "69177f4f37db9bc6c23899375988dc8b8ecedfc3" } +egglog-bridge = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "69177f4f37db9bc6c23899375988dc8b8ecedfc3" } +egglog-experimental = { git = "https://github.com/egraphs-good/egglog-experimental.git", rev = "6bfb589d1c4695232fc24e9bdf097c49af451823", default-features = false } egraph-serialize = { version = "0.3", features = ["serde", "graphviz"] } serde_json = "1" pyo3-log = "*" @@ -36,7 +35,6 @@ tracing = "0.1" tracing-opentelemetry = "0.29" tracing-subscriber = "0.3" uuid = { version = "1.18", features = ["v4"] } -rayon = "1.11" base64 = "0.22.1" # enable debug symbols for easier profiling @@ -44,8 +42,8 @@ base64 = "0.22.1" debug = true [patch."https://github.com/egraphs-good/egglog.git"] -egglog = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" } -egglog-ast = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" } -egglog-core-relations = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" } -egglog-reports = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" } -egglog-bridge = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "0476e56345f52e6fcba4b0c93f5d6d60a2b9e318" } +egglog = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "69177f4f37db9bc6c23899375988dc8b8ecedfc3" } +egglog-ast = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "69177f4f37db9bc6c23899375988dc8b8ecedfc3" } +egglog-core-relations = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "69177f4f37db9bc6c23899375988dc8b8ecedfc3" } +egglog-reports = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "69177f4f37db9bc6c23899375988dc8b8ecedfc3" } +egglog-bridge = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "69177f4f37db9bc6c23899375988dc8b8ecedfc3" } diff --git a/docs/changelog.md b/docs/changelog.md index 16b5d0ae..6212857e 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -16,9 +16,11 @@ _This project uses semantic versioning_ filenames, and batch-level command recording that omits failed parsed batches. - Add generic `Pair` and `Maybe` values, undefined-result `catch`, map - folding, map/set lengths, `f64` math - primitives and integer coercion, `i64`-to-`BigRat` coercion, and exact - `BigRat.to_i64()` conversion. + folding, map/set lengths, `f64` math primitives, `f64.is_finite()`, and + integer coercion; expand the experimental `Rational` Python API with + public `RationalLike` conversions, reflected arithmetic, powers, + `min`/`max`, and `Unit` comparisons; add `i64`-to-`BigRat` coercion and + exact `BigRat.to_i64()` conversion. - Let Python function, method, and constant bodies lower as eager primitives, preserving Python argument order for `reverse_args` callables, while bodies attached to an explicit ruleset remain rewrite-backed; allow @@ -28,6 +30,29 @@ _This project uses semantic versioning_ `unsafe-seminaive` rule evaluation when callbacks read mutable tables. - Add persistent ordinary backoff schedules, expose `RunReport.can_stop`, and keep high-level saturation running while a scheduler has deferred work. + - Configure worker threads and rule decomposition per `EGraph`, allow + `rule(..., no_decomp=True)` for individual rules, and expose getters and + setters for both settings. + - BREAKING: stop reading `RAYON_NUM_THREADS`; pass `num_threads` to + `EGraph` or call `set_num_threads(...)` instead. + - Add `tree` and true `greedy-dag` extraction modes to `extract` and + `extract_multiple`; add destructive `keep_best`; allow multi-root variant + extraction while preserving input order. + - Use the experimental dynamic cost model consistently across the default + extraction paths, including canonical `set_cost` tables. A compatible raw + cost table already occupying the canonical name is reused, while an + incompatible callable or overload occupying that name when the cost table + is created now raises instead of allocating an ignored suffix. Frozen + snapshots preserve both views of a reused raw cost table. + - Name the existing custom callback protocol `TreeCostModel`, retaining + `CostModel` as a compatibility alias, and add frozen additive + `DagCostModel` values that work with tree or greedy-DAG extraction. + The low-level `bindings.Extractor` compatibility facade now prepares + costs on each extraction call, so construction no longer invokes cost + callbacks and later e-graph mutations are observed safely. + - BREAKING: remove `GreedyDagCost`, `GreedyDagCostModel`, and + `greedy_dag_cost_model`; use `DagCostModel` with + `extractor="greedy-dag"` instead. - Improve source transcripts, diagnostics, generated-name collision handling, shared-expression factoring without leaking synthetic bindings into rules or checks, large-program AST parsing, custom-cost extraction of diff --git a/docs/reference/egglog-translation.md b/docs/reference/egglog-translation.md index 6603429c..f6ac84fd 100644 --- a/docs/reference/egglog-translation.md +++ b/docs/reference/egglog-translation.md @@ -370,13 +370,17 @@ You can also set the cost of individual values, like the egglog experimental fea egraph.register(set_cost(fib(0), 1)) ``` -This will be taken into account when extracting. Any value that can be converted to an `i64` is supported as a cost, -so dynamic costs can be created in rules. +This will be taken into account when extracting. Any value that can be +converted to an `i64` is supported, so dynamic costs can be created in rules; +the resulting cost must be nonnegative. -It does this by creating a new table for each function you set the cost for that maps the arguments to an i64. +Python creates a canonical table on demand for each backend function symbol +whose cost is set. The table maps the function's arguments to an `i64`. +Compatible aliases of that symbol share the table. Incompatible overloads are +rejected because one backend symbol cannot have multiple cost-table schemas. -_Note: Unlike in egglog, where you have to declare which functions support custom costs, in Python all functions -are automatically registered to create a custom cost table when they are constructed_ +_Note: Unlike in Egglog source, Python does not require a separate declaration +that a callable supports custom costs; calling `set_cost` creates its table._ You can also get the cost of a function with `get_cost`, which will return an `i64` if one has already been set. diff --git a/docs/reference/python-integration.md b/docs/reference/python-integration.md index 9a467ac5..f9eda07a 100644 --- a/docs/reference/python-integration.md +++ b/docs/reference/python-integration.md @@ -90,6 +90,30 @@ match MyExpr("hello"): print(f"Matched MyExpr with value: {value}") ``` +## 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. + +```{code-cell} python +from fractions import Fraction + +numeric_egraph = EGraph() +numeric_egraph.check(f64(1.0).is_finite()) +result = numeric_egraph.extract(Rational(1, 2) + Fraction(1, 3)) +assert result.value == Fraction(5, 6) +numeric_egraph.check(Rational(1, 2) < 1) +``` + +Rational comparisons return `Unit`, not a Python boolean. A false comparison +is therefore undefined, as are operations such as division by zero and powers +that the backend cannot represent. + ## Python Object Sort We define a custom "primitive sort" (i.e. a builtin type) for `PyObject`s. This allows us to store any Python object in the e-graph. @@ -775,25 +799,72 @@ Common pitfalls when authoring rules: - Ensure rules that subtract from lengths only fire when the length is proven positive. -## Custom Cost Models +## Extraction and Cost Models + +{meth}`egglog.egraph.EGraph.extract` accepts `extractor="tree"` (the +default) or `extractor="greedy-dag"`. Tree extraction charges each occurrence +of a subexpression. Greedy-DAG extraction charges shared subexpressions once +within the result; it is a heuristic rather than a globally optimal DAG +extractor. The public `ExtractionMode` alias contains these two values. +With `include_cost=True`, `extract` returns `(expression, cost)`; custom model +costs are returned directly rather than through a wrapper object. + +### Dynamic Costs + +Without a custom `cost_model`, tree and greedy-DAG extraction use the +experimental dynamic cost model. A row cost registered by `set_cost` overrides +that node's marginal cost; otherwise the model falls back to costs declared on +callables and then the backend default. The same model is used by +`extract_multiple` and `keep_best`. Dynamic row costs must be nonnegative. + +Dynamic row costs live in a canonical table named +`cost_table_`. 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. + +### Multiple Roots + +{meth}`egglog.egraph.EGraph.extract_multiple` returns up to `n` variants for +one expression. Passing a non-empty sequence performs one extraction for all +roots and returns one variant list per root in the same order: -By default, when extracting from the e-graph, we use a simple cost model, that looks at the costs assigned to each -function and any custom costs set with `set_cost`, and finds the lowest cost expression looking at the total tree size. +```{code-block} python +variants = egraph.extract_multiple(expr, 3, extractor="tree") +variants_by_root = egraph.extract_multiple([expr1, expr2], 3, extractor="greedy-dag") +``` -Custom cost models are also supported, which can be passed into `extract` as the `cost_model` keyword argument. They -are defined as functions followed the `CostModel` protocol, that take in an e-graph, an expression, and the costs of the children, and return the total cost of that expression. Costs don't have to be integers, they can be any type that supports comparison. +An inner list may be empty when its root has no extractable variant. The +variant count must be positive, and the sequence form rejects an empty input. +The roots share extraction preparation, but every root and variant is costed +independently; sharing between separate roots does not reduce either cost. +This API always uses dynamic costs; custom Python cost models are supported +only by single-root `extract`. -There are a few builtin cost models: +### Custom Tree Cost Models -- `default_cost_model`: The default cost model, which uses integer costs and sums them up. -- `greedy_dag_cost_model(inner_cost_model=default_cost_model)`: A cost model which uses a greedy DAG algorithm to find the lowest cost expression, allowing for shared sub-expressions. It takes in another cost model to use for the base costs of each expression. +A `TreeCostModel` is a callable that receives the e-graph, one expression +node, and the total costs of its immediate children, then returns the total +cost of that expression. Cost values may be any totally ordered type. The old +`CostModel` name remains as a compatibility alias for this protocol. +Models should normally return a cost no smaller than any child; non-monotone +models are responsible for avoiding cycles in the extracted term. -Note that when passed into your cost model, the expression won't be a full tree. Instead, only the top level call be present, and all of it's arguments will be opaque "value" expressions, representing e-classes in the e-graph. You can't do much with them except use them to construct other expression to pass into `egraph.lookup_function_value` to get the resulting value of a call with those arguments. The only exception is all builtin types, like ints, vecs, strings, etc. will be fully evaluated recursively, so they can be matched against. +The expression passed to a custom model contains only its top-level call. +Arguments representing e-classes are opaque value expressions, although +builtin values such as numbers and containers are reconstructed recursively. +Use {meth}`egglog.egraph.EGraph.lookup_function_value` when a model needs to +inspect a table registered before extraction using those callback arguments +directly. A same-e-graph lookup cannot evaluate a newly derived key while +extraction is holding the graph read-only; such a lookup raises `ValueError`. -For example, here is a cost model that has a boolean cost if the value is even or not: +For example, this model uses a boolean cost for whether an `i64` is even: ```{code-cell} python -def is_even_cost_model(egraph: EGraph, expr: Expr, children_costs: list[bool]) -> bool: +def is_even_cost_model(egraph: EGraph, expr: BaseExpr, children_costs: list[bool]) -> bool: from egglog import i64 # noqa: PLC0415 match expr: @@ -804,3 +875,56 @@ assert EGraph().extract(i64(10), include_cost=True, cost_model=is_even_cost_mode assert EGraph().extract(i64(5), include_cost=True, cost_model=is_even_cost_model) == (i64(5), False) ``` + +A `TreeCostModel` can only be used with the tree extractor. Passing one with +`extractor="greedy-dag"` raises `TypeError`, because a callback that returns +total child costs does not expose the marginal costs needed to account for +sharing. + +### Additive DAG Cost Models + +`DagCostModel(marginal_cost, identity)` is a frozen model that can be used with +either extraction mode. Its callback returns the cost of the current node +without its children or container elements. The backend combines those values +with Python `+`. + +Cost values must be effectively immutable and totally ordered. Addition must +be associative, commutative, and monotone, with `identity` as a two-sided +identity. Under tree extraction the values are added once per occurrence; +under greedy-DAG extraction they are added once per selected shared node. + +```{code-block} python +model = DagCostModel( + marginal_cost=lambda egraph, node: default_cost_model(egraph, node, []), + identity=0, +) + +tree_result, tree_cost = egraph.extract(expr, include_cost=True, cost_model=model) +dag_result, dag_cost = egraph.extract( + expr, + include_cost=True, + cost_model=model, + extractor="greedy-dag", +) +``` + +### Keeping Only the Best Representatives + +{meth}`egglog.egraph.EGraph.keep_best` compacts table-backed callables using +dynamic costs: + +```{code-block} python +egraph.keep_best(target, other_target, extractor="greedy-dag") +``` + +Each target must be a constructor, relation, or bodyless function with a +backend table; eager and builtin primitives are rejected. + +This operation is destructive. It clears every table in the e-graph, then +reinserts only the extracted rows of the requested callables. Declarations and +cost-table identities remain available, although their rows are cleared unless +selected. Existing handles returned by {meth}`egglog.egraph.EGraph.let` become +invalid because their rows have been cleared. Internal let caches are +invalidated, so the same `EGraph` can safely continue registering new actions +and running rules. Call it only when dropping all unselected table state is +intended. diff --git a/docs/reference/usage.md b/docs/reference/usage.md index 86aba6f4..5df6b2df 100644 --- a/docs/reference/usage.md +++ b/docs/reference/usage.md @@ -57,12 +57,28 @@ It follows [SPEC 0](https://scientific-python.org/specs/spec-0000/) in terms of ## Parallelism and threads -The underlying Rust library uses Rayon for parallelism. You can control the worker thread count via the environment variable `RAYON_NUM_THREADS`. If this variable is not set or is invalid, the Python bindings default to using a single thread (`1`). +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`. -```shell -export RAYON_NUM_THREADS=4 # use 4 threads +```python +from egglog import EGraph + +egraph = EGraph(num_threads=4) +egraph.set_num_threads(0) +assert egraph.num_threads() >= 1 ``` +## Rule decomposition + +Egglog normally decomposes rules before execution. `EGraph` defaults to +`no_decomp=False`; pass `no_decomp=True` to disable decomposition for +subsequently registered rules, and use `no_decomp()` or +`set_no_decomp(...)` to inspect or change that setting. For a single rule, +pass `no_decomp=True` to `rule(...)` instead. This is an advanced execution +control. + (community)= ## Community diff --git a/python/egglog/bindings.pyi b/python/egglog/bindings.pyi index fa51973b..d63dabd5 100644 --- a/python/egglog/bindings.pyi +++ b/python/egglog/bindings.pyi @@ -2,7 +2,7 @@ from collections.abc import Callable from datetime import timedelta from fractions import Fraction from pathlib import Path -from typing import Any, Generic, Literal, Protocol, TypeAlias, TypeVar, final +from typing import Any, Generic, Literal, Protocol, Self, TypeAlias, TypeVar, final __all__ = [ "ActionCommand", @@ -16,6 +16,7 @@ __all__ = [ "Constructor", "ContainerRebuildSpec", "CostModel", + "DagCostModel", "Datatype", "Datatypes", "DefaultPrintFunctionMode", @@ -110,6 +111,7 @@ __all__ = [ "Var", "Variant", "WithPlan", + "extract_best_with_dag_cost_model", "setup_tracing", "shutdown_tracing", ] @@ -133,7 +135,13 @@ class SerializedEGraph: @final class EGraph: def __new__( - cls, *, fact_directory: str | Path | None = None, seminaive: bool = True, record: bool = False + cls, + *, + fact_directory: str | Path | None = None, + seminaive: bool = True, + record: bool = False, + num_threads: int = 1, + no_decomp: bool = False, ) -> EGraph: ... def parse_program(self, __input: str, /, filename: str | None = None) -> list[_Command]: ... def parse_and_run_program( @@ -148,6 +156,10 @@ class EGraph: self, *commands: _Command, traceparent: str | None = None, tracestate: str | None = None ) -> list[_CommandOutput]: ... def commands(self) -> str | None: ... + def num_threads(self) -> int: ... + def set_num_threads(self, num_threads: int) -> None: ... + def no_decomp(self) -> bool: ... + def set_no_decomp(self, no_decomp: bool) -> None: ... def serialize( self, root_eclasses: list[_Expr], @@ -994,13 +1006,18 @@ class TermDag: # Extraction ## class _Cost(Protocol): - def __lt__(self, other: _Cost) -> bool: ... - def __le__(self, other: _Cost) -> bool: ... - def __gt__(self, other: _Cost) -> bool: ... - def __ge__(self, other: _Cost) -> bool: ... + def __lt__(self, other: Self) -> bool: ... + def __le__(self, other: Self) -> bool: ... + def __gt__(self, other: Self) -> bool: ... + def __ge__(self, other: Self) -> bool: ... _COST = TypeVar("_COST", bound=_Cost) +class _DagCost(_Cost, Protocol): + def __add__(self, other: Self) -> Self: ... + +_DAG_COST = TypeVar("_DAG_COST", bound=_DagCost) + _ENODE_COST = TypeVar("_ENODE_COST") @final @@ -1013,6 +1030,26 @@ class CostModel(Generic[_COST, _ENODE_COST]): base_value_cost: Callable[[str, Value], _COST], ) -> CostModel[_COST, _ENODE_COST]: ... +@final +class DagCostModel(Generic[_DAG_COST]): + def __new__( + cls, + identity: _DAG_COST, + enode_cost: Callable[[str, list[Value]], _DAG_COST], + container_cost: Callable[[str, Value], _DAG_COST], + base_value_cost: Callable[[str, Value], _DAG_COST], + ) -> DagCostModel[_DAG_COST]: ... + +def extract_best_with_dag_cost_model( + egraph: EGraph, + roots: list[tuple[str, Value]], + cost_model: DagCostModel[_DAG_COST], + *, + extractor: Literal["tree", "greedy-dag"] = "tree", + traceparent: str | None = None, + tracestate: str | None = None, +) -> tuple[TermDag, list[tuple[_DAG_COST, _TermId] | None]]: ... + @final class Extractor(Generic[_COST]): def __new__( diff --git a/python/egglog/builtins.py b/python/egglog/builtins.py index 06e4d29e..65582c74 100644 --- a/python/egglog/builtins.py +++ b/python/egglog/builtins.py @@ -52,6 +52,7 @@ "Primitive", "PyObject", "Rational", + "RationalLike", "Set", "SetLike", "String", @@ -386,6 +387,9 @@ def log(self) -> f64: ... @method(egg_fn="sqrt") def sqrt(self) -> f64: ... + @method(egg_fn="f64-is-finite") + def is_finite(self) -> Unit: ... + @method(egg_fn="<") def __lt__(self, other: f64Like) -> Unit: # type: ignore[has-type] ... @@ -808,22 +812,30 @@ def __init__(self, num: i64Like, den: i64Like) -> None: ... def to_f64(self) -> f64: ... @method(egg_fn="+") - def __add__(self, other: Rational) -> Rational: ... + def __add__(self, other: RationalLike) -> Rational: ... @method(egg_fn="-") - def __sub__(self, other: Rational) -> Rational: ... + def __sub__(self, other: RationalLike) -> Rational: ... @method(egg_fn="*") - def __mul__(self, other: Rational) -> Rational: ... + def __mul__(self, other: RationalLike) -> Rational: ... @method(egg_fn="/") - def __truediv__(self, other: Rational) -> Rational: ... + def __truediv__(self, other: RationalLike) -> Rational: ... + + def __radd__(self, other: RationalLike) -> Rational: ... + + def __rsub__(self, other: RationalLike) -> Rational: ... + + def __rmul__(self, other: RationalLike) -> Rational: ... + + def __rtruediv__(self, other: RationalLike) -> Rational: ... @method(egg_fn="min") - def min(self, other: Rational) -> Rational: ... + def min(self, other: RationalLike) -> Rational: ... @method(egg_fn="max") - def max(self, other: Rational) -> Rational: ... + def max(self, other: RationalLike) -> Rational: ... @method(egg_fn="neg") def __neg__(self) -> Rational: ... @@ -841,7 +853,9 @@ def ceil(self) -> Rational: ... def round(self) -> Rational: ... @method(egg_fn="pow") - def __pow__(self, other: Rational) -> Rational: ... + def __pow__(self, other: RationalLike) -> Rational: ... + + def __rpow__(self, other: RationalLike) -> Rational: ... @method(egg_fn="log") def log(self) -> Rational: ... @@ -860,6 +874,23 @@ def numer(self) -> i64: ... @property def denom(self) -> i64: ... + @method(egg_fn="<") + def __lt__(self, other: RationalLike) -> Unit: ... # type: ignore[has-type] + + @method(egg_fn=">") + def __gt__(self, other: RationalLike) -> Unit: ... + + @method(egg_fn="<=") + def __le__(self, other: RationalLike) -> Unit: ... # type: ignore[has-type] + + @method(egg_fn=">=") + def __ge__(self, other: RationalLike) -> Unit: ... + + +converter(i64, Rational, lambda i: Rational(i, 1)) +converter(Fraction, Rational, lambda f: Rational(f.numerator, f.denominator)) +RationalLike: TypeAlias = Rational | Fraction | i64Like + class BigInt(BuiltinExpr, egg_sort="BigInt"): @method(preserve=True) diff --git a/python/egglog/declarations.py b/python/egglog/declarations.py index f4b5c59a..a918197a 100644 --- a/python/egglog/declarations.py +++ b/python/egglog/declarations.py @@ -1203,6 +1203,7 @@ class RuleDecl: body: tuple[FactDecl, ...] name: str | None eval_mode: RuleEvalMode = "seminaive" + no_decomp: bool = False @dataclass(frozen=True) diff --git a/python/egglog/egraph.py b/python/egglog/egraph.py index 38ff4f86..88e4c5cc 100644 --- a/python/egglog/egraph.py +++ b/python/egglog/egraph.py @@ -5,7 +5,7 @@ import pathlib import sys import tempfile -from collections.abc import Callable, Generator, Iterable +from collections.abc import Callable, Generator, Iterable, Sequence from contextvars import ContextVar, Token from dataclasses import InitVar, dataclass, field, replace from functools import partial @@ -64,18 +64,20 @@ "Command", "Command", "CostModel", + "DagCostModel", "EGraph", "Expr", "ExprCallable", + "ExtractionMode", "Fact", "Fact", "GraphvizKwargs", - "GreedyDagCost", "RewriteOrRule", "RuleEvalMode", "Ruleset", "RunReport", "Schedule", + "TreeCostModel", "_BirewriteBuilder", "_EqBuilder", "_NeBuilder", @@ -95,7 +97,6 @@ "expr_parts", "function", "get_cost", - "greedy_dag_cost_model", "let", "method", "ne", @@ -1159,10 +1160,15 @@ def __init__( *actions: ActionLike, seminaive: bool = True, save_egglog_string: bool = False, + num_threads: int = 1, + no_decomp: bool = False, ) -> None: with _TRACER.start_as_current_span("create"): with _TRACER.start_as_current_span("create_bindings"): - self._state = EGraphState(bindings.EGraph(seminaive=seminaive), save_egglog_string=save_egglog_string) + self._state = EGraphState( + bindings.EGraph(seminaive=seminaive, num_threads=num_threads, no_decomp=no_decomp), + save_egglog_string=save_egglog_string, + ) self._state_stack = [] self._token_stack = [] if actions: @@ -1178,6 +1184,27 @@ def set_report_level(self, level: bindings._ReportLevel) -> None: """ self._egraph.set_report_level(level) + def num_threads(self) -> int: + """Return the number of worker threads configured for this e-graph.""" + return self._egraph.num_threads() + + def set_num_threads(self, num_threads: int) -> None: + """ + Set the number of worker threads used by this e-graph. + + Passing ``1`` keeps execution serial. Passing ``0`` uses available + parallelism. + """ + self._egraph.set_num_threads(num_threads) + + def no_decomp(self) -> bool: + """Return whether rule decomposition is disabled for this e-graph.""" + return self._egraph.no_decomp() + + def set_no_decomp(self, no_decomp: bool) -> None: + """Set whether subsequently registered rules skip decomposition.""" + self._egraph.set_no_decomp(no_decomp) + @property def as_egglog_string(self) -> str: """ @@ -1308,43 +1335,79 @@ def _facts_to_check(self, fact_likes: Iterable[FactLike]) -> bindings.Check: @overload def extract( - self, expr: BASE_EXPR, /, include_cost: Literal[False] = False, cost_model: CostModel | None = None + self, + expr: BASE_EXPR, + /, + include_cost: Literal[False] = False, + cost_model: TreeCostModel | DagCostModel | None = None, + *, + extractor: ExtractionMode = "tree", ) -> BASE_EXPR: ... @overload def extract( - self, expr: BASE_EXPR, /, include_cost: Literal[True], cost_model: None = None + self, + expr: BASE_EXPR, + /, + include_cost: Literal[True], + cost_model: None = None, + *, + extractor: ExtractionMode = "tree", ) -> tuple[BASE_EXPR, int]: ... @overload def extract( - self, expr: BASE_EXPR, /, include_cost: Literal[True], cost_model: CostModel[COST] + self, + expr: BASE_EXPR, + /, + include_cost: Literal[True], + cost_model: TreeCostModel[COST], + *, + extractor: ExtractionMode = "tree", ) -> tuple[BASE_EXPR, COST]: ... + @overload + def extract( + self, + expr: BASE_EXPR, + /, + include_cost: Literal[True], + cost_model: DagCostModel[DAG_COST], + *, + extractor: ExtractionMode = "tree", + ) -> tuple[BASE_EXPR, DAG_COST]: ... + @_TRACER.start_as_current_span("extract") def extract( - self, expr: BASE_EXPR, /, include_cost: bool = False, cost_model: CostModel[COST] | None = None - ) -> BASE_EXPR | tuple[BASE_EXPR, COST]: + self, + expr: BASE_EXPR, + /, + include_cost: bool = False, + cost_model: TreeCostModel[Any] | DagCostModel[Any] | None = None, + *, + extractor: ExtractionMode = "tree", + ) -> BASE_EXPR | tuple[BASE_EXPR, Any]: """ Extract the lowest cost expression from the egraph. """ + _extractor_options(extractor) + if extractor == "greedy-dag" and cost_model is not None and not isinstance(cost_model, DagCostModel): + msg = "The greedy-dag extractor requires a DagCostModel; a general TreeCostModel cannot be adapted" + raise TypeError(msg) runtime_expr = to_runtime_expr(expr) self._add_decls(runtime_expr) tp = runtime_expr.__egg_typed_expr__.tp if cost_model is None: - extract_report = self._run_extract(runtime_expr, 0) + extract_report = self._run_extract(runtime_expr, 0, extractor) assert isinstance(extract_report, bindings.ExtractBest) res = self._from_termdag(extract_report.termdag, extract_report.term, tp) - cost = cast("COST", extract_report.cost) + cost = extract_report.cost else: if isinstance(runtime_expr.__egg_typed_expr__.expr, CallDecl): # Register the root through the normal command path before computing costs, so shared subexpressions # use synthetic lets and the extractor sees the already-materialized root value. self.register(expr) - egg_cost_model = _CostModel(cost_model, self).to_bindings_cost_model() egg_sort = self._state.type_ref_to_egg(tp) - extractor = call_with_current_trace(bindings.Extractor, [egg_sort], self._state.egraph, egg_cost_model) - termdag = bindings.TermDag() typed_expr = runtime_expr.__egg_typed_expr__ if isinstance(typed_expr.expr, ValueDecl): # Values returned by lookup_function_value already identify an e-graph value and cannot be lowered @@ -1355,7 +1418,28 @@ def extract( # typed_expr_to_value's direct lowering for non-registering callers such as lookup_function_value. egg_expr = self._state.typed_expr_to_egg(typed_expr, expr_to_let=True) value = call_with_current_trace(self._state.egraph.eval_expr, egg_expr)[1] - cost, term = call_with_current_trace(extractor.extract_best, self._state.egraph, termdag, value, egg_sort) + if isinstance(cost_model, DagCostModel): + termdag, extracted = call_with_current_trace( + bindings.extract_best_with_dag_cost_model, + self._state.egraph, + [(egg_sort, value)], + _DagCostModel(cost_model, self).to_bindings_cost_model(), + extractor=extractor, + ) + extracted_root = extracted[0] + if extracted_root is None: + msg = "Unextractable root" + raise ValueError(msg) + cost, term = extracted_root + else: + egg_cost_model = _CostModel(cost_model, self).to_bindings_cost_model() + tree_extractor = call_with_current_trace( + bindings.Extractor, [egg_sort], self._state.egraph, egg_cost_model + ) + termdag = bindings.TermDag() + cost, term = call_with_current_trace( + tree_extractor.extract_best, self._state.egraph, termdag, value, egg_sort + ) res = self._from_termdag(termdag, term, tp) return (res, cost) if include_cost else res @@ -1363,24 +1447,113 @@ def _from_termdag(self, termdag: bindings.TermDag, term: int, tp: JustTypeRef) - (new_typed_expr,) = self._state.exprs_from_egg(termdag, [term], tp) return RuntimeExpr.__from_values__(self.__egg_decls__, new_typed_expr) - def extract_multiple(self, expr: BASE_EXPR, n: int) -> list[BASE_EXPR]: + @overload + def extract_multiple(self, expr: BASE_EXPR, n: int, *, extractor: ExtractionMode = "tree") -> list[BASE_EXPR]: ... + + @overload + def extract_multiple( + self, expr: Sequence[BASE_EXPR], n: int, *, extractor: ExtractionMode = "tree" + ) -> list[list[BASE_EXPR]]: ... + + def extract_multiple( + self, + expr: BASE_EXPR | Sequence[BASE_EXPR], + n: int, + *, + extractor: ExtractionMode = "tree", + ) -> list[BASE_EXPR] | list[list[BASE_EXPR]]: """ - Extract multiple expressions from the egraph. + Extract up to ``n`` variants of one expression or each expression in a sequence. + + Sequence results preserve the roots' order and may contain an empty + list when a root has no extractable representation. """ + if n <= 0: + msg = "The number of variants must be positive" + raise ValueError(msg) + extractor_options = _extractor_options(extractor) + if isinstance(expr, Sequence): + if not expr: + msg = "extract_multiple requires at least one expression" + raise ValueError(msg) + runtime_exprs = [to_runtime_expr(item) for item in expr] + self._add_decls(*runtime_exprs) + egg_exprs = [self._state.typed_expr_to_egg(item.__egg_typed_expr__) for item in runtime_exprs] + cmd = bindings.UserDefined( + span(2), + "multi-extract", + [bindings.Lit(span(2), bindings.Int(n)), *egg_exprs, *extractor_options], + ) + try: + outputs = self._state.run_program(cmd) + except BaseException as e: + e.add_note("while extracting: " + ", ".join(map(str, runtime_exprs))) + raise + if len(outputs) != len(runtime_exprs) or not all( + isinstance(output, bindings.ExtractVariants) for output in outputs + ): + msg = "multi-extract returned unexpected command outputs" + raise RuntimeError(msg) + results: list[list[BASE_EXPR]] = [] + for runtime_expr, output in zip(runtime_exprs, outputs, strict=True): + assert isinstance(output, bindings.ExtractVariants) + typed_exprs = self._state.exprs_from_egg( + output.termdag, output.terms, runtime_expr.__egg_typed_expr__.tp + ) + results.append([ + cast("BASE_EXPR", RuntimeExpr.__from_values__(self.__egg_decls__, typed_expr)) + for typed_expr in typed_exprs + ]) + return results + runtime_expr = to_runtime_expr(expr) self._add_decls(runtime_expr) - extract_report = self._run_extract(runtime_expr, n) + extract_report = self._run_extract(runtime_expr, n, extractor) assert isinstance(extract_report, bindings.ExtractVariants) new_exprs = self._state.exprs_from_egg( extract_report.termdag, extract_report.terms, runtime_expr.__egg_typed_expr__.tp ) return [cast("BASE_EXPR", RuntimeExpr.__from_values__(self.__egg_decls__, expr)) for expr in new_exprs] - def _run_extract(self, expr: RuntimeExpr, n: int) -> bindings._CommandOutput: + def keep_best( + self, + fn: ExprCallable, + /, + *fns: ExprCallable, + extractor: ExtractionMode = "tree", + ) -> None: + """ + Keep the best rows of selected callables and clear every other table. + + This destructively compacts the e-graph. Declarations and dynamic-cost + table identities remain available for subsequent iteration, but their + rows are cleared unless selected by the command. + """ + extractor_options = _extractor_options(extractor) + resolved = [resolve_callable(callable_) for callable_ in (fn, *fns)] + self._add_decls(*(decls for _, decls in resolved)) + for ref, _ in resolved: + self._require_table_backed(ref) + table_names = [self._state.callable_ref_to_egg(ref)[0] for ref, _ in resolved] + args: list[bindings._Expr] = [bindings.Lit(span(2), bindings.String(table_name)) for table_name in table_names] + self._state.run_program(bindings.UserDefined(span(2), "keep-best", [*args, *extractor_options])) + + # keep-best clears every table, including synthetic and user let rows. + # Do not let later lowering reuse references to those now-empty tables. + self._state.expr_to_letref_cache.clear() + self._state.expr_to_let_egg_cache.clear() + self._state.expr_to_egg_cache = { + expr: egg_expr + for expr, egg_expr in self._state.expr_to_egg_cache.items() + if not isinstance(expr, LetRefDecl) + } + + def _run_extract(self, expr: RuntimeExpr, n: int, extractor: ExtractionMode = "tree") -> bindings._CommandOutput: egg_expr = self._state.typed_expr_to_egg(expr.__egg_typed_expr__) # If we have defined any cost tables use the custom extraction - args = (egg_expr, bindings.Lit(span(2), bindings.Int(n))) - if self._state.cost_table_names: + args: tuple[bindings._Expr, ...] = (egg_expr, bindings.Lit(span(2), bindings.Int(n))) + args += _extractor_options(extractor) + if self._state.cost_table_names or extractor == "greedy-dag": cmd: bindings._Command = bindings.UserDefined(span(2), "extract", list(args)) else: cmd = bindings.Extract(span(2), *args) @@ -1662,15 +1835,48 @@ def lookup_function_value(self, expr: BASE_EXPR) -> BASE_EXPR | None: Cost lookups use their associated cost table. Eager and builtin primitive calls cannot be inspected with this method. + + During a custom cost-model callback, same-e-graph lookups may use + values supplied to that callback but cannot evaluate newly derived + arguments while extraction holds the e-graph read-only. """ runtime_expr = to_runtime_expr(expr) - self._add_decls(runtime_expr) typed_expr = runtime_expr.__egg_typed_expr__ assert isinstance(typed_expr.expr, CallDecl | GetCostDecl) + 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) + else: + self._add_decls(runtime_expr) if isinstance(typed_expr.expr, CallDecl): self._require_table_backed(typed_expr.expr.callable) egg_fn, typed_args = self._state.translate_call(typed_expr.expr) - values_args = [self._state.typed_expr_to_value(a) for a in typed_args] + if in_cost_model_callback: + assert callback_context is not None + callback_values = callback_context[1] + values_args = [] + for arg in typed_args: + if arg in callback_values: + values_args.append(callback_values[arg]) + elif isinstance(arg.expr, ValueDecl): + values_args.append(arg.expr.value) + else: + msg = ( + "Cost-model callbacks can only look up tables using values supplied to the callback; " + "evaluating new expressions would require mutating the borrowed e-graph" + ) + raise ValueError(msg) + else: + values_args = [self._state.typed_expr_to_value(arg) for arg in typed_args] possible_value = self._egraph.lookup_function(egg_fn, values_args) if possible_value is None: return None @@ -1752,12 +1958,15 @@ def append_e_class_row(output: bindings.Value, tp: JustTypeRef, call: CallDecl, output_tp = self._state.egg_sort_to_type_ref[fn.output_sort] let_bindings[name] = TypedExprDecl(output_tp, self._state.value_to_expr(output_tp, row.output)) continue - is_cost = False - if name in self._state.egg_fn_to_callable_refs: - (callable_ref,) = self._state.egg_fn_to_callable_refs[name] + is_cost = name in self._state.cost_table_names.values() + cost_callable_refs = tuple( + ref for ref, cost_name in self._state.cost_table_names.items() if name == cost_name + ) + raw_callable_refs = self._state.egg_fn_to_callable_refs.get(name, set()) + if is_cost and not raw_callable_refs: + callable_ref = cost_callable_refs[0] else: - (callable_ref,) = (ref for ref, cost_name in self._state.cost_table_names.items() if name == cost_name) - is_cost = True + (callable_ref,) = raw_callable_refs callable_decl = self.__egg_decls__.get_callable_decl(callable_ref) signature = callable_decl.signature assert isinstance(signature, FunctionSignature), ( @@ -1772,17 +1981,30 @@ def append_e_class_row(output: bindings.Value, tp: JustTypeRef, call: CallDecl, for arg_type, value in zip(signature.arg_types, row.inputs, strict=True) for tp in (arg_type.to_just(),) ) - call = CallDecl(callable_ref, arg_exprs) if is_cost: cost_tp = self._state.egg_sort_to_type_ref[fn.output_sort] cost_expr = TypedExprDecl(cost_tp, self._state.value_to_expr(cost_tp, row.output)) - match cost_expr.expr: - case LitDecl(int(value)): - costs[call] = (signature.semantic_return_type.to_just(), value) - case _: - raise TypeError(f"Expected integer cost for {callable_ref}, got {cost_expr.expr}") - continue + for cost_callable_ref in cost_callable_refs: + cost_signature = self.__egg_decls__.get_callable_decl(cost_callable_ref).signature + assert isinstance(cost_signature, FunctionSignature) + cost_arg_exprs = tuple( + TypedExprDecl(tp, self._state.value_to_expr(tp, value)) + for arg_type, value in zip(cost_signature.arg_types, row.inputs, strict=True) + for tp in (arg_type.to_just(),) + ) + cost_call = CallDecl(cost_callable_ref, cost_arg_exprs) + match cost_expr.expr: + case LitDecl(int(value)): + costs[cost_call] = (cost_signature.semantic_return_type.to_just(), value) + case _: + raise TypeError(f"Expected integer cost for {cost_callable_ref}, got {cost_expr.expr}") + # A user-declared bodyless function may intentionally own the + # canonical table that also serves as the dynamic-cost table. + # Preserve its public rows as ordinary sets as well as costs. + if not raw_callable_refs: + continue + call = CallDecl(callable_ref, arg_exprs) output_tp = signature.semantic_return_type.to_just() match callable_decl: case ConstructorDecl(): @@ -1815,21 +2037,28 @@ def append_e_class_row(output: bindings.Value, tp: JustTypeRef, call: CallDecl, ), ) - def _values_to_expr(self, args: list[bindings.Value], name: str) -> RuntimeExpr | None: + def _values_to_expr_and_callback_values( + self, args: list[bindings.Value], name: str + ) -> tuple[RuntimeExpr, dict[TypedExprDecl, bindings.Value]] | None: + """Reconstruct a callback call and map its Python-order arguments to raw backend values.""" if name not in self._state.egg_fn_to_callable_refs: return None (callable_ref,) = self._state.egg_fn_to_callable_refs[name] signature = self.__egg_decls__.get_callable_decl(callable_ref).signature assert isinstance(signature, FunctionSignature) + python_args = args[::-1] if signature.reverse_args else args arg_exprs = tuple( TypedExprDecl(tp, self._state.value_to_expr(tp, arg)) - for arg_type, arg in zip(signature.arg_types, args, strict=True) + for arg_type, arg in zip(signature.arg_types, python_args, strict=True) for tp in (arg_type.to_just(),) ) res_type = signature.semantic_return_type.to_just() - return RuntimeExpr.__from_values__( - self.__egg_decls__, - TypedExprDecl(res_type, CallDecl(callable_ref, arg_exprs)), + return ( + RuntimeExpr.__from_values__( + self.__egg_decls__, + TypedExprDecl(res_type, CallDecl(callable_ref, arg_exprs)), + ), + dict(zip(arg_exprs, python_args, strict=True)), ) @@ -2149,11 +2378,15 @@ def set_cost(expr: BaseExpr, cost: i64Like) -> Action: expr_runtime = to_runtime_expr(expr) cost_runtime = to_runtime_expr(convert(cost, i64)) typed_expr_decl = expr_runtime.__egg_typed_expr__ + cost_decl = cost_runtime.__egg_typed_expr__.expr expr_decl = typed_expr_decl.expr assert isinstance(expr_decl, CallDecl), "Can only set cost of calls, not literals or vars" + if isinstance(cost_decl, LitDecl) and isinstance(cost_decl.value, int) and cost_decl.value < 0: + msg = "Dynamic extraction costs must be nonnegative" + raise ValueError(msg) return Action( Declarations.create(expr_runtime, cost_runtime), - SetCostDecl(typed_expr_decl.tp, expr_decl, cost_runtime.__egg_typed_expr__.expr), + SetCostDecl(typed_expr_decl.tp, expr_decl, cost_decl), ) @@ -2206,9 +2439,16 @@ def rule( ruleset: None = None, name: str | None = None, eval_mode: RuleEvalMode = "seminaive", + no_decomp: bool = False, ) -> _RuleBuilder: """Create a rule with the given facts.""" - return _RuleBuilder(facts=_fact_likes(facts), name=name, ruleset=ruleset, eval_mode=eval_mode) + return _RuleBuilder( + facts=_fact_likes(facts), + name=name, + ruleset=ruleset, + eval_mode=eval_mode, + no_decomp=no_decomp, + ) def var(name: str, bound: TypeForm[T], egg_name: str | None = None) -> T: @@ -2381,6 +2621,7 @@ class _RuleBuilder: name: str | None ruleset: Ruleset | None eval_mode: RuleEvalMode + no_decomp: bool def then(self, *actions: ActionLike) -> RewriteOrRule: actions = _action_likes(actions) @@ -2391,6 +2632,7 @@ def then(self, *actions: ActionLike) -> RewriteOrRule: tuple(f.fact for f in self.facts), self.name, self.eval_mode, + self.no_decomp, ), ) if self.ruleset: @@ -2404,6 +2646,8 @@ def __str__(self) -> str: args.append(f"name={self.name!r}") if self.eval_mode != "seminaive": args.append(f"eval_mode={self.eval_mode!r}") + if self.no_decomp: + args.append("no_decomp=True") if self.ruleset is not None: args.append(f"ruleset={self.ruleset}") return f"rule({', '.join(args)})" @@ -2560,6 +2804,9 @@ def _fact_like(fact_like: FactLike) -> Fact: _CURRENT_RULESET = ContextVar[Ruleset | None]("CURRENT_RULESET", default=None) +_COST_MODEL_CALLBACK_VALUES = ContextVar[tuple[EGraph, dict[TypedExprDecl, bindings.Value]] | None]( + "COST_MODEL_CALLBACK_VALUES", default=None +) def get_current_ruleset() -> Ruleset | None: @@ -2575,6 +2822,19 @@ def set_current_ruleset(r: Ruleset | None) -> Generator[None, None, None]: _CURRENT_RULESET.reset(token) +@contextlib.contextmanager +def _cost_model_callback_values( + egraph: EGraph, + values: dict[TypedExprDecl, bindings.Value], +) -> Generator[None, None, None]: + """Make raw callback values available to read-only table lookups without evaluating expressions.""" + token = _COST_MODEL_CALLBACK_VALUES.set((egraph, values)) + try: + yield + finally: + _COST_MODEL_CALLBACK_VALUES.reset(token) + + def get_cost(expr: BaseExpr) -> i64: """ Return a lookup of the cost of an expression. If not set, won't match. @@ -2601,9 +2861,22 @@ def __ge__(self, other: Self) -> bool: ... COST = TypeVar("COST", bound=Comparable) +ExtractionMode: TypeAlias = Literal["tree", "greedy-dag"] + + +def _extractor_options(extractor: str) -> tuple[bindings._Expr, ...]: + """Validate an extraction mode and lower its optional command selector.""" + match extractor: + case "tree": + return () + case "greedy-dag": + return (bindings.Var(span(2), ":extractor"), bindings.Var(span(2), "greedy-dag")) + case _: + msg = f"Unknown extractor {extractor!r}; expected 'tree' or 'greedy-dag'" + raise ValueError(msg) -class CostModel(Protocol, Generic[COST]): +class TreeCostModel(Protocol, Generic[COST]): """ A cost model for an e-graph. Used to determine the cost of an expression based on its structure and the costs of its sub-expressions. @@ -2623,6 +2896,31 @@ def __call__(self, egraph: EGraph, expr: BaseExpr, children_costs: list[COST]) - raise NotImplementedError +CostModel: TypeAlias = TreeCostModel + + +class ComparableAdd(Comparable, Protocol): + def __add__(self, other: Self) -> Self: ... + + +DAG_COST = TypeVar("DAG_COST", bound=ComparableAdd) + + +@dataclass(frozen=True) +class DagCostModel(Generic[DAG_COST]): + """ + An additive marginal cost model for tree or greedy-DAG extraction. + + ``marginal_cost`` excludes selected children and container elements. + Costs are combined with Python ``+``. The operation must be associative, + commutative, monotone, and have ``identity`` as a two-sided identity. + Cost values must also have a total order and be effectively immutable. + """ + + marginal_cost: Callable[[EGraph, BaseExpr], DAG_COST] + identity: DAG_COST + + def default_cost_model(egraph: EGraph, expr: BaseExpr, children_costs: list[int]) -> int: """ A default cost model for an e-graph, which looks up costs set on function calls, or uses 1 as the default cost. @@ -2648,84 +2946,6 @@ def default_cost_model(egraph: EGraph, expr: BaseExpr, children_costs: list[int] return sum(children_costs, start=self_cost) -class ComparableAddSub(Comparable, Protocol): - def __add__(self, other: Self) -> Self: ... - def __sub__(self, other: Self) -> Self: ... - - -DAG_COST = TypeVar("DAG_COST", bound=ComparableAddSub) - - -@dataclass -class GreedyDagCost(Generic[DAG_COST]): - """ - Cost of a DAG, which stores children costs. Use `.total` to get the underlying cost. - """ - - total: DAG_COST - _costs: dict[TypedExprDecl, DAG_COST] = field(repr=False) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, GreedyDagCost): - return NotImplemented - return self.total == other.total - - def __lt__(self, other: Self) -> bool: - return self.total < other.total - - def __le__(self, other: Self) -> bool: - return self.total <= other.total - - def __gt__(self, other: Self) -> bool: - return self.total > other.total - - def __ge__(self, other: Self) -> bool: - return self.total >= other.total - - def __hash__(self) -> int: - return hash(self.total) - - -@dataclass -class GreedyDagCostModel(CostModel[GreedyDagCost[DAG_COST]]): - """ - A cost model which will count duplicate nodes only once. - - Should have similar behavior as https://github.com/egraphs-good/extraction-gym/blob/main/src/extract/greedy_dag.rs - but implemented as a cost model that will be used with the default extractor. - """ - - base: CostModel[DAG_COST] - - def __call__( - self, egraph: EGraph, expr: BaseExpr, children_costs: list[GreedyDagCost[DAG_COST]] - ) -> GreedyDagCost[DAG_COST]: - cost = self.base(egraph, expr, [c.total for c in children_costs]) - for c in children_costs: - cost -= c.total - costs = {} - for c in children_costs: - costs.update(c._costs) - total = sum(costs.values(), start=cost) - costs[to_runtime_expr(expr).__egg_typed_expr__] = cost - return GreedyDagCost(total, costs) - - -@overload -def greedy_dag_cost_model() -> CostModel[GreedyDagCost[int]]: ... - - -@overload -def greedy_dag_cost_model(base: CostModel[DAG_COST]) -> CostModel[GreedyDagCost[DAG_COST]]: ... - - -def greedy_dag_cost_model(base: CostModel[Any] = default_cost_model) -> CostModel[GreedyDagCost[Any]]: - """ - Creates a greedy dag cost model from a base cost model. - """ - return GreedyDagCostModel(base or default_cost_model) - - def get_callable_cost(fn: ExprCallable) -> int | None: """ Returns the cost of a callable, if it has one set. Otherwise returns None. @@ -2745,9 +2965,8 @@ class _CostModel(Generic[COST]): egraph: EGraph enode_cost_results: dict[tuple[str, tuple[bindings.Value, ...]], int] = field(default_factory=dict) enode_cost_expressions: list[RuntimeExpr] = field(default_factory=list) - fold_results: dict[tuple[int, tuple[COST, ...]], COST] = field(default_factory=dict) + enode_cost_argument_values: list[dict[TypedExprDecl, bindings.Value]] = field(default_factory=list) base_value_cost_results: dict[tuple[str, bindings.Value], COST] = field(default_factory=dict) - container_cost_results: dict[tuple[str, bindings.Value, tuple[COST, ...]], COST] = field(default_factory=dict) def call_model(self, expr: RuntimeExpr, children_costs: list[COST]) -> COST: return self.model(self.egraph, cast("BaseExpr", expr), children_costs) @@ -2758,13 +2977,9 @@ def call_model(self, expr: RuntimeExpr, children_costs: list[COST]) -> COST: # raise ValueError(msg) def fold(self, _fn: str, index: int, children_costs: list[COST]) -> COST: - try: - return self.fold_results[(index, tuple(children_costs))] - except KeyError: - pass - expr = self.enode_cost_expressions[index] - return self.call_model(expr, children_costs) + with _cost_model_callback_values(self.egraph, self.enode_cost_argument_values[index]): + return self.call_model(expr, children_costs) # enode cost is only ever called right before fold, for the head_cost def enode_cost(self, name: str, args: list[bindings.Value]) -> int: @@ -2772,12 +2987,14 @@ def enode_cost(self, name: str, args: list[bindings.Value]) -> int: return self.enode_cost_results[(name, tuple(args))] except KeyError: pass - res = self.egraph._values_to_expr(args, name) - if res is None: + callback = self.egraph._values_to_expr_and_callback_values(args, name) + if callback is None: msg = f"Cannot compute custom cost for unknown egg function {name!r}" raise ValueError(msg) + res, argument_values = callback index = len(self.enode_cost_expressions) self.enode_cost_expressions.append(res) + self.enode_cost_argument_values.append(argument_values) self.enode_cost_results[(name, tuple(args))] = index return index @@ -2791,13 +3008,54 @@ def base_value_cost(self, tp: str, value: bindings.Value) -> COST: self.egraph.__egg_decls__, TypedExprDecl(type_ref, self.egraph._state.value_to_expr(type_ref, value)), ) - res = self.call_model(expr, []) + with _cost_model_callback_values(self.egraph, {expr.__egg_typed_expr__: value}): + res = self.call_model(expr, []) self.base_value_cost_results[(tp, value)] = res return res def container_cost(self, tp: str, value: bindings.Value, element_costs: list[COST]) -> COST: + type_ref = self.egraph._state.egg_sort_to_type_ref[tp] + expr = RuntimeExpr.__from_values__( + self.egraph.__egg_decls__, + TypedExprDecl(type_ref, self.egraph._state.value_to_expr(type_ref, value)), + ) + with _cost_model_callback_values(self.egraph, {expr.__egg_typed_expr__: value}): + return self.call_model(expr, element_costs) + + def to_bindings_cost_model(self) -> bindings.CostModel[COST, int]: + return bindings.CostModel(self.fold, self.enode_cost, self.container_cost, self.base_value_cost) + + +@dataclass +class _DagCostModel(Generic[DAG_COST]): + """Adapts a high-level marginal model to the raw binding callbacks.""" + + model: DagCostModel[DAG_COST] + egraph: EGraph + enode_cost_results: dict[tuple[str, tuple[bindings.Value, ...]], DAG_COST] = field(default_factory=dict) + base_value_cost_results: dict[tuple[str, bindings.Value], DAG_COST] = field(default_factory=dict) + container_cost_results: dict[tuple[str, bindings.Value], DAG_COST] = field(default_factory=dict) + + def enode_cost(self, name: str, args: list[bindings.Value]) -> DAG_COST: + key = (name, tuple(args)) + try: + return self.enode_cost_results[key] + except KeyError: + pass + callback = self.egraph._values_to_expr_and_callback_values(args, name) + if callback is None: + msg = f"Cannot compute custom cost for unknown egg function {name!r}" + raise ValueError(msg) + expr, argument_values = callback + with _cost_model_callback_values(self.egraph, argument_values): + result = self.model.marginal_cost(self.egraph, cast("BaseExpr", expr)) + self.enode_cost_results[key] = result + return result + + def base_value_cost(self, tp: str, value: bindings.Value) -> DAG_COST: + key = (tp, value) try: - return self.container_cost_results[(tp, value, tuple(element_costs))] + return self.base_value_cost_results[key] except KeyError: pass type_ref = self.egraph._state.egg_sort_to_type_ref[tp] @@ -2805,9 +3063,31 @@ def container_cost(self, tp: str, value: bindings.Value, element_costs: list[COS self.egraph.__egg_decls__, TypedExprDecl(type_ref, self.egraph._state.value_to_expr(type_ref, value)), ) - res = self.call_model(expr, element_costs) - self.container_cost_results[(tp, value, tuple(element_costs))] = res - return res + with _cost_model_callback_values(self.egraph, {expr.__egg_typed_expr__: value}): + result = self.model.marginal_cost(self.egraph, cast("BaseExpr", expr)) + self.base_value_cost_results[key] = result + return result - def to_bindings_cost_model(self) -> bindings.CostModel[COST, int]: - return bindings.CostModel(self.fold, self.enode_cost, self.container_cost, self.base_value_cost) + def container_cost(self, tp: str, value: bindings.Value) -> DAG_COST: + key = (tp, value) + try: + return self.container_cost_results[key] + except KeyError: + pass + type_ref = self.egraph._state.egg_sort_to_type_ref[tp] + expr = RuntimeExpr.__from_values__( + self.egraph.__egg_decls__, + TypedExprDecl(type_ref, self.egraph._state.value_to_expr(type_ref, value)), + ) + with _cost_model_callback_values(self.egraph, {expr.__egg_typed_expr__: value}): + result = self.model.marginal_cost(self.egraph, cast("BaseExpr", expr)) + self.container_cost_results[key] = result + return result + + def to_bindings_cost_model(self) -> bindings.DagCostModel[DAG_COST]: + return bindings.DagCostModel( + self.model.identity, + self.enode_cost, + self.container_cost, + self.base_value_cost, + ) diff --git a/python/egglog/egraph_state.py b/python/egglog/egraph_state.py index 985ae462..c2347f20 100644 --- a/python/egglog/egraph_state.py +++ b/python/egglog/egraph_state.py @@ -496,7 +496,7 @@ def command_to_egg(self, cmd: CommandDecl, ruleset: Ident) -> bindings._Command self.rule_name_to_command_decl[f"{serialized_name}{suffix}"] = cmd self.rule_name_to_command_decl[f"{reported_name}{suffix}"] = cmd return egg_cmd - case RuleDecl(head, body, name, eval_mode): + case RuleDecl(head, body, name, eval_mode, no_decomp): if not name: name = str(self.rule_name_counter) self.rule_name_counter += 1 @@ -517,6 +517,7 @@ def command_to_egg(self, cmd: CommandDecl, ruleset: Ident) -> bindings._Command name or "", str(ruleset), binding_eval_mode, + no_decomp, ) ) case DefaultRewriteDecl(ref, expr, subsume): @@ -620,15 +621,62 @@ def create_cost_table(self, ref: CallableRef) -> str: """ if ref in self.cost_table_names: return self.cost_table_names[ref] - base_name = f"cost_table_{self.callable_ref_to_egg(ref)[0]}" - name = self._allocate_name((base_name,), self._backend_symbol_is_occupied) + name = f"cost_table_{self.callable_ref_to_egg(ref)[0]}" signature = self.__egg_decls__.get_callable_decl(ref).signature assert isinstance(signature, FunctionSignature), "Can only add cost tables for functions" - signature = replace(signature, return_type=TypeRefWithVars(Ident.builtin("i64"))) - self.run_program(bindings.FunctionCommand(span(), name, self._signature_to_egg_schema(signature), None)) + target_schema = self._signature_to_egg_schema(signature) + schema = self._signature_to_egg_schema(replace(signature, return_type=TypeRefWithVars(Ident.builtin("i64")))) + + # egglog-experimental's DynamicCostModel probes this exact canonical + # name, so choosing a generated suffix would silently ignore costs. + # Aliases of the same backend callable share a cost table, but overloaded + # callables with incompatible schemas cannot: the backend protocol names + # cost tables only by the callable's backend symbol. A user-declared raw + # cost table can also be reused, but only when it is a bodyless function + # with the protocol's input sorts and i64 output. + if not self._has_compatible_cost_table_target(name, target_schema): + existing_refs = self.egg_fn_to_callable_refs.get(name, set()) + compatible_raw_table = bool(existing_refs) + for existing_ref in existing_refs: + existing_decl = self.__egg_decls__.get_callable_decl(existing_ref) + if not ( + isinstance(existing_decl, FunctionDecl) + and not existing_decl.builtin + and existing_decl.body is None + and isinstance(existing_decl.signature, FunctionSignature) + ): + compatible_raw_table = False + break + existing_schema = self._signature_to_egg_schema(existing_decl.signature) + if existing_schema.input != schema.input or existing_schema.output != schema.output: + compatible_raw_table = False + break + if existing_refs and not compatible_raw_table: + msg = ( + f"Canonical dynamic-cost table {name!r} is already used by an incompatible callable; " + "it must be a bodyless function with the target's input sorts and i64 output" + ) + raise ValueError(msg) + if not compatible_raw_table: + if self._backend_symbol_is_occupied(name): + raise ValueError(f"Canonical dynamic-cost table name {name!r} is already in use") + self.run_program(bindings.FunctionCommand(span(), name, schema, None)) self.cost_table_names[ref] = name return name + def _has_compatible_cost_table_target(self, name: str, target_schema: bindings.Schema) -> bool: + """Validate every callable already sharing a canonical dynamic-cost table.""" + existing_cost_refs = [ref for ref, existing_name in self.cost_table_names.items() if existing_name == name] + for existing_ref in existing_cost_refs: + existing_signature = self.__egg_decls__.get_callable_decl(existing_ref).signature + assert isinstance(existing_signature, FunctionSignature) + existing_schema = self._signature_to_egg_schema(existing_signature) + if existing_schema.input != target_schema.input or existing_schema.output != target_schema.output: + raise ValueError( + f"Canonical dynamic-cost table {name!r} already serves a callable with an incompatible schema" + ) + return bool(existing_cost_refs) + def fact_to_egg(self, fact: FactDecl, *, expr_to_let: bool = False) -> bindings._Fact: match fact: case EqDecl(tp, left, right): diff --git a/python/egglog/exp/any_expr_example.ipynb b/python/egglog/exp/any_expr_example.ipynb index e65e757d..52933585 100644 --- a/python/egglog/exp/any_expr_example.ipynb +++ b/python/egglog/exp/any_expr_example.ipynb @@ -370,7 +370,7 @@ "# egraph.register(x)\n", "# egraph.run(given_ruleset.saturate())\n", "# # egraph.display()\n", - "# print(str(egraph.extract(x))) # $, cost_model=greedy_dag_cost_model())))\n", + "# print(str(egraph.extract(x, extractor=\"greedy-dag\")))\n", "\n", "# a" ] diff --git a/python/egglog/exp/array_api_jit.py b/python/egglog/exp/array_api_jit.py index 1fa9624e..3435f01c 100644 --- a/python/egglog/exp/array_api_jit.py +++ b/python/egglog/exp/array_api_jit.py @@ -5,7 +5,7 @@ import numpy as np from opentelemetry import trace -from egglog import EGraph, bindings, greedy_dag_cost_model +from egglog import EGraph, bindings from egglog.exp.array_api import NDArray, set_array_api_egraph from egglog.exp.array_api_numba import array_api_numba_schedule from egglog.exp.array_api_program_gen import EvalProgram, array_api_program_gen_schedule, ndarray_function_two_program @@ -58,7 +58,7 @@ def function_to_program(fn: Callable, save_egglog_string: bool) -> tuple[EGraph, res = fn(NDArray.var(arg1), NDArray.var(arg2)) egraph.register(res) egraph.run(array_api_numba_schedule) - res_optimized = egraph.extract(res, cost_model=greedy_dag_cost_model()) + res_optimized = egraph.extract(res, extractor="greedy-dag") return ( egraph, diff --git a/python/egglog/exp/param_eq/pipeline.py b/python/egglog/exp/param_eq/pipeline.py index 89f4fce5..69c9fb2b 100644 --- a/python/egglog/exp/param_eq/pipeline.py +++ b/python/egglog/exp/param_eq/pipeline.py @@ -19,10 +19,6 @@ BACKOFF_BAN_LENGTH = 30 -@function(builtin=True, egg_fn="f64-is-finite") -def _f64_is_finite(value: f64) -> Unit: ... - - # Keep derived map operations as explicitly typed folds in this research # module; only map_fold_kv is a backend primitive and public builtin. # Store discovered constants in a global map so semi-naive analysis can join @@ -66,22 +62,22 @@ def _f64_is_finite(value: f64) -> Unit: ... @ruleset def shared_analysis_rules(a: f64) -> Iterable[RewriteOrRule]: - yield rewrite(exp(Num(a)), subsume=True).to(Num(a.exp()), _f64_is_finite(a.exp())) - yield rewrite(log(Num(a)), subsume=True).to(Num(a.log()), a > 0.0, _f64_is_finite(a.log())) + yield rewrite(exp(Num(a)), subsume=True).to(Num(a.exp()), a.exp().is_finite()) + yield rewrite(log(Num(a)), subsume=True).to(Num(a.log()), a > 0.0, a.log().is_finite()) yield rule(log(Num(a)), a <= 0.0).then(panic("Log of non-positive number")) - yield rewrite(abs(Num(a)), subsume=True).to(Num(abs(a)), _f64_is_finite(abs(a))) + yield rewrite(abs(Num(a)), subsume=True).to(Num(abs(a)), abs(a).is_finite()) @ruleset def binary_analysis_rules(x: Num, a: f64, b: f64) -> Iterable[RewriteOrRule]: - yield rewrite(Num(a) / Num(b), subsume=True).to(Num(a / b), b != f64(0.0), _f64_is_finite(a / b)) + yield rewrite(Num(a) / Num(b), subsume=True).to(Num(a / b), b != f64(0.0), (a / b).is_finite()) yield rule(x / Num(0.0)).then(panic("Division by zero")) - yield rewrite(Num(a) - Num(b), subsume=True).to(Num(a - b), _f64_is_finite(a - b)) - yield rewrite(Num(a) * Num(b), subsume=True).to(Num(a * b), _f64_is_finite(a * b)) - yield rewrite(Num(a) + Num(b), subsume=True).to(Num(a + b), _f64_is_finite(a + b)) + yield rewrite(Num(a) - Num(b), subsume=True).to(Num(a - b), (a - b).is_finite()) + yield rewrite(Num(a) * Num(b), subsume=True).to(Num(a * b), (a * b).is_finite()) + yield rewrite(Num(a) + Num(b), subsume=True).to(Num(a + b), (a + b).is_finite()) - yield rewrite(Num(a) ** Num(b), subsume=True).to(Num(a**b), _f64_is_finite(a**b)) - yield rewrite(sqrt(Num(a)), subsume=True).to(Num(a.sqrt()), a >= 0.0, _f64_is_finite(a.sqrt())) + yield rewrite(Num(a) ** Num(b), subsume=True).to(Num(a**b), (a**b).is_finite()) + yield rewrite(sqrt(Num(a)), subsume=True).to(Num(a.sqrt()), a >= 0.0, a.sqrt().is_finite()) yield rule(sqrt(Num(a)), a < 0.0).then(panic("Sqrt of negative number")) # cancellations diff --git a/python/egglog/pretty.py b/python/egglog/pretty.py index c85d04a8..1985df9b 100644 --- a/python/egglog/pretty.py +++ b/python/egglog/pretty.py @@ -204,7 +204,7 @@ def __call__(self, decl: AllDecls, toplevel: bool = False) -> None: # noqa: C90 self(rhs) for cond in conditions: self(cond) - case RuleDecl(head, body, _, _): + case RuleDecl(head, body, _, _, _): for action in head: self(action) for fact in body: @@ -343,12 +343,14 @@ def uncached( # noqa: C901, PLR0911, PLR0912 args = ", ".join(map(self, (rhs, *conditions))) fn = "rewrite" if isinstance(decl, RewriteDecl) else "birewrite" return f"{fn}({self(lhs)}).to({args})", "rewrite" - case RuleDecl(head, body, name, eval_mode): + case RuleDecl(head, body, name, eval_mode, no_decomp): args = list(map(self, body)) if name: args.append(f"name={name!r}") if eval_mode != "seminaive": args.append(f"eval_mode={eval_mode!r}") + if no_decomp: + args.append("no_decomp=True") r = ", ".join(map(self, head)) return f"rule({', '.join(args)}).then({r})", "rule" case SetDecl(_, lhs, rhs): diff --git a/python/tests/__snapshots__/test_array_api/test_jit[lda][code].py b/python/tests/__snapshots__/test_array_api/test_jit[lda][code].py index bf59a430..c035c457 100644 --- a/python/tests/__snapshots__/test_array_api/test_jit[lda][code].py +++ b/python/tests/__snapshots__/test_array_api/test_jit[lda][code].py @@ -39,7 +39,7 @@ def __fn(X, y): _28 = _27 / np.array(_26.shape[0]) _29 = np.sqrt(_28) _30 = _29 == np.array(0) - _29[_30] = np.array((150 / 150)) + _29[_30] = np.array(float(1)) _31 = _21 / _29 _32 = _17 * _31 _33 = np.linalg.svd(_32, full_matrices=False) diff --git a/python/tests/__snapshots__/test_array_api/test_jit[lda][expr].py b/python/tests/__snapshots__/test_array_api/test_jit[lda][expr].py index cda4a169..5c1705af 100644 --- a/python/tests/__snapshots__/test_array_api/test_jit[lda][expr].py +++ b/python/tests/__snapshots__/test_array_api/test_jit[lda][expr].py @@ -33,7 +33,9 @@ _NDArray_9 = square(_NDArray_8 - expand_dims(sum(_NDArray_8, OptionalIntOrTuple.int(Int(0))) / NDArray(RecursiveValue(Value.from_int(_NDArray_8.shape[Int(0)]))))) _NDArray_10 = sqrt(sum(_NDArray_9, OptionalIntOrTuple.int(Int(0))) / NDArray(RecursiveValue(Value.from_int(_NDArray_9.shape[Int(0)])))) _NDArray_11 = copy(_NDArray_10) -_NDArray_11[IndexKey.ndarray(_NDArray_10 == NDArray(RecursiveValue(Value.from_int(Int(0)))))] = NDArray(RecursiveValue(Value.from_int(Int(150)) / Value.from_int(Int(150)))) +_NDArray_11[IndexKey.ndarray(_NDArray_10 == NDArray(RecursiveValue(Value.from_int(Int(0)))))] = NDArray( + RecursiveValue(Value.from_float(Float.rational(BigRat(BigInt.from_string("1"), BigInt.from_string("1"))))) +) _TupleNDArray_1 = svd_( sqrt( asarray( diff --git a/python/tests/param_eq/test_pipeline.py b/python/tests/param_eq/test_pipeline.py index 2e254af1..73a4ed69 100644 --- a/python/tests/param_eq/test_pipeline.py +++ b/python/tests/param_eq/test_pipeline.py @@ -90,6 +90,14 @@ def test_nonfinite_constant_results_are_not_folded(source: str) -> None: assert parse_expression(report.extracted) == parse_expression(source) +@pytest.mark.param_eq_smoke +@pytest.mark.parametrize(("source", "expected"), [("exp(1.0)", math.e), ("sqrt(4.0)", 2.0)]) +def test_finite_constant_results_are_folded(source: str, expected: float) -> None: + report = run_paper_pipeline(parse_expression(source)) + + assert parse_expression(report.extracted) == Num(expected) + + @pytest.mark.param_eq_smoke @pytest.mark.parametrize("source", ["(-1.5) ** 0.25", "1e308 * 1e308", "(1e308 * x0) * 1e308"]) def test_container_pipeline_rejects_nonfinite_coefficient_normalization(source: str) -> None: diff --git a/python/tests/test_array_api.py b/python/tests/test_array_api.py index d9eea325..463f9ee7 100644 --- a/python/tests/test_array_api.py +++ b/python/tests/test_array_api.py @@ -12,7 +12,7 @@ from sklearn import config_context, datasets from sklearn.discriminant_analysis import LinearDiscriminantAnalysis -from egglog import greedy_dag_cost_model, set_current_ruleset +from egglog import set_current_ruleset from egglog.exp.array_api import * from egglog.exp.array_api import NDArray, Value from egglog.exp.array_api_jit import function_to_program, jit @@ -379,7 +379,7 @@ def test_program_compile(program: Program, snapshot_py): egraph = EGraph() egraph.register(program) egraph.run(array_api_numba_schedule) - simplified_program = egraph.extract(program, cost_model=greedy_dag_cost_model()) + simplified_program = egraph.extract(program, extractor="greedy-dag") assert str(simplified_program) == snapshot_py(name="expr") egraph = EGraph() egraph.register(simplified_program.compile()) diff --git a/python/tests/test_bindings.py b/python/tests/test_bindings.py index 8054d848..f7fdd95e 100644 --- a/python/tests/test_bindings.py +++ b/python/tests/test_bindings.py @@ -79,6 +79,23 @@ def extract_best_term(program: str) -> str: class TestEGraph: + def test_per_egraph_configuration(self): + configured = EGraph(num_threads=2, no_decomp=True) + default = EGraph() + + assert configured.num_threads() == 2 + assert configured.no_decomp() + assert default.num_threads() == 1 + assert not default.no_decomp() + + configured.set_num_threads(1) + configured.set_no_decomp(False) + + assert configured.num_threads() == 1 + assert not configured.no_decomp() + assert default.num_threads() == 1 + assert not default.no_decomp() + def test_parse_program(self, snapshot_py): res = EGraph().parse_program( """(datatype Math @@ -422,6 +439,77 @@ def test_extract_value_reports_extraction_failure(self): with pytest.raises(EggSmolError, match="Unable to find any valid extraction"): egraph.extract_value(value, sort) + @pytest.mark.parametrize("extractor", ["tree", "greedy-dag"]) + def test_dag_cost_model_batch_extraction(self, extractor): + egraph = EGraph() + egraph.parse_and_run_program("(datatype Expr (Num i64)) (let root (Num 1)) (union root (Num 2))") + sort, value = egraph.eval_expr(Call(DUMMY_SPAN, "Num", [Lit(DUMMY_SPAN, Int(1))])) + model = DagCostModel( + 0, + lambda name, args: 1, + lambda name, value: 0, + lambda name, value: 1, + ) + + termdag, best = extract_best_with_dag_cost_model(egraph, [(sort, value)], model, extractor=extractor) + assert best[0] is not None + cost, term = best[0] + assert cost == 2 + assert termdag.to_string(term) in {"(Num 1)", "(Num 2)"} + + def test_tree_extractor_observes_post_construction_mutation(self): + egraph = EGraph() + egraph.parse_and_run_program("(datatype Expr (Num i64)) (let root (Num 1))") + sort, value = egraph.eval_expr(Call(DUMMY_SPAN, "Num", [Lit(DUMMY_SPAN, Int(1))])) + callback_count = 0 + + def enode_cost(name, args): + nonlocal callback_count + callback_count += 1 + return 0 if name == "Num" and egraph.value_to_i64(args[0]) == 2 else 1 + + model = CostModel( + lambda name, annotation, children: annotation + sum(children), + enode_cost, + lambda name, value, children: sum(children), + lambda name, value: 0, + ) + extractor = Extractor([sort], egraph, model) + + assert callback_count == 0 + first_dag = TermDag() + _, first = extractor.extract_best(egraph, first_dag, value, sort) + assert first_dag.to_string(first) == "(Num 1)" + first_callback_count = callback_count + + egraph.parse_and_run_program("(union root (Num 2))") + second_dag = TermDag() + _, second = extractor.extract_best(egraph, second_dag, value, sort) + + assert second_dag.to_string(second) == "(Num 2)" + assert callback_count > first_callback_count + + def test_tree_cost_callback_failure_does_not_mutate_termdag(self): + egraph = EGraph() + sort, value = egraph.eval_expr(Lit(DUMMY_SPAN, Int(1))) + + def fail(name, value): + msg = "base cost failed" + raise LookupError(msg) + + model = CostModel( + lambda name, annotation, children: annotation, + lambda name, args: 0, + lambda name, value, children: 0, + fail, + ) + extractor = Extractor([sort], egraph, model) + termdag = TermDag() + + with pytest.raises(LookupError, match="base cost failed"): + extractor.extract_best(egraph, termdag, value, sort) + assert termdag.size() == 0 + def test_sort_alias(self): # From map example egraph = EGraph() diff --git a/python/tests/test_high_level.py b/python/tests/test_high_level.py index 9ece37dc..c8eb6e55 100644 --- a/python/tests/test_high_level.py +++ b/python/tests/test_high_level.py @@ -7,6 +7,7 @@ import pathlib from collections.abc import Callable, Iterator from copy import copy +from dataclasses import dataclass from fractions import Fraction from functools import partial from typing import ClassVar, TypeAlias, TypeVar, cast @@ -61,6 +62,27 @@ def test_rule_eval_mode(eval_mode: RuleEvalMode) -> None: egraph.check(rel(i64(1))) +def test_per_egraph_configuration() -> None: + egraph = EGraph(num_threads=2, no_decomp=True) + + assert egraph.num_threads() == 2 + assert egraph.no_decomp() + egraph.set_num_threads(1) + egraph.set_no_decomp(False) + assert egraph.num_threads() == 1 + assert not egraph.no_decomp() + + +def test_rule_no_decomp_reaches_backend() -> None: + rel = relation("no_decomp_rel", i64) + x = var("x", i64) + egraph = EGraph(save_egglog_string=True) + + egraph.register(rule(rel(x), no_decomp=True).then(rel(x + 1))) + + assert ":no-decomp" in egraph.as_egglog_string + + def test_eqsat_basic(): egraph = EGraph() @@ -308,7 +330,7 @@ def test_generated_callable_name_avoids_an_existing_cost_table() -> None: assert state.callable_ref_to_egg(FunctionRef(conflict))[0] == "pkg_cost_cost_table_f" -def test_generated_cost_table_name_avoids_an_existing_callable() -> None: +def test_canonical_cost_table_rejects_an_incompatible_callable() -> None: state = EGraph(save_egglog_string=True)._state state.__egg_decls__ |= cast("HasDeclarations", i64) ret = Ident("CostRet", "pkg.cost") @@ -323,7 +345,27 @@ def test_generated_cost_table_name_avoids_an_existing_callable() -> None: ) assert state.callable_ref_to_egg(FunctionRef(conflict))[0] == "cost_table_f" - assert state.create_cost_table(FunctionRef(fn)) == "cost_table_f_1" + with pytest.raises(ValueError, match="already used by an incompatible callable"): + state.create_cost_table(FunctionRef(fn)) + + +def test_canonical_cost_table_reuses_a_compatible_raw_table() -> None: + state = EGraph(save_egglog_string=True)._state + state.__egg_decls__ |= cast("HasDeclarations", i64) + ret = Ident("CostRet", "pkg.cost") + fn = Ident("f", "pkg.cost") + raw_cost = Ident("cost_table_f", "pkg.cost") + state.__egg_decls__ |= Declarations( + _classes={ret: ClassDecl()}, + _functions={ + fn: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret))), + raw_cost: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(Ident.builtin("i64")))), + }, + ) + + assert state.callable_ref_to_egg(FunctionRef(raw_cost))[0] == "cost_table_f" + assert state.create_cost_table(FunctionRef(fn)) == "cost_table_f" + assert state.cost_table_names[FunctionRef(fn)] == "cost_table_f" @pytest.mark.parametrize( @@ -981,6 +1023,45 @@ def test_f64_math_primitives() -> None: assert egraph.extract(f64(4.0).sqrt()).value == pytest.approx(2.0) +def test_f64_is_finite_predicate() -> None: + egraph = EGraph() + egraph.check(f64(1.0).is_finite()) + for value in (float("nan"), float("inf"), float("-inf")): + with pytest.raises(EggSmolError): + egraph.check(f64(value).is_finite()) + + +def test_rational_like_operations() -> None: + assert "RationalLike" in egg_builtins.__all__ + assert EGraph().extract(Rational(1, 2) + Fraction(1, 3)).value == Fraction(5, 6) + assert EGraph().extract(Fraction(1, 3) + Rational(1, 2)).value == Fraction(5, 6) + assert EGraph().extract(Rational(1, 2) - 1).value == Fraction(-1, 2) + assert EGraph().extract(1 - Rational(1, 2)).value == Fraction(1, 2) + assert EGraph().extract(Rational(2, 3) * Fraction(3, 4)).value == Fraction(1, 2) + assert EGraph().extract(Fraction(3, 4) * Rational(2, 3)).value == Fraction(1, 2) + assert EGraph().extract(Rational(1, 2) / 2).value == Fraction(1, 4) + assert EGraph().extract(1 / Rational(1, 2)).value == Fraction(2, 1) + assert EGraph().extract(Rational(2, 1) ** 3).value == Fraction(8, 1) + assert EGraph().extract(2 ** Rational(3, 1)).value == Fraction(8, 1) + assert EGraph().extract(Rational(1, 2).min(Fraction(1, 3))).value == Fraction(1, 3) + assert EGraph().extract(Rational(1, 2).max(1)).value == Fraction(1, 1) + + egraph = EGraph() + egraph.check(Rational(1, 2) < Fraction(2, 3)) + egraph.check(Rational(2, 3) > Fraction(1, 2)) + egraph.check(Rational(1, 2) <= Fraction(1, 2)) + egraph.check(Rational(1, 2) >= Fraction(1, 2)) + + +def test_rational_partial_operations_remain_undefined() -> None: + with pytest.raises(EggSmolError): + EGraph().extract(Rational(1, 2) / 0) + with pytest.raises(EggSmolError): + EGraph().extract(Rational(2, 1) ** -1) + with pytest.raises(EggSmolError): + EGraph().check(Rational(2, 3) < Fraction(1, 2)) + + def test_bigrat_to_i64_is_exact_and_bounded() -> None: assert EGraph().extract(BigRat(4, 2).to_i64()) == i64(2) assert EGraph().extract(BigRat(1, 2) + i64(1)).value == Fraction(3, 2) @@ -2189,6 +2270,90 @@ def __sub__(self, other: E) -> E: ... assert egraph.extract(E(2), include_cost=True) == (E(1) + E(1), 203) egraph.register(set_cost(E(5) - E(3), 198)) assert egraph.extract(E(2), include_cost=True) == (E(5) - E(3), 202) + assert egraph.extract(E(2), include_cost=True, extractor="greedy-dag") == (E(1) + E(1), 102) + + +def test_dynamic_cost_reuses_a_compatible_canonical_table() -> None: + @function(egg_fn="canonical_cost_target") + def target(x: i64Like) -> i64: ... + + @function(egg_fn="cost_table_canonical_cost_target") + def raw_cost(x: i64Like) -> i64: ... + + egraph = EGraph() + egraph.register( + set_(raw_cost(2)).to(i64(5)), + set_(target(1)).to(i64(2)), + set_cost(target(1), 7), + ) + + assert egraph.lookup_function_value(raw_cost(1)) == i64(7) + assert egraph.lookup_function_value(raw_cost(2)) == i64(5) + assert egraph.has_custom_cost(target) + + +def test_freeze_preserves_a_reused_canonical_cost_table_as_raw_rows_and_costs() -> None: + @function(egg_fn="freeze_cost_target") + def target(x: i64Like) -> i64: ... + + @function(egg_fn="cost_table_freeze_cost_target") + def raw_cost(x: i64Like) -> i64: ... + + egraph = EGraph( + set_(raw_cost(2)).to(i64(5)), + set_(target(1)).to(i64(2)), + set_cost(target(1), 7), + ) + + rendered = str(egraph.freeze()) + assert "set_(raw_cost(2)).to(i64(5))" in rendered + assert "set_(raw_cost(1)).to(i64(7))" in rendered + assert "set_cost(target(2), 5)" in rendered + assert "set_cost(target(1), 7)" in rendered + + replayed = eval(rendered.removesuffix(".freeze()"), globals(), locals()) + assert isinstance(replayed, EGraph) + assert replayed.lookup_function_value(raw_cost(1)) == i64(7) + assert replayed.lookup_function_value(raw_cost(2)) == i64(5) + assert replayed.lookup_function_value(target(1)) == i64(2) + assert replayed.has_custom_cost(target) + + +def test_freeze_preserves_every_callable_alias_for_a_shared_cost_table() -> None: + @function(egg_fn="+", builtin=True) + def plus_alias(left: i64Like, right: i64Like) -> i64: ... + + egraph = EGraph( + set_cost(i64(1) + i64(2), 5), + set_cost(plus_alias(3, 4), 6), + ) + + rendered = str(egraph.freeze()) + replayed = eval(rendered.removesuffix(".freeze()"), globals(), locals()) + assert isinstance(replayed, EGraph) + assert replayed.has_custom_cost(i64.__add__) + assert replayed.has_custom_cost(plus_alias) + assert replayed.lookup_function_value(get_cost(i64(1) + i64(2))) == i64(5) + assert replayed.lookup_function_value(get_cost(plus_alias(3, 4))) == i64(6) + + +def test_dynamic_cost_rejects_an_incompatible_overload_without_recording_it() -> None: + egraph = EGraph() + egraph.register(set_cost(i64(1) + i64(2), 5)) + + with pytest.raises(ValueError, match="already serves a callable with an incompatible schema"): + egraph.register(set_cost(Rational(1, 2) + Rational(1, 3), 6)) + + assert egraph.has_custom_cost(i64.__add__) + assert not egraph.has_custom_cost(Rational.__add__) + + +def test_dynamic_cost_rejects_a_negative_literal() -> None: + class Costed(Expr): + def __init__(self, value: i64Like) -> None: ... + + with pytest.raises(ValueError, match="must be nonnegative"): + set_cost(Costed(1), -1) class TestScheduler: @@ -2398,6 +2563,89 @@ def ff(x: i64Like, y: i64Like) -> E: ... def gg() -> E: ... +@pytest.mark.parametrize("extractor", ["tree", "greedy-dag"]) +def test_extract_multiple_sequence_preserves_heterogeneous_roots(extractor: ExtractionMode) -> None: + class MultiRoot(Expr): + def __init__(self, value: i64Like) -> None: ... + + def __add__(self, other: MultiRoot) -> MultiRoot: ... + + @method(unextractable=True) + def opaque(self) -> MultiRoot: ... + + egraph = EGraph() + opaque = egraph.let("opaque_multi_root", MultiRoot(1).opaque()) + repeated = MultiRoot(0) + MultiRoot(0) + egraph.register( + union(MultiRoot(2)).with_(repeated), + set_cost(MultiRoot(2), 100), + set_cost(MultiRoot(0), 1), + ) + + extracted = egraph.extract_multiple([String("first"), opaque, MultiRoot(2)], 1, extractor=extractor) + + assert extracted == [[String("first")], [], [repeated]] + assert egraph.extract_multiple(i64(4), 1, extractor=extractor) == [i64(4)] + homogeneous: list[list[MultiRoot]] = egraph.extract_multiple([MultiRoot(2)], 1, extractor=extractor) + assert homogeneous == [[repeated]] + + +def test_extract_multiple_validates_batch_arguments() -> None: + with pytest.raises(ValueError, match="must be positive"): + EGraph().extract_multiple(i64(1), 0) + with pytest.raises(ValueError, match="at least one expression"): + EGraph().extract_multiple([], 1) + with pytest.raises(ValueError, match="Unknown extractor"): + EGraph().extract_multiple(i64(1), 1, extractor="unknown") # type: ignore[call-overload] + + +@pytest.mark.parametrize("extractor", ["tree", "greedy-dag"]) +def test_keep_best_compacts_and_allows_continued_iteration(extractor: ExtractionMode) -> None: + class CompactExpr(Expr): + def __init__(self, value: i64Like) -> None: ... + + def __add__(self, other: CompactExpr) -> CompactExpr: ... + + @function(merge=lambda old, new: new) + def target(key: i64Like) -> CompactExpr: ... + + @function(merge=lambda old, new: new) + def discarded(key: i64Like) -> CompactExpr: ... + + direct = CompactExpr(2) + shared = CompactExpr(1) + repeated = shared + shared + egraph = EGraph(save_egglog_string=True) + egraph.register( + union(direct).with_(repeated), + set_(target(0)).to(direct), + set_(discarded(0)).to(CompactExpr(9)), + set_cost(direct, 100), + set_cost(CompactExpr(1), 4), + ) + + egraph.keep_best(target, extractor=extractor) + + assert egraph.function_size(target) == 1 + assert egraph.function_size(discarded) == 0 + assert egraph.has_custom_cost(CompactExpr) + assert egraph.extract(target(0)) == repeated + + # Reusing an expression that was factored through a synthetic let must not + # reference the row that keep-best cleared. + egraph.register(repeated, set_(target(1)).to(repeated)) + assert egraph.function_size(target) == 2 + + +def test_keep_best_rejects_non_table_callable_before_compaction() -> None: + @function + def eager(value: i64Like) -> i64: + return cast("i64", value) + 1 + + with pytest.raises(ValueError, match="table-backed"): + EGraph().keep_best(eager) + + class TestCustomExtract: def test_literal_root(self) -> None: def is_even_cost_model(egraph: EGraph, expr: BaseExpr, children_costs: list[int]) -> int: @@ -2633,7 +2881,6 @@ def test_maybe_container_children_costs_match_python_value_order(self): assert cast("Maybe[i64]", extracted_none).value is None assert seen_none_children_costs == [] - @pytest.mark.xfail(reason="Errors dont bubble, just panic") def test_errors_bubble(self): def my_cost_model(egraph: EGraph, expr: BaseExpr, children_costs: list[int]) -> int: msg = "bad" @@ -2645,15 +2892,21 @@ def my_cost_model(egraph: EGraph, expr: BaseExpr, children_costs: list[int]) -> egraph.extract(i64(10), cost_model=my_cost_model) def test_dag_cost_model(self): + model = DagCostModel( + marginal_cost=lambda egraph, expr: default_cost_model(egraph, expr, []), + identity=0, + ) egraph = EGraph() expr = ff(1, 2) - res, cost = egraph.extract(expr, include_cost=True, cost_model=greedy_dag_cost_model()) - assert cost.total == 3 + res, cost = egraph.extract(expr, include_cost=True, cost_model=model, extractor="greedy-dag") + assert cost == 3 assert expr == res expr = ff(1, 1) - res, cost = egraph.extract(expr, include_cost=True, cost_model=greedy_dag_cost_model()) - assert cost.total == 2 + _, tree_cost = egraph.extract(expr, include_cost=True, cost_model=model) + res, cost = egraph.extract(expr, include_cost=True, cost_model=model, extractor="greedy-dag") + assert tree_cost == 3 + assert cost == 2 assert expr == res @function @@ -2663,10 +2916,276 @@ def bin(l: E, r: E) -> E: ... y = constant("y", E) expr = bin(x, bin(x, y)) egraph.register(expr) - res, cost = egraph.extract(expr, include_cost=True, cost_model=greedy_dag_cost_model()) - assert cost.total == 4 + res, cost = egraph.extract(expr, include_cost=True, cost_model=model, extractor="greedy-dag") + assert cost == 4 assert expr == res + @pytest.mark.parametrize( + ("model_kind", "extractor", "expected_cost"), + [ + pytest.param("tree", "tree", 7, id="tree-model"), + pytest.param("dag", "tree", 7, id="dag-model-tree-extractor"), + pytest.param("dag", "greedy-dag", 4, id="dag-model-greedy-dag-extractor"), + ], + ) + def test_default_cost_model_reads_dynamic_costs_during_callbacks( + self, + model_kind: str, + extractor: ExtractionMode, + expected_cost: int, + ) -> None: + class DynamicCostExpr(Expr): + def __init__(self, value: i64Like) -> None: ... + + def __add__(self, other: DynamicCostExpr) -> DynamicCostExpr: ... + + egraph = EGraph() + egraph.register( + union(DynamicCostExpr(2)).with_(DynamicCostExpr(1) + DynamicCostExpr(1)), + set_cost(DynamicCostExpr(2), 50), + set_cost(DynamicCostExpr(1), 2), + ) + + if model_kind == "tree": + result, cost = egraph.extract( + DynamicCostExpr(2), include_cost=True, cost_model=default_cost_model, extractor=extractor + ) + else: + model = DagCostModel( + marginal_cost=lambda callback_egraph, node: default_cost_model(callback_egraph, node, []), + identity=0, + ) + result, cost = egraph.extract(DynamicCostExpr(2), include_cost=True, cost_model=model, extractor=extractor) + + assert result == DynamicCostExpr(1) + DynamicCostExpr(1) + assert cost == expected_cost + + def test_tree_cost_model_can_lookup_table_with_a_primitive_callback_argument(self) -> None: + class LookupByPrimitive(Expr): + def __init__(self, value: i64Like) -> None: ... + + @function + def score(value: i64Like) -> i64: ... + + egraph = EGraph() + egraph.register(set_(score(3)).to(i64(17))) + + def lookup_cost(callback_egraph: EGraph, expr: BaseExpr, children_costs: list[int]) -> int: + if isinstance(expr, LookupByPrimitive): + args = get_callable_args(expr) + assert args is not None + value = callback_egraph.lookup_function_value(score(cast("i64", args[0]))) + assert value is not None + return int(value) + sum(children_costs) + return sum(children_costs) + + assert egraph.extract( + LookupByPrimitive(3), include_cost=True, cost_model=cast("TreeCostModel[int]", lookup_cost) + ) == ( + LookupByPrimitive(3), + 17, + ) + + def test_cost_model_callback_rejects_evaluating_new_lookup_arguments(self) -> None: + class LookupByPrimitive(Expr): + def __init__(self, value: i64Like) -> None: ... + + @function + def score(value: i64Like) -> i64: ... + + egraph = EGraph() + egraph.register(set_(score(99)).to(i64(17))) + + def lookup_cost(callback_egraph: EGraph, expr: BaseExpr, children_costs: list[int]) -> int: + if isinstance(expr, LookupByPrimitive): + callback_egraph.lookup_function_value(score(99)) + return sum(children_costs) + + with pytest.raises(ValueError, match="only look up tables using values supplied to the callback"): + egraph.extract(LookupByPrimitive(3), cost_model=cast("TreeCostModel[int]", lookup_cost)) + + assert egraph.lookup_function_value(score(99)) == i64(17) + + def test_cost_model_callback_requires_lookup_table_to_be_registered(self) -> None: + class LookupByPrimitive(Expr): + def __init__(self, value: i64Like) -> None: ... + + @function + def never_registered(value: i64Like) -> i64: ... + + def lookup_cost(callback_egraph: EGraph, expr: BaseExpr, children_costs: list[int]) -> int: + if isinstance(expr, LookupByPrimitive): + args = get_callable_args(expr) + assert args is not None + callback_egraph.lookup_function_value(never_registered(cast("i64", args[0]))) + return sum(children_costs) + + with pytest.raises(ValueError, match="must be registered before extraction starts"): + EGraph().extract(LookupByPrimitive(3), cost_model=cast("TreeCostModel[int]", lookup_cost)) + + def test_cost_model_callback_values_are_scoped_to_their_egraph(self) -> None: + class LookupByString(Expr): + def __init__(self, value: StringLike) -> None: ... + + @function + def score(value: StringLike) -> i64: ... + + lookup_egraph = EGraph() + lookup_egraph.register( + set_(score("padding")).to(i64(1)), + set_(score("needle")).to(i64(17)), + ) + + def lookup_cost(_callback_egraph: EGraph, expr: BaseExpr, children_costs: list[int]) -> int: + if isinstance(expr, LookupByString): + args = get_callable_args(expr) + assert args is not None + value = lookup_egraph.lookup_function_value(score(cast("String", args[0]))) + assert value is not None + return int(value) + sum(children_costs) + return sum(children_costs) + + source_egraph = EGraph() + assert source_egraph.extract( + LookupByString("needle"), include_cost=True, cost_model=cast("TreeCostModel[int]", lookup_cost) + ) == (LookupByString("needle"), 17) + + @pytest.mark.parametrize( + ("model_kind", "extractor"), + [ + pytest.param("tree", "tree", id="tree-model"), + pytest.param("dag", "greedy-dag", id="dag-model"), + ], + ) + def test_cost_model_callbacks_preserve_reverse_argument_order( + self, model_kind: str, extractor: ExtractionMode + ) -> None: + class ReverseCostResult(Expr): ... + + class ReverseCostSource(Expr): + def __init__(self, value: i64Like) -> None: ... + + @method(reverse_args=True) + def make(self, label: StringLike) -> ReverseCostResult: ... + + source = ReverseCostSource(3) + expr = source.make("label") + egraph = EGraph() + egraph.register(set_cost(expr, 9)) + seen_args: list[BaseExpr] = [] + + def marginal_cost(callback_egraph: EGraph, node: BaseExpr) -> int: + if get_callable_fn(node) != ReverseCostSource.make: + return 0 + args = get_callable_args(node) + assert args is not None + seen_args.extend(args) + return default_cost_model(callback_egraph, node, []) + + if model_kind == "tree": + + def tree_cost(callback_egraph: EGraph, node: BaseExpr, children_costs: list[int]) -> int: + return marginal_cost(callback_egraph, node) + sum(children_costs) + + result, cost = egraph.extract( + expr, + include_cost=True, + cost_model=cast("TreeCostModel[int]", tree_cost), + extractor=extractor, + ) + else: + result, cost = egraph.extract( + expr, + include_cost=True, + cost_model=DagCostModel(marginal_cost, 0), + extractor=extractor, + ) + + assert result == expr + assert cost == 9 + assert seen_args + assert len(seen_args) % 2 == 0 + for source_arg, label_arg in zip(seen_args[::2], seen_args[1::2], strict=True): + assert isinstance(source_arg, ReverseCostSource) + assert isinstance(label_arg, String) + assert label_arg.value == "label" + + def test_tree_cost_model_rejected_by_greedy_dag(self): + with pytest.raises(TypeError, match="requires a DagCostModel"): + EGraph().extract(i64(1), cost_model=default_cost_model, extractor="greedy-dag") + + def test_dag_marginal_error_bubbles(self): + def marginal_cost(egraph: EGraph, expr: BaseExpr) -> int: + del egraph, expr + msg = "marginal failed" + raise LookupError(msg) + + with pytest.raises(LookupError, match="marginal failed"): + EGraph().extract(i64(1), cost_model=DagCostModel(marginal_cost, 0)) + + def test_dag_add_error_bubbles(self): + @dataclass(frozen=True) + class AddErrorCost: + value: int + + def __add__(self, other: AddErrorCost) -> AddErrorCost: + del other + msg = "addition failed" + raise ArithmeticError(msg) + + def __lt__(self, other: AddErrorCost) -> bool: + return self.value < other.value + + def __le__(self, other: AddErrorCost) -> bool: + return self.value <= other.value + + def __gt__(self, other: AddErrorCost) -> bool: + return self.value > other.value + + def __ge__(self, other: AddErrorCost) -> bool: + return self.value >= other.value + + model = DagCostModel(lambda egraph, expr: AddErrorCost(1), AddErrorCost(0)) + with pytest.raises(ArithmeticError, match="addition failed"): + EGraph().extract(ff(1, 2), cost_model=model) + + def test_dag_comparison_error_bubbles(self): + @dataclass(frozen=True) + class CompareErrorCost: + value: int + + def __add__(self, other: CompareErrorCost) -> CompareErrorCost: + return CompareErrorCost(self.value + other.value) + + def __eq__(self, other: object) -> bool: + return False + + def __lt__(self, other: CompareErrorCost) -> bool: + del other + msg = "comparison failed" + raise RuntimeError(msg) + + def __le__(self, other: CompareErrorCost) -> bool: + del other + msg = "comparison failed" + raise RuntimeError(msg) + + def __gt__(self, other: CompareErrorCost) -> bool: + del other + msg = "comparison failed" + raise RuntimeError(msg) + + def __ge__(self, other: CompareErrorCost) -> bool: + del other + msg = "comparison failed" + raise RuntimeError(msg) + + model = DagCostModel(lambda egraph, expr: CompareErrorCost(1), CompareErrorCost(0)) + egraph = EGraph() + egraph.register(union(ff(1, 2)).with_(gg())) + with pytest.raises(RuntimeError, match="comparison failed"): + egraph.extract(ff(1, 2), cost_model=model, extractor="greedy-dag") + def test_class_module(): class A(Expr): diff --git a/python/tests/test_pretty.py b/python/tests/test_pretty.py index 12cdf6f5..daed9961 100644 --- a/python/tests/test_pretty.py +++ b/python/tests/test_pretty.py @@ -276,6 +276,14 @@ def test_rule_eval_mode_pretty_round_trip(eval_mode: RuleEvalMode, option: str) assert eval(rendered, globals()).decl == original.decl +def test_rule_no_decomp_pretty_round_trip() -> None: + original = rule(rel(g()), name="no decomp rule", no_decomp=True).then(rel(h())) + rendered = 'rule(rel(g()), name="no decomp rule", no_decomp=True).then(rel(h()))' + + assert str(original) == rendered + assert eval(rendered, globals()).decl == original.decl + + FREEZE_PARAMS = [ pytest.param((A(),), "EGraph(A()).freeze()", id="freeze add"), pytest.param((b,), "EGraph(b).freeze()", id="freeze constant"), diff --git a/src/egraph.rs b/src/egraph.rs index fc3548b7..6fdfdc6f 100644 --- a/src/egraph.rs +++ b/src/egraph.rs @@ -74,11 +74,19 @@ impl EGraph { #[pymethods] impl EGraph { #[new] - #[pyo3(signature = (*, fact_directory=None, seminaive=true, record=false))] - fn new(fact_directory: Option, seminaive: bool, record: bool) -> Self { + #[pyo3(signature = (*, fact_directory=None, seminaive=true, record=false, num_threads=1, no_decomp=false))] + fn new( + fact_directory: Option, + seminaive: bool, + record: bool, + num_threads: usize, + no_decomp: bool, + ) -> Self { let mut egraph = egglog_experimental::new_experimental_egraph(); egraph.fact_directory = fact_directory; egraph.seminaive = seminaive; + egraph.set_num_threads(num_threads); + egraph.no_decomp = no_decomp; add_base_sort(&mut egraph, PyObjectSort {}, span!()).unwrap(); Self { egraph, @@ -137,6 +145,26 @@ impl EGraph { self.cmds.clone() } + /// Return the number of worker threads configured for this EGraph. + fn num_threads(&self) -> usize { + self.egraph.num_threads() + } + + /// Set the number of worker threads used by this EGraph. + fn set_num_threads(&mut self, num_threads: usize) { + self.egraph.set_num_threads(num_threads); + } + + /// Return whether rule decomposition is disabled globally for this EGraph. + fn no_decomp(&self) -> bool { + self.egraph.no_decomp + } + + /// Set whether rule decomposition is disabled for subsequently registered rules. + fn set_no_decomp(&mut self, no_decomp: bool) { + self.egraph.no_decomp = no_decomp; + } + /// Serialize the EGraph to a SerializedEGraph object. #[pyo3( signature = (root_eclasses, *, max_functions=None, max_calls_per_function=None, include_temporary_functions=false, traceparent=None, tracestate=None), diff --git a/src/extract.rs b/src/extract.rs index 54306c8d..c8689335 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -1,17 +1,55 @@ -use std::cmp::Ordering; +use std::{ + cmp::Ordering, + panic::{AssertUnwindSafe, catch_unwind, resume_unwind}, + sync::Arc, +}; +use egglog::{ + Term, TermId, + extract::{ + DagCostModel as EggDagCostModel, MonoidCost, TreeCostModel as EggTreeCostModel, + TreeCostModelFromDag, TreeExtractor, + }, +}; use pyo3::{exceptions::PyValueError, prelude::*}; use crate::{egraph::EGraph, egraph::Value, termdag::TermDag, tracing_otel}; -use egglog::TermId; + +/// Private unwind payload used to cross Rust APIs whose cost traits cannot +/// return errors. Only this payload is translated back to Python; unrelated +/// Rust panics continue unwinding normally. +struct PythonCostError(PyErr); + +fn python_or_unwind(result: PyResult) -> T { + result.unwrap_or_else(|err| resume_unwind(Box::new(PythonCostError(err)))) +} + +fn catch_python_cost_error(f: impl FnOnce() -> T) -> PyResult { + match catch_unwind(AssertUnwindSafe(f)) { + Ok(value) => Ok(value), + Err(payload) => match payload.downcast::() { + Ok(error) => Err(error.0), + Err(payload) => resume_unwind(payload), + }, + } +} #[derive(Debug)] -// We have to store the result, since the cost model does not return errors -struct Cost(Py); +struct Cost(Arc>); + +impl Cost { + fn from_py(value: Py) -> Self { + Self(Arc::new(value)) + } + + fn to_py(&self, py: Python<'_>) -> Py { + self.0.as_ref().clone_ref(py) + } +} impl Ord for Cost { fn cmp(&self, other: &Self) -> Ordering { - Python::attach(|py| self.0.bind(py).compare(other.0.bind(py)).unwrap()) + Python::attach(|py| python_or_unwind(self.0.bind(py).compare(other.0.bind(py)))) } } @@ -23,7 +61,7 @@ impl PartialOrd for Cost { impl PartialEq for Cost { fn eq(&self, other: &Self) -> bool { - Python::attach(|py| self.0.bind(py).eq(other.0.bind(py))).unwrap() + self.cmp(other) == Ordering::Equal } } @@ -31,25 +69,23 @@ impl Eq for Cost {} impl Clone for Cost { fn clone(&self) -> Self { - Python::attach(|py| Cost(self.0.clone_ref(py))) + Self(self.0.clone()) } } -impl egglog::extract::Cost for Cost { - fn identity() -> Self { - panic!("Should never be called from Rust directly"); - } - - fn unit() -> Self { - panic!("Should never be called from Rust directly"); - } +#[derive(Debug)] +struct TreeEnodeCost { + head: String, + annotation: Py, +} - fn combine(self, _other: &Self) -> Self { - panic!("Should never be called from Rust directly"); - } +#[derive(Debug)] +struct TreeContainerCost { + sort: String, + value: Value, } -/// Cost model defined by Python functions. +/// Tree cost model defined by Python functions. #[derive(Debug)] #[pyclass( frozen, @@ -57,12 +93,12 @@ impl egglog::extract::Cost for Cost { )] pub struct CostModel { /// Function mapping from a term's head and its children's costs to the term's total cost. - /// (head: str, head_cost: COST, children_costs: list[COST]) -> COST + /// (head: str, head_cost: ENODE_COST, children_costs: list[COST]) -> COST fold: Py, - /// Function mapping from an expression node to its cost. - /// (func_name: str, args: list[Value]) -> COST + /// Function mapping from an expression node to an annotation consumed by `fold`. + /// (func_name: str, args: list[Value]) -> ENODE_COST enode_cost: Py, - /// Function mapping from a container value to its cost given the costs of its elements. + /// Function mapping from a container value and element costs to its total cost. /// (sort_name: str, value: Value, element_costs: list[COST]) -> COST container_cost: Py, /// Function mapping from a base value to its cost. @@ -79,7 +115,7 @@ impl CostModel { container_cost: Py, base_value_cost: Py, ) -> Self { - CostModel { + Self { fold, enode_cost, container_cost, @@ -90,7 +126,7 @@ impl CostModel { impl Clone for CostModel { fn clone(&self) -> Self { - Python::attach(|py| CostModel { + Python::attach(|py| Self { fold: self.fold.clone_ref(py), enode_cost: self.enode_cost.clone_ref(py), container_cost: self.container_cost.clone_ref(py), @@ -99,19 +135,21 @@ impl Clone for CostModel { } } -impl egglog::extract::CostModel for CostModel { - fn fold(&self, head: &str, children_cost: &[Cost], head_cost: Cost) -> Cost { - Cost(Python::attach(|py| { - let head_cost = head_cost.0.clone_ref(py); +impl EggTreeCostModel for CostModel { + type EnodeCost = TreeEnodeCost; + type ContainerCost = TreeContainerCost; + + fn fold_enode_cost(&self, enode_cost: Self::EnodeCost, children_cost: &[Cost]) -> Cost { + Python::attach(|py| { let children_cost = children_cost - .into_iter() - .cloned() - .map(|c| c.0.clone_ref(py)) + .iter() + .map(|cost| cost.to_py(py)) .collect::>(); - self.fold - .call1(py, (head, head_cost, children_cost)) - .unwrap() - })) + Cost::from_py(python_or_unwind( + self.fold + .call1(py, (enode_cost.head, enode_cost.annotation, children_cost)), + )) + }) } fn enode_cost( @@ -119,10 +157,34 @@ impl egglog::extract::CostModel for CostModel { _egraph: &egglog::EGraph, func: &egglog::Function, enode: &egglog::Enode<'_>, + ) -> Self::EnodeCost { + Python::attach(|py| { + let values = enode + .children + .iter() + .map(|value| Value(*value)) + .collect::>(); + TreeEnodeCost { + head: func.name().to_owned(), + annotation: python_or_unwind(self.enode_cost.call1(py, (func.name(), values))), + } + }) + } + + fn fold_container_cost( + &self, + container_cost: Self::ContainerCost, + element_costs: &[Cost], ) -> Cost { Python::attach(|py| { - let values = enode.children.iter().map(|v| Value(*v)).collect::>(); - Cost(self.enode_cost.call1(py, (func.name(), values)).unwrap()) + let element_costs = element_costs + .iter() + .map(|cost| cost.to_py(py)) + .collect::>(); + Cost::from_py(python_or_unwind(self.container_cost.call1( + py, + (container_cost.sort, container_cost.value, element_costs), + ))) }) } @@ -131,18 +193,11 @@ impl egglog::extract::CostModel for CostModel { _egraph: &egglog::EGraph, sort: &egglog::ArcSort, value: egglog::Value, - element_costs: &[Cost], - ) -> Cost { - Cost(Python::attach(|py| { - let element_costs = element_costs - .into_iter() - .cloned() - .map(|c| c.0.clone_ref(py)) - .collect::>(); - self.container_cost - .call1(py, (sort.name(), Value(value), element_costs)) - .unwrap() - })) + ) -> Self::ContainerCost { + TreeContainerCost { + sort: sort.name().to_owned(), + value: Value(value), + } } // https://github.com/PyO3/pyo3/issues/1190 @@ -153,34 +208,312 @@ impl egglog::extract::CostModel for CostModel { value: egglog::Value, ) -> Cost { Python::attach(|py| { - Cost( - self.base_value_cost - .call1(py, (sort.name(), Value(value))) - .unwrap(), + Cost::from_py(python_or_unwind( + self.base_value_cost.call1(py, (sort.name(), Value(value))), + )) + }) + } +} + +#[derive(Debug)] +struct DagCostContext { + identity: Arc>, +} + +#[derive(Clone, Debug)] +enum DagCost { + Identity, + Value { + value: Arc>, + context: Arc, + }, +} + +impl DagCost { + fn value(value: Py, context: Arc) -> Self { + Self::Value { + value: Arc::new(value), + context, + } + } + + fn to_py(&self, py: Python<'_>, context: &Arc) -> Py { + match self { + Self::Identity => context.identity.as_ref().clone_ref(py), + Self::Value { value, .. } => value.as_ref().clone_ref(py), + } + } + + fn compare_values(left: &Arc>, right: &Arc>) -> Ordering { + Python::attach(|py| python_or_unwind(left.bind(py).compare(right.bind(py)))) + } + + fn ensure_same_context(left: &Arc, right: &Arc) { + if !Arc::ptr_eq(left, right) { + python_or_unwind::<()>(Err(PyValueError::new_err( + "cannot combine costs from different extraction contexts", + ))); + } + } +} + +impl Ord for DagCost { + fn cmp(&self, other: &Self) -> Ordering { + match (self, other) { + (Self::Identity, Self::Identity) => Ordering::Equal, + ( + Self::Value { + value: left, + context: left_context, + }, + Self::Value { + value: right, + context: right_context, + }, + ) => { + Self::ensure_same_context(left_context, right_context); + Self::compare_values(left, right) + } + (Self::Identity, Self::Value { value, context }) => { + Self::compare_values(&context.identity, value) + } + (Self::Value { value, context }, Self::Identity) => { + Self::compare_values(value, &context.identity) + } + } + } +} + +impl PartialOrd for DagCost { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl PartialEq for DagCost { + fn eq(&self, other: &Self) -> bool { + self.cmp(other) == Ordering::Equal + } +} + +impl Eq for DagCost {} + +impl MonoidCost for DagCost { + fn identity() -> Self { + Self::Identity + } + + fn combine(self, other: &Self) -> Self { + match (self, other) { + (Self::Identity, Self::Identity) => Self::Identity, + (Self::Identity, value @ Self::Value { .. }) => value.clone(), + (value @ Self::Value { .. }, Self::Identity) => value, + ( + Self::Value { + value: left, + context: left_context, + }, + Self::Value { + value: right, + context: right_context, + }, + ) => { + Self::ensure_same_context(&left_context, right_context); + let combined = Python::attach(|py| { + python_or_unwind(left.bind(py).add(right.bind(py))).unbind() + }); + Self::value(combined, left_context) + } + } + } +} + +/// Additive marginal DAG cost model defined by Python functions. +#[derive(Debug)] +#[pyclass( + frozen, + str = "DagCostModel({identity:?}, {enode_cost:?}, {container_cost:?}, {base_value_cost:?}" +)] +pub struct DagCostModel { + identity: Py, + enode_cost: Py, + container_cost: Py, + base_value_cost: Py, +} + +#[pymethods] +impl DagCostModel { + #[new] + fn new( + identity: Py, + enode_cost: Py, + container_cost: Py, + base_value_cost: Py, + ) -> Self { + Self { + identity, + enode_cost, + container_cost, + base_value_cost, + } + } +} + +impl DagCostModel { + fn runtime(&self, py: Python<'_>) -> RuntimeDagCostModel { + RuntimeDagCostModel { + context: Arc::new(DagCostContext { + identity: Arc::new(self.identity.clone_ref(py)), + }), + enode_cost: Arc::new(self.enode_cost.clone_ref(py)), + container_cost: Arc::new(self.container_cost.clone_ref(py)), + base_value_cost: Arc::new(self.base_value_cost.clone_ref(py)), + } + } +} + +#[derive(Clone, Debug)] +struct RuntimeDagCostModel { + context: Arc, + enode_cost: Arc>, + container_cost: Arc>, + base_value_cost: Arc>, +} + +impl EggDagCostModel for RuntimeDagCostModel { + fn enode_cost( + &self, + _egraph: &egglog::EGraph, + func: &egglog::Function, + enode: &egglog::Enode<'_>, + ) -> DagCost { + Python::attach(|py| { + let values = enode + .children + .iter() + .map(|value| Value(*value)) + .collect::>(); + DagCost::value( + python_or_unwind(self.enode_cost.call1(py, (func.name(), values))), + self.context.clone(), + ) + }) + } + + fn container_cost( + &self, + _egraph: &egglog::EGraph, + sort: &egglog::ArcSort, + value: egglog::Value, + ) -> DagCost { + Python::attach(|py| { + DagCost::value( + python_or_unwind(self.container_cost.call1(py, (sort.name(), Value(value)))), + self.context.clone(), ) }) } + + fn base_value_cost( + &self, + _egraph: &egglog::EGraph, + sort: &egglog::ArcSort, + value: egglog::Value, + ) -> DagCost { + Python::attach(|py| { + DagCost::value( + python_or_unwind(self.base_value_cost.call1(py, (sort.name(), Value(value)))), + self.context.clone(), + ) + }) + } +} + +#[derive(Copy, Clone)] +enum ExtractionMode { + Tree, + GreedyDag, +} + +fn extraction_mode(name: &str) -> PyResult { + match name { + "tree" => Ok(ExtractionMode::Tree), + "greedy-dag" => Ok(ExtractionMode::GreedyDag), + _ => Err(PyValueError::new_err(format!( + "unknown extractor {name:?}; expected 'tree' or 'greedy-dag'" + ))), + } } -// TODO: Don't progress just return an error if there was an exception? +fn roots_from_names( + egraph: &egglog::EGraph, + roots: Vec<(String, Value)>, +) -> PyResult> { + roots + .into_iter() + .map(|(sort, value)| { + let arcsort = egraph + .get_sort_by_name(&sort) + .cloned() + .ok_or_else(|| PyValueError::new_err(format!("unknown sort {sort:?}")))?; + Ok((arcsort, value.0)) + }) + .collect() +} +fn rootsorts_from_names( + egraph: &egglog::EGraph, + rootsorts: Option<&[String]>, +) -> PyResult>> { + rootsorts + .map(|rootsorts| { + rootsorts + .iter() + .map(|sort| { + egraph + .get_sort_by_name(sort) + .cloned() + .ok_or_else(|| PyValueError::new_err(format!("unknown sort {sort:?}"))) + }) + .collect() + }) + .transpose() +} + +fn copy_termdag(source: &egglog::TermDag, target: &mut egglog::TermDag) -> Vec { + let mut copied = Vec::with_capacity(source.size()); + for term in 0..source.size() { + let target_term = match source.get(term).clone() { + Term::Lit(literal) => target.lit(literal), + Term::Var(variable) => target.var(variable), + Term::App(head, children) => { + let children = children.into_iter().map(|child| copied[child]).collect(); + target.app(head, children) + } + }; + copied.push(target_term); + } + copied +} + +/// Compatibility facade for the former owned core extractor. The current core +/// extractor borrows its e-graph, so this object stores preparation inputs and +/// prepares locally for each extraction call. #[pyclass(unsendable)] -pub struct Extractor(egglog::extract::Extractor); +pub struct Extractor { + rootsorts: Option>, + cost_model: CostModel, +} #[pymethods] impl Extractor { - /// Create a new extractor from the given egraph and cost model. + /// Create a new extractor facade from the given egraph and cost model. /// - /// Bulk of the computation happens at initialization time. - /// The later extractions only reuses saved results. - /// This means a new extractor must be created if the egraph changes. - /// Holding a reference to the egraph would enforce this but prevents the extractor being reused. - /// - /// For convenience, if the rootsorts is `None`, it defaults to extract all extractable rootsorts. + /// For convenience, if the rootsorts is `None`, it defaults to all + /// extractable rootsorts. #[new] #[pyo3(signature = (rootsorts, egraph, cost_model, *, traceparent=None, tracestate=None))] fn new( - py: Python<'_>, rootsorts: Option>, egraph: &EGraph, cost_model: CostModel, @@ -194,28 +527,20 @@ impl Extractor { has_rootsorts = rootsorts.is_some() ); let _entered = span.enter(); - let egraph = &egraph.egraph; - // Transforms sorts to arcsorts, returning an error if any are unknown - let rootsorts = rootsorts - .map(|rs| { - rs.into_iter() - .map(|s| egraph.get_sort_by_name(&s).cloned()) - .collect::>>() - .ok_or(PyValueError::new_err("Unknown sort in rootsorts")) - }) - .map_or(Ok(None), |r| r.map(Some))?; - let extractor = - egglog::extract::Extractor::compute_costs_from_rootsorts(rootsorts, egraph, cost_model); - if let Some(err) = PyErr::take(py) { - return Err(err); - }; - Ok(Extractor(extractor)) + if let Some(rootsorts) = &rootsorts { + for sort in rootsorts { + if egraph.egraph.get_sort_by_name(sort).is_none() { + return Err(PyValueError::new_err(format!("unknown sort {sort:?}"))); + } + } + } + Ok(Self { + rootsorts, + cost_model, + }) } /// Extract the best term of a value from a given sort. - /// - /// This function expects the sort to be already computed, - /// which can be one of the rootsorts, or reachable from rootsorts, or primitives, or containers of computed sorts. #[pyo3(signature = (egraph, termdag, value, sort, *, traceparent=None, tracestate=None))] fn extract_best( &self, @@ -234,18 +559,27 @@ impl Extractor { let sort = egraph .egraph .get_sort_by_name(&sort) - .ok_or(PyValueError::new_err("Unknown sort"))?; - let (cost, term) = self - .0 - .extract_best_with_sort(&egraph.egraph, &mut termdag.0, value.0, sort.clone()) - .ok_or(PyValueError::new_err("Unextractable root".to_string()))?; - Ok((cost.0.clone_ref(py), term)) + .cloned() + .ok_or_else(|| PyValueError::new_err("unknown sort"))?; + let rootsorts = rootsorts_from_names(&egraph.egraph, self.rootsorts.as_deref())?; + let extracted = catch_python_cost_error(|| { + let extractor = TreeExtractor::compute_costs_from_rootsorts( + rootsorts, + &egraph.egraph, + self.cost_model.clone(), + ); + let mut local_termdag = egglog::TermDag::default(); + extractor + .extract_best_with_sort(&mut local_termdag, value.0, sort) + .map(|extracted| (local_termdag, extracted.cost, extracted.term)) + })? + .ok_or_else(|| PyValueError::new_err("unextractable root"))?; + let (local_termdag, cost, term) = extracted; + let term = copy_termdag(&local_termdag, &mut termdag.0)[term]; + Ok((cost.to_py(py), term)) } /// Extract variants of an e-class. - /// - /// The variants are selected by first picking `nvariants` e-nodes with the lowest cost from the e-class - /// and then extracting a term from each e-node. #[pyo3(signature = (egraph, termdag, value, nvariants, sort, *, traceparent=None, tracestate=None))] fn extract_variants( &self, @@ -265,17 +599,65 @@ impl Extractor { let sort = egraph .egraph .get_sort_by_name(&sort) - .ok_or(PyValueError::new_err("Unknown sort"))?; - let variants = self.0.extract_variants_with_sort( - &egraph.egraph, - &mut termdag.0, - value.0, - nvariants, - sort.clone(), - ); + .cloned() + .ok_or_else(|| PyValueError::new_err("unknown sort"))?; + let rootsorts = rootsorts_from_names(&egraph.egraph, self.rootsorts.as_deref())?; + let (local_termdag, variants) = catch_python_cost_error(|| { + let extractor = TreeExtractor::compute_costs_from_rootsorts( + rootsorts, + &egraph.egraph, + self.cost_model.clone(), + ); + let mut local_termdag = egglog::TermDag::default(); + let variants = + extractor.extract_variants_with_sort(&mut local_termdag, value.0, nvariants, sort); + (local_termdag, variants) + })?; + let copied = copy_termdag(&local_termdag, &mut termdag.0); Ok(variants .into_iter() - .map(|(cost, term)| (cost.0.clone_ref(py), term)) + .map(|variant| (variant.cost.to_py(py), copied[variant.term])) .collect()) } } + +/// Extract the best term for each root with a custom additive marginal model. +#[pyfunction] +#[pyo3(signature = (egraph, roots, cost_model, *, extractor="tree", traceparent=None, tracestate=None))] +pub fn extract_best_with_dag_cost_model( + py: Python<'_>, + egraph: &EGraph, + roots: Vec<(String, Value)>, + cost_model: &DagCostModel, + extractor: &str, + traceparent: Option, + tracestate: Option, +) -> PyResult<(TermDag, Vec, TermId)>>)> { + let _context_guard = + tracing_otel::attach_parent_context(traceparent.as_deref(), tracestate.as_deref()); + let span = tracing::info_span!( + "bindings.extract_best_with_dag_cost_model", + root_count = roots.len(), + extractor + ); + let _entered = span.enter(); + let roots = roots_from_names(&egraph.egraph, roots)?; + let mode = extraction_mode(extractor)?; + let runtime = cost_model.runtime(py); + let context = runtime.context.clone(); + let extracted = catch_python_cost_error(|| match mode { + ExtractionMode::Tree => egraph + .egraph + .extract_best_with_cost_model(roots, TreeCostModelFromDag(runtime)), + ExtractionMode::GreedyDag => { + egglog_experimental::extract_best_greedy_dag(&egraph.egraph, roots, runtime) + } + })? + .map_err(crate::error::WrappedError::Egglog)?; + let terms = extracted + .terms + .into_iter() + .map(|term| term.map(|term| (term.cost.to_py(py, &context), term.term))) + .collect(); + Ok((TermDag(extracted.termdag), terms)) +} diff --git a/src/lib.rs b/src/lib.rs index 36bbb772..f4534993 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -30,16 +30,6 @@ fn shutdown_tracing(py: Python<'_>) -> PyResult<()> { /// Bindings for egglog rust library #[pymodule] fn bindings(m: &Bound<'_, PyModule>) -> PyResult<()> { - // Configure Rayon thread pool from env var, defaulting to 1 if unset/invalid. - let num_threads = std::env::var("RAYON_NUM_THREADS") - .ok() - .and_then(|s| s.parse::().ok()) - .unwrap_or(1); - rayon::ThreadPoolBuilder::new() - .num_threads(num_threads) - .build_global() - .unwrap(); - pyo3_log::init(); m.add_class::()?; @@ -51,12 +41,17 @@ fn bindings(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_function(wrap_pyfunction!(setup_tracing, m)?)?; m.add_function(wrap_pyfunction!(shutdown_tracing, m)?)?; + m.add_function(wrap_pyfunction!( + crate::extract::extract_best_with_dag_cost_model, + m + )?)?; crate::conversions::add_structs_to_module(m)?; crate::conversions::add_enums_to_module(m)?; From ddae949183c16555cb3a448ea33411f223bf85e7 Mon Sep 17 00:00:00 2001 From: Saul Shanabrook Date: Wed, 2 Sep 2026 12:58:14 -0700 Subject: [PATCH 4/6] Expose structured multi-extract outputs --- Cargo.lock | 20 ++++++++++---------- Cargo.toml | 25 +++++++++++++------------ docs/changelog.md | 4 +++- docs/reference/python-integration.md | 8 ++++++++ python/egglog/bindings.pyi | 9 ++++++++- python/egglog/egraph.py | 20 ++++++++++++-------- python/tests/test_bindings.py | 16 ++++++++++++++++ python/tests/test_high_level.py | 5 +++++ src/conversions.rs | 26 ++++++++++++++++++++++++++ src/lib.rs | 1 + 10 files changed, 102 insertions(+), 32 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1ad18c95..ce9509cd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -382,7 +382,7 @@ checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" [[package]] name = "egglog" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=69177f4f37db9bc6c23899375988dc8b8ecedfc3#69177f4f37db9bc6c23899375988dc8b8ecedfc3" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=ee58d8a537cec77e07595b1f0b577eee617f34b6#ee58d8a537cec77e07595b1f0b577eee617f34b6" dependencies = [ "csv", "dyn-clone", @@ -410,7 +410,7 @@ dependencies = [ [[package]] name = "egglog-add-primitive" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=69177f4f37db9bc6c23899375988dc8b8ecedfc3#69177f4f37db9bc6c23899375988dc8b8ecedfc3" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=ee58d8a537cec77e07595b1f0b577eee617f34b6#ee58d8a537cec77e07595b1f0b577eee617f34b6" dependencies = [ "quote", "syn 2.0.117", @@ -419,7 +419,7 @@ dependencies = [ [[package]] name = "egglog-ast" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=69177f4f37db9bc6c23899375988dc8b8ecedfc3#69177f4f37db9bc6c23899375988dc8b8ecedfc3" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=ee58d8a537cec77e07595b1f0b577eee617f34b6#ee58d8a537cec77e07595b1f0b577eee617f34b6" dependencies = [ "ordered-float", ] @@ -427,7 +427,7 @@ dependencies = [ [[package]] name = "egglog-bridge" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=69177f4f37db9bc6c23899375988dc8b8ecedfc3#69177f4f37db9bc6c23899375988dc8b8ecedfc3" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=ee58d8a537cec77e07595b1f0b577eee617f34b6#ee58d8a537cec77e07595b1f0b577eee617f34b6" dependencies = [ "anyhow", "dyn-clone", @@ -450,7 +450,7 @@ dependencies = [ [[package]] name = "egglog-concurrency" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=69177f4f37db9bc6c23899375988dc8b8ecedfc3#69177f4f37db9bc6c23899375988dc8b8ecedfc3" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=ee58d8a537cec77e07595b1f0b577eee617f34b6#ee58d8a537cec77e07595b1f0b577eee617f34b6" dependencies = [ "arc-swap", "bumpalo", @@ -462,7 +462,7 @@ dependencies = [ [[package]] name = "egglog-core-relations" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=69177f4f37db9bc6c23899375988dc8b8ecedfc3#69177f4f37db9bc6c23899375988dc8b8ecedfc3" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=ee58d8a537cec77e07595b1f0b577eee617f34b6#ee58d8a537cec77e07595b1f0b577eee617f34b6" dependencies = [ "anyhow", "bumpalo", @@ -490,7 +490,7 @@ dependencies = [ [[package]] name = "egglog-experimental" version = "3.0.0" -source = "git+https://github.com/egraphs-good/egglog-experimental.git?rev=6bfb589d1c4695232fc24e9bdf097c49af451823#6bfb589d1c4695232fc24e9bdf097c49af451823" +source = "git+https://github.com/egraphs-good/egglog-experimental.git?rev=2b4627a5806f8476bc34814ecf817e3b77c16f87#2b4627a5806f8476bc34814ecf817e3b77c16f87" dependencies = [ "egglog", "egglog-ast", @@ -506,12 +506,12 @@ dependencies = [ [[package]] name = "egglog-numeric-id" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=69177f4f37db9bc6c23899375988dc8b8ecedfc3#69177f4f37db9bc6c23899375988dc8b8ecedfc3" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=ee58d8a537cec77e07595b1f0b577eee617f34b6#ee58d8a537cec77e07595b1f0b577eee617f34b6" [[package]] name = "egglog-reports" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=69177f4f37db9bc6c23899375988dc8b8ecedfc3#69177f4f37db9bc6c23899375988dc8b8ecedfc3" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=ee58d8a537cec77e07595b1f0b577eee617f34b6#ee58d8a537cec77e07595b1f0b577eee617f34b6" dependencies = [ "clap", "hashbrown 0.16.1", @@ -525,7 +525,7 @@ dependencies = [ [[package]] name = "egglog-union-find" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=69177f4f37db9bc6c23899375988dc8b8ecedfc3#69177f4f37db9bc6c23899375988dc8b8ecedfc3" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=ee58d8a537cec77e07595b1f0b577eee617f34b6#ee58d8a537cec77e07595b1f0b577eee617f34b6" dependencies = [ "crossbeam", "egglog-concurrency", diff --git a/Cargo.toml b/Cargo.toml index 80cb6386..6f65e460 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,13 +18,14 @@ opentelemetry = "0.28" opentelemetry-otlp = { version = "0.28", features = ["http-proto", "reqwest-blocking-client", "trace"] } opentelemetry-stdout = { version = "0.28", features = ["trace"] } opentelemetry_sdk = "0.28" -# Egglog main at integration plus the rule-name round-trip fix used by this branch. -egglog = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "69177f4f37db9bc6c23899375988dc8b8ecedfc3", default-features = false } -egglog-ast = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "69177f4f37db9bc6c23899375988dc8b8ecedfc3" } -egglog-core-relations = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "69177f4f37db9bc6c23899375988dc8b8ecedfc3" } -egglog-reports = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "69177f4f37db9bc6c23899375988dc8b8ecedfc3" } -egglog-bridge = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "69177f4f37db9bc6c23899375988dc8b8ecedfc3" } -egglog-experimental = { git = "https://github.com/egraphs-good/egglog-experimental.git", rev = "6bfb589d1c4695232fc24e9bdf097c49af451823", default-features = false } +# Egglog main at integration plus the rule-name round-trip and user-output +# downcasting fixes used by this branch. +egglog = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "ee58d8a537cec77e07595b1f0b577eee617f34b6", default-features = false } +egglog-ast = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "ee58d8a537cec77e07595b1f0b577eee617f34b6" } +egglog-core-relations = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "ee58d8a537cec77e07595b1f0b577eee617f34b6" } +egglog-reports = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "ee58d8a537cec77e07595b1f0b577eee617f34b6" } +egglog-bridge = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "ee58d8a537cec77e07595b1f0b577eee617f34b6" } +egglog-experimental = { git = "https://github.com/egraphs-good/egglog-experimental.git", rev = "2b4627a5806f8476bc34814ecf817e3b77c16f87", default-features = false } egraph-serialize = { version = "0.3", features = ["serde", "graphviz"] } serde_json = "1" pyo3-log = "*" @@ -42,8 +43,8 @@ base64 = "0.22.1" debug = true [patch."https://github.com/egraphs-good/egglog.git"] -egglog = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "69177f4f37db9bc6c23899375988dc8b8ecedfc3" } -egglog-ast = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "69177f4f37db9bc6c23899375988dc8b8ecedfc3" } -egglog-core-relations = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "69177f4f37db9bc6c23899375988dc8b8ecedfc3" } -egglog-reports = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "69177f4f37db9bc6c23899375988dc8b8ecedfc3" } -egglog-bridge = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "69177f4f37db9bc6c23899375988dc8b8ecedfc3" } +egglog = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "ee58d8a537cec77e07595b1f0b577eee617f34b6" } +egglog-ast = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "ee58d8a537cec77e07595b1f0b577eee617f34b6" } +egglog-core-relations = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "ee58d8a537cec77e07595b1f0b577eee617f34b6" } +egglog-reports = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "ee58d8a537cec77e07595b1f0b577eee617f34b6" } +egglog-bridge = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "ee58d8a537cec77e07595b1f0b577eee617f34b6" } diff --git a/docs/changelog.md b/docs/changelog.md index 6212857e..0211726d 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -37,7 +37,9 @@ _This project uses semantic versioning_ `EGraph` or call `set_num_threads(...)` instead. - Add `tree` and true `greedy-dag` extraction modes to `extract` and `extract_multiple`; add destructive `keep_best`; allow multi-root variant - extraction while preserving input order. + extraction while preserving input order. Low-level bindings can recover the + structured aggregate returned by experimental `multi-extract` with + `UserDefinedCommandOutput.as_multi_extract()`. - Use the experimental dynamic cost model consistently across the default extraction paths, including canonical `set_cost` tables. A compatible raw cost table already occupying the canonical name is reused, while an diff --git a/docs/reference/python-integration.md b/docs/reference/python-integration.md index f9eda07a..831075e9 100644 --- a/docs/reference/python-integration.md +++ b/docs/reference/python-integration.md @@ -844,6 +844,14 @@ independently; sharing between separate roots does not reduce either cost. This API always uses dynamic costs; custom Python cost models are supported only by single-root `extract`. +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. + ### Custom Tree Cost Models A `TreeCostModel` is a callable that receives the e-graph, one expression diff --git a/python/egglog/bindings.pyi b/python/egglog/bindings.pyi index d63dabd5..31104d10 100644 --- a/python/egglog/bindings.pyi +++ b/python/egglog/bindings.pyi @@ -47,6 +47,7 @@ __all__ = [ "IterationReport", "Let", "Lit", + "MultiExtractOutput", "Naive", "NewSort", "Output", @@ -462,7 +463,13 @@ class IdentSort: def __new__(cls, ident: str, sort: str) -> IdentSort: ... @final -class UserDefinedCommandOutput: ... +class MultiExtractOutput: + termdag: TermDag + terms: list[list[_TermId]] + +@final +class UserDefinedCommandOutput: + def as_multi_extract(self) -> MultiExtractOutput | None: ... @final class SingleScan: diff --git a/python/egglog/egraph.py b/python/egglog/egraph.py index 88e4c5cc..2baaa115 100644 --- a/python/egglog/egraph.py +++ b/python/egglog/egraph.py @@ -1489,17 +1489,21 @@ def extract_multiple( except BaseException as e: e.add_note("while extracting: " + ", ".join(map(str, runtime_exprs))) raise - if len(outputs) != len(runtime_exprs) or not all( - isinstance(output, bindings.ExtractVariants) for output in outputs - ): + if len(outputs) != 1 or not isinstance(outputs[0], bindings.UserDefinedOutput): msg = "multi-extract returned unexpected command outputs" raise RuntimeError(msg) + output = outputs[0].output.as_multi_extract() + if output is None: + msg = "multi-extract returned an unexpected user-defined output" + raise RuntimeError(msg) + termdag = output.termdag + terms_by_root = output.terms + if len(terms_by_root) != len(runtime_exprs): + msg = "multi-extract returned an unexpected user-defined output" + raise RuntimeError(msg) results: list[list[BASE_EXPR]] = [] - for runtime_expr, output in zip(runtime_exprs, outputs, strict=True): - assert isinstance(output, bindings.ExtractVariants) - typed_exprs = self._state.exprs_from_egg( - output.termdag, output.terms, runtime_expr.__egg_typed_expr__.tp - ) + for runtime_expr, terms in zip(runtime_exprs, terms_by_root, strict=True): + typed_exprs = self._state.exprs_from_egg(termdag, terms, runtime_expr.__egg_typed_expr__.tp) results.append([ cast("BASE_EXPR", RuntimeExpr.__from_values__(self.__egg_decls__, typed_expr)) for typed_expr in typed_exprs diff --git a/python/tests/test_bindings.py b/python/tests/test_bindings.py index f7fdd95e..725eafb6 100644 --- a/python/tests/test_bindings.py +++ b/python/tests/test_bindings.py @@ -117,6 +117,22 @@ def test_parse_and_run_program(self): assert egraph.run_program(*egraph.parse_program(program)) == [] + def test_downcast_multi_extract_output(self): + egraph = EGraph() + (output,) = egraph.parse_and_run_program("(datatype Expr (Num i64)) (multi-extract 1 (Num 1) 2)") + + assert isinstance(output, UserDefinedOutput) + multi_extract = output.output.as_multi_extract() + assert isinstance(multi_extract, MultiExtractOutput) + assert [[multi_extract.termdag.to_string(term) for term in terms] for terms in multi_extract.terms] == [ + ["(Num 1)"], + ["2"], + ] + + (other_output,) = egraph.parse_and_run_program("(print-table-stats Num)") + assert isinstance(other_output, UserDefinedOutput) + assert other_output.output.as_multi_extract() is None + def test_parse_program_preserves_uf_extraction_behavior(self): program = (EGG_SMOL_FOLDER / "tests" / "uf-extraction.egg").read_text() diff --git a/python/tests/test_high_level.py b/python/tests/test_high_level.py index c8eb6e55..d2616465 100644 --- a/python/tests/test_high_level.py +++ b/python/tests/test_high_level.py @@ -2585,6 +2585,11 @@ def opaque(self) -> MultiRoot: ... extracted = egraph.extract_multiple([String("first"), opaque, MultiRoot(2)], 1, extractor=extractor) assert extracted == [[String("first")], [], [repeated]] + assert egraph.extract_multiple([String("first"), opaque, MultiRoot(2)], 2, extractor=extractor) == [ + [String("first")], + [], + [repeated, MultiRoot(2)], + ] assert egraph.extract_multiple(i64(4), 1, extractor=extractor) == [i64(4)] homogeneous: list[list[MultiRoot]] = egraph.extract_multiple([MultiRoot(2)], 1, extractor=extractor) assert homogeneous == [[repeated]] diff --git a/src/conversions.rs b/src/conversions.rs index 73fb9d42..b5a5fab5 100644 --- a/src/conversions.rs +++ b/src/conversions.rs @@ -1033,8 +1033,34 @@ mod duration_tests { #[derive(Clone)] pub struct UserDefinedCommandOutput(Arc); +#[pyclass(eq, frozen, get_all)] +#[derive(Clone, PartialEq, Eq)] +pub struct MultiExtractOutput { + termdag: TermDag, + terms: Vec>, +} + +#[pymethods] +impl MultiExtractOutput { + fn __repr__(slf: PyRef<'_, Self>, py: Python) -> PyResult { + data_repr(py, slf, vec!["termdag", "terms"]) + } +} + #[pymethods] impl UserDefinedCommandOutput { + /// Return this output as a structured experimental multi-extraction, if it is one. + fn as_multi_extract(&self) -> Option { + self.0 + .as_ref() + .as_any() + .downcast_ref::() + .map(|output| MultiExtractOutput { + termdag: TermDag(output.termdag.clone()), + terms: output.terms.clone(), + }) + } + fn __str__(&self) -> String { format!("{}", self.0) } diff --git a/src/lib.rs b/src/lib.rs index f4534993..ac8bb92f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -38,6 +38,7 @@ fn bindings(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; From 1bb25ea9c516c760979781ea19a211f075912b5e Mon Sep 17 00:00:00 2001 From: Saul Shanabrook Date: Thu, 3 Sep 2026 16:05:22 -0700 Subject: [PATCH 5/6] Finish Egglog 3 integration and Param-Eq handoff --- AGENTS.md | 6 +- Cargo.lock | 20 +- Cargo.toml | 29 +- docs/changelog.md | 106 +- docs/explanation/2026_02_containers.md | 21 +- docs/reference/bindings.md | 10 + docs/reference/egglog-translation.md | 65 +- docs/reference/python-integration.md | 132 +- docs/reference/usage.md | 11 +- experiments/param_eq/NOTES.md | 38 +- experiments/param_eq/README.md | 34 +- .../param_eq/results/paper-replication.csv | 1 - .../results/representation-comparison.csv | 1 - python/egglog/bindings.pyi | 8 +- python/egglog/builtins.py | 5 - python/egglog/declarations.py | 7 +- python/egglog/egraph.py | 563 ++++--- python/egglog/egraph_state.py | 828 ++++++--- python/egglog/exp/param_eq/domain.py | 21 +- python/egglog/exp/param_eq/pipeline.py | 94 +- python/egglog/type_constraint_solver.py | 9 +- python/tests/param_eq/test_domain.py | 30 + python/tests/param_eq/test_pipeline.py | 105 +- .../tests/param_eq/test_research_harness.py | 2 + python/tests/test_bindings.py | 49 +- python/tests/test_egraph_state.py | 379 +++++ python/tests/test_high_level.py | 1481 ++++++++++++++--- python/tests/test_py_object_sort.py | 40 + src/conversions.rs | 12 +- src/egraph.rs | 178 +- src/error.rs | 28 +- src/extract.rs | 72 +- src/py_object_sort.rs | 110 +- test-data/unit/check-high-level.test | 5 + 34 files changed, 3334 insertions(+), 1166 deletions(-) delete mode 100644 experiments/param_eq/results/paper-replication.csv delete mode 100644 experiments/param_eq/results/representation-comparison.csv create mode 100644 python/tests/test_egraph_state.py diff --git a/AGENTS.md b/AGENTS.md index 1f3775c8..d4d1a314 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -19,6 +19,7 @@ - Use the Context7 MCP server for egglog documentation instead of copying external doc summaries into this file. - Keep general workflows in the how-to guides, and keep Python-specific runtime/reference examples in `docs/reference/python-integration.md`. +- Keep high-level docs focused on current public APIs and observable behavior. Put low-level binding details in `docs/reference/bindings.md`, and migration-only details in `docs/changelog.md`. - If a PR adds or updates a changelog entry in `docs/changelog.md`, keep it aligned with the final code changes. - For a clean docs rebuild, clear `docs/_build/`; the MyST-NB execution cache lives in `docs/_build/.jupyter_cache`. @@ -28,7 +29,9 @@ - When changing public high-level APIs, update the public docs, stubs, and pretty/freeze round-trip expectations together. - Keep `builtins.py` limited to operations implemented by Egglog primitives. Compose derived expressions at their domain call sites, without private runtime plumbing in public APIs. - Pretty-print declaration structure directly; do not evaluate runtime expressions merely to canonicalize their output. -- Higher-order callable type probing should stay isolated from the live ruleset: copy declarations and run with no current ruleset so inference does not register temporary unnamed functions or rewrites. +- Higher-order callable type probing should stay isolated from the live ruleset: + copy declarations and run with no current ruleset so inference does not + register temporary unnamed declarations or eager bodies. ## Array API @@ -44,6 +47,7 @@ - Prefer the minimal code change and the minimal diff that solves the task; only broaden the change if the smaller fix is not sufficient. - High-level tests should assert observable behavior through public APIs. Avoid private `_` APIs and exact generated names unless serialized output is itself the public contract. +- Keep white-box `EGraphState` transcript, cleanup, and fault-injection tests in `python/tests/test_egraph_state.py`; keep `test_high_level.py` focused on public behavior. - Run `make mypy` for typing changes. - Run targeted pytest for touched modules. - Run `make docs` for docs or public API changes. diff --git a/Cargo.lock b/Cargo.lock index ce9509cd..67b7c956 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -382,7 +382,7 @@ checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" [[package]] name = "egglog" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=ee58d8a537cec77e07595b1f0b577eee617f34b6#ee58d8a537cec77e07595b1f0b577eee617f34b6" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=1eea60a14f214505741d22bdd6c9b501b21c04ee#1eea60a14f214505741d22bdd6c9b501b21c04ee" dependencies = [ "csv", "dyn-clone", @@ -410,7 +410,7 @@ dependencies = [ [[package]] name = "egglog-add-primitive" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=ee58d8a537cec77e07595b1f0b577eee617f34b6#ee58d8a537cec77e07595b1f0b577eee617f34b6" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=1eea60a14f214505741d22bdd6c9b501b21c04ee#1eea60a14f214505741d22bdd6c9b501b21c04ee" dependencies = [ "quote", "syn 2.0.117", @@ -419,7 +419,7 @@ dependencies = [ [[package]] name = "egglog-ast" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=ee58d8a537cec77e07595b1f0b577eee617f34b6#ee58d8a537cec77e07595b1f0b577eee617f34b6" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=1eea60a14f214505741d22bdd6c9b501b21c04ee#1eea60a14f214505741d22bdd6c9b501b21c04ee" dependencies = [ "ordered-float", ] @@ -427,7 +427,7 @@ dependencies = [ [[package]] name = "egglog-bridge" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=ee58d8a537cec77e07595b1f0b577eee617f34b6#ee58d8a537cec77e07595b1f0b577eee617f34b6" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=1eea60a14f214505741d22bdd6c9b501b21c04ee#1eea60a14f214505741d22bdd6c9b501b21c04ee" dependencies = [ "anyhow", "dyn-clone", @@ -450,7 +450,7 @@ dependencies = [ [[package]] name = "egglog-concurrency" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=ee58d8a537cec77e07595b1f0b577eee617f34b6#ee58d8a537cec77e07595b1f0b577eee617f34b6" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=1eea60a14f214505741d22bdd6c9b501b21c04ee#1eea60a14f214505741d22bdd6c9b501b21c04ee" dependencies = [ "arc-swap", "bumpalo", @@ -462,7 +462,7 @@ dependencies = [ [[package]] name = "egglog-core-relations" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=ee58d8a537cec77e07595b1f0b577eee617f34b6#ee58d8a537cec77e07595b1f0b577eee617f34b6" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=1eea60a14f214505741d22bdd6c9b501b21c04ee#1eea60a14f214505741d22bdd6c9b501b21c04ee" dependencies = [ "anyhow", "bumpalo", @@ -490,7 +490,7 @@ dependencies = [ [[package]] name = "egglog-experimental" version = "3.0.0" -source = "git+https://github.com/egraphs-good/egglog-experimental.git?rev=2b4627a5806f8476bc34814ecf817e3b77c16f87#2b4627a5806f8476bc34814ecf817e3b77c16f87" +source = "git+https://github.com/egraphs-good/egglog-experimental.git?rev=e0a20ce67bbfedf4ace91f4235293165dbb098bc#e0a20ce67bbfedf4ace91f4235293165dbb098bc" dependencies = [ "egglog", "egglog-ast", @@ -506,12 +506,12 @@ dependencies = [ [[package]] name = "egglog-numeric-id" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=ee58d8a537cec77e07595b1f0b577eee617f34b6#ee58d8a537cec77e07595b1f0b577eee617f34b6" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=1eea60a14f214505741d22bdd6c9b501b21c04ee#1eea60a14f214505741d22bdd6c9b501b21c04ee" [[package]] name = "egglog-reports" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=ee58d8a537cec77e07595b1f0b577eee617f34b6#ee58d8a537cec77e07595b1f0b577eee617f34b6" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=1eea60a14f214505741d22bdd6c9b501b21c04ee#1eea60a14f214505741d22bdd6c9b501b21c04ee" dependencies = [ "clap", "hashbrown 0.16.1", @@ -525,7 +525,7 @@ dependencies = [ [[package]] name = "egglog-union-find" version = "3.0.0" -source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=ee58d8a537cec77e07595b1f0b577eee617f34b6#ee58d8a537cec77e07595b1f0b577eee617f34b6" +source = "git+https://github.com/saulshanabrook/egg-smol.git?rev=1eea60a14f214505741d22bdd6c9b501b21c04ee#1eea60a14f214505741d22bdd6c9b501b21c04ee" dependencies = [ "crossbeam", "egglog-concurrency", diff --git a/Cargo.toml b/Cargo.toml index 6f65e460..7b9ecf4d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,7 +10,9 @@ name = "egglog" crate-type = ["cdylib"] [dependencies] -pyo3 = { version = "0.27", features = ["extension-module", "num-bigint", "num-rational", "indexmap"] } +# Maturin enables `pyo3/extension-module` through pyproject.toml. Keep it out of +# the base feature set so `cargo test` links against libpython on Unix. +pyo3 = { version = "0.27", features = ["num-bigint", "num-rational", "indexmap"] } num-bigint = "*" num-rational = "*" indexmap = "2.12" @@ -18,14 +20,13 @@ opentelemetry = "0.28" opentelemetry-otlp = { version = "0.28", features = ["http-proto", "reqwest-blocking-client", "trace"] } opentelemetry-stdout = { version = "0.28", features = ["trace"] } opentelemetry_sdk = "0.28" -# Egglog main at integration plus the rule-name round-trip and user-output -# downcasting fixes used by this branch. -egglog = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "ee58d8a537cec77e07595b1f0b577eee617f34b6", default-features = false } -egglog-ast = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "ee58d8a537cec77e07595b1f0b577eee617f34b6" } -egglog-core-relations = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "ee58d8a537cec77e07595b1f0b577eee617f34b6" } -egglog-reports = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "ee58d8a537cec77e07595b1f0b577eee617f34b6" } -egglog-bridge = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "ee58d8a537cec77e07595b1f0b577eee617f34b6" } -egglog-experimental = { git = "https://github.com/egraphs-good/egglog-experimental.git", rev = "2b4627a5806f8476bc34814ecf817e3b77c16f87", default-features = false } +# Egglog main plus the core compatibility fixes required by this branch. +egglog = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "1eea60a14f214505741d22bdd6c9b501b21c04ee", default-features = false } +egglog-ast = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "1eea60a14f214505741d22bdd6c9b501b21c04ee" } +egglog-core-relations = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "1eea60a14f214505741d22bdd6c9b501b21c04ee" } +egglog-reports = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "1eea60a14f214505741d22bdd6c9b501b21c04ee" } +egglog-bridge = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "1eea60a14f214505741d22bdd6c9b501b21c04ee" } +egglog-experimental = { git = "https://github.com/egraphs-good/egglog-experimental.git", rev = "e0a20ce67bbfedf4ace91f4235293165dbb098bc", default-features = false } egraph-serialize = { version = "0.3", features = ["serde", "graphviz"] } serde_json = "1" pyo3-log = "*" @@ -43,8 +44,8 @@ base64 = "0.22.1" debug = true [patch."https://github.com/egraphs-good/egglog.git"] -egglog = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "ee58d8a537cec77e07595b1f0b577eee617f34b6" } -egglog-ast = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "ee58d8a537cec77e07595b1f0b577eee617f34b6" } -egglog-core-relations = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "ee58d8a537cec77e07595b1f0b577eee617f34b6" } -egglog-reports = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "ee58d8a537cec77e07595b1f0b577eee617f34b6" } -egglog-bridge = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "ee58d8a537cec77e07595b1f0b577eee617f34b6" } +egglog = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "1eea60a14f214505741d22bdd6c9b501b21c04ee" } +egglog-ast = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "1eea60a14f214505741d22bdd6c9b501b21c04ee" } +egglog-core-relations = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "1eea60a14f214505741d22bdd6c9b501b21c04ee" } +egglog-reports = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "1eea60a14f214505741d22bdd6c9b501b21c04ee" } +egglog-bridge = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "1eea60a14f214505741d22bdd6c9b501b21c04ee" } diff --git a/docs/changelog.md b/docs/changelog.md index 0211726d..8d56a00b 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -4,66 +4,52 @@ _This project uses semantic versioning_ ## UNRELEASED -- Upgrade the Python extension to Egglog 3 and the matching - `egglog-experimental` APIs [#414](https://github.com/egraphs-good/egglog-python/pull/414) - - BREAKING: remove `Map.rebuild()`, `Set.rebuild()`, and `Vec.rebuild()`; - rebuilding is now handled by the backend container sorts. - - Preserve the Egglog 3 sort, constructor, function, and proof metadata that - its source syntax can round-trip in the low-level AST bindings; support - arbitrary-size `BigInt` values, generic value extraction, - constructor/relation lookup, correct run-report duration units with - microsecond resolution, direct parse-and-run execution with source - filenames, and batch-level command recording that omits failed parsed - batches. - - Add generic `Pair` and `Maybe` values, undefined-result `catch`, map - folding, map/set lengths, `f64` math primitives, `f64.is_finite()`, and - integer coercion; expand the experimental `Rational` Python API with - public `RationalLike` conversions, reflected arithmetic, powers, - `min`/`max`, and `Unit` comparisons; add `i64`-to-`BigRat` coercion and - exact `BigRat.to_i64()` conversion. - - Let Python function, method, and constant bodies lower as eager - primitives, preserving Python argument order for `reverse_args` callables, - while bodies attached to an explicit ruleset remain rewrite-backed; allow - `constant(..., merge=...)` for merged function-backed constants. - Higher-order callable probing is isolated from the active ruleset, and - `rule(..., eval_mode=...)` exposes safe `naive` and explicit - `unsafe-seminaive` rule evaluation when callbacks read mutable tables. - - Add persistent ordinary backoff schedules, expose `RunReport.can_stop`, - and keep high-level saturation running while a scheduler has deferred work. - - Configure worker threads and rule decomposition per `EGraph`, allow - `rule(..., no_decomp=True)` for individual rules, and expose getters and - setters for both settings. - - BREAKING: stop reading `RAYON_NUM_THREADS`; pass `num_threads` to - `EGraph` or call `set_num_threads(...)` instead. - - Add `tree` and true `greedy-dag` extraction modes to `extract` and - `extract_multiple`; add destructive `keep_best`; allow multi-root variant - extraction while preserving input order. Low-level bindings can recover the - structured aggregate returned by experimental `multi-extract` with - `UserDefinedCommandOutput.as_multi_extract()`. - - Use the experimental dynamic cost model consistently across the default - extraction paths, including canonical `set_cost` tables. A compatible raw - cost table already occupying the canonical name is reused, while an - incompatible callable or overload occupying that name when the cost table - is created now raises instead of allocating an ignored suffix. Frozen - snapshots preserve both views of a reused raw cost table. - - Name the existing custom callback protocol `TreeCostModel`, retaining - `CostModel` as a compatibility alias, and add frozen additive - `DagCostModel` values that work with tree or greedy-DAG extraction. - The low-level `bindings.Extractor` compatibility facade now prepares - costs on each extraction call, so construction no longer invokes cost - callbacks and later e-graph mutations are observed safely. - - BREAKING: remove `GreedyDagCost`, `GreedyDagCostModel`, and - `greedy_dag_cost_model`; use `DagCostModel` with - `extractor="greedy-dag"` instead. - - Improve source transcripts, diagnostics, generated-name collision - handling, shared-expression factoring without leaking synthetic bindings - into rules or checks, large-program AST parsing, custom-cost extraction of - literal roots, typed pretty/freeze round trips, map duplicate-key behavior, - and set iteration deduplication. - - Preserve the paused Param-Eq research as a small reusable experimental - module and CLI, three bounded public CI stress cases, and an optional - aggregate-only external-corpus harness with explicit iteration-limit, - timeout, and error accounting plus loaded-extension provenance. +- Upgrade to Egglog 3 and matching `egglog-experimental` APIs + [#414](https://github.com/egraphs-good/egglog-python/pull/414). + - BREAKING: container rebuilding is now handled by Egglog, so `Map.rebuild()`, + `Set.rebuild()`, and `Vec.rebuild()` are removed; configure threads through + `EGraph(num_threads=...)` instead of `RAYON_NUM_THREADS`; and replace the + removed greedy-DAG cost helpers with `DagCostModel` and + `extractor="greedy-dag"`. + - BREAKING: function, method, constant, and class-variable bodies now execute + eagerly unless an explicit `ruleset` keeps an eqsort body rewrite-backed. + Add merged constants, `reverse_args` support for bodies, per-rule evaluation + modes and decomposition control, and clear errors for invalid callable + option combinations and non-call top-level expression actions. + - Add generic `Pair` and `Maybe` values, `catch`, map folding, map/set lengths, + more `f64` operations including `is_finite()`, integer coercions, exact + `BigRat.to_i64()`, and an expanded overflow-safe experimental `Rational` API + with clearly partial powers, logarithms, and exact roots. Also fix duplicate + map keys and duplicate set iteration values. + - Add tree and greedy-DAG extraction modes, ordered multi-root variant + extraction, destructive `keep_best`, consistent dynamic `set_cost` support, + and custom `TreeCostModel` and additive `DagCostModel` callbacks. Custom-cost + failures now propagate without partially updating output term DAGs, and + opaque lookup values from another e-graph, a popped scope, or before + compaction are rejected. Computed negative dynamic costs now fail when + written instead of panicking during later extraction. User-declared raw + cost tables can be reused for table-backed callables, but not for eager or + builtin primitives whose valid rows cannot be recovered from a snapshot; + frozen snapshots preserve negative raw rows without replaying them as valid + dynamic costs. + - Add persistent backoff schedules, `RunReport.can_stop`, per-e-graph thread + and decomposition settings, constructor/relation table inspection, generic + value extraction, and `var()` typing for parameterized expression types. + Correct run-report durations that were previously 1,000 times too small. + Add removable source transcripts that preserve rule evaluation and + decomposition settings, use replay-safe names, report source-aware parse + errors, and reject use after non-replayable failures. + Python exceptions raised by `PyObject` primitives on worker threads now + propagate to the caller, and the low-level bindings round-trip Egglog 3 + AST/report data, parse and run source programs directly, and expose + structured experimental multi-extraction results. Generated backend names + remain fully qualified, avoid spellings parsed as Egglog syntax, and + reserve explicit names before lowering so ordinary explicit names remain + available regardless of action order within one registration batch. + - Preserve the paused Param-Eq work as a reusable experimental module and CLI, + three bounded CI stress cases, and an optional external-corpus harness with + explicit iteration-limit, timeout, memory-limit, error, and provenance + reporting. ## 13.2.0 (2026-06-03) diff --git a/docs/explanation/2026_02_containers.md b/docs/explanation/2026_02_containers.md index eadcbc48..aec4651b 100644 --- a/docs/explanation/2026_02_containers.md +++ b/docs/explanation/2026_02_containers.md @@ -154,14 +154,15 @@ two expression `f(a)` and `f(b)` in the e-graph, and then you make `a == b`, the this same property to hold for something like a vector, so if you have `Vec(a, c)` and `Vec(b, c)` in the e-graph, and you make `a == b`, then these two vecs should also be equal `Vec(a, c) == Vec(b, c)`. -We do this by implementing one additional operation on containers, rebuilding. This is called whenever we want to renormalize -the e-graph to preserve congruence. We defer it so we don't do it after every union operation, to reduce the amount of work. +Egglog's container primitives implement an internal rebuilding operation. The backend calls it whenever it renormalizes +the e-graph to preserve congruence; there is no public Python `.rebuild()` method. Rebuilding is deferred so it does not run +after every union operation, reducing the amount of work. Since containers "contain" references to other e-classes, we need to update those references. That what this rebuilding operation does, so that when its time to rebuild, the `Vec` type calls rebuilding on each of its inner values, updating them with new names for each e-class. So then when we check for equality after that, it will preserve congruence. -Egglog doesn't know anything more about the structures of containers besides how to rebuild them and any primitive functions you define on them. +Egglog doesn't know anything more about the structures of containers besides their backend rebuilding support and any primitive functions you define on them. This both makes them relatively easy to implement and add, but also limits the ability to "match" over them, which will see how to work around in the next section. @@ -191,7 +192,7 @@ z = egraph.let("z", sum_(MultiSet(x, y))) egraph.check(z == sum_(MultiSet(y, x))) ``` -We have the rebuilding property we talked about above as well, to maintain congruence. If we now union `x` with `y`, +The backend rebuilding described above also maintains congruence. If we now union `x` with `y`, the sum will reflect this to become `sum_(MultiSet(x, x))`: ```{code-cell} python @@ -240,7 +241,7 @@ def constant_fold_index(xs: MultiSet[Num], i: i64, k: i64): # Try replacing any sum with the folded version yield rule( sum_(xs), - # These are conditions for the rewrite to match: + # These are conditions for the rule to match: # Look for a multiset that contains two numbers that # are not the same one ms_num_index(xs, Num(i)), @@ -340,7 +341,7 @@ to the semantics of the your use case, compared to say a tree of binary operatio This work also highlights some of the current limitations of egglog. -The original version of this post identified primitive composition as a major +**Update (September 2026):** The original version of this post identified primitive composition as a major limitation: argument-order adapters required bespoke flipped backend primitives. The current Python bindings can lower body-defined functions and lambdas to primitives, so those adapters can now be expressed at the call site. @@ -365,9 +366,9 @@ since we can use their reference implementation in Mathematica to verify that ou ![meme from tim and eric TV show with someone miming their mind being blown, with the text "yarn = polynomials" imposed](./2026_02_yarn-polynomials.gif) -*Note that all code for this case study is reproducible in [this notebook](https://github.com/egraphs-good/egglog-python/blob/270a1876b6dbea37e441c132adbfdc8c11cbb319/docs/explanation/2026_02_containers_code.ipynb).* -*It is currently based on a branch of the Python bindings and Rust source, that adds additional multiset operations.* -For this docs version, the notebook content is reproduced later in this page in a folded appendix block. +*The [commit-linked notebook](https://github.com/egraphs-good/egglog-python/blob/270a1876b6dbea37e441c132adbfdc8c11cbb319/docs/explanation/2026_02_containers_code.ipynb) +is the historical original. The maintained reproduction appears below in +Appendix 1.* We define define a function to produce the amount of bending for a certain point over the [Python Array API Specification](https://data-apis.org/array-api/latest/API_specification/), so that it works on both concrete NumPy arrays and symbolic arrays. It takes in a number of 1D arrays and returns a 0D array. @@ -719,7 +720,7 @@ to never saturating: Instead, if we represented this as a product of a multiset, we could simply have a rule that looked for a zero element in the multiset and replaced that with zero. Then there would be no associativity needed, and so no chance for this to blow up. A `product(MultiSet(...))` -operation can handle associativity and commutativity and the rebuilding handles merges. +operation can handle associativity and commutativity while backend rebuilding handles merges. When I asked on the EGraph's Zulip for more examples, Sophia B also [shared another example with me](https://egraphs.zulipchat.com/#narrow/channel/328972-general/topic/A.2FC.20Blowup.20Example/near/573091425). If you have the rule `f(a + b) + 1 = f(a) + f(b)` plus A/C, you can derive equalities like `f(x) + (f(y) + f(z)) = f(x + (y + z)) + 2`, but it can take a large number of nodes. Instead in this system, we would have to encode that rule over multisets and add constant propagation to the sum function, to see how it could be found more directly, through normalization. diff --git a/docs/reference/bindings.md b/docs/reference/bindings.md index a677ec83..5eac0981 100644 --- a/docs/reference/bindings.md +++ b/docs/reference/bindings.md @@ -40,6 +40,16 @@ commands = egraph.parse_program(eqsat_basic) egraph.run_program(*commands) ``` +`EggSmolError.replayable_by_fail` conservatively reports whether the failing +command can be reproduced by wrapping it in Egglog's `(fail ...)` command. + +Experimental commands may return typed data through +`UserDefinedCommandOutput`. Command execution first returns a +`UserDefinedOutput` wrapper; call `output.output.as_multi_extract()` on its +payload. This returns a `MultiExtractOutput` for `multi-extract` results and +`None` for other custom outputs. Its `termdag` is shared by all results, while +`terms` contains one ordered list of variant term IDs per input root. + The commands are a representation which is close the AST of the egglog text language. We can see this by printing the commands: diff --git a/docs/reference/egglog-translation.md b/docs/reference/egglog-translation.md index f6ac84fd..e015985d 100644 --- a/docs/reference/egglog-translation.md +++ b/docs/reference/egglog-translation.md @@ -7,6 +7,16 @@ file_format: mystnb The high level bindings available at the top module (`egglog`) expose most of the functionality of the `egglog` text format. This guide explains how to translate between the two. Any EGraph can also be converted to egglog with the `egraph.as_egglog_string` property, as long as it was created with `EGraph(save_egglog_string=True)`. +Call `egraph.close()` to remove that saved transcript when it is no longer +needed. A recorded e-graph does not accept further commands after it is closed. +If a failure cannot be represented by Egglog's `(fail ...)` command—for +example, a Python exception or a parse, expansion, or typechecking error—the +transcript can no longer be guaranteed to replay the live state. Later +commands and transcript reads then raise `RuntimeError`. +Because the transcript is Egglog source, explicit callable, sort, variable, +ruleset, and `let` names must also be unambiguous Egglog symbols. Names that are +valid only through the direct AST API are rejected before executing the saved +command; direct e-graphs continue to accept them. ## Builtin Types @@ -42,12 +52,21 @@ i64(10) + 2 BigRat(1, 2) / BigRat(2, 1) ``` -The floating-point sort also exposes the backend's `exp()`, `log()`, and -`sqrt()` primitives. `BigRat.to_i64()` is partial: it is defined only when the -rational value is an integer that fits in `i64`. As with other partial +The floating-point sort also exposes the backend's `exp()`, `log()`, `sqrt()`, +and `is_finite()` primitives. `BigRat.to_i64()` is partial: it is defined only +when the rational value is an integer that fits in `i64`. As with other partial primitives, undefined use in a rule fact skips that match, while undefined use in an action is an error. +The experimental `Rational` sort stores canonical fractions with i64 +numerators and denominators and accepts `fractions.Fraction` and integer inputs +whose components fit in i64. Those values work in arithmetic, powers, +`min`/`max`, and comparisons. Construction, arithmetic, negation, and absolute +value are undefined when the canonical result cannot be represented. Powers +require nonnegative integer exponents (`0 ** 0` is undefined); square and cube +roots require exact rational results; and `log()` is defined only at one. Its +comparisons return a `Unit` fact rather than a Python `bool`. + ### `!=` Operator The `!=` function in egglog works on any two types with the same sort. In Python, this is mapped to the `ne` function: @@ -69,7 +88,9 @@ class Math(Expr): pass ``` -By default, the egg sort name is generated from the Python class name. You can override this if you wish with the `egg_sort` keyword argument: +By default, the Egglog sort name is generated from the module-qualified Python +class name and made safe for Egglog source. You can override it with the +`egg_sort` keyword argument: ```{code-cell} python class Math(Expr, egg_sort="Math2"): @@ -142,7 +163,9 @@ Note that instead of using `i64` as the argument type, we used `i64Like` which i The `function` decorator also accepts keyword arguments that map to backend features. Which ones are valid depends on how the callable lowers, as described in [Functions vs Constructors](#functions-vs-constructors): -- `egg_fn`: The name of the function in egglog. By default, this is the same as the Python function name. +- `egg_fn`: The name of the function in Egglog. By default, this is generated + from the module-qualified Python function name and made safe for Egglog + source. - `merge`: A function to merge the results of function-style declarations. This must take the old and new return values and return a single value of the same type. - `cost`: The extraction cost for constructor-style declarations. @@ -160,7 +183,8 @@ The static types on the decorator preserve the type of the underlying function, The Python bindings follow the backend split in egglog: -- primitive-returning callables use function-style lowering +- non-`Unit` primitive-returning callables use function-style lowering +- bodyless `Unit`-returning callables use relation-style lowering - eqsort-returning callables use constructor-style lowering That is not a Python-only policy choice. It comes from which backend features exist on each command: @@ -177,7 +201,8 @@ For bodies and defaults, the canonical lowering mapping is: | Python shape | Lowering | | --- | --- | -| primitive return, no body | lower to `function` | +| non-`Unit` primitive return, no body | lower to `function` | +| `Unit` return, no body | lower to `relation` | | primitive return, body | lower to eager `primitive` | | eqsort return, no body, no `merge` | lower to `constructor` | | eqsort return, no body, with `merge` | lower to `function` | @@ -212,7 +237,10 @@ For the Python ergonomics of attaching rewrite-backed bodies/defaults to an expl In egglog, the `(datatype ...)` command can also be used to declare functions. All of the functions declared in this block return the type of the declared datatype. Similarly, in Python, any methods of an `Expr` will be registered automatically. These can be either instance methods (including any supported `__` method), class methods, or the `__init__` method. The return type of these functions is inferred from the return type of the function. Additionally, any supported keyword argument for the `@function` decorator can be used here as well, by using the `@method` decorator to add values. -Note that by default, the egg name for any method is the Python class name combined with the method name. This allows us to define two classes with the same method name, with different signatures, that map to different egglog functions. +By default, a method's Egglog name is generated from its module-qualified +Python class name and method name, then made safe for Egglog source. This lets +classes define methods with the same Python name and different signatures +without mapping them to the same Egglog function. ```{code-cell} python # egg: @@ -374,13 +402,9 @@ This will be taken into account when extracting. Any value that can be converted to an `i64` is supported, so dynamic costs can be created in rules; the resulting cost must be nonnegative. -Python creates a canonical table on demand for each backend function symbol -whose cost is set. The table maps the function's arguments to an `i64`. -Compatible aliases of that symbol share the table. Incompatible overloads are -rejected because one backend symbol cannot have multiple cost-table schemas. - _Note: Unlike in Egglog source, Python does not require a separate declaration -that a callable supports custom costs; calling `set_cost` creates its table._ +that a callable supports custom costs; calling `set_cost` enables them +automatically._ You can also get the cost of a function with `get_cost`, which will return an `i64` if one has already been set. @@ -408,6 +432,13 @@ with `rule(..., eval_mode="naive")`. The third mode, `eval_mode="unsafe-seminaive"`, skips semi-naive validation and should only be used when the rule is known to be valid under that evaluation strategy. +Egglog normally decomposes rules before execution. `EGraph` defaults to +`no_decomp=False`; pass `no_decomp=True` to disable decomposition for +subsequently registered rules, and use `no_decomp()` or +`set_no_decomp(...)` to inspect or change that setting. For a single rule, +pass `no_decomp=True` to `rule(...)` instead. This is an advanced execution +control. + ### Variables Unlike in egglog, variables must be declared before being use and must be given a type. They need a type both so that they can be checked statically and also so that we know what types are used to understand what how the names of the egg functions correspond to the method names. @@ -717,9 +748,7 @@ egraph.register( # (extract y :variants 2) y = egraph.let("y", Math(6) + Math(2) * Math.var("x")) egraph.run(10) -# TODO: For some reason this is extracting temp vars -# egraph.extract_multiple(y, 2) -egraph +egraph.extract_multiple(y, 2) ``` ## Push/Pop @@ -744,7 +773,7 @@ egraph.check_fail(eq(Math(0)).to(Math(1))) The `(print-size ?)` command is translated into either `egraph.function_size(fn)` to get the number of rows in one table-backed callable or `egraph.all_function_sizes()` to list the sizes of all registered function tables. -Relations, constructors, and bodyless functions have tables; eager and builtin primitives do not: +Relations, constructors, bodyless functions, and bodyless constants have tables; eager and builtin primitives do not: ```{code-cell} python # (function-size Math) diff --git a/docs/reference/python-integration.md b/docs/reference/python-integration.md index 831075e9..ea9e45df 100644 --- a/docs/reference/python-integration.md +++ b/docs/reference/python-integration.md @@ -90,30 +90,6 @@ match MyExpr("hello"): print(f"Matched MyExpr with value: {value}") ``` -## 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. - -```{code-cell} python -from fractions import Fraction - -numeric_egraph = EGraph() -numeric_egraph.check(f64(1.0).is_finite()) -result = numeric_egraph.extract(Rational(1, 2) + Fraction(1, 3)) -assert result.value == Fraction(5, 6) -numeric_egraph.check(Rational(1, 2) < 1) -``` - -Rational comparisons return `Unit`, not a Python boolean. A false comparison -is therefore undefined, as are operations such as division by zero and powers -that the backend cannot represent. - ## Python Object Sort We define a custom "primitive sort" (i.e. a builtin type) for `PyObject`s. This allows us to store any Python object in the e-graph. @@ -369,9 +345,16 @@ if the need arises: ### "Preserved" methods -You can use the `@method(preserve=True)` decorator to mark a method as "preserved", meaning that calling it will actually execute the body of the function and a corresponding egglog function will not be created, +You can use the `@method(preserve=True)` decorator to mark a method as +"preserved", meaning that calling it executes ordinary Python behavior and no +corresponding egglog function is created. -Normally, all methods defined on a egglog `Expr` will ignore their bodies and simply build an expression object based on the arguments. +A bodyless method written with `...` builds an egglog call. A body-defined +method is lowered eagerly, or as a rewrite when an eqsort method has an +explicit ruleset, as described in the +[function translation rules](egglog-translation.md#functions-vs-constructors). +Use `preserve=True` when the method must instead return an ordinary Python +value or perform ordinary Python behavior. However, there are times in Python when you need the return type of a method to be an instance of a particular Python type, and some similar acting expression won't cut it. @@ -527,10 +510,9 @@ mutate_egraph.check(eq(incremented).to(Int(10) + Int(1))) mutate_egraph ``` -The bodyful form lowers to an eager primitive. Because this example constructs -an e-class value, bind its result through an action before using that value in -a read-only check. Use an explicit `ruleset=` when the body should remain a -rewrite instead. +The body executes eagerly, so it does not add a rewrite or require a `run()`. +The `let` action registers the resulting e-class value before the read-only +check. Use an explicit `ruleset=` when the body should remain a rewrite. Note that dunder methods such as `__setitem__` will automatically be marked as mutating their first argument. @@ -596,8 +578,6 @@ def map_add_two(x: MathList) -> MathList: check_eq(map_add_two(MathList.EMPTY.append(Math(1))), MathList.EMPTY.append(Math(1) + Math(2)), math_list_ruleset.saturate()) ``` -Generated primitive names are internal implementation details. - ## Default Replacements The full lowering matrix for functions, constructors, primitives, and defaults is documented in @@ -673,25 +653,6 @@ egraph.check(eq(x).to(WrappedMath(math_float(3.14)) + WrappedMath(math_float(3.1 egraph ``` -## Param-Eq Stress Demo - -The experimental `egglog.exp.param_eq` module preserves a bounded -parameter-reducing symbolic-regression pipeline. Its CLI runs either retained -representation and emits a JSON report: - -```{code-block} console -$ python -m egglog.exp.param_eq --expr '2.3 * (3.7*x0 + 5.1*x1) / 7.9' --variant container -``` - -Expressions use finite numeric literals, variables, Python arithmetic with -literal exponents, and `abs`, `exp`, `log`, `sqrt`, `plog`, `square`, or -`cube`. A `saturated` status means every inner schedule could stop; -`iteration_limit` means the retained 30-round boundary was reached. The rules -target real inputs where every relevant subexpression is defined, and the -included finite sample checks are regression tests rather than a proof of -universal equivalence. The container variant rejects inputs whose coefficient -normalization produces a non-finite `f64` value. - ## Debugging and Inspection When a rule does not fire or an equality appears unexpectedly, the most useful @@ -753,8 +714,8 @@ stats.num_matches_per_rule ### `function_values` Use {meth}`egglog.egraph.EGraph.function_values` to inspect the current rows in a -function table. This accepts relations, constructors, and bodyless functions; -eager and builtin primitives do not have tables to inspect: +function table. This accepts relations, constructors, bodyless functions, and +bodyless constants; eager and builtin primitives do not have tables to inspect: ```{code-cell} python egraph.function_values(score) @@ -811,20 +772,16 @@ costs are returned directly rather than through a wrapper object. ### Dynamic Costs -Without a custom `cost_model`, tree and greedy-DAG extraction use the -experimental dynamic cost model. A row cost registered by `set_cost` overrides -that node's marginal cost; otherwise the model falls back to costs declared on -callables and then the backend default. The same model is used by -`extract_multiple` and `keep_best`. Dynamic row costs must be nonnegative. - -Dynamic row costs live in a canonical table named -`cost_table_`. 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. +Ordinary single-root tree extraction uses Egglog's default tree cost model. +When no explicit `cost_model` is supplied, registering a row cost with +`set_cost`, selecting greedy-DAG extraction, or using multi-root extraction or +`keep_best` selects the experimental dynamic cost model. A row cost then +overrides that node's marginal cost; otherwise the model falls back to costs +declared on callables and then the backend default. An explicit `TreeCostModel` +or `DagCostModel` takes precedence over those dynamic row costs. +Dynamic row costs must be nonnegative. Literal negatives are rejected when +`set_cost` is constructed; a computed negative fails when the action runs and +is not stored. ### Multiple Roots @@ -840,17 +797,10 @@ variants_by_root = egraph.extract_multiple([expr1, expr2], 3, extractor="greedy- An inner list may be empty when its root has no extractable variant. The variant count must be positive, and the sequence form rejects an empty input. The roots share extraction preparation, but every root and variant is costed -independently; sharing between separate roots does not reduce either cost. -This API always uses dynamic costs; custom Python cost models are supported -only by single-root `extract`. - -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. +independently; sharing between separate roots does not reduce either cost. The +sequence form uses dynamic costs; the single-expression form follows the +ordinary `extract` selection described above. Custom Python cost models are +supported only by single-root `extract`. ### Custom Tree Cost Models @@ -868,6 +818,8 @@ Use {meth}`egglog.egraph.EGraph.lookup_function_value` when a model needs to inspect a table registered before extraction using those callback arguments directly. A same-e-graph lookup cannot evaluate a newly derived key while extraction is holding the graph read-only; such a lookup raises `ValueError`. +Opaque user-sort values returned by a lookup belong to that e-graph and are +rejected if passed to a different one. For example, this model uses a boolean cost for whether an `i64` is even: @@ -900,6 +852,8 @@ Cost values must be effectively immutable and totally ordered. Addition must be associative, commutative, and monotone, with `identity` as a two-sided identity. Under tree extraction the values are added once per occurrence; under greedy-DAG extraction they are added once per selected shared node. +When using tree extraction, recursive e-classes must not contain a +cost-improving cycle or cost computation may not terminate. ```{code-block} python model = DagCostModel( @@ -925,14 +879,16 @@ dynamic costs: egraph.keep_best(target, other_target, extractor="greedy-dag") ``` -Each target must be a constructor, relation, or bodyless function with a -backend table; eager and builtin primitives are rejected. +Each target must be a constructor, relation, bodyless function, or bodyless +constant with a backend table; eager and builtin primitives are rejected. -This operation is destructive. It clears every table in the e-graph, then -reinserts only the extracted rows of the requested callables. Declarations and -cost-table identities remain available, although their rows are cleared unless -selected. Existing handles returned by {meth}`egglog.egraph.EGraph.let` become -invalid because their rows have been cleared. Internal let caches are -invalidated, so the same `EGraph` can safely continue registering new actions -and running rules. Call it only when dropping all unselected table state is -intended. +This operation is destructive. It clears all old rows, then rebuilds the +extracted rows of the requested callables and any constructor rows needed to +represent their values. Unreferenced rows remain absent. Declarations and +cost-table identities remain available, although their rows are cleared. +Existing handles returned by {meth}`egglog.egraph.EGraph.let` become invalid +because their rows have been cleared. Opaque values previously returned by +{meth}`egglog.egraph.EGraph.lookup_function_value` are also invalid because +compaction assigns fresh backend value identities. The same `EGraph` can +continue registering new actions and running rules afterward. Call this method +only when dropping all unselected table state is intended. diff --git a/docs/reference/usage.md b/docs/reference/usage.md index 5df6b2df..95ceb10a 100644 --- a/docs/reference/usage.md +++ b/docs/reference/usage.md @@ -60,7 +60,7 @@ It follows [SPEC 0](https://scientific-python.org/specs/spec-0000/) in terms of 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`. +`num_threads`. ```python from egglog import EGraph @@ -70,15 +70,6 @@ egraph.set_num_threads(0) assert egraph.num_threads() >= 1 ``` -## Rule decomposition - -Egglog normally decomposes rules before execution. `EGraph` defaults to -`no_decomp=False`; pass `no_decomp=True` to disable decomposition for -subsequently registered rules, and use `no_decomp()` or -`set_no_decomp(...)` to inspect or change that setting. For a single rule, -pass `no_decomp=True` to `rule(...)` instead. This is an advanced execution -control. - (community)= ## Community diff --git a/experiments/param_eq/NOTES.md b/experiments/param_eq/NOTES.md index cb91afc1..88a089a5 100644 --- a/experiments/param_eq/NOTES.md +++ b/experiments/param_eq/NOTES.md @@ -37,11 +37,12 @@ Both variants now use ordinary persistent upstream backoff. This is close to, but not identical to, the prototype scheduler and is a documented fidelity boundary rather than an unverified claim of exact replication. -The binary repeated-monomial public stress case currently reaches the retained -30-round inner limit. CI still verifies every sample point of its extracted -expression, but the corpus runner records it as `iteration_limit` and excludes -it from successful aggregates. Revisit scheduler parity before resuming corpus -measurements. +The binary repeated-monomial public stress case is the known boundary for the +retained 30-round inner limit. When it reaches that limit, CI still verifies +every sample point of its extracted expression, and the corpus runner records +it as `iteration_limit` instead of including it in numeric summaries. CI also +accepts saturation if scheduler behavior improves. Revisit scheduler parity +before resuming corpus measurements if the limit still recurs. ## Semantic limits @@ -50,9 +51,17 @@ and every introduced subexpression are defined. Guards cover the literal and structural domain boundaries needed by the retained cases, but the pipeline has no general sign or interval analysis. CI compares every configured finite sample point; those checks are regression evidence, not a universal proof of -equivalence. The container variant fails explicitly if coefficient -normalization becomes non-finite. Treat external-corpus measurements the same -way. +equivalence. The binary rules omit cancellations whose only justification is +current e-class disequality: a later merge could invalidate that test while +leaving an incorrect finite result. The remaining guarded logarithm rules have +the same positive-input domain on both sides; their disequality guard only +avoids introducing an already undefined term. The container representation +still combines equal bases algebraically, so expressions with no defined +inputs, such as `(x - x) / (x - x)`, are outside the comparison contract and +may normalize differently between representations. Before broadening that +contract, add monotone per-e-class definedness/nonzero tracking. The container +variant fails explicitly if coefficient normalization becomes non-finite. +Treat external-corpus measurements the same way. ## Performance evidence worth preserving @@ -70,8 +79,9 @@ testing later: These are hypotheses supported by local probes, not causal or portable performance conclusions. The stale row-level artifacts and chronological debug -transcript were deliberately removed. A final dependency-compatible rerun is -required before publishing numerical corpus results. +transcript were deliberately removed, and no result CSVs remain checked in. A +final dependency-compatible aggregation must write and verify them before +publishing numerical corpus results. ## Rejected or parked directions @@ -92,14 +102,14 @@ required before publishing numerical corpus results. hashes the loaded native extension; the hash detects a changed executable but does not itself prove which checkout produced it. 2. Run the three public cases in both variants and keep their independent - numeric checks green. Confirm the documented status boundary: binary - `repeated_monomial` reaches `iteration_limit`, while the other reports are - `saturated`. Resolve that limit before publishing new corpus measurements. + numeric checks green. Record whether binary `repeated_monomial` still reaches + `iteration_limit`; the other reports should be `saturated`. Resolve a + recurring limit before publishing new corpus measurements. 3. Set `EGGLOG_PARAM_EQ_DATA_DIR` to the private archive and `EGGLOG_PARAM_EQ_EXPECTED_ARCHIVE_SHA256` to the stable value recovered from private research records. The runner refuses an absent or mismatched hash. 4. Run binary and container rows with the same time/memory limits and inspect - every iteration-limit/timeout/error count. + every iteration-limit/timeout/memory-limit/error count. Optionally run `make -C experiments/param_eq haskell`; this live baseline compiles its temporary runner once, forces both result counts inside the timed region, and requires Stack plus the external Haskell checkout. Its raw diff --git a/experiments/param_eq/README.md b/experiments/param_eq/README.md index 664a76ed..596daf1e 100644 --- a/experiments/param_eq/README.md +++ b/experiments/param_eq/README.md @@ -6,9 +6,9 @@ as part of `egglog`. The reusable expression domain and simplifier live in isolated row execution, resource limits, aggregation, and the research handoff. The work is paused. The bounded public demonstrations remain maintained in CI, -while the private 714-row corpus is not run automatically. -The checked-in result CSVs contain headers only as schema examples; a manifest -is created only by a final dependency-compatible aggregate run. +while the filtered private corpus is not run automatically. +No result CSVs are checked in while the work is paused. A verified, +dependency-compatible aggregate run writes the CSVs and manifest. ## Provenance and redistribution boundary @@ -19,8 +19,8 @@ Fabrício provided the original Haskell experiment repository and a separate `pandoc-symreg` archive in personal correspondence. Those private files were used for behavioral validation but are not redistributed here. The Python implementation reproduces the published method without copying the prototype's -source text; checked-in result artifacts contain only aggregate measurements -and no source expressions. +source text. Any future checked-in result artifacts must contain only aggregate +measurements and no source expressions. The supplied archive has placeholder copyright/author metadata, so attribution alone is not sufficient redistribution permission. Keep it outside this @@ -55,9 +55,9 @@ The restricted Python-like syntax accepts finite numeric literals, variables, `saturated` when every inner schedule can stop, or `iteration_limit` when the retained 30-round boundary is reached. In either case the extracted expression is available for independent checking; only saturated corpus rows contribute -to aggregates. The container variant rejects inputs whose coefficient -normalization produces a non-finite `f64` value. See `NOTES.md` for fidelity -and semantic limits. +to numeric aggregate summaries, while every row contributes to status counts. +The container variant rejects inputs whose coefficient normalization produces +a non-finite `f64` value. See `NOTES.md` for fidelity and semantic limits. ## External corpus commands @@ -72,9 +72,10 @@ make -C experiments/param_eq aggregate `binary` and `container` are useful for focused work and write expression-free row metrics only under the ignored `results/raw/` directory. `aggregate` runs the paired mode, alternating variant order by stable row hash, validates -identities, configuration, and input hashes, then replaces the tracked -aggregate CSVs and manifest. Never add files from `results/raw/`, the external -archive, source expressions, extracted expressions, or private absolute paths. +identities, configuration, and input hashes, then writes the aggregate CSVs and +manifest. Do not check in those results until that verification succeeds. Never +add files from `results/raw/`, the external archive, source expressions, +extracted expressions, or private absolute paths. The raw `external_archive_sha256` column intentionally repeats one hash of the corpus inputs, Haskell source modules, and Stack/Cabal lock/configuration files; @@ -91,11 +92,12 @@ command only when the installed extension was actually built in debug mode; the default is `release`. Full timing comparisons are single-machine exploratory measurements. The -runner isolates each row, records iteration limits, timeouts, and errors instead -of dropping them, and alternates binary/container order by stable row hash when -`--variant both` is used. Aggregate rows retain separate counts for iteration, -timeout, memory, and execution failures. Ratio summaries omit pairs whose -binary denominator is zero and expose the remaining sample as `n_ratio`. +runner isolates each row, records iteration limits, timeouts, memory limits, +and errors instead of dropping them, and alternates binary/container order by +stable row hash when `--variant both` is used. Aggregate rows retain separate +counts for iteration, timeout, memory, and execution failures. Ratio summaries +omit pairs whose binary denominator is zero and expose the remaining sample as +`n_ratio`. The optional `haskell` target compiles one temporary runner against the author-supplied implementation, then executes it once per isolated row. The diff --git a/experiments/param_eq/results/paper-replication.csv b/experiments/param_eq/results/paper-replication.csv deleted file mode 100644 index e4056960..00000000 --- a/experiments/param_eq/results/paper-replication.csv +++ /dev/null @@ -1 +0,0 @@ -implementation,dataset,algorithm,input_kind,metric,n_total,n_success,n_iteration_limit,n_timeout,n_memory_limit,n_error,value,min,q1,median,q3,max diff --git a/experiments/param_eq/results/representation-comparison.csv b/experiments/param_eq/results/representation-comparison.csv deleted file mode 100644 index 5048014f..00000000 --- a/experiments/param_eq/results/representation-comparison.csv +++ /dev/null @@ -1 +0,0 @@ -slice,metric,n_pairs,n_ratio,n_binary_missing,n_container_missing,container_better,same,container_worse,ratio_p10,ratio_p25,ratio_median,ratio_p75,ratio_p90 diff --git a/python/egglog/bindings.pyi b/python/egglog/bindings.pyi index 31104d10..c2cf50dc 100644 --- a/python/egglog/bindings.pyi +++ b/python/egglog/bindings.pyi @@ -173,7 +173,7 @@ class EGraph: ) -> SerializedEGraph: ... def set_report_level(self, level: _ReportLevel) -> None: ... def lookup_function(self, name: str, key: list[Value]) -> Value | None: ... - # `sort` must match the runtime sort returned with `value` by `eval_expr`. + # `value` must come from this EGraph's `eval_expr` and use the returned runtime sort. def extract_value(self, value: Value, sort: str) -> tuple[TermDag, int, int]: ... def eval_expr( self, expr: _Expr, *, traceparent: str | None = None, tracestate: str | None = None @@ -206,7 +206,8 @@ class Value: @final class EggSmolError(Exception): context: str - def __new__(cls, context: str) -> EggSmolError: ... + replayable_by_fail: bool + def __new__(cls, context: str, replayable_by_fail: bool = ...) -> EggSmolError: ... def __init__(self, /, *args: Any, **kwargs: Any) -> None: ... ## @@ -1047,6 +1048,7 @@ class DagCostModel(Generic[_DAG_COST]): base_value_cost: Callable[[str, Value], _DAG_COST], ) -> DagCostModel[_DAG_COST]: ... +# Each value must come from this EGraph's `eval_expr` and use the returned runtime sort. def extract_best_with_dag_cost_model( egraph: EGraph, roots: list[tuple[str, Value]], @@ -1068,6 +1070,7 @@ class Extractor(Generic[_COST]): traceparent: str | None = None, tracestate: str | None = None, ) -> Extractor[_COST]: ... + # `value` must come from this EGraph's `eval_expr` and use the returned runtime sort. def extract_best( self, egraph: EGraph, @@ -1078,6 +1081,7 @@ class Extractor(Generic[_COST]): traceparent: str | None = None, tracestate: str | None = None, ) -> tuple[_COST, _TermId]: ... + # `value` must come from this EGraph's `eval_expr` and use the returned runtime sort. def extract_variants( self, egraph: EGraph, diff --git a/python/egglog/builtins.py b/python/egglog/builtins.py index 65582c74..e628b12c 100644 --- a/python/egglog/builtins.py +++ b/python/egglog/builtins.py @@ -432,11 +432,6 @@ def to_string(self) -> String: ... class Maybe(BuiltinExpr, Generic[T], egg_sort="Maybe"): - @method(preserve=True) - @deprecated("use .value") - def eval(self) -> T | None: - return self.value - @method(preserve=True) # type: ignore[prop-decorator] @property def value(self) -> T | None: diff --git a/python/egglog/declarations.py b/python/egglog/declarations.py index a918197a..1044df29 100644 --- a/python/egglog/declarations.py +++ b/python/egglog/declarations.py @@ -32,8 +32,6 @@ __all__ = [ - "BUILTIN_EGG_FN_NAMES", - "BUILTIN_EGG_SORT_NAMES", "ActionCommandDecl", "ActionDecl", "BackOffDecl", @@ -109,8 +107,8 @@ ] -BUILTIN_EGG_FN_NAMES: set[str] = {"!="} -BUILTIN_EGG_SORT_NAMES: set[str] = set() +_BUILTIN_EGG_FN_NAMES: set[str] = {"!="} +_BUILTIN_EGG_SORT_NAMES: set[str] = set() @dataclass(match_args=False) @@ -972,6 +970,7 @@ class GetCostDecl: @dataclass(frozen=True) class ValueDecl: value: Value + owner: object = field(repr=False) ExprDecl: TypeAlias = ( diff --git a/python/egglog/egraph.py b/python/egglog/egraph.py index 2baaa115..4078251e 100644 --- a/python/egglog/egraph.py +++ b/python/egglog/egraph.py @@ -1,6 +1,7 @@ from __future__ import annotations import contextlib +import dis import inspect import pathlib import sys @@ -40,7 +41,7 @@ from .conversion import * from .conversion import convert_to_same_type, resolve_literal from .declarations import * -from .declarations import is_callable_decl_constructor +from .declarations import _BUILTIN_EGG_FN_NAMES, _BUILTIN_EGG_SORT_NAMES, is_callable_decl_constructor from .egraph_state import * from .ipython_magic import IN_IPYTHON from .pretty import pretty_decl @@ -207,6 +208,19 @@ def _resolve_merge( return resolved_merge.__egg_typed_expr__.expr +def _function_has_body(fn: FunctionType) -> bool: + """Return whether a function does more than implicitly return ``None``.""" + instructions = [ + (instruction.opname, instruction.argval) + for instruction in dis.get_instructions(fn) + if instruction.opname not in {"CACHE", "NOP", "RESUME"} + ] + return instructions not in ( + [("RETURN_CONST", None)], + [("LOAD_CONST", None), ("RETURN_VALUE", None)], + ) + + CALLABLE = TypeVar("CALLABLE", bound=Callable) CONSTRUCTOR_CALLABLE = TypeVar("CONSTRUCTOR_CALLABLE", bound=Callable[..., "Expr | None"]) @@ -299,17 +313,6 @@ def function( ) -> Callable[[Callable[P, BASE_EXPR]], Callable[P, BASE_EXPR]]: ... -# constructor -@overload -def function( - *, - egg_fn: str | None = ..., - cost: int | None = ..., - mutates_first_arg: bool = ..., - unextractable: bool = ..., -) -> Callable[[CONSTRUCTOR_CALLABLE], CONSTRUCTOR_CALLABLE]: ... - - @overload def function( *, @@ -365,8 +368,12 @@ def __new__( # type: ignore[misc] return super().__new__(cls, name, bases, namespace) builtin = BuiltinExpr in bases if builtin and egg_sort is not None: - BUILTIN_EGG_SORT_NAMES.add(egg_sort) - _register_builtin_class_egg_fns(namespace) + _BUILTIN_EGG_SORT_NAMES.add(egg_sort) + # Reserve explicit builtin names before declarations are lazily + # materialized so generated user names cannot claim them first. + for method in namespace.values(): + if isinstance(method, _WrappedMethod) and method.egg_fn is not None: + _BUILTIN_EGG_FN_NAMES.add(method.egg_fn) frame = currentframe() assert frame @@ -473,20 +480,10 @@ def _generate_class_decls( # noqa: C901,PLR0912 if has_default and ruleset is not None and not return_type_is_eqsort: msg = "Primitive-returning defaults cannot use an explicit ruleset" raise ValueError(msg) - default_mode = _normalize_callable_mode( - return_type_is_eqsort=return_type_is_eqsort, - returns_unit=type_ref == TypeRefWithVars(Ident.builtin("Unit")), - has_body=True if has_default else None, - has_ruleset=ruleset is not None, - require_body_for_ruleset=False, - has_merge=False, - builtin=False, - has_cost=False, - unextractable=False, - subsume=False, - ) resolved_default = ( - resolve_literal(type_ref, default_value, Thunk.value(decls)) if default_mode == "eager" else None + resolve_literal(type_ref, default_value, Thunk.value(decls)) + if has_default and ruleset is None + else None ) if resolved_default is not None: decls |= resolved_default @@ -494,7 +491,7 @@ def _generate_class_decls( # noqa: C901,PLR0912 type_ref.to_just(), body=resolved_default.__egg_typed_expr__ if resolved_default is not None else None, ) - if default_mode == "rewrite": + if has_default and ruleset is not None: _add_default_rewrite( decls, ClassVariableRef(cls_ident, k), type_ref, default_value, ruleset, subsume=False ) @@ -585,12 +582,6 @@ def _generate_class_decls( # noqa: C901,PLR0912 return decls -def _register_builtin_class_egg_fns(namespace: dict[str, Any]) -> None: - for method in namespace.values(): - if isinstance(method, _WrappedMethod) and method.egg_fn is not None: - BUILTIN_EGG_FN_NAMES.add(method.egg_fn) - - @dataclass class _FunctionConstructor: hint_locals: dict[str, Any] @@ -605,7 +596,7 @@ class _FunctionConstructor: def __post_init__(self) -> None: if self.builtin and self.egg_fn is not None: - BUILTIN_EGG_FN_NAMES.add(self.egg_fn) + _BUILTIN_EGG_FN_NAMES.add(self.egg_fn) def __call__(self, fn: Callable) -> RuntimeFunction: return RuntimeFunction(*split_thunk(Thunk.fn(self.create_decls, fn))) @@ -630,7 +621,7 @@ def create_decls(self, fn: Callable) -> tuple[Declarations, CallableRef]: return decls, ref -def _fn_decl( +def _fn_decl( # noqa: C901, PLR0912 decls: Declarations, egg_name: str | None, ref: FunctionRef | MethodRef | PropertyRef | ClassMethodRef | InitRef, @@ -701,8 +692,6 @@ def _fn_decl( arg_names = tuple(t.name for t in params) - merge_expr = _resolve_merge(decls, return_type, merge) - # Keep these lazy so builtin declarations do not resolve them eagerly. # Eager primitive bodies are bound in backend argument order. Rewrite-backed # bodies keep Python-order variables because their call pattern is reversed @@ -718,6 +707,8 @@ def _fn_decl( ) return_type_is_eqsort = isinstance(return_type, TypeRefWithVars) and not decls._classes[return_type.ident].builtin + has_merge = merge is not None + has_body = _function_has_body(fn) signature_ = FunctionSignature( return_type=None if mutates_first_arg else return_type, var_arg_type=var_arg_type, @@ -727,20 +718,68 @@ def _fn_decl( reverse_args=reverse_args, ) doc = fn.__doc__ - mode = _normalize_callable_mode( - return_type_is_eqsort=return_type_is_eqsort, - returns_unit=signature_.semantic_return_type == TypeRefWithVars(Ident.builtin("Unit")), - has_body=None, - has_ruleset=ruleset is not None, - require_body_for_ruleset=isinstance(ref, FunctionRef), - has_merge=merge_expr is not None, - builtin=is_builtin, - has_cost=cost is not None, - unextractable=unextractable, - subsume=subsume, - ) + if is_builtin and has_merge: + msg = "Builtin callables cannot use merge" + raise ValueError(msg) + if subsume: + if not return_type_is_eqsort: + msg = "Primitive-returning callables cannot use subsume" + raise ValueError(msg) + if ruleset is None: + msg = "subsume requires an explicit ruleset" + raise ValueError(msg) + if return_type_is_eqsort and is_builtin: + msg = "Eqsort-returning callables cannot be builtin" + raise ValueError(msg) + if not return_type_is_eqsort: + if cost is not None: + msg = "Primitive-returning callables cannot use cost" + raise ValueError(msg) + if unextractable: + msg = "Primitive-returning callables cannot be unextractable" + raise ValueError(msg) + if signature_.semantic_return_type == TypeRefWithVars(Ident.builtin("Unit")) and has_merge: + msg = "Functions that return Unit cannot use merge" + raise ValueError(msg) + elif has_merge: + # A merge makes this a function rather than a constructor, so constructor-only + # options must be rejected before evaluating a possibly effectful Python body. + if cost is not None: + msg = "Cost can only be set for constructors" + raise ValueError(msg) + if unextractable: + msg = "Unextractable can only be set for constructors" + raise ValueError(msg) + if is_builtin and ruleset is not None: + msg = "Builtin callables cannot use an explicit ruleset" + raise ValueError(msg) + if is_builtin and has_body: + msg = "Builtin callables cannot have a body" + raise ValueError(msg) + # Reject statically incompatible body options before evaluating either the + # body or merge callback. The checks in _add_default_rewrite_function stay + # as a defensive boundary for the value the body actually returns. + if has_body: + if return_type_is_eqsort: + if has_merge: + msg = "Eqsort-returning callables with bodies cannot use merge" + raise ValueError(msg) + if ruleset is None and cost is not None: + msg = "Eqsort-returning eager bodies cannot use cost" + raise ValueError(msg) + if ruleset is None and unextractable: + msg = "Eqsort-returning eager bodies cannot be unextractable" + raise ValueError(msg) + else: + if ruleset is not None: + msg = "Primitive-returning callables with bodies cannot use an explicit ruleset" + raise ValueError(msg) + if has_merge: + msg = "Primitive-returning callables with bodies cannot use merge" + raise ValueError(msg) + merge_expr = _resolve_merge(decls, return_type, merge) decl: ConstructorDecl | FunctionDecl - if mode == "constructor": + if return_type_is_eqsort and merge_expr is None: decl = ConstructorDecl(signature_, egg_name, cost, unextractable, doc) else: decl = FunctionDecl( @@ -751,11 +790,7 @@ def _fn_decl( doc=doc, ) decls.set_function_decl(ref, decl) - if is_builtin and ( - any(tp.vars for tp in arg_types) - or (var_arg_type is not None and bool(var_arg_type.vars)) - or bool(return_type.vars) - ): + if is_builtin: return lambda: None return Thunk.fn( _add_default_rewrite_function, @@ -767,6 +802,7 @@ def _fn_decl( subsume, return_type, mutates_first_arg, + has_body, context=f"creating {ref}", ) @@ -912,20 +948,21 @@ def _constant_thunk( if has_default and ruleset is not None and not return_type_is_eqsort: msg = "Primitive-returning defaults cannot use an explicit ruleset" raise ValueError(msg) + has_merge = merge is not None + if type_ref == TypeRefWithVars(Ident.builtin("Unit")) and has_merge: + msg = "Functions that return Unit cannot use merge" + raise ValueError(msg) + if has_default and has_merge: + msg = ( + "Eqsort-returning callables with bodies cannot use merge" + if return_type_is_eqsort + else "Primitive-returning callables with bodies cannot use merge" + ) + raise ValueError(msg) merge_expr = _resolve_merge(decls, type_ref, merge) - mode = _normalize_callable_mode( - return_type_is_eqsort=return_type_is_eqsort, - returns_unit=type_ref == TypeRefWithVars(Ident.builtin("Unit")), - has_body=True if has_default else None, - has_ruleset=ruleset is not None, - require_body_for_ruleset=False, - has_merge=merge_expr is not None, - builtin=False, - has_cost=False, - unextractable=False, - subsume=False, + resolved_default = ( + resolve_literal(type_ref, default_replacement, Thunk.value(decls)) if has_default and ruleset is None else None ) - resolved_default = resolve_literal(type_ref, default_replacement, Thunk.value(decls)) if mode == "eager" else None if resolved_default is not None: decls |= resolved_default decls._constants[ident] = ConstantDecl( @@ -934,91 +971,12 @@ def _constant_thunk( resolved_default.__egg_typed_expr__ if resolved_default is not None else None, merge_expr, ) - if mode == "rewrite": + if has_default and ruleset is not None: _add_default_rewrite(decls, callable_ref, type_ref, default_replacement, ruleset, subsume=False) return decls, TypedExprDecl(type_ref.to_just(), CallDecl(callable_ref)) -_CallableMode: TypeAlias = Literal["function", "constructor", "eager", "rewrite"] - - -def _normalize_callable_mode( # noqa: C901, PLR0911, PLR0912 - *, - return_type_is_eqsort: bool, - returns_unit: bool, - has_body: bool | None, - has_ruleset: bool, - require_body_for_ruleset: bool, - has_merge: bool, - builtin: bool, - has_cost: bool, - unextractable: bool, - subsume: bool, -) -> _CallableMode: - if builtin and has_merge: - msg = "Builtin callables cannot use merge" - raise ValueError(msg) - if require_body_for_ruleset and has_ruleset and has_body is False: - msg = "Explicit rulesets require a body" - raise ValueError(msg) - if subsume: - if not return_type_is_eqsort: - msg = "Primitive-returning callables cannot use subsume" - raise ValueError(msg) - if not has_ruleset: - msg = "subsume requires an explicit ruleset" - raise ValueError(msg) - if has_body is False: - msg = "subsume requires a body" - raise ValueError(msg) - - if return_type_is_eqsort: - if builtin: - msg = "Eqsort-returning callables cannot be builtin" - raise ValueError(msg) - if has_body is None: - return "function" if has_merge else "constructor" - if has_body: - if has_merge: - msg = "Eqsort-returning callables with bodies cannot use merge" - raise ValueError(msg) - if has_ruleset: - return "rewrite" - if has_cost: - msg = "Eqsort-returning eager bodies cannot use cost" - raise ValueError(msg) - if unextractable: - msg = "Eqsort-returning eager bodies cannot be unextractable" - raise ValueError(msg) - return "eager" - return "function" if has_merge else "constructor" - - if has_cost: - msg = "Primitive-returning callables cannot use cost" - raise ValueError(msg) - if unextractable: - msg = "Primitive-returning callables cannot be unextractable" - raise ValueError(msg) - if returns_unit and has_merge: - msg = "Functions that return Unit cannot use merge" - raise ValueError(msg) - if has_body is None: - return "function" - if has_body: - if has_ruleset: - msg = "Primitive-returning callables with bodies cannot use an explicit ruleset" - raise ValueError(msg) - if builtin: - msg = "Builtin callables cannot have a body" - raise ValueError(msg) - if has_merge: - msg = "Primitive-returning callables with bodies cannot use merge" - raise ValueError(msg) - return "eager" - return "function" - - -def _add_default_rewrite_function( +def _add_default_rewrite_function( # noqa: C901, PLR0912 decls: Declarations, ref: FunctionRef | MethodRef | PropertyRef | ClassMethodRef | InitRef, fn: Callable, @@ -1027,6 +985,7 @@ def _add_default_rewrite_function( subsume: bool, res_type: TypeOrVarRef, mutates_first_arg: bool, + has_body: bool, ) -> None: args = list(args) arg_exprs: list[RuntimeExpr | RuntimeClass] = [RuntimeExpr.__from_values__(decls, a) for a in args] @@ -1047,25 +1006,43 @@ def _add_default_rewrite_function( res = arg_exprs[0] decl = decls.get_callable_decl(ref) assert isinstance(decl, ConstructorDecl | FunctionDecl) - mode = _normalize_callable_mode( - return_type_is_eqsort=isinstance(res_type, TypeRefWithVars) and not decls._classes[res_type.ident].builtin, - returns_unit=res_type == TypeRefWithVars(Ident.builtin("Unit")), - has_body=res is not None, - has_ruleset=ruleset is not None, - require_body_for_ruleset=isinstance(ref, FunctionRef), - has_merge=isinstance(decl, FunctionDecl) and decl.merge is not None, - builtin=isinstance(decl, FunctionDecl) and decl.builtin, - has_cost=isinstance(decl, ConstructorDecl) and decl.cost is not None, - unextractable=isinstance(decl, ConstructorDecl) and decl.unextractable, - subsume=subsume, - ) - if mode in ("function", "constructor"): + if res is None: + if has_body: + msg = "Callable bodies must return a value" + raise ValueError(msg) + if ruleset is not None and isinstance(ref, FunctionRef): + msg = "Explicit rulesets require a body" + raise ValueError(msg) + if subsume: + msg = "subsume requires a body" + raise ValueError(msg) return - if mode == "rewrite": + return_type_is_eqsort = isinstance(res_type, TypeRefWithVars) and not decls._classes[res_type.ident].builtin + if return_type_is_eqsort and isinstance(decl, FunctionDecl) and decl.merge is not None: + msg = "Eqsort-returning callables with bodies cannot use merge" + raise ValueError(msg) + if return_type_is_eqsort and ruleset is not None: _add_default_rewrite(decls, ref, res_type, res, ruleset, subsume) return - assert res is not None + if isinstance(decl, ConstructorDecl): + if decl.cost is not None: + msg = "Eqsort-returning eager bodies cannot use cost" + raise ValueError(msg) + if decl.unextractable: + msg = "Eqsort-returning eager bodies cannot be unextractable" + raise ValueError(msg) + else: + if ruleset is not None: + msg = "Primitive-returning callables with bodies cannot use an explicit ruleset" + raise ValueError(msg) + if decl.builtin: + msg = "Builtin callables cannot have a body" + raise ValueError(msg) + if decl.merge is not None: + msg = "Primitive-returning callables with bodies cannot use merge" + raise ValueError(msg) + resolved_value = resolve_literal(res_type, res, Thunk.value(decls)) decls |= resolved_value match decl: @@ -1083,36 +1060,20 @@ def _add_default_rewrite( ref: FunctionRef | ConstantRef | MethodRef | ClassMethodRef | InitRef | ClassVariableRef | PropertyRef, type_ref: TypeOrVarRef, default_rewrite: object, - ruleset: Ruleset | None, + ruleset: Ruleset, subsume: bool, ) -> None: """ Adds a default rewrite for the callable when an explicit ruleset is provided. """ - if ruleset is None: - msg = "Default rewrites require an explicit ruleset" - raise ValueError(msg) resolved_value = resolve_literal(type_ref, default_rewrite, Thunk.value(decls)) rewrite_decl = DefaultRewriteDecl(ref, resolved_value.__egg_typed_expr__.expr, subsume) - ruleset_decls = _add_default_rewrite_inner(decls, rewrite_decl, ruleset) + ruleset_decls = ruleset._current_egg_decls + ruleset.__egg_ruleset__.rules.append(rewrite_decl) ruleset_decls |= decls ruleset_decls |= resolved_value -def _add_default_rewrite_inner( - decls: Declarations, - rewrite_decl: DefaultRewriteDecl, - ruleset: Ruleset | None, -) -> Declarations: - if ruleset is None: - msg = "Default rewrites require an explicit ruleset" - raise ValueError(msg) - ruleset_decls = ruleset._current_egg_decls - ruleset_decl = ruleset.__egg_ruleset__ - ruleset_decl.rules.append(rewrite_decl) - return ruleset_decls - - def _last_param_variable(params: list[Parameter]) -> bool: """ Checks if the last paramater is a variable arg. @@ -1154,6 +1115,8 @@ class EGraph: _state_stack: list[EGraphState] = field(default_factory=list, repr=False) # For storing the global "current" egraph _token_stack: list[EGraph] = field(default_factory=list, repr=False) + # Raw values supplied while extraction holds this e-graph read-only. + _cost_callback_values: dict[int, bindings.Value] | None = field(default=None, init=False, repr=False, compare=False) def __init__( self, @@ -1167,14 +1130,17 @@ def __init__( with _TRACER.start_as_current_span("create_bindings"): self._state = EGraphState( bindings.EGraph(seminaive=seminaive, num_threads=num_threads, no_decomp=no_decomp), + seminaive=seminaive, save_egglog_string=save_egglog_string, ) self._state_stack = [] self._token_stack = [] + self._cost_callback_values = None if actions: self.register(*actions) def _add_decls(self, *decls: DeclarationsLike) -> None: + self._state.ensure_open() for d in decls: self._state.__egg_decls__ |= d @@ -1208,12 +1174,22 @@ def set_no_decomp(self, no_decomp: bool) -> None: @property def as_egglog_string(self) -> str: """ - Returns the egglog string for this module. + Return the replayable Egglog transcript for this e-graph. + + If a failure may have partially changed the graph and cannot be + represented by a replayable command, the transcript is invalidated and + this property raises ``RuntimeError``. """ return self._state.egglog_string() def close(self) -> None: - """Close and remove the optional saved Egglog transcript.""" + """ + Close and remove the optional saved Egglog transcript. + + For an e-graph created with ``save_egglog_string=True``, subsequent + commands are rejected because they could no longer be recorded. + Otherwise this method has no effect and the e-graph remains usable. + """ self._state.close() def _ipython_display_(self) -> None: @@ -1527,11 +1503,12 @@ def keep_best( extractor: ExtractionMode = "tree", ) -> None: """ - Keep the best rows of selected callables and clear every other table. + Keep the best rows of selected callables and discard unreferenced rows. - This destructively compacts the e-graph. Declarations and dynamic-cost - table identities remain available for subsequent iteration, but their - rows are cleared unless selected by the command. + This destructively clears all old rows, then rebuilds the selected rows + and the constructor rows needed by their extracted representations. + Declarations and dynamic-cost table identities remain available for + subsequent iteration. """ extractor_options = _extractor_options(extractor) resolved = [resolve_callable(callable_) for callable_ in (fn, *fns)] @@ -1542,6 +1519,12 @@ def keep_best( args: list[bindings._Expr] = [bindings.Lit(span(2), bindings.String(table_name)) for table_name in table_names] self._state.run_program(bindings.UserDefined(span(2), "keep-best", [*args, *extractor_options])) + # Compaction rebuilds the tables with fresh raw value IDs. Reject every + # opaque value captured before it, including values inherited from a + # parent push scope; popping restores that parent's owner generation. + self._state.value_owner = object() + self._state.valid_value_owners = frozenset((self._state.value_owner,)) + # keep-best clears every table, including synthetic and user let rows. # Do not let later lowering reuse references to those now-empty tables. self._state.expr_to_letref_cache.clear() @@ -1594,7 +1577,21 @@ def __enter__(self) -> Self: return self def __exit__(self, exc_type, exc, exc_tb) -> None: - self.pop() + egglog_file_state = self._state.egglog_file_state + if egglog_file_state is None or (not egglog_file_state.poisoned and not egglog_file_state.file.closed): + self.pop() + return + + # A closed or poisoned transcript rejects ordinary commands, but the + # backend scope still has to be restored. If another exception is + # already propagating, do not replace it with a cleanup failure. + try: + call_with_current_trace(self._state.egraph.run_program, bindings.Pop(span(1), 1)) + except BaseException: + if exc_type is None: + raise + finally: + self._state = self._state_stack.pop() def _serialize( self, @@ -1762,10 +1759,20 @@ def register( def _register_commands(self, cmds: list[Command]) -> None: self._add_decls(*cmds) - egg_cmds = [egg_cmd for cmd in cmds if (egg_cmd := self._command_to_egg(cmd)) is not None] + previous_pending_let_names = self._state.pending_let_names + self._state.pending_let_names |= frozenset( + name if name.startswith("$") else f"${name}" + for cmd in cmds + if isinstance(cmd, Action) and isinstance(cmd.action, LetDecl) + for name in (cmd.action.name,) + ) + try: + egg_cmds = [egg_cmd for cmd in cmds for egg_cmd in self._commands_to_egg(cmd)] + finally: + self._state.pending_let_names = previous_pending_let_names self._state.run_program(*egg_cmds) - def _command_to_egg(self, cmd: Command) -> bindings._Command | None: + def _commands_to_egg(self, cmd: Command) -> list[bindings._Command]: ruleset_ident = Ident("") cmd_decl: CommandDecl match cmd: @@ -1776,7 +1783,7 @@ def _command_to_egg(self, cmd: Command) -> bindings._Command | None: cmd_decl = ActionCommandDecl(action) case _: assert_never(cmd) - return self._state.command_to_egg(cmd_decl, ruleset_ident) + return self._state.commands_to_egg(cmd_decl, ruleset_ident) def function_size(self, fn: ExprCallable) -> int: """ @@ -1843,12 +1850,15 @@ def lookup_function_value(self, expr: BASE_EXPR) -> BASE_EXPR | None: During a custom cost-model callback, same-e-graph lookups may use values supplied to that callback but cannot evaluate newly derived arguments while extraction holds the e-graph read-only. + + Opaque values returned for user-defined sorts belong to this EGraph + and cannot be used to query another EGraph. """ runtime_expr = to_runtime_expr(expr) typed_expr = runtime_expr.__egg_typed_expr__ assert isinstance(typed_expr.expr, CallDecl | GetCostDecl) - callback_context = _COST_MODEL_CALLBACK_VALUES.get() - in_cost_model_callback = callback_context is not None and callback_context[0] is self + callback_values = self._cost_callback_values + in_cost_model_callback = callback_values is not None if in_cost_model_callback: ref = typed_expr.expr.callable table_is_registered = ( @@ -1864,23 +1874,7 @@ def lookup_function_value(self, expr: BASE_EXPR) -> BASE_EXPR | None: if isinstance(typed_expr.expr, CallDecl): self._require_table_backed(typed_expr.expr.callable) egg_fn, typed_args = self._state.translate_call(typed_expr.expr) - if in_cost_model_callback: - assert callback_context is not None - callback_values = callback_context[1] - values_args = [] - for arg in typed_args: - if arg in callback_values: - values_args.append(callback_values[arg]) - elif isinstance(arg.expr, ValueDecl): - values_args.append(arg.expr.value) - else: - msg = ( - "Cost-model callbacks can only look up tables using values supplied to the callback; " - "evaluating new expressions would require mutating the borrowed e-graph" - ) - raise ValueError(msg) - else: - values_args = [self._state.typed_expr_to_value(arg) for arg in typed_args] + values_args = self._lookup_argument_values(typed_args, callback_values) possible_value = self._egraph.lookup_function(egg_fn, values_args) if possible_value is None: return None @@ -1892,18 +1886,44 @@ def lookup_function_value(self, expr: BASE_EXPR) -> BASE_EXPR | None: ), ) + def _lookup_argument_values( + self, + typed_args: list[TypedExprDecl], + callback_values: dict[int, bindings.Value] | None, + ) -> list[bindings.Value]: + """Resolve lookup keys without mutating an e-graph borrowed by extraction.""" + if callback_values is None: + args_to_materialize = list(typed_args) + while args_to_materialize: + arg = args_to_materialize.pop() + match arg.expr: + case CallDecl(): + self.register(cast("BaseExpr", RuntimeExpr.__from_values__(self.__egg_decls__, arg))) + case PartialCallDecl(CallDecl(args=nested_args)): + # The unstable-fn value itself has no table row. Only + # evaluating its captured calls can mutate the e-graph. + args_to_materialize.extend(nested_args) + return [self._state.typed_expr_to_value(arg) for arg in typed_args] + + values = [] + for arg in typed_args: + if id(arg) in callback_values: + values.append(callback_values[id(arg)]) + continue + if isinstance(arg.expr, ValueDecl) and arg.expr.owner in self._state.valid_value_owners: + values.append(arg.expr.value) + continue + msg = ( + "Cost-model callbacks can only look up tables using values supplied to the callback; " + "evaluating new expressions would require mutating the borrowed e-graph" + ) + raise ValueError(msg) + return values + def _require_table_backed(self, ref: CallableRef) -> None: """Reject callable shapes that Egglog lowers without a queryable table.""" - decl = self.__egg_decls__.get_callable_decl(ref) - match decl: - case RelationDecl() | ConstructorDecl() | ConstantDecl(body=None): - return - case FunctionDecl(body=None, builtin=False) if not isinstance(ref, UnnamedFunctionRef): - return - case ConstantDecl() | FunctionDecl(): - pass - case _: - assert_never(decl) + if self._state._callable_is_table_backed(ref): + return msg = ( "This operation requires a table-backed relation, constructor, or bodyless function; " "eager and builtin primitives do not have tables" @@ -1954,6 +1974,16 @@ def append_e_class_row(output: bindings.Value, tp: JustTypeRef, call: CallDecl, subsumed.append((tp, call)) synthetic_let_names = {var.name for var in self._state.expr_to_letref_cache.values()} + cost_target_inputs: dict[CallableRef, set[tuple[bindings.Value, ...]] | None] = {} + for cost_callable_ref in self._state.cost_table_names: + target_name = self._state.callable_ref_to_egg_fn[cost_callable_ref][0] + target_fn = frozen.functions.get(target_name) + # Table-backed targets need a live row before replay can set its + # cost. Eager and builtin targets have no frozen table and can be + # reconstructed directly from the cost-table arguments. + cost_target_inputs[cost_callable_ref] = ( + None if target_fn is None else {tuple(target_row.inputs) for target_row in target_fn.rows} + ) for name, fn in frozen.functions.items(): if fn.is_let_binding: if name in synthetic_let_names: @@ -1977,29 +2007,37 @@ def append_e_class_row(output: bindings.Value, tp: JustTypeRef, call: CallDecl, f"Cannot freeze special callable {callable_ref} with signature {signature}" ) assert signature.var_arg_type is None, f"Frozen calls do not support var args: {callable_ref}" - assert not signature.reverse_args, f"Frozen calls do not support reverse_args: {callable_ref}" for row in fn.rows: + python_inputs = row.inputs[::-1] if signature.reverse_args else row.inputs arg_exprs = tuple( TypedExprDecl(tp, self._state.value_to_expr(tp, value)) - for arg_type, value in zip(signature.arg_types, row.inputs, strict=True) + for arg_type, value in zip(signature.arg_types, python_inputs, strict=True) for tp in (arg_type.to_just(),) ) if is_cost: cost_tp = self._state.egg_sort_to_type_ref[fn.output_sort] cost_expr = TypedExprDecl(cost_tp, self._state.value_to_expr(cost_tp, row.output)) for cost_callable_ref in cost_callable_refs: + target_inputs = cost_target_inputs[cost_callable_ref] + if target_inputs is not None and tuple(row.inputs) not in target_inputs: + continue cost_signature = self.__egg_decls__.get_callable_decl(cost_callable_ref).signature assert isinstance(cost_signature, FunctionSignature) + cost_inputs = row.inputs[::-1] if cost_signature.reverse_args else row.inputs cost_arg_exprs = tuple( TypedExprDecl(tp, self._state.value_to_expr(tp, value)) - for arg_type, value in zip(cost_signature.arg_types, row.inputs, strict=True) + for arg_type, value in zip(cost_signature.arg_types, cost_inputs, strict=True) for tp in (arg_type.to_just(),) ) cost_call = CallDecl(cost_callable_ref, cost_arg_exprs) match cost_expr.expr: case LitDecl(int(value)): - costs[cost_call] = (cost_signature.semantic_return_type.to_just(), value) + # Raw aliases may contain negative rows. Dynamic + # extraction ignores them, and replaying them as + # validated set_cost actions would fail. + if value >= 0: + costs[cost_call] = (cost_signature.semantic_return_type.to_just(), value) case _: raise TypeError(f"Expected integer cost for {cost_callable_ref}, got {cost_expr.expr}") # A user-declared bodyless function may intentionally own the @@ -2043,7 +2081,7 @@ def append_e_class_row(output: bindings.Value, tp: JustTypeRef, call: CallDecl, def _values_to_expr_and_callback_values( self, args: list[bindings.Value], name: str - ) -> tuple[RuntimeExpr, dict[TypedExprDecl, bindings.Value]] | None: + ) -> tuple[RuntimeExpr, dict[int, bindings.Value]] | None: """Reconstruct a callback call and map its Python-order arguments to raw backend values.""" if name not in self._state.egg_fn_to_callable_refs: return None @@ -2056,13 +2094,26 @@ def _values_to_expr_and_callback_values( for arg_type, arg in zip(signature.arg_types, python_args, strict=True) for tp in (arg_type.to_just(),) ) + callback_values = dict(zip(map(id, arg_exprs), python_args, strict=True)) + stack = list(arg_exprs) + while stack: + arg = stack.pop() + match arg.expr: + case ValueDecl(value, owner): + # Reconstructed nested e-class values also belong to this + # callback borrow, but structurally equal values from + # another e-graph must not be accepted. + if owner in self._state.valid_value_owners: + callback_values[id(arg)] = value + case CallDecl(args=nested_args) | PartialCallDecl(CallDecl(args=nested_args)): + stack.extend(nested_args) res_type = signature.semantic_return_type.to_just() return ( RuntimeExpr.__from_values__( self.__egg_decls__, TypedExprDecl(res_type, CallDecl(callable_ref, arg_exprs)), ), - dict(zip(arg_exprs, python_args, strict=True)), + callback_values, ) @@ -2808,9 +2859,6 @@ def _fact_like(fact_like: FactLike) -> Fact: _CURRENT_RULESET = ContextVar[Ruleset | None]("CURRENT_RULESET", default=None) -_COST_MODEL_CALLBACK_VALUES = ContextVar[tuple[EGraph, dict[TypedExprDecl, bindings.Value]] | None]( - "COST_MODEL_CALLBACK_VALUES", default=None -) def get_current_ruleset() -> Ruleset | None: @@ -2829,14 +2877,17 @@ def set_current_ruleset(r: Ruleset | None) -> Generator[None, None, None]: @contextlib.contextmanager def _cost_model_callback_values( egraph: EGraph, - values: dict[TypedExprDecl, bindings.Value], + values: dict[int, bindings.Value], ) -> Generator[None, None, None]: """Make raw callback values available to read-only table lookups without evaluating expressions.""" - token = _COST_MODEL_CALLBACK_VALUES.set((egraph, values)) + # Values belong to the physical e-graph borrow, not Python's ambient context. + # Restore any outer callback scope if callback setup is nested. + previous = egraph._cost_callback_values + egraph._cost_callback_values = values try: yield finally: - _COST_MODEL_CALLBACK_VALUES.reset(token) + egraph._cost_callback_values = previous def get_cost(expr: BaseExpr) -> i64: @@ -2919,6 +2970,10 @@ class DagCostModel(Generic[DAG_COST]): Costs are combined with Python ``+``. The operation must be associative, commutative, monotone, and have ``identity`` as a two-sided identity. Cost values must also have a total order and be effectively immutable. + The identity is required because a generic Python ``+`` operation does not + provide a way to construct its neutral value from the cost type. + With tree extraction, recursive e-classes must not contain a + cost-improving cycle or cost computation may not terminate. """ marginal_cost: Callable[[EGraph, BaseExpr], DAG_COST] @@ -2937,8 +2992,9 @@ def default_cost_model(egraph: EGraph, expr: BaseExpr, children_costs: list[int] (callable_fn := get_callable_fn(expr)) is not None and egraph.has_custom_cost(callable_fn) and (i := egraph.lookup_function_value(get_cost(expr))) is not None + and (row_cost := int(i)) >= 0 ): - self_cost = int(i) + self_cost = row_cost # 2. Else, check if this is a callable and it has a cost set on its declaration elif callable_fn is not None and (callable_cost := get_callable_cost(callable_fn)) is not None: self_cost = callable_cost @@ -2969,7 +3025,7 @@ class _CostModel(Generic[COST]): egraph: EGraph enode_cost_results: dict[tuple[str, tuple[bindings.Value, ...]], int] = field(default_factory=dict) enode_cost_expressions: list[RuntimeExpr] = field(default_factory=list) - enode_cost_argument_values: list[dict[TypedExprDecl, bindings.Value]] = field(default_factory=list) + enode_cost_argument_values: list[dict[int, bindings.Value]] = field(default_factory=list) base_value_cost_results: dict[tuple[str, bindings.Value], COST] = field(default_factory=dict) def call_model(self, expr: RuntimeExpr, children_costs: list[COST]) -> COST: @@ -3012,7 +3068,7 @@ def base_value_cost(self, tp: str, value: bindings.Value) -> COST: self.egraph.__egg_decls__, TypedExprDecl(type_ref, self.egraph._state.value_to_expr(type_ref, value)), ) - with _cost_model_callback_values(self.egraph, {expr.__egg_typed_expr__: value}): + with _cost_model_callback_values(self.egraph, {id(expr.__egg_typed_expr__): value}): res = self.call_model(expr, []) self.base_value_cost_results[(tp, value)] = res return res @@ -3023,7 +3079,7 @@ def container_cost(self, tp: str, value: bindings.Value, element_costs: list[COS self.egraph.__egg_decls__, TypedExprDecl(type_ref, self.egraph._state.value_to_expr(type_ref, value)), ) - with _cost_model_callback_values(self.egraph, {expr.__egg_typed_expr__: value}): + with _cost_model_callback_values(self.egraph, {id(expr.__egg_typed_expr__): value}): return self.call_model(expr, element_costs) def to_bindings_cost_model(self) -> bindings.CostModel[COST, int]: @@ -3037,8 +3093,7 @@ class _DagCostModel(Generic[DAG_COST]): model: DagCostModel[DAG_COST] egraph: EGraph enode_cost_results: dict[tuple[str, tuple[bindings.Value, ...]], DAG_COST] = field(default_factory=dict) - base_value_cost_results: dict[tuple[str, bindings.Value], DAG_COST] = field(default_factory=dict) - container_cost_results: dict[tuple[str, bindings.Value], DAG_COST] = field(default_factory=dict) + value_cost_results: dict[tuple[str, bindings.Value], DAG_COST] = field(default_factory=dict) def enode_cost(self, name: str, args: list[bindings.Value]) -> DAG_COST: key = (name, tuple(args)) @@ -3056,26 +3111,10 @@ def enode_cost(self, name: str, args: list[bindings.Value]) -> DAG_COST: self.enode_cost_results[key] = result return result - def base_value_cost(self, tp: str, value: bindings.Value) -> DAG_COST: - key = (tp, value) - try: - return self.base_value_cost_results[key] - except KeyError: - pass - type_ref = self.egraph._state.egg_sort_to_type_ref[tp] - expr = RuntimeExpr.__from_values__( - self.egraph.__egg_decls__, - TypedExprDecl(type_ref, self.egraph._state.value_to_expr(type_ref, value)), - ) - with _cost_model_callback_values(self.egraph, {expr.__egg_typed_expr__: value}): - result = self.model.marginal_cost(self.egraph, cast("BaseExpr", expr)) - self.base_value_cost_results[key] = result - return result - - def container_cost(self, tp: str, value: bindings.Value) -> DAG_COST: + def value_cost(self, tp: str, value: bindings.Value) -> DAG_COST: key = (tp, value) try: - return self.container_cost_results[key] + return self.value_cost_results[key] except KeyError: pass type_ref = self.egraph._state.egg_sort_to_type_ref[tp] @@ -3083,15 +3122,15 @@ def container_cost(self, tp: str, value: bindings.Value) -> DAG_COST: self.egraph.__egg_decls__, TypedExprDecl(type_ref, self.egraph._state.value_to_expr(type_ref, value)), ) - with _cost_model_callback_values(self.egraph, {expr.__egg_typed_expr__: value}): + with _cost_model_callback_values(self.egraph, {id(expr.__egg_typed_expr__): value}): result = self.model.marginal_cost(self.egraph, cast("BaseExpr", expr)) - self.container_cost_results[key] = result + self.value_cost_results[key] = result return result def to_bindings_cost_model(self) -> bindings.DagCostModel[DAG_COST]: return bindings.DagCostModel( self.model.identity, self.enode_cost, - self.container_cost, - self.base_value_cost, + self.value_cost, + self.value_cost, ) diff --git a/python/egglog/egraph_state.py b/python/egglog/egraph_state.py index c2347f20..b68d9ba2 100644 --- a/python/egglog/egraph_state.py +++ b/python/egglog/egraph_state.py @@ -5,13 +5,14 @@ from __future__ import annotations import contextlib +import math import re import tempfile import weakref from base64 import standard_b64decode, standard_b64encode from dataclasses import InitVar, dataclass, field, replace from pathlib import Path -from typing import TYPE_CHECKING, Literal, TextIO, assert_never, cast, overload +from typing import TYPE_CHECKING, TextIO, assert_never, cast, overload from uuid import UUID import cloudpickle @@ -20,17 +21,77 @@ from . import bindings from ._tracing import call_with_current_trace from .declarations import * -from .declarations import ConstructorDecl, is_callable_decl_constructor +from .declarations import ( + _BUILTIN_EGG_FN_NAMES, + _BUILTIN_EGG_SORT_NAMES, + ConstructorDecl, + is_callable_decl_constructor, +) from .pretty import * from .type_constraint_solver import * if TYPE_CHECKING: - from collections.abc import Callable, Iterable + from collections.abc import Iterable __all__ = ["EGraphState", "span"] _TRACER = trace.get_tracer(__name__) +_VALIDATE_COST_PRIMITIVE = "@validate-dynamic-cost" + +# These heads are interpreted as syntax before Egglog falls back to parsing a +# generic top-level call. Keep this aligned with Parser::parse_command, +# Parser::parse_action, Parser::parse_fact, and the extensions installed by +# egglog_experimental::new_experimental_egraph. +_EGGLOG_RESERVED_CALL_HEADS = frozenset({ + "=", + "birewrite", + "check", + "constructor", + "datatype", + "datatype*", + "delete", + "extract", + "fail", + "for", + "function", + "include", + "input", + "keep-best", + "let", + "let-scheduler", + "multi-extract", + "output", + "panic", + "pop", + "primitive", + "print-function", + "print-size", + "print-stats", + "print-table-stats", + "prove", + "prove-exists", + "push", + "relation", + "rewrite", + "rule", + "ruleset", + "run", + "run-schedule", + "set", + "set-cost", + "sort", + "subsume", + "union", + "unstable-combined-ruleset", + "unstable-fresh!", + "with-dynamic-cost", + "with-ruleset", +}) +_EGGLOG_RESERVED_ACTION_HEADS = frozenset({"delete", "let", "panic", "set", "set-cost", "subsume", "union"}) +_EGGLOG_RESERVED_COMMAND_OR_ACTION_HEADS = _EGGLOG_RESERVED_CALL_HEADS - {"="} +_EGGLOG_LITERAL_NAMES = frozenset({"false", "NaN", "inf", "-inf", "true"}) +_EGGLOG_NUMBER = re.compile(r"[+-]?(?:[0-9]+(?:\.[0-9]*)?|\.[0-9]+)(?:[eE][+-]?[0-9]+)?") @dataclass @@ -45,13 +106,16 @@ class _SavedEgglogFile: cumulative source log. We keep the append handle open for performance and for easy post-failure - inspection. Successful commands are saved normally; failed commands are - saved as expected failures with trailing error comments. + inspection. Successful commands are saved normally; Egglog-reported + failures are saved as expected failures with trailing error comments. + Other execution failures invalidate the transcript because their partial + effects cannot be represented by a replayable command. """ path: str file: TextIO line_count: int = 0 + poisoned: bool = False _finalizer: weakref.finalize = field(init=False, repr=False) def __post_init__(self) -> None: @@ -73,6 +137,16 @@ def _normalize_global_let_name(name: str) -> str: return name if name.startswith("$") else f"${name}" +def _egg_name_is_source_safe_symbol(name: str) -> bool: + """Return whether a name is safe to emit as an ordinary Egglog symbol.""" + return ( + bool(name) + and name != "_" + and not name.startswith(('"', "@")) + and not any(c.isspace() or c in ";()" for c in name) + ) + + def _saved_egglog_failure_message(error: bindings.EggSmolError) -> str: lines = [line.strip() for line in error.context.splitlines() if line.strip()] for message in lines: @@ -85,6 +159,59 @@ def _saved_egglog_failure_message(error: bindings.EggSmolError) -> str: return str(error) +def _rule_variable_names(rule: RuleDecl) -> set[str]: # noqa: C901, PLR0912 + """Collect source-level names that compiler-generated rule lets must not shadow.""" + names: set[str] = set() + + def visit_expr(expr: ExprDecl) -> None: + match expr: + case UnboundVarDecl(name, egg_name): + emitted_name = egg_name or f"_{name}" + names.add(emitted_name.removeprefix("$")) + case LetRefDecl(name): + names.add(_normalize_global_let_name(name).removeprefix("$")) + case CallDecl(_, args) | GetCostDecl(_, args) | PartialCallDecl(CallDecl(_, args)): + for arg in args: + visit_expr(arg.expr) + case LitDecl() | PyObjectDecl() | ValueDecl() | DummyDecl(): + pass + case _: + assert_never(expr) + + for fact in rule.body: + match fact: + case EqDecl(_, left, right): + visit_expr(left) + visit_expr(right) + case ExprFactDecl(typed_expr): + visit_expr(typed_expr.expr) + case _: + assert_never(fact) + for action in rule.head: + match action: + case LetDecl(name, typed_expr): + names.add(_normalize_global_let_name(name).removeprefix("$")) + visit_expr(typed_expr.expr) + case SetDecl(_, call, rhs): + visit_expr(call) + visit_expr(rhs) + case ExprActionDecl(typed_expr): + visit_expr(typed_expr.expr) + case ChangeDecl(_, call, _): + visit_expr(call) + case UnionDecl(_, left, right): + visit_expr(left) + visit_expr(right) + case SetCostDecl(_, call, cost): + visit_expr(call) + visit_expr(cost) + case PanicDecl(): + pass + case _: + assert_never(action) + return names + + def span(frame_index: int = 0) -> bindings.RustSpan: """ Returns a span for the current file and line. @@ -108,6 +235,12 @@ class EGraphState: """ egraph: bindings.EGraph + seminaive: bool = True + # Opaque identity for values created in the current push scope. Descendant + # scopes may use ancestor values, but values from a popped scope must not be + # accepted if the backend later reuses their raw IDs. + value_owner: object = field(default_factory=object, repr=False) + valid_value_owners: frozenset[object] = field(default_factory=frozenset, repr=False) save_egglog_string: InitVar[bool] = False egglog_file_state: _SavedEgglogFile | None = field(default=None, repr=False) # The declarations we have added. @@ -146,6 +279,9 @@ class EGraphState: cost_table_names: dict[CallableRef, str] = field(default_factory=dict) # Counter for deterministic synthetic let bindings created while lowering expressions to egg. expr_to_let_counter: int = 0 + # Explicit top-level lets waiting to be lowered in the current registration batch. + # Synthetic lets must avoid them regardless of command order within that batch. + pending_let_names: frozenset[str] = field(default_factory=frozenset, repr=False) # Counter for deterministic synthetic names assigned to unnamed functions. unnamed_function_counter: int = 0 @@ -155,11 +291,13 @@ class EGraphState: rule_name_to_command_decl: dict[str, RuleDecl | BiRewriteDecl | RewriteDecl] = field(default_factory=dict) def __post_init__(self, save_egglog_string: bool) -> None: + if not self.valid_value_owners: + self.valid_value_owners = frozenset((self.value_owner,)) if save_egglog_string and self.egglog_file_state is None: # Keep one persistent temp `.egg` file per high-level egraph so parse errors # can point at a stable filename the user can open after a failure. egglog_file = tempfile.NamedTemporaryFile( # noqa: SIM115 - kept open for incremental appends - mode="w+", encoding="utf-8", suffix=".egg", delete=False + mode="w+", encoding="utf-8", newline="", suffix=".egg", delete=False ) self.egglog_file_state = _SavedEgglogFile(egglog_file.name, cast("TextIO", egglog_file)) @@ -167,8 +305,12 @@ def copy(self) -> EGraphState: """ Returns a copy of the state. The egraph reference is kept the same. Used for pushing/popping. """ + value_owner = object() return EGraphState( egraph=self.egraph, + seminaive=self.seminaive, + value_owner=value_owner, + valid_value_owners=self.valid_value_owners | {value_owner}, save_egglog_string=self.egglog_file_state is not None, egglog_file_state=self.egglog_file_state, __egg_decls__=self.__egg_decls__.copy(), @@ -183,6 +325,7 @@ def copy(self) -> EGraphState: expr_to_letref_cache=self.expr_to_letref_cache.copy(), cost_table_names=self.cost_table_names.copy(), expr_to_let_counter=self.expr_to_let_counter, + pending_let_names=self.pending_let_names, unnamed_function_counter=self.unnamed_function_counter, rule_name_counter=self.rule_name_counter, rule_name_to_command_decl=self.rule_name_to_command_decl.copy(), @@ -195,18 +338,30 @@ def egglog_string(self) -> str: if self.egglog_file_state.file.closed: msg = "Can't get egglog string after the saved transcript has been closed" raise ValueError(msg) + if self.egglog_file_state.poisoned: + msg = "Can't use the saved Egglog transcript after an execution failure whose partial effects cannot be replayed" + raise RuntimeError(msg) # The append handle stays open for execution, so flush before reading the saved source. self.egglog_file_state.file.flush() - with open(self.egglog_file_state.path, encoding="utf-8") as saved_file: + with open(self.egglog_file_state.path, encoding="utf-8", newline="") as saved_file: return saved_file.read() def close(self) -> None: if self.egglog_file_state is not None: self.egglog_file_state.close() + def ensure_open(self) -> None: + if self.egglog_file_state is not None and self.egglog_file_state.file.closed: + msg = "Cannot run commands after the saved Egglog transcript has been closed" + raise ValueError(msg) + if self.egglog_file_state is not None and self.egglog_file_state.poisoned: + msg = "Cannot run commands after an execution failure whose partial effects cannot be replayed" + raise RuntimeError(msg) + def run_program(self, *commands: bindings._Command) -> list[bindings._CommandOutput]: if not commands: return [] + self.ensure_open() if self.egglog_file_state is None: return call_with_current_trace(self.egraph.run_program, *commands) @@ -223,14 +378,36 @@ def run_program(self, *commands: bindings._Command) -> list[bindings._CommandOut self.egraph.parse_and_run_program, padded_command, filename=self.egglog_file_state.path ) except bindings.EggSmolError as error: + if not error.replayable_by_fail: + # Parsing, expansion, and typechecking can mutate backend + # metadata before they fail, while `(fail ...)` cannot + # reproduce those failures. Do not claim a replayable log. + self.egglog_file_state.poisoned = True + raise fail_command_text = str(bindings.Fail(span(), command)).rstrip("\n") saved_text = f"{fail_command_text} ; {_saved_egglog_failure_message(error)}\n" - self.egglog_file_state.file.write(saved_text) - self.egglog_file_state.file.flush() + try: + self.egglog_file_state.file.write(saved_text) + self.egglog_file_state.file.flush() + except BaseException: + self.egglog_file_state.poisoned = True + raise self.egglog_file_state.line_count += saved_text.count("\n") raise - self.egglog_file_state.file.write(command_text) - self.egglog_file_state.file.flush() + except BaseException: + # A Python primitive or unrelated runtime failure can interrupt + # a command after some actions have committed. There is no + # replayable Egglog command for that partial state. + self.egglog_file_state.poisoned = True + raise + try: + self.egglog_file_state.file.write(command_text) + self.egglog_file_state.file.flush() + except BaseException: + # Execution already succeeded, so omitting even part of this + # command would make the cumulative source diverge. + self.egglog_file_state.poisoned = True + raise self.egglog_file_state.line_count += command_text.count("\n") outputs.extend(command_outputs) return outputs @@ -243,6 +420,22 @@ def _persistent_scheduler_name(scheduler: BackOffDecl) -> str: def _local_scheduler_name(index: int) -> str: return f"_scheduler_{index}" + @staticmethod + def _back_off_scheduler_to_egg(scheduler: BackOffDecl) -> bindings.Call: + """Serialize the shared option protocol for local and persistent backoff schedulers.""" + args: list[bindings._Expr] = [] + if scheduler.match_limit is not None: + args.extend(( + bindings.Var(span(), ":match-limit"), + bindings.Lit(span(), bindings.Int(scheduler.match_limit)), + )) + if scheduler.ban_length is not None: + args.extend(( + bindings.Var(span(), ":ban-length"), + bindings.Lit(span(), bindings.Int(scheduler.ban_length)), + )) + return bindings.Call(span(), "back-off", args) + @_TRACER.start_as_current_span("run_schedule_to_egg") def run_schedule_to_egg(self, schedule: ScheduleDecl) -> bindings._Command: """ @@ -346,26 +539,20 @@ def _schedule_with_scheduler_to_egg( # noqa: C901, PLR0912 """ match schedule: case LetSchedulerDecl(scheduler, inner): - match_limit = scheduler.match_limit - ban_length = scheduler.ban_length name = self._local_scheduler_name(len(bound_schedulers)) bound_schedulers.append(scheduler) - args: list[bindings._Expr] = [] - if match_limit is not None: - args.append(bindings.Var(span(), ":match-limit")) - args.append(bindings.Lit(span(), bindings.Int(match_limit))) - if ban_length is not None: - args.append(bindings.Var(span(), ":ban-length")) - args.append(bindings.Lit(span(), bindings.Int(ban_length))) - back_off_decl = bindings.Call(span(), "back-off", args) - let_decl = bindings.Call(span(), "let-scheduler", [bindings.Var(span(), name), back_off_decl]) + let_decl = bindings.Call( + span(), + "let-scheduler", + [bindings.Var(span(), name), self._back_off_scheduler_to_egg(scheduler)], + ) try: inner_exprs = self._schedule_with_scheduler_to_egg(inner, bound_schedulers) finally: bound_schedulers.pop() return [bindings.Call(span(), "seq", [let_decl, *inner_exprs])] case RunDecl(ruleset_ident, until, scheduler): - args = [bindings.Var(span(), str(ruleset_ident))] + args: list[bindings._Expr] = [bindings.Var(span(), str(ruleset_ident))] if scheduler: name = "run-with" scheduler_name = self._persistent_scheduler_name(scheduler) @@ -411,18 +598,13 @@ def _schedule_with_scheduler_to_egg( # noqa: C901, PLR0912 assert_never(schedule) def _persistent_scheduler_to_egg(self, scheduler: BackOffDecl) -> bindings._Command: - args: list[bindings._Expr] = [] - if scheduler.match_limit is not None: - args.append(bindings.Var(span(), ":match-limit")) - args.append(bindings.Lit(span(), bindings.Int(scheduler.match_limit))) - if scheduler.ban_length is not None: - args.append(bindings.Var(span(), ":ban-length")) - args.append(bindings.Lit(span(), bindings.Int(scheduler.ban_length))) - back_off_decl = bindings.Call(span(), "back-off", args) return bindings.UserDefined( span(), "let-scheduler", - [bindings.Var(span(), self._persistent_scheduler_name(scheduler)), back_off_decl], + [ + bindings.Var(span(), self._persistent_scheduler_name(scheduler)), + self._back_off_scheduler_to_egg(scheduler), + ], ) def ruleset_to_egg(self, ident: Ident) -> None: # noqa: C901 @@ -432,6 +614,15 @@ def ruleset_to_egg(self, ident: Ident) -> None: # noqa: C901 if ident.name == "" and ident not in self.__egg_decls__._rulesets: self.rulesets.setdefault(ident, set()) return + egg_name = str(ident) + if self.egglog_file_state is not None and ( + not _egg_name_is_source_safe_symbol(egg_name) or _egg_name_is_parser_literal(egg_name) + ): + msg = ( + f"Ruleset name {egg_name!r} cannot be used with save_egglog_string=True because " + "it does not serialize as one Egglog symbol" + ) + raise ValueError(msg) match self.__egg_decls__._rulesets[ident]: case RulesetDecl(rules): if ident not in self.rulesets: @@ -443,9 +634,8 @@ def ruleset_to_egg(self, ident: Ident) -> None: # noqa: C901 for rule in rules: if rule in added_rules: continue - cmd = self.command_to_egg(rule, ident) - if cmd is not None: - self.run_program(cmd) + commands = self.commands_to_egg(rule, ident) + self.run_program(*commands) added_rules.add(rule) case CombinedRulesetDecl(rulesets): if ident in self.rulesets: @@ -455,13 +645,15 @@ def ruleset_to_egg(self, ident: Ident) -> None: # noqa: C901 self.ruleset_to_egg(ruleset) self.run_program(bindings.UnstableCombinedRuleset(span(), str(ident), list(map(str, rulesets)))) - def command_to_egg(self, cmd: CommandDecl, ruleset: Ident) -> bindings._Command | None: + def commands_to_egg( # noqa: C901, PLR0912 + self, cmd: CommandDecl, ruleset: Ident + ) -> list[bindings._Command]: match cmd: case ActionCommandDecl(action): - action_egg = self.action_to_egg(action, expr_to_let=True) - if not action_egg: - return None - return bindings.ActionCommand(action_egg) + return [ + bindings.ActionCommand(action_egg) + for action_egg in self.actions_to_egg(action, expr_to_let=True, standalone=True) + ] case RewriteDecl(tp, lhs, rhs, conditions) | BiRewriteDecl(tp, lhs, rhs, conditions): self.type_ref_to_egg(tp) name = str(self.rule_name_counter) @@ -475,51 +667,88 @@ def command_to_egg(self, cmd: CommandDecl, ruleset: Ident) -> bindings._Command ) egg_cmd: bindings._Command if isinstance(cmd, RewriteDecl): - self.rule_name_to_command_decl[name] = cmd egg_cmd = bindings.RewriteCommand(str(ruleset), rewrite, cmd.subsume) + serialized_name = str(egg_cmd) + reported_name = serialized_name.replace('"', "'") + if ( + self.egglog_file_state is not None + and (not self.seminaive or self.egraph.no_decomp()) + and (not cmd.subsume or isinstance(lhs, CallDecl)) + ): + rule = self._rewrite_to_rule_decl(tp, lhs, rhs, conditions, reported_name, cmd.subsume) + commands = self.commands_to_egg(rule, ruleset) + self.rule_name_to_command_decl[reported_name] = cmd + return commands + self.rule_name_to_command_decl[name] = cmd # Saving a transcript executes the serialized command, whose syntax does not # preserve the internal rewrite name. The engine then reports the serialized # rewrite itself as its name, so retain that alias for RunReport translation. - serialized_name = str(egg_cmd) self.rule_name_to_command_decl[serialized_name] = cmd # Backend report names render every Symbol with single quotes, while # command serialization accepts symbols as double-quoted literals. - reported_name = serialized_name.replace('"', "'") self.rule_name_to_command_decl[reported_name] = cmd else: - self.rule_name_to_command_decl[f"{name}=>"] = cmd - self.rule_name_to_command_decl[f"{name}<="] = cmd egg_cmd = bindings.BiRewriteCommand(str(ruleset), rewrite) serialized_name = str(egg_cmd) reported_name = serialized_name.replace('"', "'") + if self.egglog_file_state is not None and (not self.seminaive or self.egraph.no_decomp()): + report_names = (f"{reported_name}=>", f"{reported_name}<=") + rules = ( + self._rewrite_to_rule_decl(tp, lhs, rhs, conditions, report_names[0], False), + self._rewrite_to_rule_decl(tp, rhs, lhs, conditions, report_names[1], False), + ) + commands = [command for rule in rules for command in self.commands_to_egg(rule, ruleset)] + for report_name in report_names: + self.rule_name_to_command_decl[report_name] = cmd + return commands + self.rule_name_to_command_decl[f"{name}=>"] = cmd + self.rule_name_to_command_decl[f"{name}<="] = cmd for suffix in ("=>", "<="): self.rule_name_to_command_decl[f"{serialized_name}{suffix}"] = cmd self.rule_name_to_command_decl[f"{reported_name}{suffix}"] = cmd - return egg_cmd + return [egg_cmd] case RuleDecl(head, body, name, eval_mode, no_decomp): if not name: name = str(self.rule_name_counter) self.rule_name_counter += 1 self.rule_name_to_command_decl[name] = cmd - binding_eval_mode = cast( - "bindings.Seminaive | bindings.Naive | bindings.UnsafeSeminaive", - { - "seminaive": bindings.Seminaive(), - "naive": bindings.Naive(), - "unsafe-seminaive": bindings.UnsafeSeminaive(), - }[eval_mode], - ) - return bindings.RuleCommand( - bindings.Rule( - span(), - [self.action_to_egg(a) for a in head], - [self.fact_to_egg(f, expr_to_let=False) for f in body], - name or "", - str(ruleset), - binding_eval_mode, - no_decomp, + eval_modes = { + "seminaive": bindings.Seminaive(), + "naive": bindings.Naive(), + "unsafe-seminaive": bindings.UnsafeSeminaive(), + } + if eval_mode not in eval_modes: + msg = ( + f"Unknown rule evaluation mode {eval_mode!r}; expected " + "'seminaive', 'naive', or 'unsafe-seminaive'" + ) + raise ValueError(msg) + binding_eval_mode = ( + bindings.Naive() + if not self.seminaive + else cast( + "bindings.Seminaive | bindings.Naive | bindings.UnsafeSeminaive", + eval_modes[eval_mode], ) ) + used_variable_names = _rule_variable_names(cmd) + return [ + bindings.RuleCommand( + bindings.Rule( + span(), + [ + action_egg + for action in head + for action_egg in self.actions_to_egg(action, used_variable_names=used_variable_names) + ], + [self.fact_to_egg(f, expr_to_let=False) for f in body], + name or "", + str(ruleset), + binding_eval_mode, + no_decomp or self.egraph.no_decomp(), + ) + ) + ] case DefaultRewriteDecl(ref, expr, subsume): sig = self.__egg_decls__.get_callable_decl(ref).signature assert isinstance(sig, FunctionSignature) @@ -531,56 +760,93 @@ def command_to_egg(self, cmd: CommandDecl, ruleset: Ident) -> bindings._Command rewrite_decl = RewriteDecl( sig.semantic_return_type.to_just(), CallDecl(ref, arg_mapping), expr, (), subsume ) - return self.command_to_egg(rewrite_decl, ruleset) + return self.commands_to_egg(rewrite_decl, ruleset) case _: assert_never(cmd) - @overload - def action_to_egg(self, action: ActionDecl) -> bindings._Action: ... - - @overload - def action_to_egg( + def _rewrite_to_rule_decl( self, - action: ActionDecl, - expr_to_let: Literal[True] = ..., - ) -> bindings._Action | None: ... - - @overload - def action_to_egg(self, action: ActionDecl, expr_to_let: bool) -> bindings._Action | None: ... + tp: JustTypeRef, + lhs: ExprDecl, + rhs: ExprDecl, + conditions: tuple[FactDecl, ...], + name: str, + subsume: bool, + ) -> RuleDecl: + """Desugar a rewrite while retaining rule options that rewrite syntax cannot encode.""" + used_variable_names = _rule_variable_names(RuleDecl((), (EqDecl(tp, lhs, rhs), *conditions), None)) + fresh_name = self._allocate_synthetic_let_name().removeprefix("$") + while fresh_name in used_variable_names: + fresh_name = self._allocate_synthetic_let_name().removeprefix("$") + fresh = UnboundVarDecl(fresh_name, fresh_name) + head: list[ActionDecl] = [UnionDecl(tp, fresh, rhs)] + if subsume: + if not isinstance(lhs, CallDecl): + msg = "subsumed rewrite must have a function call on the lhs" + raise ValueError(msg) + head.append(ChangeDecl(tp, lhs, "subsume")) + return RuleDecl(tuple(head), (EqDecl(tp, fresh, lhs), *conditions), name) - def action_to_egg( # noqa: C901, PLR0911, PLR0912 + def actions_to_egg( # noqa: C901, PLR0911, PLR0912 self, action: ActionDecl, expr_to_let: bool = False, - ) -> bindings._Action | None: + *, + standalone: bool = False, + used_variable_names: set[str] | None = None, + ) -> list[bindings._Action]: match action: case LetDecl(name, typed_expr): + normalized_name = _normalize_global_let_name(name) + if self.egglog_file_state is not None and not _egg_name_is_source_safe_symbol(normalized_name): + msg = ( + f"Let name {name!r} cannot be used with save_egglog_string=True because " + "it does not serialize as one Egglog symbol" + ) + raise ValueError(msg) var_decl = LetRefDecl(name) var_egg = self._expr_to_egg(var_decl) self.expr_to_egg_cache[var_decl] = var_egg - return bindings.Let( - span(), - var_egg.name, - self.typed_expr_to_egg(typed_expr, expr_to_let=expr_to_let), - ) + return [ + bindings.Let( + span(), + var_egg.name, + self.typed_expr_to_egg(typed_expr, expr_to_let=expr_to_let), + ) + ] case SetDecl(tp, call, rhs): self.type_ref_to_egg(tp) egg_fn, typed_args = self.translate_call(call) - return bindings.Set( - span(), - egg_fn, - [self.typed_expr_to_egg(arg, expr_to_let) for arg in typed_args], - self._expr_to_egg(rhs, expr_to_let=expr_to_let), - ) + return [ + bindings.Set( + span(), + egg_fn, + [self.typed_expr_to_egg(arg, expr_to_let) for arg in typed_args], + self._expr_to_egg(rhs, expr_to_let=expr_to_let), + ) + ] case ExprActionDecl(typed_expr): if not isinstance(typed_expr.expr, CallDecl): msg = "Top-level egglog expr commands must be calls" raise ValueError(msg) # noqa: TRY004 - preserve the public validation error + callable_decl = self.__egg_decls__.get_callable_decl(typed_expr.expr.callable) + if ( + self.egglog_file_state is not None + and callable_decl.egg_name is not None + and callable_decl.egg_name + in (_EGGLOG_RESERVED_COMMAND_OR_ACTION_HEADS if standalone else _EGGLOG_RESERVED_ACTION_HEADS) + ): + context = "top-level action" if standalone else "rule action" + msg = ( + f"Explicit Egglog callable name {callable_decl.egg_name!r} cannot be used as a {context} " + "with save_egglog_string=True because it is parsed as Egglog syntax" + ) + raise ValueError(msg) egg_expr = self.typed_expr_to_egg(typed_expr, expr_to_let=expr_to_let) if isinstance(egg_expr, bindings.Var): - return None + return [] assert isinstance(egg_expr, bindings.Call) - return bindings.Expr_(span(), egg_expr) + return [bindings.Expr_(span(), egg_expr)] case ChangeDecl(tp, call, change): self.type_ref_to_egg(tp) egg_fn, typed_args = self.translate_call(call) @@ -592,26 +858,57 @@ def action_to_egg( # noqa: C901, PLR0911, PLR0912 egg_change = bindings.Subsume() case _: assert_never(change) - return bindings.Change( - span(), - egg_change, - egg_fn, - [self.typed_expr_to_egg(arg, expr_to_let) for arg in typed_args], - ) + return [ + bindings.Change( + span(), + egg_change, + egg_fn, + [self.typed_expr_to_egg(arg, expr_to_let) for arg in typed_args], + ) + ] case UnionDecl(tp, lhs, rhs): self.type_ref_to_egg(tp) - return bindings.Union( - span(), - self._expr_to_egg(lhs, expr_to_let=expr_to_let), - self._expr_to_egg(rhs, expr_to_let=expr_to_let), - ) + return [ + bindings.Union( + span(), + self._expr_to_egg(lhs, expr_to_let=expr_to_let), + self._expr_to_egg(rhs, expr_to_let=expr_to_let), + ) + ] case PanicDecl(name): - return bindings.Panic(span(), name) + return [bindings.Panic(span(), name)] case SetCostDecl(tp, expr, cost): self.type_ref_to_egg(tp) - cost_table, typed_args = self.translate_call(GetCostDecl(expr.callable, expr.args)) - args_egg = [self.typed_expr_to_egg(x, expr_to_let) for x in typed_args] - return bindings.Set(span(), cost_table, args_egg, self._expr_to_egg(cost, expr_to_let=expr_to_let)) + egg_fn, typed_args = self.translate_call(expr) + cost_table = self.create_cost_table(expr.callable) + # Match egglog-experimental's set-cost action macro: bind each + # argument once, materialize the target call, then write its + # validated cost. Structural lowering here avoids evaluating a + # cost expression early through top-level factoring. + lowered: list[bindings._Action] = [] + args_egg: list[bindings._Expr] = [] + for typed_arg in typed_args: + name = self._allocate_synthetic_let_name() + if standalone: + # Mark implementation-only globals so freeze() omits them + # without making later expressions reuse their values. + var_egg = bindings.Var(span(), name) + self.expr_to_letref_cache[LetRefDecl(name)] = var_egg + self.expr_to_egg_cache[LetRefDecl(name)] = var_egg + else: + name = name.removeprefix("$") + while used_variable_names is not None and name in used_variable_names: + name = self._allocate_synthetic_let_name().removeprefix("$") + if used_variable_names is not None: + used_variable_names.add(name) + var_egg = bindings.Var(span(), name) + lowered.append(bindings.Let(span(), name, self.typed_expr_to_egg(typed_arg, False))) + args_egg.append(var_egg) + lowered.append(bindings.Expr_(span(), bindings.Call(span(), egg_fn, args_egg))) + cost_expr = self._expr_to_egg(cost, expr_to_let=False) + validated_cost = bindings.Call(span(), _VALIDATE_COST_PRIMITIVE, [cost_expr]) + lowered.append(bindings.Set(span(), cost_table, args_egg, validated_cost)) + return lowered case _: assert_never(action) @@ -636,21 +933,14 @@ def create_cost_table(self, ref: CallableRef) -> str: # with the protocol's input sorts and i64 output. if not self._has_compatible_cost_table_target(name, target_schema): existing_refs = self.egg_fn_to_callable_refs.get(name, set()) - compatible_raw_table = bool(existing_refs) - for existing_ref in existing_refs: - existing_decl = self.__egg_decls__.get_callable_decl(existing_ref) - if not ( - isinstance(existing_decl, FunctionDecl) - and not existing_decl.builtin - and existing_decl.body is None - and isinstance(existing_decl.signature, FunctionSignature) - ): - compatible_raw_table = False - break - existing_schema = self._signature_to_egg_schema(existing_decl.signature) - if existing_schema.input != schema.input or existing_schema.output != schema.output: - compatible_raw_table = False - break + compatible_raw_table = bool(existing_refs) and all( + self._raw_cost_table_matches_schema(existing_ref, schema) for existing_ref in existing_refs + ) + if compatible_raw_table and not self._callable_is_table_backed(ref): + raise ValueError( + f"Canonical dynamic-cost table {name!r} for an eager or builtin primitive " + "cannot also be a user-declared function" + ) if existing_refs and not compatible_raw_table: msg = ( f"Canonical dynamic-cost table {name!r} is already used by an incompatible callable; " @@ -664,6 +954,33 @@ def create_cost_table(self, ref: CallableRef) -> str: self.cost_table_names[ref] = name return name + def _callable_is_table_backed(self, ref: CallableRef) -> bool: + """Return whether Egglog stores rows for this callable.""" + decl = self.__egg_decls__.get_callable_decl(ref) + match decl: + case RelationDecl() | ConstructorDecl() | ConstantDecl(body=None): + return True + case FunctionDecl(body=None, builtin=False): + return not isinstance(ref, UnnamedFunctionRef) + case ConstantDecl() | FunctionDecl(): + return False + case _: + assert_never(decl) + + def _raw_cost_table_matches_schema(self, ref: CallableRef, schema: bindings.Schema) -> bool: + """Check the complete public declaration contract for a raw dynamic-cost table.""" + decl = self.__egg_decls__.get_callable_decl(ref) + if not ( + isinstance(decl, FunctionDecl) + and not decl.builtin + and decl.body is None + and decl.merge is None + and isinstance(decl.signature, FunctionSignature) + ): + return False + existing_schema = self._signature_to_egg_schema(decl.signature) + return existing_schema.input == schema.input and existing_schema.output == schema.output + def _has_compatible_cost_table_target(self, name: str, target_schema: bindings.Schema) -> bool: """Validate every callable already sharing a canonical dynamic-cost table.""" existing_cost_refs = [ref for ref, existing_name in self.cost_table_names.items() if existing_name == name] @@ -687,6 +1004,20 @@ def fact_to_egg(self, fact: FactDecl, *, expr_to_let: bool = False) -> bindings. self._expr_to_egg(right, expr_to_let=expr_to_let), ) case ExprFactDecl(typed_expr): + if isinstance(typed_expr.expr, CallDecl) and typed_expr.expr.callable != FunctionRef( + Ident.builtin("!=") + ): + callable_decl = self.__egg_decls__.get_callable_decl(typed_expr.expr.callable) + if ( + self.egglog_file_state is not None + and callable_decl.egg_name is not None + and callable_decl.egg_name == "=" + ): + msg = ( + "Explicit Egglog callable name '=' cannot be used as a fact with " + "save_egglog_string=True because it is parsed as equality syntax" + ) + raise ValueError(msg) return bindings.Fact(self.typed_expr_to_egg(typed_expr, expr_to_let=expr_to_let)) case _: assert_never(fact) @@ -700,8 +1031,47 @@ def callable_ref_to_egg(self, ref: CallableRef) -> tuple[str, bool]: # noqa: C9 if ref in self.callable_ref_to_egg_fn: return self.callable_ref_to_egg_fn[ref] decl = self.__egg_decls__.get_callable_decl(ref) - egg_name = decl.egg_name or self._allocate_callable_egg_name(ref) - self.egg_fn_to_callable_refs.setdefault(egg_name, set()).add(ref) + if ( + self.egglog_file_state is not None + and decl.egg_name + and (not _egg_name_is_source_safe_symbol(decl.egg_name) or _egg_name_is_parser_literal(decl.egg_name)) + ): + msg = ( + f"Explicit Egglog callable name {decl.egg_name!r} cannot be used with " + "save_egglog_string=True because it is parsed as Egglog syntax" + ) + raise ValueError(msg) + egg_name = decl.egg_name or self._allocate_name( + self._generate_callable_egg_name(ref), avoid_reserved_call_heads=True + ) + cost_table_targets = tuple( + target_ref for target_ref, cost_table_name in self.cost_table_names.items() if cost_table_name == egg_name + ) + reuse_cost_table = bool(cost_table_targets) + existing_refs = self.egg_fn_to_callable_refs.get(egg_name, set()) + if existing_refs and not reuse_cost_table and not (isinstance(decl, FunctionDecl) and decl.builtin): + msg = f"Explicit Egglog callable name {egg_name!r} is already registered" + raise ValueError(msg) + for target_ref in cost_table_targets: + if not self._callable_is_table_backed(target_ref): + raise ValueError( + f"Canonical dynamic-cost table {egg_name!r} for an eager or builtin primitive " + "cannot also be a user-declared function" + ) + target_signature = self.__egg_decls__.get_callable_decl(target_ref).signature + assert isinstance(target_signature, FunctionSignature) + cost_schema = self._signature_to_egg_schema( + replace(target_signature, return_type=TypeRefWithVars(Ident.builtin("i64"))) + ) + if not self._raw_cost_table_matches_schema(ref, cost_schema): + msg = ( + f"Canonical dynamic-cost table {egg_name!r} is already used by an incompatible callable; " + "it must be a bodyless function with the target's input sorts and i64 output" + ) + raise ValueError(msg) + if reuse_cost_table and self.egg_fn_to_callable_refs.get(egg_name): + msg = f"Canonical dynamic-cost table {egg_name!r} already has a raw callable alias" + raise ValueError(msg) callable_signature = decl.signature reverse_args = callable_signature.reverse_args if isinstance(callable_signature, FunctionSignature) else False match decl: @@ -727,7 +1097,7 @@ def callable_ref_to_egg(self, ref: CallableRef) -> tuple[str, bool]: # noqa: C9 else: self.run_program(bindings.Constructor(span(), egg_name, schema, None, False)) case FunctionDecl(signature=signature, builtin=builtin, body=body, merge=merge): - if not builtin: + if not builtin and not reuse_cost_table: assert isinstance(signature, FunctionSignature), "Cannot turn special function to egg" if body is None and isinstance(ref, UnnamedFunctionRef): body = ref.res @@ -763,6 +1133,10 @@ def callable_ref_to_egg(self, ref: CallableRef) -> tuple[str, bool]: # noqa: C9 ) case _: assert_never(decl) + # Publish the reverse mapping only after any backend declaration has + # succeeded; otherwise a failed registration corrupts extraction and + # freeze by claiming an alias the backend never accepted. + self.egg_fn_to_callable_refs.setdefault(egg_name, set()).add(ref) self.callable_ref_to_egg_fn[ref] = egg_name, reverse_args return egg_name, reverse_args @@ -813,9 +1187,22 @@ def type_ref_to_egg(self, ref: JustTypeRef) -> str: except KeyError: pass decl = self.__egg_decls__._classes[ref.ident] + if ( + self.egglog_file_state is not None + and decl.egg_name + and (not _egg_name_is_source_safe_symbol(decl.egg_name) or _egg_name_is_parser_literal(decl.egg_name)) + ): + msg = ( + f"Explicit Egglog sort name {decl.egg_name!r} cannot be used with " + "save_egglog_string=True because it is not parsed as an Egglog symbol" + ) + raise ValueError(msg) + if not decl.builtin and not ref.args and decl.egg_name and self._backend_symbol_is_occupied(decl.egg_name): + msg = f"Explicit Egglog sort name {decl.egg_name!r} is already registered" + raise ValueError(msg) arg_names = [self.type_ref_to_egg(arg) for arg in ref.args] - self.type_ref_to_egg_sort[ref] = egg_name = (not ref.args and decl.egg_name) or self._allocate_type_egg_name( - ref, decl, arg_names + self.type_ref_to_egg_sort[ref] = egg_name = (not ref.args and decl.egg_name) or self._allocate_name( + self._generate_type_egg_name(ref, decl, arg_names) ) self.egg_sort_to_type_ref[egg_name] = ref @@ -899,16 +1286,17 @@ def _transform_let(self, typed_expr: TypedExprDecl) -> TypedExprDecl | None: self.__egg_decls__, self.__egg_decls__.get_callable_decl(typed_expr.expr.callable) ): return typed_expr + # A synthetic let is a top-level binding, so it cannot capture a rule + # variable (or an invalid top-level unbound variable). Leave those + # expressions inline for their enclosing command to validate instead. + if _contains_unbound_var(typed_expr): + return typed_expr if typed_expr.expr in self.expr_to_letref_cache: return None var_decl = LetRefDecl(self._allocate_synthetic_let_name()) var_egg = self._expr_to_egg(var_decl) cmd = bindings.ActionCommand(bindings.Let(span(), var_egg.name, self.typed_expr_to_egg(typed_expr, True))) - try: - self.run_program(cmd) - # errors when creating let bindings for things like `(vec-empty)` - except bindings.EggSmolError: - return typed_expr + self.run_program(cmd) self.expr_to_letref_cache[typed_expr.expr] = var_egg self.expr_to_egg_cache[var_decl] = var_egg return None @@ -941,7 +1329,19 @@ def _expr_to_egg(self, expr_decl: ExprDecl, *, expr_to_let: bool = False) -> bin case LetRefDecl(name): res = bindings.Var(span(), _normalize_global_let_name(name)) case UnboundVarDecl(name, egg_name): - res = bindings.Var(span(), egg_name or f"_{name}") + emitted_name = egg_name or f"_{name}" + if self.egglog_file_state is not None and ( + emitted_name == "_" + or emitted_name.startswith("@") + or not _egg_name_is_source_safe_symbol(emitted_name) + or _egg_name_is_parser_literal(emitted_name) + ): + msg = ( + f"Egglog variable name {emitted_name!r} cannot be used with " + "save_egglog_string=True because it is not parsed as an Egglog symbol" + ) + raise ValueError(msg) + res = bindings.Var(span(), emitted_name) case LitDecl(value): l: bindings._Literal match value: @@ -1019,93 +1419,123 @@ def _get_possible_types(self, cls_ident: Ident) -> frozenset[JustTypeRef]: """ return frozenset(tp for tp in self.type_ref_to_egg_sort if tp.ident == cls_ident) - 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, ...]: + def _generate_callable_egg_name(self, ref: CallableRef) -> str: """ - Generates short and fully-qualified egg function name candidates for a callable reference. + Generate a fully-qualified Egglog name for a callable reference. + + Qualification separates same-named Python declarations; explicit names + in the same registration batch are reserved by `_allocate_name`. """ match ref: case FunctionRef(ident): - return _name_candidates(ident.name, str(ident), sanitize=True) + return _sanitize_egg_ident(str(ident)) case ConstantRef(ident): # Prefix to avoid name collisions with local vars - return _name_candidates(f"%{ident.name}", f"%{ident}", sanitize=True) + return _sanitize_egg_ident(f"%{ident}") case ( MethodRef(cls_ident, name) | ClassMethodRef(cls_ident, name) | ClassVariableRef(cls_ident, name) | PropertyRef(cls_ident, name) ): - return _name_candidates(f"{cls_ident.name}.{name}", f"{cls_ident}.{name}", sanitize=True) + return _sanitize_egg_ident(f"{cls_ident}.{name}") case InitRef(cls_ident): - return _name_candidates(f"{cls_ident.name}.__init__", f"{cls_ident}.__init__", sanitize=True) + return _sanitize_egg_ident(f"{cls_ident}.__init__") case UnnamedFunctionRef(): name = f"_lambda_{self.unnamed_function_counter}" self.unnamed_function_counter += 1 - return (name,) + return name case _: assert_never(ref) - def _allocate_type_egg_name(self, ref: JustTypeRef, decl: ClassDecl, arg_names: list[str]) -> str: - return self._allocate_name( - self._generate_type_egg_name_candidates(ref, decl, arg_names), self._backend_symbol_is_occupied - ) - def _backend_symbol_is_occupied(self, name: str) -> bool: - """Check Egglog's shared namespace for sorts, tables, and primitives.""" + """Check Egglog's shared namespace and replay-sensitive syntax names.""" return ( - bool(self.egg_fn_to_callable_refs.get(name)) + not _egg_name_is_source_safe_symbol(name) + or _egg_name_is_parser_literal(name) + or bool(self.egg_fn_to_callable_refs.get(name)) or name in self.egg_sort_to_type_ref or name in self.cost_table_names.values() - or name in BUILTIN_EGG_FN_NAMES - or name in BUILTIN_EGG_SORT_NAMES + or name in _BUILTIN_EGG_FN_NAMES + or name in _BUILTIN_EGG_SORT_NAMES ) - def _generate_type_egg_name_candidates( - self, ref: JustTypeRef, decl: ClassDecl, arg_names: list[str] - ) -> tuple[str, ...]: - base_short = decl.egg_name or ref.ident.name - base_full = decl.egg_name or str(ref.ident) + def _generate_type_egg_name(self, ref: JustTypeRef, decl: ClassDecl, arg_names: list[str]) -> str: + if decl.egg_name: + base = decl.egg_name + else: + # Preserve the readable dotted qualification while sanitizing each + # Python identifier component into a source-safe Egglog symbol. + parts = (*ref.ident.module.split("."), ref.ident.name) if ref.ident.module else (ref.ident.name,) + base = ".".join(_sanitize_egg_ident(part) or "_" for part in parts) if not ref.args: - return _name_candidates(base_short, base_full, sanitize=False) + return base args = ",".join(arg_names) - return _name_candidates(f"{base_short}[{args}]", f"{base_full}[{args}]", sanitize=False) + return f"{base}[{args}]" def _allocate_synthetic_let_name(self) -> str: + existing_let_names = self.pending_let_names | { + egg_expr.name + for decl, egg_expr in self.expr_to_egg_cache.items() + if isinstance(decl, LetRefDecl) and isinstance(egg_expr, bindings.Var) + } while True: - name = f"$__expr_{self.expr_to_let_counter}" + candidate = f"$__expr_{self.expr_to_let_counter}" self.expr_to_let_counter += 1 - if name not in { - egg_expr.name - for decl, egg_expr in self.expr_to_egg_cache.items() - if isinstance(decl, LetRefDecl) and isinstance(egg_expr, bindings.Var) - }: + name = self._allocate_name(candidate) + if name not in existing_let_names: return name - @staticmethod - def _allocate_name(candidates: Iterable[str], is_taken: Callable[[str], bool]) -> str: - candidate_list = tuple(dict.fromkeys(candidates)) - for candidate in candidate_list: - if not is_taken(candidate): - return candidate + def _allocate_name(self, candidate: str, *, avoid_reserved_call_heads: bool = False) -> str: + # All declarations for a register(...) batch are merged before any + # command is lowered. Reserve their explicit backend names up front so + # generated names do not depend on action order within that batch. + explicit_backend_names = { + decl.egg_name + for decl in (*self.__egg_decls__._functions.values(), *self.__egg_decls__._constants.values()) + if decl.egg_name is not None + } + for class_decl in self.__egg_decls__._classes.values(): + if class_decl.egg_name is not None: + explicit_backend_names.add(class_decl.egg_name) + class_callables = ( + *class_decl.class_methods.values(), + *class_decl.class_variables.values(), + *class_decl.methods.values(), + *class_decl.properties.values(), + ) + explicit_backend_names.update(decl.egg_name for decl in class_callables if decl.egg_name is not None) + if class_decl.init is not None and class_decl.init.egg_name is not None: + explicit_backend_names.add(class_decl.init.egg_name) + + if ( + candidate not in explicit_backend_names + and not self._backend_symbol_is_occupied(candidate) + and (not avoid_reserved_call_heads or candidate not in _EGGLOG_RESERVED_CALL_HEADS) + ): + return candidate - fallback = candidate_list[-1] index = 1 - while is_taken(f"{fallback}_{index}"): + while ( + f"{candidate}_{index}" in explicit_backend_names + or self._backend_symbol_is_occupied(f"{candidate}_{index}") + or (avoid_reserved_call_heads and f"{candidate}_{index}" in _EGGLOG_RESERVED_CALL_HEADS) + ): index += 1 - return f"{fallback}_{index}" + return f"{candidate}_{index}" def typed_expr_to_value(self, typed_expr: TypedExprDecl) -> bindings.Value: if isinstance(typed_expr.expr, ValueDecl): + if typed_expr.expr.owner not in self.valid_value_owners: + msg = "Cannot use a value that belongs to a different EGraph or inactive push scope" + raise ValueError(msg) return typed_expr.expr.value egg_expr = self.typed_expr_to_egg(typed_expr, False) return call_with_current_trace(self.egraph.eval_expr, egg_expr)[1] def value_to_expr(self, tp: JustTypeRef, value: bindings.Value) -> ExprDecl: # noqa: C901, PLR0911, PLR0912 if tp.ident.module != Ident.builtin("").module: - return ValueDecl(value) + return ValueDecl(value, self.value_owner) match tp.ident.name: # Should match list in egraph bindings @@ -1214,7 +1644,7 @@ def value_to_expr(self, tp: JustTypeRef, value: bindings.Value) -> ExprDecl: # return FromEggState(self, termdag).resolve_term(term, tp).expr case _: # If this is not a builtin type, or we don't know how to convert it, just return as value - return ValueDecl(value) + return ValueDecl(value, self.value_owner) def _unstable_fn_value_to_expr( self, name: str, partial_args: list[bindings.Value], return_tp: JustTypeRef, _arg_types: list[JustTypeRef] @@ -1242,6 +1672,14 @@ def _unstable_fn_value_to_expr( _EGGLOG_INVALID_IDENT = re.compile(r"[^\w\-+*/?!=<>&|^/%]") +def _egg_name_is_parser_literal(name: str) -> bool: + if not name or name in _EGGLOG_LITERAL_NAMES: + return True + # Egglog parses every finite Rust f64 spelling as a literal rather than an + # atom. Callable names are sanitized first, but sort names can retain dots. + return bool(_EGGLOG_NUMBER.fullmatch(name)) and math.isfinite(float(name)) + + def _sanitize_egg_ident(input_string: str) -> str: """ Replaces all invalid characters in an egg identifier with an underscore. @@ -1249,13 +1687,6 @@ def _sanitize_egg_ident(input_string: str) -> str: return _EGGLOG_INVALID_IDENT.sub("_", input_string) -def _name_candidates(short: str, full: str, *, sanitize: bool) -> tuple[str, ...]: - if sanitize: - short = _sanitize_egg_ident(short) - full = _sanitize_egg_ident(full) - return short, full - - def _exprs_multiple_parents(typed_expr: TypedExprDecl) -> list[TypedExprDecl]: """ Returns all expressions that have multiple parents (a list but semantically just an ordered set). @@ -1263,25 +1694,44 @@ def _exprs_multiple_parents(typed_expr: TypedExprDecl) -> list[TypedExprDecl]: parent_counts: dict[TypedExprDecl, int] = {} traversal_order: list[TypedExprDecl] = [] traversed: set[TypedExprDecl] = set() - - def visit(node: TypedExprDecl) -> None: + stack = [typed_expr] + while stack: + node = stack.pop() if node in traversed: - return + continue traversed.add(node) + if node is not typed_expr: + traversal_order.append(node) match node.expr: case CallDecl(args=args) | PartialCallDecl(CallDecl(args=args)): for child in args: parent_counts[child] = parent_counts.get(child, 0) + 1 - if child not in traversed: - traversal_order.append(child) - visit(child) + stack.extend(reversed(args)) case _: pass - - visit(typed_expr) return [node for node in traversal_order if parent_counts[node] > 1] +def _contains_unbound_var(typed_expr: TypedExprDecl) -> bool: + """Check for an unbound variable without recursively hashing a deep expression DAG.""" + seen: set[int] = set() + stack = [typed_expr] + while stack: + node = stack.pop() + node_id = id(node) + if node_id in seen: + continue + seen.add(node_id) + match node.expr: + case UnboundVarDecl(): + return True + case CallDecl(args=args) | PartialCallDecl(CallDecl(args=args)): + stack.extend(args) + case _: + pass + return False + + @dataclass class FromEggState: """ diff --git a/python/egglog/exp/param_eq/domain.py b/python/egglog/exp/param_eq/domain.py index 91e300d9..b7c8e19b 100644 --- a/python/egglog/exp/param_eq/domain.py +++ b/python/egglog/exp/param_eq/domain.py @@ -297,7 +297,10 @@ def _binary_to_containers( # noqa: C901, PLR0911, PLR0912 rhs_mapped = _binary_to_containers(cast("Num", rhs)) lhs_is_polynomial = _is_expr_instance(lhs_mapped, ContainerPolynomial) rhs_is_polynomial = _is_expr_instance(rhs_mapped, ContainerPolynomial) - if lhs_is_polynomial: + if rhs_is_polynomial and not lhs_is_polynomial: + lhs_mapped, rhs_mapped = rhs_mapped, lhs_mapped + lhs_is_polynomial, rhs_is_polynomial = rhs_is_polynomial, lhs_is_polynomial + if lhs_is_polynomial and not rhs_is_polynomial: lhs_poly = cast("ContainerPolynomial", lhs_mapped) match get_callable_args(rhs_mapped, Num): case (f64(scalar),): @@ -306,19 +309,7 @@ def _binary_to_containers( # noqa: C901, PLR0911, PLR0912 ContainerPolynomial.empty(), lhs_poly, ) - if not rhs_is_polynomial: - return _multiply_container_polynomial_by_monomial(lhs_poly, _to_container_mono(rhs_mapped)) - if rhs_is_polynomial: - rhs_poly = cast("ContainerPolynomial", rhs_mapped) - match get_callable_args(lhs_mapped, Num): - case (f64(scalar),): - return map_fold_kv( - lambda result, mono, coef: result.insert(mono, coef * scalar), - ContainerPolynomial.empty(), - rhs_poly, - ) - if not lhs_is_polynomial: - return _multiply_container_polynomial_by_monomial(rhs_poly, _to_container_mono(lhs_mapped)) + return _multiply_container_polynomial_by_monomial(lhs_poly, _to_container_mono(rhs_mapped)) return map_fold_kv( lambda result, term, exponent: catch(lambda: result[term]).match( lambda old_exponent: result.insert(term, old_exponent + exponent), @@ -594,6 +585,7 @@ def _decoded_monomial_cost(mono: ContainerMonomial, children_costs: list[ParamCo Like _decode_container_mono_term. Assumes that if we have an empty numerator we include the 1.0 """ items = list(mono.value.items()) + # Cost-model callbacks receive Map children in Map.value.items() order: key, value, key, value, ... if len(children_costs) != len(items) * 2: msg = f"Expected {len(items) * 2} monomial child costs, got {len(children_costs)}" raise ValueError(msg) @@ -654,6 +646,7 @@ def _decoded_polynomial_cost(poly: ContainerPolynomial, children_costs: list[Par Should correspond to getting the cost from the return value of _decode_container_polynomial """ items = list(poly.value.items()) + # Cost-model callbacks receive Map children in Map.value.items() order: key, value, key, value, ... if len(children_costs) != len(items) * 2: msg = f"Expected {len(items) * 2} polynomial child costs, got {len(children_costs)}" raise ValueError(msg) diff --git a/python/egglog/exp/param_eq/pipeline.py b/python/egglog/exp/param_eq/pipeline.py index 69c9fb2b..4688cf9a 100644 --- a/python/egglog/exp/param_eq/pipeline.py +++ b/python/egglog/exp/param_eq/pipeline.py @@ -7,17 +7,36 @@ import time from collections.abc import Callable, Iterable from dataclasses import dataclass -from typing import Literal +from typing import Literal, TypeVar from egglog import * from .domain import * +# The authors' prototype applies `rewriteTree` at most twice in `FixTree.simplifyE`; +# its Hegg `runEqualitySaturation` stops after 30 inner rounds. MAX_PASSES = 2 -HASKELL_INNER_ITERATION_LIMIT = 30 +MAX_INNER_ITERATIONS = 30 + BACKOFF_MATCH_LIMIT = 1000 BACKOFF_BAN_LENGTH = 30 +_MAP_KEY = TypeVar("_MAP_KEY", bound=BaseExpr) +_MAP_VALUE = TypeVar("_MAP_VALUE", bound=BaseExpr) + + +def _left_biased_map_merge( + left: Map[_MAP_KEY, _MAP_VALUE], right: Map[_MAP_KEY, _MAP_VALUE] +) -> Map[_MAP_KEY, _MAP_VALUE]: + """Merge maps while retaining the left value for every duplicate key.""" + return map_fold_kv( + lambda result, key, value: catch(lambda: result[key]).match( + lambda old_value: result.insert(key, old_value), result.insert(key, value) + ), + left, + right, + ) + # Keep derived map operations as explicitly typed folds in this research # module; only map_fold_kv is a backend primitive and public builtin. @@ -28,13 +47,7 @@ CONSTS = constant( "CONSTS", Map[Num, f64], - merge=lambda left, right: map_fold_kv( - lambda result, key, value: catch(lambda: result[key]).match( - lambda old_value: result.insert(key, old_value), result.insert(key, value) - ), - left, - right, - ), + merge=_left_biased_map_merge, ) # Map a monomial of the form `{polynomial(P): 1}` to one representative `P`. @@ -50,13 +63,7 @@ POLYNOMIAL_MONOMIALS = constant( "POLYNOMIAL_MONOMIALS", Map[ContainerMonomial, ContainerPolynomial], - merge=lambda left, right: map_fold_kv( - lambda result, key, value: catch(lambda: result[key]).match( - lambda old_value: result.insert(key, old_value), result.insert(key, value) - ), - left, - right, - ), + merge=_left_biased_map_merge, ) @@ -80,15 +87,9 @@ def binary_analysis_rules(x: Num, a: f64, b: f64) -> Iterable[RewriteOrRule]: yield rewrite(sqrt(Num(a)), subsume=True).to(Num(a.sqrt()), a >= 0.0, a.sqrt().is_finite()) yield rule(sqrt(Num(a)), a < 0.0).then(panic("Sqrt of negative number")) - # cancellations + # Identities that do not require a nonzero assumption. yield rewrite(x - x, subsume=True).to(Num(0.0)) - yield rewrite(x / x, subsume=True).to(Num(1.0), x != Num(0.0)) - - # multiplicative of inverse - yield rewrite(x * (1 / x), subsume=True).to(Num(1.0), x != Num(0.0)) - yield rewrite(0 * x, subsume=True).to(Num(0.0)) - yield rewrite(0 / x, subsume=True).to(Num(0.0), x != Num(0.0)) @ruleset @@ -208,15 +209,16 @@ def binary_basic_rules(x: Num, y: Num, z: Num, af: f64, bf: f64, cf: f64, df: f6 yield rewrite(x * (y / z)).to((x * y) / z) # no-op yield rewrite((x * y) / z).to(x * (y / z)) # no-op yield rewrite((a * x) * (b * y)).to((a * b) * (x * y)) # no-op - yield rewrite(a * x + b).to(a * (x + b / a)) # no-op - yield rewrite(a * x - b).to(a * (x - b / a)) # no-op - yield rewrite(b - (a * x)).to(a * ((b / a) - x)) # no-op + yield rewrite(a * x + b).to(a * (x + b / a), af != f64(0.0)) # no-op + yield rewrite(a * x - b).to(a * (x - b / a), af != f64(0.0)) # no-op + yield rewrite(b - (a * x)).to(a * ((b / a) - x), af != f64(0.0)) # no-op yield rewrite(a * x + b * y).to( - a * (x + (b / a) * y) + a * (x + (b / a) * y), + af != f64(0.0), ) # factoring out one constant from one term, and dividing the others who have constant terms to compensate - yield rewrite(a * x - b * y).to(a * (x - (b / a) * y)) # same as above - yield rewrite(a * x + b / y).to(a * (x + (b / a) / y)) # same as above - yield rewrite(a * x - b / y).to(a * (x - (b / a) / y)) # same as above + yield rewrite(a * x - b * y).to(a * (x - (b / a) * y), af != f64(0.0)) # same as above + yield rewrite(a * x + b / y).to(a * (x + (b / a) / y), af != f64(0.0)) # same as above + yield rewrite(a * x - b / y).to(a * (x - (b / a) / y), af != f64(0.0)) # same as above yield rewrite(a / (b * x)).to((a / b) / x) # no-op yield rewrite(x / (b * y)).to((1 / b) * x / y) # no-op @@ -225,8 +227,8 @@ def binary_basic_rules(x: Num, y: Num, z: Num, af: f64, bf: f64, cf: f64, df: f6 yield rewrite(b - x / a).to(((b * a) - x) / a) # same as above yield rewrite(x / a + b * y).to((x + (b * a) * y) / a) # same as above yield rewrite(x / a - b * y).to((x - (b * a) * y) / a) # same as above - yield rewrite((b + a * x) / (c + d * y)).to((a / d) * (b / a + x) / (c / d + y)) - yield rewrite((b + x) / (c + d * y)).to((1 / d) * (b + x) / (c / d + y)) + yield rewrite((b + a * x) / (c + d * y)).to((a / d) * (b / a + x) / (c / d + y), af != f64(0.0), df != f64(0.0)) + yield rewrite((b + x) / (c + d * y)).to((1 / d) * (b + x) / (c / d + y), df != f64(0.0)) # identities yield rewrite(0 + x).to(x) @@ -295,6 +297,7 @@ def container_basic_rules( poly2, ).unwrap() ], + coef != f64(0.0), poly2.length() == nonconst_poly.length(), poly1 == map_fold_kv( @@ -519,7 +522,7 @@ def _run_single_pass( n = egraph.let("n", num) current_size = _graph_size(egraph) saturated = False - for _ in range(HASKELL_INNER_ITERATION_LIMIT): + for _ in range(MAX_INNER_ITERATIONS): analysis_report = egraph.run(analysis_schedule) rewrite_report = egraph.run(schedule) current_size = _graph_size(egraph) @@ -530,12 +533,13 @@ def _run_single_pass( return extracted, cost, current_size, saturated -def run_paper_pipeline( +def _run_pipeline( initial: Num, - decode: Callable[[Num], Num] = lambda x: x, - cost_model: CostModel[ParamCost] = param_cost_model, - schedule: Schedule = binary_schedule, - analysis_schedule: Schedule = binary_analysis_schedule, + *, + decode: Callable[[Num], Num], + cost_model: CostModel[ParamCost], + schedule: Schedule, + analysis_schedule: Schedule, ) -> PaperPipelineReport: current, before_cost = EGraph(save_egglog_string=False).extract(initial, include_cost=True, cost_model=cost_model) # get schedule decls so that it's pre-cached @@ -581,9 +585,21 @@ def run_paper_pipeline( ) +def run_paper_pipeline(initial: Num) -> PaperPipelineReport: + """Run the retained paper pipeline with the binary representation.""" + return _run_pipeline( + initial, + decode=lambda num: num, + cost_model=param_cost_model, + schedule=binary_schedule, + analysis_schedule=binary_analysis_schedule, + ) + + def run_paper_pipeline_container(initial: Num) -> PaperPipelineReport: + """Run the retained paper pipeline with the container representation.""" try: - return run_paper_pipeline( + return _run_pipeline( initial, decode=containers_to_binary, cost_model=container_cost_model, diff --git a/python/egglog/type_constraint_solver.py b/python/egglog/type_constraint_solver.py index 532fe439..d8e83b4d 100644 --- a/python/egglog/type_constraint_solver.py +++ b/python/egglog/type_constraint_solver.py @@ -110,11 +110,14 @@ def substitute_typevars_try_function( except TypeConstraintError: if isinstance(tp, TypeVarRef) or tp.ident != Ident.builtin("UnstableFn") or not callable(value): raise - # Probe against an isolated copy of the declarations with no ambient ruleset so any temporary - # unnamed-function rewrites created while inferring types are discarded after the probe. + # Probe against an isolated copy of the declarations with no ambient ruleset so temporary + # unnamed-function declarations and bodies are discarded after type inference. probe_decls = decls().copy() dummy_args = [ - RuntimeExpr.__from_values__(probe_decls, TypedExprDecl(self.substitute_typevars(arg_tp), DummyDecl())) + RuntimeExpr.__from_values__( + probe_decls, + TypedExprDecl(self.substitute_typevars(arg_tp), DummyDecl()), + ) for arg_tp in tp.args[1:] ] try: diff --git a/python/tests/param_eq/test_domain.py b/python/tests/param_eq/test_domain.py index 00e3eb7f..cd3a5974 100644 --- a/python/tests/param_eq/test_domain.py +++ b/python/tests/param_eq/test_domain.py @@ -69,6 +69,36 @@ def test_lowering_distributes_scalar_products_into_a_polynomial() -> None: assert all(all(get_callable_fn(term) != polynomial for term in monomial.value) for monomial in poly.value) +@pytest.mark.parametrize( + "source", + ["(x0 + x1) * 2.3", "2.3 * (x0 + x1)", "(x0 + x1) * x0", "x0 * (x0 + x1)"], +) +def test_lowering_distributes_when_exactly_one_factor_is_polynomial(source: str) -> None: + lowered = EGraph().extract(binary_to_containers(parse_expression(source))) + poly_args = get_callable_args(lowered, polynomial) + assert poly_args is not None + (poly_expr,) = poly_args + poly = EGraph().extract(cast("ContainerPolynomial", poly_expr)) + + assert len(poly.value) == 2 + assert all(all(get_callable_fn(term) != polynomial for term in monomial.value) for monomial in poly.value) + decoded = render_num(containers_to_binary(lowered)) + assert math.isclose(evaluate(decoded, x0=1.25, x1=-0.5), evaluate(source, x0=1.25, x1=-0.5)) + + +def test_lowering_keeps_polynomial_by_polynomial_products_nested() -> None: + lowered = EGraph().extract(binary_to_containers(parse_expression("(x0 + x1) * (x0 + x1)"))) + poly_args = get_callable_args(lowered, polynomial) + assert poly_args is not None + (poly_expr,) = poly_args + poly = EGraph().extract(cast("ContainerPolynomial", poly_expr)) + (monomial,) = poly.value + ((term, exponent),) = monomial.value.items() + + assert get_callable_fn(term) == polynomial + assert exponent.value == Fraction(2, 1) + + def test_parser_rejects_unsupported_calls() -> None: with pytest.raises(ValueError, match="Unsupported function call"): parse_expression("sin(x0)") diff --git a/python/tests/param_eq/test_pipeline.py b/python/tests/param_eq/test_pipeline.py index 73a4ed69..42e02d27 100644 --- a/python/tests/param_eq/test_pipeline.py +++ b/python/tests/param_eq/test_pipeline.py @@ -2,10 +2,12 @@ import json import math +from collections.abc import Callable +from inspect import signature import pytest -from egglog import EGraph, back_off, eq, rewrite, ruleset, run, var +from egglog import EGraph, back_off, eq, rewrite, ruleset, run, union, var from egglog.exp.param_eq import ( DEMO_CASES, DemoCase, @@ -17,7 +19,12 @@ run_paper_pipeline_container, ) from egglog.exp.param_eq.__main__ import main -from egglog.exp.param_eq.pipeline import container_schedule, containers_analysis_schedule +from egglog.exp.param_eq.pipeline import ( + binary_basic_rules, + container_basic_rules, + container_schedule, + containers_analysis_schedule, +) from .evaluation import evaluate @@ -31,8 +38,6 @@ def test_public_end_to_end_cases(case: DemoCase, variant: str) -> None: if variant == "binary" else run_paper_pipeline_container(binary_to_containers(parse_expression(case.source))) ) - expected_status = "iteration_limit" if (case.name, variant) == ("repeated_monomial", "binary") else "saturated" - assert report.status == expected_status assert 1 <= report.passes <= 2 assert report.extracted_params <= report.before_params assert report.extracted_params < report.before_params or report.extracted_nodes < report.before_nodes @@ -44,6 +49,14 @@ def test_public_end_to_end_cases(case: DemoCase, variant: str) -> None: assert math.isfinite(expected) assert math.isfinite(actual) assert math.isclose(actual, expected, rel_tol=1e-9, abs_tol=1e-9) + if (case.name, variant) == ("repeated_monomial", "binary") and report.status == "iteration_limit": + pytest.xfail("binary repeated-monomial case still reaches the retained inner iteration limit") + assert report.status == "saturated" + + +def test_public_pipeline_runners_only_accept_initial() -> None: + assert tuple(signature(run_paper_pipeline).parameters) == ("initial",) + assert tuple(signature(run_paper_pipeline_container).parameters) == ("initial",) @pytest.mark.param_eq_smoke @@ -90,6 +103,72 @@ def test_nonfinite_constant_results_are_not_folded(source: str) -> None: assert parse_expression(report.extracted) == parse_expression(source) +@pytest.mark.param_eq_smoke +@pytest.mark.parametrize( + ("source", "invalid_rewrite"), + [ + ("0.0 * x0 + 2.0", "0.0 * (x0 + 2.0 / 0.0)"), + ("0.0 * x0 - 2.0", "0.0 * (x0 - 2.0 / 0.0)"), + ("2.0 - 0.0 * x0", "0.0 * (2.0 / 0.0 - x0)"), + ("0.0 * x0 + 2.0 * x1", "0.0 * (x0 + (2.0 / 0.0) * x1)"), + ("0.0 * x0 - 2.0 * x1", "0.0 * (x0 - (2.0 / 0.0) * x1)"), + ("0.0 * x0 + 2.0 / x1", "0.0 * (x0 + (2.0 / 0.0) / x1)"), + ("0.0 * x0 - 2.0 / x1", "0.0 * (x0 - (2.0 / 0.0) / x1)"), + ( + "(2.0 + 0.0 * x0) / (3.0 + 4.0 * x1)", + "(0.0 / 4.0) * (2.0 / 0.0 + x0) / (3.0 / 4.0 + x1)", + ), + ( + "(2.0 + 5.0 * x0) / (3.0 + 0.0 * x1)", + "(5.0 / 0.0) * (2.0 / 5.0 + x0) / (3.0 / 0.0 + x1)", + ), + ( + "(2.0 + x0) / (3.0 + 0.0 * x1)", + "(1.0 / 0.0) * (2.0 + x0) / (3.0 / 0.0 + x1)", + ), + ], +) +def test_binary_factoring_does_not_introduce_zero_denominators(source: str, invalid_rewrite: str) -> None: + source_expr = parse_expression(source) + egraph = EGraph(source_expr, save_egglog_string=False) + + egraph.run(1, ruleset=binary_basic_rules) + + assert not egraph.check_bool(eq(source_expr).to(parse_expression(invalid_rewrite))) + + +@pytest.mark.param_eq_smoke +@pytest.mark.parametrize( + ("expression", "invalid_result"), + [ + (lambda x: x / x, Num(1.0)), + (lambda x: x * (1 / x), Num(1.0)), + (lambda x: Num(0.0) / x, Num(0.0)), + ], +) +def test_symbolic_cancellations_do_not_survive_a_zero_merge( + expression: Callable[[Num], Num], invalid_result: Num +) -> None: + x = Num.var("x0") + source = expression(x) + egraph = EGraph(source, Num(0.0), save_egglog_string=False) + + egraph.run(1, ruleset=pipeline.binary_analysis_rules) + egraph.register(union(x).with_(Num(0.0))) + + assert not egraph.check_bool(eq(source).to(invalid_result)) + + +@pytest.mark.param_eq_smoke +def test_container_factoring_skips_a_zero_coefficient() -> None: + source = binary_to_containers(parse_expression("0.0 * x0 + 2.0")) + egraph = EGraph(source, save_egglog_string=False) + + report = egraph.run(1, ruleset=container_basic_rules) + + assert not report.updated + + @pytest.mark.param_eq_smoke @pytest.mark.parametrize(("source", "expected"), [("exp(1.0)", math.e), ("sqrt(4.0)", 2.0)]) def test_finite_constant_results_are_folded(source: str, expected: float) -> None: @@ -116,22 +195,22 @@ def test_iteration_limited_pass_is_not_reported_as_saturated(monkeypatch: pytest limited_rules, scheduler=back_off(match_limit=0, ban_length=100).persistent(), ) - monkeypatch.setattr(pipeline, "HASKELL_INNER_ITERATION_LIMIT", 1) + monkeypatch.setattr(pipeline, "MAX_INNER_ITERATIONS", 1) + monkeypatch.setattr(pipeline, "binary_schedule", limited_schedule) + monkeypatch.setattr(pipeline, "binary_analysis_schedule", empty_analysis) - report = run_paper_pipeline( - parse_expression("x0 + 0.0"), - schedule=limited_schedule, - analysis_schedule=empty_analysis, - ) + report = run_paper_pipeline(parse_expression("x0 + 0.0")) assert report.status == "iteration_limit" @pytest.mark.param_eq_smoke -def test_cli_emits_complete_json(capsys: pytest.CaptureFixture[str]) -> None: - assert main(["--variant", "binary", "--expr", DEMO_CASES[2].source]) == 0 +@pytest.mark.parametrize("variant", ["binary", "container"]) +def test_cli_emits_complete_json(capsys: pytest.CaptureFixture[str], variant: str) -> None: + source = "2.3 * (3.7*x0 + 5.1*x1) / 7.9" + assert main(["--variant", variant, "--expr", source]) == 0 payload = json.loads(capsys.readouterr().out) - assert payload["variant"] == "binary" + assert payload["variant"] == variant assert payload["status"] == "saturated" assert payload["extracted"] assert payload["extracted_params"] <= payload["before_params"] diff --git a/python/tests/param_eq/test_research_harness.py b/python/tests/param_eq/test_research_harness.py index 8a233560..a8bbb8e9 100644 --- a/python/tests/param_eq/test_research_harness.py +++ b/python/tests/param_eq/test_research_harness.py @@ -144,6 +144,8 @@ def test_haskell_program_forces_results_and_parser_keeps_expression_text_out() - source_n_rank=1.0, ) program = _build_haskell_program() + # CI does not run the generated Haskell, so inspect its source to guard + # explicit forcing and expression-free metadata lookup. assert "beforeNodes <- evaluate (countNodes expr)" in program assert "beforeParams <- evaluate (recountParams (replaceConstsWithParams expr))" in program assert "afterNodes <- evaluate (countNodes simplified)" in program diff --git a/python/tests/test_bindings.py b/python/tests/test_bindings.py index 725eafb6..e74475aa 100644 --- a/python/tests/test_bindings.py +++ b/python/tests/test_bindings.py @@ -196,9 +196,34 @@ def test_parse_and_run_program_exception(self): with pytest.raises( EggSmolError, match="to have type", - ): + ) as exc_info: egraph.run_program(*egraph.parse_program(program)) + assert not exc_info.value.replayable_by_fail + + @pytest.mark.parametrize( + "program", + [ + pytest.param("(check (= 1 2))", id="check"), + pytest.param('(panic "expected")', id="action"), + ], + ) + def test_runtime_command_error_is_replayable_by_fail(self, program: str): + with pytest.raises(EggSmolError) as exc_info: + EGraph().parse_and_run_program(program) + + assert exc_info.value.replayable_by_fail + + def test_parse_error_is_not_replayable_by_fail(self): + with pytest.raises(EggSmolError) as exc_info: + EGraph().parse_and_run_program("(") + + assert not exc_info.value.replayable_by_fail + + def test_egglog_error_constructor_defaults_to_non_replayable(self): + assert not EggSmolError("expected").replayable_by_fail + assert EggSmolError("expected", True).replayable_by_fail + def test_parse_and_run_program_error_keeps_recording_transactional(self): program = """(function f (i64) i64 :no-merge) (set (f 1) 2) @@ -455,6 +480,27 @@ def test_extract_value_reports_extraction_failure(self): with pytest.raises(EggSmolError, match="Unable to find any valid extraction"): egraph.extract_value(value, sort) + def test_tree_extractor_extract_variants(self): + egraph = EGraph() + egraph.parse_and_run_program( + "(datatype Expr (Num i64)) (let root (Num 1)) (union root (Num 2)) (union root (Num 3))" + ) + sort, value = egraph.eval_expr(Call(DUMMY_SPAN, "Num", [Lit(DUMMY_SPAN, Int(1))])) + model = CostModel( + lambda _name, annotation, children: annotation + sum(children), + lambda _name, _args: 1, + lambda _name, _value, children: sum(children), + lambda _name, _value: 1, + ) + extractor = Extractor([sort], egraph, model) + termdag = TermDag() + + variants = extractor.extract_variants(egraph, termdag, value, 2, sort) + + assert len(variants) == 2 + assert {termdag.to_string(term) for _cost, term in variants} <= {"(Num 1)", "(Num 2)", "(Num 3)"} + assert all(cost == 2 for cost, _term in variants) + @pytest.mark.parametrize("extractor", ["tree", "greedy-dag"]) def test_dag_cost_model_batch_extraction(self, extractor): egraph = EGraph() @@ -466,6 +512,7 @@ def test_dag_cost_model_batch_extraction(self, extractor): lambda name, value: 0, lambda name, value: 1, ) + assert str(model).endswith(")") termdag, best = extract_best_with_dag_cost_model(egraph, [(sort, value)], model, extractor=extractor) assert best[0] is not None diff --git a/python/tests/test_egraph_state.py b/python/tests/test_egraph_state.py new file mode 100644 index 00000000..9ac8f4f5 --- /dev/null +++ b/python/tests/test_egraph_state.py @@ -0,0 +1,379 @@ +from __future__ import annotations + +import gc +import pathlib +from typing import Literal, TextIO, cast + +import pytest + +import egglog.bindings as egg_bindings +from egglog import ( + EggSmolError, + EGraph, + Map, + PyObject, + eq, + f64, + i64, + map_fold_kv, + relation, + rule, + ruleset, + run, + set_current_ruleset, + var, +) +from egglog.declarations import ( + ClassDecl, + Declarations, + FunctionDecl, + FunctionRef, + FunctionSignature, + HasDeclarations, + Ident, + JustTypeRef, + LitDecl, + TypedExprDecl, + TypeRefWithVars, +) +from egglog.egraph import get_current_ruleset +from egglog.runtime import RuntimeExpr + + +def test_saved_egglog_transcript_close_removes_backing_file() -> None: + egraph = EGraph(save_egglog_string=True) + assert egraph._state.egglog_file_state is not None + path = pathlib.Path(egraph._state.egglog_file_state.path) + + egraph.close() + + assert not path.exists() + + +def test_saved_egglog_transcript_is_removed_on_finalization() -> None: + egraph = EGraph(save_egglog_string=True) + assert egraph._state.egglog_file_state is not None + path = pathlib.Path(egraph._state.egglog_file_state.path) + assert path.exists() + + del egraph + gc.collect() + + assert not path.exists() + + +def test_closing_saved_transcript_inside_context_restores_scope() -> None: + egraph = EGraph(save_egglog_string=True) + parent_state = egraph._state + + with egraph: + egraph.close() + + assert egraph._state is parent_state + assert not egraph._state_stack + + +def test_file_backed_errors_report_saved_file_line() -> None: + egraph = EGraph(save_egglog_string=True) + egraph.let("x", i64(1)) + egraph.let("y", i64(2)) + expected_line = len(egraph.as_egglog_string.splitlines()) + 1 + assert egraph._state.egglog_file_state is not None + path = egraph._state.egglog_file_state.path + + with pytest.raises(EggSmolError) as exc_info: + egraph.check(eq(i64(1)).to(i64(2))) + + error_text = exc_info.value.context + assert exc_info.value.replayable_by_fail + assert path in error_text + assert f"In {expected_line}:" in error_text + lines = egraph.as_egglog_string.splitlines() + assert "(fail (check (= 1 2))) ; Check failed:" in lines + assert "(check (= 1 2))" not in lines + + +def test_non_replayable_egglog_error_invalidates_saved_transcript() -> None: + egraph = EGraph(save_egglog_string=True) + command = egg_bindings.Check( + egg_bindings.RustSpan(__name__, 0, 0), + [ + egg_bindings.Eq( + egg_bindings.RustSpan(__name__, 0, 0), + egg_bindings.Lit(egg_bindings.RustSpan(__name__, 0, 0), egg_bindings.Int(1)), + egg_bindings.Lit(egg_bindings.RustSpan(__name__, 0, 0), egg_bindings.Float(1.0)), + ) + ], + ) + + with pytest.raises(EggSmolError) as exc_info: + egraph._state.run_program(command) + + assert not exc_info.value.replayable_by_fail + with pytest.raises(RuntimeError, match="partial effects cannot be replayed"): + _ = egraph.as_egglog_string + with pytest.raises(RuntimeError, match="partial effects cannot be replayed"): + egraph.register(relation("after_non_replayable_failure")()) + + +def _raise_after_partial_write(_: object) -> object: + message = "callback failed" + raise ValueError(message) + + +def test_non_egglog_failure_invalidates_saved_transcript() -> None: + trigger = relation("transcript_failure_trigger", i64) + done = relation("transcript_failure_done", i64) + x = var("x", i64) + failing_rules = ruleset( + rule(trigger(x)).then(done(x), PyObject(_raise_after_partial_write)(PyObject(None))), + name="transcript_failure_rules", + ) + + unrecorded = EGraph(trigger(i64(1))) + with pytest.raises(ValueError, match="callback failed"): + unrecorded.run(run(failing_rules)) + unrecorded.check(done(i64(1))) + + recorded = EGraph(trigger(i64(1)), save_egglog_string=True) + with pytest.raises(ValueError, match="callback failed"): + recorded.run(run(failing_rules)) + with pytest.raises(RuntimeError, match="partial effects cannot be replayed"): + recorded.check(done(i64(1))) + with pytest.raises(RuntimeError, match="partial effects cannot be replayed"): + _ = recorded.as_egglog_string + recorded.close() + + scoped = EGraph(trigger(i64(1)), save_egglog_string=True) + parent_state = scoped._state + with pytest.raises(ValueError, match="callback failed"), scoped: + scoped.run(run(failing_rules)) + assert scoped._state is parent_state + assert not scoped._state_stack + with pytest.raises(RuntimeError, match="partial effects cannot be replayed"): + scoped.check(done(i64(1))) + scoped.close() + + +class _FaultingWriter: + def __init__(self, file: TextIO, fail_on: Literal["write", "flush"]) -> None: + self.file = file + self.fail_on = fail_on + + @property + def closed(self) -> bool: + return self.file.closed + + def write(self, text: str) -> int: + if self.fail_on == "write": + message = "transcript write failed" + raise OSError(message) + return self.file.write(text) + + def flush(self) -> None: + if self.fail_on == "flush": + message = "transcript flush failed" + raise OSError(message) + self.file.flush() + + +@pytest.mark.parametrize("fail_on", ["write", "flush"]) +@pytest.mark.parametrize("command_succeeds", [False, True], ids=["egglog-error", "success"]) +def test_saved_transcript_io_failure_invalidates_transcript( + fail_on: Literal["write", "flush"], *, command_succeeds: bool +) -> None: + marker = relation(f"transcript_io_{fail_on}_{command_succeeds}", i64) + egraph = EGraph(marker(i64(0)), save_egglog_string=True) + assert egraph._state.egglog_file_state is not None + file = egraph._state.egglog_file_state.file + egraph._state.egglog_file_state.file = cast("TextIO", _FaultingWriter(file, fail_on)) + + def execute_command() -> None: + if command_succeeds: + egraph.register(marker(i64(1))) + else: + egraph.check(eq(i64(1)).to(i64(2))) + + with pytest.raises(OSError, match=f"transcript {fail_on} failed"): + execute_command() + with pytest.raises(RuntimeError, match="partial effects cannot be replayed"): + _ = egraph.as_egglog_string + egraph.close() + + +def test_higher_order_callable_inference_does_not_mutate_ambient_ruleset() -> None: + ambient = ruleset(name="hof-inference-ambient") + initial_rules = tuple(ambient.__egg_ruleset__.rules) + + with set_current_ruleset(ambient): + initial: Map[i64, f64] = Map[i64, f64].empty() + expr = map_fold_kv( + lambda result, key, value: result.insert(key, -value), + initial, + Map[i64, f64].empty().insert(i64(1), f64(2.0)), + ) + _ = cast("RuntimeExpr", expr).__egg_decls__ + + assert tuple(ambient.__egg_ruleset__.rules) == initial_rules + + +def test_set_current_ruleset_restores_nested_contexts() -> None: + outer = ruleset(name="current-ruleset-outer") + inner = ruleset(name="current-ruleset-inner") + initial = get_current_ruleset() + + with set_current_ruleset(outer): + assert get_current_ruleset() is outer + with set_current_ruleset(inner): + assert get_current_ruleset() is inner + assert get_current_ruleset() is outer + + assert get_current_ruleset() is initial + + +def test_generated_names_are_fully_qualified() -> None: + state = EGraph(save_egglog_string=True)._state + ret1 = Ident("Ret", "pkg.one") + ret2 = Ident("Ret", "pkg.two") + fn1 = Ident("make", "pkg.one") + fn2 = Ident("make", "pkg.two") + state.__egg_decls__ |= Declarations( + _classes={ret1: ClassDecl(), ret2: ClassDecl()}, + _functions={ + fn1: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret1))), + fn2: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret2))), + }, + ) + + assert state.callable_ref_to_egg(FunctionRef(fn1))[0] == "pkg_one_make" + assert state.callable_ref_to_egg(FunctionRef(fn2))[0] == "pkg_two_make" + assert state.type_ref_to_egg(JustTypeRef(ret1)) == "pkg.one.Ret" + assert state.type_ref_to_egg(JustTypeRef(ret2)) == "pkg.two.Ret" + + +def test_missing_function_lookup_does_not_reserve_generated_name() -> None: + state = EGraph(save_egglog_string=True)._state + ret = Ident("LookupRet", "pkg.lookup") + fn = Ident("lookup_short_name", "pkg.lookup") + state.__egg_decls__ |= Declarations( + _classes={ret: ClassDecl()}, + _functions={fn: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret)))}, + ) + + qualified_name = "pkg_lookup_lookup_short_name" + assert list(state.possible_egglog_functions([qualified_name])) == [] + assert state.callable_ref_to_egg(FunctionRef(fn))[0] == qualified_name + + +def test_generated_names_fall_back_from_builtin_names() -> None: + state = EGraph(save_egglog_string=True)._state + ret = Ident("BuiltinConflictRet", "pkg.builtin_conflict") + fn = Ident("exp", "pkg.builtin_conflict") + sort = Ident("Map", "pkg.builtin_conflict") + state.__egg_decls__ |= Declarations( + _classes={ret: ClassDecl(), sort: ClassDecl()}, + _functions={fn: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret)))}, + ) + + assert state.callable_ref_to_egg(FunctionRef(fn))[0] == "pkg_builtin_conflict_exp" + assert state.type_ref_to_egg(JustTypeRef(sort)) == "pkg.builtin_conflict.Map" + + +def test_generated_callable_name_avoids_an_existing_cost_table() -> None: + state = EGraph(save_egglog_string=True)._state + state.__egg_decls__ |= cast("HasDeclarations", i64) + ret = Ident("CostRet", "pkg.cost") + fn = Ident("f", "pkg.cost") + conflict = Ident("cost_table_pkg_cost_f") + state.__egg_decls__ |= Declarations( + _classes={ret: ClassDecl()}, + _functions={ + fn: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret))), + conflict: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret))), + }, + ) + + assert state.create_cost_table(FunctionRef(fn)) == "cost_table_pkg_cost_f" + assert state.callable_ref_to_egg(FunctionRef(conflict))[0] == "cost_table_pkg_cost_f_1" + + +def test_canonical_cost_table_rejects_an_incompatible_callable() -> None: + state = EGraph(save_egglog_string=True)._state + state.__egg_decls__ |= cast("HasDeclarations", i64) + ret = Ident("CostRet", "pkg.cost") + fn = Ident("f", "pkg.cost") + conflict = Ident("raw_cost", "pkg.cost") + state.__egg_decls__ |= Declarations( + _classes={ret: ClassDecl()}, + _functions={ + fn: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret))), + conflict: FunctionDecl( + signature=FunctionSignature(return_type=TypeRefWithVars(ret)), + egg_name="cost_table_pkg_cost_f", + ), + }, + ) + + assert state.callable_ref_to_egg(FunctionRef(conflict))[0] == "cost_table_pkg_cost_f" + with pytest.raises(ValueError, match="already used by an incompatible callable"): + state.create_cost_table(FunctionRef(fn)) + + +def test_canonical_cost_table_reuses_a_compatible_raw_table() -> None: + state = EGraph(save_egglog_string=True)._state + state.__egg_decls__ |= cast("HasDeclarations", i64) + ret = Ident("CostRet", "pkg.cost") + fn = Ident("f", "pkg.cost") + raw_cost = Ident("raw_cost", "pkg.cost") + state.__egg_decls__ |= Declarations( + _classes={ret: ClassDecl()}, + _functions={ + fn: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret))), + raw_cost: FunctionDecl( + signature=FunctionSignature(return_type=TypeRefWithVars(Ident.builtin("i64"))), + egg_name="cost_table_pkg_cost_f", + ), + }, + ) + + assert state.callable_ref_to_egg(FunctionRef(raw_cost))[0] == "cost_table_pkg_cost_f" + assert state.create_cost_table(FunctionRef(fn)) == "cost_table_pkg_cost_f" + assert state.cost_table_names[FunctionRef(fn)] == "cost_table_pkg_cost_f" + + +@pytest.mark.parametrize( + "body", + [ + pytest.param(None, id="bodyless-function"), + pytest.param( + TypedExprDecl(JustTypeRef(Ident.builtin("i64")), LitDecl(1)), + id="eager-primitive", + ), + ], +) +@pytest.mark.parametrize("sort_first", [True, False], ids=["sort-first", "callable-first"]) +def test_generated_names_share_the_backend_sort_and_callable_namespace( + body: TypedExprDecl | None, *, sort_first: bool +) -> None: + state = EGraph(save_egglog_string=True)._state + state.__egg_decls__ |= cast("HasDeclarations", i64) + sort_ident = Ident("Node") + fn_ident = Ident("Node") + sort_ref = JustTypeRef(sort_ident) + fn_ref = FunctionRef(fn_ident) + state.__egg_decls__ |= Declarations( + _classes={sort_ident: ClassDecl()}, + _functions={ + fn_ident: FunctionDecl( + signature=FunctionSignature(return_type=TypeRefWithVars(Ident.builtin("i64"))), + body=body, + ) + }, + ) + + if sort_first: + assert state.type_ref_to_egg(sort_ref) == "Node" + assert state.callable_ref_to_egg(fn_ref)[0] == "Node_1" + else: + assert state.callable_ref_to_egg(fn_ref)[0] == "Node" + assert state.type_ref_to_egg(sort_ref) == "Node_1" diff --git a/python/tests/test_high_level.py b/python/tests/test_high_level.py index d2616465..83e641cf 100644 --- a/python/tests/test_high_level.py +++ b/python/tests/test_high_level.py @@ -1,7 +1,6 @@ # mypy: disable-error-code="empty-body" from __future__ import annotations -import gc import importlib import math import pathlib @@ -15,29 +14,23 @@ import pytest +import egglog.bindings as egg_bindings import egglog.builtins as egg_builtins from egglog import * from egglog.declarations import ( - BUILTIN_EGG_FN_NAMES, - BUILTIN_EGG_SORT_NAMES, - CallableDecl, + BiRewriteDecl, CallDecl, - ClassDecl, - Declarations, - FunctionDecl, FunctionRef, - FunctionSignature, - HasDeclarations, Ident, JustTypeRef, - LitDecl, MethodRef, + RewriteDecl, TypedExprDecl, - TypeRefWithVars, ) -from egglog.egraph import get_current_ruleset from egglog.runtime import RuntimeExpr, RuntimeFunction +_BuiltinExprT = TypeVar("_BuiltinExprT", bound=BaseExpr) + class TestExprStr: def test_unwrap_lit(self): @@ -62,6 +55,16 @@ def test_rule_eval_mode(eval_mode: RuleEvalMode) -> None: egraph.check(rel(i64(1))) +@pytest.mark.parametrize("seminaive", [False, True]) +def test_rule_eval_mode_rejects_unknown_value(seminaive: bool) -> None: + rel = relation(f"invalid_eval_mode_{seminaive}", i64) + x = var("x", i64) + invalid = cast("RuleEvalMode", "unknown") + + with pytest.raises(ValueError, match="Unknown rule evaluation mode"): + EGraph(seminaive=seminaive).register(rule(rel(x), eval_mode=invalid).then(rel(x))) + + def test_per_egraph_configuration() -> None: egraph = EGraph(num_threads=2, no_decomp=True) @@ -73,16 +76,311 @@ def test_per_egraph_configuration() -> None: assert not egraph.no_decomp() -def test_rule_no_decomp_reaches_backend() -> None: - rel = relation("no_decomp_rel", i64) +@pytest.mark.parametrize("use_setter", [False, True], ids=["constructor", "setter"]) +def test_zero_threads_uses_available_parallelism(*, use_setter: bool) -> None: + egraph = EGraph(num_threads=1 if use_setter else 0) + if use_setter: + egraph.set_num_threads(0) + + assert egraph.num_threads() >= 1 + + +@pytest.mark.parametrize("use_setter", [False, True], ids=["constructor", "setter"]) +def test_egraph_no_decomp_reaches_saved_rules(*, use_setter: bool) -> None: + rel = relation(f"no_decomp_rel_{use_setter}", i64) x = var("x", i64) - egraph = EGraph(save_egglog_string=True) + egraph = EGraph(save_egglog_string=True, no_decomp=not use_setter) + if use_setter: + egraph.set_no_decomp(True) - egraph.register(rule(rel(x), no_decomp=True).then(rel(x + 1))) + egraph.register(rule(rel(x)).then(rel(x + 1))) assert ":no-decomp" in egraph.as_egglog_string +def test_global_naive_mode_reaches_saved_rules() -> None: + rel = relation("global_naive_saved_rule", i64) + x = var("x", i64) + egraph = EGraph(save_egglog_string=True, seminaive=False) + + egraph.register(rule(rel(x)).then(rel(x + 1))) + + assert ":naive" in egraph.as_egglog_string + + +@pytest.mark.parametrize("bidirectional", [False, True], ids=["rewrite", "birewrite"]) +def test_global_rule_options_reach_saved_rewrites(*, bidirectional: bool) -> None: + class RewriteConfigExpr(Expr): + def __init__(self, value: i64Like) -> None: ... + + @function + def rewrite_config_lookup(value: i64) -> i64: ... + + @function + def rewrite_config_read(value: i64) -> i64: + return rewrite_config_lookup(value) + + grounded = relation("rewrite_config_grounded", i64) + x = var("x", i64, egg_name="__expr_0") + rewrite_builder = birewrite if bidirectional else rewrite + registered = rewrite_builder(RewriteConfigExpr(x)).to(RewriteConfigExpr(rewrite_config_read(x)), grounded(x)) + egraph = EGraph( + set_(rewrite_config_lookup(i64(1))).to(i64(2)), + grounded(i64(1)), + RewriteConfigExpr(i64(1)), + save_egglog_string=True, + seminaive=False, + no_decomp=True, + ) + + egraph.register(registered) + report = egraph.run(1) + egraph.check(eq(RewriteConfigExpr(i64(1))).to(RewriteConfigExpr(i64(2)))) + + transcript = egraph.as_egglog_string + expected_rule_count = 2 if bidirectional else 1 + expected_decl_type = BiRewriteDecl if bidirectional else RewriteDecl + assert report.num_matches_per_rule + assert all(isinstance(decl, expected_decl_type) for decl in report.num_matches_per_rule) + assert transcript.count(":naive") == expected_rule_count + assert transcript.count(":no-decomp") == expected_rule_count + egg_bindings.EGraph().parse_and_run_program(transcript) + + +@pytest.mark.parametrize("egg_name", ["check", "true", "_", "@internal"]) +def test_saved_transcript_rejects_parser_sensitive_explicit_callable_name(egg_name: str) -> None: + rel = relation(f"python_name_{egg_name}", egg_fn=egg_name) + egraph = EGraph(save_egglog_string=True) + + with pytest.raises(ValueError, match="Explicit Egglog callable name"): + egraph.register(rel()) + + assert not egraph.as_egglog_string + + +@pytest.mark.parametrize("save_egglog_string", [False, True], ids=["direct", "saved"]) +def test_empty_explicit_backend_names_are_treated_as_unspecified(*, save_egglog_string: bool) -> None: + class EmptyNamed(Expr, egg_sort=""): + def __init__(self) -> None: ... + + marker = relation("empty_explicit_backend_name", EmptyNamed, egg_fn="") + egraph = EGraph(marker(EmptyNamed()), save_egglog_string=save_egglog_string) + + egraph.check(marker(EmptyNamed())) + if save_egglog_string: + egg_bindings.EGraph().parse_and_run_program(egraph.as_egglog_string) + + +@pytest.mark.parametrize( + ("value", "normalized"), + [ + pytest.param("a\rb", "a\nb", id="cr"), + pytest.param("a\r\nb", "a\nb", id="crlf"), + ], +) +def test_saved_transcript_preserves_string_newlines(value: str, normalized: str) -> None: + marker = relation(f"transcript_newline_{value.count(chr(13))}_{len(value)}", String) + egraph = EGraph(marker(String(value)), save_egglog_string=True) + + with pytest.raises(EggSmolError): + egraph.check(marker(String(normalized))) + + transcript = egraph.as_egglog_string + assert value in transcript + egg_bindings.EGraph().parse_and_run_program(transcript) + + +@pytest.mark.parametrize("message", ['bad "quote"', "bad \\ slash"]) +def test_saved_transcript_escapes_panic_message(message: str) -> None: + trigger = relation("escaped_panic_message_trigger") + panic_rules = ruleset(rule(trigger()).then(panic(message)), name="escaped-panic-message") + egraph = EGraph(trigger(), save_egglog_string=True) + + with pytest.raises(EggSmolError) as exc_info: + egraph.run(panic_rules) + + assert message in str(exc_info.value) + egg_bindings.EGraph().parse_and_run_program(egraph.as_egglog_string) + + +def test_saved_transcript_allows_reserved_callable_name_in_nested_expression() -> None: + @function(egg_fn="check") + def nested_check(value: i64Like) -> i64: ... + + egraph = EGraph(save_egglog_string=True) + egraph.register(set_(nested_check(1)).to(i64(7))) + + assert egraph.lookup_function_value(nested_check(1)) == i64(7) + egg_bindings.EGraph().parse_and_run_program(egraph.as_egglog_string + "\n(check (= (check 1) 7))") + + +def test_saved_transcript_allows_command_head_name_in_rule_action() -> None: + source = relation("rule_action_command_head_source") + target = relation("rule_action_command_head_target", egg_fn="check") + action_rules = ruleset(rule(source()).then(target()), name="rule-action-command-head") + egraph = EGraph(source(), save_egglog_string=True) + + egraph.run(action_rules) + + egraph.check(target()) + egg_bindings.EGraph().parse_and_run_program(egraph.as_egglog_string) + + +def test_saved_transcript_rejects_equality_head_name_in_rule_fact() -> None: + source = relation("equality_head_source", i64, i64) + claimed = relation("equality_head_claimed", i64, i64, egg_fn="=") # type: ignore[call-overload] + target = relation("equality_head_target") + x, y = vars_("x y", i64) + equality_head_rules = ruleset(rule(source(x, y), claimed(x, y)).then(target()), name="equality-head") + + direct = EGraph(source(i64(1), i64(2))) + direct.run(equality_head_rules) + assert not direct.check_bool(target()) + + recorded = EGraph(source(i64(1), i64(2)), save_egglog_string=True) + with pytest.raises(ValueError, match="cannot be used as a fact"): + recorded.run(equality_head_rules) + + +def test_saved_transcript_allows_equality_head_name_as_standalone_action() -> None: + claimed = relation("standalone_equality_head", i64, i64, egg_fn="=") # type: ignore[call-overload] + egraph = EGraph(save_egglog_string=True) + + egraph.register(claimed(i64(1), i64(2))) + + assert egraph.function_size(claimed) == 1 + egg_bindings.EGraph().parse_and_run_program(egraph.as_egglog_string) + + +def test_saved_transcript_rejects_literal_explicit_sort_name() -> None: + class LiteralNamed(Expr, egg_sort="true"): + def __init__(self) -> None: ... + + egraph = EGraph(save_egglog_string=True) + + with pytest.raises(ValueError, match="Explicit Egglog sort name"): + egraph.register(LiteralNamed()) + + assert not egraph.as_egglog_string + + +def test_saved_transcript_rejects_unparseable_explicit_sort_name() -> None: + class UnparseableNamed(Expr, egg_sort="has space"): + def __init__(self) -> None: ... + + egraph = EGraph(save_egglog_string=True) + + with pytest.raises(ValueError, match="Explicit Egglog sort name"): + egraph.register(UnparseableNamed()) + + assert not egraph.as_egglog_string + + +@pytest.mark.parametrize("egg_name", ["_", "@internal"]) +def test_saved_transcript_rejects_context_sensitive_explicit_sort_name(egg_name: str) -> None: + class ContextSensitiveNamed(Expr, egg_sort=egg_name): + def __init__(self) -> None: ... + + egraph = EGraph(save_egglog_string=True) + + with pytest.raises(ValueError, match="Explicit Egglog sort name"): + egraph.register(Map[ContextSensitiveNamed, i64].empty()) + + assert not egraph.as_egglog_string + + +@pytest.mark.parametrize( + "name", + [ + pytest.param("has space", id="space"), + pytest.param("x)", id="parenthesis"), + pytest.param("x;y", id="comment"), + pytest.param("x\ny", id="newline"), + ], +) +def test_saved_transcript_rejects_unparseable_let_name(name: str) -> None: + egraph = EGraph(save_egglog_string=True) + + with pytest.raises(ValueError, match="does not serialize as one Egglog symbol"): + egraph.let(name, i64(1)) + + assert not egraph.as_egglog_string + + +def test_direct_egraph_still_accepts_non_source_let_name() -> None: + egraph = EGraph() + + value = egraph.let("has space", i64(1)) + + egraph.check(eq(value).to(i64(1))) + + +@pytest.mark.parametrize("name", ["has space", "x;y", "x\ny"]) +def test_saved_transcript_rejects_unparseable_ruleset_name(name: str) -> None: + named_ruleset = ruleset(name=name) + + EGraph().run(named_ruleset) + + recorded = EGraph(save_egglog_string=True) + with pytest.raises(ValueError, match="Ruleset name"): + recorded.run(named_ruleset) + assert not recorded.as_egglog_string + + +def test_saved_transcript_rejects_literal_explicit_variable_name() -> None: + source = relation("literal_variable_source", i64) + destination = relation("literal_variable_destination", i64) + x = var("x", i64, egg_name="0") + + direct = EGraph(source(i64(7))) + direct.register(rule(source(x)).then(destination(x))) + direct.run(1) + direct.check(destination(i64(7))) + + recorded = EGraph(source(i64(7)), save_egglog_string=True) + with pytest.raises(ValueError, match="Egglog variable name"): + recorded.register(rule(source(x)).then(destination(x))) + + +@pytest.mark.parametrize( + ("python_name", "egg_name"), + [ + pytest.param("", None, id="default-wildcard"), + pytest.param("has space", None, id="default-space"), + pytest.param("x", "@x", id="internal"), + ], +) +def test_saved_transcript_rejects_non_source_variable_name(python_name: str, egg_name: str | None) -> None: + source = relation("non_source_variable_source", i64) + destination = relation("non_source_variable_destination", i64) + x = var(python_name, i64, egg_name=egg_name) + + direct = EGraph(source(i64(7))) + direct.register(rule(source(x)).then(destination(x))) + direct.run(1) + direct.check(destination(i64(7))) + + recorded = EGraph(source(i64(7)), save_egglog_string=True) + with pytest.raises(ValueError, match="Egglog variable name"): + recorded.register(rule(source(x)).then(destination(x))) + + +def test_saved_transcript_rejects_variable_wildcard_semantic_drift() -> None: + source = relation("wildcard_variable_source", i64, i64) + destination = relation("wildcard_variable_destination") + x = var("x", i64, egg_name="_") + repeated_variable_rule = rule(source(x, x)).then(destination()) + + direct = EGraph(source(i64(1), i64(2))) + direct.register(repeated_variable_rule) + direct.run(1) + assert not direct.check_bool(destination()) + + recorded = EGraph(source(i64(1), i64(2)), save_egglog_string=True) + with pytest.raises(ValueError, match="Egglog variable name"): + recorded.register(repeated_variable_rule) + + def test_eqsat_basic(): egraph = EGraph() @@ -184,6 +482,32 @@ def pair(cls, left: LetConflictNum, right: LetConflictNum) -> LetConflictNum: .. assert egraph.function_size(LetConflictNum.pair) == 1 +@pytest.mark.parametrize("save_egglog_string", [False, True], ids=["direct", "saved"]) +@pytest.mark.parametrize("explicit_first", [False, True], ids=["generated-first", "explicit-first"]) +def test_synthetic_let_names_reserve_explicit_lets_in_a_batch( + *, save_egglog_string: bool, explicit_first: bool +) -> None: + class BatchedLetConflict(Expr): + @classmethod + def leaf(cls, value: i64Like) -> BatchedLetConflict: ... + + @classmethod + def pair(cls, left: BatchedLetConflict, right: BatchedLetConflict) -> BatchedLetConflict: ... + + explicit = let("__expr_0", BatchedLetConflict.leaf(2)) + shared = BatchedLetConflict.leaf(1) + pair = BatchedLetConflict.pair(shared, shared) + actions = (explicit, pair) if explicit_first else (pair, explicit) + egraph = EGraph(save_egglog_string=save_egglog_string) + + egraph.register(*actions) + + assert egraph.function_size(BatchedLetConflict.leaf) == 2 + assert egraph.function_size(BatchedLetConflict.pair) == 1 + if save_egglog_string: + egg_bindings.EGraph().parse_and_run_program(egraph.as_egglog_string) + + def test_synthetic_let_names_do_not_shadow_default_rewrite_variables() -> None: default_ruleset = ruleset(name="synthetic-let-shadow-default-rewrite") @@ -205,6 +529,28 @@ def pair(cls, left: LetShadowDefaultNum, right: LetShadowDefaultNum) -> LetShado egraph.check(eq(LetShadowDefaultNum.make(i64(1))).to(LetShadowDefaultNum(i64(1)))) +def test_shared_expression_discovery_handles_deep_dags() -> None: + class DeepDag(Expr): + @classmethod + def leaf(cls) -> DeepDag: ... + + @classmethod + def wrap(cls, value: DeepDag) -> DeepDag: ... + + @classmethod + def pair(cls, left: DeepDag, right: DeepDag) -> DeepDag: ... + + expr = DeepDag.leaf() + for index in range(1_050): + expr = DeepDag.wrap(expr) + if index % 100 == 99: + expr = DeepDag.pair(expr, expr) + + egraph = EGraph(expr) + + assert egraph.function_size(DeepDag.wrap) == 1_050 + + def test_save_egglog_string_defaults_to_disabled() -> None: egraph = EGraph() @@ -212,31 +558,270 @@ def test_save_egglog_string_defaults_to_disabled() -> None: _ = egraph.as_egglog_string +@pytest.mark.parametrize( + "name", + [ + "", + "true", + "false", + "123", + "-1", + "1e3", + "NaN", + "inf", + "-inf", + "=", + "sort", + "datatype", + "datatype*", + "function", + "constructor", + "relation", + "ruleset", + "unstable-combined-ruleset", + "rule", + "rewrite", + "birewrite", + "run", + "run-schedule", + "extract", + "check", + "prove", + "prove-exists", + "push", + "pop", + "print-stats", + "print-function", + "print-size", + "input", + "output", + "include", + "fail", + "let", + "set", + "delete", + "subsume", + "union", + "panic", + "for", + "with-ruleset", + "with-dynamic-cost", + "set-cost", + "let-scheduler", + "multi-extract", + "keep-best", + "print-table-stats", + "primitive", + ], +) +def test_generated_relation_names_avoid_egglog_parser_tokens(name: str) -> None: + relation_with_parser_token_name = relation(name) + egraph = EGraph(save_egglog_string=True) + + egraph.register(relation_with_parser_token_name()) + + egraph.check(relation_with_parser_token_name()) + + +@pytest.mark.parametrize("generated_first", [True, False], ids=["generated-first", "explicit-first"]) +def test_generated_relation_names_reserve_explicit_backend_names(generated_first: bool) -> None: + generated = relation("batch_generated_backend_name", i64) + generated_candidate = f"{__name__.replace('.', '_')}_batch_generated_backend_name" + explicit = relation("explicit_relation", i64, egg_fn=generated_candidate) + generated_action = generated(i64(1)) + explicit_action = explicit(i64(2)) + actions = (generated_action, explicit_action) if generated_first else (explicit_action, generated_action) + egraph = EGraph(save_egglog_string=True) + + egraph.register(*actions) + + egraph.check(generated(i64(1)), explicit(i64(2))) + egraph.check_fail(generated(i64(2))) + egraph.check_fail(explicit(i64(1))) + + +@pytest.mark.parametrize("save_egglog_string", [False, True], ids=["direct", "saved"]) +@pytest.mark.parametrize("explicit_first", [False, True], ids=["generated-first", "explicit-first"]) +def test_synthetic_let_names_reserve_explicit_backend_names(*, save_egglog_string: bool, explicit_first: bool) -> None: + class SyntheticLetNode(Expr): + @classmethod + def leaf(cls, value: i64Like) -> SyntheticLetNode: ... + + @classmethod + def pair(cls, left: SyntheticLetNode, right: SyntheticLetNode) -> SyntheticLetNode: ... + + explicit = relation("synthetic_let_explicit", egg_fn="$__expr_0") + shared = SyntheticLetNode.leaf(1) + pair = SyntheticLetNode.pair(shared, shared) + actions = (explicit(), pair) if explicit_first else (pair, explicit()) + egraph = EGraph(save_egglog_string=save_egglog_string) + + egraph.register(*actions) + + egraph.check(explicit()) + assert egraph.function_size(SyntheticLetNode.pair) == 1 + if save_egglog_string: + egg_bindings.EGraph().parse_and_run_program(egraph.as_egglog_string) + + +@pytest.mark.parametrize("save_egglog_string", [False, True], ids=["direct", "saved"]) +def test_synthetic_lets_do_not_capture_unbound_variables(*, save_egglog_string: bool) -> None: + class SyntheticLetVariable(Expr): + @classmethod + def leaf(cls, value: i64Like) -> SyntheticLetVariable: ... + + @classmethod + def pair(cls, left: SyntheticLetVariable, right: SyntheticLetVariable) -> SyntheticLetVariable: ... + + x = var("synthetic_let_x", i64) + shared = SyntheticLetVariable.leaf(x) + egraph = EGraph(save_egglog_string=save_egglog_string) + + with pytest.raises(EggSmolError, match="Unbound symbol"): + egraph.register(SyntheticLetVariable.pair(shared, shared)) + + if save_egglog_string: + with pytest.raises(RuntimeError, match="partial effects cannot be replayed"): + egraph.register(SyntheticLetVariable.leaf(1)) + else: + egraph.register(SyntheticLetVariable.leaf(1)) + + +@pytest.mark.parametrize("save_egglog_string", [False, True], ids=["direct", "saved"]) +def test_late_explicit_backend_name_cannot_replace_registered_generated_relation(*, save_egglog_string: bool) -> None: + generated = relation("late_explicit_backend_name", i64) + egraph = EGraph(save_egglog_string=save_egglog_string) + egraph.register(generated(i64(1))) + transcript = egraph.as_egglog_string if save_egglog_string else None + + generated_name = f"{__name__.replace('.', '_')}_late_explicit_backend_name" + explicit = relation("late_explicit_relation", i64, egg_fn=generated_name) + with pytest.raises(ValueError, match="already registered"): + egraph.register(explicit(i64(2))) + + egraph.check(generated(i64(1))) + if transcript is not None: + assert egraph.as_egglog_string.startswith(transcript) + + +@pytest.mark.parametrize("save_egglog_string", [False, True], ids=["direct", "saved"]) +def test_duplicate_explicit_backend_name_does_not_publish_a_failed_alias(*, save_egglog_string: bool) -> None: + @function(egg_fn="duplicate_explicit_backend") + def first(value: i64Like) -> i64: ... + + @function(egg_fn="duplicate_explicit_backend") + def second(value: i64Like) -> i64: ... + + egraph = EGraph(set_(first(1)).to(i64(10)), save_egglog_string=save_egglog_string) + transcript = egraph.as_egglog_string if save_egglog_string else None + + with pytest.raises(ValueError, match="already registered"): + egraph.register(set_(second(2)).to(i64(20))) + + assert egraph.lookup_function_value(first(1)) == i64(10) + assert "first" in str(egraph.freeze()) + if transcript is not None: + assert egraph.as_egglog_string == transcript + egg_bindings.EGraph().parse_and_run_program(transcript) + + +@pytest.mark.parametrize("save_egglog_string", [False, True], ids=["direct", "saved"]) +def test_late_explicit_backend_name_cannot_replace_registered_generated_sort(*, save_egglog_string: bool) -> None: + class GeneratedLateSort(Expr): + def __init__(self) -> None: ... + + egraph = EGraph(save_egglog_string=save_egglog_string) + egraph.register(GeneratedLateSort()) + transcript = egraph.as_egglog_string if save_egglog_string else None + + class ExplicitLateSort(Expr, egg_sort=f"{__name__}.GeneratedLateSort"): + def __init__(self) -> None: ... + + with pytest.raises(ValueError, match="already registered"): + egraph.register(ExplicitLateSort()) + + egraph.check(GeneratedLateSort()) + if transcript is not None: + assert egraph.as_egglog_string.startswith(transcript) + + +def test_generated_relation_names_avoid_command_macro_heads() -> None: + source = relation("command_macro_source") + generated = relation("unstable-fresh!") + copy = ruleset(rule(source()).then(generated())) + egraph = EGraph(source(), save_egglog_string=True) + + egraph.run(copy) + + egraph.check(generated()) + + +@pytest.mark.parametrize("name", ["1.0", "-1.0", "+.5", "1.", "1E+3"]) +def test_generated_sort_names_avoid_egglog_float_literals(name: str) -> None: + def init(_self: object) -> None: ... + + numeric_name_expr = type(name, (Expr,), {"__module__": __name__, "__init__": init}) + expr = numeric_name_expr() + egraph = EGraph(save_egglog_string=True) + + egraph.register(expr) + + egraph.check(expr) + + +@pytest.mark.parametrize( + "name", + [ + pytest.param("has space", id="space"), + pytest.param("x;y", id="comment"), + pytest.param("x(y)", id="parentheses"), + pytest.param("x\ny", id="newline"), + ], +) +def test_generated_sort_names_are_safe_in_saved_source(name: str) -> None: + def init(_self: object) -> None: ... + + generated_expr = type(name, (Expr,), {"__module__": __name__, "__init__": init}) + expr = generated_expr() + egraph = EGraph(expr, save_egglog_string=True) + + egraph.check(expr) + egg_bindings.EGraph().parse_and_run_program(egraph.as_egglog_string) + + def test_saved_egglog_transcript_close_is_idempotent() -> None: egraph = EGraph(save_egglog_string=True) - assert egraph._state.egglog_file_state is not None - path = pathlib.Path(egraph._state.egglog_file_state.path) egraph.let("x", i64(1)) assert egraph.as_egglog_string egraph.close() egraph.close() - assert not path.exists() with pytest.raises(ValueError, match="has been closed"): _ = egraph.as_egglog_string -def test_saved_egglog_transcript_is_removed_on_finalization() -> None: - egraph = EGraph(save_egglog_string=True) - assert egraph._state.egglog_file_state is not None - path = pathlib.Path(egraph._state.egglog_file_state.path) - assert path.exists() +def test_saved_egglog_transcript_close_rejects_commands_before_mutation() -> None: + class ClosedExpr(Expr): + def __init__(self, value: i64Like) -> None: ... + + egraph = EGraph(ClosedExpr(1), save_egglog_string=True) + before = egraph.freeze() + egraph.close() + + with pytest.raises(ValueError, match="transcript has been closed"): + egraph.register(ClosedExpr(2)) + + assert egraph.freeze() == before - del egraph - gc.collect() - assert not path.exists() +def test_close_without_a_saved_transcript_is_a_noop() -> None: + egraph = EGraph() + + egraph.close() + value = egraph.let("after_close", i64(1)) + + egraph.check(eq(value).to(i64(1))) def test_saved_egglog_transcript_is_shared_across_push_and_pop() -> None: @@ -251,7 +836,7 @@ def test_saved_egglog_transcript_is_shared_across_push_and_pop() -> None: assert "(let $inner 2)" in egraph.as_egglog_string -def test_saved_egglog_string_uses_short_generated_sort_and_function_names() -> None: +def test_saved_egglog_string_uses_qualified_generated_sort_and_function_names() -> None: class Num(Expr): @classmethod def var(cls, v: StringLike) -> Num: ... @@ -260,183 +845,24 @@ def var(cls, v: StringLike) -> Num: ... egraph.register(Num.var("x")) egglog_string = egraph.as_egglog_string - assert "(sort Num)" in egglog_string - assert "(constructor Num_var (String) Num)" in egglog_string - assert "test_high_level" not in egglog_string - - -def test_generated_names_fall_back_to_full_name_on_conflict() -> None: - state = EGraph(save_egglog_string=True)._state - ret1 = Ident("Ret", "pkg.one") - ret2 = Ident("Ret", "pkg.two") - fn1 = Ident("make", "pkg.one") - fn2 = Ident("make", "pkg.two") - state.__egg_decls__ |= Declarations( - _classes={ret1: ClassDecl(), ret2: ClassDecl()}, - _functions={ - fn1: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret1))), - fn2: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret2))), - }, - ) - - assert state.callable_ref_to_egg(FunctionRef(fn1))[0] == "make" - assert state.callable_ref_to_egg(FunctionRef(fn2))[0] == "pkg_two_make" - assert state.type_ref_to_egg(JustTypeRef(ret1)) == "Ret" - assert state.type_ref_to_egg(JustTypeRef(ret2)) == "pkg.two.Ret" - - -def test_missing_function_lookup_does_not_reserve_generated_name() -> None: - state = EGraph(save_egglog_string=True)._state - ret = Ident("LookupRet", "pkg.lookup") - fn = Ident("lookup_short_name", "pkg.lookup") - state.__egg_decls__ |= Declarations( - _classes={ret: ClassDecl()}, - _functions={fn: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret)))}, - ) - - assert list(state.possible_egglog_functions(["lookup_short_name"])) == [] - assert state.callable_ref_to_egg(FunctionRef(fn))[0] == "lookup_short_name" - - -def test_generated_names_fall_back_from_builtin_names() -> None: - state = EGraph(save_egglog_string=True)._state - ret = Ident("BuiltinConflictRet", "pkg.builtin_conflict") - fn = Ident("exp", "pkg.builtin_conflict") - sort = Ident("Map", "pkg.builtin_conflict") - state.__egg_decls__ |= Declarations( - _classes={ret: ClassDecl(), sort: ClassDecl()}, - _functions={fn: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret)))}, - ) - - assert state.callable_ref_to_egg(FunctionRef(fn))[0] == "pkg_builtin_conflict_exp" - assert state.type_ref_to_egg(JustTypeRef(sort)) == "pkg.builtin_conflict.Map" - - -def test_generated_callable_name_avoids_an_existing_cost_table() -> None: - state = EGraph(save_egglog_string=True)._state - state.__egg_decls__ |= cast("HasDeclarations", i64) - ret = Ident("CostRet", "pkg.cost") - fn = Ident("f", "pkg.cost") - conflict = Ident("cost_table_f", "pkg.cost") - state.__egg_decls__ |= Declarations( - _classes={ret: ClassDecl()}, - _functions={ - fn: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret))), - conflict: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret))), - }, - ) + qualified_sort = f"{__name__}.Num" + qualified_var = f"{__name__.replace('.', '_')}_Num_var" + assert f"(sort {qualified_sort})" in egglog_string + assert f"(constructor {qualified_var} (String) {qualified_sort})" in egglog_string - assert state.create_cost_table(FunctionRef(fn)) == "cost_table_f" - assert state.callable_ref_to_egg(FunctionRef(conflict))[0] == "pkg_cost_cost_table_f" - - -def test_canonical_cost_table_rejects_an_incompatible_callable() -> None: - state = EGraph(save_egglog_string=True)._state - state.__egg_decls__ |= cast("HasDeclarations", i64) - ret = Ident("CostRet", "pkg.cost") - fn = Ident("f", "pkg.cost") - conflict = Ident("cost_table_f", "pkg.cost") - state.__egg_decls__ |= Declarations( - _classes={ret: ClassDecl()}, - _functions={ - fn: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret))), - conflict: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret))), - }, - ) - assert state.callable_ref_to_egg(FunctionRef(conflict))[0] == "cost_table_f" - with pytest.raises(ValueError, match="already used by an incompatible callable"): - state.create_cost_table(FunctionRef(fn)) - - -def test_canonical_cost_table_reuses_a_compatible_raw_table() -> None: - state = EGraph(save_egglog_string=True)._state - state.__egg_decls__ |= cast("HasDeclarations", i64) - ret = Ident("CostRet", "pkg.cost") - fn = Ident("f", "pkg.cost") - raw_cost = Ident("cost_table_f", "pkg.cost") - state.__egg_decls__ |= Declarations( - _classes={ret: ClassDecl()}, - _functions={ - fn: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(ret))), - raw_cost: FunctionDecl(signature=FunctionSignature(return_type=TypeRefWithVars(Ident.builtin("i64")))), - }, - ) - - assert state.callable_ref_to_egg(FunctionRef(raw_cost))[0] == "cost_table_f" - assert state.create_cost_table(FunctionRef(fn)) == "cost_table_f" - assert state.cost_table_names[FunctionRef(fn)] == "cost_table_f" +def test_generated_names_do_not_collide_with_lazy_builtin_declarations() -> None: + class String(Expr): + def __init__(self, value: i64Like) -> None: ... + @function + def log(value: i64Like) -> i64: ... -@pytest.mark.parametrize( - "body", - [ - pytest.param(None, id="bodyless-function"), - pytest.param( - TypedExprDecl(JustTypeRef(Ident.builtin("i64")), LitDecl(1)), - id="eager-primitive", - ), - ], -) -@pytest.mark.parametrize("sort_first", [True, False], ids=["sort-first", "callable-first"]) -def test_generated_names_share_the_backend_sort_and_callable_namespace( - body: TypedExprDecl | None, *, sort_first: bool -) -> None: - state = EGraph(save_egglog_string=True)._state - state.__egg_decls__ |= cast("HasDeclarations", i64) - sort_ident = Ident("Node", "pkg.sort") - fn_ident = Ident("Node", "pkg.fn") - sort_ref = JustTypeRef(sort_ident) - fn_ref = FunctionRef(fn_ident) - state.__egg_decls__ |= Declarations( - _classes={sort_ident: ClassDecl()}, - _functions={ - fn_ident: FunctionDecl( - signature=FunctionSignature(return_type=TypeRefWithVars(Ident.builtin("i64"))), - body=body, - ) - }, - ) + egraph = EGraph() + user_string = String(1) + egraph.register(user_string, set_(log(1)).to(i64(2))) - if sort_first: - assert state.type_ref_to_egg(sort_ref) == "Node" - assert state.callable_ref_to_egg(fn_ref)[0] == "pkg_fn_Node" - else: - assert state.callable_ref_to_egg(fn_ref)[0] == "Node" - assert state.type_ref_to_egg(sort_ref) == "pkg.sort.Node" - - -def test_builtin_name_reservations_cover_builtins_module_declarations() -> None: - expected_fn_names = set[str]() - expected_sort_names = set[str]() - - def add_callable_name(decl: CallableDecl | None) -> None: - if decl is not None and decl.egg_name is not None: - expected_fn_names.add(decl.egg_name) - - for name in egg_builtins.__all__: - obj = getattr(egg_builtins, name, None) - if not isinstance(obj, HasDeclarations): - continue - decls = obj.__egg_decls__ - for decl in decls._functions.values(): - add_callable_name(decl) - for decl in decls._constants.values(): - add_callable_name(decl) - for decl in decls._classes.values(): - if decl.builtin and decl.egg_name is not None: - expected_sort_names.add(decl.egg_name) - add_callable_name(decl.init) - for callable_decl in ( - *decl.class_methods.values(), - *decl.class_variables.values(), - *decl.methods.values(), - *decl.properties.values(), - ): - add_callable_name(callable_decl) - - assert expected_fn_names <= BUILTIN_EGG_FN_NAMES - assert expected_sort_names <= BUILTIN_EGG_SORT_NAMES + egraph.check(eq(user_string).to(String(1)), eq(log(1)).to(i64(2))) def test_parameterized_sort_names_use_allocated_argument_names() -> None: @@ -498,11 +924,16 @@ def pair(cls, left: DuplicateEdge, right: DuplicateEdge) -> DuplicateEdge: ... egraph.register(DuplicateEdge.pair(second_pair, second_pair)) - transcript = egraph.as_egglog_string - assert transcript.count('(let $__expr_0 (DuplicateEdge_leaf "shared"))') == 1 - assert transcript.count("(let $__expr_1 (DuplicateEdge_pair $__expr_0 $__expr_0))") == 1 - assert transcript.count("(let $__expr_2 (DuplicateEdge_pair $__expr_1 $__expr_1))") == 1 - assert "(DuplicateEdge_pair $__expr_2 $__expr_2)" in transcript + lines = egraph.as_egglog_string.splitlines() + let_lines = [line for line in lines if line.startswith("(let $__expr_")] + assert len(let_lines) == 3 + assert let_lines[0].startswith("(let $__expr_0 (") + assert let_lines[0].endswith(' "shared"))') + assert let_lines[1].startswith("(let $__expr_1 (") + assert let_lines[1].endswith(" $__expr_0 $__expr_0))") + assert let_lines[2].startswith("(let $__expr_2 (") + assert let_lines[2].endswith(" $__expr_1 $__expr_1))") + assert sum(line.endswith(" $__expr_2 $__expr_2)") for line in lines) == 1 def test_freeze_omits_synthetic_let_bindings() -> None: @@ -566,55 +997,6 @@ def test_map_and_set_length_primitives() -> None: check_eq(Set(i64(1), i64(2)).length(), i64(2)) -def test_higher_order_callable_inference_does_not_mutate_ambient_ruleset() -> None: - ambient = ruleset(name="hof-inference-ambient") - initial_rules = tuple(ambient.__egg_ruleset__.rules) - - with set_current_ruleset(ambient): - initial: Map[i64, f64] = Map[i64, f64].empty() - expr = map_fold_kv( - lambda result, key, value: result.insert(key, -value), - initial, - Map[i64, f64].empty().insert(i64(1), f64(2.0)), - ) - _ = cast("RuntimeExpr", expr).__egg_decls__ - - assert tuple(ambient.__egg_ruleset__.rules) == initial_rules - - -def test_set_current_ruleset_restores_nested_contexts() -> None: - outer = ruleset(name="current-ruleset-outer") - inner = ruleset(name="current-ruleset-inner") - initial = get_current_ruleset() - - with set_current_ruleset(outer): - assert get_current_ruleset() is outer - with set_current_ruleset(inner): - assert get_current_ruleset() is inner - assert get_current_ruleset() is outer - - assert get_current_ruleset() is initial - - -def test_file_backed_errors_report_saved_file_line() -> None: - egraph = EGraph(save_egglog_string=True) - egraph.let("x", i64(1)) - egraph.let("y", i64(2)) - expected_line = len(egraph.as_egglog_string.splitlines()) + 1 - assert egraph._state.egglog_file_state is not None - path = egraph._state.egglog_file_state.path - - with pytest.raises(EggSmolError) as exc_info: - egraph.check(eq(i64(1)).to(i64(2))) - - error_text = exc_info.value.context - assert path in error_text - assert f"In {expected_line}:" in error_text - lines = egraph.as_egglog_string.splitlines() - assert "(fail (check (= 1 2))) ; Check failed:" in lines - assert "(check (= 1 2))" not in lines - - def test_unnamed_lambda_returning_builtin_is_eager() -> None: check_eq( map_fold_kv( @@ -1045,6 +1427,12 @@ def test_rational_like_operations() -> None: assert EGraph().extract(2 ** Rational(3, 1)).value == Fraction(8, 1) assert EGraph().extract(Rational(1, 2).min(Fraction(1, 3))).value == Fraction(1, 3) assert EGraph().extract(Rational(1, 2).max(1)).value == Fraction(1, 1) + assert EGraph().extract(Rational(1, 1).log()).value == Fraction(0, 1) + assert EGraph().extract(Rational(0, 1).sqrt()).value == Fraction(0, 1) + assert EGraph().extract(Rational(4, 9).sqrt()).value == Fraction(2, 3) + assert EGraph().extract(Rational(-8, 27).cbrt()).value == Fraction(-2, 3) + assert EGraph().extract(Rational(-(2**63), 1).floor()).value == Fraction(-(2**63), 1) + assert EGraph().extract(Rational(2**63 - 1, 1).ceil()).value == Fraction(2**63 - 1, 1) egraph = EGraph() egraph.check(Rational(1, 2) < Fraction(2, 3)) @@ -1054,10 +1442,22 @@ def test_rational_like_operations() -> None: def test_rational_partial_operations_remain_undefined() -> None: - with pytest.raises(EggSmolError): - EGraph().extract(Rational(1, 2) / 0) - with pytest.raises(EggSmolError): - EGraph().extract(Rational(2, 1) ** -1) + expressions = [ + Rational(1, 0), + Rational(1, -(2**63)), + Rational(-(2**63), -1), + -Rational(-(2**63), 1), + abs(Rational(-(2**63), 1)), + Rational(1, 2) / 0, + Rational(2, 1) ** -1, + Rational(0, 1) ** Fraction(1, 2), + Rational(2, 1).log(), + Rational(2, 1).sqrt(), + Rational(2, 1).cbrt(), + ] + for expression in expressions: + with pytest.raises(EggSmolError): + EGraph().extract(expression) with pytest.raises(EggSmolError): EGraph().check(Rational(2, 3) < Fraction(1, 2)) @@ -1699,6 +2099,41 @@ def __init__(self) -> None: ... class TestCallableValidation: + @pytest.mark.parametrize( + ("options", "message"), + [ + pytest.param({"cost": 1}, "Cost can only be set for constructors", id="cost"), + pytest.param({"unextractable": True}, "Unextractable can only be set for constructors", id="unextractable"), + pytest.param( + {"cost": 1, "unextractable": True}, + "Cost can only be set for constructors", + id="cost-before-unextractable", + ), + ], + ) + def test_eqsort_merge_rejects_constructor_options_before_body( + self, options: dict[str, object], message: str + ) -> None: + body_calls = 0 + merge_calls = 0 + + def merge(old: A, _new: A) -> A: + nonlocal merge_calls + merge_calls += 1 + return old + + @function(merge=merge, **options) # type: ignore[call-overload] + def f() -> A: + nonlocal body_calls + body_calls += 1 + return A() + + with pytest.raises(ValueError, match=message): + f() + + assert body_calls == 0 + assert merge_calls == 0 + def test_primitive_function_ruleset_subsume_rejected(self): r = ruleset() @@ -1733,10 +2168,18 @@ def test_primitive_constant_default_cannot_use_explicit_ruleset(self): ) def test_eqsort_constant_default_cannot_use_merge(self): + merge_calls = 0 + + def merge(old: A, _new: A) -> A: + nonlocal merge_calls + merge_calls += 1 + return old + with pytest.raises(ValueError, match="Eqsort-returning callables with bodies cannot use merge"): EGraph().register( - constant("default_merge", A, A(), merge=lambda old, _new: old) # type: ignore[call-overload] + constant("default_merge", A, A(), merge=merge) # type: ignore[call-overload] ) + assert merge_calls == 0 def test_primitive_constant_default_cannot_use_merge(self): with pytest.raises(ValueError, match="Primitive-returning callables with bodies cannot use merge"): @@ -1751,12 +2194,24 @@ def test_unit_constant_cannot_use_merge(self): EGraph().register(constant("unit_merge", Unit, merge=lambda old, _new: old)) def test_eqsort_eager_body_cannot_use_merge(self): - @function(merge=lambda old, new: old) + body_calls = 0 + merge_calls = 0 + + def merge(old: A, _new: A) -> A: + nonlocal merge_calls + merge_calls += 1 + return old + + @function(merge=merge) def f() -> A: + nonlocal body_calls + body_calls += 1 return A() with pytest.raises(ValueError, match="Eqsort-returning callables with bodies cannot use merge"): f() + assert body_calls == 0 + assert merge_calls == 0 def test_primitive_returning_functions_cannot_use_cost(self): @function(cost=1) # type: ignore[type-var] # Deliberately invalid runtime API call. @@ -1780,12 +2235,39 @@ def f() -> i64: ... f() def test_primitive_body_cannot_use_builtin(self): + body_calls = 0 + @function(builtin=True) def f() -> i64: + nonlocal body_calls + body_calls += 1 return i64(1) with pytest.raises(ValueError, match="Builtin callables cannot have a body"): f() + assert body_calls == 0 + + def test_generic_primitive_body_cannot_use_builtin(self): + body_calls = 0 + + @function(builtin=True) + def f(value: Vec[_BuiltinExprT]) -> i64: + nonlocal body_calls + body_calls += 1 + return value.length() + + with pytest.raises(ValueError, match="Builtin callables cannot have a body"): + f(Vec(i64(1))) + assert body_calls == 0 + + def test_generic_builtin_cannot_use_explicit_ruleset(self): + r = ruleset() + + @function(builtin=True, ruleset=r) # type: ignore[call-overload] # Deliberately invalid runtime API call. + def f(value: Vec[_BuiltinExprT]) -> i64: ... + + with pytest.raises(ValueError, match="Builtin callables cannot use an explicit ruleset"): + f(Vec(i64(1))) def test_primitive_body_cannot_use_merge(self): @function(merge=lambda old, new: old) @@ -1807,6 +2289,32 @@ def f() -> i64: ): f() + def test_primitive_body_must_return_a_value(self): + body_calls = 0 + + @function + def f() -> i64: + nonlocal body_calls + body_calls += 1 + return None # type: ignore[return-value] # Exercise a genuine body with a missing symbolic result. + + with pytest.raises(ValueError, match="Callable bodies must return a value"): + f() + assert body_calls == 1 + + def test_eqsort_body_must_return_a_value(self): + body_calls = 0 + + @function + def f() -> A: + nonlocal body_calls + body_calls += 1 + return None # type: ignore[return-value] # Exercise a genuine body with a missing symbolic result. + + with pytest.raises(ValueError, match="Callable bodies must return a value"): + f() + assert body_calls == 1 + def test_eqsort_body_cannot_use_merge(self): r = ruleset() @@ -2219,6 +2727,41 @@ def f(x: i64Like) -> i64: ... assert values == {f(i64(1)): i64(2)} +def test_lookup_function_value_records_materialized_constructor_arguments() -> None: + class LookupReplayValue(Expr, egg_sort="LookupReplayValueSort"): + @method(egg_fn="LookupReplayValue") + def __init__(self, value: i64Like) -> None: ... + + @function(egg_fn="lookup-replay-result") + def result(value: LookupReplayValue) -> i64: ... + + egraph = EGraph(save_egglog_string=True) + egraph.register(set_(result(LookupReplayValue(1))).to(i64(10))) + + assert egraph.lookup_function_value(result(LookupReplayValue(2))) is None + assert egraph.function_size(LookupReplayValue) == 2 + + replayed = egg_bindings.EGraph() + outputs = replayed.parse_and_run_program(egraph.as_egglog_string + "\n(print-size)") + sizes = outputs[-1] + assert isinstance(sizes, egg_bindings.PrintAllFunctionsSize) + assert dict(sizes.sizes)["LookupReplayValue"] == 2 + + +@pytest.mark.parametrize("save_egglog_string", [False, True], ids=["direct", "saved"]) +def test_lookup_function_value_with_unstable_fn_argument(*, save_egglog_string: bool) -> None: + @function + def add_one(value: i64) -> i64: ... + + @function + def score(fn: UnstableFn[i64, i64]) -> i64: ... + + fn = UnstableFn[i64, i64](add_one) + egraph = EGraph(set_(score(fn)).to(i64(7)), save_egglog_string=save_egglog_string) + + assert egraph.lookup_function_value(score(fn)) == i64(7) + + def test_table_inspection_rejects_eager_primitives() -> None: @function def eager_plus_one(x: i64Like) -> i64: @@ -2242,6 +2785,24 @@ def eager_text() -> String: egraph.input(eager_text, "unused.csv") +@pytest.mark.parametrize( + "filename", + ['input\r"\\.csv', "input\u0300.csv"], + ids=["control-and-escapes", "combining-unicode"], +) +def test_saved_transcript_preserves_input_path(tmp_path: pathlib.Path, filename: str) -> None: + @function + def loaded() -> String: ... + + path = tmp_path / filename + path.write_text("value\n") + recorded = EGraph(save_egglog_string=True) + + recorded.input(loaded, str(path)) + recorded.check(eq(loaded()).to(String("value"))) + egg_bindings.EGraph().parse_and_run_program(recorded.as_egglog_string) + + def test_dynamic_cost(): """ https://github.com/egraphs-good/egglog-experimental/blob/6d07a34ac76deec751f86f70d9b9358cd3e236ca/tests/integration_test.rs#L5-L35 @@ -2273,7 +2834,83 @@ def __sub__(self, other: E) -> E: ... assert egraph.extract(E(2), include_cost=True, extractor="greedy-dag") == (E(1) + E(1), 102) -def test_dynamic_cost_reuses_a_compatible_canonical_table() -> None: +@pytest.mark.parametrize("save_egglog_string", [False, True], ids=["direct", "saved"]) +def test_dynamic_cost_materializes_constructor_before_storing_cost(*, save_egglog_string: bool) -> None: + class MaterializedCost(Expr): + def __init__(self, value: i64Like) -> None: ... + + expression = MaterializedCost(7) + egraph = EGraph(save_egglog_string=save_egglog_string) + + egraph.register(set_cost(expression, 9)) + + assert egraph.function_size(MaterializedCost) == 1 + assert egraph.lookup_function_value(get_cost(expression)) == i64(9) + if save_egglog_string: + egg_bindings.EGraph().parse_and_run_program(egraph.as_egglog_string) + + +@pytest.mark.parametrize("save_egglog_string", [False, True], ids=["direct", "saved"]) +def test_dynamic_cost_rule_materializes_constructor_without_shadowing(*, save_egglog_string: bool) -> None: + class RuleMaterializedCost(Expr): + def __init__(self, value: i64Like) -> None: ... + + source = relation("dynamic_cost_materialize_source", i64) + # Exercise a name used by compiler-generated rule lets: lowering must + # allocate a distinct binding instead of shadowing this variable. + value = var("value", i64, egg_name="__expr_0") + materialize = ruleset( + rule(source(value)).then(set_cost(RuleMaterializedCost(value), 11)), + name="dynamic-cost-materialize", + ) + expression = RuleMaterializedCost(8) + egraph = EGraph(source(i64(8)), save_egglog_string=save_egglog_string) + + egraph.run(materialize) + + assert egraph.function_size(RuleMaterializedCost) == 1 + assert egraph.lookup_function_value(get_cost(expression)) == i64(11) + if save_egglog_string: + egg_bindings.EGraph().parse_and_run_program(egraph.as_egglog_string) + + +@pytest.mark.parametrize("save_egglog_string", [False, True], ids=["direct", "saved"]) +def test_dynamic_cost_does_not_create_a_missing_function_row(*, save_egglog_string: bool) -> None: + @function + def missing_target(value: i64Like) -> i64: ... + + expression = missing_target(1) + egraph = EGraph(save_egglog_string=save_egglog_string) + + with pytest.raises(EggSmolError, match=r"lookup .* failed"): + egraph.register(set_cost(expression, 7)) + + assert egraph.function_size(missing_target) == 0 + assert egraph.lookup_function_value(get_cost(expression)) is None + if save_egglog_string: + egg_bindings.EGraph().parse_and_run_program(egraph.as_egglog_string) + + +def test_dynamic_cost_evaluates_each_argument_once(tmp_path: pathlib.Path) -> None: + marker = tmp_path / "set-cost-calls" + + def record_call(value: int) -> int: + marker.write_text(marker.read_text() + "x" if marker.exists() else "x") + return value + + class SingleEvaluationCost(Expr): + def __init__(self, value: PyObject) -> None: ... + + egraph = EGraph(save_egglog_string=True) + egraph.register(set_cost(SingleEvaluationCost(PyObject(record_call)(PyObject(7))), 9)) + + assert marker.read_text() == "x" + egg_bindings.EGraph().parse_and_run_program(egraph.as_egglog_string) + assert marker.read_text() == "xx" + + +@pytest.mark.parametrize("cost_first", [False, True], ids=["raw-table-first", "cost-first"]) +def test_dynamic_cost_reuses_a_compatible_canonical_table(cost_first: bool) -> None: @function(egg_fn="canonical_cost_target") def target(x: i64Like) -> i64: ... @@ -2281,10 +2918,12 @@ def target(x: i64Like) -> i64: ... def raw_cost(x: i64Like) -> i64: ... egraph = EGraph() + cost_action = set_cost(target(1), 7) + raw_table_action = set_(raw_cost(2)).to(i64(5)) + cost_actions = (cost_action, raw_table_action) if cost_first else (raw_table_action, cost_action) egraph.register( - set_(raw_cost(2)).to(i64(5)), set_(target(1)).to(i64(2)), - set_cost(target(1), 7), + *cost_actions, ) assert egraph.lookup_function_value(raw_cost(1)) == i64(7) @@ -2292,6 +2931,53 @@ def raw_cost(x: i64Like) -> i64: ... assert egraph.has_custom_cost(target) +@pytest.mark.parametrize("cost_first", [False, True], ids=["raw-table-first", "cost-first"]) +def test_dynamic_cost_rejects_a_merged_canonical_table(cost_first: bool) -> None: + @function(egg_fn="merged_cost_target") + def target(x: i64Like) -> i64: ... + + @function(egg_fn="cost_table_merged_cost_target", merge=lambda old, _new: old) + def raw_cost(x: i64Like) -> i64: ... + + cost_action = set_cost(target(1), 7) + raw_table_action = set_(raw_cost(2)).to(i64(5)) + actions = (cost_action, raw_table_action) if cost_first else (raw_table_action, cost_action) + + with pytest.raises(ValueError, match="already used by an incompatible callable"): + EGraph().register(*actions) + + +def test_dynamic_cost_rejects_multiple_raw_canonical_table_aliases() -> None: + @function(egg_fn="aliased_cost_target") + def target(x: i64Like) -> i64: ... + + @function(egg_fn="cost_table_aliased_cost_target") + def first_raw_cost(x: i64Like) -> i64: ... + + @function(egg_fn="cost_table_aliased_cost_target") + def second_raw_cost(x: i64Like) -> i64: ... + + with pytest.raises(ValueError, match="already has a raw callable alias"): + EGraph().register( + set_cost(target(1), 7), + set_(first_raw_cost(2)).to(i64(5)), + set_(second_raw_cost(3)).to(i64(6)), + ) + + +@pytest.mark.parametrize("cost_first", [False, True], ids=["raw-table-first", "cost-first"]) +def test_dynamic_cost_rejects_a_raw_table_alias_for_a_primitive(cost_first: bool) -> None: + @function(egg_fn="cost_table_/") + def raw_cost(left: i64Like, right: i64Like) -> i64: ... + + cost_action = set_cost(i64(4) / i64(2), 7) + raw_table_action = set_(raw_cost(1, 0)).to(i64(5)) + actions = (cost_action, raw_table_action) if cost_first else (raw_table_action, cost_action) + + with pytest.raises(ValueError, match="eager or builtin primitive"): + EGraph(*actions) + + def test_freeze_preserves_a_reused_canonical_cost_table_as_raw_rows_and_costs() -> None: @function(egg_fn="freeze_cost_target") def target(x: i64Like) -> i64: ... @@ -2308,7 +2994,7 @@ def raw_cost(x: i64Like) -> i64: ... rendered = str(egraph.freeze()) assert "set_(raw_cost(2)).to(i64(5))" in rendered assert "set_(raw_cost(1)).to(i64(7))" in rendered - assert "set_cost(target(2), 5)" in rendered + assert "set_cost(target(2), 5)" not in rendered assert "set_cost(target(1), 7)" in rendered replayed = eval(rendered.removesuffix(".freeze()"), globals(), locals()) @@ -2337,6 +3023,23 @@ def plus_alias(left: i64Like, right: i64Like) -> i64: ... assert replayed.lookup_function_value(get_cost(plus_alias(3, 4))) == i64(6) +def test_freeze_preserves_reverse_argument_order_for_rows_and_costs() -> None: + class ReverseFreeze(Expr): + def __init__(self, value: i64Like) -> None: ... + + @method(reverse_args=True) + def label(self, value: StringLike) -> i64: ... + + expression = ReverseFreeze(1).label("key") + egraph = EGraph(set_(expression).to(i64(3)), set_cost(expression, 7)) + + replayed = eval(str(egraph.freeze()).removesuffix(".freeze()"), globals(), locals()) + + assert isinstance(replayed, EGraph) + assert replayed.lookup_function_value(expression) == i64(3) + assert replayed.lookup_function_value(get_cost(expression)) == i64(7) + + def test_dynamic_cost_rejects_an_incompatible_overload_without_recording_it() -> None: egraph = EGraph() egraph.register(set_cost(i64(1) + i64(2), 5)) @@ -2356,6 +3059,62 @@ def __init__(self, value: i64Like) -> None: ... set_cost(Costed(1), -1) +@pytest.mark.parametrize("save_egglog_string", [False, True], ids=["direct", "saved"]) +def test_dynamic_cost_rejects_a_negative_computed_value_without_storing_it(*, save_egglog_string: bool) -> None: + class Costed(Expr): + def __init__(self, value: i64Like) -> None: ... + + egraph = EGraph(save_egglog_string=save_egglog_string) + with pytest.raises(EggSmolError, match="@validate-dynamic-cost"): + egraph.register(set_cost(Costed(1), i64(0) - i64(1))) + + assert egraph.function_size(Costed) == 1 + assert egraph.lookup_function_value(get_cost(Costed(1))) is None + assert egraph.extract(Costed(1)) == Costed(1) + if save_egglog_string: + egg_bindings.EGraph().parse_and_run_program(egraph.as_egglog_string) + + +def test_cost_models_ignore_negative_values_written_through_a_raw_cost_table() -> None: + class RawNegativeCost(Expr): + @method(egg_fn="raw-negative-cost-node") + def __init__(self, value: i64Like) -> None: ... + + @function(egg_fn="cost_table_raw-negative-cost-node") + def raw_cost(value: i64Like) -> i64: ... + + expression = RawNegativeCost(1) + egraph = EGraph( + set_(raw_cost(1)).to(i64(-7)), + set_cost(RawNegativeCost(2), 5), + expression, + ) + expected = expression, 2 + + assert egraph.extract(expression, include_cost=True) == expected + assert egraph.extract(expression, include_cost=True, cost_model=default_cost_model) == expected + assert ( + egraph.extract( + expression, + include_cost=True, + cost_model=DagCostModel( + marginal_cost=lambda callback_egraph, node: default_cost_model(callback_egraph, node, []), + identity=0, + ), + extractor="greedy-dag", + ) + == expected + ) + + rendered = str(egraph.freeze()) + assert "set_(raw_cost(1)).to(i64(-7))" in rendered + assert "set_cost(RawNegativeCost(1), -7)" not in rendered + replayed = eval(rendered.removesuffix(".freeze()"), globals(), locals()) + assert isinstance(replayed, EGraph) + assert replayed.lookup_function_value(raw_cost(1)) == i64(-7) + assert replayed.extract(expression, include_cost=True) == expected + + class TestScheduler: def test_seq_schedule_decls_track_ruleset_updates(self): egraph = EGraph() @@ -2363,7 +3122,7 @@ def test_seq_schedule_decls_track_ruleset_updates(self): rel = relation("rel_live", i64) live_rules = ruleset(name="live-rules") schedule = seq(live_rules, run()).saturate() - _ = schedule.__egg_decls__ + _ = str(schedule) live_rules.register(rule(rel(i64(0))).then(rel(i64(1)))) @@ -2512,13 +3271,6 @@ def test_persistent_scheduler_is_saved_once_across_runs(self): assert len(scheduler_lines) == 1 assert len(run_with_lines) == 2 - def test_persistent_scheduler_gets_a_fresh_identity(self): - scheduler = back_off(match_limit=2, ban_length=2) - - persistent = scheduler.persistent() - - assert persistent.scheduler.id != scheduler.scheduler.id - def test_scheduler_scope_does_not_leak_to_sequence_sibling(self): r = ruleset(name="scheduler-lexical-scope") scheduler = back_off(match_limit=2, ban_length=2) @@ -2528,7 +3280,7 @@ def test_scheduler_scope_does_not_leak_to_sequence_sibling(self): run_schedule = next(line for line in egraph.as_egglog_string.splitlines() if line.startswith("(run-schedule ")) assert run_schedule.count("(let-scheduler ") == 2 - ruleset_name = str(r.__egg_ident__) + ruleset_name = f"{__name__}.scheduler-lexical-scope" assert f"(run-with _scheduler_0 {ruleset_name})" in run_schedule assert f"(run-with _scheduler_1 {ruleset_name})" in run_schedule @@ -2651,6 +3403,36 @@ def eager(value: i64Like) -> i64: EGraph().keep_best(eager) +def test_keep_best_invalidates_opaque_values_until_a_parent_scope_is_restored() -> None: + class CompactKey(Expr): + def __init__(self, value: i64Like) -> None: ... + + @function + def score(key: CompactKey) -> i64: ... + + egraph = EGraph(CompactKey(1), set_(score(CompactKey(1))).to(i64(11))) + parent_key = egraph.lookup_function_value(CompactKey(1)) + assert parent_key is not None + + egraph.push() + egraph.keep_best(score) + fresh_key = egraph.lookup_function_value(CompactKey(1)) + assert fresh_key is not None + assert egraph.lookup_function_value(score(fresh_key)) == i64(11) + with pytest.raises(ValueError, match="inactive push scope"): + egraph.lookup_function_value(score(parent_key)) + + egraph.pop() + assert egraph.lookup_function_value(score(parent_key)) == i64(11) + + egraph.keep_best(score) + with pytest.raises(ValueError, match="inactive push scope"): + egraph.lookup_function_value(score(parent_key)) + root_fresh_key = egraph.lookup_function_value(CompactKey(1)) + assert root_fresh_key is not None + assert egraph.lookup_function_value(score(root_fresh_key)) == i64(11) + + class TestCustomExtract: def test_literal_root(self) -> None: def is_even_cost_model(egraph: EGraph, expr: BaseExpr, children_costs: list[int]) -> int: @@ -2925,6 +3707,27 @@ def bin(l: E, r: E) -> E: ... assert cost == 4 assert expr == res + def test_dag_cost_model_supports_nonzero_typed_identity(self) -> None: + @dataclass(frozen=True, order=True) + class OffsetCost: + value: int + + def __add__(self, other: OffsetCost) -> OffsetCost: + return OffsetCost(self.value + other.value - 10) + + model = DagCostModel( + marginal_cost=lambda _egraph, _expr: OffsetCost(11), + identity=OffsetCost(10), + ) + egraph = EGraph() + expr = ff(1, 1) + + assert egraph.extract(expr, include_cost=True, cost_model=model) == (expr, OffsetCost(13)) + assert egraph.extract(expr, include_cost=True, cost_model=model, extractor="greedy-dag") == ( + expr, + OffsetCost(12), + ) + @pytest.mark.parametrize( ("model_kind", "extractor", "expected_cost"), [ @@ -3028,6 +3831,100 @@ def lookup_cost(callback_egraph: EGraph, expr: BaseExpr, children_costs: list[in with pytest.raises(ValueError, match="must be registered before extraction starts"): EGraph().extract(LookupByPrimitive(3), cost_model=cast("TreeCostModel[int]", lookup_cost)) + def test_cost_model_callback_rejects_value_from_another_egraph(self) -> None: + class Key(Expr): + def __init__(self, value: i64Like) -> None: ... + + class Root(Expr): + def __init__(self) -> None: ... + + @function + def score(value: Key) -> i64: ... + + foreign_egraph = EGraph(Key(3)) + foreign_key = foreign_egraph.lookup_function_value(Key(3)) + assert foreign_key is not None + + egraph = EGraph() + egraph.register(set_(score(Key(3))).to(i64(17))) + + def lookup_cost(callback_egraph: EGraph, expr: BaseExpr, children_costs: list[int]) -> int: + if isinstance(expr, Root): + callback_egraph.lookup_function_value(score(foreign_key)) + return sum(children_costs) + + with pytest.raises(ValueError, match="only look up tables using values supplied to the callback"): + egraph.extract(Root(), cost_model=cast("TreeCostModel[int]", lookup_cost)) + + def test_lookup_rejects_value_from_another_egraph(self) -> None: + class Key(Expr): + def __init__(self, value: i64Like) -> None: ... + + @function + def score(value: Key) -> i64: ... + + foreign_egraph = EGraph(Key(3)) + foreign_key = foreign_egraph.lookup_function_value(Key(3)) + assert foreign_key is not None + + egraph = EGraph() + egraph.register(set_(score(Key(3))).to(i64(17))) + + with pytest.raises(ValueError, match="belongs to a different EGraph"): + egraph.lookup_function_value(score(foreign_key)) + + def test_lookup_values_follow_push_scope_lifetimes(self) -> None: + class Key(Expr): + def __init__(self, value: i64Like) -> None: ... + + @function + def score(value: Key) -> i64: ... + + egraph = EGraph(Key(1), set_(score(Key(1))).to(i64(11))) + parent_key = egraph.lookup_function_value(Key(1)) + assert parent_key is not None + + egraph.push() + assert egraph.lookup_function_value(score(parent_key)) == i64(11) + egraph.register(Key(2), set_(score(Key(2))).to(i64(22))) + child_key = egraph.lookup_function_value(Key(2)) + assert child_key is not None + egraph.pop() + + assert egraph.lookup_function_value(score(parent_key)) == i64(11) + egraph.register(Key(3), set_(score(Key(3))).to(i64(33))) + with pytest.raises(ValueError, match="inactive push scope"): + egraph.lookup_function_value(score(child_key)) + + def test_cost_model_callback_can_lookup_nested_value_from_container_argument(self) -> None: + class Key(Expr): + def __init__(self, value: i64Like) -> None: ... + + class Root(Expr): + def __init__(self, values: Map[Key, i64]) -> None: ... + + @function + def score(value: Key) -> i64: ... + + egraph = EGraph() + egraph.register(set_(score(Key(3))).to(i64(17))) + + def lookup_cost(callback_egraph: EGraph, expr: BaseExpr, children_costs: list[int]) -> int: + if not isinstance(expr, Root): + return sum(children_costs) + args = get_callable_args(expr) + assert args is not None + (key,) = cast("Map[Key, i64]", args[0]).value + value = callback_egraph.lookup_function_value(score(key)) + assert value is not None + return int(value) + sum(children_costs) + + expr = Root(Map[Key, i64].empty().insert(Key(3), i64(0))) + assert egraph.extract(expr, include_cost=True, cost_model=cast("TreeCostModel[int]", lookup_cost)) == ( + expr, + 17, + ) + def test_cost_model_callback_values_are_scoped_to_their_egraph(self) -> None: class LookupByString(Expr): def __init__(self, value: StringLike) -> None: ... @@ -3055,6 +3952,44 @@ def lookup_cost(_callback_egraph: EGraph, expr: BaseExpr, children_costs: list[i LookupByString("needle"), include_cost=True, cost_model=cast("TreeCostModel[int]", lookup_cost) ) == (LookupByString("needle"), 17) + def test_nested_cost_model_callback_preserves_outer_egraph_values(self) -> None: + class Outer(Expr): + def __init__(self, value: i64Like) -> None: ... + + class Inner(Expr): + def __init__(self) -> None: ... + + @function + def score(value: i64Like) -> i64: ... + + outer_egraph = EGraph() + outer_egraph.register(set_(score(3)).to(i64(17))) + inner_egraph = EGraph() + + def outer_cost(callback_egraph: EGraph, expr: BaseExpr, children_costs: list[int]) -> int: + if not isinstance(expr, Outer): + return sum(children_costs) + args = get_callable_args(expr) + assert args is not None + outer_value = cast("i64", args[0]) + + def inner_cost(_inner_egraph: EGraph, inner_expr: BaseExpr, inner_children: list[int]) -> int: + if not isinstance(inner_expr, Inner): + return sum(inner_children) + value = callback_egraph.lookup_function_value(score(outer_value)) + assert value is not None + return int(value) + + _, cost = inner_egraph.extract( + Inner(), include_cost=True, cost_model=cast("TreeCostModel[int]", inner_cost) + ) + return cost + sum(children_costs) + + assert outer_egraph.extract(Outer(3), include_cost=True, cost_model=cast("TreeCostModel[int]", outer_cost)) == ( + Outer(3), + 17, + ) + @pytest.mark.parametrize( ("model_kind", "extractor"), [ diff --git a/python/tests/test_py_object_sort.py b/python/tests/test_py_object_sort.py index 18cb7277..79a48f16 100644 --- a/python/tests/test_py_object_sort.py +++ b/python/tests/test_py_object_sort.py @@ -2,9 +2,11 @@ import dataclasses import json +import threading from base64 import standard_b64decode, standard_b64encode from typing import TYPE_CHECKING +import pytest from cloudpickle import dumps, loads from egglog.bindings import * @@ -73,6 +75,14 @@ def my_add(a, b): return a + b +def _raise_worker_error(caller_thread_id: int) -> object: + if threading.get_ident() == caller_thread_id: + message = "primitive did not run on a worker thread" + raise AssertionError(message) + message = "parallel primitive failed" + raise ValueError(message) + + class TestEval: def test_eval(self): egraph = EGraph() @@ -150,6 +160,36 @@ def test_call(): ) +@pytest.mark.parametrize( + "rule", + [ + "(rule ((target fn arg) (= value (py-call fn arg))) ((hit arg)))", + "(rule ((target fn arg)) ((hit (py-call fn arg))))", + ], + ids=["query", "action"], +) +def test_parallel_py_object_primitive_error_is_propagated(rule: str): + egraph = EGraph(num_threads=2) + # Cross Egglog's default 10,000-row cutoff so rule evaluation uses its worker pool. + noise = "\n".join(f"(noise {i})" for i in range(10_001)) + egraph.parse_and_run_program( + f""" + (relation noise (i64)) + {noise} + (relation target (PyObject PyObject)) + (target {py_object_to_expr(_raise_worker_error)} {py_object_to_expr(threading.get_ident())}) + (relation hit (PyObject)) + """ + ) + + with pytest.raises(ValueError, match="parallel primitive failed"): + egraph.parse_and_run_program(f"{rule} (run 1)") + + # Draining the worker error leaves the same graph usable by later Python primitives. + recovered = py_object_to_expr("recovered") + egraph.parse_and_run_program(f'(check (= (py-to-string {recovered}) "recovered"))') + + def test_serialize_string(): """ Verify that when serializing the e-graph, PyObjects are turned into their repr diff --git a/src/conversions.rs b/src/conversions.rs index b5a5fab5..14fa44d6 100644 --- a/src/conversions.rs +++ b/src/conversions.rs @@ -383,7 +383,7 @@ convert_enums!( PrintOverallStatistics(span: Span, file: Option) c -> egglog::ast::Command::PrintOverallStatistics( c.span.clone().into(), - c.file.as_ref().map(|f| f.clone().into()) + c.file.clone() ), egglog::ast::Command::PrintOverallStatistics(span, file) => PrintOverallStatistics { span: span.into(), @@ -764,7 +764,7 @@ convert_struct!( }, r -> RuleReport { plan: r.plan.as_ref().map(|p| p.clone().into()), - search_and_apply_time: r.search_and_apply_time.clone().into(), + search_and_apply_time: r.search_and_apply_time.into(), num_matches: r.num_matches }; egglog_reports::RuleSetReport: "{:?}" => RuleSetReport( @@ -800,8 +800,8 @@ convert_struct!( ) }) .collect(), - search_and_apply_time: r.search_and_apply_time.clone().into(), - merge_time: r.merge_time.clone().into(), + search_and_apply_time: r.search_and_apply_time.into(), + merge_time: r.merge_time.into(), }; egglog_reports::IterationReport: "{:?}" => IterationReport( rule_set_report: RuleSetReport, @@ -813,7 +813,7 @@ convert_struct!( }, r -> IterationReport { rule_set_report: (&r.rule_set_report).into(), - rebuild_time: r.rebuild_time.clone().into() + rebuild_time: r.rebuild_time.into() }; egglog_reports::RunReport: "{:?}" => RunReport( iterations: Vec, @@ -1051,6 +1051,8 @@ impl MultiExtractOutput { impl UserDefinedCommandOutput { /// Return this output as a structured experimental multi-extraction, if it is one. fn as_multi_extract(&self) -> Option { + // Dispatch through the wrapped trait object (rather than Arc's blanket + // Any implementation), then clone so the Python result outlives this borrow. self.0 .as_ref() .as_any() diff --git a/src/egraph.rs b/src/egraph.rs index 6fdfdc6f..10511af6 100644 --- a/src/egraph.rs +++ b/src/egraph.rs @@ -3,7 +3,7 @@ use crate::conversions::*; use crate::error::{EggResult, WrappedError}; use crate::freeze::FrozenEGraph; -use crate::py_object_sort::{PyObjectSort, PyPickledValue, load}; +use crate::py_object_sort::{PyObjectErrorState, PyObjectSort, PyPickledValue, load}; use crate::serialize::SerializedEGraph; use crate::termdag::TermDag; use crate::tracing_otel; @@ -14,8 +14,48 @@ use log::info; use num_rational::{BigRational, Rational64}; use pyo3::prelude::*; use std::collections::{BTreeMap, BTreeSet}; +use std::panic::{AssertUnwindSafe, catch_unwind, resume_unwind}; use std::path::PathBuf; +fn run_with_py_error(py_error: &PyObjectErrorState, f: impl FnOnce() -> T) -> PyResult { + py_error.lock().unwrap().take(); + let result = catch_unwind(AssertUnwindSafe(f)); + if let Some(error) = py_error.lock().unwrap().take() { + return Err(error); + } + match result { + Ok(value) => Ok(value), + Err(error) => resume_unwind(error), + } +} + +fn error_is_replayable_by_fail(command: &egglog::ast::Command, error: &egglog::Error) -> bool { + use egglog::Error; + use egglog::ast::Command; + + matches!( + (command, error), + ( + Command::Check(..), + Error::CheckError(..) | Error::BackendError(_) + ) | ( + Command::Action(_), + Error::BackendError(_) | Error::SubsumeMergeError(..) + ) | ( + Command::RunSchedule(_), + Error::BackendError(_) | Error::NoSuchRuleset(..) | Error::ParseError(_) + ) | ( + Command::Extract(..), + Error::ExtractError(_) | Error::BackendError(_) + ) | (Command::Pop(..), Error::Pop(_)) + | (Command::Fail(..), Error::ExpectFail(_)) + // Registered commands are resolved and typechecked by their + // implementation during `run_command`, so any error from a sole + // user-defined command can be reproduced by wrapping it in `fail`. + | (Command::UserDefined(..), _) + ) +} + /// EGraph() /// -- /// @@ -24,6 +64,7 @@ use std::path::PathBuf; pub struct EGraph { pub(crate) egraph: egglog::EGraph, cmds: Option, + py_error: PyObjectErrorState, } impl EGraph { @@ -33,34 +74,40 @@ impl EGraph { commands: Vec, parsed_from_source: bool, ) -> EggResult> { + let failed_command = (commands.len() == 1).then(|| commands[0].clone()); let cmds_str = commands .iter() .map(|command| format!("{command}\n")) .collect::(); - let res = if parsed_from_source { - let span = tracing::info_span!( + let span = if parsed_from_source { + tracing::info_span!( "bindings.parse_and_run_program", command_count = commands.len(), commands = tracing::field::display(cmds_str.trim_end()) - ); - let _entered = span.enter(); - info!("Running commands:\n{}", cmds_str); - py.detach(|| self.egraph.run_program(commands)) + ) } else { - let span = tracing::info_span!( + tracing::info_span!( "bindings.run_program", command_count = commands.len(), commands = tracing::field::display(cmds_str.trim_end()) - ); - let _entered = span.enter(); - info!("Running commands:\n{}", cmds_str); - py.detach(|| self.egraph.run_program(commands)) + ) }; - if let Some(err) = PyErr::take(py) { - return Err(WrappedError::Py(err)); - } + let _entered = span.enter(); + info!("Running commands:\n{}", cmds_str); + let res = run_with_py_error(&self.py_error, || { + py.detach(|| self.egraph.run_program(commands)) + })?; match res { - Err(e) => Err(WrappedError::Egglog(e)), + Err(error) => { + if failed_command + .as_ref() + .is_some_and(|command| error_is_replayable_by_fail(command, &error)) + { + Err(WrappedError::ReplayableEgglog(error)) + } else { + Err(WrappedError::Egglog(error)) + } + } Ok(outputs) => { if let Some(cmds) = &mut self.cmds { cmds.push_str(&cmds_str); @@ -87,10 +134,19 @@ impl EGraph { egraph.seminaive = seminaive; egraph.set_num_threads(num_threads); egraph.no_decomp = no_decomp; - add_base_sort(&mut egraph, PyObjectSort {}, span!()).unwrap(); + let py_error = PyObjectErrorState::default(); + add_base_sort( + &mut egraph, + PyObjectSort { + py_error: py_error.clone(), + }, + span!(), + ) + .unwrap(); Self { egraph, cmds: record.then(String::new), + py_error, } } @@ -178,7 +234,7 @@ impl EGraph { include_temporary_functions: bool, traceparent: Option, tracestate: Option, - ) -> SerializedEGraph { + ) -> EggResult { let _context_guard = tracing_otel::attach_parent_context(traceparent.as_deref(), tracestate.as_deref()); let span = tracing::info_span!( @@ -186,24 +242,26 @@ impl EGraph { root_eclass_count = root_eclasses.len() ); let _entered = span.enter(); - Python::attach(|py| { - py.detach(|| { - let root_eclasses: Vec<_> = root_eclasses - .into_iter() - .map(|x| self.egraph.eval_expr(&egglog::ast::Expr::from(x)).unwrap()) - .collect(); - let res = self.egraph.serialize(SerializeConfig { - max_functions, - max_calls_per_function, - include_temporary_functions, - root_eclasses, - }); - SerializedEGraph { - egraph: res.egraph, - truncated_functions: res.truncated_functions, - discarded_functions: res.discarded_functions, - } + let res = Python::attach(|py| { + run_with_py_error(&self.py_error, || { + py.detach(|| -> Result<_, egglog::Error> { + let root_eclasses: Vec<_> = root_eclasses + .into_iter() + .map(|x| self.egraph.eval_expr(&egglog::ast::Expr::from(x))) + .collect::>()?; + Ok(self.egraph.serialize(SerializeConfig { + max_functions, + max_calls_per_function, + include_temporary_functions, + root_eclasses, + })) + }) }) + })??; + Ok(SerializedEGraph { + egraph: res.egraph, + truncated_functions: res.truncated_functions, + discarded_functions: res.discarded_functions, }) } @@ -226,8 +284,9 @@ impl EGraph { Ok(value.map(Value)) } - /// Extract `value` using its runtime sort. `sort` must match the sort returned with `value` - /// by `eval_expr`; passing a different existing sort is unsupported. + /// Extract `value` using its runtime sort. The value must come from `eval_expr` + /// on this e-graph, and `sort` must be the sort returned with it. Values from + /// another e-graph or paired with another existing sort are unsupported. fn extract_value(&self, value: Value, sort: &str) -> EggResult<(TermDag, usize, u64)> { let sort = self.egraph.get_sort_by_name(sort).ok_or_else(|| { WrappedError::Egglog(egglog::TypeError::UndefinedSort(sort.to_owned(), span!()).into()) @@ -249,16 +308,10 @@ impl EGraph { let span = tracing::info_span!("bindings.eval_expr"); let _entered = span.enter(); let expr: egglog::ast::Expr = expr.into(); - let res = py.detach(|| { - self.egraph - .eval_expr(&expr) - .map(|(s, v)| (s.name().to_string(), Value(v))) - .map_err(|e| WrappedError::Egglog(e)) - }); - if let Some(err) = PyErr::take(py) { - return Err(WrappedError::Py(err)); - } - res + let res = run_with_py_error(&self.py_error, || { + py.detach(|| self.egraph.eval_expr(&expr)) + })??; + Ok((res.0.name().to_string(), Value(res.1))) } fn value_to_i64(&self, v: Value) -> i64 { @@ -353,3 +406,34 @@ impl EGraph { #[derive(Ord, PartialOrd, Eq, PartialEq, Hash, Debug, Clone)] #[pyclass(eq, frozen, ord, hash, str = "{0:?}")] pub struct Value(pub egglog::Value); + +#[cfg(test)] +mod tests { + use super::*; + use pyo3::exceptions::PyValueError; + + #[test] + fn python_error_from_worker_precedes_backend_panic() { + Python::initialize(); + let error = PyObjectErrorState::default(); + *error.lock().unwrap() = Some(PyValueError::new_err("stale")); + + let worker_error = error.clone(); + let result = run_with_py_error(&error, || { + assert!(worker_error.lock().unwrap().is_none()); + std::thread::spawn(move || { + *worker_error.lock().unwrap() = Some(PyValueError::new_err("worker boom")); + }) + .join() + .unwrap(); + panic!("secondary backend panic"); + }); + + let captured = result.unwrap_err(); + assert!(error.lock().unwrap().is_none()); + Python::attach(|py| { + assert!(captured.is_instance_of::(py)); + assert_eq!(captured.value(py).to_string(), "worker boom"); + }); + } +} diff --git a/src/error.rs b/src/error.rs index 5e225975..1fff9946 100644 --- a/src/error.rs +++ b/src/error.rs @@ -6,13 +6,25 @@ use pyo3::prelude::*; pub struct EggSmolError { #[pyo3(get)] context: String, + /// Whether replaying the failed command inside `(fail ...)` preserves the + /// command's partial effects. This is deliberately conservative. + #[pyo3(get)] + replayable_by_fail: bool, } #[pymethods] impl EggSmolError { #[new] - fn new(context: String) -> Self { - EggSmolError { context } + #[pyo3(signature = (context, replayable_by_fail=false))] + fn new(context: String, replayable_by_fail: bool) -> Self { + EggSmolError { + context, + replayable_by_fail, + } + } + + fn __str__(&self) -> &str { + &self.context } } @@ -22,6 +34,7 @@ impl EggSmolError { // TODO: Create classes for each of these errors pub enum WrappedError { Egglog(egglog::Error), + ReplayableEgglog(egglog::Error), ParseError(egglog::ast::ParseError), Py(PyErr), } @@ -30,9 +43,16 @@ pub enum WrappedError { impl From for PyErr { fn from(error: WrappedError) -> Self { match error { - WrappedError::Egglog(error) => PyErr::new::(error.to_string()), + WrappedError::Egglog(error) => { + PyErr::new::((error.to_string(), false)) + } + WrappedError::ReplayableEgglog(error) => { + PyErr::new::((error.to_string(), true)) + } WrappedError::Py(error) => error, - WrappedError::ParseError(error) => PyErr::new::(error.to_string()), + WrappedError::ParseError(error) => { + PyErr::new::((error.to_string(), false)) + } } } } diff --git a/src/extract.rs b/src/extract.rs index c8689335..d2ff0884 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -37,16 +37,6 @@ fn catch_python_cost_error(f: impl FnOnce() -> T) -> PyResult { #[derive(Debug)] struct Cost(Arc>); -impl Cost { - fn from_py(value: Py) -> Self { - Self(Arc::new(value)) - } - - fn to_py(&self, py: Python<'_>) -> Py { - self.0.as_ref().clone_ref(py) - } -} - impl Ord for Cost { fn cmp(&self, other: &Self) -> Ordering { Python::attach(|py| python_or_unwind(self.0.bind(py).compare(other.0.bind(py)))) @@ -143,12 +133,12 @@ impl EggTreeCostModel for CostModel { Python::attach(|py| { let children_cost = children_cost .iter() - .map(|cost| cost.to_py(py)) + .map(|cost| cost.0.as_ref().clone_ref(py)) .collect::>(); - Cost::from_py(python_or_unwind( - self.fold - .call1(py, (enode_cost.head, enode_cost.annotation, children_cost)), - )) + Cost(Arc::new(python_or_unwind(self.fold.call1( + py, + (enode_cost.head, enode_cost.annotation, children_cost), + )))) }) } @@ -179,12 +169,12 @@ impl EggTreeCostModel for CostModel { Python::attach(|py| { let element_costs = element_costs .iter() - .map(|cost| cost.to_py(py)) + .map(|cost| cost.0.as_ref().clone_ref(py)) .collect::>(); - Cost::from_py(python_or_unwind(self.container_cost.call1( + Cost(Arc::new(python_or_unwind(self.container_cost.call1( py, (container_cost.sort, container_cost.value, element_costs), - ))) + )))) }) } @@ -208,16 +198,16 @@ impl EggTreeCostModel for CostModel { value: egglog::Value, ) -> Cost { Python::attach(|py| { - Cost::from_py(python_or_unwind( + Cost(Arc::new(python_or_unwind( self.base_value_cost.call1(py, (sort.name(), Value(value))), - )) + ))) }) } } #[derive(Debug)] struct DagCostContext { - identity: Arc>, + identity: Py, } #[derive(Clone, Debug)] @@ -239,15 +229,11 @@ impl DagCost { fn to_py(&self, py: Python<'_>, context: &Arc) -> Py { match self { - Self::Identity => context.identity.as_ref().clone_ref(py), + Self::Identity => context.identity.clone_ref(py), Self::Value { value, .. } => value.as_ref().clone_ref(py), } } - fn compare_values(left: &Arc>, right: &Arc>) -> Ordering { - Python::attach(|py| python_or_unwind(left.bind(py).compare(right.bind(py)))) - } - fn ensure_same_context(left: &Arc, right: &Arc) { if !Arc::ptr_eq(left, right) { python_or_unwind::<()>(Err(PyValueError::new_err( @@ -272,14 +258,14 @@ impl Ord for DagCost { }, ) => { Self::ensure_same_context(left_context, right_context); - Self::compare_values(left, right) - } - (Self::Identity, Self::Value { value, context }) => { - Self::compare_values(&context.identity, value) - } - (Self::Value { value, context }, Self::Identity) => { - Self::compare_values(value, &context.identity) + Python::attach(|py| python_or_unwind(left.bind(py).compare(right.bind(py)))) } + (Self::Identity, Self::Value { value, context }) => Python::attach(|py| { + python_or_unwind(context.identity.bind(py).compare(value.bind(py))) + }), + (Self::Value { value, context }, Self::Identity) => Python::attach(|py| { + python_or_unwind(value.bind(py).compare(context.identity.bind(py))) + }), } } } @@ -332,7 +318,7 @@ impl MonoidCost for DagCost { #[derive(Debug)] #[pyclass( frozen, - str = "DagCostModel({identity:?}, {enode_cost:?}, {container_cost:?}, {base_value_cost:?}" + str = "DagCostModel({identity:?}, {enode_cost:?}, {container_cost:?}, {base_value_cost:?})" )] pub struct DagCostModel { identity: Py, @@ -363,7 +349,7 @@ impl DagCostModel { fn runtime(&self, py: Python<'_>) -> RuntimeDagCostModel { RuntimeDagCostModel { context: Arc::new(DagCostContext { - identity: Arc::new(self.identity.clone_ref(py)), + identity: self.identity.clone_ref(py), }), enode_cost: Arc::new(self.enode_cost.clone_ref(py)), container_cost: Arc::new(self.container_cost.clone_ref(py)), @@ -541,6 +527,10 @@ impl Extractor { } /// Extract the best term of a value from a given sort. + /// + /// `value` must come from `EGraph.eval_expr` on the supplied e-graph, and + /// `sort` must be the sort returned with it. Values from another e-graph or + /// paired with another existing sort are unsupported. #[pyo3(signature = (egraph, termdag, value, sort, *, traceparent=None, tracestate=None))] fn extract_best( &self, @@ -576,10 +566,14 @@ impl Extractor { .ok_or_else(|| PyValueError::new_err("unextractable root"))?; let (local_termdag, cost, term) = extracted; let term = copy_termdag(&local_termdag, &mut termdag.0)[term]; - Ok((cost.to_py(py), term)) + Ok((cost.0.as_ref().clone_ref(py), term)) } /// Extract variants of an e-class. + /// + /// `value` must come from `EGraph.eval_expr` on the supplied e-graph, and + /// `sort` must be the sort returned with it. Values from another e-graph or + /// paired with another existing sort are unsupported. #[pyo3(signature = (egraph, termdag, value, nvariants, sort, *, traceparent=None, tracestate=None))] fn extract_variants( &self, @@ -616,12 +610,16 @@ impl Extractor { let copied = copy_termdag(&local_termdag, &mut termdag.0); Ok(variants .into_iter() - .map(|variant| (variant.cost.to_py(py), copied[variant.term])) + .map(|variant| (variant.cost.0.as_ref().clone_ref(py), copied[variant.term])) .collect()) } } /// Extract the best term for each root with a custom additive marginal model. +/// +/// Every value must come from `EGraph.eval_expr` on the supplied e-graph and +/// be paired with the sort returned with it. Values from another e-graph or +/// paired with another existing sort are unsupported. #[pyfunction] #[pyo3(signature = (egraph, roots, cost_model, *, extractor="tree", traceparent=None, tracestate=None))] pub fn extract_best_with_dag_cost_model( diff --git a/src/py_object_sort.rs b/src/py_object_sort.rs index 1c5e78df..65a3755a 100644 --- a/src/py_object_sort.rs +++ b/src/py_object_sort.rs @@ -20,6 +20,7 @@ use std::{ fmt::Debug, fs::File, io::Write, + sync::{Arc, Mutex}, }; use uuid::Uuid; @@ -53,8 +54,17 @@ pub fn load<'py>(py: Python<'py>, pickled: &PyPickledValue) -> PyResult>>; + +pub struct PyObjectSort { + pub(crate) py_error: PyObjectErrorState, +} + +impl Debug for PyObjectSort { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("PyObjectSort").finish_non_exhaustive() + } +} impl BaseSort for PyObjectSort { type Base = PyPickledValue; @@ -72,8 +82,8 @@ impl BaseSort for PyObjectSort { } ); // Supports calling (py-eval ) - add_primitive!(eg, "py-eval" = |code: S, globals: PyPickledValue, locals: PyPickledValue| -?> PyPickledValue { - attach("py-eval", |py| { + add_primitive!(eg, "py-eval" = {self.py_error.clone(): PyObjectErrorState} |code: S, globals: PyPickledValue, locals: PyPickledValue| -?> PyPickledValue { + attach(&self.ctx, "py-eval", |py| { dump(py.eval( CString::new(code.to_string()).unwrap().as_c_str(), Some(load(py, &globals)?.cast::()?), @@ -86,8 +96,9 @@ impl BaseSort for PyObjectSort { add_primitive!( eg, "py-exec" = + {self.py_error.clone(): PyObjectErrorState} |code: S, globals: PyPickledValue, locals: PyPickledValue| -?> PyPickledValue { - attach("py-exec", |py| { + attach(&self.ctx, "py-exec", |py| { let locals = load(py, &locals)?; // Copy code into temporary file // Keep it around so that if errors occur we can debug them after the program exits @@ -109,8 +120,8 @@ impl BaseSort for PyObjectSort { } ); // (py-call []*) - add_primitive!(eg, "py-call" = [xs: PyPickledValue] -?> PyPickledValue { - attach("py-call", |py| { + add_primitive!(eg, "py-call" = {self.py_error.clone(): PyObjectErrorState} [xs: PyPickledValue] -?> PyPickledValue { + attach(&self.ctx, "py-call", |py| { let xs = xs.map(|x| load(py, &x)).collect::>>().map_err(|e| {e.add_note(py, "Loading arguments").unwrap(); e})?; let fn_obj = &xs[0]; let args = PyTuple::new(py, xs[1..].to_vec()).map_err(|e| {e.add_note(py, "Creating tuple").unwrap(); e})?; @@ -119,8 +130,8 @@ impl BaseSort for PyObjectSort { }); // (py-call-extended ) - add_primitive!(eg, "py-call-extended" = |fn_: PyPickledValue, args: PyPickledValue, kwargs: PyPickledValue| -?> PyPickledValue { - attach("py-call-extended", |py| { + add_primitive!(eg, "py-call-extended" = {self.py_error.clone(): PyObjectErrorState} |fn_: PyPickledValue, args: PyPickledValue, kwargs: PyPickledValue| -?> PyPickledValue { + attach(&self.ctx, "py-call-extended", |py| { let fn_ = load(py, &fn_)?; let args = load(py, &args)?; let kwargs = load(py, &kwargs)?; @@ -129,8 +140,8 @@ impl BaseSort for PyObjectSort { }); // (py-dict [ ]*) - add_primitive!(eg, "py-dict" = [xs: PyPickledValue] -?> PyPickledValue { - attach("py-dict", |py| { + add_primitive!(eg, "py-dict" = {self.py_error.clone(): PyObjectErrorState} [xs: PyPickledValue] -?> PyPickledValue { + attach(&self.ctx, "py-dict", |py| { let dict = PyDict::new(py); for i in xs.map(|x| load(py, &x)).collect::>>()?.chunks_exact(2) { dict.set_item(i[0].clone(), i[1].clone())?; @@ -139,8 +150,8 @@ impl BaseSort for PyObjectSort { }) }); // Supports calling (py-dict-update [ ]*) - add_primitive!(eg, "py-dict-update" = [xs: PyPickledValue] -?> PyPickledValue {{ - attach("py-dict-update", |py| { + add_primitive!(eg, "py-dict-update" = {self.py_error.clone(): PyObjectErrorState} [xs: PyPickledValue] -?> PyPickledValue {{ + attach(&self.ctx, "py-dict-update", |py| { let xs = xs.map(|x| load(py, &x)).collect::>>()?; // Copy the dict so we can mutate it and return it let dict = xs[0].cast::()?; @@ -154,9 +165,9 @@ impl BaseSort for PyObjectSort { // (py-to-string ) add_primitive!( eg, - "py-to-string" = |x: PyPickledValue| -?> S { + "py-to-string" = {self.py_error.clone(): PyObjectErrorState} |x: PyPickledValue| -?> S { { - let s: String = attach("py-to-string", move |py| load(py, &x)?.extract())?; + let s: String = attach(&self.ctx, "py-to-string", move |py| load(py, &x)?.extract())?; Some(s.into()) } } @@ -164,17 +175,17 @@ impl BaseSort for PyObjectSort { // (py-to-bool ) add_primitive!( eg, - "py-to-bool" = |x: PyPickledValue| -?> bool { + "py-to-bool" = {self.py_error.clone(): PyObjectErrorState} |x: PyPickledValue| -?> bool { { - attach("py-to-bool", move |py| load(py, &x)?.extract()) + attach(&self.ctx, "py-to-bool", move |py| load(py, &x)?.extract()) } } ); // (py-from-string ) add_primitive!( eg, - "py-from-string" = |x: S| -?> PyPickledValue { - attach("py-from-string", |py| { + "py-from-string" = {self.py_error.clone(): PyObjectErrorState} |x: S| -?> PyPickledValue { + attach(&self.ctx, "py-from-string", |py| { dump(x.to_string().into_pyobject(py)?) }) } @@ -182,8 +193,8 @@ impl BaseSort for PyObjectSort { // (py-from-int ) add_primitive!( eg, - "py-from-int" = |x: i64| -?> PyPickledValue { - attach("py-from-int", |py| { + "py-from-int" = {self.py_error.clone(): PyObjectErrorState} |x: i64| -?> PyPickledValue { + attach(&self.ctx, "py-from-int", |py| { dump(x.into_pyobject(py)?) }) } @@ -204,13 +215,13 @@ impl BaseSort for PyObjectSort { /// Attaches to the Python interpreter and runs the given closure. /// -/// Also handles errors, by saving them on the interpreter and returning None. -fn attach(name: &str, f: F) -> Option +/// Also handles errors, by saving the first one for the binding thread and returning None. +fn attach(py_error: &PyObjectErrorState, name: &str, f: F) -> Option where F: for<'py> FnOnce(Python<'py>) -> PyResult, { Python::attach(|py| { - if PyErr::occurred(py) { + if py_error.lock().unwrap().is_some() { return None; }; match f(py) { @@ -218,7 +229,10 @@ where Err(err) => { err.add_note(py, format!("While calling primitive '{}'", name)) .unwrap(); - err.restore(py); + let mut first_error = py_error.lock().unwrap(); + if first_error.is_none() { + *first_error = Some(err); + } None } } @@ -239,3 +253,49 @@ fn run_path<'py>( debug_assert!(obj.is_none()); }) } + +#[cfg(test)] +mod tests { + use super::*; + use pyo3::exceptions::PyValueError; + use std::sync::atomic::{AtomicBool, Ordering}; + + #[test] + fn shares_only_the_first_primitive_error_across_threads() { + Python::initialize(); + let error = PyObjectErrorState::default(); + let worker_error = error.clone(); + let called_after_error = Arc::new(AtomicBool::new(false)); + let worker_called_after_error = called_after_error.clone(); + + std::thread::spawn(move || { + let result: Option<()> = attach(&worker_error, "py-call", |_py| { + Err(PyValueError::new_err("worker boom")) + }); + assert!(result.is_none()); + + let result = attach(&worker_error, "py-call", |_py| { + worker_called_after_error.store(true, Ordering::Relaxed); + Ok(()) + }); + assert!(result.is_none()); + Python::attach(|py| assert!(!PyErr::occurred(py))); + }) + .join() + .unwrap(); + + assert!(!called_after_error.load(Ordering::Relaxed)); + let captured = error.lock().unwrap().take().unwrap(); + Python::attach(|py| { + assert!(captured.is_instance_of::(py)); + assert_eq!(captured.value(py).to_string(), "worker boom"); + let notes: Vec = captured + .value(py) + .getattr("__notes__") + .unwrap() + .extract() + .unwrap(); + assert_eq!(notes, ["While calling primitive 'py-call'"]); + }); + } +} diff --git a/test-data/unit/check-high-level.test b/test-data/unit/check-high-level.test index 020b2062..b78f0e0c 100644 --- a/test-data/unit/check-high-level.test +++ b/test-data/unit/check-high-level.test @@ -6,6 +6,11 @@ _ = i64(0) == i64(0) from egglog import * _ = i64(0) != i64(0) +[case parameterizedVar] +from egglog import * +value = var("value", Vec[i64]) +reveal_type(value) # N: Revealed type is "egglog.builtins.Vec[egglog.builtins.i64]" + [case eqToAllowed] from egglog import * _ = eq(i64(0)).to(i64(0)) From 1e81090581dafce13351b16af404d746dba017d0 Mon Sep 17 00:00:00 2001 From: Saul Shanabrook Date: Thu, 3 Sep 2026 16:49:43 -0700 Subject: [PATCH 6/6] Cache explicit backend names during declaration merges --- python/egglog/egraph.py | 3 +- python/egglog/egraph_state.py | 68 ++++++++++++++++++++++--------- python/tests/test_egraph_state.py | 23 +++++++++++ 3 files changed, 72 insertions(+), 22 deletions(-) diff --git a/python/egglog/egraph.py b/python/egglog/egraph.py index 4078251e..a2331e8c 100644 --- a/python/egglog/egraph.py +++ b/python/egglog/egraph.py @@ -1141,8 +1141,7 @@ def __init__( def _add_decls(self, *decls: DeclarationsLike) -> None: self._state.ensure_open() - for d in decls: - self._state.__egg_decls__ |= d + self._state.add_declarations(*decls) def set_report_level(self, level: bindings._ReportLevel) -> None: """ diff --git a/python/egglog/egraph_state.py b/python/egglog/egraph_state.py index b68d9ba2..c3f26ee1 100644 --- a/python/egglog/egraph_state.py +++ b/python/egglog/egraph_state.py @@ -137,6 +137,28 @@ def _normalize_global_let_name(name: str) -> str: return name if name.startswith("$") else f"${name}" +def _collect_explicit_backend_names(declarations: Declarations) -> set[str]: + """Collect names that declarations explicitly reserve in Egglog's shared symbol namespace.""" + names = { + decl.egg_name + for decl in (*declarations._functions.values(), *declarations._constants.values()) + if decl.egg_name is not None + } + for class_decl in declarations._classes.values(): + if class_decl.egg_name is not None: + names.add(class_decl.egg_name) + class_callables = ( + *class_decl.class_methods.values(), + *class_decl.class_variables.values(), + *class_decl.methods.values(), + *class_decl.properties.values(), + ) + names.update(decl.egg_name for decl in class_callables if decl.egg_name is not None) + if class_decl.init is not None and class_decl.init.egg_name is not None: + names.add(class_decl.init.egg_name) + return names + + def _egg_name_is_source_safe_symbol(name: str) -> bool: """Return whether a name is safe to emit as an ordinary Egglog symbol.""" return ( @@ -245,6 +267,11 @@ class EGraphState: egglog_file_state: _SavedEgglogFile | None = field(default=None, repr=False) # The declarations we have added. __egg_decls__: Declarations = field(default_factory=Declarations) + # Explicit names are cached as declarations are merged. Generated names + # consult this set so a registration batch reserves all user-selected + # backend symbols without rescanning the complete declaration graph for + # every generated callable, sort, cost table, and synthetic let. + _explicit_backend_names: set[str] = field(default_factory=set, init=False, repr=False) # Mapping of added rulesets to the added rules rulesets: dict[Ident, set[RewriteOrRuleDecl]] = field(default_factory=dict) # Persistent schedulers live outside a single run-schedule command; only emit their let once per active scope. @@ -293,6 +320,7 @@ class EGraphState: def __post_init__(self, save_egglog_string: bool) -> None: if not self.valid_value_owners: self.valid_value_owners = frozenset((self.value_owner,)) + self._explicit_backend_names = _collect_explicit_backend_names(self.__egg_decls__) if save_egglog_string and self.egglog_file_state is None: # Keep one persistent temp `.egg` file per high-level egraph so parse errors # can point at a stable filename the user can open after a failure. @@ -331,6 +359,24 @@ def copy(self) -> EGraphState: rule_name_to_command_decl=self.rule_name_to_command_decl.copy(), ) + def add_declarations(self, *declarations_like: DeclarationsLike) -> None: + """Merge declarations while maintaining the explicit backend-name index.""" + attempted_update = False + try: + for declarations in declarations_like: + if declarations is None: + continue + attempted_update = True + self.__egg_decls__ |= declarations + finally: + if attempted_update: + # Rebuild after the batch rather than accumulating names: declaration + # merges may replace an earlier declaration and release its explicit + # name for later generated symbols. The finally path also keeps the + # index aligned when a later lazy declaration fails to resolve after + # an earlier declaration has already been merged. + self._explicit_backend_names = _collect_explicit_backend_names(self.__egg_decls__) + def egglog_string(self) -> str: if self.egglog_file_state is None: msg = "Can't get egglog string unless EGraph created with save_egglog_string=True" @@ -1490,26 +1536,8 @@ def _allocate_name(self, candidate: str, *, avoid_reserved_call_heads: bool = Fa # All declarations for a register(...) batch are merged before any # command is lowered. Reserve their explicit backend names up front so # generated names do not depend on action order within that batch. - explicit_backend_names = { - decl.egg_name - for decl in (*self.__egg_decls__._functions.values(), *self.__egg_decls__._constants.values()) - if decl.egg_name is not None - } - for class_decl in self.__egg_decls__._classes.values(): - if class_decl.egg_name is not None: - explicit_backend_names.add(class_decl.egg_name) - class_callables = ( - *class_decl.class_methods.values(), - *class_decl.class_variables.values(), - *class_decl.methods.values(), - *class_decl.properties.values(), - ) - explicit_backend_names.update(decl.egg_name for decl in class_callables if decl.egg_name is not None) - if class_decl.init is not None and class_decl.init.egg_name is not None: - explicit_backend_names.add(class_decl.init.egg_name) - if ( - candidate not in explicit_backend_names + candidate not in self._explicit_backend_names and not self._backend_symbol_is_occupied(candidate) and (not avoid_reserved_call_heads or candidate not in _EGGLOG_RESERVED_CALL_HEADS) ): @@ -1517,7 +1545,7 @@ def _allocate_name(self, candidate: str, *, avoid_reserved_call_heads: bool = Fa index = 1 while ( - f"{candidate}_{index}" in explicit_backend_names + f"{candidate}_{index}" in self._explicit_backend_names or self._backend_symbol_is_occupied(f"{candidate}_{index}") or (avoid_reserved_call_heads and f"{candidate}_{index}" in _EGGLOG_RESERVED_CALL_HEADS) ): diff --git a/python/tests/test_egraph_state.py b/python/tests/test_egraph_state.py index 9ac8f4f5..b7c224d4 100644 --- a/python/tests/test_egraph_state.py +++ b/python/tests/test_egraph_state.py @@ -26,6 +26,7 @@ from egglog.declarations import ( ClassDecl, Declarations, + DelayedDeclarations, FunctionDecl, FunctionRef, FunctionSignature, @@ -251,6 +252,28 @@ def test_generated_names_are_fully_qualified() -> None: assert state.type_ref_to_egg(JustTypeRef(ret2)) == "pkg.two.Ret" +def test_explicit_backend_names_remain_reserved_after_a_later_declaration_fails() -> None: + state = EGraph(save_egglog_string=True)._state + reserved_name = "reserved_before_declaration_failure" + declarations = Declarations( + _functions={ + Ident("explicit_before_failure"): FunctionDecl( + signature=FunctionSignature(return_type=TypeRefWithVars(Ident.builtin("i64"))), + egg_name=reserved_name, + ) + } + ) + + def fail_to_resolve() -> Declarations: + msg = "declaration resolution failed" + raise RuntimeError(msg) + + with pytest.raises(RuntimeError, match="declaration resolution failed"): + state.add_declarations(declarations, DelayedDeclarations(fail_to_resolve)) + + assert state._allocate_name(reserved_name) == f"{reserved_name}_1" + + def test_missing_function_lookup_does_not_reserve_generated_name() -> None: state = EGraph(save_egglog_string=True)._state ret = Ident("LookupRet", "pkg.lookup")