diff --git a/.gitignore b/.gitignore index 4f82135..383ea93 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ /target/ /.reports*.jsonl +/.profiles/ /.bench-worktrees/ /.venv/ /.mypy_cache/ diff --git a/Cargo.lock b/Cargo.lock index f66efe3..0da4312 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -97,6 +97,15 @@ version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" +[[package]] +name = "bincode" +version = "1.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f45e9417d87227c7a56d22e471c6206462cba514c7590c09aff4cf6d1ddcad" +dependencies = [ + "serde", +] + [[package]] name = "bit-set" version = "0.8.0" @@ -148,6 +157,18 @@ version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" + +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + [[package]] name = "cc" version = "1.2.65" @@ -296,6 +317,38 @@ dependencies = [ "windows-sys 0.59.0", ] +[[package]] +name = "columnar" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94063ab5cdb61ed8e5a8294c3c7cc8466c1a9eafc7b9ca02b2a521ea06e52d92" +dependencies = [ + "bytemuck", + "columnar_derive", + "smallvec", +] + +[[package]] +name = "columnar_derive" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8d11c5dbac40214de1027e4783272d192198c986d94b65ae43c0886b4fb58ec5" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "columnation" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b90b1614014f6958477dcdb77a2d489555db48ca61efe94c5cde2b0446933ed1" +dependencies = [ + "paste", + "smallvec", +] + [[package]] name = "condtype" version = "1.3.0" @@ -454,6 +507,21 @@ dependencies = [ "thiserror", ] +[[package]] +name = "differential-dataflow" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "befdc46c04f5671a2dfe32a9d06ce8efe44f650f0f6f974d190b9a6ee1c17665" +dependencies = [ + "columnar", + "columnation", + "fnv", + "paste", + "serde", + "smallvec", + "timely", +] + [[package]] name = "digest" version = "0.10.7" @@ -507,6 +575,7 @@ dependencies = [ "dyn-clone", "egglog-add-primitive", "egglog-ast", + "egglog-backend-trait", "egglog-bridge", "egglog-core-relations", "egglog-numeric-id", @@ -549,6 +618,17 @@ dependencies = [ "ordered-float", ] +[[package]] +name = "egglog-backend-trait" +version = "2.0.0" +dependencies = [ + "anyhow", + "egglog-ast", + "egglog-bridge", + "egglog-core-relations", + "egglog-reports", +] + [[package]] name = "egglog-bridge" version = "2.0.0" @@ -624,6 +704,7 @@ dependencies = [ "codspeed-divan-compat", "egglog", "egglog-ast", + "egglog-experimental-dd", "egglog-reports", "glob", "insta", @@ -633,6 +714,26 @@ dependencies = [ "num", ] +[[package]] +name = "egglog-experimental-dd" +version = "0.1.0" +dependencies = [ + "anyhow", + "differential-dataflow", + "egglog", + "egglog-ast", + "egglog-backend-trait", + "egglog-core-relations", + "egglog-experimental", + "egglog-numeric-id", + "hashbrown 0.16.1", + "insta", + "libtest-mimic", + "log", + "serde", + "timely", +] + [[package]] name = "egglog-numeric-id" version = "2.0.0" @@ -840,6 +941,15 @@ dependencies = [ "version_check", ] +[[package]] +name = "getopts" +version = "0.2.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfe4fbac503b8d1f88e6676011885f34b7174f46e59956bba534ba83abded4df" +dependencies = [ + "unicode-width", +] + [[package]] name = "getrandom" version = "0.2.17" @@ -1246,6 +1356,12 @@ dependencies = [ "windows-link", ] +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + [[package]] name = "pest" version = "2.8.6" @@ -1675,6 +1791,9 @@ name = "smallvec" version = "1.15.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" +dependencies = [ + "serde", +] [[package]] name = "statrs" @@ -1766,6 +1885,62 @@ dependencies = [ "syn 2.0.118", ] +[[package]] +name = "timely" +version = "0.30.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c5af6859fde675a87f12da672e0bdf08bfc809e68d8c6728f50942d09c0480a3" +dependencies = [ + "bincode", + "byteorder", + "columnar", + "columnation", + "getopts", + "itertools", + "serde", + "smallvec", + "timely_bytes", + "timely_communication", + "timely_container", + "timely_logging", +] + +[[package]] +name = "timely_bytes" +version = "0.30.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b427b919d60d815c9c30b52d8080ebde599380db18c7c98a8c99fc172601ca1" + +[[package]] +name = "timely_communication" +version = "0.30.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2137e50f9c368c04d9971b7909c6fceadedb24c6e4fa4173bf55bf6d415329db" +dependencies = [ + "byteorder", + "columnar", + "getopts", + "serde", + "timely_bytes", + "timely_container", + "timely_logging", +] + +[[package]] +name = "timely_container" +version = "0.30.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b45d41809553b99db9abc2ad2c5f5b8f755df1f2191051527deaadeea12d5a7e" + +[[package]] +name = "timely_logging" +version = "0.30.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00582f3cfc652e3173eb9a140981133f359029e6b746667414a8d334285163d6" +dependencies = [ + "timely_container", +] + [[package]] name = "toml_datetime" version = "1.1.1+spec-1.1.0" @@ -1820,6 +1995,12 @@ version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" +[[package]] +name = "unicode-width" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" + [[package]] name = "utf8parse" version = "0.2.2" diff --git a/Cargo.toml b/Cargo.toml index 34895e3..ef12034 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,6 +5,7 @@ members = [ "egglog/concurrency", "egglog/core-relations", "egglog/egglog-ast", + "egglog/egglog-backend-trait", "egglog/egglog-bridge", "egglog/egglog-reports", "egglog/numeric-id", @@ -12,6 +13,7 @@ members = [ "egglog/src/sort/add_primitive", "egglog/wasm-example", "egglog-experimental", + "egglog-experimental/dd", ] [workspace.package] @@ -81,6 +83,7 @@ egglog-numeric-id = { path = "egglog/numeric-id", version = "2.0.0" } egglog-union-find = { path = "egglog/union-find", version = "2.0.0" } egglog-add-primitive = { path = "egglog/src/sort/add_primitive", version = "2.0.0" } egglog-reports = { path = "egglog/egglog-reports", version = "2.0.0" } +egglog-backend-trait = { path = "egglog/egglog-backend-trait", version = "2.0.0" } egglog = { path = "egglog", default-features = false, version = "2.0.0" } ###################### diff --git a/bench.py b/bench.py index 93db46f..6dbbc94 100755 --- a/bench.py +++ b/bench.py @@ -7,11 +7,14 @@ import os import re import resource +import shlex +import shutil import subprocess import sys import tempfile import time -from collections.abc import Mapping, Sequence +import uuid +from collections.abc import Callable, Mapping, Sequence from contextlib import suppress from dataclasses import dataclass, replace from datetime import UTC, datetime @@ -30,6 +33,7 @@ MofNCompleteColumn, Progress, SpinnerColumn, + TaskID, TextColumn, TimeElapsedColumn, ) @@ -38,16 +42,43 @@ from rich.tree import Tree from scipy import stats +import samply_analysis + Status = Literal["success", "timed-out", "failure"] +Backend = str Treatment = Literal["off", "term", "proofs"] +BuildProfile = Literal["release", "profiling"] +OutputFormat = Literal["rich", "markdown"] +TableAlignment = Literal["left", "right"] + + +@dataclass(frozen=True) +class BackendSpec: + display_name: str + treatments: tuple[Treatment, ...] + flags: tuple[str, ...] + cargo_features: tuple[str, ...] = () + + +BACKEND_SPECS: dict[Backend, BackendSpec] = { + "main": BackendSpec("main", ("off", "term", "proofs"), ()), + "dd": BackendSpec("DD", ("term", "proofs"), ("--backend", "dd"), ("dd-backend",)), +} DEFAULT_REPORT = ".reports.jsonl" +DEFAULT_PROFILES_DIR = ".profiles" DEFAULT_ROUNDS = 6 DEFAULT_TIMEOUT_SEC = 120 +DEFAULT_PROFILE_SECONDS = 10 +DEFAULT_PROFILE_TOP = 15 +MAX_PROFILE_ITERATIONS = 10_000 +PROFILE_CACHE_VERSION = "v1" +PROFILE_SAMPLY_RATE_HZ = 1000 TARGET_STARTUP_WARMUP_SUBPROCESSES = 1 +DEFAULT_BACKENDS: tuple[Backend, ...] = ("main",) DEFAULT_TREATMENTS: tuple[Treatment, ...] = ("off", "term", "proofs") DEFAULT_FILES = ( - "egglog/tests/math-microbenchmark.egg", + "egglog/tests/math-microbenchmark-mini.egg", "egglog/tests/web-demo/rw-analysis.egg", "egglog/tests/integer_math.egg", "egglog/tests/web-demo/resolution.egg", @@ -61,7 +92,15 @@ "Ratio is target / baseline. Values below 1x use less peak RSS; above 1x use more. " "RSS change is derived from the ratio; negative uses less memory. Intervals are 95% CIs." ) -PROOF_OVERHEAD_CAPTION = "Within-target proof overhead. This is separate from target-vs-baseline wall-time change." +BACKEND_WALL_TIME_CAPTION = ( + "Ratio is candidate backend / baseline backend for the same target and treatment. " + "Values below 1x are faster; above 1x are slower. Intervals are 95% CIs." +) +BACKEND_PEAK_RSS_CAPTION = ( + "Ratio is candidate backend / baseline backend for the same target and treatment. " + "Values below 1x use less peak RSS; above 1x use more. Intervals are 95% CIs." +) +PROOF_OVERHEAD_CAPTION = "Within-backend proof overhead. This is not backend-vs-main performance." RESULT_STYLES = { "descriptive": "dim", "established": "green", @@ -119,6 +158,7 @@ class Config: binary_sha256: Series[str] file_path: Series[str] file_sha256: Series[str] + backend: Series[str] = pa.Field(isin=list(BACKEND_SPECS)) treatment: Series[str] = pa.Field(isin=["off", "term", "proofs"]) timeout_sec: Series[int] = pa.Field(gt=0) wall_sec: Series[float] = pa.Field(nullable=True, ge=0) @@ -161,6 +201,42 @@ class BenchmarkSpec: treatments: tuple[Treatment, ...] rounds: int timeout_sec: int + backends: tuple[Backend, ...] = DEFAULT_BACKENDS + + +@dataclass(frozen=True) +class BenchmarkCell: + backend: Backend + treatment: Treatment + + +@dataclass(frozen=True) +class ProfileMode: + iterations: int | None + profile_seconds: int | None + + @property + def cache_label(self) -> str: + if self.iterations is not None: + return f"i{self.iterations}" + assert self.profile_seconds is not None + return f"auto{self.profile_seconds}s" + + +@dataclass(frozen=True) +class ProfileRequest: + file: FileSpec + target_request: TargetRequest + backend: Backend + treatment: Treatment + timeout_sec: int + profiles_dir: Path + mode: ProfileMode + open_after: bool + force_run: bool + top: int = DEFAULT_PROFILE_TOP + show_summary: bool = True + output_format: OutputFormat = "rich" @dataclass(frozen=True) @@ -215,12 +291,14 @@ class EstimateKey: file_sha256: str treatment: Treatment timeout_sec: int + backend: Backend = "main" @dataclass(frozen=True) class CellPlan: target: ResolvedTarget file: FileSpec + backend: Backend treatment: Treatment required_rows: int selected_cached_rows: DataFrame[ReportFrame] @@ -270,7 +348,7 @@ def ok(self) -> bool: return self.issue is None and self.mean is not None -CellMap = dict[tuple[str, Treatment], CellSummary] +CellMap = dict[tuple[str, Backend, Treatment], CellSummary] TargetCellMaps = dict[ResolvedTarget, CellMap] @@ -292,7 +370,33 @@ def ok(self) -> bool: return self.issue is None and self.point is not None +@dataclass(frozen=True) +class ReportTableData: + title: str + headers: tuple[str, ...] + rows: tuple[tuple[str, ...], ...] + caption: str | None = None + alignments: tuple[TableAlignment, ...] | None = None + + +@dataclass(frozen=True) +class MetricSpec: + title: str + caption: str + file_count_heading: str + format_total: Callable[[float | None], str] + format_change: Callable[[RatioSummary], str] + format_result: Callable[[RatioSummary], str] + omit_without_samples: bool = False + + def parse_args(argv: Sequence[str]) -> argparse.Namespace: + if argv and argv[0] == "profile": + return parse_profile_args(argv[1:]) + return parse_benchmark_args(argv) + + +def parse_benchmark_args(argv: Sequence[str]) -> argparse.Namespace: parser = argparse.ArgumentParser( description="Collect or reuse egglog-experimental benchmark reports.", ) @@ -308,6 +412,12 @@ def parse_args(argv: Sequence[str]) -> argparse.Namespace: default=DEFAULT_REPORT, help=f"JSONL report/cache path, or - for stdout (default: {DEFAULT_REPORT})", ) + parser.add_argument( + "--format", + choices=("rich", "markdown"), + default="rich", + help="final report format: rich to stderr, or markdown to stdout (default: rich)", + ) parser.add_argument( "--rounds", type=positive_int, @@ -320,6 +430,11 @@ def parse_args(argv: Sequence[str]) -> argparse.Namespace: default=DEFAULT_TIMEOUT_SEC, help=f"per-process timeout in seconds (default: {DEFAULT_TIMEOUT_SEC})", ) + parser.add_argument( + "--backend", + default=",".join(DEFAULT_BACKENDS), + help=f"comma-separated backends: {', '.join(BACKEND_SPECS)} (default: {','.join(DEFAULT_BACKENDS)})", + ) parser.add_argument( "--treatments", default=",".join(DEFAULT_TREATMENTS), @@ -330,7 +445,89 @@ def parse_args(argv: Sequence[str]) -> argparse.Namespace: action="store_true", help="append fresh rows even when enough cached rows exist", ) - return parser.parse_args(argv) + args = parser.parse_args(argv) + if args.format == "markdown" and args.report == "-": + parser.error("--format markdown cannot be combined with --report -") + args.command = "benchmark" + return args + + +def parse_profile_args(argv: Sequence[str]) -> argparse.Namespace: + parser = argparse.ArgumentParser( + prog=f"{Path(sys.argv[0]).name} profile", + description="Record or reuse a cached Samply CPU profile for one egglog workload.", + ) + parser.add_argument("file", help="egglog file to profile") + parser.add_argument( + "--target", + action="append", + default=None, + help="target source: ., /path, @git-ref, #pr, or label=source; cache-only label= is not supported", + ) + parser.add_argument( + "--backend", + default="main", + help="single backend to profile (default: main)", + ) + parser.add_argument( + "--treatment", + default="proofs", + help="single treatment to profile (default: proofs)", + ) + iteration_group = parser.add_mutually_exclusive_group() + iteration_group.add_argument( + "--iterations", + type=positive_int, + default=None, + help="explicit Samply iteration count", + ) + iteration_group.add_argument( + "--profile-seconds", + type=positive_int, + default=None, + help=f"target duration for automatic iteration selection (default: {DEFAULT_PROFILE_SECONDS})", + ) + parser.add_argument( + "--profiles-dir", + default=DEFAULT_PROFILES_DIR, + help=f"profile cache directory (default: {DEFAULT_PROFILES_DIR})", + ) + parser.add_argument( + "--top", + type=positive_int, + default=DEFAULT_PROFILE_TOP, + help=f"application functions to show in the macOS CPU summary (default: {DEFAULT_PROFILE_TOP})", + ) + parser.add_argument( + "--no-summary", + action="store_true", + help="print only the profile artifact path", + ) + parser.add_argument( + "--format", + choices=("rich", "markdown"), + default="rich", + help="summary format: rich to stderr, or markdown to stdout (default: rich)", + ) + parser.add_argument( + "--open", + action="store_true", + help="open the profile with samply load after recording or cache hit", + ) + parser.add_argument( + "--force-run", + action="store_true", + help="record again and atomically replace the cached profile", + ) + parser.add_argument( + "--timeout-sec", + type=positive_int, + default=DEFAULT_TIMEOUT_SEC, + help=f"per-workload timeout in seconds for calibration and profiling watchdog (default: {DEFAULT_TIMEOUT_SEC})", + ) + args = parser.parse_args(argv) + args.command = "profile" + return args def positive_int(value: str) -> int: @@ -358,6 +555,86 @@ def parse_treatments(value: str) -> tuple[Treatment, ...]: return tuple(treatments) +def parse_backends(value: str) -> tuple[Backend, ...]: + backends: list[Backend] = [] + seen: set[str] = set() + for raw in value.split(","): + item = raw.strip() + if not item: + continue + if item not in BACKEND_SPECS: + raise ValueError(f"unknown backend: {item}") + if item in seen: + raise ValueError(f"duplicate backend: {item}") + seen.add(item) + backends.append(item) + if not backends: + raise ValueError("at least one backend is required") + return tuple(backends) + + +def backend_spec(backend: Backend) -> BackendSpec: + try: + return BACKEND_SPECS[backend] + except KeyError as error: + raise ValueError(f"unknown backend: {backend}") from error + + +def backend_supports_treatment(backend: Backend, treatment: Treatment) -> bool: + return treatment in backend_spec(backend).treatments + + +def supported_treatments(backend: Backend) -> tuple[Treatment, ...]: + return backend_spec(backend).treatments + + +def backend_cargo_features(backends: Sequence[Backend]) -> tuple[str, ...]: + return tuple(dict.fromkeys(feature for backend in backends for feature in backend_spec(backend).cargo_features)) + + +def backend_treatment_cells( + backends: Sequence[Backend], + treatments: Sequence[Treatment], +) -> tuple[BenchmarkCell, ...]: + cells: list[BenchmarkCell] = [] + requested = ",".join(treatments) + for backend in backends: + backend_cells = tuple( + BenchmarkCell(backend, treatment) + for treatment in treatments + if backend_supports_treatment(backend, treatment) + ) + if not backend_cells: + supported = ",".join(supported_treatments(backend)) + raise ValueError( + f"backend {backend} has no supported treatments in requested set {requested}; " + f"supported treatments: {supported}" + ) + cells.extend(backend_cells) + return tuple(cells) + + +def benchmark_cells(spec: BenchmarkSpec) -> tuple[BenchmarkCell, ...]: + return backend_treatment_cells(spec.backends, spec.treatments) + + +def backend_has_treatment(spec: BenchmarkSpec, backend: Backend, treatment: Treatment) -> bool: + return any(cell.backend == backend and cell.treatment == treatment for cell in benchmark_cells(spec)) + + +def shared_backend_treatments( + spec: BenchmarkSpec, + baseline_backend: Backend, + candidate_backend: Backend, +) -> tuple[Treatment, ...]: + return tuple( + treatment + for treatment in spec.treatments + if backend_supports_treatment(baseline_backend, treatment) + and backend_supports_treatment(candidate_backend, treatment) + ) + + def parse_target(raw: str) -> TargetRequest: if "=" in raw: label, source = raw.split("=", 1) @@ -453,6 +730,7 @@ def file_contains_executable_prove_command(path: Path) -> bool: def validate_spec(spec: BenchmarkSpec) -> None: + benchmark_cells(spec) for file_spec in spec.files: if file_contains_executable_prove_command(file_spec.absolute_path): raise ValueError( @@ -461,6 +739,61 @@ def validate_spec(spec: BenchmarkSpec) -> None: ) +def parse_single_backend(value: str) -> Backend: + backends = parse_backends(value) + if len(backends) != 1: + raise ValueError("profile mode requires exactly one backend") + return backends[0] + + +def parse_single_treatment(value: str) -> Treatment: + treatments = parse_treatments(value) + if len(treatments) != 1: + raise ValueError("profile mode requires exactly one treatment") + return treatments[0] + + +def resolve_profile_request(args: argparse.Namespace, invocation_cwd: Path) -> ProfileRequest: + files = resolve_files([str(args.file)], invocation_cwd) + backend = parse_single_backend(str(args.backend)) + treatment = parse_single_treatment(str(args.treatment)) + target_specs = args.target if args.target is not None else ["."] + if len(target_specs) != 1: + raise ValueError("profile mode requires exactly one target") + if args.iterations is not None: + mode = ProfileMode(iterations=args.iterations, profile_seconds=None) + else: + profile_seconds = args.profile_seconds if args.profile_seconds is not None else DEFAULT_PROFILE_SECONDS + mode = ProfileMode(iterations=None, profile_seconds=profile_seconds) + profiles_dir = Path(str(args.profiles_dir)).expanduser() + if not profiles_dir.is_absolute(): + profiles_dir = invocation_cwd / profiles_dir + request = ProfileRequest( + file=files[0], + target_request=parse_target(str(target_specs[0])), + backend=backend, + treatment=treatment, + timeout_sec=int(args.timeout_sec), + profiles_dir=profiles_dir, + mode=mode, + open_after=bool(args.open), + force_run=bool(args.force_run), + top=int(args.top), + show_summary=not bool(args.no_summary), + output_format=cast(OutputFormat, str(args.format)), + ) + validate_spec( + BenchmarkSpec( + files=(request.file,), + treatments=(request.treatment,), + rounds=1, + timeout_sec=request.timeout_sec, + backends=(request.backend,), + ) + ) + return request + + def resolve_report_destination(raw_report: str, invocation_cwd: Path) -> ReportDestination: if raw_report == "-": return ReportDestination(path=None, stream=sys.stdout) @@ -502,10 +835,13 @@ def normalize_report_frame(frame: pd.DataFrame) -> DataFrame[ReportFrame]: extra = sorted(set(normalized.columns) - set(columns)) if extra: raise ValueError(f"unexpected report column(s): {', '.join(extra)}") + if "backend" not in normalized.columns: + normalized["backend"] = "main" for column in columns: if column not in normalized.columns: normalized[column] = pd.NA normalized = normalized.loc[:, columns] + normalized["backend"] = normalized["backend"].fillna("main") normalized["started_at"] = pd.to_datetime(normalized["started_at"], utc=True, errors="raise") for column in ["wall_sec", "user_sec", "system_sec", "cpu_wall_ratio"]: normalized[column] = pd.to_numeric(normalized[column], errors="raise") @@ -573,6 +909,7 @@ def from_rows(cls, rows: DataFrame[ReportFrame]) -> EstimateModel: file_sha256=str(record["file_sha256"]), treatment=cast(Treatment, str(record["treatment"])), timeout_sec=int(record["timeout_sec"]), + backend=cast(Backend, str(record["backend"])), ) samples.setdefault(key, []).append(float(record["wall_sec"])) return cls(samples) @@ -626,6 +963,7 @@ def selected_rows( matches = rows.loc[ rows["binary_sha256"].eq(key.binary_sha256) & rows["file_sha256"].eq(key.file_sha256) + & rows["backend"].eq(key.backend) & rows["treatment"].eq(key.treatment) & rows["timeout_sec"].eq(key.timeout_sec) ] @@ -641,6 +979,7 @@ def status_counts_for_rows(rows: DataFrame[ReportFrame]) -> dict[str, int]: def estimate_key_for( target: ResolvedTarget, file_spec: FileSpec, + backend: Backend, treatment: Treatment, timeout_sec: int, ) -> EstimateKey: @@ -649,6 +988,7 @@ def estimate_key_for( file_sha256=file_spec.sha256, treatment=treatment, timeout_sec=timeout_sec, + backend=backend, ) @@ -660,8 +1000,8 @@ def build_collection_plan( ) -> CollectionPlan: cells: list[CellPlan] = [] for file_spec in spec.files: - for treatment in spec.treatments: - estimate_key = estimate_key_for(target, file_spec, treatment, spec.timeout_sec) + for cell in benchmark_cells(spec): + estimate_key = estimate_key_for(target, file_spec, cell.backend, cell.treatment, spec.timeout_sec) cached = selected_rows( rows, estimate_key, @@ -672,7 +1012,8 @@ def build_collection_plan( CellPlan( target=target, file=file_spec, - treatment=treatment, + backend=cell.backend, + treatment=cell.treatment, required_rows=spec.rounds, selected_cached_rows=cached, missing_observations=missing, @@ -786,24 +1127,62 @@ def target_row_for_request( ) -def build_target(row: TargetRow, output: RunnerOutput) -> tuple[Path, str]: +def build_target( + row: TargetRow, + output: RunnerOutput, + build_profile: BuildProfile = "release", + cargo_features: Sequence[str] = (), +) -> tuple[Path, str]: checkout_path = Path(row.path) output.build_start(row) + if build_profile == "release": + build_args = ["cargo", "build", "--release", "-p", "egglog-experimental"] + else: + build_args = ["cargo", "build", "--profile", "profiling", "-p", "egglog-experimental"] + if cargo_features: + build_args.extend(("--features", ",".join(cargo_features))) subprocess.run( - ["cargo", "build", "--release", "-p", "egglog-experimental"], + build_args, cwd=checkout_path, check=True, stdout=sys.stderr, stderr=sys.stderr, ) binary_name = "egglog-experimental.exe" if os.name == "nt" else "egglog-experimental" - binary_path = checkout_path / "target" / "release" / binary_name + binary_path = checkout_path / "target" / build_profile / binary_name if not binary_path.is_file(): - raise FileNotFoundError(f"release binary was not produced: {binary_path}") + raise FileNotFoundError(f"{build_profile} binary was not produced: {binary_path}") binary_sha256 = sha256_file(binary_path) return (binary_path, binary_sha256) +def materialize_target_request(request: TargetRequest, invocation_cwd: Path, repo_root: Path) -> TargetRow: + if request.is_label_lookup: + raise ValueError("cache-only target labels cannot be materialized without a cached report row") + if request.source.startswith("@"): + ref = request.source[1:] + if not ref: + raise ValueError(f"git target is missing a ref: {request.raw}") + checkout_path, git_ref_value = materialize_git_ref(repo_root, ref, request.label or ref) + elif parse_pr_number(request.source) is not None: + checkout_path, git_ref_value = materialize_pr_target(repo_root, request.source, request.label) + else: + checkout_path, git_ref_value = resolve_path_target(request.source, invocation_cwd) + return target_row_for_request(request, checkout_path, git_ref_value) + + +def build_resolved_target( + request: TargetRequest, + row: TargetRow, + output: RunnerOutput, + build_profile: BuildProfile, + cargo_features: Sequence[str], +) -> ResolvedTarget: + binary_path, binary_sha256 = build_target(row, output, build_profile, cargo_features) + row = replace(row, is_dirty=git_dirty(Path(row.path))) + return ResolvedTarget(request=request, row=row, binary_sha256=binary_sha256, binary_path=binary_path) + + def resolve_target( request: TargetRequest, rows: DataFrame[ReportFrame], @@ -835,26 +1214,23 @@ def resolve_target( ) checkout_path, resolved_sha = materialize_git_ref(repo_root, str(pointer["target_git_sha"]), request.label) row = target_row_for_request(request, checkout_path, resolved_sha) - binary_path, binary_sha256 = build_target(row, output) - row = replace(row, is_dirty=git_dirty(checkout_path)) - return ResolvedTarget(request=request, row=row, binary_sha256=binary_sha256, binary_path=binary_path) + return build_resolved_target(request, row, output, "release", backend_cargo_features(spec.backends)) - if request.source.startswith("@"): - ref = request.source[1:] - if not ref: - raise ValueError(f"git target is missing a ref: {request.raw}") - checkout_path, resolved_sha = materialize_git_ref(repo_root, ref, request.label or ref) - row = target_row_for_request(request, checkout_path, resolved_sha) - elif parse_pr_number(request.source) is not None: - checkout_path, resolved_sha = materialize_pr_target(repo_root, request.source, request.label) - row = target_row_for_request(request, checkout_path, resolved_sha) - else: - checkout_path, git_ref_value = resolve_path_target(request.source, invocation_cwd) - row = target_row_for_request(request, checkout_path, git_ref_value) + row = materialize_target_request(request, invocation_cwd, repo_root) + return build_resolved_target(request, row, output, "release", backend_cargo_features(spec.backends)) - binary_path, binary_sha256 = build_target(row, output) - row = replace(row, is_dirty=git_dirty(Path(row.path))) - return ResolvedTarget(request=request, row=row, binary_sha256=binary_sha256, binary_path=binary_path) + +def resolve_profile_target( + request: TargetRequest, + backend: Backend, + invocation_cwd: Path, + repo_root: Path, + output: RunnerOutput, +) -> ResolvedTarget: + if request.is_label_lookup: + raise ValueError("profile mode does not support cache-only label= targets; use label=SOURCE") + row = materialize_target_request(request, invocation_cwd, repo_root) + return build_resolved_target(request, row, output, "profiling", backend_cargo_features((backend,))) def label_has_enough_rows( @@ -863,10 +1239,10 @@ def label_has_enough_rows( spec: BenchmarkSpec, ) -> bool: for file_spec in spec.files: - for treatment in spec.treatments: + for cell in benchmark_cells(spec): matches = selected_rows( rows, - EstimateKey(binary_sha256, file_spec.sha256, treatment, spec.timeout_sec), + EstimateKey(binary_sha256, file_spec.sha256, cell.treatment, spec.timeout_sec, cell.backend), spec.rounds, ) if len(matches) < spec.rounds: @@ -919,29 +1295,196 @@ def treatment_flags(treatment: Treatment) -> list[str]: return ["--proofs"] -def run_process( +def backend_flags(backend: Backend) -> list[str]: + return list(backend_spec(backend).flags) + + +def workload_command( binary_path: Path, - checkout_path: Path, file_spec: FileSpec, + backend: Backend, treatment: Treatment, - timeout_sec: int, -) -> TimingResult: - command = [ +) -> list[str]: + return [ str(binary_path), "--mode", "no-messages", "-j", "1", + *backend_flags(backend), *treatment_flags(treatment), str(file_spec.absolute_path), ] - return run_command(command, checkout_path, timeout_sec) + + +def run_process( + binary_path: Path, + checkout_path: Path, + file_spec: FileSpec, + backend: Backend, + treatment: Treatment, + timeout_sec: int, +) -> TimingResult: + return run_command(workload_command(binary_path, file_spec, backend, treatment), checkout_path, timeout_sec) def run_startup_warmup(binary_path: Path, checkout_path: Path, timeout_sec: int) -> TimingResult: return run_command([str(binary_path), "--help"], checkout_path, timeout_sec) +def profile_hash_component(value: str) -> str: + return value.removeprefix("sha256:") + + +def profile_cache_path( + profiles_dir: Path, + binary_sha256: str, + file_sha256: str, + backend: Backend, + treatment: Treatment, + mode: ProfileMode, +) -> Path: + return ( + profiles_dir + / PROFILE_CACHE_VERSION + / profile_hash_component(binary_sha256) + / profile_hash_component(file_sha256) + / f"{backend}-{treatment}-{mode.cache_label}.json.gz" + ) + + +def profile_cache_hit(path: Path) -> bool: + return path.is_file() and path.stat().st_size > 0 + + +def profile_display_path(path: Path, invocation_cwd: Path) -> Path: + resolved = path.resolve() + try: + return resolved.relative_to(invocation_cwd.resolve()) + except ValueError: + return resolved + + +def calculate_profile_iterations( + elapsed_seconds: float, + profile_seconds: int, + max_iterations: int = MAX_PROFILE_ITERATIONS, +) -> tuple[int, bool]: + if elapsed_seconds <= 0: + return (max_iterations, True) + uncapped = max(1, math.ceil(profile_seconds * 1.10 / elapsed_seconds)) + if uncapped > max_iterations: + return (max_iterations, True) + return (uncapped, False) + + +def profile_name( + file_spec: FileSpec, + backend: Backend, + treatment: Treatment, + mode: ProfileMode, + iterations: int, + binary_sha256: str, +) -> str: + stem = Path(file_spec.display_path).stem + file_hash = profile_hash_component(file_spec.sha256)[:8] + binary_hash = profile_hash_component(binary_sha256)[:8] + mode_label = mode.cache_label if mode.iterations is not None else f"auto>={mode.profile_seconds}s" + return f"{stem} {backend}/{treatment} {mode_label} x{iterations} [bin:{binary_hash} file:{file_hash}]" + + +def profile_temp_path(artifact: Path) -> Path: + base_name = artifact.name.removesuffix(".json.gz") + return artifact.with_name(f".{base_name}.tmp-{uuid.uuid4().hex}.json.gz") + + +def samply_executable() -> str: + executable = shutil.which("samply") + if executable is None: + raise FileNotFoundError("Install Samply with: cargo install --locked samply") + return executable + + +def samply_record_command( + samply: str, + artifact: Path, + name: str, + iterations: int, + workload: Sequence[str], +) -> list[str]: + return [ + samply, + "record", + "--save-only", + "--rate", + str(PROFILE_SAMPLY_RATE_HZ), + "--reuse-threads", + "--iteration-count", + str(iterations), + "--profile-name", + name, + "--output", + str(artifact), + "--", + *workload, + ] + + +def profile_record_timeout(timeout_sec: int, iterations: int) -> int: + return max(timeout_sec + 60, timeout_sec * iterations + 60) + + +def run_samply_record( + *, + artifact: Path, + name: str, + iterations: int, + workload: Sequence[str], + checkout_path: Path, + timeout_sec: int, +) -> dict[str, Any]: + temp_artifact = profile_temp_path(artifact) + temp_artifact.parent.mkdir(parents=True, exist_ok=True) + with suppress(FileNotFoundError): + temp_artifact.unlink() + command = samply_record_command(samply_executable(), temp_artifact, name, iterations, workload) + env = os.environ.copy() + env["RUST_LOG"] = "error" + try: + subprocess.run( + command, + cwd=checkout_path, + env=env, + check=True, + timeout=profile_record_timeout(timeout_sec, iterations), + stdout=sys.stderr, + stderr=sys.stderr, + ) + if not profile_cache_hit(temp_artifact): + raise ValueError(f"Samply did not produce a nonempty profile artifact: {temp_artifact}") + profile = samply_analysis.read_artifact(temp_artifact) + artifact.parent.mkdir(parents=True, exist_ok=True) + os.replace(temp_artifact, artifact) + return profile + except BaseException: + with suppress(FileNotFoundError): + temp_artifact.unlink() + raise + + +def open_samply_profile(artifact: Path, checkout_path: Path) -> None: + try: + subprocess.run( + [samply_executable(), "load", str(artifact)], + cwd=checkout_path, + check=True, + stdout=sys.stderr, + stderr=sys.stderr, + ) + except KeyboardInterrupt: + return + + def run_command(command: Sequence[str], checkout_path: Path, timeout_sec: int) -> TimingResult: env = os.environ.copy() env["RUST_LOG"] = "error" @@ -1028,12 +1571,13 @@ def ru_maxrss_to_bytes(ru_maxrss: int, platform: str = sys.platform) -> int | No def collection_label( file_spec: FileSpec, + backend: Backend, treatment: Treatment, round_index: int, rounds: int, ) -> str: filename = Path(file_spec.display_path).name - return f"{filename} {treatment} {round_index + 1}/{rounds}" + return f"{filename} {backend}/{treatment} {round_index + 1}/{rounds}" def format_timing_result(result: TimingResult) -> str: @@ -1086,8 +1630,8 @@ def emit_collection_plan( ) -> None: total_estimate = collection_plan_estimate(plan, estimate_model) table = Table(title=f"{plan.target.display_label}: cache and estimate plan") - table.add_column("File") - table.add_column("Treatment") + table.add_column("File", overflow="fold") + table.add_column("Cell", no_wrap=True) table.add_column("Cached") table.add_column("Missing") table.add_column("Statuses") @@ -1100,7 +1644,7 @@ def emit_collection_plan( statuses = ", ".join(f"{status}:{count}" for status, count in sorted(status_counts.items())) or "-" table.add_row( cell.file.display_path, - cell.treatment, + f"{cell.backend}/{cell.treatment}", f"{len(cell.selected_cached_rows)}/{cell.required_rows}", str(cell.missing_observations), statuses, @@ -1133,6 +1677,7 @@ def flat_report_record( "binary_sha256": target.binary_sha256, "file_path": cell.file.display_path, "file_sha256": cell.file.sha256, + "backend": cell.backend, "treatment": cell.treatment, "timeout_sec": spec.timeout_sec, "wall_sec": result.timing.wall_sec, @@ -1185,22 +1730,21 @@ def decrement_remaining(key: EstimateKey, count: int = 1) -> None: else: remaining_processes.pop(key, None) - def run_loop(progress: Progress | None = None, process_task: Any | None = None) -> None: + def run_loop(progress: Progress, process_task: TaskID) -> None: nonlocal next_index, completed_observations for round_index in range(max_deficit): for cell in plan.cells: if round_index >= cell.missing_observations: continue cell_file = cell.file + cell_backend = cell.backend cell_treatment = cell.treatment cell_key = cell.estimate_key required_rounds = cell.missing_observations observation_number = completed_observations + 1 - label = collection_label(cell_file, cell_treatment, round_index, required_rounds) + label = collection_label(cell_file, cell_backend, cell_treatment, round_index, required_rounds) def update_progress(current: str, advance: int = 0) -> None: - if progress is None or process_task is None: - return progress.update( process_task, advance=advance, @@ -1210,7 +1754,14 @@ def update_progress(current: str, advance: int = 0) -> None: update_progress(f"row {observation_number}/{total_observations} timed") started_at = now_iso() - result = run_process(binary_path, Path(target.row.path), cell_file, cell_treatment, spec.timeout_sec) + result = run_process( + binary_path, + Path(target.row.path), + cell_file, + cell_backend, + cell_treatment, + spec.timeout_sec, + ) estimate_model.record_process(cell_key, result) decrement_remaining(cell_key) update_progress( @@ -1230,7 +1781,6 @@ def update_progress(current: str, advance: int = 0) -> None: fresh_frames.append(fresh_frame) next_index += 1 completed_observations += 1 - assert progress is not None progress.console.print(f" {label}: fresh {format_timing_result(result)}") update_progress( f"rows {completed_observations}/{total_observations}; last {format_timing_result(result)}" @@ -1443,20 +1993,22 @@ def target_cell_summaries( rows: DataFrame[ReportFrame], target: ResolvedTarget, spec: BenchmarkSpec, + backends: Sequence[Backend] | None = None, treatments: Sequence[Treatment] | None = None, ) -> CellMap: + chosen_backends = spec.backends if backends is None else backends chosen_treatments = spec.treatments if treatments is None else treatments return { - (file_spec.sha256, treatment): summarize_cell( + (file_spec.sha256, cell.backend, cell.treatment): summarize_cell( selected_rows( rows, - estimate_key_for(target, file_spec, treatment, spec.timeout_sec), + estimate_key_for(target, file_spec, cell.backend, cell.treatment, spec.timeout_sec), spec.rounds, ), spec.rounds, ) for file_spec in spec.files - for treatment in chosen_treatments + for cell in backend_treatment_cells(chosen_backends, chosen_treatments) } @@ -1464,20 +2016,22 @@ def target_rss_cell_summaries( rows: DataFrame[ReportFrame], target: ResolvedTarget, spec: BenchmarkSpec, + backends: Sequence[Backend] | None = None, treatments: Sequence[Treatment] | None = None, ) -> CellMap: + chosen_backends = spec.backends if backends is None else backends chosen_treatments = spec.treatments if treatments is None else treatments return { - (file_spec.sha256, treatment): summarize_rss_cell( + (file_spec.sha256, cell.backend, cell.treatment): summarize_rss_cell( selected_rows( rows, - estimate_key_for(target, file_spec, treatment, spec.timeout_sec), + estimate_key_for(target, file_spec, cell.backend, cell.treatment, spec.timeout_sec), spec.rounds, ), spec.rounds, ) for file_spec in spec.files - for treatment in chosen_treatments + for cell in backend_treatment_cells(chosen_backends, chosen_treatments) } @@ -1486,12 +2040,13 @@ def target_suite_treatment_ratio( candidate_cells: CellMap, spec: BenchmarkSpec, treatment: Treatment, + backend: Backend = "main", ) -> RatioSummary: return suite_ratio( [ ( - baseline_cells[(file_spec.sha256, treatment)], - candidate_cells[(file_spec.sha256, treatment)], + baseline_cells[(file_spec.sha256, backend, treatment)], + candidate_cells[(file_spec.sha256, backend, treatment)], ) for file_spec in spec.files ] @@ -1501,14 +2056,15 @@ def target_suite_treatment_ratio( def treatment_file_cells( cell_map: CellMap, spec: BenchmarkSpec, + backend: Backend, baseline_treatment: Treatment, candidate_treatment: Treatment, ) -> list[tuple[FileSpec, CellSummary, CellSummary]]: return [ ( file_spec, - cell_map[(file_spec.sha256, baseline_treatment)], - cell_map[(file_spec.sha256, candidate_treatment)], + cell_map[(file_spec.sha256, backend, baseline_treatment)], + cell_map[(file_spec.sha256, backend, candidate_treatment)], ) for file_spec in spec.files ] @@ -1518,13 +2074,31 @@ def target_treatment_file_cells( baseline_cells: CellMap, candidate_cells: CellMap, spec: BenchmarkSpec, + backend: Backend, + treatment: Treatment, +) -> list[tuple[FileSpec, CellSummary, CellSummary]]: + return [ + ( + file_spec, + baseline_cells[(file_spec.sha256, backend, treatment)], + candidate_cells[(file_spec.sha256, backend, treatment)], + ) + for file_spec in spec.files + ] + + +def backend_treatment_file_cells( + cell_map: CellMap, + spec: BenchmarkSpec, + baseline_backend: Backend, + candidate_backend: Backend, treatment: Treatment, ) -> list[tuple[FileSpec, CellSummary, CellSummary]]: return [ ( file_spec, - baseline_cells[(file_spec.sha256, treatment)], - candidate_cells[(file_spec.sha256, treatment)], + cell_map[(file_spec.sha256, baseline_backend, treatment)], + cell_map[(file_spec.sha256, candidate_backend, treatment)], ) for file_spec in spec.files ] @@ -1550,6 +2124,37 @@ def worst_file_ratio(ratios: Sequence[tuple[FileSpec, RatioSummary]]) -> tuple[F return max(valid, key=lambda item: item[1].point or 0.0) +def best_file_ratio(ratios: Sequence[tuple[FileSpec, RatioSummary]]) -> tuple[FileSpec | None, RatioSummary]: + if not ratios: + return (None, RatioSummary(None, None, None, "no files")) + valid = [(file_spec, ratio) for file_spec, ratio in ratios if ratio.point is not None] + if valid: + return min(valid, key=lambda item: item[1].point or math.inf) + return ratios[0] + + +def count_better_files(ratios: Sequence[tuple[FileSpec, RatioSummary]]) -> tuple[int, int]: + valid = [ratio for _, ratio in ratios if ratio.point is not None] + better = sum(1 for ratio in valid if ratio.ci_high is not None and ratio.ci_high < 1.0) + return (better, len(valid)) + + +def format_better_file_count(count: tuple[int, int]) -> str: + better, total = count + if total == 0: + return "-" + return f"{better}/{total}" + + +def suite_total_mean(cells: Sequence[CellSummary]) -> float | None: + total = 0.0 + for cell in cells: + if not cell.ok or cell.mean is None: + return None + total += cell.mean + return total + + def format_estimate_or_interval( point: float | None, low: float | None, @@ -1569,6 +2174,10 @@ def format_seconds_summary(summary: CellSummary) -> str: return format_estimate_or_interval(summary.mean, summary.ci_low, summary.ci_high, "s", 4) +def format_seconds_value(value: float | None) -> str: + return "-" if value is None else f"{value:.4f}s" + + def format_bytes(value: float | None) -> str: if value is None: return "-" @@ -1635,6 +2244,10 @@ def format_comparison_result(summary: RatioSummary) -> Text: return styled_status(result) +def comparison_result_text(summary: RatioSummary) -> str: + return format_comparison_result(summary).plain + + def lower_is_better_result(summary: RatioSummary) -> str: if summary.point is None: return "invalid" @@ -1654,6 +2267,10 @@ def format_lower_is_better_result(summary: RatioSummary) -> Text: return styled_status(result) +def lower_is_better_result_text(summary: RatioSummary) -> str: + return format_lower_is_better_result(summary).plain + + def proof_gate_result(summary: RatioSummary) -> tuple[str, str]: if summary.point is None: return ("invalid", f"invalid: {summary.issue or 'unavailable'}") @@ -1669,6 +2286,14 @@ def format_proof_gate_result(summary: RatioSummary) -> Text: return styled_status(status, text) +def proof_gate_result_text(summary: RatioSummary) -> str: + return proof_gate_result(summary)[1] + + +def report_row(*values: object) -> tuple[str, ...]: + return tuple(str(value) for value in values) + + def report_table(title: str, *, caption: str | None = None, show_lines: bool = False) -> Table: return Table( title=title, @@ -1682,6 +2307,69 @@ def report_table(title: str, *, caption: str | None = None, show_lines: bool = F ) +def rich_result_cell(value: str) -> str | Text: + if value.startswith("invalid:"): + return styled_status("invalid", value) + if value == "<2x established": + return styled_status("established", value) + if value == "<2x not established": + return styled_status("not established", value) + if value in RESULT_STYLES: + return styled_status(value) + return value + + +def render_rich_table(table_data: ReportTableData) -> Table: + table = report_table(table_data.title, caption=table_data.caption) + alignments = table_data.alignments or tuple("left" for _ in table_data.headers) + for header, alignment in zip(table_data.headers, alignments, strict=True): + table.add_column(header, no_wrap=alignment == "right", justify="right" if alignment == "right" else "left") + result_columns = {"Result", "Best result"} + for row in table_data.rows: + values: list[str | Text] = [] + for header, value in zip(table_data.headers, row, strict=True): + values.append(rich_result_cell(value) if header in result_columns else value) + table.add_row(*values) + return table + + +def markdown_escape_cell(value: object) -> str: + text = str(value) + return text.replace("\\", "\\\\").replace("|", "\\|").replace("\r\n", "\n").replace("\n", "
") + + +def render_markdown_table(table_data: ReportTableData, *, heading_level: int = 3) -> str: + alignments = table_data.alignments or tuple("left" for _ in table_data.headers) + separator_cells = tuple("---:" if alignment == "right" else "---" for alignment in alignments) + lines = [ + f"{'#' * heading_level} {table_data.title}", + "", + "| " + " | ".join(markdown_escape_cell(header) for header in table_data.headers) + " |", + "| " + " | ".join(separator_cells) + " |", + ] + for row in table_data.rows: + lines.append("| " + " | ".join(markdown_escape_cell(value) for value in row) + " |") + if table_data.caption: + lines.extend(["", f"*{markdown_escape_cell(table_data.caption)}*"]) + return "\n".join(lines) + + +def benchmark_command_block(argv: Sequence[str]) -> str: + return "```shell\n$ " + shlex.join(["./bench.py", *argv]) + "\n```" + + +def cell_label(spec: BenchmarkSpec, backend: Backend, treatment: Treatment) -> str: + if len(spec.backends) == 1: + return treatment + return f"{backend}/{treatment}" + + +def backend_metric_label(spec: BenchmarkSpec, backend: Backend, label: str) -> str: + if len(spec.backends) == 1: + return label + return f"{backend} {label}" + + def render_report( console: Console, report_destination: ReportDestination, @@ -1700,6 +2388,9 @@ def render_report( if len(targets) > 1: render_per_file_wall_time_change(console, cell_maps, targets, spec) render_per_file_peak_rss_change(console, rss_cell_maps, targets, spec) + if len(spec.backends) > 1: + render_per_file_backend_wall_time_change(console, cell_maps, targets, spec) + render_per_file_backend_peak_rss_change(console, rss_cell_maps, targets, spec) for target in targets: render_target_diagnostics(console, cell_maps[target], rss_cell_maps[target], target, spec) @@ -1707,6 +2398,62 @@ def render_report( render_benchmark_summary(console, cell_maps, rss_cell_maps, targets, spec) +def render_markdown_report( + report_destination: ReportDestination, + rows: DataFrame[ReportFrame], + targets: Sequence[ResolvedTarget], + spec: BenchmarkSpec, + command_argv: Sequence[str] | None = None, +) -> str: + cell_maps = {target: target_cell_summaries(rows, target, spec) for target in targets} + rss_cell_maps = {target: target_rss_cell_summaries(rows, target, spec) for target in targets} + parts = [] + if command_argv is not None: + parts.append(benchmark_command_block(command_argv)) + parts.extend( + [ + "# Benchmark Report", + f"- Report: `{report_destination.display_path}`\n- Selected rows per cell: `{spec.rounds}`", + render_markdown_table(targets_table_data(targets), heading_level=2), + ] + ) + + comparison_tables: list[ReportTableData] = [] + if len(targets) > 1: + comparison_tables.extend(per_file_wall_time_change_tables(cell_maps, targets, spec)) + comparison_tables.extend(per_file_peak_rss_change_tables(rss_cell_maps, targets, spec)) + if len(spec.backends) > 1: + comparison_tables.extend(per_file_backend_wall_time_change_tables(cell_maps, targets, spec)) + comparison_tables.extend(per_file_backend_peak_rss_change_tables(rss_cell_maps, targets, spec)) + if comparison_tables: + parts.append("## Comparisons") + parts.extend(render_markdown_table(table_data) for table_data in comparison_tables) + + diagnostic_tables: list[ReportTableData] = [] + for target in targets: + overhead = overhead_ratios_table_data(cell_maps[target], target, spec) + if overhead is not None: + diagnostic_tables.append(overhead) + diagnostic_tables.append(target_wall_time_table_data(cell_maps[target], target, spec)) + peak_rss = target_peak_rss_table_data(rss_cell_maps[target], target, spec) + if peak_rss is not None: + diagnostic_tables.append(peak_rss) + if diagnostic_tables: + parts.append("## Target Diagnostics") + parts.extend(render_markdown_table(table_data) for table_data in diagnostic_tables) + + summary_tables: list[ReportTableData] = list(backend_summary_tables(cell_maps, rss_cell_maps, targets, spec)) + if len(targets) == 1: + summary_tables.append( + single_target_summary_table_data(cell_maps[targets[0]], rss_cell_maps[targets[0]], targets[0], spec) + ) + else: + summary_tables.extend(multi_target_summary_tables(cell_maps, rss_cell_maps, targets, spec)) + parts.append("## Benchmark Summary") + parts.extend(render_markdown_table(table_data) for table_data in summary_tables) + return "\n\n".join(part.strip() for part in parts if part.strip()) + + def render_targets_tree(console: Console, targets: Sequence[ResolvedTarget]) -> None: tree = Tree("[bold]Targets[/bold]", guide_style="dim") for index, target in enumerate(targets): @@ -1725,75 +2472,236 @@ def render_targets_tree(console: Console, targets: Sequence[ResolvedTarget]) -> console.print(tree) -def render_per_file_wall_time_change( - console: Console, +def targets_table_data(targets: Sequence[ResolvedTarget]) -> ReportTableData: + rows: list[tuple[str, ...]] = [] + for index, target in enumerate(targets): + role = "target" + if len(targets) > 1: + role = "baseline" if index == 0 else "candidate" + git = target.row.git_sha[:12] + if target.row.git_ref != "HEAD": + git = f"{git} ({target.row.git_ref})" + rows.append( + report_row( + role, + target.display_label, + git, + "yes" if target.row.is_dirty else "no", + target.binary_sha256.removeprefix("sha256:")[:12], + target.row.path, + ) + ) + return ReportTableData( + title="Targets", + headers=("Role", "Label", "Git", "Dirty", "Binary", "Path"), + rows=tuple(rows), + ) + + +def per_file_wall_time_change_tables( cell_maps: TargetCellMaps, targets: Sequence[ResolvedTarget], spec: BenchmarkSpec, -) -> None: +) -> tuple[ReportTableData, ...]: baseline = targets[0] + cells = benchmark_cells(spec) + tables: list[ReportTableData] = [] for target in targets[1:]: - table = report_table( - f"Per-file wall-time change vs {baseline.display_label}: {target.display_label}", - caption=TARGET_WALL_TIME_CAPTION, - ) - table.add_column("File") - table.add_column("Treatment") - table.add_column("Time ratio", no_wrap=True) - table.add_column("Wall-time change", no_wrap=True) - table.add_column("Result") + rows: list[tuple[str, ...]] = [] for file_spec in spec.files: - for treatment_index, treatment in enumerate(spec.treatments): + for cell_index, cell in enumerate(cells): ratio = ratio_summary( - cell_maps[baseline][(file_spec.sha256, treatment)], - cell_maps[target][(file_spec.sha256, treatment)], + cell_maps[baseline][(file_spec.sha256, cell.backend, cell.treatment)], + cell_maps[target][(file_spec.sha256, cell.backend, cell.treatment)], ) - table.add_row( - file_spec.display_path if treatment_index == 0 else "", - treatment, - format_ratio_summary(ratio), - format_wall_time_change(ratio), - format_comparison_result(ratio), - end_section=treatment_index == len(spec.treatments) - 1, + rows.append( + report_row( + file_spec.display_path if cell_index == 0 else "", + cell.backend, + cell.treatment, + format_ratio_summary(ratio), + format_wall_time_change(ratio), + comparison_result_text(ratio), + ) ) - console.print(table) + tables.append( + ReportTableData( + title=f"Per-file wall-time change vs {baseline.display_label}: {target.display_label}", + headers=("File", "Backend", "Treatment", "Time ratio", "Wall-time change", "Result"), + rows=tuple(rows), + caption=TARGET_WALL_TIME_CAPTION, + alignments=("left", "left", "left", "right", "right", "left"), + ) + ) + return tuple(tables) -def render_per_file_peak_rss_change( +def render_per_file_wall_time_change( console: Console, - rss_cell_maps: TargetCellMaps, + cell_maps: TargetCellMaps, targets: Sequence[ResolvedTarget], spec: BenchmarkSpec, ) -> None: + for table_data in per_file_wall_time_change_tables(cell_maps, targets, spec): + console.print(render_rich_table(table_data)) + + +def per_file_peak_rss_change_tables( + rss_cell_maps: TargetCellMaps, + targets: Sequence[ResolvedTarget], + spec: BenchmarkSpec, +) -> tuple[ReportTableData, ...]: baseline = targets[0] baseline_has_rss = any(cell.samples for cell in rss_cell_maps[baseline].values()) + cells = benchmark_cells(spec) + tables: list[ReportTableData] = [] for target in targets[1:]: if not baseline_has_rss and not any(cell.samples for cell in rss_cell_maps[target].values()): continue - table = report_table( - f"Per-file peak RSS change vs {baseline.display_label}: {target.display_label}", - caption=TARGET_PEAK_RSS_CAPTION, - ) - table.add_column("File") - table.add_column("Treatment") - table.add_column("RSS ratio", no_wrap=True) - table.add_column("RSS change", no_wrap=True) - table.add_column("Result") + rows: list[tuple[str, ...]] = [] for file_spec in spec.files: - for treatment_index, treatment in enumerate(spec.treatments): + for cell_index, cell in enumerate(cells): ratio = ratio_summary( - rss_cell_maps[baseline][(file_spec.sha256, treatment)], - rss_cell_maps[target][(file_spec.sha256, treatment)], + rss_cell_maps[baseline][(file_spec.sha256, cell.backend, cell.treatment)], + rss_cell_maps[target][(file_spec.sha256, cell.backend, cell.treatment)], ) - table.add_row( - file_spec.display_path if treatment_index == 0 else "", - treatment, - format_ratio_summary(ratio), - format_percent_change(ratio), - format_lower_is_better_result(ratio), - end_section=treatment_index == len(spec.treatments) - 1, + rows.append( + report_row( + file_spec.display_path if cell_index == 0 else "", + cell.backend, + cell.treatment, + format_ratio_summary(ratio), + format_percent_change(ratio), + lower_is_better_result_text(ratio), + ) + ) + tables.append( + ReportTableData( + title=f"Per-file peak RSS change vs {baseline.display_label}: {target.display_label}", + headers=("File", "Backend", "Treatment", "RSS ratio", "RSS change", "Result"), + rows=tuple(rows), + caption=TARGET_PEAK_RSS_CAPTION, + alignments=("left", "left", "left", "right", "right", "left"), + ) + ) + return tuple(tables) + + +def render_per_file_peak_rss_change( + console: Console, + rss_cell_maps: TargetCellMaps, + targets: Sequence[ResolvedTarget], + spec: BenchmarkSpec, +) -> None: + for table_data in per_file_peak_rss_change_tables(rss_cell_maps, targets, spec): + console.print(render_rich_table(table_data)) + + +def per_file_backend_wall_time_change_tables( + cell_maps: TargetCellMaps, + targets: Sequence[ResolvedTarget], + spec: BenchmarkSpec, +) -> tuple[ReportTableData, ...]: + baseline_backend = spec.backends[0] + tables: list[ReportTableData] = [] + for target in targets: + for backend in spec.backends[1:]: + shared_treatments = shared_backend_treatments(spec, baseline_backend, backend) + if not shared_treatments: + continue + rows: list[tuple[str, ...]] = [] + for file_spec in spec.files: + for treatment_index, treatment in enumerate(shared_treatments): + ratio = ratio_summary( + cell_maps[target][(file_spec.sha256, baseline_backend, treatment)], + cell_maps[target][(file_spec.sha256, backend, treatment)], + ) + rows.append( + report_row( + file_spec.display_path if treatment_index == 0 else "", + treatment, + format_ratio_summary(ratio), + format_wall_time_change(ratio), + comparison_result_text(ratio), + ) + ) + tables.append( + ReportTableData( + title=f"Per-file backend wall-time change vs {baseline_backend}: {target.display_label} {backend}", + headers=("File", "Treatment", "Time ratio", "Wall-time change", "Result"), + rows=tuple(rows), + caption=BACKEND_WALL_TIME_CAPTION, + alignments=("left", "left", "right", "right", "left"), ) - console.print(table) + ) + return tuple(tables) + + +def render_per_file_backend_wall_time_change( + console: Console, + cell_maps: TargetCellMaps, + targets: Sequence[ResolvedTarget], + spec: BenchmarkSpec, +) -> None: + for table_data in per_file_backend_wall_time_change_tables(cell_maps, targets, spec): + console.print(render_rich_table(table_data)) + + +def per_file_backend_peak_rss_change_tables( + rss_cell_maps: TargetCellMaps, + targets: Sequence[ResolvedTarget], + spec: BenchmarkSpec, +) -> tuple[ReportTableData, ...]: + baseline_backend = spec.backends[0] + baseline_has_rss = { + target: any(cell.samples for key, cell in rss_cell_maps[target].items() if key[1] == baseline_backend) + for target in targets + } + tables: list[ReportTableData] = [] + for target in targets: + for backend in spec.backends[1:]: + shared_treatments = shared_backend_treatments(spec, baseline_backend, backend) + if not shared_treatments: + continue + candidate_has_rss = any(cell.samples for key, cell in rss_cell_maps[target].items() if key[1] == backend) + if not baseline_has_rss[target] and not candidate_has_rss: + continue + rows: list[tuple[str, ...]] = [] + for file_spec in spec.files: + for treatment_index, treatment in enumerate(shared_treatments): + ratio = ratio_summary( + rss_cell_maps[target][(file_spec.sha256, baseline_backend, treatment)], + rss_cell_maps[target][(file_spec.sha256, backend, treatment)], + ) + rows.append( + report_row( + file_spec.display_path if treatment_index == 0 else "", + treatment, + format_ratio_summary(ratio), + format_percent_change(ratio), + lower_is_better_result_text(ratio), + ) + ) + tables.append( + ReportTableData( + title=f"Per-file backend peak RSS change vs {baseline_backend}: {target.display_label} {backend}", + headers=("File", "Treatment", "RSS ratio", "RSS change", "Result"), + rows=tuple(rows), + caption=BACKEND_PEAK_RSS_CAPTION, + alignments=("left", "left", "right", "right", "left"), + ) + ) + return tuple(tables) + + +def render_per_file_backend_peak_rss_change( + console: Console, + rss_cell_maps: TargetCellMaps, + targets: Sequence[ResolvedTarget], + spec: BenchmarkSpec, +) -> None: + for table_data in per_file_backend_peak_rss_change_tables(rss_cell_maps, targets, spec): + console.print(render_rich_table(table_data)) def render_benchmark_summary( @@ -1804,12 +2712,88 @@ def render_benchmark_summary( spec: BenchmarkSpec, ) -> None: console.rule("[bold]Benchmark summary[/bold]") + render_backend_summary(console, cell_maps, rss_cell_maps, targets, spec) if len(targets) == 1: render_single_target_summary(console, cell_maps[targets[0]], rss_cell_maps[targets[0]], targets[0], spec) else: render_multi_target_summary(console, cell_maps, rss_cell_maps, targets, spec) +def single_target_summary_table_data( + cell_map: CellMap, + rss_cell_map: CellMap, + target: ResolvedTarget, + spec: BenchmarkSpec, +) -> ReportTableData: + rows: list[tuple[str, ...]] = [] + for backend in spec.backends: + has_off_proofs = backend_has_treatment(spec, backend, "off") and backend_has_treatment(spec, backend, "proofs") + has_term_proofs = backend_has_treatment(spec, backend, "term") and backend_has_treatment( + spec, backend, "proofs" + ) + if has_off_proofs: + rows.extend( + within_target_wall_summary_rows( + cell_map, + spec, + backend, + "off", + "proofs", + backend_metric_label(spec, backend, "wall proofs/off"), + ) + ) + rows.append( + within_target_rss_summary_row( + rss_cell_map, + spec, + backend, + "off", + "proofs", + backend_metric_label(spec, backend, "peak RSS proofs/off"), + ) + ) + elif has_term_proofs: + rows.extend( + within_target_wall_summary_rows( + cell_map, + spec, + backend, + "term", + "proofs", + backend_metric_label(spec, backend, "wall proofs/term"), + ) + ) + else: + rows.append( + report_row( + backend_metric_label(spec, backend, "no proof baseline"), + "-", + "-", + "-", + "-", + "select off and proofs", + ) + ) + if has_term_proofs: + rows.append( + within_target_rss_summary_row( + rss_cell_map, + spec, + backend, + "term", + "proofs", + backend_metric_label(spec, backend, "peak RSS proofs/term"), + ) + ) + return ReportTableData( + title=f"{target.display_label}: proof overhead summary", + headers=("Metric", "Ratio", "Change", "Worst file", "Worst ratio", "Result"), + rows=tuple(rows), + caption=PROOF_OVERHEAD_CAPTION, + alignments=("left", "right", "right", "left", "right", "left"), + ) + + def render_single_target_summary( console: Console, cell_map: CellMap, @@ -1817,143 +2801,306 @@ def render_single_target_summary( target: ResolvedTarget, spec: BenchmarkSpec, ) -> None: - table = report_table(f"{target.display_label}: proof overhead summary", caption=PROOF_OVERHEAD_CAPTION) - table.add_column("Metric") - table.add_column("Ratio", no_wrap=True) - table.add_column("Change", no_wrap=True) - table.add_column("Worst file") - table.add_column("Worst ratio", no_wrap=True) - table.add_column("Result") - if "off" in spec.treatments and "proofs" in spec.treatments: - add_within_target_wall_summary_row(table, cell_map, spec, "off", "proofs", "wall proofs/off") - add_within_target_rss_summary_row(table, rss_cell_map, spec, "off", "proofs", "peak RSS proofs/off") - else: - table.add_row("no proof baseline", "-", "-", "-", "-", styled_status("descriptive", "select off and proofs")) - if "term" in spec.treatments and "proofs" in spec.treatments: - add_within_target_rss_summary_row(table, rss_cell_map, spec, "term", "proofs", "peak RSS proofs/term") - console.print(table) + console.print(render_rich_table(single_target_summary_table_data(cell_map, rss_cell_map, target, spec))) -def add_within_target_wall_summary_row( - table: Table, +def within_target_wall_summary_rows( cell_map: CellMap, spec: BenchmarkSpec, + backend: Backend, baseline_treatment: Treatment, candidate_treatment: Treatment, label: str, -) -> None: - file_cells = treatment_file_cells(cell_map, spec, baseline_treatment, candidate_treatment) +) -> tuple[tuple[str, ...], tuple[str, ...]]: + file_cells = treatment_file_cells(cell_map, spec, backend, baseline_treatment, candidate_treatment) pairs = summary_pairs(file_cells) summary = suite_ratio(pairs) geometric = geometric_mean_ratio(pairs) worst_file, worst = worst_file_ratio(ratio_pairs(file_cells)) - table.add_row( - label, - format_ratio_summary(summary), - format_percent_change(summary), - format_worst_file(worst_file), - format_ratio_summary(worst), - format_proof_gate_result(summary) + result = ( + proof_gate_result_text(summary) if baseline_treatment == "off" and candidate_treatment == "proofs" - else format_comparison_result(summary), + else comparison_result_text(summary) ) - table.add_row( - f"{label} geomean", - format_ratio_summary(geometric), - format_percent_change(geometric), - "-", - "-", - styled_status("descriptive", "descriptive"), + return ( + report_row( + label, + format_ratio_summary(summary), + format_percent_change(summary), + format_worst_file(worst_file), + format_ratio_summary(worst), + result, + ), + report_row( + f"{label} geomean", + format_ratio_summary(geometric), + format_percent_change(geometric), + "-", + "-", + "descriptive", + ), ) -def add_within_target_rss_summary_row( - table: Table, +def within_target_rss_summary_row( rss_cell_map: CellMap, spec: BenchmarkSpec, + backend: Backend, baseline_treatment: Treatment, candidate_treatment: Treatment, label: str, -) -> None: - file_cells = treatment_file_cells(rss_cell_map, spec, baseline_treatment, candidate_treatment) +) -> tuple[str, ...]: + file_cells = treatment_file_cells(rss_cell_map, spec, backend, baseline_treatment, candidate_treatment) pairs = summary_pairs(file_cells) summary = suite_ratio(pairs) worst_file, worst = worst_file_ratio(ratio_pairs(file_cells)) - table.add_row( + return report_row( label, format_ratio_summary(summary), format_percent_change(summary), format_worst_file(worst_file), format_ratio_summary(worst), - format_lower_is_better_result(summary), + lower_is_better_result_text(summary), ) -def render_multi_target_summary( - console: Console, +def multi_target_summary_tables( cell_maps: TargetCellMaps, rss_cell_maps: TargetCellMaps, targets: Sequence[ResolvedTarget], spec: BenchmarkSpec, -) -> None: +) -> tuple[ReportTableData, ReportTableData]: baseline = targets[0] - wall_table = report_table( - f"Wall-time summary vs {baseline.display_label}", - caption=TARGET_WALL_TIME_CAPTION, - ) - wall_table.add_column("Target") - wall_table.add_column("Treatment") - wall_table.add_column("Wall-time change", no_wrap=True) - wall_table.add_column("Geomean", no_wrap=True) - wall_table.add_column("Worst file") - wall_table.add_column("Worst ratio", no_wrap=True) - wall_table.add_column("Result") + cells = benchmark_cells(spec) + wall_rows: list[tuple[str, ...]] = [] for target in targets[1:]: - for treatment in spec.treatments: - file_cells = target_treatment_file_cells(cell_maps[baseline], cell_maps[target], spec, treatment) + for cell in cells: + file_cells = target_treatment_file_cells( + cell_maps[baseline], cell_maps[target], spec, cell.backend, cell.treatment + ) pairs = summary_pairs(file_cells) suite = suite_ratio(pairs) geometric = geometric_mean_ratio(pairs) worst_file, worst = worst_file_ratio(ratio_pairs(file_cells)) - wall_table.add_row( - target.display_label, - treatment, - format_wall_time_change(suite), - format_ratio_summary(geometric), - format_worst_file(worst_file), - format_ratio_summary(worst), - format_comparison_result(suite), + wall_rows.append( + report_row( + target.display_label, + cell.backend, + cell.treatment, + format_wall_time_change(suite), + format_ratio_summary(geometric), + format_worst_file(worst_file), + format_ratio_summary(worst), + comparison_result_text(suite), + ) ) - console.print(wall_table) - rss_table = report_table( - f"Peak RSS summary vs {baseline.display_label}", - caption=TARGET_PEAK_RSS_CAPTION, - ) - rss_table.add_column("Target") - rss_table.add_column("Treatment") - rss_table.add_column("RSS change", no_wrap=True) - rss_table.add_column("Geomean", no_wrap=True) - rss_table.add_column("Worst file") - rss_table.add_column("Worst ratio", no_wrap=True) - rss_table.add_column("Result") + rss_rows: list[tuple[str, ...]] = [] for target in targets[1:]: - for treatment in spec.treatments: - file_cells = target_treatment_file_cells(rss_cell_maps[baseline], rss_cell_maps[target], spec, treatment) + for cell in cells: + file_cells = target_treatment_file_cells( + rss_cell_maps[baseline], + rss_cell_maps[target], + spec, + cell.backend, + cell.treatment, + ) pairs = summary_pairs(file_cells) summary = suite_ratio(pairs) geometric = geometric_mean_ratio(pairs) worst_file, worst = worst_file_ratio(ratio_pairs(file_cells)) - rss_table.add_row( - target.display_label, - treatment, - format_percent_change(summary), - format_ratio_summary(geometric), - format_worst_file(worst_file), - format_ratio_summary(worst), - format_lower_is_better_result(summary), + rss_rows.append( + report_row( + target.display_label, + cell.backend, + cell.treatment, + format_percent_change(summary), + format_ratio_summary(geometric), + format_worst_file(worst_file), + format_ratio_summary(worst), + lower_is_better_result_text(summary), + ) ) - console.print(rss_table) + return ( + ReportTableData( + title=f"Wall-time summary vs {baseline.display_label}", + headers=( + "Target", + "Backend", + "Treatment", + "Wall-time change", + "Geomean", + "Worst file", + "Worst ratio", + "Result", + ), + rows=tuple(wall_rows), + caption=TARGET_WALL_TIME_CAPTION, + alignments=("left", "left", "left", "right", "right", "left", "right", "left"), + ), + ReportTableData( + title=f"Peak RSS summary vs {baseline.display_label}", + headers=("Target", "Backend", "Treatment", "RSS change", "Geomean", "Worst file", "Worst ratio", "Result"), + rows=tuple(rss_rows), + caption=TARGET_PEAK_RSS_CAPTION, + alignments=("left", "left", "left", "right", "right", "left", "right", "left"), + ), + ) + + +def render_multi_target_summary( + console: Console, + cell_maps: TargetCellMaps, + rss_cell_maps: TargetCellMaps, + targets: Sequence[ResolvedTarget], + spec: BenchmarkSpec, +) -> None: + for table_data in multi_target_summary_tables(cell_maps, rss_cell_maps, targets, spec): + console.print(render_rich_table(table_data)) + + +def display_backend(backend: Backend) -> str: + return backend_spec(backend).display_name + + +def backend_summary_table( + cell_maps: TargetCellMaps, + targets: Sequence[ResolvedTarget], + spec: BenchmarkSpec, + metric: MetricSpec, +) -> ReportTableData: + baseline_backend = spec.backends[0] + baseline_name = display_backend(baseline_backend) + include_target = len(targets) > 1 + if len(spec.backends) == 2: + candidate_name = display_backend(spec.backends[1]) + title = f"{candidate_name} vs {baseline_name} {metric.title}" + else: + candidate_name = "Candidate" + title = f"Backend {metric.title} summary vs {baseline_name}" + ratio_heading = f"{candidate_name}/{baseline_name}" + candidate_total_heading = f"{candidate_name} total" + headers: list[str] = [] + if include_target: + headers.append("Target") + headers.extend( + [ + "Backend", + "Mode", + f"{baseline_name} total", + candidate_total_heading, + ratio_heading, + "Change", + "File geomean", + metric.file_count_heading, + "Best file", + "Best ratio", + "Best result", + ] + ) + rows: list[tuple[str, ...]] = [] + for target in targets: + for backend in spec.backends[1:]: + for treatment in shared_backend_treatments(spec, baseline_backend, backend): + file_cells = backend_treatment_file_cells( + cell_maps[target], + spec, + baseline_backend, + backend, + treatment, + ) + pairs = summary_pairs(file_cells) + summary = suite_ratio(pairs) + geometric = geometric_mean_ratio(pairs) + ratios = ratio_pairs(file_cells) + best_file, best = best_file_ratio(ratios) + row_values: list[object] = [] + if include_target: + row_values.append(target.display_label) + row_values.extend( + [ + backend, + treatment, + metric.format_total(suite_total_mean([baseline for baseline, _ in pairs])), + metric.format_total(suite_total_mean([candidate for _, candidate in pairs])), + format_ratio_summary(summary), + metric.format_change(summary), + format_ratio_summary(geometric), + format_better_file_count(count_better_files(ratios)), + format_worst_file(best_file), + format_ratio_summary(best), + metric.format_result(best), + ] + ) + rows.append(report_row(*row_values)) + right_aligned = { + f"{baseline_name} total", + candidate_total_heading, + ratio_heading, + "Change", + "File geomean", + metric.file_count_heading, + "Best ratio", + } + return ReportTableData( + title=title, + headers=tuple(headers), + rows=tuple(rows), + caption=metric.caption, + alignments=tuple("right" if header in right_aligned else "left" for header in headers), + ) + + +def backend_summary_tables( + cell_maps: TargetCellMaps, + rss_cell_maps: TargetCellMaps, + targets: Sequence[ResolvedTarget], + spec: BenchmarkSpec, +) -> tuple[ReportTableData, ...]: + if len(spec.backends) <= 1: + return () + metrics = ( + ( + cell_maps, + MetricSpec( + title="wall time", + caption=BACKEND_WALL_TIME_CAPTION, + file_count_heading="Faster files", + format_total=format_seconds_value, + format_change=format_wall_time_change, + format_result=comparison_result_text, + ), + ), + ( + rss_cell_maps, + MetricSpec( + title="peak RSS", + caption=BACKEND_PEAK_RSS_CAPTION, + file_count_heading="Lower-RSS files", + format_total=format_bytes, + format_change=format_percent_change, + format_result=lower_is_better_result_text, + omit_without_samples=True, + ), + ), + ) + return tuple( + backend_summary_table(metric_cell_maps, targets, spec, metric) + for metric_cell_maps, metric in metrics + if not metric.omit_without_samples + or any(cell.samples for target in targets for cell in metric_cell_maps[target].values()) + ) + + +def render_backend_summary( + console: Console, + cell_maps: TargetCellMaps, + rss_cell_maps: TargetCellMaps, + targets: Sequence[ResolvedTarget], + spec: BenchmarkSpec, +) -> None: + for table_data in backend_summary_tables(cell_maps, rss_cell_maps, targets, spec): + console.print(render_rich_table(table_data)) def format_worst_file(file_spec: FileSpec | None) -> str: @@ -1968,36 +3115,42 @@ def render_target_diagnostics( spec: BenchmarkSpec, ) -> None: render_overhead_ratios(console, cell_map, target, spec) + console.print(render_rich_table(target_wall_time_table_data(cell_map, target, spec))) + render_peak_rss_diagnostics(console, rss_cell_map, target, spec) - means_table = report_table( - f"{target.display_label}: per-file wall time", - caption="Within-target wall-time estimates. These are not target-vs-baseline ratios.", - ) - means_table.add_column("File") - for treatment in spec.treatments: - means_table.add_column(treatment, no_wrap=True) - means_rows: list[tuple[list[str], str]] = [] + +def target_wall_time_table_data(cell_map: CellMap, target: ResolvedTarget, spec: BenchmarkSpec) -> ReportTableData: + cells = benchmark_cells(spec) + headers = ["File", *(cell_label(spec, cell.backend, cell.treatment) for cell in cells)] + rows_with_issue: list[tuple[list[str], str]] = [] has_mean_issues = False for file_spec in spec.files: issue_parts: list[str] = [] row_values = [file_spec.display_path] - for treatment in spec.treatments: - cell = cell_map[(file_spec.sha256, treatment)] - row_values.append(format_seconds_summary(cell)) - if cell.issue is not None: - issue_parts.append(f"{treatment}: {cell.issue}") + for cell in cells: + summary = cell_map[(file_spec.sha256, cell.backend, cell.treatment)] + row_values.append(format_seconds_summary(summary)) + if summary.issue is not None: + issue_parts.append(f"{cell_label(spec, cell.backend, cell.treatment)}: {summary.issue}") issue_text = "; ".join(issue_parts) has_mean_issues = has_mean_issues or bool(issue_text) - means_rows.append((row_values, issue_text)) + rows_with_issue.append((row_values, issue_text)) if has_mean_issues: - means_table.add_column("Issue") - for row_values, issue_text in means_rows: + headers.append("Issue") + rows: list[tuple[str, ...]] = [] + for row_values, issue_text in rows_with_issue: if has_mean_issues: row_values.append(issue_text) - means_table.add_row(*row_values) - console.print(means_table) - - render_peak_rss_diagnostics(console, rss_cell_map, target, spec) + rows.append(tuple(row_values)) + return ReportTableData( + title=f"{target.display_label}: per-file wall time", + headers=tuple(headers), + rows=tuple(rows), + caption="Within-target wall-time estimates. These are not target-vs-baseline ratios.", + alignments=tuple( + "left" if index == 0 or header == "Issue" else "right" for index, header in enumerate(headers) + ), + ) def render_overhead_ratios( @@ -2006,26 +3159,39 @@ def render_overhead_ratios( target: ResolvedTarget, spec: BenchmarkSpec, ) -> None: - ratio_columns = ratio_specs(spec.treatments) - if not ratio_columns: - return - ratio_table = report_table( - f"{target.display_label}: overhead ratios", - caption="Within-target treatment ratios. These are not target-vs-baseline wall-time change.", - ) - ratio_table.add_column("File") - for _, _, ratio_name in ratio_columns: - ratio_table.add_column(ratio_name, no_wrap=True) + table_data = overhead_ratios_table_data(cell_map, target, spec) + if table_data is not None: + console.print(render_rich_table(table_data)) + + +def overhead_ratios_table_data( + cell_map: CellMap, target: ResolvedTarget, spec: BenchmarkSpec +) -> ReportTableData | None: + ratio_columns = {backend: ratio_specs_for_backend(spec, backend) for backend in spec.backends} + if not any(ratio_columns.values()): + return None + headers = ["File"] + for backend in spec.backends: + for _, _, ratio_name in ratio_columns[backend]: + headers.append(backend_metric_label(spec, backend, ratio_name)) + rows: list[tuple[str, ...]] = [] for file_spec in spec.files: row_values = [file_spec.display_path] - for baseline_treatment, candidate_treatment, _ in ratio_columns: - ratio = ratio_summary( - cell_map[(file_spec.sha256, baseline_treatment)], - cell_map[(file_spec.sha256, candidate_treatment)], - ) - row_values.append(format_ratio_summary(ratio)) - ratio_table.add_row(*row_values) - console.print(ratio_table) + for backend in spec.backends: + for baseline_treatment, candidate_treatment, _ in ratio_columns[backend]: + ratio = ratio_summary( + cell_map[(file_spec.sha256, backend, baseline_treatment)], + cell_map[(file_spec.sha256, backend, candidate_treatment)], + ) + row_values.append(format_ratio_summary(ratio)) + rows.append(tuple(row_values)) + return ReportTableData( + title=f"{target.display_label}: overhead ratios", + headers=tuple(headers), + rows=tuple(rows), + caption="Within-target treatment ratios. These are not target-vs-baseline wall-time change.", + alignments=tuple("left" if index == 0 else "right" for index, _ in enumerate(headers)), + ) def render_peak_rss_diagnostics( @@ -2034,35 +3200,49 @@ def render_peak_rss_diagnostics( target: ResolvedTarget, spec: BenchmarkSpec, ) -> None: + table_data = target_peak_rss_table_data(rss_cell_map, target, spec) + if table_data is not None: + console.print(render_rich_table(table_data)) + + +def target_peak_rss_table_data( + rss_cell_map: CellMap, + target: ResolvedTarget, + spec: BenchmarkSpec, +) -> ReportTableData | None: if not any(cell.samples for cell in rss_cell_map.values()): - return - rss_table = report_table( - f"{target.display_label}: per-file peak RSS", - caption="Within-target peak resident set size estimates. These are separate from wall-time ratios.", - ) - rss_table.add_column("File") - for treatment in spec.treatments: - rss_table.add_column(treatment, no_wrap=True) + return None + cells = benchmark_cells(spec) + headers = ["File", *(cell_label(spec, cell.backend, cell.treatment) for cell in cells)] rss_rows: list[tuple[list[str], str]] = [] has_rss_issues = False for file_spec in spec.files: issue_parts: list[str] = [] row_values = [file_spec.display_path] - for treatment in spec.treatments: - cell = rss_cell_map[(file_spec.sha256, treatment)] - row_values.append(format_bytes_summary(cell)) - if cell.issue is not None: - issue_parts.append(f"{treatment}: {cell.issue}") + for cell in cells: + summary = rss_cell_map[(file_spec.sha256, cell.backend, cell.treatment)] + row_values.append(format_bytes_summary(summary)) + if summary.issue is not None: + issue_parts.append(f"{cell_label(spec, cell.backend, cell.treatment)}: {summary.issue}") issue_text = "; ".join(issue_parts) has_rss_issues = has_rss_issues or bool(issue_text) rss_rows.append((row_values, issue_text)) if has_rss_issues: - rss_table.add_column("Issue") + headers.append("Issue") + rows: list[tuple[str, ...]] = [] for row_values, issue_text in rss_rows: if has_rss_issues: row_values.append(issue_text) - rss_table.add_row(*row_values) - console.print(rss_table) + rows.append(tuple(row_values)) + return ReportTableData( + title=f"{target.display_label}: per-file peak RSS", + headers=tuple(headers), + rows=tuple(rows), + caption="Within-target peak resident set size estimates. These are separate from wall-time ratios.", + alignments=tuple( + "left" if index == 0 or header == "Issue" else "right" for index, header in enumerate(headers) + ), + ) def ratio_specs(treatments: Sequence[Treatment]) -> tuple[tuple[Treatment, Treatment, str], ...]: @@ -2076,14 +3256,139 @@ def ratio_specs(treatments: Sequence[Treatment]) -> tuple[tuple[Treatment, Treat return tuple(specs) +def ratio_specs_for_backend( + spec: BenchmarkSpec, + backend: Backend, +) -> tuple[tuple[Treatment, Treatment, str], ...]: + return tuple( + ratio_spec + for ratio_spec in ratio_specs(spec.treatments) + if backend_has_treatment(spec, backend, ratio_spec[0]) and backend_has_treatment(spec, backend, ratio_spec[1]) + ) + + +def run_profile(args: argparse.Namespace, output: RunnerOutput, invocation_cwd: Path, repo_root: Path) -> None: + request = resolve_profile_request(args, invocation_cwd) + target = resolve_profile_target(request.target_request, request.backend, invocation_cwd, repo_root, output) + if target.binary_path is None: + raise ValueError(f"target {target.display_label} needs a profiling binary") + checkout_path = Path(target.row.path) + workload = workload_command(target.binary_path, request.file, request.backend, request.treatment) + artifact = profile_cache_path( + request.profiles_dir, + target.binary_sha256, + request.file.sha256, + request.backend, + request.treatment, + request.mode, + ) + profile: dict[str, Any] | None = None + cache_status: Literal["hit", "recorded"] = "recorded" + if profile_cache_hit(artifact) and not request.force_run: + try: + profile = samply_analysis.read_artifact(artifact) + except ValueError as error: + output.console.print(f"[yellow]warning:[/yellow] ignoring invalid profile cache entry: {error}") + else: + cache_status = "hit" + output.console.print(f"[bold]Profile cache hit[/bold] {artifact}") + + if profile is None: + iterations = request.mode.iterations + if iterations is None: + assert request.mode.profile_seconds is not None + output.console.print( + f"[bold]Calibrating[/bold] {request.file.display_path} " + f"{request.backend}/{request.treatment} for {request.mode.profile_seconds}s" + ) + calibration = run_command(workload, checkout_path, request.timeout_sec) + if calibration.status != "success" or calibration.timing.wall_sec is None: + detail = calibration.error.message if calibration.error is not None else calibration.status + raise ValueError(f"profile calibration failed: {detail}") + iterations, capped = calculate_profile_iterations( + calibration.timing.wall_sec, + request.mode.profile_seconds, + ) + output.console.print( + f" calibration: {calibration.timing.wall_sec:.3f}s; recording {iterations} Samply iteration(s)" + ) + if capped: + output.console.print( + "[yellow]warning:[/yellow] maximum profile iterations reached; " + "the profile may be shorter than the requested duration" + ) + + assert iterations is not None + name = profile_name( + request.file, + request.backend, + request.treatment, + request.mode, + iterations, + target.binary_sha256, + ) + output.console.print(f"[bold]Recording profile[/bold] {artifact}") + profile = run_samply_record( + artifact=artifact, + name=name, + iterations=iterations, + workload=workload, + checkout_path=checkout_path, + timeout_sec=request.timeout_sec, + ) + output.console.print(f"[bold]Profile written[/bold] {artifact}") + + if request.show_summary: + display_artifact = profile_display_path(artifact, invocation_cwd) + summary: samply_analysis.ProfileCpuSummary | None = None + if sys.platform == "darwin": + try: + summary = samply_analysis.summarize(profile, target.binary_path) + except ValueError as error: + output.console.print(f"[yellow]warning:[/yellow] CPU profile summary unavailable: {error}") + if summary is not None: + for warning in summary.warnings: + output.console.print(f"[yellow]warning:[/yellow] {warning}") + else: + output.console.print( + "[yellow]warning:[/yellow] CPU profile summaries are currently available on macOS only; " + "the Samply artifact was created normally." + ) + report = samply_analysis.ProfileReport( + artifact=display_artifact, + cache_status=cache_status, + workload=request.file.display_path, + backend=request.backend, + treatment=request.treatment, + top=request.top, + cpu_summary=summary, + ) + if request.output_format == "markdown": + rendered = samply_analysis.render_markdown(report) + sys.stdout.write(rendered + "\n") + sys.stdout.flush() + else: + samply_analysis.render_rich(output.console, report) + else: + sys.stdout.write(str(artifact.resolve()) + "\n") + sys.stdout.flush() + if request.open_after: + open_samply_profile(artifact, checkout_path) + + def main(argv: Sequence[str] | None = None) -> int: - args = parse_args(sys.argv[1:] if argv is None else argv) + raw_argv = tuple(sys.argv[1:] if argv is None else argv) + args = parse_args(raw_argv) output = RunnerOutput() try: script_root = Path(__file__).resolve().parent - treatments = parse_treatments(args.treatments) invocation_cwd = Path.cwd() repo_root = git_root_for_path(script_root) + if args.command == "profile": + run_profile(args, output, invocation_cwd, repo_root) + return 0 + backends = parse_backends(args.backend) + treatments = parse_treatments(args.treatments) report_destination = resolve_report_destination(args.report, invocation_cwd) files = resolve_files(args.files, invocation_cwd) spec = BenchmarkSpec( @@ -2091,6 +3396,7 @@ def main(argv: Sequence[str] | None = None) -> int: treatments=treatments, rounds=args.rounds, timeout_sec=args.timeout_sec, + backends=backends, ) validate_spec(spec) target_specs = args.target if args.target is not None else ["."] @@ -2119,14 +3425,17 @@ def main(argv: Sequence[str] | None = None) -> int: emit_collection_plan(output, plan, estimate_model) collection = collect_rows(rows, report_destination, plan, spec, output, estimate_model) rows = collection.rows - render_report( - output.console, - report_destination, - rows, - targets, - spec, - ) - except (FileNotFoundError, ValueError, subprocess.CalledProcessError) as error: + if args.format == "markdown": + sys.stdout.write(render_markdown_report(report_destination, rows, targets, spec, raw_argv) + "\n") + else: + render_report( + output.console, + report_destination, + rows, + targets, + spec, + ) + except (FileNotFoundError, ValueError, subprocess.CalledProcessError, subprocess.TimeoutExpired) as error: output.print_error(error) return 2 return 0 diff --git a/egglog-experimental/Cargo.toml b/egglog-experimental/Cargo.toml index 739105f..4c25352 100644 --- a/egglog-experimental/Cargo.toml +++ b/egglog-experimental/Cargo.toml @@ -19,11 +19,13 @@ name = "files" [features] default = ["bin"] bin = ["egglog/bin"] +dd-backend = ["dep:egglog-experimental-dd"] [dependencies] egglog = { workspace = true } egglog-ast = { workspace = true } egglog-reports = { workspace = true } +egglog-experimental-dd = { path = "dd", optional = true } num = "0.4.3" lazy_static = "1.4" diff --git a/egglog-experimental/dd/Cargo.toml b/egglog-experimental/dd/Cargo.toml new file mode 100644 index 0000000..616b3d6 --- /dev/null +++ b/egglog-experimental/dd/Cargo.toml @@ -0,0 +1,38 @@ +# An experimental differential-dataflow backend for egglog. It depends only on +# the public backend SPI plus the neutral shared rule AST and core relations, +# so it can serve as an out-of-tree backend example. +[package] +name = "egglog-experimental-dd" +version = "0.1.0" +edition = "2021" +description = "An experimental differential-dataflow-backed egglog Backend, as an out-of-tree example of the egglog-backend-trait SPI." +license = "MIT" +publish = false + +[dependencies] +anyhow = "1" +egglog-backend-trait = { path = "../../egglog/egglog-backend-trait" } +egglog-ast = { path = "../../egglog/egglog-ast" } +egglog-core-relations = { path = "../../egglog/core-relations" } +egglog-numeric-id = { path = "../../egglog/numeric-id" } +hashbrown = "0.16" +log = "0.4" +# Raw differential-dataflow + timely, used in-process by `dd_native.rs`. +differential-dataflow = "=0.24.0" +timely = "=0.30.0" +serde = "1" + +[[test]] +# Parity harness over egglog's .egg corpus: checks the DD backend against +# egglog's OWN shared snapshots (`egglog/tests/snapshots/files__shared_snapshot_ +# *.snap`), so egglog is the single source of truth. Custom libtest-mimic main. +name = "files" +harness = false + +[dev-dependencies] +# The frontend, so the integration tests can run real `.egg` files on the +# Differential Dataflow backend via `egglog::EGraph::with_backend(..)`. +egglog = { path = "../../egglog", default-features = false } +egglog-experimental = { path = "..", default-features = false } +insta = "1.40" +libtest-mimic = "0.8" diff --git a/egglog-experimental/dd/src/compile.rs b/egglog-experimental/dd/src/compile.rs new file mode 100644 index 0000000..289052f --- /dev/null +++ b/egglog-experimental/dd/src/compile.rs @@ -0,0 +1,54 @@ +//! DD-specific physical values and merge plans. + +use egglog_ast::core::GenericAtomTerm; +use egglog_backend_trait::{ExternalFunctionId, FunctionId, ReadMode, RuleValue, RuleVar}; +use egglog_numeric_id::NumericId; + +/// Variable-width row stored in the host-side relation mirror. +pub type Row = Box<[u32]>; + +/// How a function resolves a functional-dependency conflict. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum MergeMode { + Relation, + Old, + New, + Min, + Computed, +} + +/// An evaluatable computed merge expression. +#[derive(Clone, Debug)] +pub enum MergeTree { + Old, + New, + Const(u32), + Prim(ExternalFunctionId, Vec), + Func(FunctionId, Vec), +} + +/// A physical column operand used by a DD join plan. +#[derive(Clone, Debug)] +pub enum Slot { + Var(u32), + Const(u32), +} + +impl Slot { + pub fn from_term(term: &GenericAtomTerm) -> Result { + match term { + GenericAtomTerm::Var(_, variable) => Ok(Self::Var(variable.id)), + GenericAtomTerm::Literal(_, value) => Ok(Self::Const(value.value.rep())), + GenericAtomTerm::Global(..) => { + Err("globals must be desugared before DD rule lowering".into()) + } + } + } +} + +/// A distinct DD input stream for one table read view. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct ReadKey { + pub func: FunctionId, + pub mode: ReadMode, +} diff --git a/egglog-experimental/dd/src/dd_native.rs b/egglog-experimental/dd/src/dd_native.rs new file mode 100644 index 0000000..0ac38a7 --- /dev/null +++ b/egglog-experimental/dd/src/dd_native.rs @@ -0,0 +1,991 @@ +//! In-process, build-once, epoch-driven incremental body join on RAW +//! `differential-dataflow` + `timely`. +//! +//! This is the only join path for the Differential Dataflow backend (driven by +//! [`crate::interpret::run_iteration`]); there is no host nested-loop fallback. +//! Unsupported shapes are reported to the caller; see [`plan_join`]. +//! +//! ## The join: one fused worker per ruleset +//! +//! [`FusedDdJoin`] is the path the interpreter drives: ONE shared timely +//! `Worker` hosts ONE dataflow for a whole RULESET, every distinct body relation +//! is a single SHARED input `Collection`, and each rule is a join sub-stream +//! reading those shared collections. It feeds only signed DELTAS into +//! never-cleared InputSessions, so the arrangements persist across epochs = +//! incremental. +//! +//! ## The base architecture +//! +//! Each ruleset owns one dataflow in a single-threaded timely [`Worker`] that is +//! stepped across host calls. Every body relation view has one shared +//! [`InputSession`], while each atom-bearing rule lowers to a left-deep join +//! subgraph. Table constraints run in the dataflow; body primitives run later +//! in the host interpreter. Binding deltas flow through `.inspect_batch` into +//! per-rule capture buffers. +//! +//! The host feeds only per-relation signed deltas (`+1` insert, `-1` retract) +//! into the input sessions, advances the timely timestamp for each nonempty +//! delta batch, drives the worker to that epoch's frontier, and drains the +//! output binding deltas. One egglog iteration may submit a removal batch and +//! then an insertion batch to preserve reinsertion freshness. The +//! InputSessions are NEVER cleared — the DD arrangements persist across epochs, +//! which is what makes the join genuinely incremental (epoch K does only +//! delta·integral work, not a full recompute) — the whole point of the design. +//! +//! ## Fixpoint structure +//! +//! We use EXTERNAL epoch-drive (the host loop advances epochs and feeds head +//! outputs back as the next epoch's inputs), NOT an in-dataflow `iterate()` +//! scope. This matches egglog's bounded `(run N)` fire->rebuild->repeat model and +//! sidesteps DD `iterate()`'s monotonicity constraints under retraction (a +//! rebuild RETRACTS non-canonical rows, which `iterate()` cannot express +//! cleanly). The dataflow itself is NON-recursive: one epoch = one bounded hop. + +use std::cell::RefCell; +use std::rc::Rc; + +use anyhow::{anyhow, bail, Result}; +use differential_dataflow::input::InputSession; +use egglog_ast::core::{GenericAtom, GenericAtomTerm, GenericCoreAction}; +use egglog_backend_trait::{RuleActionCall, RuleBodyCall, RuleSpec, RuleValue, RuleVar}; +use hashbrown::HashMap; +use timely::communication::allocator::thread::Thread; +use timely::communication::allocator::Allocator; +use timely::dataflow::operators::probe::Handle as ProbeHandle; +use timely::dataflow::operators::{Inspect, Probe}; +use timely::worker::Worker; +use timely::WorkerConfig; + +use crate::compile::{ReadKey, Slot}; + +/// A signed `(row, weight)` delta for one relation (`+1` inserted, `-1` +/// retracted), with rows as plain `Vec`. +type SignedDelta = Vec<(Vec, isize)>; +/// Per-relation-view input deltas fed into one [`FusedDdJoin::step`]. +type DeltaMap = HashMap; +/// One `step`'s captured output deltas, parallel to the fused join's rule list. +type StepOutput = Vec; +/// A per-rule output-capture buffer shared with the DD closure (fixed-width +/// [`Row`] rows). +type CaptureBuf = Rc>>; + +/// Fixed binding-row width (DD `Data` needs a `Sized + Ord + Hash` type; an +/// array gives us that). Set to 48 to cover the widest live-variable frontier +/// in the backend corpus: `luminal-llama`'s `@rebuild_rule34` uses 35 distinct +/// body vars in a wide congruence-closure rebuild. A larger live frontier is +/// reported as a row-width-cap wall. Raising `W` extends coverage at a cost of +/// `W * 4` bytes per relation or binding row. +pub const W: usize = 48; + +/// A fixed-width relation or binding row flowing through the DD dataflow. Input +/// rows store relation columns in the low slots. Intermediate binding columns +/// are assigned by the current `ProjectionPlan` stage, and captured outputs +/// repack surviving variables into the low slots in [`JoinPlan::var_order`]. +/// +/// A NEWTYPE over `[u32; W]` (rather than the bare array) because timely's +/// `ExchangeData` bound required by DD joins is +/// `Serialize + Deserialize`, and `serde` only derives those for arrays up to +/// length 32. The hand-written serde impl (serialize as a fixed-length seq of +/// `W` `u32`s) lifts that cap so `W` can exceed 32 (the corpus needs 35). All +/// other derives (`Ord`/`Hash`/`Clone`/`Copy`) are auto for any array size. +#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)] +pub struct Row([u32; W]); + +impl std::ops::Index for Row { + type Output = u32; + #[inline] + fn index(&self, i: usize) -> &u32 { + &self.0[i] + } +} + +impl std::ops::IndexMut for Row { + #[inline] + fn index_mut(&mut self, i: usize) -> &mut u32 { + &mut self.0[i] + } +} + +impl serde::Serialize for Row { + fn serialize(&self, s: S) -> Result { + use serde::ser::SerializeTuple; + // Fixed-length tuple of W u32s — bincode-friendly, no length prefix + // needed (the deserializer knows W). Sidesteps serde's 32-array cap. + let mut t = s.serialize_tuple(W)?; + for v in &self.0 { + t.serialize_element(v)?; + } + t.end() + } +} + +impl<'de> serde::Deserialize<'de> for Row { + fn deserialize>(d: D) -> Result { + struct RowVisitor; + impl<'de> serde::de::Visitor<'de> for RowVisitor { + type Value = Row; + fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "a tuple of {W} u32s") + } + fn visit_seq>(self, mut seq: A) -> Result { + let mut a = [0u32; W]; + for (i, slot) in a.iter_mut().enumerate() { + *slot = seq + .next_element()? + .ok_or_else(|| serde::de::Error::invalid_length(i, &self))?; + } + Ok(Row(a)) + } + } + d.deserialize_tuple(W, RowVisitor) + } +} + +impl Default for Row { + fn default() -> Self { + Row([0; W]) + } +} + +/// A planned DD join: table atoms plus a bounded live-variable layout. +pub struct JoinPlan { + /// Body table atoms in emission order. + atoms: Vec, + /// Per-step column allocation for variables live at each join stage. + projection: ProjectionPlan, +} + +struct PlanAtom { + read_key: ReadKey, + slots: Vec, +} + +impl JoinPlan { + /// The variable id at each captured binding-row column, in column order. + #[cfg(test)] + pub fn var_order(&self) -> Vec { + self.projection.head_vars.clone() + } +} + +/// Build the join plan for `rule`, or `Err(reason)` if the DD dataflow cannot +/// support its shape. +/// Supported rules have one or more table atoms, atom arity at most [`W`], and +/// no more than [`W`] simultaneously live body variables. Body primitives are +/// evaluated later by the host interpreter, so they do not become DD operators. +pub fn plan_join(rule: &RuleSpec) -> Result { + use hashbrown::HashSet; + + let mut body_vars = HashSet::new(); + let mut atoms: Vec = Vec::new(); + + for atom in &rule.core.body.atoms { + match atom.head { + RuleBodyCall::Table { id, read } => { + if atom.args.len() > W { + return Err(format!("atom arity {} > W {}", atom.args.len(), W)); + } + let slots = atom + .args + .iter() + .map(Slot::from_term) + .collect::, _>>()?; + for s in &slots { + if let Slot::Var(v) = s { + body_vars.insert(*v); + } + } + atoms.push(PlanAtom { + read_key: ReadKey { + func: id, + mode: read, + }, + slots, + }); + } + // The host replays body primitives over each completed table + // binding. Their variables remain live through the projection. + RuleBodyCall::Primitive { .. } => {} + } + } + + if atoms.is_empty() { + return Err("no body table atoms (atom-less rule)".to_string()); + } + let projection = build_projection(&atoms, &rule.core.head.0, &rule.core.body.atoms) + .ok_or_else(|| { + format!( + "live body-variable frontier exceeds W {W} ({} distinct body variables)", + body_vars.len() + ) + })?; + Ok(JoinPlan { atoms, projection }) +} + +/// Per-step binding columns produced by linear-scan allocation over body-atom +/// liveness. Reusing dead variables' columns lets some rules with more than +/// [`W`] total variables fit in a fixed-width [`Row`]. +#[derive(Clone, Debug)] +struct ProjectionPlan { + /// `step_col[i]` maps each variable LIVE during step `i` (atom `i`'s join) to + /// its binding-row column at that step. A var's column is stable from its + /// birth step to its death step but may differ across non-overlapping vars. + step_col: Vec>, + /// The surviving (head/body-prim-relevant) variables and their FINAL columns, + /// in a deterministic order. Drives the reduced head-scatter `var_order`. + head_vars: Vec, + /// Final column of each surviving var (parallel to `head_vars`). + head_cols: Vec, +} + +/// Build a column-reusing layout for the body atoms, or `None` if the reused +/// frontier still exceeds `W`. Liveness is first-use..last-use over the EMITTED +/// atom order, EXTENDED so any variable read by the head or a body prim stays +/// live to the end (it must survive into the captured row). Linear-scan slot +/// assignment: a column freed by a dead var is reused by a later var's birth. +fn build_projection( + atoms: &[PlanAtom], + head: &[GenericCoreAction], + body: &[GenericAtom], +) -> Option { + use hashbrown::HashSet; + let n = atoms.len(); + + // Variables the head / body prims read: these must survive to the end. + let mut survivor: HashSet = HashSet::new(); + collect_head_vars(head, &mut survivor); + for atom in body { + if matches!(atom.head, RuleBodyCall::Primitive { .. }) { + for term in &atom.args { + if let GenericAtomTerm::Var(_, variable) = term { + survivor.insert(variable.id); + } + } + } + } + + // first[v]/last[v] = first/last atom index where v appears. Survivors get + // last = n-1 (live to the final step) so they are never freed mid-chain. + let mut first: HashMap = HashMap::new(); + let mut last: HashMap = HashMap::new(); + for (i, a) in atoms.iter().enumerate() { + for s in &a.slots { + if let Slot::Var(v) = s { + first.entry(*v).or_insert(i); + last.insert(*v, i); + } + } + } + for v in &survivor { + if let Some(l) = last.get_mut(v) { + *l = n - 1; + } + } + + // Vars born / dying at each step (in deterministic var-id order so the layout + // is reproducible). + let mut births: Vec> = vec![Vec::new(); n]; + let mut deaths: Vec> = vec![Vec::new(); n]; + let mut vars: Vec = first.keys().copied().collect(); + vars.sort_unstable(); + for v in vars { + births[first[&v]].push(v); + deaths[last[&v]].push(v); + } + + // Linear-scan register allocation. `col_of` is the current column of each + // live var; `free` is a stack of reusable columns (lowest first, so the + // layout is deterministic); `next` is the high-water column allocator. + let mut col_of: HashMap = HashMap::new(); + let mut free: Vec = Vec::new(); + let mut next: usize = 0; + let mut step_col: Vec> = Vec::with_capacity(n); + for i in 0..n { + // Births first (so this atom's fresh vars get columns before we snapshot). + for &v in &births[i] { + let c = if let Some(c) = free.pop() { + c + } else { + let c = next; + next += 1; + c + }; + col_of.insert(v, c); + } + if col_of.len() > W { + return None; + } + // Snapshot the layout AT this step (after births, before deaths) — every + // var the atom touches plus all still-live carried vars are present. + step_col.push(col_of.clone()); + // Deaths free their columns for reuse by later births. + for &v in &deaths[i] { + if let Some(c) = col_of.remove(&v) { + free.push(c); + } + // freeing low-first keeps assignment deterministic + free.sort_unstable_by(|a, b| b.cmp(a)); + } + } + + // Surviving vars + their FINAL columns (the columns they hold at the last + // step, where they are guaranteed live). Deterministic order by var id. + let final_layout = &step_col[n - 1]; + let mut head_vars: Vec = survivor + .into_iter() + .filter(|v| final_layout.contains_key(v)) + .collect(); + head_vars.sort_unstable(); + let head_cols: Vec = head_vars.iter().map(|v| final_layout[v]).collect(); + + Some(ProjectionPlan { + step_col, + head_vars, + head_cols, + }) +} + +/// Collect every variable a head action references into `out`. +fn collect_head_vars( + head: &[GenericCoreAction], + out: &mut hashbrown::HashSet, +) { + let add = |term: &GenericAtomTerm, out: &mut hashbrown::HashSet| { + if let GenericAtomTerm::Var(_, variable) = term { + out.insert(variable.id); + } + }; + for action in head { + match action { + GenericCoreAction::Let(_, variable, _, arguments) => { + for argument in arguments { + add(argument, out); + } + out.insert(variable.id); + } + GenericCoreAction::LetAtomTerm(_, variable, term) => { + add(term, out); + out.insert(variable.id); + } + GenericCoreAction::Set(_, _, arguments, value) => { + for argument in arguments { + add(argument, out); + } + add(value, out); + } + GenericCoreAction::Change(_, _, _, arguments) => { + for argument in arguments { + add(argument, out); + } + } + GenericCoreAction::Union(_, lhs, rhs) => { + add(lhs, out); + add(rhs, out); + } + GenericCoreAction::Panic(..) => {} + } + } +} + +/// A fused, delta-fed body join for a WHOLE ruleset on a single shared timely +/// `Worker`. Built once via [`FusedDdJoin::build`]; driven across epochs via +/// [`FusedDdJoin::step`] with a SINGLE `worker.step_while` per call. +pub struct FusedDdJoin { + worker: Worker, + /// One shared input session per DISTINCT body relation read view across all + /// rules. + inputs: HashMap>, + /// Single probe on all rule outputs (they share the dataflow scope, so one + /// probe gates the whole epoch's fixpoint). + probe: ProbeHandle, + /// Distinct relation read views across the whole ruleset, in first-use + /// order. This is the authoritative host feed order for the fused worker. + reads: Vec, + /// The fused rules in caller-supplied build order. The sorted rule-index list + /// identifies the ruleset cache entry but does not reorder these outputs. + rules: Vec, + /// Current epoch (monotonic; advanced once per [`step`]). + epoch: u32, +} + +/// One rule's lowering inside a [`FusedDdJoin`]: its rule index (for routing +/// bindings to its head), its per-epoch output capture buffer, and the variable +/// order used to unpack captured rows. +struct FusedRule { + idx: usize, + /// This rule's per-epoch output binding-delta capture (`inspect_batch` + /// appends `(row, weight)`; drained by [`FusedDdJoin::step`]). + captured: CaptureBuf, + /// Variable ids packed into each captured row, in capture-column order. + var_order: Vec, +} + +impl FusedDdJoin { + /// Build ONE worker + ONE dataflow for the whole ruleset. `plans` pairs each + /// rule's index with its [`JoinPlan`], in the order they should fire. Every + /// rule — congruence, user, and canonicalization — runs through the same + /// general fused join. + pub fn build(plans: &[(usize, JoinPlan)]) -> Result { + if plans.is_empty() { + bail!("cannot build a fused DD join without any rule plans"); + } + let alloc = Allocator::Thread(Thread::default()); + let mut worker = Worker::new( + WorkerConfig::default(), + alloc, + Some(std::time::Instant::now()), + ); + + // Distinct body relation read views across all rules → one shared input + // each. + let mut reads: Vec = Vec::new(); + for (_, plan) in plans { + for a in &plan.atoms { + if !reads.contains(&a.read_key) { + reads.push(a.read_key); + } + } + } + + // Owned per-rule plan snapshots so the `move` dataflow closure is 'static. + struct RulePlan { + idx: usize, + atoms: Vec>, + atom_reads: Vec, + var_order: Vec, + projection: ProjectionPlan, + } + let rule_plans: Vec = plans + .iter() + .map(|(idx, plan)| { + let atom_reads: Vec = plan.atoms.iter().map(|a| a.read_key).collect(); + RulePlan { + idx: *idx, + atoms: plan.atoms.iter().map(|a| a.slots.clone()).collect(), + atom_reads, + var_order: plan.projection.head_vars.clone(), + projection: plan.projection.clone(), + } + }) + .collect(); + + let probe: ProbeHandle = ProbeHandle::new(); + let probe_in = probe.clone(); + // Per-rule capture buffers, allocated outside the closure so we can keep a + // clone here and route each rule's output to its head after `step`. + let captures: Vec = rule_plans + .iter() + .map(|_| Rc::new(RefCell::new(Vec::new()))) + .collect(); + let captures_in = captures.clone(); + let reads_in = reads.clone(); + // The per-rule metadata `FusedRule` needs (kept here; the closure consumes + // `rule_plans` for the dataflow build). + let rule_meta: Vec<(usize, Vec)> = rule_plans + .iter() + .map(|rp| (rp.idx, rp.var_order.clone())) + .collect(); + + // The per-epoch input delta is already set-semantic: it comes from the + // versioned current-vs-fed map diff in `interpret::fused_bindings`, so + // each row appears at most once with weight ±1. The input integral + // therefore stays 0/1 per row without `.distinct()`. Adding that operator + // would perform redundant consolidation over the large relation integrals. + let inputs = worker.dataflow::(move |scope| { + // ONE shared input + base collection per distinct relation, shared by + // every atom occurrence (in every rule) that reads it. + let mut inputs: HashMap> = HashMap::new(); + let mut rel_coll: HashMap = HashMap::new(); + for &read in &reads_in { + let mut session: InputSession = InputSession::new(); + let coll = session.to_collection(scope); + inputs.insert(read, session); + rel_coll.insert(read, coll); + } + // A collection-level join arranges both inputs at every call site. + // Share base-relation arrangements with the same key projection. + let mut arranged_right = HashMap::new(); + + for (rp, cap) in rule_plans.iter().zip(captures_in.iter()) { + // This rule's per-atom collection vector, from the SHARED relation + // collections (cloning a DD collection is just a handle copy). + let n_atoms = rp.atoms.len(); + let atom_slots = &rp.atoms; + let proj = &rp.projection; + let step_col = &proj.step_col; + let slots0 = atom_slots[0].clone(); + let sc0 = step_col[0].clone(); + let mut cur = rel_coll[&rp.atom_reads[0]] + .clone() + .flat_map(move |r: Row| bind_atom(&r, &slots0, &sc0)); + + for i in 1..n_atoms { + let slots = atom_slots[i].clone(); + let prev = &step_col[i - 1]; + let next = &step_col[i]; + let shared: Vec = atom_vars(&slots) + .into_iter() + .filter(|v| prev.contains_key(v)) + .collect(); + let shared_cols_left: Vec = shared.iter().map(|v| prev[v]).collect(); + let shared_atom_cols: Vec = shared + .iter() + .map(|v| { + slots + .iter() + .position(|s| matches!(s, Slot::Var(x) if x == v)) + .expect( + "DD join invariant: a shared variable must occur in the joined atom", + ) + }) + .collect(); + + let left_cols = shared_cols_left.clone(); + let left = cur.map(move |b: Row| (pack_key(&b, &left_cols), b)); + let arrangement_key = (rp.atom_reads[i], shared_atom_cols.clone()); + let right = arranged_right + .entry(arrangement_key) + .or_insert_with(|| { + let right_cols = shared_atom_cols.clone(); + rel_coll[&rp.atom_reads[i]] + .clone() + .map(move |r: Row| (pack_key(&r, &right_cols), r)) + .arrange_by_key() + }) + .clone(); + + let previous_layout = prev.clone(); + let next_layout = next.clone(); + cur = left.join_core(right, move |_key, binding, row| { + remap_merge_atom_into(binding, row, &slots, &previous_layout, &next_layout) + }); + } + + // Pack the variables needed by body primitives and head actions + // into the capture columns expected by `var_order()`. + let head_cols = proj.head_cols.clone(); + let cur = cur.map(move |binding: Row| pack_key(&binding, &head_cols)); + + let cap = Rc::clone(cap); + // `step` accumulates captured deltas by binding row before + // interpreting their sign, so consolidating this stream in DD + // would duplicate the same per-key aggregation. + cur.inner + .inspect_batch(move |_t, batch| { + let mut buf = cap.borrow_mut(); + for (row, _time, w) in batch.iter() { + buf.push((*row, *w)); + } + }) + .probe_with(&probe_in); + } + + inputs + }); + + let rules: Vec = rule_meta + .into_iter() + .zip(captures) + .map(|((idx, var_order), captured)| FusedRule { + idx, + captured, + var_order, + }) + .collect(); + + Ok(FusedDdJoin { + worker, + inputs, + probe, + reads, + rules, + epoch: 0, + }) + } + + /// The rule indices this fused worker serves (build order). + pub fn rule_indices(&self) -> Vec { + self.rules.iter().map(|r| r.idx).collect() + } + + /// Distinct body relation read views across the ruleset, in first-use order. + pub fn read_keys(&self) -> &[ReadKey] { + &self.reads + } + + /// Captured variable order for the fused rule at build position `pos`. + pub fn rule_var_order(&self, pos: usize) -> &[u32] { + &self.rules[pos].var_order + } + + /// Feed one epoch of signed relation deltas into the SHARED inputs, advance + /// the timestamp, run the SINGLE worker to this epoch's fixpoint, and return + /// per-rule binding deltas. The outer `Vec` is in [`Self::rule_indices`] + /// order; each inner `Vec` is `(binding_row_as_var_order_vec, weight)`. + /// + /// CRUCIAL: the InputSessions are NEVER cleared — only the delta is pushed, so + /// the DD arrangements persist and the join is genuinely incremental. + pub fn step(&mut self, deltas: &DeltaMap) -> Result { + let mut pushed = false; + for (read, rows) in deltas { + let inp = self.inputs.get_mut(read).ok_or_else(|| { + anyhow!("DD invariant: received deltas for unplanned read view {read:?}") + })?; + for (row, w) in rows { + inp.update(pack_row(row)?, *w); + pushed = true; + } + } + let next_epoch = self + .epoch + .checked_add(1) + .ok_or_else(|| anyhow!("DD worker epoch overflow"))?; + if !pushed { + self.epoch = next_epoch; + return Ok(vec![Vec::new(); self.rules.len()]); + } + + for rule in &self.rules { + rule.captured.borrow_mut().clear(); + } + self.drive_to(next_epoch); + self.epoch = next_epoch; + + let mut accs: Vec, isize>> = + (0..self.rules.len()).map(|_| HashMap::new()).collect(); + for (rule, acc) in self.rules.iter().zip(accs.iter_mut()) { + for (row, weight) in rule.captured.borrow_mut().drain(..) { + let key = (0..rule.var_order.len()).map(|i| row[i]).collect(); + *acc.entry(key).or_insert(0) += weight; + } + } + Ok(accs + .into_iter() + .map(|acc| acc.into_iter().filter(|(_, w)| *w != 0).collect()) + .collect()) + } + + /// Step the SINGLE worker to `epoch`. + fn drive_to(&mut self, epoch: u32) { + for inp in self.inputs.values_mut() { + inp.advance_to(epoch); + inp.flush(); + } + let probe = self.probe.clone(); + self.worker.step_while(|| probe.less_than(&epoch)); + } +} + +/// Pack a slice of column values into a fixed-width row (0-padded). +fn pack_row(vals: &[u32]) -> Result { + if vals.len() > W { + bail!( + "DD input row has {} columns, exceeding fixed row width {W}", + vals.len() + ); + } + let mut a = Row::default(); + for (i, v) in vals.iter().enumerate() { + a[i] = *v; + } + Ok(a) +} + +/// Build a join key from selected columns (packed into the low slots). +fn pack_key(r: &Row, cols: &[usize]) -> Row { + let mut a = Row::default(); + for (i, &c) in cols.iter().enumerate() { + a[i] = r[c]; + } + a +} + +/// Distinct variables appearing in an atom (column order). +fn atom_vars(slots: &[Slot]) -> Vec { + let mut out = Vec::new(); + for s in slots { + if let Slot::Var(v) = s { + if !out.contains(v) { + out.push(*v); + } + } + } + out +} + +/// Match the first atom's relation row against its slots, producing the initial +/// binding row under the first-step layout (or no row if a constraint +/// fails). Returns a `Vec` for `flat_map`. +fn bind_atom(r: &Row, slots: &[Slot], layout: &HashMap) -> Vec { + let mut out = Row::default(); + let mut local: HashMap = HashMap::new(); + for (i, s) in slots.iter().enumerate() { + let val = r[i]; + match s { + Slot::Const(c) => { + if *c != val { + return Vec::new(); + } + } + Slot::Var(v) => { + if let Some(&prev) = local.get(v) { + if prev != val { + return Vec::new(); + } + } else { + local.insert(*v, val); + out[layout[v]] = val; + } + } + } + } + vec![out] +} + +/// Merge one atom into a binding while changing column layouts. `prev` is the +/// left-row layout (step `i-1`), `cur` is the output layout (step `i`). Carried +/// vars (live in both layouts but not produced here) are copied `prev[v]→cur[v]`; +/// the atom's vars are validated (shared) or written (fresh) at `cur[v]`; every +/// other output column is left zeroed. This is what reuses freed columns and so +/// keeps the frontier within `W`. Returns no row on a constraint failure. +fn remap_merge_atom_into( + b: &Row, + r: &Row, + slots: &[Slot], + prev: &HashMap, + cur: &HashMap, +) -> Vec { + let mut out = Row::default(); + // Carry over every still-live var (present in `cur`) that the left row + // already holds (present in `prev`). Atom-fresh vars are absent from `prev`, + // so they are NOT copied here — they are written from `r` below. + for (&v, &pc) in prev { + if let Some(&cc) = cur.get(&v) { + out[cc] = b[pc]; + } + } + let mut local: HashMap = HashMap::new(); + for (i, s) in slots.iter().enumerate() { + let val = r[i]; + match s { + Slot::Const(c) => { + if *c != val { + return Vec::new(); + } + } + Slot::Var(v) => { + if let Some(&prior) = local.get(v) { + if prior != val { + return Vec::new(); + } + continue; + } + local.insert(*v, val); + // Every atom var is live at this step ⇒ present in `cur`. + let cc = cur[v]; + if prev.contains_key(v) { + // Shared (already bound): the carried value must agree. (The + // join key already enforces this for the key columns; this + // also covers shared vars not in the key, if any.) + if out[cc] != val { + return Vec::new(); + } + } else { + // Fresh var born at this atom: write it. + out[cc] = val; + } + } + } + } + vec![out] +} + +#[cfg(test)] +mod tests { + use super::*; + use egglog_ast::{ + core::{GenericAtom, GenericAtomTerm, GenericCoreAction, GenericCoreRule, Query}, + span::Span, + }; + use egglog_backend_trait::{ + ColumnTy, FunctionId, ReadMode, RuleActionCall, RuleBodyCall, RuleSpec, RuleVar, + }; + use egglog_numeric_id::NumericId; + + fn live(func: FunctionId) -> ReadKey { + ReadKey { + func, + mode: ReadMode::Live, + } + } + + fn variable(id: u32) -> GenericAtomTerm { + GenericAtomTerm::Var( + Span::Panic, + RuleVar { + id, + name: format!("v{id}").into_boxed_str(), + ty: ColumnTy::Id, + }, + ) + } + + fn body_atom( + func: FunctionId, + vars: impl IntoIterator, + ) -> GenericAtom { + GenericAtom { + span: Span::Panic, + head: RuleBodyCall::Table { + id: func, + read: ReadMode::Live, + }, + args: vars.into_iter().map(variable).collect(), + } + } + + fn test_rule( + name: &str, + body: Vec>, + output: FunctionId, + vars: impl IntoIterator, + ) -> RuleSpec { + let mut terms = vars.into_iter().map(variable).collect::>(); + let value = terms.pop().expect("test output includes a value"); + RuleSpec { + name: name.to_owned(), + seminaive: true, + no_decomp: false, + core: GenericCoreRule { + span: Span::Panic, + body: Query { atoms: body }, + head: egglog_ast::core::GenericCoreActions::new(vec![GenericCoreAction::Set( + Span::Panic, + RuleActionCall::Table { + id: output, + name: "output".into(), + }, + terms, + value, + )]), + }, + } + } + + /// Build a `JoinPlan` from `(func, vars)` atoms, preserving every variable + /// in the captured output through a synthetic head action. + fn plan_of(atoms: &[(FunctionId, &[u32])]) -> JoinPlan { + let body = atoms + .iter() + .map(|(func, vars)| body_atom(*func, vars.iter().copied())) + .collect(); + let mut vars: Vec = atoms + .iter() + .flat_map(|(_, vars)| vars.iter().copied()) + .collect(); + vars.sort_unstable(); + vars.dedup(); + plan_join(&test_rule( + "test join", + body, + FunctionId::new(u32::MAX), + vars, + )) + .unwrap() + } + + #[test] + fn pack_row_checks_the_fixed_width_boundary() { + let values = (0..W as u32).collect::>(); + let packed = pack_row(&values).expect("a width-W row must fit"); + assert_eq!(&packed.0[..], values.as_slice()); + + let error = pack_row(&[0; W + 1]).unwrap_err(); + assert!(error.to_string().contains(&format!("{} columns", W + 1))); + assert!(error.to_string().contains(&format!("fixed row width {W}"))); + } + + #[test] + fn plan_reuses_columns_across_wide_variable_chain() { + let func = FunctionId::new(0); + let final_var = W as u32 + 1; + let body = (0..=W as u32) + .map(|var| body_atom(func, [var, var + 1])) + .collect(); + let plan = plan_join(&test_rule("wide chain", body, func, [final_var])).unwrap(); + + assert_eq!(plan.var_order(), vec![final_var]); + assert!(plan + .projection + .step_col + .iter() + .all(|layout| layout.len() <= W)); + } + + #[test] + fn plan_rejects_live_variable_frontier_wider_than_row() { + let func = FunctionId::new(0); + let rule = test_rule( + "wide frontier", + vec![body_atom(func, 0..W as u32), body_atom(func, [0, W as u32])], + func, + 0..=W as u32, + ); + + let error = plan_join(&rule) + .err() + .expect("wide live frontier must fail"); + assert!(error.contains("live body-variable frontier exceeds W")); + } + + /// Feed one all-at-once delta of `rows` per func into a fresh + /// `FusedDdJoin`, returning the rule's sorted output bindings. + fn run_once( + atoms: &[(FunctionId, &[u32])], + rows: &HashMap, + ) -> Vec<(Vec, isize)> { + let plan = plan_of(atoms); + let mut j = FusedDdJoin::build(&[(0usize, plan)]).unwrap(); + let deltas: DeltaMap = rows + .iter() + .map(|(&f, rows)| (live(f), rows.clone())) + .collect(); + let mut out = j.step(&deltas).unwrap().remove(0); + out.sort(); + out + } + + #[test] + fn binary_join_handles_multi_atom_self_join() { + let f = FunctionId::new(0); + let g = FunctionId::new(1); + let atoms: &[(FunctionId, &[u32])] = &[(f, &[0, 1, 2]), (g, &[0, 3]), (g, &[1, 4])]; + let mut rows = HashMap::new(); + rows.insert( + f, + vec![ + (vec![1, 2, 100], 1), + (vec![1, 3, 101], 1), + (vec![4, 2, 102], 1), + ], + ); + rows.insert( + g, + vec![ + (vec![1, 10], 1), + (vec![2, 20], 1), + (vec![3, 21], 1), + (vec![4, 40], 1), + ], + ); + + assert_eq!( + run_once(atoms, &rows), + vec![ + (vec![1, 2, 100, 10, 20], 1), + (vec![1, 3, 101, 10, 21], 1), + (vec![4, 2, 102, 40, 20], 1), + ] + ); + } +} diff --git a/egglog-experimental/dd/src/interpret.rs b/egglog-experimental/dd/src/interpret.rs new file mode 100644 index 0000000..6a3f490 --- /dev/null +++ b/egglog-experimental/dd/src/interpret.rs @@ -0,0 +1,638 @@ +//! Host-side iteration driver for the Differential Dataflow backend. +//! +//! One `run_rules` call = **one bounded egglog iteration**. The body join runs +//! on the in-process, build-once, epoch-driven raw differential-dataflow +//! dataflow (`crate::dd_native`); this module owns the orchestration around it: +//! +//! 1. compute every rule's matches against the same pre-iteration relation +//! state, following egglog's "match, then apply" model for one bounded hop; +//! 2. drive the ruleset's persistent fused DD dataflow with the per-relation +//! signed delta vs. what that dataflow was last fed, then re-run body primitives +//! (`!=` guards, value-computing prims) host-side over the produced bindings; +//! 3. for every surviving binding, execute the head ops in order — `set` / +//! `remove` / `subsume` writes, RHS `lookup` (eq-sort constructor: create on +//! miss), RHS primitive `call`, `union`, `panic`; +//! 4. apply all collected writes/removes to the mirror, folding each merge +//! function's `set`s into a single value per key by its merge mode. +//! +//! ## The engine split +//! +//! Relational table-atom joins are the only operations on the DD engine, and +//! there is no host nested-loop fallback. Any rule the DD plan cannot lower (a +//! binding row exceeding the fixed width cap [`crate::dd_native::W`], or any +//! shape `plan_join` rejects) is reported as an error. Body primitives and head +//! actions are applied host-side here. +//! +//! Primitives are invoked through `Database::with_execution_state`, so they see +//! the same interned base `Value`s the frontend created and preserve bit-for-bit +//! value parity with the reference backend. + +use anyhow::{anyhow, Result}; +use egglog_ast::core::{GenericAtom, GenericAtomTerm, GenericCoreAction, GenericCoreActions}; +use egglog_ast::generic_ast::Change; +use egglog_backend_trait::{ + FunctionId, ReadMode, RuleActionCall, RuleBodyCall, RuleSpec, RuleValue, RuleVar, Value, +}; +use egglog_numeric_id::NumericId; +use hashbrown::{HashMap, HashSet}; + +use crate::compile::{MergeMode, ReadKey, Row}; +use crate::EGraph; + +/// Binding environment: variable id → bound `u32` value. +pub(crate) type Env = HashMap; + +type DdDeltaRows = HashMap, isize)>>; + +/// Retractions batched per function: the key length plus the set of keys to +/// remove, so one `retain` pass drops them all. +type RemovesByFunc = HashMap>)>; + +/// Map the fused worker's rule order back to positions in the current caller. +/// The worker cache key identifies a set of rule ids, so a cache hit does not +/// imply that the caller supplied those ids in the original build order. +fn fused_caller_positions( + atom_positions: &[usize], + atom_rule_idxs: &[usize], + fused_rule_idxs: &[usize], +) -> Result> { + if atom_positions.len() != atom_rule_idxs.len() || atom_rule_idxs.len() != fused_rule_idxs.len() + { + return Err(anyhow!( + "fused rule mapping length mismatch: caller positions={}, caller ids={}, fused ids={}", + atom_positions.len(), + atom_rule_idxs.len(), + fused_rule_idxs.len() + )); + } + + let mut caller_by_rule = HashMap::with_capacity(atom_rule_idxs.len()); + for (&rule_idx, &position) in atom_rule_idxs.iter().zip(atom_positions) { + if caller_by_rule.insert(rule_idx, position).is_some() { + return Err(anyhow!( + "duplicate caller rule index {rule_idx} in fused ruleset" + )); + } + } + + let mut seen = HashSet::with_capacity(fused_rule_idxs.len()); + let mut positions = Vec::with_capacity(fused_rule_idxs.len()); + for &rule_idx in fused_rule_idxs { + if !seen.insert(rule_idx) { + return Err(anyhow!( + "duplicate cached rule index {rule_idx} in fused ruleset" + )); + } + let position = caller_by_rule.get(&rule_idx).copied().ok_or_else(|| { + anyhow!("cached fused rule index {rule_idx} is absent from the caller ruleset") + })?; + positions.push(position); + } + + Ok(positions) +} + +/// A pending write to apply after all matches are computed. +enum Write { + /// Insert/overwrite a full row. + Set(FunctionId, Row), + /// Retract by key (the slots address inputs for a function, whole row for a + /// relation). + Remove(FunctionId, Vec), + /// Soft-delete: move the row(s) whose leading columns match this prefix into + /// the `subsumed` side-set (hidden from ordinary matching, still present). + Subsume(FunctionId, Vec), +} + +/// One bounded egglog iteration with the body join running on the in-process, +/// build-once, epoch-driven raw differential-dataflow dataflow +/// (`crate::dd_native`). +/// +/// Compute the signed `+/-` delta of each distinct body-relation view vs. the +/// rows last fed to the ruleset's persistent fused DD dataflow. Each nonempty +/// removal or insertion phase steps the shared worker, feeding only deltas into +/// never-cleared `InputSession`s. Positive per-rule binding deltas become envs; +/// body primitives and head actions then run host-side. Writes and FD merges are +/// applied afterward so results are bit-exact. +pub fn run_iteration(eg: &mut EGraph, rules: &[(usize, RuleSpec)]) -> Result { + // Every rule — including the term encoder's canonicalization rules — lowers + // as an ordinary rule and joins on the fused DD worker; there is no special + // casing for any rule kind. + + // Snapshot the fresh-id counter: any hash-cons (`lookup_or_create`) this + // call advances it, the O(1) signal that a new term row was created. + let next_id_at_start = eg.next_id; + + let mut writes: Vec = Vec::new(); + // Iteration-scoped `key -> output` index for `lookup_or_create` (eq-sort + // constructor hash-cons). Built lazily per function so repeated lookups in + // one iteration are O(1) instead of rescanning the growing mirror each time. + let mut lookup_index: HashMap, u32>> = HashMap::new(); + + // Compute every rule's binding envs FIRST (so the whole atom-bearing ruleset + // runs on one fused DD worker via `fused_bindings`), THEN + // apply head actions in the original rule firing order. Atom-less rules + // (`(rule () …)`) have no input relation to drive the DD dataflow, so they + // stay host-side (fire once); they are computed inline below. + let envs_by_rule = fused_bindings(eg, rules)?; + + for ((_, rule), envs) in rules.iter().zip(envs_by_rule.into_iter()) { + for mut env in envs { + apply_head( + eg, + &rule.core.head, + &mut env, + &mut writes, + &mut lookup_index, + )?; + } + } + + // Apply collected writes to the mirror. + // + // Removes are BATCHED per function: applying each `Write::Remove` with its + // own `set.retain` scan is O(|removes| · |state|) — quadratic. We collect all + // retraction keys per function into a hash set, then do a SINGLE `retain` + // pass per touched function: O(|state|) total. Removes are applied FIRST + // (batched), then Sets — preserving the term encoder's `(@uf)` "delete old + // leader, set new leader" delete-then-set ordering. + // + // `changed` is computed INCREMENTALLY as writes land (O(delta)), not via a + // full before/after content compare. A hash-cons in `lookup_or_create` + // always allocates a fresh id, so any term created this call advances + // `next_id` — that alone is a real mirror change. + let mut changed = eg.next_id != next_id_at_start; + let mut removes_by_func: RemovesByFunc = HashMap::new(); + let mut sets: Vec<(FunctionId, Row)> = Vec::new(); + let mut subsumes: Vec<(FunctionId, Vec)> = Vec::new(); + for w in writes { + match w { + Write::Set(f, row) => sets.push((f, row)), + Write::Remove(f, key) => { + let entry = removes_by_func + .entry(f) + .or_insert_with(|| (key.len(), HashSet::new())); + entry.1.insert(key.into_boxed_slice()); + } + Write::Subsume(f, prefix) => subsumes.push((f, prefix)), + } + } + for (f, (keylen, keys)) in removes_by_func { + // A `delete`/rebuild retraction clears the row from BOTH the live mirror + // and the subsumed side-set (a rebuilt-away subsumed row must not linger). + changed |= eg.remove_matching_keys(f, keylen, &keys); + } + // Apply sets. A plain relation (whole-row key) just inserts. A merge function + // (Old/New/Min) folds each new value against the CURRENT value for its key by + // the merge mode; the new values are folded in EMISSION order so `New`/`Old` + // are insertion-order-correct — the pre-set mirror value is the "old" one and + // the value set this call is the "new" one. (The mirror is an unordered + // `HashSet`, so a fold that sorted the candidate rows would pick the wrong + // winner for `New`/`Old`: e.g. a `New` merge of `@UF_Mf` where a leader + // changes from `3` to `1` must keep `1`, not the sort-larger old value `3`.) + let mut merge_by_func: HashMap> = HashMap::new(); + for (f, row) in sets { + let merge = eg.merge_mode(f); + if merge == MergeMode::Relation { + changed |= eg.insert_live_row(f, row); + } else { + merge_by_func.entry(f).or_default().push(row); + } + } + for (f, rows) in merge_by_func { + changed |= eg.apply_merge_sets(f, &rows)?; + } + // Subsumes last: a row `set` this iteration can then be subsumed, and the + // move reads the just-updated live mirror. + for (f, prefix) in subsumes { + changed |= eg.subsume_rows(f, &prefix); + } + + Ok(changed) +} + +/// Compute every rule's binding envs in ONE fused pass: the whole atom-bearing +/// ruleset's body joins run on a SINGLE shared timely worker +/// ([`dd_native::FusedDdJoin`]) clocked once this iteration, then each rule's +/// host-side body primitives are re-run over its own bindings. Atom-less rules +/// (`(rule () …)`) have no input relation to drive the DD dataflow, so they are +/// fired once host-side. Returns a `Vec>` parallel to `rules` (same +/// order), ready for `apply_head`. +fn fused_bindings(eg: &mut EGraph, rules: &[(usize, RuleSpec)]) -> Result>> { + use crate::dd_native; + + let mut out: Vec> = vec![Vec::new(); rules.len()]; + + // Partition: atom-bearing rules drive the fused DD worker; atom-less rules + // fire once host-side. Record each atom-bearing rule's POSITION in `rules` so + // we can scatter the fused output back into `out` in the caller's order. + let mut atom_positions: Vec = Vec::new(); + let mut atom_rule_idxs: Vec = Vec::new(); + for (pos, (idx, rule)) in rules.iter().enumerate() { + let has_atoms = rule + .core + .body + .atoms + .iter() + .any(|atom| matches!(atom.head, RuleBodyCall::Table { .. })); + if has_atoms { + atom_positions.push(pos); + atom_rule_idxs.push(*idx); + } else { + // Atom-less rule: fire once (presence in `seen` = already fired). + if eg.seen.contains_key(idx) { + continue; + } + eg.seen.insert(*idx, ()); + let mut envs: Vec = vec![Env::new()]; + for atom in &rule.core.body.atoms { + if matches!(atom.head, RuleBodyCall::Primitive { .. }) { + envs = step_prim(eg, atom, envs)?; + if envs.is_empty() { + break; + } + } + } + out[pos] = envs; + } + } + + if atom_positions.is_empty() { + return Ok(out); + } + + // The fused join is keyed by the SORTED atom-bearing rule-index list (the + // ruleset identity). Build it ONCE (lazily) per distinct ruleset, planning + // each rule. Any shape `plan_join` rejects is returned as an error (there is + // no host fallback; the DD dataflow is the only join path). + let mut key: Vec = atom_rule_idxs.clone(); + key.sort_unstable(); + + if !eg.dd_fused.contains_key(&key) { + // Plan in the SAME order as `atom_positions` so the fused build order + // matches our scatter order (the fused join preserves plan order). + let mut plans: Vec<(usize, dd_native::JoinPlan)> = Vec::with_capacity(atom_positions.len()); + for (&pos, &idx) in atom_positions.iter().zip(atom_rule_idxs.iter()) { + let rule = &rules[pos].1; + let plan = dd_native::plan_join(rule).map_err(|reason| { + anyhow!( + "Differential Dataflow join cannot lower rule {:?}: {reason} \ + (no host fallback; the DD dataflow is the only join path)", + rule.name + ) + })?; + plans.push((idx, plan)); + } + let fused = dd_native::FusedDdJoin::build(&plans)?; + eg.dd_fused.insert(key.clone(), fused); + } + + // Capture the fused worker's build order and map each output slot back to + // the current caller by stable rule id. The sorted cache key only proves set + // equality; it does not prove that this call uses the original order. + let (fused_rule_idxs, var_orders, all_reads): (Vec, Vec>, Vec) = { + let fused = eg + .dd_fused + .get(&key) + .expect("DD cache invariant: fused join was inserted for this ruleset key"); + let rule_idxs = fused.rule_indices(); + let var_orders = (0..rule_idxs.len()) + .map(|p| fused.rule_var_order(p).to_vec()) + .collect(); + (rule_idxs, var_orders, fused.read_keys().to_vec()) + }; + let fused_positions = + fused_caller_positions(&atom_positions, &atom_rule_idxs, &fused_rule_idxs)?; + + // Distinct relation read views across the whole ruleset. We diff versioned + // current snapshots rather than plain row sets: if another ruleset removes and + // reinserts the same row between two invocations, the row is present in both + // end states but has a fresh version. Feeding `-row` then `+row` as separate + // DD steps preserves the reference backend's hidden-timestamp behavior + // without replaying transient rows that are no longer visible. + let mut removals_batch: DdDeltaRows = HashMap::new(); + let mut insertions_batch: DdDeltaRows = HashMap::new(); + { + let fed = eg.dd_fused_fed_versions.entry(key.clone()).or_default(); + for &read in &all_reads { + let cur = match read.mode { + ReadMode::Live => eg.live_versions.get(&read.func), + ReadMode::Subsumed => eg.subsumed_versions.get(&read.func), + ReadMode::All => eg.all_versions.get(&read.func), + }; + let cur_empty: HashMap = HashMap::new(); + let cur = cur.unwrap_or(&cur_empty); + let prev = fed.entry(read).or_default(); + if prev == cur { + continue; + } + + let mut removals = Vec::new(); + let mut insertions = Vec::new(); + for row in prev.keys() { + if !cur.contains_key(row) { + removals.push((row.to_vec(), -1)); + } + } + for (row, version) in cur { + match prev.get(row) { + None => insertions.push((row.to_vec(), 1)), + Some(prev_version) if prev_version != version => { + removals.push((row.to_vec(), -1)); + insertions.push((row.to_vec(), 1)); + } + Some(_) => {} + } + } + if !removals.is_empty() { + removals_batch.insert(read, removals); + } + if !insertions.is_empty() { + insertions_batch.insert(read, insertions); + } + *prev = cur.clone(); + } + } + let mut delta_batches: Vec = Vec::new(); + if !removals_batch.is_empty() { + delta_batches.push(removals_batch); + } + if !insertions_batch.is_empty() { + delta_batches.push(insertions_batch); + } + + // Step the shared worker once per nonempty signed-delta phase. A version + // change may require a removal phase followed by an insertion phase. + let mut per_rule_bindings = vec![Vec::new(); fused_positions.len()]; + { + let fused = eg + .dd_fused + .get_mut(&key) + .expect("DD cache invariant: fused join was inserted for this ruleset key"); + for delta in &delta_batches { + let stepped = fused.step(delta)?; + for (acc, rows) in per_rule_bindings.iter_mut().zip(stepped) { + acc.extend(rows); + } + } + } + + // Turn each rule's positive binding deltas into envs; re-run its body prims + // host-side. Negative weights are integral bookkeeping (a body row retracted) + // — egglog heads are monotone-fire, so we do NOT re-fire on disappearance. + for (fpos, bindings) in per_rule_bindings.into_iter().enumerate() { + let caller_pos = fused_positions[fpos]; + let rule = &rules[caller_pos].1; + let var_order = &var_orders[fpos]; + let mut envs: Vec = Vec::new(); + for (bind, w) in &bindings { + if *w <= 0 { + continue; + } + let mut env: Env = Env::new(); + for (i, &v) in var_order.iter().enumerate() { + env.insert(v, bind[i]); + } + let mut es: Vec = vec![env]; + for atom in &rule.core.body.atoms { + if matches!(atom.head, RuleBodyCall::Primitive { .. }) { + es = step_prim(eg, atom, es)?; + } + } + envs.extend(es); + } + out[caller_pos] = envs; + } + + Ok(out) +} + +/// Evaluate a primitive body op over each binding env, returning the new list of +/// envs. A value-computing prim binds (or checks) its return var; a guard prim +/// (`!=`) that fails prunes the env. Table atoms are NOT handled here — they run +/// on the DD dataflow; this is only the host-side primitive phase. +pub(crate) fn step_prim( + eg: &mut EGraph, + atom: &GenericAtom, + envs: Vec, +) -> Result> { + let RuleBodyCall::Primitive { id, .. } = atom.head else { + unreachable!("step_prim called on a non-primitive body op"); + }; + let (ret, args) = atom + .args + .split_last() + .ok_or_else(|| anyhow!("body primitive has no return term"))?; + let mut out = Vec::new(); + for env in envs { + let resolved: Option> = args + .iter() + .map(|term| term_value(term, &env).map(Value::new)) + .collect(); + let Some(argv) = resolved else { continue }; + let result = eg.eval_prim_internal(id, &argv)?; + let Some(result) = result else { + // Primitive failed (e.g. `!=` of equal args) — prune. + continue; + }; + match ret { + GenericAtomTerm::Var(_, variable) => { + let mut next = env.clone(); + match next.get(&variable.id) { + Some(&existing) if existing != result.rep() => continue, + _ => { + next.insert(variable.id, result.rep()); + } + } + out.push(next); + } + GenericAtomTerm::Literal(_, constant) => { + if constant.value.rep() == result.rep() { + out.push(env); + } + } + GenericAtomTerm::Global(..) => { + unreachable!("DD rule validation must reject global primitive return terms") + } + } + } + Ok(out) +} + +/// Execute the head ops for one binding, accumulating writes. +fn apply_head( + eg: &mut EGraph, + head: &GenericCoreActions, + env: &mut Env, + writes: &mut Vec, + lookup_index: &mut HashMap, u32>>, +) -> Result<()> { + for action in &head.0 { + match action { + GenericCoreAction::Let(_, variable, call, arguments) => { + let arguments = resolve_terms(arguments, env)?; + let result = match call { + RuleActionCall::Table { id, .. } => { + let key = arguments + .iter() + .copied() + .map(Value::new) + .collect::>(); + if eg.info(*id).lookup_mints { + Some(lookup_or_create(eg, *id, &key, lookup_index)) + } else { + Some(lookup_existing(eg, *id, &key, lookup_index).ok_or_else(|| { + anyhow!( + "lookup on `{}` failed in rule action", + eg.relation_name(*id) + ) + })?) + } + } + RuleActionCall::Primitive { id, .. } => { + let arguments = arguments + .iter() + .copied() + .map(Value::new) + .collect::>(); + eg.eval_prim_internal(*id, &arguments)? + } + }; + if let Some(result) = result { + env.insert(variable.id, result.rep()); + } + } + GenericCoreAction::LetAtomTerm(_, variable, term) => { + env.insert(variable.id, resolve_term(term, env)?); + } + GenericCoreAction::Set(_, call, arguments, value) => { + let RuleActionCall::Table { id, .. } = call else { + return Err(anyhow!("DD backend cannot set a primitive")); + }; + let mut row = resolve_terms(arguments, env)?; + row.push(resolve_term(value, env)?); + writes.push(Write::Set(*id, row.into_boxed_slice())); + } + GenericCoreAction::Change(_, change, call, arguments) => { + let RuleActionCall::Table { id, .. } = call else { + return Err(anyhow!("DD backend cannot delete or subsume a primitive")); + }; + let values = resolve_terms(arguments, env)?; + match change { + Change::Delete => writes.push(Write::Remove(*id, values)), + Change::Subsume => writes.push(Write::Subsume(*id, values)), + } + } + GenericCoreAction::Union(..) => { + // Term encoding lowers unions to `(set (@uf ...))` writes. Do not + // silently discard a native union if this backend is driven + // outside that required frontend mode. + return Err(anyhow!( + "DD backend received a native union; term encoding must lower unions to @uf writes" + )); + } + GenericCoreAction::Panic(_, msg) => { + return Err(anyhow!("{msg}")); + } + } + } + Ok(()) +} + +fn term_value(term: &GenericAtomTerm, env: &Env) -> Option { + match term { + GenericAtomTerm::Var(_, variable) => env.get(&variable.id).copied(), + GenericAtomTerm::Literal(_, value) => Some(value.value.rep()), + GenericAtomTerm::Global(..) => { + unreachable!("DD rule validation must reject residual global terms") + } + } +} + +fn resolve_term(term: &GenericAtomTerm, env: &Env) -> Result { + term_value(term, env).ok_or_else(|| anyhow!("unbound term {term:?} in rule head")) +} + +fn resolve_terms(terms: &[GenericAtomTerm], env: &Env) -> Result> { + terms.iter().map(|term| resolve_term(term, env)).collect() +} + +/// Look up the output of `func` for input `key`. If absent, create the row with +/// a fresh id (eq-sort constructor semantics — mirrors `add_term`). The created +/// row is written directly into the mirror so subsequent lookups in the same +/// iteration see it (hash-cons). +pub(crate) fn lookup_or_create( + eg: &mut EGraph, + func: FunctionId, + key: &[Value], + index: &mut HashMap, u32>>, +) -> Value { + if let Some(value) = lookup_existing(eg, func, key, index) { + return value; + } + let k: Box<[u32]> = key.iter().map(|v| v.rep()).collect(); + let id = eg.fresh_id_internal(); + index.entry(func).or_default().insert(k, id); + let mut full: Vec = key.iter().map(|v| v.rep()).collect(); + full.push(id); + let row: Row = full.into_boxed_slice(); + eg.insert_live_row(func, row); + Value::new(id) +} + +/// Look up the current output of `func` for input `key` without creating a row. +pub(crate) fn lookup_existing( + eg: &EGraph, + func: FunctionId, + key: &[Value], + index: &mut HashMap, u32>>, +) -> Option { + let info = eg.info(func); + let inputs_len = info.arity.saturating_sub(1); + // Lazily build the key->output index for this function from live ∪ subsumed + // rows so repeated lookups within one iteration are O(1) instead of O(state) + // scans. A subsumed constructor row is still the current table row in the + // reference backend; looking it up must not mint a fresh visible row. + let idx = index.entry(func).or_insert_with(|| { + let live_len = eg.mirror.get(&func).map_or(0, |rows| rows.len()); + let subsumed_len = eg.subsumed.get(&func).map_or(0, |rows| rows.len()); + let mut m: HashMap, u32> = HashMap::with_capacity(live_len + subsumed_len); + if let Some(set) = eg.mirror.get(&func) { + for row in set.iter() { + let k: Box<[u32]> = (0..inputs_len).map(|i| row[i]).collect(); + m.insert(k, row[inputs_len]); + } + } + if let Some(set) = eg.subsumed.get(&func) { + for row in set.iter() { + let k: Box<[u32]> = (0..inputs_len).map(|i| row[i]).collect(); + m.entry(k).or_insert(row[inputs_len]); + } + } + m + }); + let k: Box<[u32]> = key.iter().map(|v| v.rep()).collect(); + idx.get(&k).copied().map(Value::new) +} + +#[cfg(test)] +mod tests { + use super::fused_caller_positions; + + #[test] + fn fused_rule_mapping_uses_rule_ids_instead_of_caller_order() { + assert_eq!( + fused_caller_positions(&[4, 9], &[10, 20], &[20, 10]).unwrap(), + [9, 4] + ); + } + + #[test] + fn fused_rule_mapping_rejects_inconsistent_cache_entries() { + let missing = fused_caller_positions(&[4, 9], &[10, 20], &[10, 30]).unwrap_err(); + assert!(missing.to_string().contains("absent from the caller")); + + let duplicate = fused_caller_positions(&[4, 9], &[10, 20], &[10, 10]).unwrap_err(); + assert!(duplicate.to_string().contains("duplicate cached rule")); + } +} diff --git a/egglog-experimental/dd/src/lib.rs b/egglog-experimental/dd/src/lib.rs new file mode 100644 index 0000000..b48f06d --- /dev/null +++ b/egglog-experimental/dd/src/lib.rs @@ -0,0 +1,1965 @@ +//! # egglog-experimental-dd +//! +//! A differential-dataflow-backed implementation of egglog's +//! [`egglog_backend_trait::Backend`] interface. +//! +//! One `run_rules` call is one bounded egglog iteration. Each rule's body +//! table-atom join runs on an in-process differential-dataflow dataflow +//! (`dd_native`); body primitives and head actions are applied +//! host-side (`interpret`) into a Rust-side materialized mirror of every +//! relation. `for_each` / `lookup_id` / `table_size` read that mirror. +//! +//! ## Why this is a DD backend +//! +//! This crate constructs Timely and Differential Dataflow operators directly; +//! it does not compile through a higher-level dataflow language or runtime. That +//! keeps this prototype focused on the backend SPI and the costs of persistent +//! incremental joins. Evaluating a higher-level compiler/runtime, including its +//! tuple generation, planning, and stratification choices, would be a separate +//! backend experiment rather than an interchangeable implementation detail here. + +use std::{ + any::{Any, TypeId}, + sync::{Arc, Mutex}, +}; + +use anyhow::{anyhow, bail, Result}; +use egglog_ast::core::{GenericAtomTerm, GenericCoreAction}; +use egglog_backend_trait::{ + Backend, BaseValues, ColumnTy, ContainerMergeFn, ContainerValues, CounterId, DefaultVal, + ExecutionState, ExternalFunction, ExternalFunctionId, FunctionConfig, FunctionId, + IterationReport, MergeFn, ReportLevel, RuleActionCall, RuleBodyCall, RuleId, RuleSetRun, + RuleSpec, RuleValue, RuleVar, ScanEntry, Value, +}; +use egglog_core_relations::Database; +use egglog_numeric_id::NumericId; +use hashbrown::{HashMap, HashSet}; + +mod compile; +mod dd_native; +mod interpret; + +use compile::{MergeMode, MergeTree, ReadKey, Row}; + +type LocatedValue = (u32, RowLocation); +type RowReplacement = (Box<[u32]>, LocatedValue); + +mod dd_workers { + use hashbrown::HashMap; + + use crate::dd_native; + + /// Owns the single-threaded Timely workers behind an exclusive-access + /// boundary. The worker map is private to this module, and every accessor + /// requires `&mut self`, so immutable backend operations cannot touch + /// Timely's `Rc`/`RefCell` state. + #[derive(Default)] + pub(super) struct DdWorkers { + inner: HashMap, dd_native::FusedDdJoin>, + } + + impl DdWorkers { + pub(super) fn contains_key(&mut self, key: &[usize]) -> bool { + self.inner.contains_key(key) + } + + pub(super) fn insert(&mut self, key: Vec, value: dd_native::FusedDdJoin) { + self.inner.insert(key, value); + } + + pub(super) fn get(&mut self, key: &[usize]) -> Option<&dd_native::FusedDdJoin> { + self.inner.get(key) + } + + pub(super) fn get_mut(&mut self, key: &[usize]) -> Option<&mut dd_native::FusedDdJoin> { + self.inner.get_mut(key) + } + + pub(super) fn retain(&mut self, mut keep: impl FnMut(&Vec) -> bool) { + self.inner.retain(|key, _| keep(key)); + } + + #[cfg(test)] + pub(super) fn is_empty(&mut self) -> bool { + self.inner.is_empty() + } + } + + // SAFETY: `FusedDdJoin` does not expose or clone its Timely handles outside + // this private owner. Moving `DdWorkers` moves every related handle together, + // and the only worker accessors require `&mut self`. In particular, `Sync` + // cannot be used to reach Timely state through a shared reference. + unsafe impl Send for DdWorkers {} + unsafe impl Sync for DdWorkers {} +} + +use dd_workers::DdWorkers; + +// --------------------------------------------------------------------------- +// Relation metadata +// --------------------------------------------------------------------------- + +/// What we remember about each registered relation/function. +#[derive(Clone)] +pub(crate) struct RelationInfo { + name: String, + /// Number of columns (including the output column for functions). + pub(crate) arity: usize, + /// Whether RHS lookups may create a fresh output row on miss. + pub(crate) lookup_mints: bool, + /// How functional-dependency conflicts are resolved. + pub(crate) merge: MergeMode, + /// The evaluatable merge tree, when `merge` is [`MergeMode::Computed`]. + pub(crate) merge_tree: Option, +} + +// --------------------------------------------------------------------------- +// EGraph +// --------------------------------------------------------------------------- + +/// The experimental differential-dataflow-backed egraph. +pub struct EGraph { + relations: Vec, + /// Rule slots; `None` = freed. + pub(crate) rules: Vec>, + /// Rust-side materialized mirror: the accumulated contents of each relation. + /// This is what `for_each` / `lookup_id` / `table_size` read. + /// + pub(crate) mirror: HashMap>, + /// Subsumed rows, moved OUT of `mirror` by `(subsume …)` (a "soft delete"). + /// Ordinary [`egglog_backend_trait::ReadMode::Live`] queries read only + /// `mirror`, so subsumed rows are excluded for free; include-subsumed rules + /// and `for_each` read `mirror` union `subsumed`. Keyed like `mirror`, by + /// `FunctionId`. + pub(crate) subsumed: HashMap>, + /// A core-relations [`Database`] used purely as the base-value / primitive + /// engine, so `Value`s are bit-for-bit identical to the reference backend. + db: Database, + /// Monotonic fresh-id counter for `fresh_id` / `add_term`. + pub(crate) next_id: u32, + /// Atom-less rules (`(rule () …)`) fire ONCE; an entry here marks a rule + /// index as already fired. The DD dataflow has no input relation to drive an + /// atom-less body, so this fired-marker is the one piece of seminaive + /// bookkeeping the DD path reuses (see `interpret::fused_bindings`). + /// `free_rule` removes the entry so a re-installed rule can fire again. + pub(crate) seen: HashMap, + /// Per-RULESET fused DD join (one shared timely worker + one dataflow for the + /// whole ruleset), keyed by the sorted live rule-index list. This is the + /// join path the interpreter drives. + pub(crate) dd_fused: DdWorkers, + /// Monotone version assigned whenever a row becomes visible in one of the + /// DD read views. This stands in for the reference backend's hidden + /// timestamp: removing and reinserting the same row gives it a fresh version, + /// so seminaive rules can fire on it again. + pub(crate) next_row_version: u64, + pub(crate) live_versions: HashMap>, + pub(crate) all_versions: HashMap>, + pub(crate) subsumed_versions: HashMap>, + /// Per-ruleset, per-function version snapshot last fed to the fused DD join. + pub(crate) dd_fused_fed_versions: HashMap, HashMap>>, + /// Deferred error channel used by panic primitives. The embedded database's + /// cloned external functions share this channel, matching the reference + /// bridge's panic-function behavior. + panic_message: Arc>>, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum RowLocation { + Live, + Subsumed, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +struct CurrentValue { + value: u32, + location: RowLocation, + rows_for_key: usize, +} + +impl Default for EGraph { + fn default() -> Self { + Self::new() + } +} + +impl EGraph { + /// Construct a fresh Differential Dataflow backend. Rule bodies run on the + /// in-process DD join; body primitives and head actions are applied + /// host-side into the mirror. Pass this backend to + /// `egglog::EGraph::with_backend` through the frontend integration crate. + pub fn new() -> Self { + let mut db = Database::new(); + // Pre-register the `()` (Unit) base type so `add_table`'s relation-vs- + // function detection can always resolve the Unit `BaseValueId`. + // `register_type` is idempotent, so a later frontend registration is a + // no-op that returns the same id. + db.base_values_mut().register_type::<()>(); + EGraph { + relations: Vec::new(), + rules: Vec::new(), + mirror: HashMap::new(), + subsumed: HashMap::new(), + db, + // Start at 1 so id 0 stays a "null"/padding sentinel. + next_id: 1, + seen: HashMap::new(), + dd_fused: DdWorkers::default(), + next_row_version: 1, + live_versions: HashMap::new(), + all_versions: HashMap::new(), + subsumed_versions: HashMap::new(), + dd_fused_fed_versions: HashMap::new(), + panic_message: Default::default(), + } + } + + pub(crate) fn info(&self, f: FunctionId) -> &RelationInfo { + self.relations + .get(f.rep() as usize) + .unwrap_or_else(|| panic!("FunctionId({}) not registered", f.rep())) + } + + /// Relation name for `f`, used in diagnostics and unsupported-shape errors. + pub(crate) fn relation_name(&self, f: FunctionId) -> &str { + &self.info(f).name + } + + /// Validate the backend-facing rule form once, before assigning a rule id. + /// The DD backend retains the accepted [`RuleSpec`] directly; this pass does + /// not introduce a second rule IR. + fn validate_rule(&self, rule: &RuleSpec) -> Result<()> { + let validate_terms = + |terms: &[GenericAtomTerm], location: &str| -> Result<()> { + if let Some(GenericAtomTerm::Global(_, global)) = terms + .iter() + .find(|term| matches!(term, GenericAtomTerm::Global(..))) + { + bail!( + "DD backend cannot add rule {:?}: residual global {:?} in {location}", + rule.name, + global.name + ); + } + Ok(()) + }; + let relation = |id: FunctionId, location: &str| -> Result<&RelationInfo> { + self.relations.get(id.rep() as usize).ok_or_else(|| { + anyhow!( + "DD backend cannot add rule {:?}: unregistered FunctionId({}) in {location}", + rule.name, + id.rep() + ) + }) + }; + + for atom in &rule.core.body.atoms { + match atom.head { + RuleBodyCall::Table { id, .. } => { + if atom.args.len() > dd_native::W { + bail!( + "DD backend cannot add rule {:?}: body table atom has {} columns, exceeding fixed row width {}", + rule.name, + atom.args.len(), + dd_native::W + ); + } + let info = relation(id, "rule body")?; + if atom.args.len() != info.arity { + bail!( + "DD backend cannot add rule {:?}: body table atom for {:?} has {} columns, but relation `{}` has arity {}", + rule.name, + id, + atom.args.len(), + info.name, + info.arity + ); + } + validate_terms(&atom.args, "rule body")?; + } + RuleBodyCall::Primitive { .. } => { + if atom.args.is_empty() { + bail!( + "DD backend cannot add rule {:?}: body primitive has no return term", + rule.name + ); + } + validate_terms(&atom.args, "rule body primitive")?; + } + } + } + + for action in &rule.core.head.0 { + match action { + GenericCoreAction::Let(_, _, call, arguments) => { + validate_terms(arguments, "rule action")?; + if let RuleActionCall::Table { id, .. } = call { + let info = relation(*id, "table lookup action")?; + let Some(expected) = info.arity.checked_sub(1) else { + bail!( + "DD backend cannot add rule {:?}: cannot look up nullary relation `{}` as a function", + rule.name, + info.name + ); + }; + if arguments.len() != expected { + bail!( + "DD backend cannot add rule {:?}: lookup on `{}` has {} arguments, expected {}", + rule.name, + info.name, + arguments.len(), + expected + ); + } + } + } + GenericCoreAction::LetAtomTerm(_, _, term) => { + validate_terms(std::slice::from_ref(term), "rule let action")?; + } + GenericCoreAction::Set(_, call, arguments, value) => { + let RuleActionCall::Table { id, .. } = call else { + bail!( + "DD backend cannot add rule {:?}: cannot set a primitive", + rule.name + ); + }; + validate_terms(arguments, "rule set action")?; + validate_terms(std::slice::from_ref(value), "rule set action")?; + let info = relation(*id, "set action")?; + if arguments.len() + 1 != info.arity { + bail!( + "DD backend cannot add rule {:?}: set on `{}` writes {} columns, expected {}", + rule.name, + info.name, + arguments.len() + 1, + info.arity + ); + } + } + GenericCoreAction::Change(_, _, call, arguments) => { + let RuleActionCall::Table { id, .. } = call else { + bail!( + "DD backend cannot add rule {:?}: cannot delete or subsume a primitive", + rule.name + ); + }; + validate_terms(arguments, "rule change action")?; + let info = relation(*id, "change action")?; + if arguments.len() > info.arity { + bail!( + "DD backend cannot add rule {:?}: change on `{}` addresses {} columns, exceeding arity {}", + rule.name, + info.name, + arguments.len(), + info.arity + ); + } + } + GenericCoreAction::Union(..) => { + bail!( + "DD backend cannot add rule {:?}: received a native union; term encoding must lower unions to @uf writes", + rule.name + ); + } + GenericCoreAction::Panic(..) => {} + } + } + + Ok(()) + } + + /// The functional-dependency merge mode of a function (from `add_table`). + pub(crate) fn merge_mode(&self, f: FunctionId) -> MergeMode { + self.info(f).merge + } + + /// Evaluate a primitive through the embedded `Database` (the base-value / + /// primitive engine). Both the host interpreter and the DD join's + /// host-side primitive phase call this, so `Value`s are bit-for-bit + /// identical to the reference backend. + pub(crate) fn eval_prim_internal( + &self, + id: ExternalFunctionId, + args: &[Value], + ) -> Result> { + let result = self + .db + .with_execution_state(|st| st.call_external_func(id, args)); + if let Some(message) = self.take_panic_message() { + Err(anyhow!(message)) + } else { + Ok(result) + } + } + + fn take_panic_message(&self) -> Option { + self.panic_message + .lock() + .expect("DD panic-message side channel must not be poisoned") + .take() + } + + /// Allocate a fresh id (the interpreter's eq-sort constructor hash-cons uses + /// the same counter the trait's `fresh_id` advances). + pub(crate) fn fresh_id_internal(&mut self) -> u32 { + let id = self.next_id; + self.next_id += 1; + id + } + + /// Apply this iteration's `set`s to a merge function `f`, folding each new + /// output value against the current value for its key by the merge mode. + /// Returns whether the mirror changed. + /// + /// `rows` is in EMISSION order, so the fold is insertion-order-correct: the + /// pre-set mirror value is the "old" value and each `rows` entry is a "new" + /// value applied in turn (`New` keeps the last, `Old` keeps the first/existing, + /// `Min` keeps the smallest, `Computed` evaluates the retained merge tree). + /// Current rows are read from live ∪ subsumed state, matching the reference + /// backend's single-table representation with a merged subsumed bit. A later + /// `set` of a subsumed key must merge with that row and remain subsumed; it + /// must not resurrect the row into ordinary matching. + pub(crate) fn apply_merge_sets(&mut self, f: FunctionId, rows: &[Row]) -> Result { + if matches!(self.merge_mode(f), MergeMode::Computed) { + return self.apply_computed_merge(f, rows); + } + let arity = self.info(f).arity; + let merge = self.merge_mode(f); + let inputs_len = arity - 1; + let mut cur = self.current_values_by_key(f, inputs_len); + // Fold new values in emission order, remembering each touched key's + // original value so the stale row can be retracted. + let mut orig: HashMap, Option> = HashMap::new(); + for row in rows { + let key: Box<[u32]> = row[..inputs_len].into(); + let nv = row[inputs_len]; + let existing = cur.get(&key).copied(); + orig.entry(key.clone()).or_insert(existing); + let existing_value = existing.map(|current| current.value); + let merged = match (existing_value, merge) { + (None, _) => nv, + (Some(_), MergeMode::New) => nv, + (Some(c), MergeMode::Old) => c, + (Some(c), MergeMode::Min) => c.min(nv), + // Relation is filtered by the caller; Computed took the branch above. + (Some(c), MergeMode::Relation | MergeMode::Computed) => c, + }; + let location = existing + .map(|current| current.location) + .unwrap_or(RowLocation::Live); + let rows_for_key = existing.map(|current| current.rows_for_key).unwrap_or(0); + cur.insert( + key, + CurrentValue { + value: merged, + location, + rows_for_key, + }, + ); + } + // Apply the net change per touched key. The initial scan above also + // counted duplicate rows, so normalization does not need another full + // relation scan for every key. + let mut replacements = Vec::new(); + for (key, old) in orig { + let current = cur[&key]; + let new = (current.value, current.location); + let already_normalized = + old.is_some_and(|old| old.rows_for_key == 1 && (old.value, old.location) == new); + if !already_normalized { + replacements.push((key, new)); + } + } + Ok(self.replace_located_rows(f, inputs_len, replacements)) + } + + /// The [`MergeMode::Computed`] case of [`apply_merge_sets`]: fold each key's + /// conflicting values by EVALUATING the retained [`MergeTree`] (a primitive + /// like `(or old new)`, or a constructor that builds a term). Evaluation needs + /// `&mut self` (host-side primitive calls, e-node minting), so — unlike the + /// scalar path — it gathers the candidate values first, then reconciles the + /// mirror after the fold. + fn apply_computed_merge(&mut self, f: FunctionId, rows: &[Row]) -> Result { + let inputs_len = self.info(f).arity - 1; + let next_id_before = self.next_id; + let cur = self.current_values_by_key(f, inputs_len); + // Group new values by key in emission order. + let mut order: Vec> = Vec::new(); + let mut newv: HashMap, Vec> = HashMap::new(); + for row in rows { + let key: Box<[u32]> = row[..inputs_len].into(); + let v = row[inputs_len]; + match newv.get_mut(&key) { + Some(vs) => vs.push(v), + None => { + order.push(key.clone()); + newv.insert(key, vec![v]); + } + } + } + // The merge tree, cloned out so the evaluator can borrow `&mut self`. + let tree = self.info(f).merge_tree.clone().ok_or_else(|| { + anyhow!( + "DD relation metadata invariant: computed merge for `{}` has no merge tree", + self.relation_name(f) + ) + })?; + let mut lookup_index: HashMap, u32>> = HashMap::new(); + // Fold each touched key: start from the current value, fold in each new + // value via the tree (`old` = accumulator, `new` = the value). + let mut updates: Vec = Vec::new(); + for key in order { + let old = cur.get(&key).copied(); + let location = old + .map(|current| current.location) + .unwrap_or(RowLocation::Live); + let mut acc = old.map(|current| current.value); + for &nv in &newv[&key] { + acc = Some(match acc { + None => nv, + Some(c) => self.eval_merge_tree(&tree, c, nv, &mut lookup_index)?, + }); + } + let new_val = acc.expect("DD computed-merge invariant: touched key has a value"); + let new = (new_val, location); + if Some(new) != old.map(|old| (old.value, old.location)) { + updates.push((key, new)); + } + } + // Reconcile the mirror. A `Func` node that minted an e-node advanced + // `next_id` — that is itself a real change even if no value flipped. + let mut changed = self.next_id != next_id_before; + changed |= self.replace_located_rows(f, inputs_len, updates); + Ok(changed) + } + + /// Evaluate a [`MergeTree`] against the two conflicting output values `old` + /// (accumulated) and `new` (incoming), returning the merged value. Primitives + /// run host-side via `eval_prim_internal`; a `Func` node hash-cons / mints a + /// constructor e-node via `lookup_or_create`, so a merge that builds a term + /// works. + fn eval_merge_tree( + &mut self, + node: &MergeTree, + old: u32, + new: u32, + index: &mut HashMap, u32>>, + ) -> Result { + match node { + MergeTree::Old => Ok(old), + MergeTree::New => Ok(new), + MergeTree::Const(v) => Ok(*v), + MergeTree::Prim(id, args) => { + let argv: Vec = args + .iter() + .map(|a| self.eval_merge_tree(a, old, new, index).map(Value::new)) + .collect::>()?; + Ok(self + .eval_prim_internal(*id, &argv)? + .map(|v| v.rep()) + .unwrap_or(old)) + } + MergeTree::Func(func, args) => { + let key: Vec = args + .iter() + .map(|a| self.eval_merge_tree(a, old, new, index).map(Value::new)) + .collect::>()?; + if self.info(*func).lookup_mints { + Ok(interpret::lookup_or_create(self, *func, &key, index).rep()) + } else { + interpret::lookup_existing(self, *func, &key, index) + .map(|value| value.rep()) + .ok_or_else(|| { + anyhow::anyhow!( + "lookup on `{}` failed in merge function", + self.relation_name(*func) + ) + }) + } + } + } + } + + /// Move every live row of `f` whose LEADING columns equal `prefix` into the + /// `subsumed` side-set (a "soft delete"): the row stays in the table — still + /// counted by `table_size`, still visible to `for_each` and to + /// `:internal-include-subsumed` rules — but is hidden from ordinary query + /// matching. Returns whether anything moved. The subsume action addresses a + /// view row by its children+output columns (not the trailing epoch), hence a + /// prefix match. + pub(crate) fn subsume_rows(&mut self, f: FunctionId, prefix: &[u32]) -> bool { + let Some(set) = self.mirror.get(&f) else { + return false; + }; + let moved: Vec = set + .iter() + .filter(|r| r.len() >= prefix.len() && r[..prefix.len()] == *prefix) + .cloned() + .collect(); + if moved.is_empty() { + return false; + } + { + let live = self + .mirror + .get_mut(&f) + .expect("DD mirror invariant: rows selected for subsumption remain registered"); + for r in &moved { + live.remove(r); + } + } + for r in &moved { + self.record_row_event(f, r.clone(), -1, 0, 1); + } + let subs = self.subsumed.entry(f).or_default(); + for r in moved { + subs.insert(r); + } + true + } + + pub(crate) fn insert_live_row(&mut self, f: FunctionId, row: Row) -> bool { + if self + .subsumed + .get(&f) + .is_some_and(|subsumed| subsumed.contains(&row)) + { + return false; + } + let inserted = self.mirror.entry(f).or_default().insert(row.clone()); + if inserted { + self.record_row_event(f, row, 1, 1, 0); + } + inserted + } + + fn current_values_by_key( + &self, + f: FunctionId, + inputs_len: usize, + ) -> HashMap, CurrentValue> { + let live_len = self.mirror.get(&f).map(|set| set.len()).unwrap_or(0); + let subsumed_len = self.subsumed.get(&f).map(|set| set.len()).unwrap_or(0); + let mut cur = HashMap::with_capacity(live_len + subsumed_len); + for (location, store) in [ + (RowLocation::Live, &self.mirror), + (RowLocation::Subsumed, &self.subsumed), + ] { + if let Some(rows) = store.get(&f) { + for row in rows.iter() { + let value = row[inputs_len]; + let key: Box<[u32]> = row[..inputs_len].into(); + cur.entry(key) + .and_modify(|current: &mut CurrentValue| { + current.value = value; + current.location = location; + current.rows_for_key += 1; + }) + .or_insert(CurrentValue { + value, + location, + rows_for_key: 1, + }); + } + } + } + cur + } + + fn replace_located_rows( + &mut self, + f: FunctionId, + keylen: usize, + replacements: Vec, + ) -> bool { + if replacements.is_empty() { + return false; + } + let keys = replacements.iter().map(|(key, _)| key.clone()).collect(); + self.remove_matching_keys(f, keylen, &keys); + for (key, (value, location)) in replacements { + let mut row = key.into_vec(); + row.push(value); + self.insert_located_row(f, location, row.into_boxed_slice()); + } + true + } + + fn insert_located_row(&mut self, f: FunctionId, location: RowLocation, row: Row) -> bool { + let store = match location { + RowLocation::Live => &mut self.mirror, + RowLocation::Subsumed => &mut self.subsumed, + }; + let inserted = store.entry(f).or_default().insert(row.clone()); + if inserted { + match location { + RowLocation::Live => self.record_row_event(f, row, 1, 1, 0), + RowLocation::Subsumed => self.record_row_event(f, row, 0, 1, 1), + } + } + inserted + } + + pub(crate) fn remove_matching_keys( + &mut self, + f: FunctionId, + keylen: usize, + keys: &HashSet>, + ) -> bool { + let mut changed = false; + let removed_live = remove_keys_from_store(&mut self.mirror, f, keylen, keys); + changed |= !removed_live.is_empty(); + for row in removed_live { + self.record_row_event(f, row, -1, -1, 0); + } + let removed_subsumed = remove_keys_from_store(&mut self.subsumed, f, keylen, keys); + changed |= !removed_subsumed.is_empty(); + for row in removed_subsumed { + self.record_row_event(f, row, 0, -1, -1); + } + changed + } + + fn record_row_event( + &mut self, + func: FunctionId, + row: Row, + live_delta: isize, + all_delta: isize, + subsumed_delta: isize, + ) { + let version = if live_delta > 0 || all_delta > 0 || subsumed_delta > 0 { + let version = self.next_row_version; + self.next_row_version += 1; + Some(version) + } else { + None + }; + update_version_map(&mut self.live_versions, func, &row, live_delta, version); + update_version_map(&mut self.all_versions, func, &row, all_delta, version); + update_version_map( + &mut self.subsumed_versions, + func, + &row, + subsumed_delta, + version, + ); + } +} + +fn remove_keys_from_store( + store: &mut HashMap>, + f: FunctionId, + keylen: usize, + keys: &HashSet>, +) -> Vec { + let mut removed = Vec::new(); + if let Some(rows) = store.get_mut(&f) { + rows.retain(|row| { + if keys.contains(&row[..keylen]) { + removed.push(row.clone()); + false + } else { + true + } + }); + } + removed +} + +fn update_version_map( + versions: &mut HashMap>, + func: FunctionId, + row: &Row, + delta: isize, + version: Option, +) { + match delta.cmp(&0) { + std::cmp::Ordering::Greater => { + versions.entry(func).or_default().insert( + row.clone(), + version.expect("positive row event needs a version"), + ); + } + std::cmp::Ordering::Less => { + if let Some(rows) = versions.get_mut(&func) { + rows.remove(row); + } + } + std::cmp::Ordering::Equal => {} + } +} + +#[cfg(test)] +#[allow(clippy::items_after_test_module)] +mod tests { + use super::*; + use std::panic::{catch_unwind, AssertUnwindSafe}; + + use egglog_ast::{ + core::{GenericAtom, GenericAtomTerm, GenericCoreAction, GenericCoreRule, Query}, + generic_ast::Change, + span::Span, + }; + use egglog_backend_trait::{ + Backend, BaseValueId, ColumnTy, DefaultVal, ExternalFunctionId, FunctionConfig, MergeFn, + ReadMode, RuleActionCall, RuleBodyCall, RuleId, RuleSetRun, RuleSpec, RuleValue, RuleVar, + Value, + }; + use egglog_numeric_id::NumericId; + + type RuleTerm = GenericAtomTerm; + + struct TestRule { + spec: RuleSpec, + next_var: u32, + } + + impl TestRule { + fn new(name: &str) -> Self { + Self { + spec: RuleSpec { + name: name.to_owned(), + seminaive: true, + no_decomp: false, + core: GenericCoreRule { + span: Span::Panic, + body: Query::default(), + head: Default::default(), + }, + }, + next_var: 0, + } + } + + fn new_var(&mut self, ty: ColumnTy) -> RuleTerm { + let id = self.next_var; + self.next_var += 1; + GenericAtomTerm::Var( + Span::Panic, + RuleVar { + id, + name: format!("v{id}").into_boxed_str(), + ty, + }, + ) + } + + fn query_table(&mut self, id: FunctionId, args: &[RuleTerm], is_subsumed: Option) { + self.spec.core.body.atoms.push(GenericAtom { + span: Span::Panic, + head: RuleBodyCall::Table { + id, + read: match is_subsumed { + Some(false) => ReadMode::Live, + Some(true) => ReadMode::Subsumed, + None => ReadMode::All, + }, + }, + args: args.to_vec(), + }); + } + + fn query_prim(&mut self, id: ExternalFunctionId, args: &[RuleTerm], output: ColumnTy) { + self.spec.core.body.atoms.push(GenericAtom { + span: Span::Panic, + head: RuleBodyCall::Primitive { + id, + name: "external".into(), + output, + }, + args: args.to_vec(), + }); + } + + fn set(&mut self, id: FunctionId, entries: &[RuleTerm]) { + let (value, args) = entries.split_last().expect("set has a value"); + self.spec.core.head.0.push(GenericCoreAction::Set( + Span::Panic, + RuleActionCall::Table { + id, + name: "table".into(), + }, + args.to_vec(), + value.clone(), + )); + } + + fn change(&mut self, id: FunctionId, entries: &[RuleTerm], change: Change) { + self.spec.core.head.0.push(GenericCoreAction::Change( + Span::Panic, + change, + RuleActionCall::Table { + id, + name: "table".into(), + }, + entries.to_vec(), + )); + } + + fn remove(&mut self, id: FunctionId, entries: &[RuleTerm]) { + self.change(id, entries, Change::Delete); + } + + fn subsume(&mut self, id: FunctionId, entries: &[RuleTerm]) { + self.change(id, entries, Change::Subsume); + } + + fn union(&mut self, lhs: RuleTerm, rhs: RuleTerm) { + self.spec + .core + .head + .0 + .push(GenericCoreAction::Union(Span::Panic, lhs, rhs)); + } + + fn build(self, egraph: &mut EGraph) -> RuleId { + Backend::add_rule(egraph, self.spec).unwrap() + } + } + + fn row(vals: &[u32]) -> Row { + vals.to_vec().into_boxed_slice() + } + + fn run_rules(egraph: &mut EGraph, rules: &[RuleId]) -> Result { + Backend::run_rules( + egraph, + RuleSetRun { + name: Some("test"), + rules, + }, + ) + } + + fn constant(value: u32, ty: ColumnTy) -> RuleTerm { + GenericAtomTerm::Literal( + Span::Panic, + RuleValue { + value: Value::new(value), + ty, + }, + ) + } + + fn id_function(eg: &mut EGraph, name: &str, merge: MergeFn) -> FunctionId { + Backend::add_table( + eg, + FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Id], + default: DefaultVal::Fail, + merge, + name: name.to_string(), + can_subsume: true, + }, + ) + } + + fn id_table(eg: &mut EGraph, name: &str, arity: usize) -> FunctionId { + Backend::add_table( + eg, + FunctionConfig { + schema: vec![ColumnTy::Id; arity], + default: DefaultVal::Fail, + merge: MergeFn::Old, + name: name.to_string(), + can_subsume: false, + }, + ) + } + + fn neq_primitive(eg: &mut EGraph) -> ExternalFunctionId { + Backend::register_external_func( + eg, + Box::new(egglog_core_relations::make_external_func(|_, args| { + (args[0] != args[1]).then_some(Value::new(0)) + })), + ) + } + + fn max_primitive(eg: &mut EGraph) -> ExternalFunctionId { + Backend::register_external_func( + eg, + Box::new(egglog_core_relations::make_external_func(|_, args| { + Some(args[0].max(args[1])) + })), + ) + } + + fn self_join_tables(eg: &mut EGraph) -> (BaseValueId, FunctionId, FunctionId) { + let unit_ty = eg.db.base_values().get_ty::<()>(); + let relation = Backend::add_table( + eg, + FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Base(unit_ty)], + default: DefaultVal::Fail, + merge: MergeFn::Old, + name: "R".to_string(), + can_subsume: false, + }, + ); + let out = Backend::add_table( + eg, + FunctionConfig { + schema: vec![ + ColumnTy::Id, + ColumnTy::Id, + ColumnTy::Id, + ColumnTy::Base(unit_ty), + ], + default: DefaultVal::Fail, + merge: MergeFn::Old, + name: "Out".to_string(), + can_subsume: false, + }, + ); + eg.insert_live_row(relation, row(&[1, 2, 0])); + eg.insert_live_row(relation, row(&[1, 3, 0])); + (unit_ty, relation, out) + } + + fn path_compression_fixture(eg: &mut EGraph) -> (FunctionId, RuleId) { + let neq = neq_primitive(eg); + let unit_ty = eg.db.base_values().get_ty::<()>(); + let uf = Backend::add_table( + eg, + FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Id, ColumnTy::Base(unit_ty)], + default: DefaultVal::Fail, + merge: MergeFn::Old, + name: "UF".to_string(), + can_subsume: false, + }, + ); + let mut rb = TestRule::new("path_compress"); + let a = rb.new_var(ColumnTy::Id); + let b = rb.new_var(ColumnTy::Id); + let c = rb.new_var(ColumnTy::Id); + let unit = constant(0, ColumnTy::Base(unit_ty)); + rb.query_table(uf, &[a.clone(), b.clone(), unit.clone()], Some(false)); + rb.query_table(uf, &[b.clone(), c.clone(), unit.clone()], Some(false)); + rb.query_prim( + neq, + &[b.clone(), c.clone(), unit.clone()], + ColumnTy::Base(unit_ty), + ); + rb.remove(uf, &[a.clone(), b]); + rb.set(uf, &[a, c, unit]); + (uf, rb.build(eg)) + } + + fn write_order_tables(eg: &mut EGraph) -> (BaseValueId, FunctionId, FunctionId) { + let unit_ty = eg.db.base_values().get_ty::<()>(); + let trigger = Backend::add_table( + eg, + FunctionConfig { + schema: vec![ColumnTy::Base(unit_ty)], + default: DefaultVal::Fail, + merge: MergeFn::Old, + name: "Trigger".to_string(), + can_subsume: false, + }, + ); + let f = Backend::add_table( + eg, + FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Id], + default: DefaultVal::Fail, + merge: MergeFn::New, + name: "F".to_string(), + can_subsume: true, + }, + ); + eg.insert_live_row(trigger, row(&[0])); + (unit_ty, trigger, f) + } + + #[test] + fn native_union_is_rejected_when_rule_is_added() { + let mut eg = EGraph::new(); + let left = constant(1, ColumnTy::Id); + let right = constant(2, ColumnTy::Id); + let mut builder = TestRule::new("native union"); + builder.union(left, right); + + let error = Backend::add_rule(&mut eg, builder.spec).unwrap_err(); + assert!(error.to_string().contains("received a native union")); + } + + #[test] + fn table_and_rule_arity_match_fixed_dd_row_width() { + let mut eg = EGraph::new(); + let wide = id_table(&mut eg, "wide", dd_native::W); + + let mut boundary = TestRule::new("width boundary"); + let boundary_args = (0..dd_native::W) + .map(|_| boundary.new_var(ColumnTy::Id)) + .collect::>(); + boundary.query_table(wide, &boundary_args, Some(false)); + Backend::add_rule(&mut eg, boundary.spec).expect("width-W rule must be admitted"); + + let mut too_wide_rule = TestRule::new("too wide"); + let too_wide_args = (0..=dd_native::W) + .map(|_| too_wide_rule.new_var(ColumnTy::Id)) + .collect::>(); + too_wide_rule.query_table(wide, &too_wide_args, Some(false)); + let error = Backend::add_rule(&mut eg, too_wide_rule.spec).unwrap_err(); + assert!(error + .to_string() + .contains(&format!("{} columns", dd_native::W + 1))); + assert!(error + .to_string() + .contains(&format!("fixed row width {}", dd_native::W))); + + let relation_count = eg.relations.len(); + let oversized_table = catch_unwind(AssertUnwindSafe(|| { + id_table(&mut eg, "oversized", dd_native::W + 1) + })); + assert!(oversized_table.is_err()); + assert_eq!(eg.relations.len(), relation_count); + } + + #[test] + fn add_rule_rejects_residual_globals() { + let mut eg = EGraph::new(); + let input = id_table(&mut eg, "input", 1); + let mut rule = TestRule::new("global"); + let global = GenericAtomTerm::Global( + Span::Panic, + RuleVar { + id: 0, + name: "global".into(), + ty: ColumnTy::Id, + }, + ); + rule.query_table(input, &[global], Some(false)); + + let error = Backend::add_rule(&mut eg, rule.spec).unwrap_err(); + assert!(error.to_string().contains("residual global")); + } + + #[test] + fn join_planning_failure_is_returned_without_unwinding() { + let mut eg = EGraph::new(); + let wide = id_table(&mut eg, "wide", dd_native::W); + let edge = id_table(&mut eg, "edge", 2); + let primitive = max_primitive(&mut eg); + let mut rule = TestRule::new("wide live frontier"); + let vars = (0..=dd_native::W) + .map(|_| rule.new_var(ColumnTy::Id)) + .collect::>(); + rule.query_table(wide, &vars[..dd_native::W], Some(false)); + rule.query_table( + edge, + &[vars[0].clone(), vars[dd_native::W].clone()], + Some(false), + ); + rule.query_prim(primitive, &vars, ColumnTy::Id); + let rule = rule.build(&mut eg); + + let outcome = catch_unwind(AssertUnwindSafe(|| run_rules(&mut eg, &[rule]))); + let error = outcome + .expect("unsupported join planning must not unwind") + .unwrap_err(); + assert!(error + .to_string() + .contains("live body-variable frontier exceeds W 48")); + } + + #[test] + fn merge_set_preserves_subsumed_status() { + let mut eg = EGraph::new(); + let f = id_function(&mut eg, "f", MergeFn::New); + eg.subsumed.entry(f).or_default().insert(row(&[1, 10])); + + assert!(eg.apply_merge_sets(f, &[row(&[1, 11])]).unwrap()); + + assert!(!eg.mirror[&f].contains(&row(&[1, 11]))); + assert!(!eg.subsumed[&f].contains(&row(&[1, 10]))); + assert!(eg.subsumed[&f].contains(&row(&[1, 11]))); + } + + #[test] + fn lookup_or_create_finds_subsumed_rows() { + let mut eg = EGraph::new(); + let f = id_function(&mut eg, "f", MergeFn::New); + eg.next_id = 100; + eg.subsumed.entry(f).or_default().insert(row(&[42, 7])); + + let mut lookup_index = HashMap::new(); + let value = interpret::lookup_or_create(&mut eg, f, &[Value::new(42)], &mut lookup_index); + + assert_eq!(value.rep(), 7); + assert_eq!(eg.next_id, 100); + assert!(eg.mirror[&f].is_empty()); + } + + #[test] + fn merge_set_collapses_live_subsumed_key_duplicate() { + let mut eg = EGraph::new(); + let f = id_function(&mut eg, "f", MergeFn::New); + eg.mirror.entry(f).or_default().insert(row(&[1, 10])); + eg.subsumed.entry(f).or_default().insert(row(&[1, 11])); + + assert!(eg.apply_merge_sets(f, &[row(&[1, 12])]).unwrap()); + + assert!(eg.mirror[&f].is_empty()); + assert_eq!(eg.subsumed[&f].len(), 1); + assert!(eg.subsumed[&f].contains(&row(&[1, 12]))); + } + + #[test] + fn same_table_self_join_produces_cross_pairs() { + let mut eg = EGraph::new(); + let (unit_ty, relation, out) = self_join_tables(&mut eg); + + let rule = { + let mut rb = TestRule::new("self_join"); + let a = rb.new_var(ColumnTy::Id); + let b = rb.new_var(ColumnTy::Id); + let c = rb.new_var(ColumnTy::Id); + let unit = constant(0, ColumnTy::Base(unit_ty)); + rb.query_table(relation, &[a.clone(), b.clone(), unit.clone()], Some(false)); + rb.query_table(relation, &[a.clone(), c.clone(), unit.clone()], Some(false)); + rb.set(out, &[a, b, c, unit]); + rb.build(&mut eg) + }; + + run_rules(&mut eg, &[rule]).unwrap(); + + assert!(eg.mirror[&out].contains(&row(&[1, 2, 3, 0]))); + assert!(eg.mirror[&out].contains(&row(&[1, 3, 2, 0]))); + } + + #[test] + fn fused_ruleset_allows_mixed_read_modes_same_relation() { + let mut eg = EGraph::new(); + let unit_ty = eg.db.base_values().get_ty::<()>(); + let relation = Backend::add_table( + &mut eg, + FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Base(unit_ty)], + default: DefaultVal::Fail, + merge: MergeFn::Old, + name: "R".to_string(), + can_subsume: true, + }, + ); + let live_out = Backend::add_table( + &mut eg, + FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Base(unit_ty)], + default: DefaultVal::Fail, + merge: MergeFn::Old, + name: "LiveOut".to_string(), + can_subsume: false, + }, + ); + let all_out = Backend::add_table( + &mut eg, + FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Base(unit_ty)], + default: DefaultVal::Fail, + merge: MergeFn::Old, + name: "AllOut".to_string(), + can_subsume: false, + }, + ); + eg.insert_live_row(relation, row(&[1, 0])); + eg.insert_located_row(relation, RowLocation::Subsumed, row(&[2, 0])); + + let live_rule = { + let mut rb = TestRule::new("live_read"); + let x = rb.new_var(ColumnTy::Id); + let unit = constant(0, ColumnTy::Base(unit_ty)); + rb.query_table(relation, &[x.clone(), unit.clone()], Some(false)); + rb.set(live_out, &[x, unit]); + rb.build(&mut eg) + }; + let all_rule = { + let mut rb = TestRule::new("all_read"); + let x = rb.new_var(ColumnTy::Id); + let unit = constant(0, ColumnTy::Base(unit_ty)); + rb.query_table(relation, &[x.clone(), unit.clone()], None); + rb.set(all_out, &[x, unit]); + rb.build(&mut eg) + }; + + run_rules(&mut eg, &[live_rule, all_rule]).unwrap(); + + assert!(eg.mirror[&live_out].contains(&row(&[1, 0]))); + assert!(!eg.mirror[&live_out].contains(&row(&[2, 0]))); + assert!(eg.mirror[&all_out].contains(&row(&[1, 0]))); + assert!(eg.mirror[&all_out].contains(&row(&[2, 0]))); + } + + #[test] + fn fused_ruleset_cache_maps_reordered_rules_by_id() { + let mut eg = EGraph::new(); + let input_a = id_table(&mut eg, "InputA", 2); + let input_b = id_table(&mut eg, "InputB", 3); + let output_a = id_table(&mut eg, "OutputA", 2); + let output_b = id_table(&mut eg, "OutputB", 2); + + let rule_a = { + let mut rb = TestRule::new("rule_a"); + let x = rb.new_var(ColumnTy::Id); + let y = rb.new_var(ColumnTy::Id); + rb.query_table(input_a, &[x.clone(), y.clone()], Some(false)); + rb.set(output_a, &[y, x]); + rb.build(&mut eg) + }; + let rule_b = { + let mut rb = TestRule::new("rule_b"); + let p = rb.new_var(ColumnTy::Id); + let q = rb.new_var(ColumnTy::Id); + let r = rb.new_var(ColumnTy::Id); + rb.query_table(input_b, &[r.clone(), p.clone(), q], Some(false)); + rb.set(output_b, &[p, r]); + rb.build(&mut eg) + }; + + eg.insert_live_row(input_a, row(&[1, 2])); + eg.insert_live_row(input_b, row(&[5, 3, 4])); + run_rules(&mut eg, &[rule_a, rule_b]).unwrap(); + + eg.insert_live_row(input_a, row(&[10, 20])); + eg.insert_live_row(input_b, row(&[50, 30, 40])); + run_rules(&mut eg, &[rule_b, rule_a]).unwrap(); + + assert!(eg.mirror[&output_a].contains(&row(&[20, 10]))); + assert!(eg.mirror[&output_b].contains(&row(&[30, 50]))); + } + + #[test] + fn same_table_self_join_applies_primitive_guards() { + let mut eg = EGraph::new(); + let neq = neq_primitive(&mut eg); + let ordering_max = max_primitive(&mut eg); + let (unit_ty, relation, out) = self_join_tables(&mut eg); + + let rule = { + let mut rb = TestRule::new("guarded_self_join"); + let a = rb.new_var(ColumnTy::Id); + let b = rb.new_var(ColumnTy::Id); + let c = rb.new_var(ColumnTy::Id); + let unit = constant(0, ColumnTy::Base(unit_ty)); + rb.query_table(relation, &[a.clone(), b.clone(), unit.clone()], Some(false)); + rb.query_table(relation, &[a.clone(), c.clone(), unit.clone()], Some(false)); + rb.query_prim( + neq, + &[b.clone(), c.clone(), unit.clone()], + ColumnTy::Base(unit_ty), + ); + rb.query_prim( + ordering_max, + &[b.clone(), c.clone(), b.clone()], + ColumnTy::Id, + ); + rb.set(out, &[a, b, c, unit]); + rb.build(&mut eg) + }; + + run_rules(&mut eg, &[rule]).unwrap(); + + assert!(eg.mirror[&out].contains(&row(&[1, 3, 2, 0]))); + assert!(!eg.mirror[&out].contains(&row(&[1, 2, 3, 0]))); + } + + #[test] + fn same_table_self_join_allows_independent_unit_outputs() { + let mut eg = EGraph::new(); + let neq = neq_primitive(&mut eg); + let ordering_max = max_primitive(&mut eg); + let (unit_ty, relation, out) = self_join_tables(&mut eg); + + let rule = { + let mut rb = TestRule::new("unit_var_self_join"); + let a = rb.new_var(ColumnTy::Id); + let b = rb.new_var(ColumnTy::Id); + let c = rb.new_var(ColumnTy::Id); + let unit1 = rb.new_var(ColumnTy::Base(unit_ty)); + let unit2 = rb.new_var(ColumnTy::Base(unit_ty)); + let unit = constant(0, ColumnTy::Base(unit_ty)); + rb.query_table(relation, &[a.clone(), b.clone(), unit1], Some(false)); + rb.query_table(relation, &[a.clone(), c.clone(), unit2], Some(false)); + rb.query_prim( + neq, + &[b.clone(), c.clone(), unit.clone()], + ColumnTy::Base(unit_ty), + ); + rb.query_prim( + ordering_max, + &[b.clone(), c.clone(), b.clone()], + ColumnTy::Id, + ); + rb.set(out, &[a, b, c, unit]); + rb.build(&mut eg) + }; + + run_rules(&mut eg, &[rule]).unwrap(); + + assert!(eg.mirror[&out].contains(&row(&[1, 3, 2, 0]))); + assert!(!eg.mirror[&out].contains(&row(&[1, 2, 3, 0]))); + } + + #[test] + fn incremental_self_join_retracts_old_path_edge() { + let mut eg = EGraph::new(); + let (uf, rule) = path_compression_fixture(&mut eg); + + eg.insert_live_row(uf, row(&[40, 4, 0])); + run_rules(&mut eg, &[rule]).unwrap(); + eg.insert_live_row(uf, row(&[75, 40, 0])); + run_rules(&mut eg, &[rule]).unwrap(); + + assert!(!eg.mirror[&uf].contains(&row(&[75, 40, 0]))); + assert!(eg.mirror[&uf].contains(&row(&[75, 4, 0]))); + } + + #[test] + fn reinserted_same_row_refires_incremental_join() { + let mut eg = EGraph::new(); + let (uf, rule) = path_compression_fixture(&mut eg); + + eg.insert_live_row(uf, row(&[40, 4, 0])); + eg.insert_live_row(uf, row(&[75, 40, 0])); + run_rules(&mut eg, &[rule]).unwrap(); + assert!(!eg.mirror[&uf].contains(&row(&[75, 40, 0]))); + assert!(eg.mirror[&uf].contains(&row(&[75, 4, 0]))); + + eg.insert_live_row(uf, row(&[75, 40, 0])); + run_rules(&mut eg, &[rule]).unwrap(); + + assert!(!eg.mirror[&uf].contains(&row(&[75, 40, 0]))); + assert!(eg.mirror[&uf].contains(&row(&[75, 4, 0]))); + } + + #[test] + fn clone_preserves_authoritative_state_and_rebuilds_transient_join() { + let mut eg = EGraph::new(); + let (uf, rule) = path_compression_fixture(&mut eg); + eg.insert_live_row(uf, row(&[40, 4, 0])); + eg.insert_live_row(uf, row(&[75, 40, 0])); + run_rules(&mut eg, &[rule]).unwrap(); + assert!(!eg.dd_fused.is_empty()); + assert!(!eg.dd_fused_fed_versions.is_empty()); + let panic = Backend::new_panic(&mut eg, "cloned panic".to_owned()); + + let mut cloned = Backend::clone_boxed(&eg); + let cloned = cloned + .as_any_mut() + .downcast_mut::() + .expect("DD clone must retain its concrete backend type"); + + assert!(cloned.dd_fused.is_empty()); + assert!(cloned.dd_fused_fed_versions.is_empty()); + assert_eq!(cloned.relations.len(), eg.relations.len()); + assert_eq!(cloned.rules.len(), eg.rules.len()); + assert_eq!(cloned.mirror, eg.mirror); + assert_eq!(cloned.subsumed, eg.subsumed); + assert_eq!(cloned.next_id, eg.next_id); + assert_eq!(cloned.next_row_version, eg.next_row_version); + assert_eq!(cloned.live_versions, eg.live_versions); + assert_eq!(cloned.all_versions, eg.all_versions); + assert_eq!(cloned.subsumed_versions, eg.subsumed_versions); + + let outcome = catch_unwind(AssertUnwindSafe(|| { + cloned.eval_prim_internal(panic, &[Value::new(7)]) + })); + let error = outcome + .expect("a cloned panic primitive must not unwind") + .unwrap_err(); + assert_eq!(error.to_string(), "cloned panic"); + + cloned.insert_live_row(uf, row(&[90, 75, 0])); + run_rules(cloned, &[rule]).unwrap(); + assert!(cloned.mirror[&uf].contains(&row(&[90, 4, 0]))); + assert!(!eg.mirror[&uf].contains(&row(&[90, 4, 0]))); + assert!(!cloned.dd_fused.is_empty()); + assert!(!cloned.dd_fused_fed_versions.is_empty()); + } + + #[test] + fn same_iteration_remove_then_set_keeps_set_value() { + let mut eg = EGraph::new(); + let (unit_ty, trigger, f) = write_order_tables(&mut eg); + eg.insert_live_row(f, row(&[1, 9])); + + let rule = { + let mut rb = TestRule::new("remove_then_set"); + let unit = constant(0, ColumnTy::Base(unit_ty)); + rb.query_table(trigger, std::slice::from_ref(&unit), Some(false)); + let key = constant(1, ColumnTy::Id); + let new = constant(2, ColumnTy::Id); + rb.remove(f, std::slice::from_ref(&key)); + rb.set(f, &[key, new]); + rb.build(&mut eg) + }; + + run_rules(&mut eg, &[rule]).unwrap(); + + assert!(!eg.mirror[&f].contains(&row(&[1, 9]))); + assert!(eg.mirror[&f].contains(&row(&[1, 2]))); + } + + #[test] + fn same_iteration_set_then_subsume_ends_subsumed() { + let mut eg = EGraph::new(); + let (unit_ty, trigger, f) = write_order_tables(&mut eg); + + let rule = { + let mut rb = TestRule::new("set_then_subsume"); + let unit = constant(0, ColumnTy::Base(unit_ty)); + rb.query_table(trigger, std::slice::from_ref(&unit), Some(false)); + let key = constant(1, ColumnTy::Id); + let value = constant(2, ColumnTy::Id); + rb.set(f, &[key.clone(), value]); + rb.subsume(f, &[key]); + rb.build(&mut eg) + }; + + run_rules(&mut eg, &[rule]).unwrap(); + + assert!(!eg.mirror[&f].contains(&row(&[1, 2]))); + assert!(eg.subsumed[&f].contains(&row(&[1, 2]))); + } + + #[test] + fn fused_delta_feed_does_not_expose_mixed_old_new_snapshots() { + let mut eg = EGraph::new(); + let neq = neq_primitive(&mut eg); + let unit_ty = eg.db.base_values().get_ty::<()>(); + let view = Backend::add_table( + &mut eg, + FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Base(unit_ty)], + default: DefaultVal::Fail, + merge: MergeFn::Old, + name: "View".to_string(), + can_subsume: false, + }, + ); + let current = Backend::add_table( + &mut eg, + FunctionConfig { + schema: vec![ColumnTy::Id], + default: DefaultVal::Fail, + merge: MergeFn::New, + name: "Current".to_string(), + can_subsume: false, + }, + ); + let dummy = Backend::add_table( + &mut eg, + FunctionConfig { + schema: vec![ColumnTy::Id, ColumnTy::Base(unit_ty)], + default: DefaultVal::Fail, + merge: MergeFn::Old, + name: "Dummy".to_string(), + can_subsume: false, + }, + ); + eg.insert_live_row(view, row(&[10, 0])); + assert!(eg.apply_merge_sets(current, &[row(&[10])]).unwrap()); + + let view_first_rule = { + let mut rb = TestRule::new("view_first"); + let x = rb.new_var(ColumnTy::Id); + let unit = constant(0, ColumnTy::Base(unit_ty)); + rb.query_table(view, &[x.clone(), unit.clone()], Some(false)); + rb.set(dummy, &[x, unit]); + rb.build(&mut eg) + }; + let cleanup_rule = { + let mut rb = TestRule::new("current_cleanup"); + let selected = rb.new_var(ColumnTy::Id); + let old = rb.new_var(ColumnTy::Id); + let unit = constant(0, ColumnTy::Base(unit_ty)); + rb.query_table(current, std::slice::from_ref(&selected), Some(false)); + rb.query_table(view, &[selected.clone(), unit.clone()], Some(false)); + rb.query_table(view, &[old.clone(), unit.clone()], Some(false)); + rb.query_prim( + neq, + &[selected.clone(), old.clone(), unit.clone()], + ColumnTy::Base(unit_ty), + ); + rb.remove(view, &[old]); + rb.build(&mut eg) + }; + + run_rules(&mut eg, &[view_first_rule, cleanup_rule]).unwrap(); + assert!(eg.mirror[&view].contains(&row(&[10, 0]))); + + eg.insert_live_row(view, row(&[4, 0])); + assert!(eg.apply_merge_sets(current, &[row(&[4])]).unwrap()); + + run_rules(&mut eg, &[view_first_rule, cleanup_rule]).unwrap(); + + assert!(eg.mirror[&view].contains(&row(&[4, 0]))); + assert!(!eg.mirror[&view].contains(&row(&[10, 0]))); + } +} + +// --------------------------------------------------------------------------- +// Merge-mode recognition (shared by `add_table`) +// --------------------------------------------------------------------------- + +/// Map a single-output-column [`MergeFn`] to this backend's [`MergeMode`] (plus, for +/// [`MergeMode::Computed`], the evaluatable [`MergeTree`]) that resolves its FD +/// conflict: +/// - `AssertEq` / `Old` => keep the old value +/// - `New` => keep the new value +/// - `UnionId` => lattice-min (the union-find leader) +/// - `Primitive` / `Function` / `Const` => `Computed`: fold by evaluating the +/// retained tree (`or`/`max`/`+`, or a constructor that builds a term) +fn merge_mode_for(merge: &MergeFn) -> (MergeMode, Option) { + match merge { + MergeFn::AssertEq | MergeFn::Old => (MergeMode::Old, None), + MergeFn::New => (MergeMode::New, None), + MergeFn::UnionId => (MergeMode::Min, None), + MergeFn::Primitive(_, _) | MergeFn::Function(_, _) | MergeFn::Const(_) => { + (MergeMode::Computed, translate_merge_tree(merge)) + } + } +} + +/// Translate a trait [`MergeFn`] into the evaluatable [`MergeTree`]. `UnionId` +/// has no tree form (it is `MergeMode::Min`); it only appears nested if the term +/// encoder builds one, in which case lattice-min is the faithful lowering. +fn translate_merge_tree(merge: &MergeFn) -> Option { + Some(match merge { + MergeFn::Old | MergeFn::AssertEq => MergeTree::Old, + MergeFn::New => MergeTree::New, + MergeFn::Const(v) => MergeTree::Const(v.rep()), + MergeFn::Primitive(id, args) => MergeTree::Prim( + *id, + args.iter() + .map(translate_merge_tree) + .collect::>()?, + ), + MergeFn::Function(f, args) => MergeTree::Func( + *f, + args.iter() + .map(translate_merge_tree) + .collect::>()?, + ), + MergeFn::UnionId => return None, + }) +} + +// --------------------------------------------------------------------------- +// impl Backend +// --------------------------------------------------------------------------- + +impl Backend for EGraph { + // -- table lifecycle ---------------------------------------------------- + + fn add_table(&mut self, config: FunctionConfig) -> FunctionId { + let id = FunctionId::new(self.relations.len() as u32); + let arity = config.schema.len(); + assert!( + arity <= dd_native::W, + "DD backend fixed row width supports relations of arity <= {} (got {} for `{}`)", + dd_native::W, + arity, + config.name + ); + // Relation vs function: a table is a **relation** (whole row is the key, + // no output column to merge) iff it is nullary OR its last column is + // `Unit` — the term encoder's view-table pattern + // `(function @XView (...) Unit :merge old)` AND ordinary relations. + // Otherwise the last column is a function OUTPUT, resolved by the merge + // mode. Detected via Unit — NOT `DefaultVal`, which is `Fail` for every + // custom function regardless of whether it has a real output column. + let output_is_unit = config.schema.last().is_some_and(|t| match t { + ColumnTy::Base(bv) => { + // `()` is pre-registered in `new`, so this never panics. + *bv == self + .db + .base_values() + .get_ty_by_id(std::any::TypeId::of::<()>()) + } + _ => false, + }); + let has_output = arity > 0 && !output_is_unit; + let (merge, merge_tree) = if !has_output { + (MergeMode::Relation, None) + } else { + merge_mode_for(&config.merge) + }; + self.relations.push(RelationInfo { + name: config.name, + arity, + lookup_mints: matches!(config.default, DefaultVal::FreshId), + merge, + merge_tree, + }); + self.mirror.insert(id, HashSet::new()); + id + } + + fn table_size(&self, table: FunctionId) -> usize { + // Subsumed rows still count (they remain in the table, just hidden from + // ordinary matching), so `(print-size)` matches the reference backend. + let live = self.mirror.get(&table).map(|s| s.len()).unwrap_or(0); + let subsumed = self.subsumed.get(&table).map(|s| s.len()).unwrap_or(0); + live + subsumed + } + + // -- iteration ---------------------------------------------------------- + + fn for_each_while_dyn( + &self, + table: FunctionId, + f: &mut dyn for<'r> FnMut(ScanEntry<'r>) -> bool, + ) { + let arity = self.info(table).arity; + // Live rows first, then subsumed rows (reported with `subsumed: true`), so + // extraction/serialization see the whole table with the flag. + let live = self.mirror.get(&table).into_iter().flat_map(|s| s.iter()); + let subs = self.subsumed.get(&table).into_iter().flat_map(|s| s.iter()); + for (row, subsumed) in live.map(|r| (r, false)).chain(subs.map(|r| (r, true))) { + let vals = row[..arity] + .iter() + .copied() + .map(Value::new) + .collect::>(); + let entry = ScanEntry { + vals: &vals, + subsumed, + }; + if !f(entry) { + break; + } + } + } + + // -- direct access ------------------------------------------------------ + + fn get_canon_repr(&self, val: Value, _ty: ColumnTy) -> Value { + // No union-find; canonicalization is the identity. + val + } + + fn clear_table(&mut self, func: FunctionId) { + if let Some(set) = self.mirror.get_mut(&func) { + let removed: Vec = set.iter().cloned().collect(); + set.clear(); + for row in removed { + self.record_row_event(func, row, -1, -1, 0); + } + } + if let Some(set) = self.subsumed.get_mut(&func) { + let removed: Vec = set.iter().cloned().collect(); + set.clear(); + for row in removed { + self.record_row_event(func, row, 0, -1, -1); + } + } + } + + fn base_values(&self) -> &BaseValues { + self.db.base_values() + } + + fn base_values_mut(&mut self) -> &mut BaseValues { + self.db.base_values_mut() + } + + fn container_values(&self) -> &ContainerValues { + self.db.container_values() + } + + fn lookup_id(&self, func: FunctionId, key: &[Value]) -> Option { + let info = self.info(func); + if info.merge == MergeMode::Relation || key.len() + 1 != info.arity { + return None; + } + let key: Vec = key.iter().map(|value| value.rep()).collect(); + let find = |rows: Option<&HashSet>| { + rows? + .iter() + .find_map(|row| (row[..key.len()] == key[..]).then(|| Value::new(row[key.len()]))) + }; + find(self.mirror.get(&func)).or_else(|| find(self.subsumed.get(&func))) + } + + fn container_values_mut_dyn(&mut self) -> Option<&mut ContainerValues> { + Some(self.db.container_values_mut()) + } + + fn new_container_id_counter(&mut self) -> Option { + Some(self.db.add_counter()) + } + + fn container_merge_fn(&self, _container_type: TypeId) -> Option { + // The supported proof/term subset interns container values but does not + // rely on merging distinct ids for one rebuilt value. Keep that subset's + // handles deterministic. Faithful conflicting-id rebuild would also need + // to stage equality into the authoritative DD mirror, which is outside + // this backend's current container contract. + Some(Arc::new(|_state, old, new| std::cmp::min(old, new))) + } + + fn with_execution_state_tracked_dyn(&self, f: &mut dyn FnMut(&mut ExecutionState<'_>)) -> bool { + self.db.with_execution_state_tracked(|st| f(st)).1 + } + + // -- rule management ---------------------------------------------------- + + fn add_rule(&mut self, rule: RuleSpec) -> Result { + self.validate_rule(&rule)?; + let id = RuleId::new(self.rules.len() as u32); + self.rules.push(Some(rule)); + Ok(id) + } + + fn free_rule(&mut self, id: RuleId) { + if let Some(slot) = self.rules.get_mut(id.rep() as usize) { + *slot = None; + let i = id.rep() as usize; + self.seen.remove(&i); + // Any fused ruleset that included this rule is now stale: drop it so + // it is rebuilt (without the freed rule) on the next `run_rules`. + self.dd_fused.retain(|key| !key.contains(&i)); + self.dd_fused_fed_versions + .retain(|key, _| !key.contains(&i)); + } + } + + fn run_rules(&mut self, run: RuleSetRun<'_>) -> Result { + // One `run_rules` call = one bounded egglog iteration. The frontend calls + // this N times for `(run N)`. + if let Some(message) = self.take_panic_message() { + return Err(anyhow!(message)); + } + if run.rules.is_empty() { + return Ok(IterationReport::default()); + } + let rules: Vec<(usize, RuleSpec)> = run + .rules + .iter() + .map(|r| r.rep() as usize) + .filter_map(|i| { + self.rules + .get(i) + .and_then(Option::as_ref) + .cloned() + .map(|rule| (i, rule)) + }) + .collect(); + if rules.is_empty() { + return Ok(IterationReport::default()); + } + + let changed = interpret::run_iteration(self, &rules); + if let Some(message) = self.take_panic_message() { + return Err(anyhow!(message)); + } + let changed = changed?; + + let mut report = IterationReport::default(); + report.rule_set_report.changed = changed; + Ok(report) + } + + fn flush_updates(&mut self) -> bool { + // Seed inserts land in the mirror immediately; the DD join runs inside + // `run_rules`. No separate flush. + false + } + + // -- primitives --------------------------------------------------------- + + fn register_external_func( + &mut self, + func: Box, + ) -> ExternalFunctionId { + self.db.add_external_function(func) + } + + fn free_external_func(&mut self, func: ExternalFunctionId) { + self.db.free_external_function(func); + } + + fn new_panic(&mut self, message: String) -> ExternalFunctionId { + let panic_message = Arc::clone(&self.panic_message); + self.db + .add_external_function(Box::new(egglog_core_relations::make_external_func( + move |state, _args| { + state.trigger_early_stop(); + let mut pending = panic_message + .lock() + .expect("DD panic-message side channel must not be poisoned"); + if pending.is_none() { + *pending = Some(message.clone()); + } + None + }, + ))) + } + + // -- capability flags --------------------------------------------------- + + fn requires_term_encoding(&self) -> bool { + // This backend has no native union-find: congruence and rebuild are + // lowered to ordinary rules over `@uf` tables by the term encoder. The + // frontend refuses to run this backend without that encoding, and rule + // registration defensively rejects a native `GenericCoreAction::Union`. + true + } + + fn supports_containers(&self) -> bool { + true + } + + // -- diagnostics -------------------------------------------------------- + + fn set_report_level(&mut self, _level: ReportLevel) {} + + fn dump_debug_info(&self) { + for (i, info) in self.relations.iter().enumerate() { + let f = FunctionId::new(i as u32); + let n = self.table_size(f); + if n == 0 { + continue; + } + log::info!("== DD relation `{}` ({} rows) ==", info.name, n); + } + } + + // -- cloning ------------------------------------------------------------ + + fn clone_boxed(&self) -> Box { + // Timely workers are transient and cannot be cloned. The authoritative + // relation/rule/database/mirror/version state is copied, while workers + // and their fed-version snapshots start empty. The next `run_rules` + // lazily rebuilds the fused dataflow and feeds the cloned current state. + Box::new(Self { + relations: self.relations.clone(), + rules: self.rules.clone(), + mirror: self.mirror.clone(), + subsumed: self.subsumed.clone(), + db: self.db.clone(), + next_id: self.next_id, + seen: self.seen.clone(), + dd_fused: DdWorkers::default(), + next_row_version: self.next_row_version, + live_versions: self.live_versions.clone(), + all_versions: self.all_versions.clone(), + subsumed_versions: self.subsumed_versions.clone(), + dd_fused_fed_versions: HashMap::new(), + panic_message: Arc::clone(&self.panic_message), + }) + } + + // -- bridge-only escape hatch ------------------------------------------ + + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } +} diff --git a/egglog-experimental/dd/tests/files.rs b/egglog-experimental/dd/tests/files.rs new file mode 100644 index 0000000..412fbf1 --- /dev/null +++ b/egglog-experimental/dd/tests/files.rs @@ -0,0 +1,221 @@ +//! Runs egglog's `.egg` corpus through the DD backend and checks each +//! result against egglog's OWN golden snapshot. +//! +//! egglog's `tests/files.rs` appends `(print-size)` to every corpus file and, for +//! each file, stores ONE snapshot (`egglog/tests/snapshots/ +//! files__shared_snapshot_.snap`) whose content is normalized to be +//! identical across the native / term-encoding / proofs treatments +//! (`CommandOutput::snapshot_stable_under_proof_encoding`). DD is +//! term-encoded, so its normalized output must equal that shared snapshot. This +//! test reuses those files directly — egglog is the single source of truth, and +//! DD keeps no snapshots of its own. +//! +//! egglog's snapshot is only meaningful for DD when it is term-encoding +//! output: egglog runs the term-encoding treatment (and so the shared snapshot is +//! stable across native + term encoding) exactly when `file_supports_proofs`. For +//! files that don't support proofs the snapshot is native-only, which a +//! term-encoded backend needn't reproduce, so those are skipped. +//! +//! Per (proof-supporting) corpus file: +//! - DD can't run it (an unsupported FEATURE — see [`KNOWN_UNSUPPORTED`]) → +//! passes only if the returned error contains that file's expected diagnostic; +//! panics, new failures, and listed files that start working all fail the +//! harness so the list stays honest; +//! - DD runs it and its normalized output equals egglog's snapshot → pass; +//! - otherwise the normal insta assertion reports the mismatch. +//! +//! Run with `--release`: differential-dataflow is impractically slow in debug +//! builds. Under `cargo test` (debug) this harness covers only a fast +//! [`DEBUG_SUBSET`]; `cargo test --release` runs the whole corpus. + +use egglog::{file_supports_proofs, CommandOutput, EGraph}; +use libtest_mimic::{Arguments, Failed, Trial}; +use std::panic::{catch_unwind, AssertUnwindSafe}; +use std::path::PathBuf; + +/// Files that saturate too slowly on the differential-dataflow path to belong in +/// a test suite. Skipped entirely (not run on DD). +const HANG: &[&str] = &[ + "eggcc-2mm.egg", + "eqsolve.egg", + "math-microbenchmark.egg", + "rectangle.egg", + "repro-herbie-vanilla.egg", + "resolution.egg", + "until.egg", +]; + +/// A fast, representative slice of the corpus run under debug builds (where +/// differential-dataflow is too slow for the full set). `--release` runs +/// everything. All of these match egglog's snapshot (including `eqsat-basic`, +/// which exercises multi-step congruence closure). +const DEBUG_SUBSET: &[&str] = &[ + "i64.egg", + "bitwise.egg", + "string.egg", + "path.egg", + "birewrite.egg", + "eqsat-basic.egg", +]; + +/// Files egglog snapshots but DD cannot run, because they use a FEATURE this +/// simplified backend does not implement (as opposed to a wrong answer, of which +/// there are none). Listed explicitly so the harness fails if the set changes in +/// either direction. Reasons: +/// +/// - container rebuild read primitives (registered through egglog's +/// `ActionRegistry`; DD has direct container storage but no registry +/// execution state for read primitives over its term-encoded mirror): +/// `container-proofs`, `datatypes`, `nested-container-dirty-propagation`, +/// `repro-querybug3`. +const KNOWN_UNSUPPORTED: &[(&str, &str)] = &[ + ("container-proofs.egg", "requires a backend action registry"), + ("datatypes.egg", "requires a backend action registry"), + ( + "nested-container-dirty-propagation.egg", + "requires a backend action registry", + ), + ("repro-querybug3.egg", "requires a backend action registry"), +]; + +/// Run `program` on DD and render its outputs with the SAME normalization +/// egglog's snapshot used, so the two are directly comparable. +fn run_dd(file: &str, program: &str) -> Result { + let mut eg = + EGraph::with_backend(Box::new(egglog_experimental_dd::EGraph::new())).with_term_encoding(); + match eg.parse_and_run_program(Some(file.to_string()), program) { + Ok(outs) => Ok(CommandOutput::snapshot_stable_under_proof_encoding(&outs) + .trim_end() + .to_string()), + Err(e) => Err(e.to_string()), + } +} + +fn check(path: &PathBuf) -> Result<(), Failed> { + let name = path + .file_name() + .ok_or_else(|| format!("corpus path has no file name: {}", path.display()))? + .to_string_lossy() + .to_string(); + let stem = path + .file_stem() + .ok_or_else(|| format!("corpus path has no file stem: {}", path.display()))? + .to_string_lossy() + .replace(['.', '-', ' '], "_"); + + // egglog only produces a term-encoding-valid shared snapshot for files that + // support proofs (it runs the term-encoding treatment then). Otherwise the + // snapshot is native-only, which a term-encoded backend needn't reproduce. + if !file_supports_proofs(path) { + return Ok(()); + } + let src = std::fs::read_to_string(path).map_err(|e| format!("read {name}: {e}"))?; + let program = format!("{src}\n(print-size)"); + let file = path.display().to_string(); + let dd = catch_unwind(AssertUnwindSafe(|| run_dd(&file, &program))); + + let expected_unsupported = KNOWN_UNSUPPORTED + .iter() + .find_map(|&(file, expected)| (file == name).then_some(expected)); + match dd { + Ok(Err(err)) => match expected_unsupported { + Some(expected) if err.contains(expected) => Ok(()), + Some(expected) => Err(format!( + "{name}: DD returned the wrong controlled error for a known unsupported feature\n\ + expected diagnostic containing: {expected}\n\ + actual error: {err}" + ) + .into()), + None => Err(format!( + "{name}: DD failed but is not in KNOWN_UNSUPPORTED — \ + add it with an expected controlled diagnostic or investigate the regression\n\ + error: {err}" + ) + .into()), + }, + Err(panic) => { + let panic = panic + .downcast_ref::() + .map(String::as_str) + .or_else(|| panic.downcast_ref::<&str>().copied()) + .unwrap_or(""); + Err(format!( + "{name}: DD panicked; unsupported cases must return their documented controlled error\n\ + panic: {panic}" + ) + .into()) + } + Ok(Ok(_)) if expected_unsupported.is_some() => Err(format!( + "{name} is listed in KNOWN_UNSUPPORTED but now runs — remove it from that list" + ) + .into()), + Ok(Ok(out)) => { + let mut settings = insta::Settings::clone_current(); + settings.set_snapshot_path("../../../egglog/tests/snapshots"); + settings.bind(|| { + insta::assert_snapshot!(format!("shared_snapshot_{stem}"), out); + }); + Ok(()) + } + } +} + +fn corpus_from_dirs(dirs: &[&str]) -> Result, String> { + let mut v = vec![]; + for dir in dirs { + let rd = + std::fs::read_dir(dir).map_err(|error| format!("read corpus dir {dir}: {error}"))?; + for entry in rd { + let entry = + entry.map_err(|error| format!("read entry in corpus dir {dir}: {error}"))?; + let p = entry.path(); + if p.extension().and_then(|s| s.to_str()) == Some("egg") { + let name = p + .file_name() + .ok_or_else(|| format!("corpus path has no file name: {}", p.display()))? + .to_string_lossy() + .to_string(); + let debug_ok = !cfg!(debug_assertions) || DEBUG_SUBSET.contains(&name.as_str()); + if !HANG.contains(&name.as_str()) && debug_ok { + v.push(p); + } + } + } + } + v.sort(); + if v.is_empty() { + return Err("DD corpus selection is empty".to_string()); + } + Ok(v) +} + +fn corpus() -> Result, String> { + corpus_from_dirs(&["../../egglog/tests", "../../egglog/tests/web-demo"]) +} + +fn main() { + let args = Arguments::from_args(); + let mut trials = vec![ + Trial::test("missing_corpus_directory_is_an_error", || { + corpus_from_dirs(&["__egglog_dd_missing_corpus_directory__"]) + .expect_err("missing corpus directory must fail discovery"); + Ok(()) + }), + Trial::test("empty_corpus_selection_is_an_error", || { + corpus_from_dirs(&[]).expect_err("empty corpus selection must fail discovery"); + Ok(()) + }), + ]; + match corpus() { + Ok(paths) => trials.extend(paths.into_iter().map(|path| { + let name = path + .file_stem() + .expect("validated corpus path must have a file stem") + .to_string_lossy() + .replace(['-', '.'], "_"); + Trial::test(name, move || check(&path)) + })), + Err(error) => trials.push(Trial::test("corpus_discovery", move || Err(error.into()))), + } + libtest_mimic::run(&args, trials).exit(); +} diff --git a/egglog-experimental/dd/tests/rewrite_join.rs b/egglog-experimental/dd/tests/rewrite_join.rs new file mode 100644 index 0000000..26185bb --- /dev/null +++ b/egglog-experimental/dd/tests/rewrite_join.rs @@ -0,0 +1,45 @@ +// Confirms that a real rewrite rule fires through the Differential Dataflow +// join: without the body producing a match and the head applying it (then +// congruence closing), the final `(check ...)` would fail. +use egglog::EGraph; + +fn dd_egraph() -> EGraph { + EGraph::with_backend(Box::new(egglog_experimental_dd::EGraph::new())).with_term_encoding() +} + +#[test] +fn commutativity_fires_through_dd_join() { + let mut eg = dd_egraph(); + eg.parse_and_run_program( + None, + r#" + (datatype Math (Num i64) (Add Math Math)) + (rewrite (Add a b) (Add b a)) + (let t (Add (Num 1) (Num 2))) + (run 3) + ; only holds if (Add a b) matched on the DD join, the head built + ; (Add b a), and congruence unified the two: + (check (= (Add (Num 1) (Num 2)) (Add (Num 2) (Num 1)))) + "#, + ) + .expect("commutativity rewrite should fire and unify via the DD join"); +} + +#[test] +fn multi_atom_join_associativity() { + // A two-atom body join ((Add a b) matched against the rewrite LHS decomposed + // into view atoms) exercising a wider DD join than a single atom. + let mut eg = dd_egraph(); + eg.parse_and_run_program( + None, + r#" + (datatype Math (Num i64) (Add Math Math) (Mul Math Math)) + (rewrite (Mul a (Add b c)) (Add (Mul a b) (Mul a c))) + (let e (Mul (Num 2) (Add (Num 3) (Num 4)))) + (run 3) + (check (= (Mul (Num 2) (Add (Num 3) (Num 4))) + (Add (Mul (Num 2) (Num 3)) (Mul (Num 2) (Num 4))))) + "#, + ) + .expect("distributivity rewrite should fire via the multi-atom DD join"); +} diff --git a/egglog-experimental/dd/tests/smoke.rs b/egglog-experimental/dd/tests/smoke.rs new file mode 100644 index 0000000..a3b2857 --- /dev/null +++ b/egglog-experimental/dd/tests/smoke.rs @@ -0,0 +1,218 @@ +use std::panic::{catch_unwind, AssertUnwindSafe}; +use std::sync::Arc; + +use egglog::ast::Span; +use egglog::constraint::{SimpleTypeConstraint, TypeConstraint}; +use egglog::scheduler::{Matches, Scheduler}; +use egglog::sort::I64Sort; +use egglog::{prelude::*, EGraph, Primitive, PrimitiveValidator, Value, WritePrim, WriteState}; + +#[derive(Clone)] +struct ChooseAllScheduler; + +impl Scheduler for ChooseAllScheduler { + fn filter_matches(&mut self, _rule: &str, _ruleset: &str, matches: &mut Matches) -> bool { + matches.choose_all(); + true + } +} + +#[derive(Clone)] +struct RegistryWrite; + +impl Primitive for RegistryWrite { + fn name(&self) -> &str { + "registry-write" + } + + fn get_type_constraints(&self, span: &Span) -> Box { + SimpleTypeConstraint::new( + self.name(), + vec![I64Sort.to_arcsort(), I64Sort.to_arcsort()], + span.clone(), + ) + .into_box() + } +} + +impl WritePrim for RegistryWrite { + fn apply<'a, 'db>(&self, _state: WriteState<'a, 'db>, args: &[Value]) -> Option { + Some(args[0]) + } +} + +#[test] +fn dd_runs_basic_egg() { + let backend = Box::new(egglog_experimental_dd::EGraph::new()); + let mut eg = EGraph::with_backend(backend).with_term_encoding(); + eg.parse_and_run_program( + None, + "(datatype Math (Num i64) (Add Math Math))\n(Add (Num 1) (Num 2))\n(run 1)\n(print-size Add)", + ) + .unwrap(); +} + +#[test] +fn dd_runs_proof_mode_pair_container_side_condition() { + let backend = Box::new(egglog_experimental_dd::EGraph::new()); + let mut eg = egglog_experimental::new_experimental_egraph_with_backend_and_proofs(backend); + eg.parse_and_run_program( + None, + r#" + (datatype Expr (A)) + (sort Cost (Pair Expr i64)) + (relation Seed (Expr)) + (relation Seen ()) + + (Seed (A)) + + (rule ((Seed e) + (= c (pair e 1))) + ((Seen)) + :name "pair-side-condition") + + (run 1) + (prove (Seen)) + "#, + ) + .unwrap(); +} + +/// DD has no native union-find, so it must be paired with term encoding. +/// Without it, the frontend refuses to run rather than silently drop `union`s. +#[test] +fn dd_without_term_encoding_errors() { + let backend = Box::new(egglog_experimental_dd::EGraph::new()); + let mut eg = EGraph::with_backend(backend); // no `.with_term_encoding()` + let err = eg + .parse_and_run_program(None, "(datatype Math (Num i64))\n(run 1)") + .unwrap_err(); + assert!( + err.to_string().contains("term encoding"), + "expected a term-encoding-required error, got: {err}" + ); +} + +#[test] +fn dd_custom_scheduler_returns_a_backend_capability_error() { + let backend = Box::new(egglog_experimental_dd::EGraph::new()); + let mut eg = EGraph::with_backend(backend).with_term_encoding(); + eg.parse_and_run_program( + None, + r#" + (ruleset scheduled) + (relation Input (i64)) + (relation Output (i64)) + (rule ((Input x)) ((Output x)) :ruleset scheduled) + "#, + ) + .unwrap(); + let scheduler = eg.add_scheduler(Box::new(ChooseAllScheduler)); + + let error = eg + .step_rules_with_scheduler(scheduler, "scheduled") + .expect_err("DD cannot instantiate custom-scheduler matches through the bridge"); + + assert!( + error.to_string().contains("reference bridge backend"), + "unexpected scheduler capability error: {error}" + ); +} + +#[test] +fn registry_primitive_returns_run_rules_error_without_unwinding() { + let backend = Box::new(egglog_experimental_dd::EGraph::new()); + let mut eg = EGraph::with_backend(backend).with_term_encoding(); + let validator: PrimitiveValidator = Arc::new(|_, args| args.first().copied()); + eg.add_write_primitive(RegistryWrite, Some(validator)); + + let outcome = catch_unwind(AssertUnwindSafe(|| { + eg.parse_and_run_program( + None, + r#" + (relation Input (i64)) + (relation Output (i64)) + (Input 1) + (rule ((Input x)) + ((let y (registry-write x)) + (Output y))) + (run 1) + "#, + ) + })); + let error = outcome + .expect("registry-backed primitive failure must not unwind") + .unwrap_err(); + assert!( + error + .to_string() + .contains("requires a backend action registry"), + "unexpected registry capability error: {error}" + ); +} + +#[test] +fn dd_frontend_push_pop_rebuilds_cloned_join_state() { + let backend = Box::new(egglog_experimental_dd::EGraph::new()); + let mut eg = EGraph::with_backend(backend).with_term_encoding(); + + eg.parse_and_run_program( + None, + r#" + (relation Input (i64)) + (relation Output (i64)) + (rule ((Input x)) ((Output x))) + + (Input 1) + (run 1) + (check (Output 1)) + + (push) + (Input 2) + (run 1) + (check (Output 2)) + (pop) + + (check (Output 1)) + (fail (check (Output 2))) + (Input 3) + (run 1) + (check (Output 3)) + "#, + ) + .expect("DD push/pop must restore state and rebuild the transient join"); +} + +#[test] +fn backend_generic_reads_work_without_an_action_registry() { + let backend = Box::new(egglog_experimental_dd::EGraph::new()); + let mut eg = EGraph::with_backend(backend).with_term_encoding(); + eg.parse_and_run_program( + None, + r#" + (relation R (i64)) + (R 3) + "#, + ) + .unwrap(); + + let mut relation_entry = None; + eg.constructor_enodes("R", |enode| { + relation_entry = Some(eg.value_to_base::(enode.children[0])); + }) + .unwrap(); + assert_eq!(relation_entry, Some(3)); + assert!(eg.function_entries("R", |_| {}).is_err()); + + let mut called = false; + let error = eg + .update(|_| { + called = true; + Ok(()) + }) + .unwrap_err(); + assert!(!called); + assert!(error + .to_string() + .contains("EGraph::update requires a backend action registry")); +} diff --git a/egglog-experimental/src/lib.rs b/egglog-experimental/src/lib.rs index 0bc6cde..7b46aff 100644 --- a/egglog-experimental/src/lib.rs +++ b/egglog-experimental/src/lib.rs @@ -89,6 +89,17 @@ pub fn new_experimental_egraph_with_proofs() -> EGraph { new_experimental_egraph_for_proofs().with_proofs_enabled() } +pub fn new_experimental_egraph_with_backend_for_proofs(backend: Box) -> EGraph { + let mut egraph = EGraph::with_backend(backend); + add_experimental_extensions(&mut egraph, false); + let typechecker = new_experimental_egraph_for_proofs(); + egraph.with_term_encoding_typechecker(typechecker) +} + +pub fn new_experimental_egraph_with_backend_and_proofs(backend: Box) -> EGraph { + new_experimental_egraph_with_backend_for_proofs(backend).with_proofs_enabled() +} + pub fn new_experimental_egraph_with_proof_testing() -> EGraph { new_experimental_egraph_with_proofs().with_proof_testing() } diff --git a/egglog-experimental/src/main.rs b/egglog-experimental/src/main.rs index 9839e7b..8269e4e 100644 --- a/egglog-experimental/src/main.rs +++ b/egglog-experimental/src/main.rs @@ -1,14 +1,116 @@ +use std::ffi::OsString; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum Backend { + Main, + #[cfg(feature = "dd-backend")] + Dd, +} + fn main() { - let proof_mode = std::env::args().any(|arg| { + let (backend, args) = extract_backend_arg(std::env::args_os()).unwrap_or_else(|err| { + eprintln!("error: {err}"); + std::process::exit(2); + }); + let proof_mode = args.iter().any(|arg| { matches!( - arg.as_str(), - "--proofs" | "--proof-testing" | "--term-encoding" + arg.to_str(), + Some("--proofs" | "--proof-testing" | "--term-encoding") ) }); - let egraph = if proof_mode { - egglog_experimental::new_experimental_egraph_for_proofs() + #[cfg(feature = "dd-backend")] + let args = if matches!(backend, Backend::Dd) { + args.into_iter() + .filter(|arg| arg.to_str() != Some("--term-encoding")) + .collect() } else { - egglog_experimental::new_experimental_egraph() + args + }; + let egraph = match backend { + Backend::Main if proof_mode => egglog_experimental::new_experimental_egraph_for_proofs(), + Backend::Main => egglog_experimental::new_experimental_egraph(), + #[cfg(feature = "dd-backend")] + Backend::Dd => egglog_experimental::new_experimental_egraph_with_backend_for_proofs( + Box::new(egglog_experimental_dd::EGraph::new()), + ), }; - egglog::cli(egraph) + egglog::cli_with_args(egraph, args) +} + +fn extract_backend_arg(args: I) -> Result<(Backend, Vec), String> +where + I: IntoIterator, +{ + let mut backend = Backend::Main; + let mut saw_backend = false; + let mut filtered = Vec::new(); + let mut iter = args.into_iter(); + while let Some(arg) = iter.next() { + match arg.to_str() { + Some("--backend") => { + if saw_backend { + return Err("--backend may only be passed once".to_string()); + } + let Some(value) = iter.next() else { + return Err("--backend requires one of: main, dd".to_string()); + }; + backend = parse_backend(value.to_str())?; + saw_backend = true; + } + Some(value) if value.starts_with("--backend=") => { + if saw_backend { + return Err("--backend may only be passed once".to_string()); + } + backend = parse_backend(value.strip_prefix("--backend="))?; + saw_backend = true; + } + _ => filtered.push(arg), + } + } + Ok((backend, filtered)) +} + +fn parse_backend(value: Option<&str>) -> Result { + match value { + Some("main") => Ok(Backend::Main), + #[cfg(feature = "dd-backend")] + Some("dd") => Ok(Backend::Dd), + #[cfg(not(feature = "dd-backend"))] + Some("dd") => Err( + "backend \"dd\" requires building egglog-experimental with --features dd-backend" + .to_string(), + ), + Some(other) => Err(format!( + "unknown backend {other:?}; expected one of: main, dd" + )), + None => Err("backend value must be valid UTF-8".to_string()), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_main_backend() { + assert_eq!(parse_backend(Some("main")), Ok(Backend::Main)); + } + + #[cfg(not(feature = "dd-backend"))] + #[test] + fn explains_how_to_enable_dd_backend() { + assert_eq!( + parse_backend(Some("dd")), + Err( + "backend \"dd\" requires building egglog-experimental with --features dd-backend" + .to_string() + ) + ); + } + + #[cfg(feature = "dd-backend")] + #[test] + fn parses_enabled_dd_backend() { + assert_eq!(parse_backend(Some("dd")), Ok(Backend::Dd)); + } } diff --git a/egglog/CHANGELOG.md b/egglog/CHANGELOG.md index dff0897..2860488 100644 --- a/egglog/CHANGELOG.md +++ b/egglog/CHANGELOG.md @@ -2,6 +2,8 @@ ## [Unreleased] - ReleaseDate +- Fix user-defined primitives (registered through the Rust API after construction) being reported as unbound under term encoding / proofs: primitive registration now also reaches the term-encoding typechecker, so the encoder can typecheck the encoded program. Previously callers had to manually register the primitive on `proof_state.original_typechecking` as well. +- **Pluggable backend SPI.** `EGraph::with_backend(Box)` lets a third party drive the egglog frontend with their own backend (see the `egglog-backend-trait` crate and the `egglog-experimental-dd` example). A backend without a native union-find declares `Backend::requires_term_encoding()`; `EGraph::with_term_encoding()` opts such an e-graph into the term-encoding pipeline (congruence and rebuild lower to rules over `@uf` tables), and running a term-encoding-only backend without it now errors with `Error::BackendRequiresTermEncoding` instead of silently dropping `union`s. - Add `make nightly` and `scripts/nightly_bench.py`, a hyperfine-based benchmark harness that measures every `tests/**/*.egg` program at 1/2/4/8 threads and (where supported) in proof-testing mode, caps each run at a 2-minute timeout, skips sub-50ms programs, and emits an HTML dashboard (one row per benchmark, one column per configuration) for nightly.cs.washington.edu. The dashboard uses [eval-live](https://github.com/oflatt/eval-live) for interactive filtering and sorting. - Add typed `EGraph` extension state that clones with `EGraph` and is restored by `push`/`pop`. - Fix custom scheduler queries so subsumed rows are not offered as fresh matches. diff --git a/egglog/Cargo.toml b/egglog/Cargo.toml index 98eae30..23db9fd 100644 --- a/egglog/Cargo.toml +++ b/egglog/Cargo.toml @@ -68,6 +68,7 @@ egglog-bridge = { workspace = true } egglog-numeric-id = { workspace = true } egglog-add-primitive = { workspace = true } egglog-reports = { workspace = true } +egglog-backend-trait = { workspace = true } rayon = { workspace = true } serde_json = { workspace = true } diff --git a/egglog/core-relations/src/dependency_graph.rs b/egglog/core-relations/src/dependency_graph.rs index 3620d02..424e9bc 100644 --- a/egglog/core-relations/src/dependency_graph.rs +++ b/egglog/core-relations/src/dependency_graph.rs @@ -52,6 +52,15 @@ impl DependencyGraph { self.levels.get_or_default(level).insert(table); } + pub(crate) fn remove_last_table(&mut self, table: TableId) { + let level = self + .to_level + .take(table) + .expect("removed table must be registered in the dependency graph"); + self.levels[level].shift_remove(&table); + self.write_deps.take(table); + } + pub(crate) fn strata(&self) -> impl Iterator> { self.levels.iter().map(|(_, tables)| tables) } diff --git a/egglog/core-relations/src/free_join/mod.rs b/egglog/core-relations/src/free_join/mod.rs index d5f055d..579d6b3 100644 --- a/egglog/core-relations/src/free_join/mod.rs +++ b/egglog/core-relations/src/free_join/mod.rs @@ -827,6 +827,24 @@ impl Database { self.total_size_estimate = self.total_size_estimate.wrapping_sub(prev_len); } + /// Remove the most recently added table. + /// + /// This is intended for rolling back a failed higher-level registration. + /// It returns `false` without changing the database if `table` is not the + /// final table id. + pub fn remove_last_table(&mut self, table: TableId) -> bool { + let Some(info) = self.tables.pop_last(table) else { + return false; + }; + let notified = self.notification_list.reset(); + for other in notified.into_iter().filter(|other| *other != table) { + self.notification_list.notify(other); + } + self.deps.remove_last_table(table); + self.total_size_estimate = self.total_size_estimate.wrapping_sub(info.table.len()); + true + } + pub(crate) fn plan_query(&mut self, query: Query) -> Plan { plan::plan_query(query, ColumnCardEst::new(self)) } diff --git a/egglog/egglog-ast/src/core.rs b/egglog/egglog-ast/src/core.rs new file mode 100644 index 0000000..216ce48 --- /dev/null +++ b/egglog/egglog-ast/src/core.rs @@ -0,0 +1,258 @@ +use std::{ + fmt::Display, + hash::{Hash, Hasher}, + ops::AddAssign, +}; + +use crate::{ + generic_ast::{Change, GenericExpr, Literal}, + span::Span, + util::ListDisplay, +}; + +/// A variable, literal, or global in lowered rule syntax. +#[derive(Debug, Clone)] +pub enum GenericAtomTerm { + Var(Span, Leaf), + Literal(Span, Constant), + Global(Span, Leaf), +} + +// Source annotations do not affect equality or hashing. +impl PartialEq for GenericAtomTerm { + fn eq(&self, other: &Self) -> bool { + match (self, other) { + (Self::Var(_, lhs), Self::Var(_, rhs)) => lhs == rhs, + (Self::Literal(_, lhs), Self::Literal(_, rhs)) => lhs == rhs, + (Self::Global(_, lhs), Self::Global(_, rhs)) => lhs == rhs, + _ => false, + } + } +} + +impl Eq for GenericAtomTerm {} + +impl Hash for GenericAtomTerm { + fn hash(&self, state: &mut H) { + match self { + Self::Var(_, value) | Self::Global(_, value) => value.hash(state), + Self::Literal(_, value) => value.hash(state), + } + } +} + +impl GenericAtomTerm { + pub fn span(&self) -> &Span { + match self { + Self::Var(span, _) | Self::Literal(span, _) | Self::Global(span, _) => span, + } + } +} + +impl GenericAtomTerm { + pub fn to_expr(&self) -> GenericExpr { + match self { + Self::Var(span, value) | Self::Global(span, value) => { + GenericExpr::Var(span.clone(), value.clone()) + } + Self::Literal(span, literal) => GenericExpr::Lit(span.clone(), literal.clone()), + } + } +} + +impl Display for GenericAtomTerm { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Var(_, value) | Self::Global(_, value) => write!(f, "{value}"), + Self::Literal(_, literal) => write!(f, "{literal}"), + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct GenericAtom { + pub span: Span, + pub head: Head, + pub args: Vec>, +} + +impl Display for GenericAtom { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "({} {}) ", self.head, ListDisplay(&self.args, " ")) + } +} + +impl GenericAtom +where + Leaf: Clone, +{ + pub fn vars(&self) -> impl Iterator + '_ { + self.args.iter().filter_map(|term| match term { + GenericAtomTerm::Var(_, value) => Some(value.clone()), + GenericAtomTerm::Literal(..) | GenericAtomTerm::Global(..) => None, + }) + } + + pub fn substitute_with( + &mut self, + substitute: &mut impl FnMut(&Leaf) -> Option>, + ) where + Constant: Clone, + { + for argument in &mut self.args { + if let GenericAtomTerm::Var(_, variable) = argument + && let Some(replacement) = substitute(variable) + { + *argument = replacement; + } + } + } +} + +impl GenericAtom { + pub fn to_expr(&self) -> GenericExpr { + let input_count = self + .args + .len() + .checked_sub(1) + .expect("an atom must include an output argument"); + GenericExpr::Call( + self.span.clone(), + self.head.clone(), + self.args[..input_count] + .iter() + .map(GenericAtomTerm::to_expr) + .collect(), + ) + } +} + +#[derive(Debug, Clone)] +pub struct Query { + pub atoms: Vec>, +} + +impl Default for Query { + fn default() -> Self { + Self { atoms: Vec::new() } + } +} + +impl Query { + pub fn vars(&self) -> impl Iterator + '_ { + self.atoms.iter().flat_map(GenericAtom::vars) + } +} + +impl AddAssign for Query { + fn add_assign(&mut self, rhs: Self) { + self.atoms.extend(rhs.atoms); + } +} + +impl Display for Query { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + for atom in &self.atoms { + writeln!(f, "{atom}")?; + } + Ok(()) + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub enum GenericCoreAction { + Let(Span, Leaf, Head, Vec>), + LetAtomTerm(Span, Leaf, GenericAtomTerm), + Set( + Span, + Head, + Vec>, + GenericAtomTerm, + ), + Change(Span, Change, Head, Vec>), + Union( + Span, + GenericAtomTerm, + GenericAtomTerm, + ), + Panic(Span, String), +} + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct GenericCoreActions( + pub Vec>, +); + +impl Default for GenericCoreActions { + fn default() -> Self { + Self(Vec::new()) + } +} + +impl GenericCoreActions { + pub fn new(actions: Vec>) -> Self { + Self(actions) + } + + pub fn free_vars(&self) -> Vec + where + Leaf: Clone + Eq, + { + fn insert_unique(values: &mut Vec, value: T) { + if !values.contains(&value) { + values.push(value); + } + } + + fn add_term( + free: &mut Vec, + term: &GenericAtomTerm, + ) { + if let GenericAtomTerm::Var(_, variable) = term { + insert_unique(free, variable.clone()); + } + } + + let mut free = Vec::new(); + + for action in self.0.iter().rev() { + match action { + GenericCoreAction::Let(_, variable, _, arguments) => { + for argument in arguments { + add_term(&mut free, argument); + } + free.retain(|candidate| candidate != variable); + } + GenericCoreAction::LetAtomTerm(_, variable, term) => { + add_term(&mut free, term); + free.retain(|candidate| candidate != variable); + } + GenericCoreAction::Set(_, _, arguments, value) => { + for argument in arguments { + add_term(&mut free, argument); + } + add_term(&mut free, value); + } + GenericCoreAction::Change(_, _, _, arguments) => { + for argument in arguments { + add_term(&mut free, argument); + } + } + GenericCoreAction::Union(_, lhs, rhs) => { + add_term(&mut free, lhs); + add_term(&mut free, rhs); + } + GenericCoreAction::Panic(..) => {} + } + } + free + } +} + +/// Lowered rule syntax with potentially distinct body and action call types. +#[derive(Debug, Clone)] +pub struct GenericCoreRule { + pub span: Span, + pub body: Query, + pub head: GenericCoreActions, +} diff --git a/egglog/egglog-ast/src/lib.rs b/egglog/egglog-ast/src/lib.rs index 544982b..6caf1a6 100644 --- a/egglog/egglog-ast/src/lib.rs +++ b/egglog/egglog-ast/src/lib.rs @@ -1,3 +1,4 @@ +pub mod core; pub mod generic_ast; pub mod generic_ast_helpers; pub mod span; diff --git a/egglog/egglog-backend-trait/Cargo.toml b/egglog/egglog-backend-trait/Cargo.toml new file mode 100644 index 0000000..36158c9 --- /dev/null +++ b/egglog/egglog-backend-trait/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "egglog-backend-trait" +version.workspace = true +edition.workspace = true +description = "A backend-agnostic interface (the `Backend` trait) for driving an egglog e-graph, so third parties can plug in their own execution engine." +repository.workspace = true +keywords.workspace = true +license.workspace = true + +[dependencies] +anyhow = { workspace = true } +egglog-bridge = { workspace = true } +egglog-ast = { workspace = true } +egglog-core-relations = { workspace = true } +egglog-reports = { workspace = true } diff --git a/egglog/egglog-backend-trait/src/backend_impl.rs b/egglog/egglog-backend-trait/src/backend_impl.rs new file mode 100644 index 0000000..a8ce775 --- /dev/null +++ b/egglog/egglog-backend-trait/src/backend_impl.rs @@ -0,0 +1,244 @@ +//! `impl Backend for egglog_bridge::EGraph` — the in-memory reference backend. +//! +//! Every method is a thin passthrough to an inherent method on the bridge +//! `EGraph`. Lives in this crate (not in `egglog-bridge`) so the bridge stays +//! free of any dependency on the trait; the orphan rule permits it because the +//! [`Backend`] trait is local here. + +use std::any::Any; +use std::collections::BTreeMap; +use std::sync::{Arc, RwLock}; + +use anyhow::{Result, bail}; +use egglog_bridge::{ActionRegistry, EGraph, QueryEntry, RuleBuilder}; + +use egglog_ast::core::{GenericAtomTerm, GenericCoreAction}; + +use crate::{ + Backend, BaseValues, ColumnTy, ContainerValues, ExecutionState, ExternalFunction, + ExternalFunctionId, FunctionConfig, FunctionId, IterationReport, ReportLevel, RuleActionCall, + RuleBodyCall, RuleId, RuleSetRun, RuleSpec, RuleValue, RuleVar, ScanEntry, Value, +}; + +fn rule_entry( + builder: &mut RuleBuilder<'_>, + variables: &mut BTreeMap, + term: &GenericAtomTerm, +) -> Result { + match term { + GenericAtomTerm::Var(_, variable) => Ok(variables + .entry(variable.id) + .or_insert_with(|| builder.new_var_named(variable.ty, &variable.name)) + .clone()), + GenericAtomTerm::Literal(_, constant) => Ok(QueryEntry::Const { + val: constant.value, + ty: constant.ty, + }), + GenericAtomTerm::Global(..) => bail!("globals must be desugared before backend lowering"), + } +} + +fn rule_entries( + builder: &mut RuleBuilder<'_>, + variables: &mut BTreeMap, + terms: &[GenericAtomTerm], +) -> Result> { + terms + .iter() + .map(|term| rule_entry(builder, variables, term)) + .collect() +} + +fn build_rule(egraph: &mut EGraph, rule: RuleSpec) -> Result { + let RuleSpec { + name, + seminaive, + no_decomp, + core, + } = rule; + let mut builder = egraph.new_rule(&name, seminaive); + builder.set_no_decomp(no_decomp); + let mut variables = BTreeMap::new(); + + for atom in &core.body.atoms { + let entries = rule_entries(&mut builder, &mut variables, &atom.args)?; + match &atom.head { + RuleBodyCall::Table { id, read } => { + builder.query_table(*id, &entries, read.is_subsumed())?; + } + RuleBodyCall::Primitive { id, output, .. } => { + builder.query_prim(*id, &entries, *output)?; + } + } + } + + for action in &core.head.0 { + match action { + GenericCoreAction::Let(span, variable, call, arguments) => { + let entries = rule_entries(&mut builder, &mut variables, arguments)?; + let result: QueryEntry = match call { + RuleActionCall::Table { id, name } => { + let span = span.clone(); + let name = name.clone(); + builder + .lookup(*id, &entries, move || { + format!("{span}: lookup of function {name} failed") + }) + .into() + } + RuleActionCall::Primitive { + id, name, output, .. + } => { + let span = span.clone(); + let name = name.clone(); + builder + .call_external_func(*id, &entries, *output, move || { + format!("{span}: call of primitive {name} failed") + }) + .into() + } + }; + variables.insert(variable.id, result); + } + GenericCoreAction::LetAtomTerm(_, variable, term) => { + let entry = rule_entry(&mut builder, &mut variables, term)?; + variables.insert(variable.id, entry); + } + GenericCoreAction::Set(_, call, arguments, value) => { + let RuleActionCall::Table { id, .. } = call else { + bail!("cannot set a primitive") + }; + let mut entries = rule_entries(&mut builder, &mut variables, arguments)?; + entries.push(rule_entry(&mut builder, &mut variables, value)?); + builder.set(*id, &entries); + } + GenericCoreAction::Change(_, change, call, arguments) => { + let RuleActionCall::Table { id, .. } = call else { + bail!("cannot delete or subsume a primitive") + }; + let entries = rule_entries(&mut builder, &mut variables, arguments)?; + match change { + egglog_ast::generic_ast::Change::Delete => builder.remove(*id, &entries), + egglog_ast::generic_ast::Change::Subsume => builder.subsume(*id, &entries), + } + } + GenericCoreAction::Union(_, lhs, rhs) => { + let lhs = rule_entry(&mut builder, &mut variables, lhs)?; + let rhs = rule_entry(&mut builder, &mut variables, rhs)?; + builder.union(lhs, rhs); + } + GenericCoreAction::Panic(_, message) => builder.panic(message.clone()), + } + } + + Ok(builder.build()) +} + +// --------------------------------------------------------------------------- +// Backend for the bridge EGraph +// --------------------------------------------------------------------------- + +impl Backend for EGraph { + fn add_table(&mut self, config: FunctionConfig) -> FunctionId { + EGraph::add_table(self, config) + } + + fn table_size(&self, table: FunctionId) -> usize { + EGraph::table_size(self, table) + } + + fn clear_table(&mut self, func: FunctionId) { + EGraph::clear_table(self, func); + } + + fn for_each_while_dyn( + &self, + table: FunctionId, + f: &mut dyn for<'r> FnMut(ScanEntry<'r>) -> bool, + ) { + EGraph::for_each_while(self, table, f); + } + + fn get_canon_repr(&self, val: Value, ty: ColumnTy) -> Value { + EGraph::get_canon_repr(self, val, ty) + } + + fn base_values(&self) -> &BaseValues { + EGraph::base_values(self) + } + + fn base_values_mut(&mut self) -> &mut BaseValues { + EGraph::base_values_mut(self) + } + + fn container_values(&self) -> &ContainerValues { + EGraph::container_values(self) + } + + fn lookup_id(&self, func: FunctionId, key: &[Value]) -> Option { + EGraph::lookup_id(self, func, key) + } + + fn with_execution_state_tracked_dyn(&self, f: &mut dyn FnMut(&mut ExecutionState<'_>)) -> bool { + EGraph::with_execution_state_tracked(self, |es| f(es)).1 + } + + fn add_rule(&mut self, rule: RuleSpec) -> Result { + build_rule(self, rule) + } + + fn free_rule(&mut self, id: RuleId) { + EGraph::free_rule(self, id); + } + + fn run_rules(&mut self, run: RuleSetRun<'_>) -> Result { + EGraph::run_rules(self, run.rules) + } + + fn flush_updates(&mut self) -> bool { + EGraph::flush_updates(self) + } + + fn register_external_func( + &mut self, + func: Box, + ) -> ExternalFunctionId { + EGraph::register_external_func(self, func) + } + + fn free_external_func(&mut self, func: ExternalFunctionId) { + EGraph::free_external_func(self, func); + } + + fn new_panic(&mut self, message: String) -> ExternalFunctionId { + EGraph::new_panic(self, message) + } + + fn set_report_level(&mut self, level: ReportLevel) { + EGraph::set_report_level(self, level); + } + + fn dump_debug_info(&self) { + EGraph::dump_debug_info(self); + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } + + fn action_registry(&self) -> Option<&Arc>> { + Some(EGraph::action_registry(self)) + } + + fn supports_containers(&self) -> bool { + true + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } +} diff --git a/egglog/egglog-backend-trait/src/lib.rs b/egglog/egglog-backend-trait/src/lib.rs new file mode 100644 index 0000000..da73371 --- /dev/null +++ b/egglog/egglog-backend-trait/src/lib.rs @@ -0,0 +1,460 @@ +//! # egglog-backend-trait +//! +//! A small, backend-agnostic interface for driving an egglog e-graph. The +//! frontend `egglog::EGraph` holds a `Box` and performs all state +//! access through the [`Backend`] trait, so a third party can plug in their own +//! execution engine (a relational database, a dataflow system, …) without +//! forking the frontend. +//! +//! ## What a backend must provide +//! +//! [`Backend`] is deliberately small. The *required* surface is the handful of +//! operations the frontend actually needs: register functions ([`Backend::add_table`]), +//! register and run rules ([`Backend::add_rule`] / [`Backend::run_rules`]), read +//! tables back ([`Backend::for_each_while_dyn`] / [`Backend::table_size`]), +//! canonicalize ids ([`Backend::get_canon_repr`]), reach the value registries +//! ([`Backend::base_values`] / [`Backend::container_values`]), and register +//! primitives ([`Backend::register_external_func`]). +//! +//! Rules cross the backend boundary as a complete [`RuleSpec`]. This lets each +//! backend lower or retain the shared logical rule without replaying a callback +//! builder API or reconstructing an equivalent IR. +//! +//! ## Rule execution contract +//! +//! A [`Backend::run_rules`] call is one bounded rule-iteration over the +//! backend's state. RHS effects are staged while matches are being evaluated and +//! are flushed into tables according to the same semantics as the reference +//! bridge backend. Rule bodies observe a stable read view for the iteration; a +//! row produced by an RHS `set`, `lookup`, `remove`, or `subsume` does not create +//! new matches for another rule body until a later `run_rules` call, although +//! RHS lookups in the same action stream may observe earlier staged +//! lookup-or-insert predictions. +//! +//! During the flush, deletions are applied before insertions/sets for the same +//! bounded iteration. This is observable for rebuild-style actions that emit +//! `delete old; set new`: if both target the same function key, the set is the +//! row retained after the flush. Subsumption actions are applied after sets, so a +//! row written and subsumed in the same iteration ends the iteration subsumed +//! rather than live. +//! +//! For function tables, `set` observes the current row for the function's input +//! key, if any, and folds conflicting output values through the configured +//! [`MergeFn`]. In each fold, `old` is the currently retained output value and +//! `new` is the incoming staged value. Implementations may choose their own +//! physical representation, but they must not derive `old`/`new` from unordered +//! table iteration; non-monotone merge expressions are already user-visible +//! undefined behavior, but backend-specific container ordering should not add an +//! extra source of divergence. +//! +//! Seminaive freshness is part of the observable execution model. If a row is +//! removed and later reinserted with the same logical columns, later rule +//! iterations must be able to treat that row as newly available, just as the +//! reference backend does with its hidden row timestamp. Backends that feed an +//! incremental engine must therefore track row generations or equivalent event +//! identity, not only end-state set membership. +//! +//! Subsumption is semantically a bit on a row, not a separate table. Body atoms +//! using [`ReadMode::Live`] see only live rows, [`ReadMode::All`] sees live and +//! subsumed rows, and [`ReadMode::Subsumed`] sees only subsumed rows. A subsumed +//! row remains the current row for lookup and merge purposes, so a later lookup +//! or `set` for the same key must find/merge with that row without making it +//! live again unless an explicit backend operation says otherwise. +//! +//! ## Advanced features are optional +//! +//! Capabilities that not every backend can offer — the seminaive-safe +//! [`ActionRegistry`] ([`Backend::action_registry`]) and container sorts — are +//! exposed as methods with **default** bodies. An implementer overrides only +//! what it supports; the frontend gates on the returned capability. +//! +//! ## Ergonomic sugar +//! +//! The object-safe [`Backend`] methods take erased closures (`&mut dyn FnMut(..)`) +//! and carry a `_dyn` suffix. Generic, ergonomic wrappers ([`BackendExt::for_each`], +//! [`BackendExt::with_execution_state`], …) live on the [`BackendExt`] extension +//! trait (a blanket impl over every `Backend`), so call sites read naturally — +//! import `BackendExt` to use them. +//! +//! ## Relationship to `egglog-bridge` +//! +//! This crate depends on `egglog-bridge` and re-exports its vocabulary types +//! ([`FunctionConfig`], [`MergeFn`], [`ColumnTy`], [`Value`], …) so an +//! implementer imports them from here. It also provides the reference +//! `impl Backend for egglog_bridge::EGraph` (see `backend_impl`). + +use std::any::{Any, TypeId}; +use std::sync::{Arc, RwLock}; + +use anyhow::Result; + +mod backend_impl; + +// --------------------------------------------------------------------------- +// Re-exports: the shared vocabulary a backend implementer works with. +// --------------------------------------------------------------------------- + +pub use egglog_bridge::{ + ActionRegistry, ColumnTy, DefaultVal, FunctionConfig, FunctionId, MergeFn, RuleId, ScanEntry, +}; +pub use egglog_core_relations::{ + BaseValue, BaseValueId, BaseValues, ContainerValue, ContainerValues, CounterId, ExecutionState, + ExternalFunction, ExternalFunctionId, Value, +}; +pub use egglog_reports::{IterationReport, ReportLevel}; + +/// Which subsumption view a table atom reads. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub enum ReadMode { + Live, + Subsumed, + All, +} + +impl ReadMode { + pub(crate) fn is_subsumed(self) -> Option { + match self { + Self::Live => Some(false), + Self::Subsumed => Some(true), + Self::All => None, + } + } +} + +/// A typed variable in a backend rule. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct RuleVar { + pub id: u32, + pub name: Box, + pub ty: ColumnTy, +} + +/// A typed constant in a backend rule. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct RuleValue { + pub value: Value, + pub ty: ColumnTy, +} + +/// A call in a rule body. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub enum RuleBodyCall { + Table { + id: FunctionId, + read: ReadMode, + }, + Primitive { + id: ExternalFunctionId, + name: Box, + output: ColumnTy, + }, +} + +/// A call in a rule action stream. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub enum RuleActionCall { + Table { + id: FunctionId, + name: Box, + }, + Primitive { + id: ExternalFunctionId, + name: Box, + output: ColumnTy, + }, +} + +pub type BackendCoreRule = + egglog_ast::core::GenericCoreRule; + +/// A complete logical rule supplied to a backend. +#[derive(Clone, Debug)] +pub struct RuleSpec { + pub name: String, + pub seminaive: bool, + pub no_decomp: bool, + pub core: BackendCoreRule, +} + +/// One bounded invocation of a logical ruleset. +#[derive(Clone, Copy, Debug)] +pub struct RuleSetRun<'a> { + pub name: Option<&'a str>, + pub rules: &'a [RuleId], +} + +/// Backend-selected policy for reconciling two ids that rebuild to the same +/// container value. The container type is supplied separately to +/// [`Backend::container_merge_fn`], so a backend may choose a different policy +/// for each registered [`ContainerValue`] type. +pub type ContainerMergeFn = + Arc Fn(&mut ExecutionState<'a>, Value, Value) -> Value + Send + Sync>; + +// --------------------------------------------------------------------------- +// The `Backend` trait +// --------------------------------------------------------------------------- + +/// A backend that stores an egglog e-graph and runs rules over it. +/// +/// The frontend holds a `Box` and drives everything through this +/// trait. Implementations: the in-memory reference `egglog_bridge::EGraph` +/// (this crate's `backend_impl`), and separate-crate engines such as the +/// Differential Dataflow backend in `egglog-experimental/dd`. +/// +/// See the crate docs for which methods are required vs. optional and how the +/// ergonomic sugar on `dyn Backend` relates to the `_dyn` methods here. +pub trait Backend: Send + Sync { + // -- table lifecycle ---------------------------------------------------- + + /// Register a function / relation / constructor and return its handle. + fn add_table(&mut self, config: FunctionConfig) -> FunctionId; + + /// Number of rows currently in the given function's table. + fn table_size(&self, table: FunctionId) -> usize; + + /// Remove every row from the given function's table. + fn clear_table(&mut self, func: FunctionId); + + // -- iteration (object-safe; see `for_each` / `for_each_while` sugar) ---- + + /// Iterate rows in `table`, stopping early when `f` returns `false`. + /// Object-safe form; prefer the [`BackendExt::for_each_while`] sugar. + fn for_each_while_dyn( + &self, + table: FunctionId, + f: &mut dyn for<'r> FnMut(ScanEntry<'r>) -> bool, + ); + + // -- direct access ------------------------------------------------------ + + /// Canonical representative of `val` under `ty` (union-find find for + /// [`ColumnTy::Id`]; identity for a base value). + fn get_canon_repr(&self, val: Value, ty: ColumnTy) -> Value; + + /// The backend's base-value (primitive) registry. + fn base_values(&self) -> &BaseValues; + + /// Mutable base-value registry (used to register new base types). + fn base_values_mut(&mut self) -> &mut BaseValues; + + /// The backend's container-value registry. + fn container_values(&self) -> &ContainerValues; + + /// Mutable access to the backend's container-value registry, for backends + /// that can execute container primitives without the bridge. + fn container_values_mut_dyn(&mut self) -> Option<&mut ContainerValues> { + None + } + + /// Read the output value for `key` from a function table without inserting + /// a missing row. Subsumed rows remain lookup-visible. + fn lookup_id(&self, func: FunctionId, key: &[Value]) -> Option; + + /// Allocate a fresh counter for container-value ids, for backends that can + /// execute container primitives without the bridge. + fn new_container_id_counter(&mut self) -> Option { + None + } + + /// Select the merge policy for a registered container type. Backends that + /// advertise container support outside the reference bridge must provide + /// this alongside a container registry and id counter. + fn container_merge_fn(&self, _container_type: TypeId) -> Option { + None + } + + // -- execution state (object-safe; see `with_execution_state` sugar) ----- + + /// Run `f` against a fresh execution state and return whether it staged any + /// mutation. Object-safe form; prefer + /// [`BackendExt::with_execution_state_tracked`]. + fn with_execution_state_tracked_dyn(&self, f: &mut dyn FnMut(&mut ExecutionState<'_>)) -> bool; + + // -- rule management ---------------------------------------------------- + + /// Register a complete logical rule. + fn add_rule(&mut self, rule: RuleSpec) -> Result; + + /// Drop a registered rule. + fn free_rule(&mut self, id: RuleId); + + /// Run one bounded iteration of the given rule set. + /// + /// Implementations should evaluate rule bodies against a stable read view + /// for this iteration and stage RHS effects until the rule firing is + /// applied. The externally visible behavior must match the rule execution + /// contract in the crate docs, especially for function `set`/`merge`, + /// `remove`, `subsume`, and constructor lookup-or-insert effects. + fn run_rules(&mut self, run: RuleSetRun<'_>) -> Result; + + /// Drain staged inserts and rebuild if the union-find changed. Returns + /// whether the database changed. + fn flush_updates(&mut self) -> bool; + + // -- primitives --------------------------------------------------------- + + /// Register a user-defined primitive (`ExternalFunction`). + fn register_external_func( + &mut self, + func: Box, + ) -> ExternalFunctionId; + + /// Drop a user-defined primitive. + fn free_external_func(&mut self, func: ExternalFunctionId); + + /// Register a deferred-panic primitive and return its id. + /// + /// The registered function must ignore any arguments supplied by its call + /// site, trigger the backend's normal early-stop path, and make the next + /// [`Backend::run_rules`] return the provided message as an error. + fn new_panic(&mut self, message: String) -> ExternalFunctionId; + + // -- diagnostics -------------------------------------------------------- + + /// Set the verbosity of the per-iteration timing report. + fn set_report_level(&mut self, level: ReportLevel); + + /// Dump the database state to the log channel (debug only). + fn dump_debug_info(&self); + + // -- cloning ------------------------------------------------------------ + + /// Deep-clone the backend (used for push/pop snapshots). + fn clone_boxed(&self) -> Box; + + // -- optional / advanced (default-provided) ----------------------------- + + /// The seminaive-safe [`ActionRegistry`] used by registry-backed + /// primitives, when this backend provides one. + /// + /// Backends without an in-memory action registry return `None`; the + /// frontend routes unsupported primitive calls through + /// [`Backend::new_panic`] so execution fails through the backend's normal + /// error channel. + fn action_registry(&self) -> Option<&Arc>> { + None + } + + /// Whether this backend supports `Vec` / `Set` / `Map` / `MultiSet` + /// container sorts. + fn supports_containers(&self) -> bool { + false + } + + /// Whether this backend needs the frontend's term-encoding pipeline to be + /// enabled. A backend without a native union-find, such as the experimental + /// Differential Dataflow backend, relies on term encoding to lower + /// congruence and rebuild to ordinary rules over `@uf` tables. Native mode + /// has no correct representation for those unions, so the frontend refuses + /// to run such a backend unless the e-graph was built with term encoding + /// (via `EGraph::with_term_encoding`). + fn requires_term_encoding(&self) -> bool { + false + } + + // -- concrete-backend access (escape hatch) ----------------------------- + + /// `&self` as `&dyn Any`, to downcast to the concrete backend type. Used + /// by the container-registration sugar to reach the reference bridge. + /// Implementations return `self`. + fn as_any(&self) -> &dyn Any; + + /// Mutable counterpart of [`Backend::as_any`]. Implementations return + /// `self`. + fn as_any_mut(&mut self) -> &mut dyn Any; +} + +impl Clone for Box { + fn clone(&self) -> Self { + self.clone_boxed() + } +} + +// --------------------------------------------------------------------------- +// Ergonomic sugar: `BackendExt` +// --------------------------------------------------------------------------- + +/// Ergonomic, generic wrappers over the object-safe [`Backend`] surface. +/// +/// Provided by a blanket impl for every `B: Backend + ?Sized`, so the sugar is +/// available on `dyn Backend`, `Box`, and any concrete backend. +/// Import this trait to call `for_each` / `with_execution_state` / +/// `register_container_ty` on a backend. +pub trait BackendExt { + /// Call `f` on every row in `table`. + fn for_each(&self, table: FunctionId, f: impl for<'r> FnMut(ScanEntry<'r>)); + + /// Iterate rows in `table`, stopping when `f` returns `false`. + fn for_each_while(&self, table: FunctionId, f: impl for<'r> FnMut(ScanEntry<'r>) -> bool); + + /// Run `f` against a fresh execution state and return its result. + fn with_execution_state(&self, f: impl FnOnce(&mut ExecutionState<'_>) -> R) -> R; + + /// Run `f` against a fresh execution state, also reporting whether `f` + /// staged any mutation. + fn with_execution_state_tracked( + &self, + f: impl FnOnce(&mut ExecutionState<'_>) -> R, + ) -> (R, bool); + + /// Register a container-value type `C`. + /// + /// Container registration wires the container into the backend's rebuild + /// loop, which requires backend-internal state. The reference bridge handles + /// this directly; another container-capable backend supplies its registry, + /// id counter, and per-container merge policy through the capability methods + /// on [`Backend`]. Registration is a no-op for backends that do not advertise + /// container support. + fn register_container_ty(&mut self); +} + +impl BackendExt for B { + fn for_each(&self, table: FunctionId, mut f: impl for<'r> FnMut(ScanEntry<'r>)) { + self.for_each_while_dyn(table, &mut |row| { + f(row); + true + }); + } + + fn for_each_while(&self, table: FunctionId, mut f: impl for<'r> FnMut(ScanEntry<'r>) -> bool) { + self.for_each_while_dyn(table, &mut f); + } + + fn with_execution_state(&self, f: impl FnOnce(&mut ExecutionState<'_>) -> R) -> R { + self.with_execution_state_tracked(f).0 + } + + fn with_execution_state_tracked( + &self, + f: impl FnOnce(&mut ExecutionState<'_>) -> R, + ) -> (R, bool) { + let mut f = Some(f); + let mut out: Option = None; + let mutated = self.with_execution_state_tracked_dyn(&mut |es| { + let f = f + .take() + .expect("with_execution_state_tracked closure called once"); + out = Some(f(es)); + }); + ( + out.expect("with_execution_state_tracked_dyn must invoke its closure exactly once"), + mutated, + ) + } + + fn register_container_ty(&mut self) { + if let Some(bridge) = self.as_any_mut().downcast_mut::() { + bridge.register_container_ty::(); + } else if self.supports_containers() { + let counter = self + .new_container_id_counter() + .expect("backend advertises container support without an id counter"); + let merge_fn = self + .container_merge_fn(TypeId::of::()) + .expect("backend advertises container support without merge semantics"); + self.container_values_mut_dyn() + .expect("backend returned a container counter without a container registry") + .register_type::(counter, move |state, old, new| merge_fn(state, old, new)); + } + } +} diff --git a/egglog/egglog-bridge/src/lib.rs b/egglog/egglog-bridge/src/lib.rs index ed42721..d1156d9 100644 --- a/egglog/egglog-bridge/src/lib.rs +++ b/egglog/egglog-bridge/src/lib.rs @@ -38,7 +38,7 @@ pub(crate) mod rule; #[cfg(test)] mod tests; -pub use rule::{Function, QueryEntry, RuleBuilder}; +pub use rule::{Function, QueryEntry, RuleBuilder, Variable, VariableId}; use thiserror::Error; /// A live registry of action handles for use by typed primitives. @@ -125,7 +125,7 @@ pub struct EGraph { /// around allows us to cache external function ids with repeat panic messages and they can /// also serve as a debugging tool in the case that the number of panic messages grows without /// bound. - panic_funcs: HashMap, + panic_funcs: HashMap, report_level: ReportLevel, /// Live registry of name-indexed action handles. Shared (via /// `Arc>`) with state wrappers and primitive callbacks @@ -137,6 +137,12 @@ pub struct EGraph { pub type Result = std::result::Result; +#[derive(Clone, Copy)] +struct CachedPanic { + id: ExternalFunctionId, + references: usize, +} + impl Default for EGraph { fn default() -> Self { let mut db = Database::new(); @@ -156,13 +162,19 @@ impl Default for EGraph { // also seeds `panic_funcs` so a later `new_panic` with the // same message reuses the id. let panic_message: SideChannel = Default::default(); - let mut panic_funcs: HashMap = Default::default(); + let mut panic_funcs: HashMap = Default::default(); let default_panic_msg = "primitive panicked".to_string(); let default_panic_id = db.add_external_function(Box::new(Panic( default_panic_msg.clone(), panic_message.clone(), ))); - panic_funcs.insert(default_panic_msg, default_panic_id); + panic_funcs.insert( + default_panic_msg, + CachedPanic { + id: default_panic_id, + references: 1, + }, + ); let union_action = UnionAction { table: uf_table, @@ -291,7 +303,21 @@ impl EGraph { } pub fn free_external_func(&mut self, func: ExternalFunctionId) { - self.db.free_external_function(func) + let mut free = true; + self.panic_funcs.retain(|_, cached| { + if cached.id != func { + true + } else if cached.references > 1 { + cached.references -= 1; + free = false; + true + } else { + false + } + }); + if free { + self.db.free_external_function(func); + } } /// Generate a fresh id. @@ -423,6 +449,50 @@ impl EGraph { self.db.clear_table(table_id); } + /// Remove the most recently registered function and its backing table. + /// + /// This is a rollback operation for callers whose later registration work + /// failed. It rejects older functions because later rules or merge + /// callbacks may already depend on them. + pub fn remove_last_table(&mut self, func: FunctionId) -> Result<()> { + anyhow::ensure!( + self.funcs.get(func).is_some(), + "cannot remove unknown function {func:?}" + ); + let removed_action = TableAction::new(self, func); + let info = self.funcs.pop_last(func).ok_or_else(|| { + anyhow::anyhow!("can only remove the most recently registered function") + })?; + if !self.db.remove_last_table(info.table) { + self.funcs.insert(func, info); + anyhow::bail!("function's backing table was not the most recently registered table"); + } + + self.free_rule(info.nonincremental_rebuild_rule); + for rule in info.incremental_rebuild_rules.iter().rev() { + self.free_rule(*rule); + } + + let replacement = self + .funcs + .iter() + .filter(|(_, candidate)| candidate.name == info.name) + .map(|(id, _)| TableAction::new(self, id)) + .last(); + let mut registry = self.action_registry.write().unwrap(); + if registry.table_actions.get(info.name.as_ref()) == Some(&removed_action) { + match replacement { + Some(action) => { + registry.register_table(info.name.to_string(), action); + } + None => { + registry.table_actions.remove(info.name.as_ref()); + } + } + } + Ok(()) + } + /// Read the contents of the given function. /// /// The callback `f` is called with each row and its subsumption status. @@ -979,6 +1049,7 @@ impl EGraph { struct RuleInfo { last_run_at: Timestamp, query: rule::Query, + owned_external_funcs: Vec, cached_plan: Option, desc: Arc, } @@ -1562,13 +1633,15 @@ struct Panic(String, SideChannel); impl EGraph { /// Create a new `ExternalFunction` that panics with the given message. pub fn new_panic(&mut self, message: String) -> ExternalFunctionId { - *self - .panic_funcs - .entry(message.to_string()) - .or_insert_with(|| { - let panic = Panic(message, self.panic_message.clone()); - self.db.add_external_function(Box::new(panic)) - }) + if let Some(cached) = self.panic_funcs.get_mut(&message) { + cached.references += 1; + return cached.id; + } + let panic = Panic(message.clone(), self.panic_message.clone()); + let id = self.db.add_external_function(Box::new(panic)); + self.panic_funcs + .insert(message, CachedPanic { id, references: 1 }); + id } pub fn new_panic_lazy( @@ -1582,10 +1655,8 @@ impl EGraph { } impl ExternalFunction for Panic { - fn invoke(&self, state: &mut core_relations::ExecutionState, args: &[Value]) -> Option { + fn invoke(&self, state: &mut core_relations::ExecutionState, _args: &[Value]) -> Option { // TODO (egglog feature): change this to support interpolating panic messages - assert!(args.is_empty()); - state.trigger_early_stop(); let mut guard = self.1.lock().unwrap(); if guard.is_none() { diff --git a/egglog/egglog-bridge/src/rule.rs b/egglog/egglog-bridge/src/rule.rs index 8a4a7da..0cbda7a 100644 --- a/egglog/egglog-bridge/src/rule.rs +++ b/egglog/egglog-bridge/src/rule.rs @@ -102,7 +102,6 @@ pub(crate) struct Query { uf_table: TableId, id_counter: CounterId, ts_counter: CounterId, - rule_id: RuleId, vars: DenseIdMap, atoms: Vec<(TableId, Vec, SchemaMath)>, /// The builders for queries in this module essentially wrap the lower-level @@ -123,11 +122,24 @@ pub(crate) struct Query { } pub struct RuleBuilder<'a> { - egraph: &'a mut EGraph, + resources: RuleResources<'a>, desc: Arc, query: Query, } +struct RuleResources<'a> { + egraph: &'a mut EGraph, + external_funcs: Vec, +} + +impl Drop for RuleResources<'_> { + fn drop(&mut self) { + for func in std::mem::take(&mut self.external_funcs) { + self.egraph.free_external_func(func); + } + } +} + impl EGraph { /// Add a rewrite rule for this [`EGraph`] using a [`RuleBuilder`]. /// If you aren't sure, use `egraph.new_rule("", true)`. @@ -135,15 +147,16 @@ impl EGraph { let uf_table = self.uf_table; let id_counter = self.id_counter; let ts_counter = self.timestamp_counter; - let rule_id = self.rules.reserve_slot(); RuleBuilder { - egraph: self, + resources: RuleResources { + egraph: self, + external_funcs: Vec::new(), + }, desc: Arc::from(desc), query: Query { uf_table, id_counter, ts_counter, - rule_id, seminaive, sole_focus: None, vars: Default::default(), @@ -157,18 +170,31 @@ impl EGraph { /// Remove a rewrite rule from this [`EGraph`]. pub fn free_rule(&mut self, id: RuleId) { - self.rules.take(id); + if let Some(info) = self.rules.take(id) { + for func in info.owned_external_funcs { + self.free_external_func(func); + } + } } } impl RuleBuilder<'_> { + fn new_panic_lazy( + &mut self, + message: impl FnOnce() -> String + Send + 'static, + ) -> ExternalFunctionId { + let panic = self.resources.egraph.new_panic_lazy(message); + self.resources.external_funcs.push(panic); + panic + } + fn add_callback(&mut self, cb: impl Brc + 'static) { self.query.add_rule.push(Box::new(cb)); } /// Access the underlying egraph within the builder. pub fn egraph(&self) -> &EGraph { - self.egraph + self.resources.egraph } /// Register a runtime panic with a custom message and return its @@ -179,7 +205,9 @@ impl RuleBuilder<'_> { /// /// [`call_external_func`]: Self::call_external_func pub fn new_panic(&mut self, message: String) -> crate::ExternalFunctionId { - self.egraph.new_panic(message) + let panic = self.resources.egraph.new_panic(message); + self.resources.external_funcs.push(panic); + panic } pub(crate) fn set_plan_strategy(&mut self, strategy: PlanStrategy) { @@ -295,15 +323,16 @@ impl RuleBuilder<'_> { if self.query.atoms.len() == 1 { self.query.plan_strategy = PlanStrategy::MinCover; } - let res = self.query.rule_id; + let res = self.resources.egraph.rules.reserve_slot(); let info = RuleInfo { last_run_at: Timestamp::new(0), query: self.query, + owned_external_funcs: std::mem::take(&mut self.resources.external_funcs), cached_plan: None, desc: self.desc, }; debug!("created rule {res:?} / {}", info.desc); - self.egraph.rules.insert(res, info); + self.resources.egraph.rules.insert(res, info); res } @@ -355,7 +384,7 @@ impl RuleBuilder<'_> { ) -> AtomId { let mut atom = entries.to_vec(); let schema_math = if let Some(func) = func { - let info = &self.egraph.funcs[func]; + let info = &self.resources.egraph.funcs[func]; assert_eq!(info.schema.len(), entries.len()); SchemaMath { subsume: info.can_subsume, @@ -394,7 +423,7 @@ impl RuleBuilder<'_> { let args = args.to_vec(); let res = self.new_var(ret_ty); // External functions that fail on the RHS of a rule should cause a panic. - let panic_fn = self.egraph.new_panic_lazy(panic_msg); + let panic_fn = self.new_panic_lazy(panic_msg); self.query.add_rule.push(Box::new(move |inner, rb| { let args = inner.convert_all(&args); let var = rb.call_external_with_fallback(func, &args, panic_fn, &[])?; @@ -413,7 +442,7 @@ impl RuleBuilder<'_> { entries: &[QueryEntry], is_subsumed: Option, ) -> Result { - let info = &self.egraph.funcs[func]; + let info = &self.resources.egraph.funcs[func]; let schema = &info.schema; if schema.len() != entries.len() { return Err(anyhow::Error::from(RuleBuilderError::ArityMismatch { @@ -483,7 +512,7 @@ impl RuleBuilder<'_> { }, || "subsumed a nonextestent row!".to_string(), ); - let info = &self.egraph.funcs[func]; + let info = &self.resources.egraph.funcs[func]; let schema_math = SchemaMath { subsume: info.can_subsume, func_cols: info.schema.len(), @@ -529,7 +558,7 @@ impl RuleBuilder<'_> { panic_msg: impl FnOnce() -> String + Send + 'static, ) -> Variable { let entries = entries.to_vec(); - let info = &self.egraph.funcs[func]; + let info = &self.resources.egraph.funcs[func]; let res = self .query .vars @@ -581,7 +610,7 @@ impl RuleBuilder<'_> { }) } DefaultVal::Fail => { - let panic_func = self.egraph.new_panic_lazy(panic_msg); + let panic_func = self.new_panic_lazy(panic_msg); Box::new(move |inner, rb| { let dst_vars = inner.convert_all(&entries); let var = rb.lookup_with_fallback( @@ -670,7 +699,7 @@ impl RuleBuilder<'_> { entries: &[QueryEntry], subsume_entry: QueryEntry, ) { - let info = &self.egraph.funcs[func]; + let info = &self.resources.egraph.funcs[func]; let table = info.table; let entries = entries.to_vec(); let schema_math = SchemaMath { @@ -693,7 +722,7 @@ impl RuleBuilder<'_> { /// Remove the value of a function from the database. pub fn remove(&mut self, table: FunctionId, entries: &[QueryEntry]) { - let table = self.egraph.funcs[table].table; + let table = self.resources.egraph.funcs[table].table; let entries = entries.to_vec(); let cb: BuildRuleCallback = Box::new(move |inner, rb| { let dst_vars = inner.convert_all(&entries); @@ -704,7 +733,7 @@ impl RuleBuilder<'_> { /// Panic with a given message. pub fn panic(&mut self, message: String) { - let panic = self.egraph.new_panic(message.clone()); + let panic = self.new_panic(message); let ret_ty = ColumnTy::Id; let res = self.new_var(ret_ty); self.query.add_rule.push(Box::new(move |inner, rb| { diff --git a/egglog/egglog-bridge/src/tests.rs b/egglog/egglog-bridge/src/tests.rs index 54df39f..b2d90c7 100644 --- a/egglog/egglog-bridge/src/tests.rs +++ b/egglog/egglog-bridge/src/tests.rs @@ -1,7 +1,7 @@ use std::{ fmt::Debug, hash::Hash, - slice, + iter, slice, sync::{ Arc, atomic::{AtomicUsize, Ordering}, @@ -11,7 +11,8 @@ use std::{ use crate::core_relations; use crate::core_relations::{ - ContainerValue, ExternalFunctionId, Value, ValueRebuilder, make_external_func, + ContainerValue, ExternalFunctionId, SortedWritesTable, Value, ValueRebuilder, + make_external_func, }; use crate::numeric_id::NumericId; use log::debug; @@ -19,10 +20,192 @@ use num_rational::Rational64; use once_cell::sync::Lazy; use crate::{ - ColumnTy, DefaultVal, EGraph, FunctionConfig, FunctionId, MergeFn, QueryEntry, add_expressions, - define_rule, + ColumnTy, DefaultVal, EGraph, FunctionConfig, FunctionId, MergeFn, QueryEntry, TableAction, + add_expressions, define_rule, }; +#[test] +fn dropped_rule_builder_releases_panics() { + let mut egraph = EGraph::default(); + let unit_type = egraph.base_values_mut().register_type::<()>(); + let register = |egraph: &mut EGraph| { + egraph.register_external_func(Box::new(make_external_func( + |state: &mut core_relations::ExecutionState<'_>, _args: &[Value]| { + Some(state.base_values().get(())) + }, + ))) + }; + let target = register(&mut egraph); + let reusable = register(&mut egraph); + egraph.free_external_func(reusable); + + let mut builder = egraph.new_rule("dropped", false); + builder.call_external_func(target, &[], ColumnTy::Base(unit_type), || "failed".into()); + drop(builder); + + assert_eq!(register(&mut egraph), reusable); + egraph.free_external_func(reusable); + + let mut builder = egraph.new_rule("dropped", false); + assert_eq!(builder.new_panic("direct panic".into()), reusable); + drop(builder); + assert_eq!(register(&mut egraph), reusable); + egraph.free_external_func(reusable); + + let mut builder = egraph.new_rule("dropped", false); + builder.panic("rule panic".into()); + drop(builder); + assert_eq!(register(&mut egraph), reusable); + egraph.free_external_func(reusable); + + let mut builder = egraph.new_rule("dropped", false); + let first = builder.new_panic("shared panic".into()); + let second = builder.new_panic("shared panic".into()); + assert_eq!(first, reusable); + assert_eq!(second, reusable); + drop(builder); + assert_eq!(register(&mut egraph), reusable); + egraph.free_external_func(reusable); + + let rule = { + let mut builder = egraph.new_rule("built", false); + builder.panic("built panic".into()); + builder.build() + }; + let occupied = register(&mut egraph); + assert_ne!(occupied, reusable); + egraph.free_external_func(occupied); + egraph.free_rule(rule); + assert_eq!(register(&mut egraph), reusable); +} + +#[test] +fn removing_last_table_restores_shadowed_registry_entry_and_id() { + let mut egraph = EGraph::default(); + let unit_type = egraph.base_values_mut().register_type::<()>(); + let unit = egraph.base_values().get(()); + let config = || FunctionConfig { + schema: vec![ColumnTy::Base(unit_type)], + default: DefaultVal::Const(unit), + merge: MergeFn::AssertEq, + name: "temporary".into(), + can_subsume: false, + }; + + let original = egraph.add_table(config()); + let original_action = egraph + .action_registry() + .read() + .unwrap() + .lookup_table("temporary") + .unwrap() + .clone(); + let temporary = egraph.add_table(config()); + + egraph.remove_last_table(temporary).unwrap(); + assert_eq!( + egraph + .action_registry() + .read() + .unwrap() + .lookup_table("temporary"), + Some(&original_action) + ); + assert_eq!(egraph.add_table(config()), temporary); + assert_ne!(original, temporary); +} + +#[test] +fn removing_unknown_or_non_last_function_preserves_registrations() { + let mut egraph = EGraph::default(); + let unit_type = egraph.base_values_mut().register_type::<()>(); + let unit = egraph.base_values().get(()); + let config = |name: &str| FunctionConfig { + schema: vec![ColumnTy::Base(unit_type)], + default: DefaultVal::Const(unit), + merge: MergeFn::AssertEq, + name: name.into(), + can_subsume: false, + }; + let first = egraph.add_table(config("first")); + let second = egraph.add_table(config("second")); + let first_action = TableAction::new(&egraph, first); + let second_action = TableAction::new(&egraph, second); + + assert!(egraph.remove_last_table(FunctionId::new(u32::MAX)).is_err()); + assert_eq!(egraph.table_size(first), 0); + assert_eq!(egraph.table_size(second), 0); + assert_eq!( + egraph + .action_registry() + .read() + .unwrap() + .lookup_table("first"), + Some(&first_action) + ); + assert_eq!( + egraph + .action_registry() + .read() + .unwrap() + .lookup_table("second"), + Some(&second_action) + ); + + assert!(egraph.remove_last_table(first).is_err()); + assert_eq!(egraph.table_size(first), 0); + assert_eq!(egraph.table_size(second), 0); + assert_eq!( + egraph + .action_registry() + .read() + .unwrap() + .lookup_table("first"), + Some(&first_action) + ); + assert_eq!( + egraph + .action_registry() + .read() + .unwrap() + .lookup_table("second"), + Some(&second_action) + ); +} + +#[test] +fn failed_database_table_removal_restores_function_registration() { + let mut egraph = EGraph::default(); + let unit_type = egraph.base_values_mut().register_type::<()>(); + let unit = egraph.base_values().get(()); + let function = egraph.add_table(FunctionConfig { + schema: vec![ColumnTy::Base(unit_type)], + default: DefaultVal::Const(unit), + merge: MergeFn::AssertEq, + name: "function".into(), + can_subsume: false, + }); + let action = TableAction::new(&egraph, function); + let blocker = egraph.db.add_table( + SortedWritesTable::new(1, 1, None, vec![], Box::new(|_, _, _, _| false)), + iter::empty(), + iter::empty(), + ); + + assert!(egraph.remove_last_table(function).is_err()); + assert_eq!(egraph.table_size(function), 0); + assert_eq!( + egraph + .action_registry() + .read() + .unwrap() + .lookup_table("function"), + Some(&action) + ); + assert!(egraph.db.remove_last_table(blocker)); + egraph.remove_last_table(function).unwrap(); +} + /// Run a simple associativity/commutativity test. /// /// The `can_subsume` argument is only used to enable subsumption on the underlying tables created @@ -1485,7 +1668,7 @@ fn panic_functions_trigger_early_stop() { let panic_fn = super::Panic("panic".to_string(), channel.clone()); let stopped = db.with_execution_state(|state| { assert!(!state.should_stop()); - let res = core_relations::ExternalFunction::invoke(&panic_fn, state, &[]); + let res = core_relations::ExternalFunction::invoke(&panic_fn, state, &[Value::new(1)]); assert!(res.is_none()); state.should_stop() }); diff --git a/egglog/numeric-id/src/lib.rs b/egglog/numeric-id/src/lib.rs index cfbe7d7..97d7225 100644 --- a/egglog/numeric-id/src/lib.rs +++ b/egglog/numeric-id/src/lib.rs @@ -149,6 +149,14 @@ impl DenseIdMap { self.data.get_mut(key.index()).unwrap().take() } + /// Remove `key` if it is the final slot in the map. + pub fn pop_last(&mut self, key: K) -> Option { + if key.inc() != self.next_id() { + return None; + } + self.data.pop().flatten() + } + /// Get the current mapping for `key` in the table, or insert the value /// returned by `f` and return a mutable reference to it. pub fn get_or_insert(&mut self, key: K, f: impl FnOnce() -> V) -> &mut V { diff --git a/egglog/numeric-id/src/tests.rs b/egglog/numeric-id/src/tests.rs index 6e086d8..0788906 100644 --- a/egglog/numeric-id/src/tests.rs +++ b/egglog/numeric-id/src/tests.rs @@ -44,6 +44,17 @@ fn basic_id_map() { ); } +#[test] +fn pop_last_only_removes_the_final_slot() { + let mut map = DenseIdMap::::new(); + let first = map.push("zero"); + let second = map.push("one"); + + assert_eq!(map.pop_last(first), None); + assert_eq!(map.pop_last(second), Some("one")); + assert_eq!(map.next_id(), second); +} + #[test] fn get_or_insert() { let mut map = DenseIdMap::::new(); diff --git a/egglog/src/cli.rs b/egglog/src/cli.rs index 3e39e4d..833dcd1 100644 --- a/egglog/src/cli.rs +++ b/egglog/src/cli.rs @@ -1,5 +1,7 @@ use crate::*; +use std::ffi::OsString; use std::io::{self, BufRead, BufReader, IsTerminal, Read, Write}; +use std::str::FromStr; use clap::Parser; use env_logger::Env; @@ -81,14 +83,27 @@ struct Args { /// This is what vanilla egglog uses, and custom egglog builds (i.e., "egglog batteries included") /// should also call this function. #[allow(clippy::disallowed_macros)] -pub fn cli(mut egraph: EGraph) { +pub fn cli(egraph: EGraph) { + cli_with_args(egraph, std::env::args_os()); +} + +/// Start a command-line interface with an explicit argv. +/// +/// Custom binaries can pre-parse their own flags and pass the remaining +/// arguments here while still using egglog's standard CLI behavior. +#[allow(clippy::disallowed_macros)] +pub fn cli_with_args(mut egraph: EGraph, args: I) +where + I: IntoIterator, + T: Into + Clone, +{ env_logger::Builder::from_env(Env::default().default_filter_or("warn")) .format_timestamp(None) .format_target(false) .parse_default_env() .init(); - let args = Args::parse(); + let args = Args::parse_from(args); if args.term_encoding { egraph = egraph.with_term_encoding_enabled(); diff --git a/egglog/src/constraint.rs b/egglog/src/constraint.rs index 3267a58..7ce89c2 100644 --- a/egglog/src/constraint.rs +++ b/egglog/src/constraint.rs @@ -1,7 +1,7 @@ use crate::{ core::{ Atom, CoreAction, CoreRule, GenericCoreActions, GenericCoreRule, HeadOrEq, Query, - StringOrEq, + QueryConstraints, StringOrEq, }, *, }; @@ -805,8 +805,17 @@ impl Problem { } } -impl CoreAction { - pub(crate) fn get_constraints( +trait CoreActionConstraints { + fn get_constraints( + &self, + typeinfo: &TypeInfo, + symbol_gen: &mut SymbolGen, + ctx: crate::Context, + ) -> Result>>, TypeError>; +} + +impl CoreActionConstraints for CoreAction { + fn get_constraints( &self, typeinfo: &TypeInfo, symbol_gen: &mut SymbolGen, @@ -875,8 +884,16 @@ impl CoreAction { } } -impl Atom { - pub(crate) fn get_constraints( +pub(crate) trait AtomConstraints { + fn get_constraints( + &self, + type_info: &TypeInfo, + ctx: crate::Context, + ) -> Result>>, TypeError>; +} + +impl AtomConstraints for Atom { + fn get_constraints( &self, type_info: &TypeInfo, ctx: crate::Context, @@ -1169,8 +1186,7 @@ pub(crate) fn grounded_check( let range = rule .body - .get_vars() - .into_iter() + .vars() .map(|v| ResolvedAtomTerm::Var(rule.span.clone(), v)) .collect(); let mut problem: Problem = Problem { diff --git a/egglog/src/core.rs b/egglog/src/core.rs index d0dd46d..8bdf2b5 100644 --- a/egglog/src/core.rs +++ b/egglog/src/core.rs @@ -12,12 +12,16 @@ //! Most compiler-time optimizations are expected to be done over CoreRule format. use std::hash::Hasher; use std::marker::PhantomData; -use std::ops::AddAssign; -use crate::{constraint::grounded_check, *}; -use egglog_ast::generic_ast::{Change, GenericAction, GenericActions, GenericExpr}; +use crate::{ + constraint::{AtomConstraints, grounded_check}, + *, +}; +pub use egglog_ast::core::{ + GenericAtom, GenericAtomTerm, GenericCoreAction, GenericCoreActions, GenericCoreRule, Query, +}; +use egglog_ast::generic_ast::{GenericAction, GenericActions, GenericExpr}; use egglog_ast::span::Span; -use egglog_ast::util::ListDisplay; use typechecking::{FuncType, PrimitiveValidator, PrimitiveWithId, TypeError}; #[derive(Debug, Clone, PartialEq, Eq)] @@ -27,6 +31,9 @@ pub(crate) enum HeadOrEq { } pub(crate) type StringOrEq = HeadOrEq; +pub type AtomTerm = GenericAtomTerm; +pub type Atom = GenericAtom; +pub type CoreAction = GenericCoreAction; impl From for StringOrEq { fn from(value: String) -> Self { @@ -234,157 +241,29 @@ impl IsFunc for String { } } -#[derive(Debug, Clone)] -pub enum GenericAtomTerm { - Var(Span, Leaf), - Literal(Span, Literal), - Global(Span, Leaf), -} - -// Ignores annotations for equality and hasing -impl PartialEq for GenericAtomTerm -where - Leaf: PartialEq, -{ - fn eq(&self, other: &Self) -> bool { - match (self, other) { - (GenericAtomTerm::Var(_, v1), GenericAtomTerm::Var(_, v2)) => v1 == v2, - (GenericAtomTerm::Literal(_, l1), GenericAtomTerm::Literal(_, l2)) => l1 == l2, - (GenericAtomTerm::Global(_, g1), GenericAtomTerm::Global(_, g2)) => g1 == g2, - _ => false, - } - } -} - -impl Eq for GenericAtomTerm where Leaf: Eq {} - -impl Hash for GenericAtomTerm -where - Leaf: Hash, -{ - fn hash(&self, state: &mut H) { - match self { - GenericAtomTerm::Var(_, v) => v.hash(state), - GenericAtomTerm::Literal(_, l) => l.hash(state), - GenericAtomTerm::Global(_, g) => g.hash(state), - } - } -} - -pub type AtomTerm = GenericAtomTerm; pub type ResolvedAtomTerm = GenericAtomTerm; -impl GenericAtomTerm { - pub fn span(&self) -> &Span { - match self { - GenericAtomTerm::Var(span, _) => span, - GenericAtomTerm::Literal(span, _) => span, - GenericAtomTerm::Global(span, _) => span, - } - } -} - -impl GenericAtomTerm { - pub fn to_expr(&self) -> GenericExpr { - match self { - GenericAtomTerm::Var(span, v) => GenericExpr::Var(span.clone(), v.clone()), - GenericAtomTerm::Literal(span, l) => GenericExpr::Lit(span.clone(), l.clone()), - GenericAtomTerm::Global(span, v) => GenericExpr::Var(span.clone(), v.clone()), - } - } -} - -impl ResolvedAtomTerm { - pub fn output(&self) -> ArcSort { - match self { - ResolvedAtomTerm::Var(_, v) => v.sort.clone(), - ResolvedAtomTerm::Literal(_, l) => literal_sort(l), - ResolvedAtomTerm::Global(_, v) => v.sort.clone(), - } - } -} - -impl std::fmt::Display for AtomTerm { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - AtomTerm::Var(_, v) => write!(f, "{v}"), - AtomTerm::Literal(_, lit) => write!(f, "{lit}"), - AtomTerm::Global(_, g) => write!(f, "{g}"), +fn atom_term_sort(term: &ResolvedAtomTerm) -> ArcSort { + match term { + GenericAtomTerm::Var(_, variable) | GenericAtomTerm::Global(_, variable) => { + variable.sort.clone() } + GenericAtomTerm::Literal(_, literal) => literal_sort(literal), } } -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct GenericAtom { - pub span: Span, - pub head: Head, - pub args: Vec>, -} - -pub type Atom = GenericAtom; - -impl std::fmt::Display for Atom { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "({} {}) ", self.head, ListDisplay(&self.args, " ")) - } -} - -impl GenericAtom -where - Leaf: Clone + Eq + Hash, - Head: Clone, -{ - pub fn vars(&self) -> impl Iterator + '_ { - self.args.iter().filter_map(|t| match t { - GenericAtomTerm::Var(_, v) => Some(v.clone()), - GenericAtomTerm::Literal(..) => None, - GenericAtomTerm::Global(..) => None, - }) - } - - fn subst(&mut self, subst: &HashMap>) { - for arg in self.args.iter_mut() { - match arg { - GenericAtomTerm::Var(_, v) => { - if let Some(at) = subst.get(v) { - *arg = at.clone(); - } - } - GenericAtomTerm::Literal(..) => (), - GenericAtomTerm::Global(..) => (), - } - } - } -} -impl Atom { - pub(crate) fn to_expr(&self) -> Expr { - let n = self.args.len(); - Expr::Call( - self.span.clone(), - self.head.clone(), - self.args[0..n - 1] - .iter() - .map(|arg| arg.to_expr()) - .collect(), - ) - } -} - -#[derive(Debug, Clone)] -pub struct Query { - pub atoms: Vec>, -} +pub(crate) trait QueryConstraints { + fn get_constraints( + &self, + type_info: &TypeInfo, + ctx: crate::Context, + ) -> Result>>, TypeError>; -impl Default for Query { - fn default() -> Self { - Self { - atoms: Default::default(), - } - } + fn atom_terms(&self) -> HashSet; } -impl Query { - pub fn get_constraints( +impl QueryConstraints for Query { + fn get_constraints( &self, type_info: &TypeInfo, ctx: crate::Context, @@ -396,7 +275,7 @@ impl Query { Ok(constraints) } - pub(crate) fn atom_terms(&self) -> HashSet { + fn atom_terms(&self) -> HashSet { self.atoms .iter() .flat_map(|atom| atom.args.iter().cloned()) @@ -404,174 +283,8 @@ impl Query { } } -impl Query -where - Leaf: Eq + Clone + Hash, - Head: Clone, -{ - pub(crate) fn get_vars(&self) -> IndexSet { - self.atoms - .iter() - .flat_map(|atom| atom.vars()) - .collect::>() - } -} - -impl AddAssign for Query { - fn add_assign(&mut self, rhs: Self) { - self.atoms.extend(rhs.atoms); - } -} - -impl std::fmt::Display for Query { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - for atom in &self.atoms { - writeln!(f, "{atom}")?; - } - Ok(()) - } -} - -impl std::fmt::Display for Query { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - for atom in self.funcs() { - writeln!(f, "{atom}")?; - } - let filters: Vec<_> = self.filters().collect(); - if !filters.is_empty() { - writeln!(f, "where ")?; - for filter in &filters { - writeln!( - f, - "({} {})", - filter.head.prim_with_id.primitive.name(), - ListDisplay(&filter.args, " ") - )?; - } - } - Ok(()) - } -} - -impl Query { - pub fn filters(&self) -> impl Iterator> + '_ { - self.atoms.iter().filter_map(|atom| match &atom.head { - ResolvedCall::Func(_) => None, - ResolvedCall::Primitive(head) => Some(GenericAtom { - span: atom.span.clone(), - head: head.clone(), - args: atom.args.clone(), - }), - }) - } - - pub fn funcs(&self) -> impl Iterator> + '_ { - self.atoms.iter().filter_map(|atom| match &atom.head { - ResolvedCall::Func(head) => Some(GenericAtom { - span: atom.span.clone(), - head: head.name.clone(), - args: atom.args.clone(), - }), - ResolvedCall::Primitive(_) => None, - }) - } -} - -#[derive(Clone, Debug, PartialEq, Eq, Hash)] -pub enum GenericCoreAction { - Let(Span, Leaf, Head, Vec>), - LetAtomTerm(Span, Leaf, GenericAtomTerm), - Set( - Span, - Head, - Vec>, - GenericAtomTerm, - ), - Change(Span, Change, Head, Vec>), - Union(Span, GenericAtomTerm, GenericAtomTerm), - Panic(Span, String), -} - -pub type CoreAction = GenericCoreAction; - -#[derive(Clone, Debug, PartialEq, Eq, Hash)] -pub struct GenericCoreActions(pub(crate) Vec>); pub(crate) type ResolvedCoreActions = GenericCoreActions; -impl Default for GenericCoreActions { - fn default() -> Self { - Self(vec![]) - } -} - -impl GenericCoreActions -where - Leaf: Clone + Eq + Hash, -{ - pub(crate) fn subst(&mut self, subst: &HashMap>) { - let actions = subst.iter().map(|(symbol, atom_term)| { - GenericCoreAction::LetAtomTerm( - atom_term.span().clone(), - symbol.clone(), - atom_term.clone(), - ) - }); - let existing_actions = std::mem::take(&mut self.0); - self.0 = actions.chain(existing_actions).collect(); - } - - fn new(actions: Vec>) -> GenericCoreActions { - Self(actions) - } - - pub(crate) fn get_free_vars(&self) -> HashSet { - let at_free_var = |at: &GenericAtomTerm| match at { - GenericAtomTerm::Var(_, v) => Some(v.clone()), - GenericAtomTerm::Literal(..) => None, - GenericAtomTerm::Global(..) => None, - }; - - let add_from_atom = |free_vars: &mut HashSet, at: &GenericAtomTerm| { - if let Some(v) = at_free_var(at) { - free_vars.insert(v); - } - }; - - let add_from_atoms = |free_vars: &mut HashSet, ats: &[GenericAtomTerm]| { - ats.iter().flat_map(&at_free_var).for_each(|v| { - free_vars.insert(v); - }); - }; - - let mut free_vars = HashSet::default(); - for action in self.0.iter().rev() { - match action { - GenericCoreAction::Let(_span, v, _, ats) => { - add_from_atoms(&mut free_vars, ats); - free_vars.remove(v); - } - GenericCoreAction::LetAtomTerm(_span, v, at) => { - add_from_atom(&mut free_vars, at); - free_vars.remove(v); - } - GenericCoreAction::Set(_span, _, ats, at) => { - add_from_atoms(&mut free_vars, ats); - add_from_atom(&mut free_vars, at); - } - GenericCoreAction::Change(_span, _change, _, ats) => { - add_from_atoms(&mut free_vars, ats); - } - GenericCoreAction::Union(_span, at, at1) => { - add_from_atom(&mut free_vars, at); - add_from_atom(&mut free_vars, at1); - } - GenericCoreAction::Panic(_span, _) => {} - } - } - free_vars - } -} - /// Shared state that threads through lowering from surface actions to core actions. /// pub(crate) struct CoreActionContext<'a, Head, Leaf, FG> { @@ -890,38 +603,37 @@ where } } -/// A [`GenericCoreRule`] represents a generalization of lowered form of a rule. -/// Unlike other `Generic`-prefixed types, [`GenericCoreRule`] takes two `Head` -/// parameters instead of one. This is because the `Head` parameter of `body` and -/// `head` can be different. In particular, early in the compilation pipeline, -/// `body` can contain `Eq` atoms, which denotes equality constraints, so the `Head` -/// for `body` needs to be a `HeadOrEq`, while `head` does not have equality -/// constraints. -#[derive(Debug, Clone)] -pub struct GenericCoreRule { - pub span: Span, - pub body: Query, - pub head: GenericCoreActions, -} - pub(crate) type CoreRule = GenericCoreRule; pub(crate) type ResolvedCoreRule = GenericCoreRule; -impl GenericCoreRule +trait CoreRuleSubst { + fn subst(&mut self, subst: &HashMap>); +} + +impl CoreRuleSubst for GenericCoreRule where - Head1: Clone, - Head2: Clone, Leaf: Clone + Eq + Hash, { - pub fn subst(&mut self, subst: &HashMap>) { + fn subst(&mut self, subst: &HashMap>) { for atom in &mut self.body.atoms { - atom.subst(subst); + atom.substitute_with(&mut |variable| subst.get(variable).cloned()); } - self.head.subst(subst); + let substitutions = subst.iter().map(|(variable, term)| { + GenericCoreAction::LetAtomTerm(term.span().clone(), variable.clone(), term.clone()) + }); + let actions = std::mem::take(&mut self.head.0); + self.head.0 = substitutions.chain(actions).collect(); } } -impl GenericCoreRule, Head, Leaf> +trait CanonicalizeCoreRule { + fn canonicalize( + self, + value_eq: impl Fn(&GenericAtomTerm, &GenericAtomTerm) -> Head, + ) -> GenericCoreRule; +} + +impl CanonicalizeCoreRule for GenericCoreRule, Head, Leaf> where Leaf: Eq + Clone + Hash + Debug, Head: Clone, @@ -930,7 +642,7 @@ where /// In particular, it removes equality checks between variables and /// other arguments, and turns equality checks between non-variable arguments /// into a primitive equality check `value-eq`. - pub(crate) fn canonicalize( + fn canonicalize( self, // Users need to pass in a substitute for equality constraints. value_eq: impl Fn(&GenericAtomTerm, &GenericAtomTerm) -> Head, @@ -1029,7 +741,14 @@ where eq_constraints } -impl GenericCoreRule +trait RemoveDuplicateVars { + fn remove_dup_vars( + self, + value_eq: impl Fn(&GenericAtomTerm, &GenericAtomTerm) -> Head, + ) -> Self; +} + +impl RemoveDuplicateVars for GenericCoreRule where Leaf: Eq + Clone + Hash + Debug, Head: Clone + Eq + Hash, @@ -1039,7 +758,7 @@ where /// For example, if we have two atoms `R(x, y, z1)` and `R(x, y, z2)`, /// then we can remove one of them and add an equality constraint `z1 = z2`. /// This is done until fixpoint, so it is kind of like rebuilding. - pub(crate) fn remove_dup_vars( + fn remove_dup_vars( mut self, value_eq: impl Fn(&GenericAtomTerm, &GenericAtomTerm) -> Head, ) -> Self { @@ -1109,7 +828,7 @@ where Leaf: Clone + PartialEq + Eq + Display + Hash + Debug, { let (body, _correspondence) = Facts(self.body.clone()).to_query(typeinfo, fresh_gen); - let mut binding = body.get_vars(); + let mut binding = body.vars().collect::>(); let mut ctx = CoreActionContext::new(typeinfo, &mut binding, fresh_gen, union_to_set_optimization); let (head, _correspondence) = self.head.to_core_actions(&mut ctx)?; @@ -1141,7 +860,7 @@ impl ResolvedRuleExt for ResolvedRule { let value_eq = |at1: &ResolvedAtomTerm, at2: &ResolvedAtomTerm| { ResolvedCall::Primitive(SpecializedPrimitive { prim_with_id: value_eq.clone(), - input: vec![at1.output(), at2.output()], + input: vec![atom_term_sort(at1), atom_term_sort(at2)], output: UnitSort.to_arcsort(), }) }; diff --git a/egglog/src/lib.rs b/egglog/src/lib.rs index 917890b..e21bdd1 100644 --- a/egglog/src/lib.rs +++ b/egglog/src/lib.rs @@ -39,7 +39,13 @@ pub use egglog_add_primitive::add_primitive_with_validator; use egglog_ast::generic_ast::{Change, GenericExpr, Literal}; use egglog_ast::span::Span; use egglog_ast::util::ListDisplay; -use egglog_bridge::{ColumnTy, QueryEntry}; +/// The pluggable backend interface. Re-exported so downstream crates can +/// implement their own backend (see [`EGraph::with_backend`]). +pub use egglog_backend_trait::{Backend, BackendExt}; +use egglog_backend_trait::{ + ReadMode, RuleActionCall, RuleBodyCall, RuleSetRun, RuleSpec, RuleValue, RuleVar, +}; +use egglog_bridge::ColumnTy; use egglog_core_relations as core_relations; use egglog_numeric_id as numeric_id; use egglog_reports::{ReportLevel, RunReport}; @@ -70,7 +76,6 @@ use std::io::{Read as _, Write as _}; use std::iter::once; use std::ops::Deref; use std::path::PathBuf; -use std::str::FromStr; use std::sync::Arc; pub use termdag::{OrdTerm, Term, TermDag, TermId}; use thiserror::Error; @@ -310,7 +315,7 @@ dyn_clone::clone_trait_object!(ExtensionStateValue); #[derive(Clone)] pub struct EGraph { - backend: egglog_bridge::EGraph, + backend: Box, pub parser: Parser, names: check_shadowing::Names, /// pushed_egraph forms a linked list of pushed egraphs. @@ -426,10 +431,22 @@ impl Debug for Function { impl Default for EGraph { fn default() -> Self { + Self::with_backend(Box::new(egglog_bridge::EGraph::default())) + } +} + +impl EGraph { + /// Construct an `EGraph` backed by the given [`Backend`] implementation. + /// + /// [`EGraph::default`] uses the in-memory reference backend + /// (`egglog_bridge::EGraph`); downstream crates can supply their own + /// backend (e.g. a differential-dataflow engine) by implementing + /// [`Backend`] and passing it here. + pub fn with_backend(backend: Box) -> Self { let mut parser = Parser::default(); let proof_state = EncodingState::new(&mut parser.symbol_gen); let mut eg = Self { - backend: Default::default(), + backend, parser, names: Default::default(), pushed_egraph: Default::default(), @@ -559,12 +576,34 @@ impl EGraph { self } + /// Enable the term-encoding pipeline for a custom backend. + /// + /// Relational backends without a native union-find require this: + /// congruence and rebuild are lowered to ordinary + /// rules over `@uf` tables instead of relying on the backend's own + /// union-find. Re-typechecking after the encoder runs uses a default + /// (bridge-backed) e-graph, so this backend need not implement typechecking. + pub fn with_term_encoding(mut self) -> Self { + self.proof_state.original_typechecking = Some(Box::new(EGraph::default())); + self + } + + /// Enable the term-encoding pipeline for a custom backend, using the given + /// bridge-backed e-graph for parsing/typechecking before instrumentation. + #[doc(hidden)] + pub fn with_term_encoding_typechecker(mut self, typechecker: EGraph) -> Self { + self.proof_state.original_typechecking = Some(Box::new(typechecker)); + self + } + /// Enable proof generation on this e-graph. /// TODO proofs should be turned on during creation of the e-graph, not afterwards. /// This method is to support the current CLI implementation with egglog-experimental (https://github.com/egraphs-good/egglog/issues/768) #[doc(hidden)] pub fn with_proofs_enabled(mut self) -> Self { - self = self.with_term_encoding_enabled(); + if self.proof_state.original_typechecking.is_none() { + self = self.with_term_encoding_enabled(); + } self.proof_state.proofs_enabled = true; self } @@ -779,7 +818,7 @@ impl EGraph { ) -> Result { match expr { GenericExpr::Lit(_, literal) => { - let val = literal_to_value(&self.backend, literal); + let val = literal_to_value(self.backend.base_values(), literal); Ok(egglog_bridge::MergeFn::Const(val)) } GenericExpr::Var(span, resolved_var) => match resolved_var.name.as_str() { @@ -809,17 +848,24 @@ impl EGraph { "expected string literal after `unstable-fn`".into(), )); }; + let bridge = self + .backend + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::BackendError( + "`unstable-fn` merge expressions require the reference bridge backend" + .into(), + ) + })?; + let panic_id = bridge.action_registry().read().unwrap().default_panic_id(); let resolved = resolve_function_container_target_with_context( - &self.backend, + bridge, &self.functions, &self.type_info, name, p, - self.backend - .action_registry() - .read() - .unwrap() - .default_panic_id(), + panic_id, )?; translated_args[0] = egglog_bridge::MergeFn::Const(self.backend.base_values().get(resolved)); @@ -860,7 +906,7 @@ impl EGraph { schema: input .iter() .chain([&output]) - .map(|sort| sort.column_ty(&self.backend)) + .map(|sort| sort.column_ty(self.backend.base_values())) .collect(), default: match decl.subtype { FunctionSubtype::Constructor => DefaultVal::FreshId, @@ -1160,22 +1206,16 @@ impl EGraph { let iteration_report = self .backend - .run_rules(&rule_ids) + .run_rules(RuleSetRun { + name: Some(ruleset), + rules: &rule_ids, + }) .map_err(|e| Error::BackendError(e.to_string()))?; Ok(RunReport::singleton(ruleset, iteration_report)) } fn add_rule(&mut self, rule: ast::ResolvedRule) -> Result { - // Disable union_to_set optimization in proof or term encoding mode, since - // it expects only `union` on constructors (not set). - let core_rule = rule.to_canonicalized_core_rule( - &self.type_info, - &mut self.parser.symbol_gen, - self.proof_state.original_typechecking.is_none(), - )?; - let (query, actions) = (&core_rule.body, &core_rule.head); - // The `:naive` rule option opts a single rule out of seminaive // evaluation. This widens primitive-context selection from // Pure/Write to Read/Full, so primitives that read or write the @@ -1191,33 +1231,54 @@ impl EGraph { RuleEvalMode::Naive | RuleEvalMode::UnsafeSeminaive ); + // Disable union_to_set optimization in proof or term encoding mode, since + // it expects only `union` on constructors (not set). + let union_to_set = self.proof_state.original_typechecking.is_none(); + + match self.rulesets.get(&rule.ruleset) { + Some(Ruleset::Rules(_)) => {} + Some(Ruleset::Combined(_)) => { + return Err(Error::CombinedRulesetError( + rule.ruleset.clone(), + rule.span.clone(), + )); + } + None => { + return Err(Error::NoSuchRuleset( + rule.ruleset.clone(), + rule.span.clone(), + )); + } + } + + let core_rule = rule.to_canonicalized_core_rule( + &self.type_info, + &mut self.parser.symbol_gen, + union_to_set, + )?; + let (query, actions) = (&core_rule.body, &core_rule.head); let rule_id = { - let mut rb = self.backend.new_rule(&rule.name, seminaive); - rb.set_no_decomp(no_decomp); - let mut translator = - BackendRule::new(rb, &self.functions, &self.type_info, requires_read_context); - translator.query(query, rule.include_subsumed); + let mut translator = BackendRule::new( + &mut *self.backend, + &self.functions, + &self.type_info, + requires_read_context, + ); + translator.query(query, rule.include_subsumed)?; translator.actions(actions)?; - translator.build() + translator.try_build(&rule.name, seminaive, no_decomp, core_rule.span.clone())? }; - if let Some(rules) = self.rulesets.get_mut(&rule.ruleset) { - match rules { - Ruleset::Rules(rules) => { - match rules.entry(rule.name.clone()) { - indexmap::map::Entry::Occupied(_) => { - let name = rule.name; - panic!("Rule '{name}' was already present") - } - indexmap::map::Entry::Vacant(e) => e.insert((core_rule, rule_id)), - }; - Ok(rule.name) - } - Ruleset::Combined(_) => Err(Error::CombinedRulesetError(rule.ruleset, rule.span)), + let Some(Ruleset::Rules(rules)) = self.rulesets.get_mut(&rule.ruleset) else { + unreachable!("ruleset was validated before compiling the rule") + }; + match rules.entry(rule.name.clone()) { + indexmap::map::Entry::Occupied(_) => { + panic!("Rule '{}' was already present", rule.name) } - } else { - Err(Error::NoSuchRuleset(rule.ruleset, rule.span)) - } + indexmap::map::Entry::Vacant(e) => e.insert((core_rule, rule_id)), + }; + Ok(rule.name) } fn eval_actions(&mut self, actions: &ResolvedActions) -> Result<(), Error> { @@ -1231,14 +1292,17 @@ impl EGraph { let (actions, _) = actions.to_core_actions(&mut ctx)?; let mut translator = BackendRule::new( - self.backend.new_rule("eval_actions", false), + &mut *self.backend, &self.functions, &self.type_info, true, // global action: Read/Full contexts (may read the DB) ); translator.actions(&actions)?; - let id = translator.build(); - let result = self.backend.run_rules(&[id]); + let id = translator.try_build("eval_actions", false, false, Span::Panic)?; + let result = self.backend.run_rules(RuleSetRun { + name: None, + rules: &[id], + }); self.backend.free_rule(id); match result { @@ -1262,8 +1326,18 @@ impl EGraph { /// a [`ReadState`], so it can read but not write. Because this /// borrows `&self`, the closure and its callbacks may also call other /// `&self` methods such as [`EGraph::value_to_base`]. + /// + /// # Panics + /// + /// Panics if the selected backend does not provide an action registry. + /// Use the fallible top-level table iteration methods for backend-generic + /// reads. pub fn read(&self, f: impl FnOnce(ReadState<'_, '_>) -> R) -> R { - let registry = self.backend.action_registry().clone(); + let registry = self + .backend + .action_registry() + .cloned() + .expect("EGraph::read requires a backend action registry"); let guard = registry.read().unwrap(); self.backend .with_execution_state_tracked(|es| f(ReadState::wrap(es, &guard, Context::Read))) @@ -1276,34 +1350,94 @@ impl EGraph { pub fn function_entries( &self, name: &str, - f: impl FnMut(FunctionEntry<'_>), + mut f: impl FnMut(FunctionEntry<'_>), ) -> Result<(), Error> { - self.read(|rs| rs.function_entries(name, f)) + self.function_entries_while(name, |entry| { + f(entry); + true + }) } /// Like [`EGraph::function_entries`], but stops when `f` returns `false`. pub fn function_entries_while( &self, name: &str, - f: impl FnMut(FunctionEntry<'_>) -> bool, + mut f: impl FnMut(FunctionEntry<'_>) -> bool, ) -> Result<(), Error> { - self.read(|rs| rs.function_entries_while(name, f)) + let function = + self.functions + .get(name) + .ok_or_else(|| crate::api::ApiError::MissingTable { + name: name.to_owned(), + })?; + if function.subtype() != FunctionSubtype::Custom { + return Err(crate::api::ApiError::WrongSubtype { + name: name.to_owned(), + expected: "function", + actual: "constructor", + } + .into()); + } + self.backend.for_each_while(function.backend_id, |row| { + let (output, inputs) = row + .vals + .split_last() + .expect("function row has at least an output column"); + f(FunctionEntry { + inputs, + output: *output, + subsumed: row.subsumed, + }) + }); + Ok(()) } /// Call `f` on each [`Enode`] of a constructor / relation table. /// Top-level form of [`Read::constructor_enodes`]; errors if `name` /// is a function or unregistered. - pub fn constructor_enodes(&self, name: &str, f: impl FnMut(Enode<'_>)) -> Result<(), Error> { - self.read(|rs| rs.constructor_enodes(name, f)) + pub fn constructor_enodes( + &self, + name: &str, + mut f: impl FnMut(Enode<'_>), + ) -> Result<(), Error> { + self.constructor_enodes_while(name, |enode| { + f(enode); + true + }) } /// Like [`EGraph::constructor_enodes`], but stops when `f` returns `false`. pub fn constructor_enodes_while( &self, name: &str, - f: impl FnMut(Enode<'_>) -> bool, + mut f: impl FnMut(Enode<'_>) -> bool, ) -> Result<(), Error> { - self.read(|rs| rs.constructor_enodes_while(name, f)) + let function = + self.functions + .get(name) + .ok_or_else(|| crate::api::ApiError::MissingTable { + name: name.to_owned(), + })?; + if function.subtype() != FunctionSubtype::Constructor { + return Err(crate::api::ApiError::WrongSubtype { + name: name.to_owned(), + expected: "constructor", + actual: "function", + } + .into()); + } + self.backend.for_each_while(function.backend_id, |row| { + let (eclass, children) = row + .vals + .split_last() + .expect("constructor row has at least an eclass column"); + f(Enode { + children, + eclass: *eclass, + subsumed: row.subsumed, + }) + }); + Ok(()) } /// Remove every row from the named function in bulk. @@ -1431,18 +1565,43 @@ impl EGraph { .unwrap_or_else(|| Span::Panic) ))); }; + if !self.backend.as_any().is::() { + return Err(Error::BackendError( + "`unstable-fn` is only supported on the reference bridge backend" + .into(), + )); + } let panic_id = self.backend.new_panic(format!( "unstable-fn over `{name}` was applied in a context where its wrapped \ function is not valid for this call site, if in a rule, add :naive." )); - let resolved_function = resolve_function_container_target_with_context( - &self.backend, - &self.functions, - &self.type_info, - name, - prim, - panic_id, - )?; + let resolved_function = { + let bridge = self + .backend + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::BackendError( + "`unstable-fn` is only supported on the reference bridge backend" + .into(), + ) + })?; + resolve_function_container_target_with_context( + bridge, + &self.functions, + &self.type_info, + name, + prim, + panic_id, + ) + }; + let resolved_function = match resolved_function { + Ok(resolved) => resolved, + Err(error) => { + self.backend.free_external_func(panic_id); + return Err(error); + } + }; let fn_value = self.backend.base_values().get(resolved_function); let binding_name = self.parser.symbol_gen.fresh("unstable_fn_target"); bindings.push((binding_name.clone(), fn_value)); @@ -1495,11 +1654,12 @@ impl EGraph { }))); let mut translator = BackendRule::new( - self.backend.new_rule("eval_resolved_expr", false), + &mut *self.backend, &self.functions, &self.type_info, true, // global action: Read/Full contexts (may read the DB) ); + translator.rollback_external_funcs.push(ext_id); let result_var = ResolvedVar { name: self.parser.symbol_gen.fresh("eval_resolved_expr"), @@ -1521,16 +1681,20 @@ impl EGraph { let actions = actions.to_core_actions(&mut ctx)?.0; translator.actions(&actions)?; - let arg = translator.entry(&ResolvedAtomTerm::Var(span.clone(), result_var)); - translator.rb.call_external_func( + let arg = translator.entry(&ResolvedAtomTerm::Var(span.clone(), result_var))?; + translator.call_external_func( + span.clone(), ext_id, - &[arg], + "eval_resolved_expr_result", + vec![arg], egglog_bridge::ColumnTy::Base(unit_id), - || "this function will never panic".to_string(), ); - let id = translator.build(); - let rule_result = self.backend.run_rules(&[id]); + let id = translator.try_build("eval_resolved_expr", false, false, span)?; + let rule_result = self.backend.run_rules(RuleSetRun { + name: None, + rules: &[id], + }); self.backend.free_rule(id); self.backend.free_external_func(ext_id); let _ = rule_result.map_err(|e| { @@ -1585,19 +1749,25 @@ impl EGraph { }))); let mut translator = BackendRule::new( - self.backend.new_rule("check_facts", false), + &mut *self.backend, &self.functions, &self.type_info, true, // global query: Read context (may read the DB) ); - translator.query(&query, true); - translator - .rb - .call_external_func(ext_id, &[], egglog_bridge::ColumnTy::Id, || { - "this function will never panic".to_string() - }); - let id = translator.build(); - let run_result = self.backend.run_rules(&[id]); + translator.rollback_external_funcs.push(ext_id); + translator.query(&query, true)?; + translator.call_external_func( + span.clone(), + ext_id, + "check_facts_match", + Vec::new(), + egglog_bridge::ColumnTy::Id, + ); + let id = translator.try_build("check_facts", false, false, span.clone())?; + let run_result = self.backend.run_rules(RuleSetRun { + name: None, + rules: &[id], + }); self.backend.free_rule(id); self.backend.free_external_func(ext_id); run_result.map_err(|e| Error::BackendError(e.to_string()))?; @@ -1961,7 +2131,16 @@ impl EGraph { let num_facts = parsed_contents.len(); - let table_action = egglog_bridge::TableAction::new(&self.backend, func.backend_id); + let bridge = self + .backend + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::BackendError( + "loading facts from a file requires the reference bridge backend".into(), + ) + })?; + let table_action = egglog_bridge::TableAction::new(bridge, func.backend_id); if function_type.subtype != FunctionSubtype::Constructor { self.backend.with_execution_state(|es| { @@ -2151,6 +2330,10 @@ impl EGraph { /// Run a program, represented as an AST. /// Return a list of messages. pub fn run_program(&mut self, program: Vec) -> Result, Error> { + if self.backend.requires_term_encoding() && self.proof_state.original_typechecking.is_none() + { + return Err(Error::BackendRequiresTermEncoding); + } let res = self.process_program_internal(program, true)?; Ok(res.outputs) } @@ -2292,6 +2475,12 @@ impl EGraph { self.functions.get(name) } + /// Whether this e-graph's backend exposes the in-memory action registry + /// used by registry-backed primitives. + pub fn supports_action_registry(&self) -> bool { + self.backend.action_registry().is_some() + } + /// Returns `true` if a user-defined command with the given name is /// registered in this e-graph. pub fn has_command(&self, name: &str) -> bool { @@ -2382,7 +2571,9 @@ impl EGraph { &mut self, f: impl FnOnce(FullState<'_, '_>) -> Result, ) -> Result { - let registry = self.backend.action_registry().clone(); + let registry = self.backend.action_registry().cloned().ok_or_else(|| { + Error::BackendError("EGraph::update requires a backend action registry".into()) + })?; let guard = registry.read().unwrap(); let (result, changed) = self .backend @@ -2591,8 +2782,12 @@ fn resolve_function_container_target_with_context( } struct BackendRule<'a> { - rb: egglog_bridge::RuleBuilder<'a>, - entries: HashMap, + backend: &'a mut dyn Backend, + entries: HashMap>, + next_var: u32, + body: core::Query, + head: core::GenericCoreActions, + rollback_external_funcs: Vec, functions: &'a IndexMap, type_info: &'a TypeInfo, /// Whether primitives may read the database. When true the per-phase @@ -2604,17 +2799,21 @@ struct BackendRule<'a> { impl<'a> BackendRule<'a> { fn new( - rb: egglog_bridge::RuleBuilder<'a>, + backend: &'a mut dyn Backend, functions: &'a IndexMap, type_info: &'a TypeInfo, requires_read_context: bool, ) -> BackendRule<'a> { BackendRule { - rb, + backend, functions, type_info, requires_read_context, entries: Default::default(), + next_var: 0, + body: Default::default(), + head: Default::default(), + rollback_external_funcs: Vec::new(), } } @@ -2645,19 +2844,40 @@ impl<'a> BackendRule<'a> { } } - fn entry(&mut self, x: &core::ResolvedAtomTerm) -> QueryEntry { - self.entries - .entry(x.clone()) - .or_insert_with(|| match x { - core::GenericAtomTerm::Var(_, v) => self - .rb - .new_var_named(v.sort.column_ty(self.rb.egraph()), &v.name), - core::GenericAtomTerm::Literal(_, l) => literal_to_entry(self.rb.egraph(), l), - core::GenericAtomTerm::Global(..) => { - panic!("Globals should have been desugared") - } - }) - .clone() + fn fresh_var(&mut self, variable: &ResolvedVar) -> RuleVar { + let id = self.next_var; + self.next_var += 1; + RuleVar { + id, + name: variable.name.clone().into_boxed_str(), + ty: variable.sort.column_ty(self.backend.base_values()), + } + } + + fn entry( + &mut self, + term: &core::ResolvedAtomTerm, + ) -> Result, Error> { + if let Some(entry) = self.entries.get(term) { + return Ok(entry.clone()); + } + let entry = match term { + core::GenericAtomTerm::Var(span, variable) => { + core::GenericAtomTerm::Var(span.clone(), self.fresh_var(variable)) + } + core::GenericAtomTerm::Literal(span, literal) => core::GenericAtomTerm::Literal( + span.clone(), + literal_to_rule_value(self.backend.base_values(), literal), + ), + core::GenericAtomTerm::Global(span, variable) => { + return Err(Error::BackendError(format!( + "{span}: global `{}` was not desugared before backend lowering", + variable.name + ))); + } + }; + self.entries.insert(term.clone(), entry.clone()); + Ok(entry) } fn func(&self, f: &typechecking::FuncType) -> egglog_bridge::FunctionId { @@ -2669,164 +2889,352 @@ impl<'a> BackendRule<'a> { prim: &core::SpecializedPrimitive, args: &[core::ResolvedAtomTerm], ctx: crate::Context, - ) -> (ExternalFunctionId, Vec, ColumnTy) { + ) -> Result< + ( + ExternalFunctionId, + Vec>, + ColumnTy, + ), + Error, + > { // The typechecker has already checked that this primitive is // valid in `ctx`; pick the runtime id that stamps the same ctx // onto the state wrapper when invoked. let resolved_id = prim.external_id(ctx); - let mut qe_args = self.args(args); + let mut rule_args = self.args(args)?; if prim.name() == "unstable-fn" { - let core::ResolvedAtomTerm::Literal(_, Literal::String(ref name)) = args[0] else { - panic!("expected string literal after `unstable-fn`") + let Some(core::ResolvedAtomTerm::Literal(_, Literal::String(name))) = args.first() + else { + return Err(Error::BackendError( + "expected string literal after `unstable-fn`".into(), + )); }; + if self + .backend + .as_any() + .downcast_ref::() + .is_none() + { + return Err(Error::BackendError( + "`unstable-fn` is only supported on the reference bridge backend".into(), + )); + } // Pre-register a panic id used by `FunctionContainer::apply` // when the wrapped function is applied in a context that // doesn't admit it. Triggered at runtime via the egglog // panic side channel so misuse surfaces as an `Err` from // `run_rules` rather than a thread unwind. - let panic_id = self.rb.new_panic(format!( + let panic_id = self.backend.new_panic(format!( "unstable-fn over `{name}` was applied in a context where its wrapped \ function is not valid for this call site, if in a rule, add :naive." )); + self.rollback_external_funcs.push(panic_id); + let bridge = self + .backend + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::BackendError( + "`unstable-fn` is only supported on the reference bridge backend".into(), + ) + })?; let resolved = resolve_function_container_target_with_context( - self.rb.egraph(), + bridge, self.functions, self.type_info, name, prim, panic_id, - ) - .unwrap_or_else(|err| panic!("{err}")); - - qe_args[0] = self.rb.egraph().base_value_constant(resolved); + ); + let resolved = resolved?; + rule_args[0] = core::GenericAtomTerm::Literal( + args[0].span().clone(), + base_rule_value(self.backend.base_values(), resolved), + ); } - ( - resolved_id, - qe_args, - prim.output().column_ty(self.rb.egraph()), - ) + let output_ty = prim.output().column_ty(self.backend.base_values()); + Ok((resolved_id, rule_args, output_ty)) } fn args<'b>( &mut self, args: impl IntoIterator, - ) -> Vec { - args.into_iter().map(|x| self.entry(x)).collect() + ) -> Result>, Error> { + args.into_iter().map(|term| self.entry(term)).collect() } - fn query(&mut self, query: &core::Query, include_subsumed: bool) { + fn query( + &mut self, + query: &core::Query, + include_subsumed: bool, + ) -> Result<(), Error> { for atom in &query.atoms { - match &atom.head { - ResolvedCall::Func(f) => { - let f = self.func(f); - let args = self.args(&atom.args); - let is_subsumed = match include_subsumed { - true => None, - false => Some(false), - }; - self.rb.query_table(f, &args, is_subsumed).unwrap(); - } + let (head, args) = match &atom.head { + ResolvedCall::Func(f) => ( + RuleBodyCall::Table { + id: self.func(f), + read: if include_subsumed { + ReadMode::All + } else { + ReadMode::Live + }, + }, + self.args(&atom.args)?, + ), ResolvedCall::Primitive(p) => { let ctx = self.query_context(); - let (p, args, ty) = self.prim(p, &atom.args, ctx); - self.rb.query_prim(p, &args, ty).unwrap() + let (id, args, output) = self.prim(p, &atom.args, ctx)?; + ( + RuleBodyCall::Primitive { + id, + name: p.name().into(), + output, + }, + args, + ) } - } + }; + self.body.atoms.push(core::GenericAtom { + span: atom.span.clone(), + head, + args, + }); } + Ok(()) } fn actions(&mut self, actions: &core::ResolvedCoreActions) -> Result<(), Error> { for action in &actions.0 { match action { core::GenericCoreAction::Let(span, v, f, args) => { - let v = core::GenericAtomTerm::Var(span.clone(), v.clone()); - let y = match f { - ResolvedCall::Func(f) => { - let name = f.name.clone(); - let f = self.func(f); - let args = self.args(args); - let span = span.clone(); - self.rb.lookup(f, &args, move || { - format!("{span}: lookup of function {name} failed") - }) - } + let (call, args) = match f { + ResolvedCall::Func(f) => ( + RuleActionCall::Table { + id: self.func(f), + name: f.name.clone().into_boxed_str(), + }, + self.args(args)?, + ), ResolvedCall::Primitive(p) => { - let name = p.name().to_owned(); let ctx = self.action_context(); - let (p, args, ty) = self.prim(p, args, ctx); - let span = span.clone(); - self.rb.call_external_func(p, &args, ty, move || { - format!("{span}: call of primitive {name} failed") - }) + let (id, args, output) = self.prim(p, args, ctx)?; + ( + RuleActionCall::Primitive { + id, + name: p.name().into(), + output, + }, + args, + ) } }; - self.entries.insert(v, y.into()); + let variable = self.fresh_var(v); + self.head.0.push(core::GenericCoreAction::Let( + span.clone(), + variable.clone(), + call, + args, + )); + self.entries.insert( + core::GenericAtomTerm::Var(span.clone(), v.clone()), + core::GenericAtomTerm::Var(span.clone(), variable), + ); } core::GenericCoreAction::LetAtomTerm(span, v, x) => { - let v = core::GenericAtomTerm::Var(span.clone(), v.clone()); - let x = self.entry(x); - self.entries.insert(v, x); + let value = self.entry(x)?; + let variable = self.fresh_var(v); + self.head.0.push(core::GenericCoreAction::LetAtomTerm( + span.clone(), + variable.clone(), + value, + )); + self.entries.insert( + core::GenericAtomTerm::Var(span.clone(), v.clone()), + core::GenericAtomTerm::Var(span.clone(), variable), + ); } - core::GenericCoreAction::Set(_, f, xs, y) => match f { - ResolvedCall::Primitive(..) => panic!("runtime primitive set!"), + core::GenericCoreAction::Set(span, f, xs, y) => match f { + ResolvedCall::Primitive(..) => { + return Err(Error::BackendError("cannot set a primitive".into())); + } ResolvedCall::Func(f) => { - let f = self.func(f); - let args = self.args(xs.iter().chain([y])); - self.rb.set(f, &args) + let arguments = self.args(xs)?; + let value = self.entry(y)?; + self.head.0.push(core::GenericCoreAction::Set( + span.clone(), + RuleActionCall::Table { + id: self.func(f), + name: f.name.clone().into_boxed_str(), + }, + arguments, + value, + )); } }, core::GenericCoreAction::Change(span, change, f, args) => match f { - ResolvedCall::Primitive(..) => panic!("runtime primitive change!"), + ResolvedCall::Primitive(..) => { + return Err(Error::BackendError( + "cannot delete or subsume a primitive".into(), + )); + } ResolvedCall::Func(f) => { let name = f.name.clone(); let can_subsume = self.functions[&f.name].can_subsume; - let f = self.func(f); - let args = self.args(args); - match change { - Change::Delete => self.rb.remove(f, &args), - Change::Subsume if can_subsume => self.rb.subsume(f, &args), - Change::Subsume => { - return Err(Error::SubsumeMergeError(name, span.clone())); - } + if matches!(change, Change::Subsume) && !can_subsume { + return Err(Error::SubsumeMergeError(name, span.clone())); } + let arguments = self.args(args)?; + self.head.0.push(core::GenericCoreAction::Change( + span.clone(), + *change, + RuleActionCall::Table { + id: self.func(f), + name: f.name.clone().into_boxed_str(), + }, + arguments, + )); } }, - core::GenericCoreAction::Union(_, x, y) => { - let x = self.entry(x); - let y = self.entry(y); - self.rb.union(x, y) + core::GenericCoreAction::Union(span, x, y) => { + let x = self.entry(x)?; + let y = self.entry(y)?; + self.head + .0 + .push(core::GenericCoreAction::Union(span.clone(), x, y)); + } + core::GenericCoreAction::Panic(span, message) => { + self.head.0.push(core::GenericCoreAction::Panic( + span.clone(), + message.clone(), + )); } - core::GenericCoreAction::Panic(_, message) => self.rb.panic(message.clone()), } } Ok(()) } - fn build(self) -> egglog_bridge::RuleId { - self.rb.build() + fn query_table( + &mut self, + span: Span, + table: egglog_bridge::FunctionId, + entries: Vec>, + read: ReadMode, + ) { + self.body.atoms.push(core::GenericAtom { + span, + head: RuleBodyCall::Table { id: table, read }, + args: entries, + }); + } + + fn call_external_func( + &mut self, + span: Span, + id: ExternalFunctionId, + name: &str, + arguments: Vec>, + output: ColumnTy, + ) -> core::GenericAtomTerm { + let variable = RuleVar { + id: self.next_var, + name: format!("@{name}").into_boxed_str(), + ty: output, + }; + self.next_var += 1; + self.head.0.push(core::GenericCoreAction::Let( + span.clone(), + variable.clone(), + RuleActionCall::Primitive { + id, + name: name.into(), + output, + }, + arguments, + )); + core::GenericAtomTerm::Var(span, variable) + } + + fn remove( + &mut self, + span: Span, + table: egglog_bridge::FunctionId, + name: &str, + arguments: Vec>, + ) { + self.head.0.push(core::GenericCoreAction::Change( + span, + Change::Delete, + RuleActionCall::Table { + id: table, + name: name.into(), + }, + arguments, + )); + } + + fn try_build( + mut self, + name: &str, + seminaive: bool, + no_decomp: bool, + span: Span, + ) -> Result { + let spec = RuleSpec { + name: name.to_owned(), + seminaive, + no_decomp, + core: core::GenericCoreRule { + span, + body: std::mem::take(&mut self.body), + head: std::mem::take(&mut self.head), + }, + }; + let result = self + .backend + .add_rule(spec) + .map_err(|error| Error::BackendError(error.to_string())); + if result.is_ok() { + self.rollback_external_funcs.clear(); + } + result + } +} + +impl Drop for BackendRule<'_> { + fn drop(&mut self) { + for id in self.rollback_external_funcs.drain(..) { + self.backend.free_external_func(id); + } } } -fn literal_to_entry(egraph: &egglog_bridge::EGraph, l: &Literal) -> QueryEntry { +fn base_rule_value(base_values: &BaseValues, x: T) -> RuleValue { + RuleValue { + value: base_values.get(x), + ty: ColumnTy::Base(base_values.get_ty::()), + } +} + +fn literal_to_rule_value(base_values: &BaseValues, l: &Literal) -> RuleValue { match l { - Literal::Int(x) => egraph.base_value_constant::(*x), - Literal::Float(x) => egraph.base_value_constant::(x.into()), - Literal::String(x) => egraph.base_value_constant::(sort::S::new(x.clone())), - Literal::Bool(x) => egraph.base_value_constant::(*x), - Literal::Unit => egraph.base_value_constant::<()>(()), + Literal::Int(x) => base_rule_value::(base_values, *x), + Literal::Float(x) => base_rule_value::(base_values, x.into()), + Literal::String(x) => base_rule_value::(base_values, sort::S::new(x.clone())), + Literal::Bool(x) => base_rule_value::(base_values, *x), + Literal::Unit => base_rule_value::<()>(base_values, ()), } } -fn literal_to_value(egraph: &egglog_bridge::EGraph, l: &Literal) -> Value { +fn literal_to_value(base_values: &BaseValues, l: &Literal) -> Value { match l { - Literal::Int(x) => egraph.base_values().get::(*x), - Literal::Float(x) => egraph.base_values().get::(x.into()), - Literal::String(x) => egraph.base_values().get::(sort::S::new(x.clone())), - Literal::Bool(x) => egraph.base_values().get::(*x), - Literal::Unit => egraph.base_values().get::<()>(()), + Literal::Int(x) => base_values.get::(*x), + Literal::Float(x) => base_values.get::(x.into()), + Literal::String(x) => base_values.get::(sort::S::new(x.clone())), + Literal::Bool(x) => base_values.get::(*x), + Literal::Unit => base_values.get::<()>(()), } } @@ -2852,6 +3260,10 @@ pub enum Error { CombinedRulesetError(String, Span), #[error("{0}")] BackendError(String), + #[error( + "This backend requires term encoding. Build the e-graph with `EGraph::with_backend(..).with_term_encoding()`." + )] + BackendRequiresTermEncoding, #[error("{0}\nTried to pop too much")] Pop(Span), #[error("{0}\nCommand should have failed.")] @@ -2967,6 +3379,41 @@ mod tests { } } + #[test] + fn unstable_fn_resolution_error_releases_panic_registration() { + let mut egraph = EGraph::default(); + egraph + .parse_and_run_program(None, "(sort Fn (UnstableFn (i64) i64))") + .unwrap(); + let register_probe = |egraph: &mut EGraph| { + egraph + .backend + .register_external_func(Box::new(core_relations::make_external_func( + |_state: &mut core_relations::ExecutionState<'_>, _args: &[Value]| None, + ))) + }; + let error = egraph + .parse_and_run_program(None, "(let $first (unstable-fn \"missing\"))") + .expect_err("missing unstable-fn target should return an error"); + assert!(error.to_string().contains("No resolution for \"missing\"")); + let reused = register_probe(&mut egraph); + egraph.backend.free_external_func(reused); + + let message = "unstable-fn over `missing` was applied in a context where its wrapped \ + function is not valid for this call site, if in a rule, add :naive." + .to_string(); + let shared = egraph.backend.new_panic(message); + assert_eq!(shared, reused); + let error = egraph + .parse_and_run_program(None, "(let $second (unstable-fn \"missing\"))") + .expect_err("missing unstable-fn target should return an error"); + assert!(error.to_string().contains("No resolution for \"missing\"")); + let occupied = register_probe(&mut egraph); + assert_ne!(occupied, shared); + egraph.backend.free_external_func(shared); + assert_eq!(register_probe(&mut egraph), shared); + } + #[test] fn test_user_defined_primitive() { let mut egraph = EGraph::default(); diff --git a/egglog/src/prelude.rs b/egglog/src/prelude.rs index 30194d0..c2fdaa5 100644 --- a/egglog/src/prelude.rs +++ b/egglog/src/prelude.rs @@ -862,11 +862,11 @@ impl Sort for BaseSortImpl { self.0.name() } - fn column_ty(&self, backend: &egglog_bridge::EGraph) -> ColumnTy { - ColumnTy::Base(backend.base_values().get_ty::()) + fn column_ty(&self, base_values: &BaseValues) -> ColumnTy { + ColumnTy::Base(base_values.get_ty::()) } - fn register_type(&self, backend: &mut egglog_bridge::EGraph) { + fn register_type(&self, backend: &mut dyn egglog_backend_trait::Backend) { backend.base_values_mut().register_type::(); } @@ -939,11 +939,11 @@ impl Sort for ContainerSortImpl { self.0.name() } - fn column_ty(&self, _backend: &egglog_bridge::EGraph) -> ColumnTy { + fn column_ty(&self, _base_values: &BaseValues) -> ColumnTy { ColumnTy::Id } - fn register_type(&self, backend: &mut egglog_bridge::EGraph) { + fn register_type(&self, backend: &mut dyn egglog_backend_trait::Backend) { backend.register_container_ty::(); } @@ -1040,6 +1040,23 @@ mod tests { Ok(egraph) } + #[test] + fn add_primitive_under_term_encoding() -> Result<(), Error> { + // A primitive registered through the Rust API AFTER construction must + // reach the term-encoding typechecker (a separate e-graph), not just the + // running e-graph — otherwise the encoder reports it as unbound. + let mut egraph = EGraph::new_with_term_encoding(); + add_literal_prim!(&mut egraph, "tripleit" = |x: i64| -> i64 { x * 3 }); + egraph.parse_and_run_program( + None, + "(function f (i64) i64 :merge (min old new)) +(set (f 0) (tripleit 7)) +(run 1) +(check (= (f 0) 21))", + )?; + Ok(()) + } + #[test] fn rust_api_query() -> Result<(), Error> { let mut egraph = build_test_database()?; diff --git a/egglog/src/proofs/proof_extraction.rs b/egglog/src/proofs/proof_extraction.rs index d885107..cb4cdb8 100644 --- a/egglog/src/proofs/proof_extraction.rs +++ b/egglog/src/proofs/proof_extraction.rs @@ -2,7 +2,8 @@ use crate::ast::FunctionSubtype; use crate::proofs::proof_encoding::ProofInstrumentor; use crate::proofs::proof_extractor::extract_root; use crate::proofs::proof_format::{Justification, ProofId, ProofStore, proof_store_from_term}; -use crate::{RawValues, Read, ResolvedCall, TermDag}; +use crate::{ResolvedCall, TermDag}; +use egglog_backend_trait::BackendExt; use thiserror::Error; #[derive(Debug, Error)] @@ -70,7 +71,7 @@ impl ProofInstrumentor<'_> { ) }) .clone(); - let proof_sort = self + let proof_function = self .egraph .functions .get(&proof_function_name) @@ -79,14 +80,12 @@ impl ProofInstrumentor<'_> { "proof table {proof_function_name} for constructor {} was not declared", func.name ) - }) - .schema - .output - .clone(); + }); + let proof_sort = proof_function.schema.output.clone(); let proof_value = self .egraph - .update_unchecked(|fs| fs.lookup(&proof_function_name, RawValues(vec![witness_value]))) - .unwrap() + .backend + .lookup_id(proof_function.backend_id, &[witness_value]) .unwrap_or_else(|| panic!("no proof recorded for constructor {}", func.name)); let proof_term_id = extract_root(self.egraph, &mut termdag, proof_value, proof_sort) diff --git a/egglog/src/proofs/proof_extractor.rs b/egglog/src/proofs/proof_extractor.rs index a7a5823..67a2dfb 100644 --- a/egglog/src/proofs/proof_extractor.rs +++ b/egglog/src/proofs/proof_extractor.rs @@ -2,6 +2,7 @@ use crate::ast::FunctionSubtype; use crate::termdag::{TermDag, TermId}; use crate::util::{HashMap, HashSet}; use crate::{ArcSort, EGraph, Value}; +use egglog_backend_trait::BackendExt; /// Root-directed extraction for proof terms. /// @@ -173,23 +174,29 @@ mod tests { use super::*; fn first_output(egraph: &EGraph, function_name: &str) -> Value { - let func = egraph.functions.get(function_name).unwrap(); + let func = egraph + .functions + .get(function_name) + .unwrap_or_else(|| panic!("function `{function_name}` was not declared")); let mut value = None; egraph.backend.for_each_while(func.backend_id, |row| { value = Some(row.vals[func.extraction_output_index()]); false }); - value.unwrap() + value.unwrap_or_else(|| panic!("function `{function_name}` has no rows")) } fn first_input(egraph: &EGraph, function_name: &str, input: usize) -> Value { - let func = egraph.functions.get(function_name).unwrap(); + let func = egraph + .functions + .get(function_name) + .unwrap_or_else(|| panic!("function `{function_name}` was not declared")); let mut value = None; egraph.backend.for_each_while(func.backend_id, |row| { value = Some(row.vals[input]); false }); - value.unwrap() + value.unwrap_or_else(|| panic!("function `{function_name}` has no rows")) } fn extract_to_string(egraph: &EGraph, value: Value, sort_name: &str) -> String { diff --git a/egglog/src/proofs/proof_tests.rs b/egglog/src/proofs/proof_tests.rs index 3103ffe..3f96877 100644 --- a/egglog/src/proofs/proof_tests.rs +++ b/egglog/src/proofs/proof_tests.rs @@ -207,14 +207,6 @@ mod tests { "proof-container-id" = |x: # (eq_container_sort)| -> # (eq_container_sort) { x }, validator ); - let validator = - |_: &mut TermDag, args: &[TermId]| -> Option { args.first().copied() }; - let original_typechecking = egraph.proof_state.original_typechecking.as_mut().unwrap(); - add_primitive_with_validator!( - &mut **original_typechecking, - "proof-container-id" = |x: # (eq_container_sort)| -> # (eq_container_sort) { x }, - validator - ); egraph .parse_and_run_program( diff --git a/egglog/src/scheduler.rs b/egglog/src/scheduler.rs index 8222c36..43d2155 100644 --- a/egglog/src/scheduler.rs +++ b/egglog/src/scheduler.rs @@ -1,7 +1,8 @@ use std::sync::Arc; use std::sync::Mutex; -use core_relations::{ExecutionState, ExternalFunction, Value}; +use core_relations::{ExecutionState, ExternalFunction, ExternalFunctionId, Value}; +use egglog_backend_trait::{BackendExt, ReadMode, RuleSetRun, RuleValue}; use egglog_bridge::{ ColumnTy, DefaultVal, FunctionConfig, FunctionId, MergeFn, RuleId, TableAction, }; @@ -188,90 +189,114 @@ impl EGraph { } } + if !self.backend.as_any().is::() { + return Err(Error::BackendError( + "scheduler match instantiation requires the reference bridge backend".into(), + )); + } + let mut rules = Vec::new(); let rulesets = std::mem::take(&mut self.rulesets); collect_rules(ruleset, &rulesets, &mut rules); let mut schedulers = std::mem::take(&mut self.schedulers); - - // Step 1: build all the query/action rules and worklist if have not already - let record = &mut schedulers[scheduler_id]; - rules.iter().for_each(|(id, rule)| { - record - .rule_info - .entry((*id).to_owned()) - .or_insert_with(|| SchedulerRuleInfo::new(self, rule, id)); - }); - - // Step 2: run all the queries for one iteration - let query_rules = rules - .iter() - .filter_map(|(rule_id, _rule)| { - let rule_info = record.rule_info.get(rule_id).unwrap(); - - if rule_info.should_seek { - Some(rule_info.query_rule) - } else { - None + let result = (|| -> Result { + // Step 1: build all the query/action rules and worklist if have not already + let record = &mut schedulers[scheduler_id]; + for (id, rule) in &rules { + if !record.rule_info.contains_key(id) { + let info = SchedulerRuleInfo::new(self, rule, id)?; + record.rule_info.insert(id.clone(), info); } - }) - .collect::>(); - - let query_iter_report = self - .backend - .run_rules(&query_rules) - .map_err(|e| Error::BackendError(e.to_string()))?; - - // Step 3: let the scheduler decide which matches need to be kept - self.backend.with_execution_state(|state| { - for (rule_id, _rule) in rules.iter() { - let rule_info = record.rule_info.get_mut(rule_id).unwrap(); - - let matches: Vec = - std::mem::take(rule_info.matches.lock().unwrap().as_mut()); - let mut matches = Matches::new(matches, rule_info.free_vars.clone()); - rule_info.should_seek = - record - .scheduler - .filter_matches(rule_id, ruleset, &mut matches); - let table_action = TableAction::new(&self.backend, rule_info.decided); - *rule_info.matches.lock().unwrap() = matches.instantiate(state, &table_action); } - }); - self.backend.flush_updates(); - // Step 4: run the action rules - let action_rules = rules - .iter() - .map(|(rule_id, _rule)| { - let rule_info = record.rule_info.get(rule_id).unwrap(); - rule_info.action_rule - }) - .collect::>(); - let action_iter_report = self - .backend - .run_rules(&action_rules) - .map_err(|e| Error::BackendError(e.to_string()))?; - - // Step 5: combine the reports - let mut query_report = RunReport::singleton(ruleset, query_iter_report); - let mut action_report = RunReport::singleton(ruleset, action_iter_report); - - // query matches don't count - query_report.updated = false; - query_report.num_matches_per_rule.clear(); - // Scheduler state should not count as database progress. Instead it - // determines whether a no-op iteration can be treated as fully stopped. - action_report.can_stop = !action_report.updated && { - let rule_ids = rules.iter().map(|(id, _)| id.as_str()).collect::>(); - record.scheduler.can_stop(&rule_ids, ruleset) - }; + // Step 2: run all the queries for one iteration + let query_rules = rules + .iter() + .filter_map(|(rule_id, _rule)| { + let rule_info = record.rule_info.get(rule_id).unwrap(); - query_report.union(action_report); + if rule_info.should_seek { + Some(rule_info.query_rule) + } else { + None + } + }) + .collect::>(); + + let query_iter_report = self + .backend + .run_rules(RuleSetRun { + name: Some(ruleset), + rules: &query_rules, + }) + .map_err(|e| Error::BackendError(e.to_string()))?; + + // Step 3: let the scheduler decide which matches need to be kept + let bridge = self + .backend + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::BackendError( + "scheduler match instantiation requires the reference bridge backend" + .into(), + ) + })?; + self.backend.with_execution_state(|state| { + for (rule_id, _rule) in &rules { + let rule_info = record.rule_info.get_mut(rule_id).unwrap(); + + let matches: Vec = + std::mem::take(rule_info.matches.lock().unwrap().as_mut()); + let mut matches = Matches::new(matches, rule_info.free_vars.clone()); + rule_info.should_seek = + record + .scheduler + .filter_matches(rule_id, ruleset, &mut matches); + let table_action = TableAction::new(bridge, rule_info.decided); + *rule_info.matches.lock().unwrap() = matches.instantiate(state, &table_action); + } + }); + self.backend.flush_updates(); + + // Step 4: run the action rules + let action_rules = rules + .iter() + .map(|(rule_id, _rule)| { + let rule_info = record.rule_info.get(rule_id).unwrap(); + rule_info.action_rule + }) + .collect::>(); + let action_iter_report = self + .backend + .run_rules(RuleSetRun { + name: Some(ruleset), + rules: &action_rules, + }) + .map_err(|e| Error::BackendError(e.to_string()))?; + + // Step 5: combine the reports + let mut query_report = RunReport::singleton(ruleset, query_iter_report); + let mut action_report = RunReport::singleton(ruleset, action_iter_report); + + // query matches don't count + query_report.updated = false; + query_report.num_matches_per_rule.clear(); + // Scheduler state should not count as database progress. Instead it + // determines whether a no-op iteration can be treated as fully stopped. + action_report.can_stop = !action_report.updated && { + let rule_ids = rules.iter().map(|(id, _)| id.as_str()).collect::>(); + record.scheduler.can_stop(&rule_ids, ruleset) + }; + + query_report.union(action_report); + Ok(query_report) + })(); self.rulesets = rulesets; self.schedulers = schedulers; - Ok(query_report) + result } } @@ -295,6 +320,12 @@ struct SchedulerRuleInfo { free_vars: Vec, } +struct SchedulerRuleBuild { + collect_matches: ExternalFunctionId, + query_rule: Option, + decided: Option, +} + struct CollectMatches { matches: Arc>>, } @@ -320,86 +351,266 @@ impl ExternalFunction for CollectMatches { } } +impl SchedulerRuleBuild { + fn rollback(self, egraph: &mut EGraph, error: Error) -> Error { + let table_result = self.decided.map_or(Ok(()), |table| { + let bridge = egraph + .backend + .as_any_mut() + .downcast_mut::() + .ok_or_else(|| { + Error::BackendError( + "scheduler rollback requires the reference bridge backend".into(), + ) + })?; + bridge + .remove_last_table(table) + .map_err(|error| Error::BackendError(error.to_string())) + }); + if let Some(rule) = self.query_rule { + egraph.backend.free_rule(rule); + } + egraph.backend.free_external_func(self.collect_matches); + + match table_result { + Ok(()) => error, + Err(rollback_error) => Error::BackendError(format!( + "{error}; scheduler rule rollback also failed: {rollback_error}" + )), + } + } +} + impl SchedulerRuleInfo { - fn new(egraph: &mut EGraph, rule: &ResolvedCoreRule, name: &str) -> SchedulerRuleInfo { - let free_vars = rule.head.get_free_vars().into_iter().collect::>(); + fn new( + egraph: &mut EGraph, + rule: &ResolvedCoreRule, + name: &str, + ) -> Result { + let free_vars = rule.head.free_vars(); let unit_type = egraph.backend.base_values().get_ty::<()>(); let unit = egraph.backend.base_values().get(()); - let unit_entry = egraph.backend.base_value_constant(()); + let unit_entry = GenericAtomTerm::Literal( + rule.span.clone(), + RuleValue { + value: unit, + ty: ColumnTy::Base(unit_type), + }, + ); let matches = Arc::new(Mutex::new(Vec::new())); let collect_matches = egraph .backend .register_external_func(Box::new(CollectMatches::new(matches.clone()))); + let mut build = SchedulerRuleBuild { + collect_matches, + query_rule: None, + decided: None, + }; let schema = free_vars .iter() - .map(|v| v.sort.column_ty(&egraph.backend)) + .map(|v| v.sort.column_ty(egraph.backend.base_values())) .chain(std::iter::once(ColumnTy::Base(unit_type))) .collect(); - let decided = egraph.backend.add_table(FunctionConfig { - schema, - default: DefaultVal::Const(unit), - merge: MergeFn::AssertEq, - name: "backend".to_string(), - can_subsume: false, - }); - // Step 1: build the query rule let mut qrule_builder = BackendRule::new( - egraph.backend.new_rule(name, true), + &mut *egraph.backend, &egraph.functions, &egraph.type_info, false, // seminaive query: Pure/Write contexts ); - qrule_builder.query(&rule.body, false); + if let Err(error) = qrule_builder.query(&rule.body, false) { + drop(qrule_builder); + return Err(build.rollback(egraph, error)); + } let entries = free_vars .iter() .map(|fv| qrule_builder.entry(&GenericAtomTerm::Var(span!(), fv.clone()))) - .collect::>(); - let _var = qrule_builder.rb.call_external_func( + .collect::, _>>(); + let entries = match entries { + Ok(entries) => entries, + Err(error) => { + drop(qrule_builder); + return Err(build.rollback(egraph, error)); + } + }; + qrule_builder.call_external_func( + rule.span.clone(), collect_matches, - &entries, + "collect_matches", + entries, ColumnTy::Base(unit_type), - || "collect_matches".to_string(), ); - let qrule_id = qrule_builder.build(); + let qrule_id = match qrule_builder.try_build(name, true, false, rule.span.clone()) { + Ok(rule) => rule, + Err(error) => return Err(build.rollback(egraph, error)), + }; + build.query_rule = Some(qrule_id); + + let decided = egraph.backend.add_table(FunctionConfig { + schema, + default: DefaultVal::Const(unit), + merge: MergeFn::Old, + name: "backend".to_string(), + can_subsume: false, + }); + build.decided = Some(decided); // Step 2: build the action rule let mut arule_builder = BackendRule::new( - egraph.backend.new_rule(name, false), + &mut *egraph.backend, &egraph.functions, &egraph.type_info, true, // action rule reads the DB: Read/Full contexts ); - let mut entries = free_vars + let entries = free_vars .iter() .map(|fv| arule_builder.entry(&GenericAtomTerm::Var(span!(), fv.clone()))) - .collect::>(); + .collect::, _>>(); + let mut entries = match entries { + Ok(entries) => entries, + Err(error) => { + drop(arule_builder); + return Err(build.rollback(egraph, error)); + } + }; entries.push(unit_entry); - arule_builder - .rb - .query_table(decided, &entries, None) - .unwrap(); - arule_builder.actions(&rule.head).unwrap(); + arule_builder.query_table(rule.span.clone(), decided, entries.clone(), ReadMode::All); + if let Err(error) = arule_builder.actions(&rule.head) { + drop(arule_builder); + return Err(build.rollback(egraph, error)); + } // Remove the entry as it's now done entries.pop(); - arule_builder.rb.remove(decided, &entries); - let arule_id = arule_builder.build(); + arule_builder.remove(rule.span.clone(), decided, "backend", entries); + let arule_id = match arule_builder.try_build(name, false, false, rule.span.clone()) { + Ok(rule) => rule, + Err(error) => return Err(build.rollback(egraph, error)), + }; - SchedulerRuleInfo { + Ok(SchedulerRuleInfo { free_vars, query_rule: qrule_id, action_rule: arule_id, matches, decided, should_seek: true, - } + }) } } #[cfg(test)] mod test { use super::*; + use egglog_backend_trait::RuleSpec; + + fn scheduler_rule_fixture() -> (EGraph, ResolvedCoreRule) { + let mut egraph = EGraph::default(); + egraph + .parse_and_run_program( + None, + r#" + (ruleset test) + (relation R (i64)) + (function F (i64) i64 :no-merge) + (rule ((R x)) ((set (F x) x)) :ruleset test :name "test-rule") + "#, + ) + .unwrap(); + let Ruleset::Rules(rules) = &egraph.rulesets["test"] else { + unreachable!() + }; + let rule = rules["test-rule"].0.clone(); + (egraph, rule) + } + + fn backend_probe(egraph: &mut EGraph) -> ([ExternalFunctionId; 2], [RuleId; 2], FunctionId) { + let external = std::array::from_fn(|_| { + egraph + .backend + .register_external_func(Box::new(core_relations::make_external_func( + |state: &mut ExecutionState<'_>, _args: &[Value]| { + Some(state.base_values().get(())) + }, + ))) + }); + let rules = std::array::from_fn(|_| { + egraph + .backend + .add_rule(RuleSpec { + name: "probe".into(), + seminaive: false, + no_decomp: false, + core: egglog_ast::core::GenericCoreRule { + span: span!(), + body: Default::default(), + head: Default::default(), + }, + }) + .unwrap() + }); + let unit_type = egraph.backend.base_values().get_ty::<()>(); + let unit = egraph.backend.base_values().get(()); + let table = egraph.backend.add_table(FunctionConfig { + schema: vec![ColumnTy::Base(unit_type)], + default: DefaultVal::Const(unit), + merge: MergeFn::AssertEq, + name: "probe".into(), + can_subsume: false, + }); + (external, rules, table) + } + + fn assert_failed_construction_restores_backend(mut failed: EGraph, rule: &ResolvedCoreRule) { + let (mut baseline, _) = scheduler_rule_fixture(); + assert!(SchedulerRuleInfo::new(&mut failed, rule, "test-rule").is_err()); + + let bridge = failed + .backend + .as_any() + .downcast_ref::() + .unwrap(); + assert!( + bridge + .action_registry() + .read() + .unwrap() + .lookup_table("backend") + .is_none() + ); + assert_eq!(backend_probe(&mut baseline), backend_probe(&mut failed)); + } + + #[test] + fn scheduler_query_failure_rolls_back_backend_resources() { + let (egraph, mut rule) = scheduler_rule_fixture(); + rule.body.atoms[0].args.pop(); + assert_failed_construction_restores_backend(egraph, &rule); + } + + #[test] + fn scheduler_action_failure_rolls_back_backend_resources() { + let (egraph, mut rule) = scheduler_rule_fixture(); + let (span, function, args) = rule + .head + .0 + .iter() + .find_map(|action| match action { + crate::core::GenericCoreAction::Set(span, function, args, _) => { + Some((span.clone(), function.clone(), args.clone())) + } + _ => None, + }) + .unwrap(); + rule.head.0 = vec![crate::core::GenericCoreAction::Change( + span, + crate::ast::Change::Subsume, + function, + args, + )]; + + assert_failed_construction_restores_backend(egraph, &rule); + } #[derive(Clone)] struct FirstNScheduler { diff --git a/egglog/src/serialize.rs b/egglog/src/serialize.rs index 83ab829..80ce4b5 100644 --- a/egglog/src/serialize.rs +++ b/egglog/src/serialize.rs @@ -241,7 +241,7 @@ impl EGraph { // Canonicalize the value first so that we always use the canonical e-class ID let value = self .backend - .get_canon_repr(value, sort.column_ty(&self.backend)); + .get_canon_repr(value, sort.column_ty(self.backend.base_values())); assert!( !sort.name().to_string().contains('-'), "Tag cannot contain '-' when serializing" diff --git a/egglog/src/sort/fn.rs b/egglog/src/sort/fn.rs index 0f035bf..8255c57 100644 --- a/egglog/src/sort/fn.rs +++ b/egglog/src/sort/fn.rs @@ -13,6 +13,7 @@ use std::any::TypeId; use std::sync::Mutex; use crate::exec_state::Internal; +use egglog_backend_trait::BackendExt; use enum_map::EnumMap; use super::*; @@ -152,11 +153,11 @@ impl Sort for FunctionSort { &self.name } - fn column_ty(&self, _backend: &egglog_bridge::EGraph) -> ColumnTy { + fn column_ty(&self, _base_values: &BaseValues) -> ColumnTy { ColumnTy::Id } - fn register_type(&self, backend: &mut egglog_bridge::EGraph) { + fn register_type(&self, backend: &mut dyn egglog_backend_trait::Backend) { backend.register_container_ty::(); backend .base_values_mut() diff --git a/egglog/src/sort/mod.rs b/egglog/src/sort/mod.rs index 1bcc67f..939c58a 100644 --- a/egglog/src/sort/mod.rs +++ b/egglog/src/sort/mod.rs @@ -52,7 +52,7 @@ pub trait Sort: Any + Send + Sync + Debug { fn name(&self) -> &str; /// Returns the backend-specific column type. See [`ColumnTy`]. - fn column_ty(&self, backend: &egglog_bridge::EGraph) -> ColumnTy; + fn column_ty(&self, base_values: &BaseValues) -> ColumnTy; /// return the inner sorts if a container sort /// remember that containers can contain containers @@ -65,7 +65,7 @@ pub trait Sort: Any + Send + Sync + Debug { } } - fn register_type(&self, backend: &mut egglog_bridge::EGraph); + fn register_type(&self, backend: &mut dyn egglog_backend_trait::Backend); fn as_arc_any(self: Arc) -> Arc; @@ -177,11 +177,11 @@ impl Sort for EqSort { &self.name } - fn column_ty(&self, _backend: &egglog_bridge::EGraph) -> ColumnTy { + fn column_ty(&self, _base_values: &BaseValues) -> ColumnTy { ColumnTy::Id } - fn register_type(&self, _backend: &mut egglog_bridge::EGraph) {} + fn register_type(&self, _backend: &mut dyn egglog_backend_trait::Backend) {} fn as_arc_any(self: Arc) -> Arc { self diff --git a/egglog/src/typechecking.rs b/egglog/src/typechecking.rs index 39f3749..fc03e0a 100644 --- a/egglog/src/typechecking.rs +++ b/egglog/src/typechecking.rs @@ -3,7 +3,7 @@ use std::hash::Hasher; use crate::Context; use crate::proofs::proof_container_rebuild::register_container_rebuild_from_spec; use crate::{ - core::{CoreActionContext, CoreRule, GenericActionsExt, ResolvedCall}, + core::{CoreActionContext, CoreRule, GenericActionsExt, QueryConstraints, ResolvedCall}, *, }; use ast::{ @@ -257,14 +257,22 @@ impl EGraph { /// Add a user-defined sort to the e-graph. pub fn add_arcsort(&mut self, sort: ArcSort, span: Span) -> Result<(), TypeError> { - sort.register_type(&mut self.backend); + sort.register_type(self.backend.as_mut()); let name = sort.name(); match self.type_info.sorts.entry(name.to_owned()) { HEntry::Occupied(_) => Err(TypeError::SortAlreadyBound(name.to_owned(), span)), HEntry::Vacant(e) => { e.insert(sort.clone()); + // A sort's primitives already reach the term-encoding typechecker + // through its OWN `add_arcsort` when it typechecks the sort + // command, so don't propagate them again from here (that would + // double-register and make primitive resolution ambiguous). + // Detach the typechecker while the sort registers, so only direct + // `add_*_primitive` calls propagate to it. + let saved = self.proof_state.original_typechecking.take(); sort.register_primitives(self); + self.proof_state.original_typechecking = saved; Ok(()) } } @@ -282,9 +290,14 @@ impl EGraph { where T: PurePrim + Clone, { - self.register_per_context(x, validator, PureState::valid_contexts(), |x, ctx| { - Box::new(PurePrimWrapper { prim: x, ctx }) - }); + self.register_per_context( + x, + validator, + PureState::valid_contexts(), + |backend, x, ctx| { + backend.register_external_func(Box::new(PurePrimWrapper { prim: x, ctx })) + }, + ); } /// Register a [`WritePrim`]. Pass `None` for the validator if not @@ -327,14 +340,20 @@ impl EGraph { T: Primitive + Clone, S: RegistryWrap + 'static, { - let registry = self.backend.action_registry().clone(); - self.register_per_context(x, validator, valid_ctxs, move |x, ctx| { - Box::new(RegistryPrimWrapper:: { - prim: x, - registry: registry.clone(), - ctx, - _wrap: std::marker::PhantomData, - }) + self.register_per_context(x, validator, valid_ctxs, |backend, x, ctx| { + if let Some(registry) = backend.action_registry().cloned() { + backend.register_external_func(Box::new(RegistryPrimWrapper:: { + prim: x, + registry, + ctx, + _wrap: std::marker::PhantomData, + })) + } else { + let name = x.name().to_owned(); + backend.new_panic(format!( + "primitive {name} in {ctx:?} context requires a backend action registry" + )) + } }); } @@ -355,25 +374,39 @@ impl EGraph { mut build_wrapper: F, ) where T: Primitive + Clone, - F: FnMut(T, Context) -> Box, + F: FnMut(&mut dyn Backend, T, Context) -> ExternalFunctionId, { - let primitive: Arc = Arc::new(x.clone()); - let name = primitive.name().to_owned(); - let context_ids = EnumMap::from_fn(|ctx| { - valid_ctxs.contains(&ctx).then(|| { - self.backend - .register_external_func(build_wrapper(x.clone(), ctx)) - }) - }); - self.type_info - .primitives - .entry(name) - .or_default() - .push(PrimitiveWithId { - primitive, - validator, - context_ids, + // Register on this e-graph AND every term-encoding typechecker down the + // chain. Each typechecker is a separate e-graph that typechecks the + // encoded program (see `typecheck_program`); a primitive added after + // construction is otherwise unknown to it and reported as unbound. A + // typechecker only typechecks and never evaluates, so the wrapper's + // runtime state is irrelevant there, and — since both e-graphs register + // the same built-ins during construction — a primitive added to both + // gets the same `ExternalFunctionId`. + let mut eg: &mut EGraph = self; + loop { + let primitive: Arc = Arc::new(x.clone()); + let name = primitive.name().to_owned(); + let context_ids = EnumMap::from_fn(|ctx| { + valid_ctxs + .contains(&ctx) + .then(|| build_wrapper(eg.backend.as_mut(), x.clone(), ctx)) }); + eg.type_info + .primitives + .entry(name) + .or_default() + .push(PrimitiveWithId { + primitive, + validator: validator.clone(), + context_ids, + }); + match eg.proof_state.original_typechecking.as_deref_mut() { + Some(next) => eg = next, + None => break, + } + } } } @@ -902,7 +935,7 @@ impl TypeInfo { let (query, mapped_query) = Facts(body.clone()).to_query(self, symbol_gen); constraints.extend(query.get_constraints(self, query_ctx)?); - let mut binding = query.get_vars(); + let mut binding = query.vars().collect::>(); // We lower to core actions with `union_to_set_optimization` // later in the pipeline. For typechecking we do not need it. let mut ctx = CoreActionContext::new(self, &mut binding, symbol_gen, false); diff --git a/egglog/tests/math-microbenchmark-mini.egg b/egglog/tests/math-microbenchmark-mini.egg new file mode 100644 index 0000000..fc49f2a --- /dev/null +++ b/egglog/tests/math-microbenchmark-mini.egg @@ -0,0 +1,75 @@ +;; The full benchmark runs 11 rounds. Eight preserves its rule and seed shape +;; while leaving enough headroom for the DD backend. +(datatype Math + (Diff Math Math) + (Integral Math Math) + + (Add Math Math) + (Sub Math Math) + (Mul Math Math) + (Div Math Math) + (Pow Math Math) + (Ln Math) + (Sqrt Math) + + (Sin Math) + (Cos Math) + + (Const i64) + (Var String)) + +(rewrite (Add a b) (Add b a)) +(rewrite (Mul a b) (Mul b a)) +(rewrite (Add a (Add b c)) (Add (Add a b) c)) +(rewrite (Mul a (Mul b c)) (Mul (Mul a b) c)) + +(rewrite (Sub a b) (Add a (Mul (Const -1) b))) +;; (rewrite (Div a b) (Mul a (Pow b (Const -1))) :when ((is-not-zero b))) + +(rewrite (Add a (Const 0)) a) +(rewrite (Mul a (Const 0)) (Const 0)) +(rewrite (Mul a (Const 1)) a) + +(rewrite (Sub a a) (Const 0)) + +(rewrite (Mul a (Add b c)) (Add (Mul a b) (Mul a c))) +(rewrite (Add (Mul a b) (Mul a c)) (Mul a (Add b c))) + +(rewrite (Mul (Pow a b) (Pow a c)) (Pow a (Add b c))) +(rewrite (Pow x (Const 1)) x) +(rewrite (Pow x (Const 2)) (Mul x x)) + +(rewrite (Diff x (Add a b)) (Add (Diff x a) (Diff x b))) +(rewrite (Diff x (Mul a b)) (Add (Mul a (Diff x b)) (Mul b (Diff x a)))) + +(rewrite (Diff x (Sin x)) (Cos x)) +(rewrite (Diff x (Cos x)) (Mul (Const -1) (Sin x))) + +(rewrite (Integral (Const 1) x) x) +(rewrite (Integral (Cos x) x) (Sin x)) +(rewrite (Integral (Sin x) x) (Mul (Const -1) (Cos x))) +(rewrite (Integral (Add f g) x) (Add (Integral f x) (Integral g x))) +(rewrite (Integral (Sub f g) x) (Sub (Integral f x) (Integral g x))) +(rewrite (Integral (Mul a b) x) +(Sub (Mul a (Integral b x)) + (Integral (Mul (Diff x a) (Integral b x)) x))) +(Integral (Ln (Var "x")) (Var "x")) +(Integral (Add (Var "x") (Cos (Var "x"))) (Var "x")) +(Integral (Mul (Cos (Var "x")) (Var "x")) (Var "x")) +(Diff (Var "x") (Add (Const 1) (Mul (Const 2) (Var "x")))) +(Diff (Var "x") (Sub (Pow (Var "x") (Const 3)) (Mul (Const 7) (Pow (Var "x") (Const 2))))) +(Add (Mul (Var "y") (Add (Var "x") (Var "y"))) (Sub (Add (Var "x") (Const 2)) (Add (Var "x") (Var "x")))) +(Div (Const 1) + (Sub (Div (Add (Const 1) + (Sqrt (Var "five"))) + (Const 2)) + (Div (Sub (Const 1) + (Sqrt (Var "five"))) + (Const 2)))) +(run 8) +(print-size Add) +(print-size Mul) + +(print-size) + +(print-stats) diff --git a/egglog/tests/snapshots/files__shared_snapshot_math_microbenchmark_mini.snap b/egglog/tests/snapshots/files__shared_snapshot_math_microbenchmark_mini.snap new file mode 100644 index 0000000..5aca55b --- /dev/null +++ b/egglog/tests/snapshots/files__shared_snapshot_math_microbenchmark_mini.snap @@ -0,0 +1,32 @@ +--- +source: tests/files.rs +expression: snapshot_content_across_treatments +--- +2977 +3516 +((Add 2977) + (Const 5) + (Cos 1) + (Diff 338) + (Div 3) + (Integral 782) + (Ln 1) + (Mul 3516) + (Pow 2) + (Sin 1) + (Sqrt 1) + (Sub 483) + (Var 3)) +((Add 2977) + (Const 5) + (Cos 1) + (Diff 338) + (Div 3) + (Integral 782) + (Ln 1) + (Mul 3516) + (Pow 2) + (Sin 1) + (Sqrt 1) + (Sub 483) + (Var 3)) diff --git a/egglog/tests/typed_primitive.rs b/egglog/tests/typed_primitive.rs index 159e2c9..5a48322 100644 --- a/egglog/tests/typed_primitive.rs +++ b/egglog/tests/typed_primitive.rs @@ -514,18 +514,20 @@ fn missing_sort_panics_with_type_name() { /// per application context, and duplicate same-signature registrations are /// ambiguous for every context where more than one runtime id matches. #[test] -#[should_panic(expected = "Ambiguous primitive resolution")] -fn unstable_fn_duplicate_primitive_registration_panics_on_build() { +fn unstable_fn_duplicate_primitive_registration_returns_error() { let mut egraph = EGraph::default(); egraph.add_pure_primitive(PureEcho("dup-echo"), None); egraph.add_pure_primitive(PureEcho("dup-echo"), None); - let _ = egraph.parse_and_run_program( - None, - "(sort Fn (UnstableFn (i64) i64))\n\ + let error = egraph + .parse_and_run_program( + None, + "(sort Fn (UnstableFn (i64) i64))\n\ (let $f (unstable-fn \"dup-echo\"))\n\ (check (= (unstable-app $f 7) 7))", - ); + ) + .expect_err("duplicate primitive resolution should return an error"); + assert!(error.to_string().contains("Ambiguous primitive resolution")); } // --- 4x4 unstable-app dispatch matrix --- diff --git a/samply_analysis.py b/samply_analysis.py new file mode 100644 index 0000000..ef0763b --- /dev/null +++ b/samply_analysis.py @@ -0,0 +1,447 @@ +"""Read, summarize, and present saved Samply profile artifacts. + +Profile execution, caching, and viewer lifecycle remain owned by ``bench.py``. +""" + +from __future__ import annotations + +import gzip +import json +import os +import re +import shlex +import shutil +import subprocess +from collections import defaultdict +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal, cast + +from rich import box +from rich.console import Console +from rich.table import Table +from rich.text import Text + +ATOS_BATCH_SIZE = 512 + + +@dataclass(frozen=True) +class ProfileFunctionCpu: + name: str + cpu_seconds: float + + +@dataclass(frozen=True) +class ProfileCpuSummary: + observed_cpu_seconds: float + application_cpu_seconds: float + other_library_cpu_seconds: float + unattributed_cpu_seconds: float + symbolized_application_cpu_seconds: float + functions: tuple[ProfileFunctionCpu, ...] + warnings: tuple[str, ...] + + +@dataclass(frozen=True) +class ProfileReport: + artifact: Path + cache_status: Literal["hit", "recorded"] + workload: str + backend: str + treatment: str + top: int + cpu_summary: ProfileCpuSummary | None + + +@dataclass(frozen=True) +class LeafSample: + cpu_seconds: float + library_index: int | None + function_name: str | None + relative_address: int | None + + +def read_artifact(path: Path) -> dict[str, Any]: + try: + if not path.is_file() or path.stat().st_size == 0: + raise ValueError(f"profile artifact is missing or empty: {path}") + with path.open("rb") as handle: + if handle.read(2) != b"\x1f\x8b": + raise ValueError(f"profile artifact is not gzip-compressed: {path}") + except OSError as error: + raise ValueError(f"could not read profile artifact {path}: {error}") from error + try: + with gzip.open(path, "rt", encoding="utf-8") as handle: + value = json.load(handle) + except (OSError, UnicodeDecodeError, json.JSONDecodeError) as error: + raise ValueError(f"could not parse profile artifact {path}: {error}") from error + expected_sections = {"meta": dict, "libs": list, "threads": list} + if not isinstance(value, dict) or not all( + isinstance(value.get(key), expected) for key, expected in expected_sections.items() + ): + raise ValueError(f"profile artifact has an unsupported structure: {path}") + return cast(dict[str, Any], value) + + +def _index(values: Sequence[Any], index: object) -> Any | None: + if isinstance(index, bool) or not isinstance(index, int) or index < 0 or index >= len(values): + return None + return values[index] + + +def _cpu_seconds_scale(profile: Mapping[str, Any]) -> float: + try: + unit = profile["meta"]["sampleUnits"]["threadCPUDelta"] + except (KeyError, TypeError) as error: + raise ValueError("profile has no threadCPUDelta unit") from error + scales = { + "ns": 1e-9, + "nanoseconds": 1e-9, + "us": 1e-6, + "\N{MICRO SIGN}s": 1e-6, + "microseconds": 1e-6, + "ms": 1e-3, + "milliseconds": 1e-3, + "s": 1.0, + "seconds": 1.0, + } + if not isinstance(unit, str) or unit not in scales: + raise ValueError(f"unsupported profile threadCPUDelta unit: {unit!r}") + return scales[unit] + + +def _application_library(profile: Mapping[str, Any], binary: Path) -> tuple[int, str | None]: + libraries = cast(list[Any], profile["libs"]) + binary_path = binary.resolve() + exact_matches: list[tuple[int, Mapping[str, Any]]] = [] + name_matches: list[tuple[int, Mapping[str, Any]]] = [] + for index, value in enumerate(libraries): + if not isinstance(value, dict): + continue + library = cast(Mapping[str, Any], value) + path = library.get("path") + name = library.get("name") + if isinstance(path, str) and Path(path).resolve() == binary_path: + exact_matches.append((index, library)) + if name == binary.name or (isinstance(path, str) and Path(path).name == binary.name): + name_matches.append((index, library)) + matches = exact_matches if exact_matches else name_matches + if len(matches) != 1: + raise ValueError(f"could not identify application library for {binary}") + index, library = matches[0] + arch = library.get("arch") + return (index, arch if isinstance(arch, str) and arch else None) + + +def _thread_leaf_samples( + thread: Mapping[str, Any], + scale: float, +) -> Iterator[LeafSample]: + samples = thread["samples"] + stacks = thread["stackTable"] + frames = thread["frameTable"] + functions = thread["funcTable"] + resources = thread["resourceTable"] + strings = thread["stringArray"] + for stack, delta in zip(samples["stack"][1:], samples["threadCPUDelta"][1:], strict=False): + if isinstance(delta, bool) or not isinstance(delta, (int, float)) or delta <= 0: + continue + frame = _index(stacks["frame"], stack) + function = _index(frames["func"], frame) + resource = _index(functions["resource"], function) + library = _index(resources["lib"], resource) + name = _index(strings, _index(functions["name"], function)) + address = _index(frames["address"], frame) + yield LeafSample( + cpu_seconds=float(delta) * scale, + library_index=library if isinstance(library, int) and not isinstance(library, bool) else None, + function_name=name if isinstance(name, str) else None, + relative_address=( + address if isinstance(address, int) and not isinstance(address, bool) and address >= 0 else None + ), + ) + + +def _normalize_rust_name(name: str) -> str: + for encoded, decoded in { + "$SP$": "@", + "$BP$": "*", + "$RF$": "&", + "$LT$": "<", + "$GT$": ">", + "$LP$": "(", + "$RP$": ")", + "$C$": ",", + }.items(): + name = name.replace(encoded, decoded) + + def decode_unicode(match: re.Match[str]) -> str: + try: + return chr(int(match.group(1), 16)) + except (ValueError, OverflowError): + return match.group(0) + + name = re.sub(r"\$u([0-9a-fA-F]+)\$", decode_unicode, name).replace("..", "::") + return name[1:] if name.startswith("_<") else name + + +def _normalize_atos_symbol(value: str, image_name: str) -> str | None: + name = value.strip().split(f" (in {image_name})", 1)[0] + if not name or name == "" or re.fullmatch(r"0x[0-9a-fA-F]+", name): + return None + return _normalize_rust_name(name) + + +def demangle_rust_v0_symbols(symbols: Mapping[int, str]) -> tuple[dict[int, str], tuple[str, ...]]: + demangled = dict(symbols) + raw_symbols = [(address, symbol) for address, symbol in symbols.items() if re.match(r"^_+R", symbol)] + if not raw_symbols: + return (demangled, ()) + xcrun = shutil.which("xcrun") + if xcrun is None: + for address, _ in raw_symbols: + demangled.pop(address, None) + return (demangled, ("llvm-cxxfilt is unavailable; Rust v0 symbols are incomplete",)) + try: + completed = subprocess.run( + [xcrun, "llvm-cxxfilt"], + input="\n".join(symbol for _, symbol in raw_symbols) + "\n", + check=True, + capture_output=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError) as error: + for address, _ in raw_symbols: + demangled.pop(address, None) + return (demangled, (f"Rust v0 symbol demangling failed: {error}",)) + for (address, raw_symbol), line in zip(raw_symbols, completed.stdout.splitlines(), strict=False): + symbol = _normalize_rust_name(line.strip()) + if symbol and symbol != raw_symbol and not re.match(r"^_+R", symbol): + demangled[address] = symbol + else: + demangled.pop(address, None) + return (demangled, ()) + + +def _symbolicate( + binary: Path, + arch: str | None, + addresses: Sequence[int], +) -> tuple[dict[int, str], tuple[str, ...]]: + unique_addresses = sorted(set(addresses)) + if not unique_addresses: + return ({}, ()) + atos = shutil.which("atos") + if atos is None: + return ({}, ("atos is unavailable; application symbols are incomplete",)) + + symbols: dict[int, str] = {} + warnings: list[str] = [] + for start in range(0, len(unique_addresses), ATOS_BATCH_SIZE): + batch = unique_addresses[start : start + ATOS_BATCH_SIZE] + command = [atos, "-offset", "-o", str(binary)] + if arch is not None: + command.extend(["-arch", arch]) + command.extend(hex(address) for address in batch) + try: + completed = subprocess.run(command, check=True, capture_output=True, text=True) + except (OSError, subprocess.CalledProcessError) as error: + warnings.append(f"atos symbolication failed: {error}") + continue + for address, line in zip(batch, completed.stdout.splitlines(), strict=False): + symbol = _normalize_atos_symbol(line, binary.name) + if symbol is not None: + symbols[address] = symbol + symbols, demangle_warnings = demangle_rust_v0_symbols(symbols) + warnings.extend(demangle_warnings) + unresolved = len(unique_addresses) - len(symbols) + if unresolved: + warnings.append(f"could not resolve {unresolved} application frame address(es)") + return (symbols, tuple(warnings)) + + +def summarize(profile: Mapping[str, Any], binary: Path) -> ProfileCpuSummary: + application_library, arch = _application_library(profile, binary) + scale = _cpu_seconds_scale(profile) + libraries = cast(list[Any], profile["libs"]) + threads = cast(list[Any], profile["threads"]) + observed_cpu = 0.0 + application_cpu = 0.0 + other_library_cpu = 0.0 + unattributed_cpu = 0.0 + named_cpu: defaultdict[str, float] = defaultdict(float) + address_cpu: defaultdict[int, float] = defaultdict(float) + warnings: list[str] = [] + + for thread_index, value in enumerate(threads): + if not isinstance(value, dict): + warnings.append(f"profile thread {thread_index} is not an object") + continue + try: + for sample in _thread_leaf_samples(value, scale): + observed_cpu += sample.cpu_seconds + if sample.library_index is None or not 0 <= sample.library_index < len(libraries): + unattributed_cpu += sample.cpu_seconds + elif sample.library_index != application_library: + other_library_cpu += sample.cpu_seconds + else: + application_cpu += sample.cpu_seconds + normalized_name = ( + _normalize_atos_symbol(sample.function_name, binary.name) + if sample.function_name is not None + else None + ) + if normalized_name is not None and not re.match(r"^_+R", normalized_name): + named_cpu[normalized_name] += sample.cpu_seconds + elif sample.relative_address is not None: + address_cpu[sample.relative_address] += sample.cpu_seconds + except (KeyError, TypeError) as error: + warnings.append(f"could not read profile thread {thread_index}: {error}") + + symbols, symbol_warnings = _symbolicate(binary, arch, tuple(address_cpu)) + warnings.extend(symbol_warnings) + for address, cpu_seconds in address_cpu.items(): + if symbol := symbols.get(address): + named_cpu[symbol] += cpu_seconds + functions = tuple( + ProfileFunctionCpu(name, cpu_seconds) + for name, cpu_seconds in sorted(named_cpu.items(), key=lambda item: (-item[1], item[0])) + ) + return ProfileCpuSummary( + observed_cpu_seconds=observed_cpu, + application_cpu_seconds=application_cpu, + other_library_cpu_seconds=other_library_cpu, + unattributed_cpu_seconds=unattributed_cpu, + symbolized_application_cpu_seconds=sum(named_cpu.values()), + functions=functions, + warnings=tuple(warnings), + ) + + +def load_command(artifact: Path, *, os_name: str | None = None) -> str: + command = ["samply", "load", str(artifact)] + if (os.name if os_name is None else os_name) == "nt": + return subprocess.list2cmdline(command) + return shlex.join(command) + + +def _percentage(value: float, total: float) -> float: + return 0.0 if total <= 0 else value * 100.0 / total + + +def _metadata_rows(report: ProfileReport) -> tuple[tuple[str, str], ...]: + return ( + ("Artifact", str(report.artifact)), + ("Cache", report.cache_status), + ("Workload", report.workload), + ("Backend", report.backend), + ("Treatment", report.treatment), + ) + + +def _cpu_rows(summary: ProfileCpuSummary) -> tuple[tuple[str, str, str], ...]: + observed = summary.observed_cpu_seconds + application = summary.application_cpu_seconds + return ( + ("Observed thread", f"{observed:.2f} s", "100.0%"), + ("Application leaf", f"{application:.2f} s", f"{_percentage(application, observed):.1f}%"), + ( + "Other library", + f"{summary.other_library_cpu_seconds:.2f} s", + f"{_percentage(summary.other_library_cpu_seconds, observed):.1f}%", + ), + ( + "Unattributed", + f"{summary.unattributed_cpu_seconds:.2f} s", + f"{_percentage(summary.unattributed_cpu_seconds, observed):.1f}%", + ), + ( + "Application symbols", + "-", + f"{_percentage(summary.symbolized_application_cpu_seconds, application):.1f}%" + if application > 0 + else "n/a", + ), + ) + + +def _function_rows(summary: ProfileCpuSummary, top: int) -> tuple[tuple[str, str, str], ...]: + return tuple( + ( + f"{function.cpu_seconds:.3f} s", + f"{_percentage(function.cpu_seconds, summary.application_cpu_seconds):.1f}%", + function.name, + ) + for function in summary.functions[:top] + ) + + +def render_rich( + console: Console, + report: ProfileReport, +) -> None: + summary = report.cpu_summary + console.rule("CPU Profile Summary" if summary is not None else "Profile Ready") + metadata = Table.grid(padding=(0, 1)) + metadata.add_column(style="bold") + metadata.add_column(overflow="fold") + for label, value in _metadata_rows(report): + metadata.add_row(label, Text(value)) + console.print(metadata) + if summary is not None: + cpu_table = Table(title="CPU breakdown", box=box.SIMPLE_HEAD) + cpu_table.add_column("Metric") + cpu_table.add_column("CPU", justify="right") + cpu_table.add_column("Share", justify="right") + for row in _cpu_rows(summary): + cpu_table.add_row(*row) + console.print(cpu_table) + + function_rows = _function_rows(summary, report.top) + if function_rows: + functions = Table(title="Top functions by self CPU", box=box.SIMPLE_HEAD) + functions.add_column("CPU", justify="right") + functions.add_column("App %", justify="right") + functions.add_column("Function", overflow="fold") + for cpu, percentage, function in function_rows: + functions.add_row(cpu, percentage, Text(function)) + console.print(functions) + else: + console.print("[dim]No resolved application functions.[/dim]") + console.print("[bold]Visualize[/bold]") + console.print(Text(load_command(report.artifact), style="cyan")) + + +def _markdown_cell(value: str) -> str: + normalized = value.replace("\r\n", "\n").replace("\r", "\n") + return ( + normalized.replace("&", "&") + .replace("<", "<") + .replace(">", ">") + .replace("\\", "\\\\") + .replace("|", "\\|") + .replace("\n", "
") + ) + + +def render_markdown(report: ProfileReport) -> str: + summary = report.cpu_summary + lines = [f"## {'CPU Profile Summary' if summary is not None else 'Profile Ready'}", ""] + lines.extend(["| Field | Value |", "| --- | --- |"]) + for label, value in _metadata_rows(report): + lines.append(f"| {_markdown_cell(label)} | {_markdown_cell(value)} |") + if summary is not None: + lines.extend(["", "### CPU Breakdown", "", "| Metric | CPU | Share |", "| --- | ---: | ---: |"]) + lines.extend(f"| {_markdown_cell(metric)} | {cpu} | {share} |" for metric, cpu, share in _cpu_rows(summary)) + lines.extend(["", "### Top Functions by Self CPU", ""]) + function_rows = _function_rows(summary, report.top) + if function_rows: + lines.extend(["| CPU | App % | Function |", "| ---: | ---: | --- |"]) + lines.extend( + f"| {cpu} | {percentage} | {_markdown_cell(function)} |" for cpu, percentage, function in function_rows + ) + else: + lines.append("No resolved application functions.") + lines.extend(["", "### Visualize", "", "```shell", load_command(report.artifact), "```"]) + return "\n".join(lines) diff --git a/test_bench.py b/test_bench.py index 7ff1576..2f2ad18 100644 --- a/test_bench.py +++ b/test_bench.py @@ -1,9 +1,12 @@ from __future__ import annotations +import argparse +import gzip import io import json import resource import signal +import subprocess import sys from pathlib import Path from typing import Any @@ -13,6 +16,7 @@ from rich.console import Console import bench +import samply_analysis ROOT = Path(__file__).resolve().parent @@ -26,6 +30,7 @@ def make_record( max_rss_bytes: int | None = None, binary_sha256: str = "sha256:bin", file_sha256: str = "sha256:file", + backend: bench.Backend = "main", treatment: bench.Treatment = "off", timeout_sec: int = 120, target_label: str | None = None, @@ -43,6 +48,7 @@ def make_record( "binary_sha256": binary_sha256, "file_path": "file.egg", "file_sha256": file_sha256, + "backend": backend, "treatment": treatment, "timeout_sec": timeout_sec, "wall_sec": None if status == "timed-out" else wall_sec, @@ -89,6 +95,116 @@ def make_target( ) +def make_profile_data() -> dict[str, Any]: + return { + "meta": {"sampleUnits": {"threadCPUDelta": "us"}}, + "libs": [], + "threads": [], + } + + +def make_cpu_profile_data(binary_path: Path) -> dict[str, Any]: + thread = { + "stringArray": ["0x10", "system_call", "application", "system"], + "samples": { + "length": 4, + "stack": [0, 0, 1, None], + "threadCPUDelta": [9_000_000, 100_000, 200_000, 300_000], + }, + "stackTable": {"length": 2, "frame": [0, 1], "prefix": [None, None]}, + "frameTable": {"length": 2, "address": [0x10, 0x20], "func": [0, 1]}, + "funcTable": {"length": 2, "name": [0, 1], "resource": [0, 1]}, + "resourceTable": {"length": 2, "name": [2, 3], "lib": [0, 1]}, + } + parked_thread = { + **thread, + "samples": { + "length": 2, + "stack": [0, 0], + "threadCPUDelta": [5_000_000, 0], + }, + } + return { + "meta": {"sampleUnits": {"threadCPUDelta": "us"}}, + "libs": [ + {"name": binary_path.name, "path": str(binary_path), "arch": "arm64"}, + {"name": "libsystem_kernel.dylib", "path": "/usr/lib/system/libsystem_kernel.dylib", "arch": "arm64"}, + ], + "threads": [thread, parked_thread], + } + + +def write_profile(path: Path, profile: dict[str, Any] | None = None, *, compressed: bool = True) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + data = make_profile_data() if profile is None else profile + if compressed: + with gzip.open(path, "wt", encoding="utf-8") as handle: + json.dump(data, handle) + else: + path.write_text(json.dumps(data), encoding="utf-8") + + +def make_cpu_summary() -> samply_analysis.ProfileCpuSummary: + return samply_analysis.ProfileCpuSummary( + observed_cpu_seconds=1.0, + application_cpu_seconds=0.8, + other_library_cpu_seconds=0.15, + unattributed_cpu_seconds=0.05, + symbolized_application_cpu_seconds=0.6, + functions=( + samply_analysis.ProfileFunctionCpu("first", 0.4), + samply_analysis.ProfileFunctionCpu("second", 0.2), + ), + warnings=(), + ) + + +def make_profile_case( + tmp_path: Path, + *, + mode: bench.ProfileMode | None = None, + open_after: bool = False, + force_run: bool = False, + top: int = 15, + show_summary: bool = True, + output_format: bench.OutputFormat = "rich", +) -> tuple[bench.ProfileRequest, bench.ResolvedTarget, Path]: + file_spec = bench.FileSpec("file.egg", ROOT / "file.egg", "sha256:" + "b" * 64) + request = bench.ProfileRequest( + file=file_spec, + target_request=bench.TargetRequest(raw=".", source=".", label=None), + backend="main", + treatment="proofs", + timeout_sec=120, + profiles_dir=tmp_path, + mode=bench.ProfileMode(None, 10) if mode is None else mode, + open_after=open_after, + force_run=force_run, + top=top, + show_summary=show_summary, + output_format=output_format, + ) + target = make_target(binary_sha256="sha256:" + "a" * 64, binary_path=ROOT / "egglog-experimental") + artifact = bench.profile_cache_path( + tmp_path, + target.binary_sha256, + file_spec.sha256, + request.backend, + request.treatment, + request.mode, + ) + return (request, target, artifact) + + +def mock_profile_resolution( + monkeypatch: pytest.MonkeyPatch, + request: bench.ProfileRequest, + target: bench.ResolvedTarget, +) -> None: + monkeypatch.setattr(bench, "resolve_profile_request", lambda *args: request) + monkeypatch.setattr(bench, "resolve_profile_target", lambda *args: target) + + def test_selected_rows_uses_latest_timestamp_then_jsonl_order() -> None: rows = make_rows( make_record(0, started_at="2026-07-04T12:00:00Z", wall_sec=1.0), @@ -262,6 +378,116 @@ def test_parse_treatments_rejects_duplicates() -> None: bench.parse_treatments("off,term,off") +def test_parse_backends_accepts_single_and_comma_separated_values() -> None: + assert bench.parse_backends("main") == ("main",) + assert bench.parse_backends("main,dd") == ("main", "dd") + assert bench.parse_backends(" dd , main ") == ("dd", "main") + + +def test_parse_backends_rejects_duplicates_unknowns_and_empty_values() -> None: + with pytest.raises(ValueError, match="duplicate backend: main"): + bench.parse_backends("main,dd,main") + with pytest.raises(ValueError, match="unknown backend: bogus"): + bench.parse_backends("bogus") + with pytest.raises(ValueError, match="at least one backend"): + bench.parse_backends(",,") + + +def test_backend_registry_drives_parsing_capabilities_flags_and_display( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setitem( + bench.BACKEND_SPECS, + "future", + bench.BackendSpec("Future", ("proofs",), ("--backend", "future"), ("future-backend",)), + ) + + assert bench.parse_backends("future") == ("future",) + assert bench.supported_treatments("future") == ("proofs",) + assert bench.backend_flags("future") == ["--backend", "future"] + assert bench.backend_cargo_features(("main", "future")) == ("future-backend",) + assert bench.display_backend("future") == "Future" + assert bench.backend_treatment_cells(("future",), ("off", "proofs")) == (bench.BenchmarkCell("future", "proofs"),) + + +def test_parse_args_dispatches_profile_without_changing_benchmark_defaults() -> None: + benchmark_args = bench.parse_args(["--rounds", "1", "file.egg"]) + profile_args = bench.parse_args(["profile", "file.egg"]) + + assert benchmark_args.command == "benchmark" + assert benchmark_args.files == ["file.egg"] + assert benchmark_args.rounds == 1 + assert benchmark_args.format == "rich" + assert profile_args.command == "profile" + assert profile_args.file == "file.egg" + assert profile_args.backend == "main" + assert profile_args.treatment == "proofs" + assert profile_args.top == 15 + assert not profile_args.no_summary + assert profile_args.format == "rich" + + +def test_parse_profile_args_accepts_presentation_options() -> None: + args = bench.parse_args(["profile", "file.egg", "--top", "7", "--no-summary", "--format", "markdown", "--open"]) + + assert args.top == 7 + assert args.no_summary + assert args.format == "markdown" + assert args.open + + +def test_parse_args_rejects_markdown_report_dash() -> None: + with pytest.raises(SystemExit): + bench.parse_args(["--format", "markdown", "--report", "-", "file.egg"]) + + +def test_parse_profile_args_rejects_iterations_with_profile_seconds() -> None: + with pytest.raises(SystemExit): + bench.parse_args(["profile", "file.egg", "--iterations", "1", "--profile-seconds", "1"]) + + +def test_benchmark_cells_filter_off_for_non_main_backends() -> None: + file_spec = bench.FileSpec("file.egg", ROOT / "file.egg", "sha256:file") + spec = bench.BenchmarkSpec( + files=(file_spec,), + treatments=("off", "term", "proofs"), + rounds=1, + timeout_sec=120, + backends=("main", "dd"), + ) + + assert bench.benchmark_cells(spec) == ( + bench.BenchmarkCell("main", "off"), + bench.BenchmarkCell("main", "term"), + bench.BenchmarkCell("main", "proofs"), + bench.BenchmarkCell("dd", "term"), + bench.BenchmarkCell("dd", "proofs"), + ) + + +def test_validate_spec_rejects_backend_with_no_supported_treatments() -> None: + file_spec = bench.FileSpec("file.egg", ROOT / "file.egg", "sha256:file") + spec = bench.BenchmarkSpec( + files=(file_spec,), + treatments=("off",), + rounds=1, + timeout_sec=120, + backends=("dd",), + ) + + with pytest.raises(ValueError, match="backend dd has no supported treatments"): + bench.validate_spec(spec) + + +def test_resolve_profile_request_reuses_backend_treatment_validation(tmp_path: Path) -> None: + file_path = tmp_path / "file.egg" + file_path.write_text("(check (= 1 1))\n", encoding="utf-8") + args = bench.parse_args(["profile", str(file_path), "--backend", "dd", "--treatment", "off"]) + + with pytest.raises(ValueError, match="backend dd has no supported treatments"): + bench.resolve_profile_request(args, ROOT) + + def test_validate_spec_rejects_executable_prove_benchmark_file(tmp_path: Path) -> None: prove_file = tmp_path / "prove.egg" prove_file.write_text( @@ -295,6 +521,13 @@ def test_validate_spec_allows_prove_mentions_in_comments(tmp_path: Path) -> None bench.validate_spec(spec) +def test_default_files_use_dd_safe_math_microbenchmark() -> None: + display_paths = tuple(file.display_path for file in bench.resolve_files([], ROOT)) + + assert "egglog/tests/math-microbenchmark-mini.egg" in display_paths + assert "egglog/tests/math-microbenchmark.egg" not in display_paths + + def test_estimate_model_is_exact_only_and_updates_from_successful_processes() -> None: rows = make_rows( make_record(0, started_at="2026-07-04T12:00:00Z", wall_sec=2.0), @@ -356,6 +589,59 @@ def fake_git_sha(cwd: Path, ref: str = "HEAD") -> str: ] +def test_target_resolvers_share_materialization_and_select_build_profile( + monkeypatch: pytest.MonkeyPatch, +) -> None: + request = bench.parse_target(".") + row = bench.TargetRow(".", str(ROOT), "HEAD", "abc123", False) + target = make_target(binary_path=ROOT / "egglog-experimental") + materialized: list[bench.TargetRequest] = [] + build_profiles: list[bench.BuildProfile] = [] + build_features: list[tuple[str, ...]] = [] + + def fake_materialize( + target_request: bench.TargetRequest, + invocation_cwd: Path, + repo_root: Path, + ) -> bench.TargetRow: + materialized.append(target_request) + return row + + def fake_build( + target_request: bench.TargetRequest, + target_row: bench.TargetRow, + output: bench.RunnerOutput, + build_profile: bench.BuildProfile, + cargo_features: tuple[str, ...], + ) -> bench.ResolvedTarget: + assert target_request == request + assert target_row == row + build_profiles.append(build_profile) + build_features.append(cargo_features) + return target + + monkeypatch.setattr(bench, "materialize_target_request", fake_materialize) + monkeypatch.setattr(bench, "build_resolved_target", fake_build) + file_spec = bench.FileSpec("file.egg", ROOT / "file.egg", "sha256:file") + + benchmark_target = bench.resolve_target( + request, + bench.empty_report_frame(), + make_spec(file_spec), + False, + ROOT, + ROOT, + bench.RunnerOutput(), + ) + profile_target = bench.resolve_profile_target(request, "dd", ROOT, ROOT, bench.RunnerOutput()) + + assert benchmark_target == target + assert profile_target == target + assert materialized == [request, request] + assert build_profiles == ["release", "profiling"] + assert build_features == [(), ("dd-backend",)] + + def test_collection_plan_counts_cache_and_missing_rows() -> None: rows = make_rows(make_record(0, started_at="2026-07-04T12:00:00Z", wall_sec=1.0)) target = make_target() @@ -372,6 +658,27 @@ def test_collection_plan_counts_cache_and_missing_rows() -> None: assert force_plan.total_planned_processes == 3 +def test_collection_plan_does_not_reuse_main_rows_for_dd() -> None: + rows = make_rows(make_record(0, started_at="2026-07-04T12:00:00Z", backend="main", treatment="term", wall_sec=1.0)) + target = make_target() + file_spec = bench.FileSpec("file.egg", ROOT / "file.egg", "sha256:file") + spec = bench.BenchmarkSpec( + files=(file_spec,), + treatments=("term",), + rounds=1, + timeout_sec=120, + backends=("main", "dd"), + ) + + plan = bench.build_collection_plan(rows, target, spec, False) + + main_cell, dd_cell = plan.cells + assert main_cell.backend == "main" + assert main_cell.missing_observations == 0 + assert dd_cell.backend == "dd" + assert dd_cell.missing_observations == 1 + + def test_parse_args_rejects_removed_output_mode() -> None: with pytest.raises(SystemExit): bench.parse_args(["--output", "jsonl"]) @@ -502,6 +809,107 @@ def test_render_report_compares_multiple_targets_before_bottom_summary() -> None assert output.index("base: per-file wall time") < output.index("Benchmark summary") +def test_render_report_compares_backends_for_single_target() -> None: + rows = make_rows( + make_record(0, started_at="2026-07-04T12:00:00Z", backend="main", treatment="term", wall_sec=1.0), + make_record(1, started_at="2026-07-04T12:00:01Z", backend="main", treatment="term", wall_sec=1.1), + make_record(2, started_at="2026-07-04T12:00:00Z", backend="dd", treatment="term", wall_sec=2.0), + make_record(3, started_at="2026-07-04T12:00:01Z", backend="dd", treatment="term", wall_sec=2.1), + ) + target = make_target() + file_spec = bench.FileSpec("file.egg", ROOT / "file.egg", "sha256:file") + stream = io.StringIO() + console = Console(file=stream, width=240, color_system=None) + + bench.render_report( + console, + bench.ReportDestination(path=None, stream=io.StringIO()), + rows, + [target], + bench.BenchmarkSpec( + files=(file_spec,), + treatments=("term",), + rounds=2, + timeout_sec=120, + backends=("main", "dd"), + ), + ) + + output = stream.getvalue() + assert "Per-file backend wall-time change vs main" in output + assert "DD vs main wall time" in output + assert "DD/main" in output + assert "Best file" in output + assert "Best ratio" in output + assert "Faster files" in output + assert "main/term" in output + assert "dd/term" in output + assert "dd/off" not in output + assert "1.952x" in output + assert output.index("DD vs main wall time") < output.index("proof overhead summary") + + +def test_render_report_backend_summary_highlights_best_file() -> None: + records: list[dict[str, Any]] = [] + + def add_cell( + *, + file_sha256: str, + backend: bench.Backend, + wall_sec: float, + max_rss_bytes: int, + ) -> None: + for _ in range(2): + records.append( + make_record( + len(records), + started_at=f"2026-07-04T12:00:{len(records):02d}Z", + file_sha256=file_sha256, + backend=backend, + treatment="term", + wall_sec=wall_sec, + max_rss_bytes=max_rss_bytes, + ) + ) + + mib = 1024 * 1024 + add_cell(file_sha256="sha256:slow", backend="main", wall_sec=1.0, max_rss_bytes=100 * mib) + add_cell(file_sha256="sha256:slow", backend="dd", wall_sec=2.0, max_rss_bytes=200 * mib) + add_cell(file_sha256="sha256:fast", backend="main", wall_sec=1.0, max_rss_bytes=100 * mib) + add_cell(file_sha256="sha256:fast", backend="dd", wall_sec=0.5, max_rss_bytes=80 * mib) + rows = make_rows(*records) + target = make_target() + slow_file = bench.FileSpec("slow.egg", ROOT / "slow.egg", "sha256:slow") + fast_file = bench.FileSpec("fast.egg", ROOT / "fast.egg", "sha256:fast") + stream = io.StringIO() + console = Console(file=stream, width=260, color_system=None) + + bench.render_report( + console, + bench.ReportDestination(path=None, stream=io.StringIO()), + rows, + [target], + bench.BenchmarkSpec( + files=(slow_file, fast_file), + treatments=("term",), + rounds=2, + timeout_sec=120, + backends=("main", "dd"), + ), + ) + + output = stream.getvalue() + assert "DD vs main wall time" in output + assert "DD vs main peak RSS" in output + assert "Faster files" in output + assert "Lower-RSS files" in output + assert "fast.egg" in output + assert "1/2" in output + assert "1.250x" in output + assert "0.500x" in output + assert "0.800x" in output + + def test_render_report_compares_proofs_only_targets_with_percent_change() -> None: rows = make_rows( make_record( @@ -644,6 +1052,259 @@ def test_render_report_marks_invalid_multi_target_wall_time_cells() -> None: assert "off" in output +def test_markdown_escape_cell_handles_pipes_backslashes_and_multiline() -> None: + assert bench.markdown_escape_cell("a|b\\c\r\nnext\nlast") == "a\\|b\\\\c
next
last" + + +def test_render_markdown_table_uses_github_pipe_table() -> None: + table = bench.ReportTableData( + title="Example", + headers=("A|B", "Count"), + rows=(("x\\y", "one\ntwo"),), + caption="caption | text", + alignments=("left", "right"), + ) + + output = bench.render_markdown_table(table) + + assert output == ("### Example\n\n| A\\|B | Count |\n| --- | ---: |\n| x\\\\y | one
two |\n\n*caption \\| text*") + + +def test_benchmark_command_block_uses_fixed_entrypoint_and_shell_quoting() -> None: + assert bench.benchmark_command_block(("--report", "/tmp/report path.jsonl", "pipe|file.egg")) == ( + "```shell\n$ ./bench.py --report '/tmp/report path.jsonl' 'pipe|file.egg'\n```" + ) + + +def test_render_markdown_report_deterministic_golden() -> None: + rows = make_rows(make_record(0, started_at="2026-07-04T12:00:00Z", wall_sec=1.0)) + target = make_target() + file_spec = bench.FileSpec("dir/file.egg", ROOT / "file.egg", "sha256:file") + + output = bench.render_markdown_report( + bench.ReportDestination(path=Path("reports.jsonl")), + rows, + [target], + bench.BenchmarkSpec(files=(file_spec,), treatments=("off",), rounds=1, timeout_sec=120), + ("--rounds", "1", "--format", "markdown", "dir/file.egg"), + ) + + assert output == ( + "```shell\n" + "$ ./bench.py --rounds 1 --format markdown dir/file.egg\n" + "```\n" + "\n" + "# Benchmark Report\n" + "\n" + "- Report: `reports.jsonl`\n" + "- Selected rows per cell: `1`\n" + "\n" + "## Targets\n" + "\n" + "| Role | Label | Git | Dirty | Binary | Path |\n" + "| --- | --- | --- | --- | --- | --- |\n" + f"| target | abc123 | abc123 | no | bin | {ROOT} |\n" + "\n" + "## Target Diagnostics\n" + "\n" + "### abc123: per-file wall time\n" + "\n" + "| File | off |\n" + "| --- | ---: |\n" + "| dir/file.egg | 1.0000s |\n" + "\n" + "*Within-target wall-time estimates. These are not target-vs-baseline ratios.*\n" + "\n" + "## Benchmark Summary\n" + "\n" + "### abc123: proof overhead summary\n" + "\n" + "| Metric | Ratio | Change | Worst file | Worst ratio | Result |\n" + "| --- | ---: | ---: | --- | ---: | --- |\n" + "| no proof baseline | - | - | - | - | select off and proofs |\n" + "\n" + "*Within-backend proof overhead. This is not backend-vs-main performance.*" + ) + + +def test_markdown_report_has_no_rich_markup_ansi_or_box_drawing() -> None: + rows = make_rows(make_record(0, started_at="2026-07-04T12:00:00Z", wall_sec=1.0)) + target = make_target() + file_spec = bench.FileSpec("file.egg", ROOT / "file.egg", "sha256:file") + + output = bench.render_markdown_report( + bench.ReportDestination(path=Path("reports.jsonl")), + rows, + [target], + bench.BenchmarkSpec(files=(file_spec,), treatments=("off",), rounds=1, timeout_sec=120), + ) + + assert "\x1b[" not in output + assert "[bold]" not in output + assert not any(character in output for character in "━─┏┓└┘│┃┡┩╇") + + +def test_backend_summary_table_data_feeds_rich_and_markdown_renderers() -> None: + records: list[dict[str, Any]] = [] + for file_sha256, main_time, dd_time in (("sha256:slow", 1.0, 2.0), ("sha256:fast", 1.0, 0.5)): + records.append( + make_record( + len(records), + started_at=f"2026-07-04T12:00:{len(records):02d}Z", + file_sha256=file_sha256, + backend="main", + treatment="term", + wall_sec=main_time, + ) + ) + records.append( + make_record( + len(records), + started_at=f"2026-07-04T12:00:{len(records):02d}Z", + file_sha256=file_sha256, + backend="dd", + treatment="term", + wall_sec=dd_time, + ) + ) + rows = make_rows(*records) + target = make_target() + slow_file = bench.FileSpec("slow.egg", ROOT / "slow.egg", "sha256:slow") + fast_file = bench.FileSpec("fast.egg", ROOT / "fast.egg", "sha256:fast") + spec = bench.BenchmarkSpec( + files=(slow_file, fast_file), + treatments=("term",), + rounds=1, + timeout_sec=120, + backends=("main", "dd"), + ) + cell_maps = {target: bench.target_cell_summaries(rows, target, spec)} + rss_cell_maps = {target: bench.target_rss_cell_summaries(rows, target, spec)} + table_data = bench.backend_summary_tables(cell_maps, rss_cell_maps, [target], spec)[0] + rich_stream = io.StringIO() + rich_console = Console(file=rich_stream, width=260, color_system=None) + + rich_console.print(bench.render_rich_table(table_data)) + markdown = bench.render_markdown_table(table_data) + + assert table_data.headers[-3:] == ("Best file", "Best ratio", "Best result") + assert table_data.rows == ( + ( + "dd", + "term", + "2.0000s", + "2.5000s", + "1.250x", + "25.0%", + "1.000x", + "0/2", + "fast.egg", + "0.500x", + "point only", + ), + ) + for value in table_data.rows[0]: + assert value in markdown + assert value in rich_stream.getvalue() + + +def test_main_markdown_report_goes_to_stdout_and_status_to_stderr( + monkeypatch: pytest.MonkeyPatch, + capsys: Any, + tmp_path: Path, +) -> None: + benchmark_file = tmp_path / "file.egg" + benchmark_file.write_text("(check (= 1 1))\n", encoding="utf-8") + file_spec = bench.FileSpec("file.egg", benchmark_file, "sha256:file") + rows = make_rows(make_record(0, started_at="2026-07-04T12:00:00Z", wall_sec=1.0)) + target = make_target() + + monkeypatch.setattr(bench, "git_root_for_path", lambda path: ROOT) + monkeypatch.setattr(bench, "resolve_files", lambda raw_files, invocation_cwd: (file_spec,)) + monkeypatch.setattr(bench, "load_report", lambda destination: rows) + monkeypatch.setattr(bench, "resolve_target", lambda *args: target) + monkeypatch.setattr(bench, "build_collection_plan", lambda *args: object()) + monkeypatch.setattr( + bench, + "emit_collection_plan", + lambda output, plan, estimate_model: output.console.print("status line"), + ) + monkeypatch.setattr( + bench, + "collect_rows", + lambda current_rows, *args: bench.CollectionResult(current_rows, bench.empty_report_frame()), + ) + argv = [ + "--format", + "markdown", + "--report", + str(tmp_path / "reports.jsonl"), + "--rounds", + "1", + "--treatments", + "off", + "file.egg", + ] + + result = bench.main(argv) + + captured = capsys.readouterr() + assert result == 0 + assert captured.out.startswith(bench.benchmark_command_block(argv) + "\n\n# Benchmark Report\n") + assert captured.out.endswith("\n") + assert not captured.out.endswith("\n\n") + assert "\x1b[" not in captured.out + assert "[bold]" not in captured.out + assert "status line" in captured.err + assert "# Benchmark Report" not in captured.err + + +def test_main_rich_report_remains_on_stderr_by_default( + monkeypatch: pytest.MonkeyPatch, + capsys: Any, + tmp_path: Path, +) -> None: + benchmark_file = tmp_path / "file.egg" + benchmark_file.write_text("(check (= 1 1))\n", encoding="utf-8") + file_spec = bench.FileSpec("file.egg", benchmark_file, "sha256:file") + rows = make_rows(make_record(0, started_at="2026-07-04T12:00:00Z", wall_sec=1.0)) + target = make_target() + + monkeypatch.setattr(bench, "git_root_for_path", lambda path: ROOT) + monkeypatch.setattr(bench, "resolve_files", lambda raw_files, invocation_cwd: (file_spec,)) + monkeypatch.setattr(bench, "load_report", lambda destination: rows) + monkeypatch.setattr(bench, "resolve_target", lambda *args: target) + monkeypatch.setattr(bench, "build_collection_plan", lambda *args: object()) + monkeypatch.setattr( + bench, + "emit_collection_plan", + lambda output, plan, estimate_model: output.console.print("status line"), + ) + monkeypatch.setattr( + bench, + "collect_rows", + lambda current_rows, *args: bench.CollectionResult(current_rows, bench.empty_report_frame()), + ) + + result = bench.main( + [ + "--report", + str(tmp_path / "reports.jsonl"), + "--rounds", + "1", + "--treatments", + "off", + "file.egg", + ] + ) + + captured = capsys.readouterr() + assert result == 0 + assert captured.out == "" + assert "status line" in captured.err + assert "Benchmark report" in captured.err + + def test_report_dash_writes_rows_to_stream_and_does_not_load_cache() -> None: stream = io.StringIO() destination = bench.ReportDestination(path=None, stream=stream) @@ -659,6 +1320,7 @@ def test_report_dash_writes_rows_to_stream_and_does_not_load_cache() -> None: assert records[0]["wall_sec"] == 1.0 assert records[0]["target_source"] == "." assert records[0]["file_path"] == "file.egg" + assert records[0]["backend"] == "main" assert "row_index" not in records[0] assert "warmup_rounds" not in records[0] assert "target" not in records[0] @@ -676,15 +1338,29 @@ def test_flat_jsonl_roundtrips_through_report_frame(tmp_path: Path) -> None: raw_record = json.loads(report.read_text(encoding="utf-8")) assert raw_record["target_label"] == "mine" assert raw_record["wall_sec"] == 1.0 + assert raw_record["backend"] == "main" assert "row_index" not in raw_record assert "warmup_rounds" not in raw_record assert "target" not in raw_record assert loaded["row_index"].tolist() == [0] assert loaded["target_label"].tolist() == ["mine"] + assert loaded["backend"].tolist() == ["main"] assert loaded["wall_sec"].tolist() == [1.0] +def test_old_flat_jsonl_without_backend_loads_as_main(tmp_path: Path) -> None: + report = tmp_path / "reports.jsonl" + record = make_record(0, started_at="2026-07-04T12:00:00Z", wall_sec=1.0) + record.pop("row_index") + record.pop("backend") + report.write_text(json.dumps(record) + "\n", encoding="utf-8") + + loaded = bench.load_report(bench.ReportDestination(path=report)) + + assert loaded["backend"].tolist() == ["main"] + + def test_report_frame_rejects_success_without_wall_time() -> None: record = make_record(0, started_at="2026-07-04T12:00:00Z", wall_sec=None) @@ -727,6 +1403,651 @@ def test_run_command_records_signal_separately_from_exit_code() -> None: assert result.error.signal == signal.SIGTERM +def test_run_process_passes_backend_flag_only_for_dd(monkeypatch: pytest.MonkeyPatch) -> None: + commands: list[list[str]] = [] + file_spec = bench.FileSpec("file.egg", ROOT / "file.egg", "sha256:file") + + def fake_run_command(command: list[str], checkout_path: Path, timeout_sec: int) -> bench.TimingResult: + commands.append(command) + assert checkout_path == ROOT + assert timeout_sec == 120 + return bench.TimingResult("success", bench.TimingRow(wall_sec=1.0), None) + + monkeypatch.setattr(bench, "run_command", fake_run_command) + + bench.run_process(ROOT / "egglog-experimental", ROOT, file_spec, "main", "off", 120) + bench.run_process(ROOT / "egglog-experimental", ROOT, file_spec, "dd", "proofs", 120) + + assert "--backend" not in commands[0] + assert commands[1][commands[1].index("--backend") : commands[1].index("--backend") + 2] == [ + "--backend", + "dd", + ] + assert "--proofs" in commands[1] + + +def test_workload_command_matches_benchmark_behavior() -> None: + file_spec = bench.FileSpec("file.egg", ROOT / "file.egg", "sha256:file") + + assert bench.workload_command(ROOT / "egglog-experimental", file_spec, "main", "off") == [ + str(ROOT / "egglog-experimental"), + "--mode", + "no-messages", + "-j", + "1", + str(file_spec.absolute_path), + ] + assert bench.workload_command(ROOT / "egglog-experimental", file_spec, "dd", "proofs") == [ + str(ROOT / "egglog-experimental"), + "--mode", + "no-messages", + "-j", + "1", + "--backend", + "dd", + "--proofs", + str(file_spec.absolute_path), + ] + + +def test_build_target_uses_profiling_profile(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + commands: list[list[str]] = [] + binary = tmp_path / "target" / "profiling" / "egglog-experimental" + binary.parent.mkdir(parents=True) + binary.write_text("binary", encoding="utf-8") + row = bench.TargetRow( + source=".", + path=str(tmp_path), + git_ref="HEAD", + git_sha="abc123", + is_dirty=False, + ) + + def fake_run(command: list[str], **kwargs: Any) -> None: + commands.append(command) + + monkeypatch.setattr(bench.subprocess, "run", fake_run) + monkeypatch.setattr(bench, "sha256_file", lambda path: "sha256:bin") + + binary_path, binary_sha256 = bench.build_target(row, bench.RunnerOutput(), "profiling") + + assert commands == [["cargo", "build", "--profile", "profiling", "-p", "egglog-experimental"]] + assert binary_path == binary + assert binary_sha256 == "sha256:bin" + + +def test_build_target_enables_requested_backend_features(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + commands: list[list[str]] = [] + binary = tmp_path / "target" / "release" / "egglog-experimental" + binary.parent.mkdir(parents=True) + binary.write_text("binary", encoding="utf-8") + row = bench.TargetRow(".", str(tmp_path), "HEAD", "abc123", False) + monkeypatch.setattr(bench.subprocess, "run", lambda command, **kwargs: commands.append(command)) + monkeypatch.setattr(bench, "sha256_file", lambda path: "sha256:bin") + + bench.build_target(row, bench.RunnerOutput(), cargo_features=bench.backend_cargo_features(("main", "dd"))) + + assert commands == [["cargo", "build", "--release", "-p", "egglog-experimental", "--features", "dd-backend"]] + + +def test_profile_cache_path_uses_full_binary_and_file_hashes() -> None: + binary_hash = "sha256:" + "a" * 64 + file_hash = "sha256:" + "b" * 64 + + explicit = bench.profile_cache_path( + Path(".profiles"), binary_hash, file_hash, "main", "proofs", bench.ProfileMode(5, None) + ) + automatic = bench.profile_cache_path( + Path(".profiles"), + binary_hash, + file_hash, + "main", + "proofs", + bench.ProfileMode(None, 10), + ) + + assert explicit == Path(".profiles") / "v1" / ("a" * 64) / ("b" * 64) / "main-proofs-i5.json.gz" + assert automatic == Path(".profiles") / "v1" / ("a" * 64) / ("b" * 64) / "main-proofs-auto10s.json.gz" + + +def test_profile_display_path_is_relative_inside_invocation_directory(tmp_path: Path) -> None: + artifact = tmp_path / ".profiles" / "v1" / "profile.json.gz" + + assert bench.profile_display_path(artifact, tmp_path) == Path(".profiles/v1/profile.json.gz") + + +def test_calculate_profile_iterations_uses_margin_and_cap() -> None: + assert bench.calculate_profile_iterations(2.0, 10) == (6, False) + assert bench.calculate_profile_iterations(20.0, 10) == (1, False) + assert bench.calculate_profile_iterations(0.00001, 10, max_iterations=7) == (7, True) + assert bench.calculate_profile_iterations(0.0, 10, max_iterations=7) == (7, True) + + +def test_samply_record_uses_fixed_flags_and_replaces_artifact( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + artifact = tmp_path / "profile.json.gz" + artifact.write_bytes(b"old") + profile = make_profile_data() + commands: list[list[str]] = [] + + def fake_run(command: list[str], **kwargs: Any) -> None: + commands.append(command) + output = Path(command[command.index("--output") + 1]) + write_profile(output, profile) + + monkeypatch.setattr(bench, "samply_executable", lambda: "samply") + monkeypatch.setattr(bench.subprocess, "run", fake_run) + + recorded_profile = bench.run_samply_record( + artifact=artifact, + name="profile", + iterations=3, + workload=["workload"], + checkout_path=ROOT, + timeout_sec=120, + ) + + assert recorded_profile == profile + assert artifact.read_bytes()[:2] == b"\x1f\x8b" + command = commands[0] + temporary_output = Path(command[command.index("--output") + 1]) + assert temporary_output.name.endswith(".json.gz") + assert command[:7] == ["samply", "record", "--save-only", "--rate", "1000", "--reuse-threads", "--iteration-count"] + assert command[command.index("--iteration-count") + 1] == "3" + assert command[-2:] == ["--", "workload"] + + +def test_samply_failure_leaves_no_new_artifact(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + artifact = tmp_path / "profile.json.gz" + temp_artifact = tmp_path / ".profile.tmp-test.json.gz" + + def fake_run(command: list[str], **kwargs: Any) -> None: + temp_artifact.write_bytes(b"partial") + raise bench.subprocess.CalledProcessError(1, command) + + monkeypatch.setattr(bench, "samply_executable", lambda: "samply") + monkeypatch.setattr(bench, "profile_temp_path", lambda path: temp_artifact) + monkeypatch.setattr(bench.subprocess, "run", fake_run) + + with pytest.raises(bench.subprocess.CalledProcessError): + bench.run_samply_record( + artifact=artifact, + name="profile", + iterations=1, + workload=["workload"], + checkout_path=ROOT, + timeout_sec=120, + ) + + assert not artifact.exists() + assert not temp_artifact.exists() + + +def test_read_profile_artifact_rejects_plain_json(tmp_path: Path) -> None: + artifact = tmp_path / "profile.json.gz" + write_profile(artifact, compressed=False) + + with pytest.raises(ValueError, match="not gzip-compressed"): + samply_analysis.read_artifact(artifact) + + +def test_read_profile_artifact_rejects_malformed_json(tmp_path: Path) -> None: + artifact = tmp_path / "profile.json.gz" + with gzip.open(artifact, "wt", encoding="utf-8") as handle: + handle.write("not-json") + + with pytest.raises(ValueError, match="could not parse profile artifact"): + samply_analysis.read_artifact(artifact) + + +def test_read_profile_artifact_normalizes_os_errors( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + artifact = tmp_path / "profile.json.gz" + artifact.write_bytes(b"\x1f\x8b") + original_open = Path.open + + def denied_open(path: Path, *args: Any, **kwargs: Any) -> Any: + if path == artifact: + raise PermissionError("denied") + return original_open(path, *args, **kwargs) + + monkeypatch.setattr(Path, "open", denied_open) + + with pytest.raises(ValueError, match="could not read profile artifact"): + samply_analysis.read_artifact(artifact) + + +def test_profile_leaf_samples_have_named_fields(tmp_path: Path) -> None: + binary = tmp_path / "egglog-experimental" + profile = make_cpu_profile_data(binary) + + samples = tuple(samply_analysis._thread_leaf_samples(profile["threads"][0], 1e-6)) + + assert samples[0].cpu_seconds == pytest.approx(0.1) + assert samples[0].library_index == 0 + assert samples[0].function_name == "0x10" + assert samples[0].relative_address == 0x10 + + +def test_samply_plain_json_output_is_not_promoted(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + artifact = tmp_path / "profile.json.gz" + artifact.write_bytes(b"old") + temporary_paths: list[Path] = [] + + def fake_run(command: list[str], **kwargs: Any) -> None: + output = Path(command[command.index("--output") + 1]) + temporary_paths.append(output) + write_profile(output, compressed=False) + + monkeypatch.setattr(bench, "samply_executable", lambda: "samply") + monkeypatch.setattr(bench.subprocess, "run", fake_run) + + with pytest.raises(ValueError, match="gzip"): + bench.run_samply_record( + artifact=artifact, + name="profile", + iterations=1, + workload=["workload"], + checkout_path=ROOT, + timeout_sec=120, + ) + + assert artifact.read_bytes() == b"old" + assert temporary_paths and not temporary_paths[0].exists() + + +def test_missing_samply_reports_install_command(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(bench.shutil, "which", lambda name: None) + + with pytest.raises(FileNotFoundError, match="cargo install --locked samply"): + bench.samply_executable() + + +def test_macos_profile_summary_uses_cpu_deltas_and_atos_offsets( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + binary = tmp_path / "egglog-experimental" + binary.write_bytes(b"binary") + profile = make_cpu_profile_data(binary) + commands: list[list[str]] = [] + + def fake_run(command: list[str], **kwargs: Any) -> subprocess.CompletedProcess[str]: + commands.append(command) + return subprocess.CompletedProcess( + command, + 0, + stdout="crate..Thing$LT$T$GT$::run::h1234 (in egglog-experimental) + 4\n", + stderr="", + ) + + monkeypatch.setattr(samply_analysis.shutil, "which", lambda name: "/usr/bin/atos" if name == "atos" else None) + monkeypatch.setattr(samply_analysis.subprocess, "run", fake_run) + + summary = samply_analysis.summarize(profile, binary) + + assert summary.observed_cpu_seconds == pytest.approx(0.6) + assert summary.application_cpu_seconds == pytest.approx(0.1) + assert summary.other_library_cpu_seconds == pytest.approx(0.2) + assert summary.unattributed_cpu_seconds == pytest.approx(0.3) + assert summary.symbolized_application_cpu_seconds == pytest.approx(0.1) + assert len(summary.functions) == 1 + assert summary.functions[0].name == "crate::Thing::run::h1234" + assert summary.functions[0].cpu_seconds == pytest.approx(0.1) + assert len(commands) == 1 + assert "-offset" in commands[0] + assert "0x10" in commands[0] + assert "0x100000010" not in commands[0] + + +def test_macos_profile_summary_is_partial_without_atos( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + binary = tmp_path / "egglog-experimental" + binary.write_bytes(b"binary") + monkeypatch.setattr(samply_analysis.shutil, "which", lambda name: None) + monkeypatch.setattr( + samply_analysis.subprocess, + "run", + lambda *args, **kwargs: pytest.fail("symbolizer should not run without atos"), + ) + + summary = samply_analysis.summarize(make_cpu_profile_data(binary), binary) + + assert summary.observed_cpu_seconds == pytest.approx(0.6) + assert summary.application_cpu_seconds == pytest.approx(0.1) + assert summary.symbolized_application_cpu_seconds == 0 + assert summary.functions == () + assert any("atos is unavailable" in warning for warning in summary.warnings) + + +def test_macos_rust_v0_symbols_use_llvm_cxxfilt(monkeypatch: pytest.MonkeyPatch) -> None: + commands: list[tuple[list[str], str | None]] = [] + + def fake_run(command: list[str], **kwargs: Any) -> subprocess.CompletedProcess[str]: + commands.append((command, kwargs.get("input"))) + return subprocess.CompletedProcess(command, 0, stdout="__rustc::__rdl_alloc\n", stderr="") + + monkeypatch.setattr( + samply_analysis.shutil, + "which", + lambda name: "/usr/bin/xcrun" if name == "xcrun" else None, + ) + monkeypatch.setattr(samply_analysis.subprocess, "run", fake_run) + + symbols, warnings = samply_analysis.demangle_rust_v0_symbols({0x10: "_RNvExample"}) + + assert symbols == {0x10: "__rustc::__rdl_alloc"} + assert warnings == () + assert commands == [(["/usr/bin/xcrun", "llvm-cxxfilt"], "_RNvExample\n")] + + +def test_profile_load_command_quotes_paths_for_posix_and_windows(tmp_path: Path) -> None: + artifact = tmp_path / "profiles with spaces" / "profile.json.gz" + + assert samply_analysis.load_command(artifact, os_name="posix") == f"samply load '{artifact.resolve()}'" + assert samply_analysis.load_command(artifact, os_name="nt") == f'samply load "{artifact.resolve()}"' + assert samply_analysis.load_command(Path(".profiles/profile.json.gz"), os_name="posix") == ( + "samply load .profiles/profile.json.gz" + ) + + +def test_profile_rich_and_markdown_render_same_summary(tmp_path: Path) -> None: + artifact = tmp_path / "profiles with spaces" / "profile.json.gz" + summary = make_cpu_summary() + report = samply_analysis.ProfileReport( + artifact=artifact, + cache_status="hit", + workload="dir\\name|file\nnext.egg", + backend="main", + treatment="proofs", + top=1, + cpu_summary=summary, + ) + rich_output = io.StringIO() + console = Console(file=rich_output, force_terminal=False, width=120) + + samply_analysis.render_rich(console, report) + markdown = samply_analysis.render_markdown(report) + + for value in ("CPU Profile Summary", "Observed thread", "Application symbols", "first", "samply load"): + assert value in rich_output.getvalue() + assert value in markdown + assert "second" not in rich_output.getvalue() + assert "second" not in markdown + assert "first<T>" in markdown + assert "dir\\\\name\\|file
next.egg" in markdown + assert "| Metric | CPU | Share |" in markdown + + +def test_open_samply_profile_redirects_viewer_output_and_handles_interrupt( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + artifact = tmp_path / "profile.json.gz" + calls: list[dict[str, Any]] = [] + + def fake_run(command: list[str], **kwargs: Any) -> None: + calls.append(kwargs) + raise KeyboardInterrupt + + monkeypatch.setattr(bench, "samply_executable", lambda: "samply") + monkeypatch.setattr(bench.subprocess, "run", fake_run) + + bench.open_samply_profile(artifact, ROOT) + + assert calls == [{"cwd": ROOT, "check": True, "stdout": sys.stderr, "stderr": sys.stderr}] + + +@pytest.mark.parametrize("platform", ["linux", "win32"]) +def test_profile_cache_hit_on_non_macos_skips_cpu_summary_and_workload( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + capsys: Any, + platform: str, +) -> None: + request, target, artifact = make_profile_case(tmp_path) + write_profile(artifact) + monkeypatch.setattr(bench.sys, "platform", platform) + mock_profile_resolution(monkeypatch, request, target) + monkeypatch.setattr(bench, "run_command", lambda *args, **kwargs: pytest.fail("calibration should not run")) + monkeypatch.setattr(bench, "run_samply_record", lambda **kwargs: pytest.fail("samply should not run")) + monkeypatch.setattr( + samply_analysis, + "summarize", + lambda *args, **kwargs: pytest.fail("macOS summary should not run"), + ) + + bench.run_profile(argparse.Namespace(), bench.RunnerOutput(), ROOT, ROOT) + + captured = capsys.readouterr() + assert captured.out == "" + assert "Profile Ready" in captured.err + assert "Artifact" in captured.err + assert ".json.gz" in captured.err + assert "CPU Profile Summary" not in captured.err + assert "samply load" in captured.err + assert "profiler.firefox.com" not in captured.err + assert "currently available on macOS only" in captured.err + + +def test_profile_cache_hit_on_non_macos_prints_markdown_handoff( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + capsys: Any, +) -> None: + request, target, artifact = make_profile_case(tmp_path, output_format="markdown") + write_profile(artifact) + monkeypatch.setattr(bench.sys, "platform", "linux") + mock_profile_resolution(monkeypatch, request, target) + monkeypatch.setattr(samply_analysis, "summarize", lambda *args: pytest.fail("summary should not run")) + + bench.run_profile(argparse.Namespace(), bench.RunnerOutput(), ROOT, ROOT) + + captured = capsys.readouterr() + assert captured.out.startswith("## Profile Ready\n") + assert "| Field | Value |" in captured.out + assert "```shell\nsamply load " in captured.out + assert "CPU Breakdown" not in captured.out + assert "currently available on macOS only" in captured.err + + +def test_profile_cache_hit_on_macos_prints_cpu_summary( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + capsys: Any, +) -> None: + request, target, artifact = make_profile_case(tmp_path, top=1) + write_profile(artifact) + summary = make_cpu_summary() + monkeypatch.setattr(bench.sys, "platform", "darwin") + mock_profile_resolution(monkeypatch, request, target) + monkeypatch.setattr(samply_analysis, "summarize", lambda profile, binary: summary) + + bench.run_profile(argparse.Namespace(), bench.RunnerOutput(), ROOT, ROOT) + + captured = capsys.readouterr() + assert captured.out == "" + assert "CPU Profile Summary" in captured.err + assert "CPU breakdown" in captured.err + assert "Observed thread" in captured.err + assert "Application symbols" in captured.err + assert "Top functions by self CPU" in captured.err + assert "first" in captured.err + assert "second" not in captured.err + assert "samply load" in captured.err + + +def test_profile_cache_hit_prints_github_markdown_summary( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + capsys: Any, +) -> None: + request, target, artifact = make_profile_case(tmp_path, top=1, output_format="markdown") + write_profile(artifact) + monkeypatch.setattr(bench.sys, "platform", "darwin") + mock_profile_resolution(monkeypatch, request, target) + monkeypatch.setattr(samply_analysis, "summarize", lambda profile, binary: make_cpu_summary()) + + bench.run_profile(argparse.Namespace(), bench.RunnerOutput(), ROOT, ROOT) + + captured = capsys.readouterr() + assert captured.out.startswith("## CPU Profile Summary\n") + assert "| Field | Value |" in captured.out + assert "### CPU Breakdown" in captured.out + assert "| Metric | CPU | Share |" in captured.out + assert "### Top Functions by Self CPU" in captured.out + assert "first" in captured.out + assert "second" not in captured.out + assert "```shell\nsamply load " in captured.out + assert "\x1b[" not in captured.out + assert not any(character in captured.out for character in "┏┓┗┛━┃") + assert captured.out.endswith("\n") and not captured.out.endswith("\n\n") + assert "Profile cache hit" in captured.err + + +def test_profile_no_summary_prints_only_absolute_artifact_path( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + capsys: Any, +) -> None: + request, target, artifact = make_profile_case( + tmp_path, + mode=bench.ProfileMode(1, None), + show_summary=False, + ) + write_profile(artifact) + monkeypatch.setattr(bench.sys, "platform", "linux") + mock_profile_resolution(monkeypatch, request, target) + + bench.run_profile(argparse.Namespace(), bench.RunnerOutput(), ROOT, ROOT) + + captured = capsys.readouterr() + assert captured.out == f"{artifact.resolve()}\n" + assert "currently available on macOS only" not in captured.err + + +def test_profile_cache_hit_can_open_profile(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + request, target, artifact = make_profile_case( + tmp_path, + mode=bench.ProfileMode(1, None), + open_after=True, + show_summary=False, + ) + write_profile(artifact) + opened: list[Path] = [] + mock_profile_resolution(monkeypatch, request, target) + monkeypatch.setattr(bench, "open_samply_profile", lambda artifact, checkout_path: opened.append(artifact)) + + bench.run_profile(argparse.Namespace(), bench.RunnerOutput(), ROOT, ROOT) + + assert opened == [artifact] + + +def test_profile_explicit_iterations_skip_calibration_and_bypass_cache( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + request, target, artifact = make_profile_case( + tmp_path, + mode=bench.ProfileMode(3, None), + force_run=True, + show_summary=False, + ) + artifact.parent.mkdir(parents=True) + artifact.write_bytes(b"old") + recorded: list[tuple[Path, int]] = [] + + def fake_record(**kwargs: Any) -> dict[str, Any]: + recorded.append((kwargs["artifact"], kwargs["iterations"])) + return make_profile_data() + + mock_profile_resolution(monkeypatch, request, target) + monkeypatch.setattr(bench, "run_command", lambda *args, **kwargs: pytest.fail("calibration should not run")) + monkeypatch.setattr(bench, "run_samply_record", fake_record) + + bench.run_profile(argparse.Namespace(), bench.RunnerOutput(), ROOT, ROOT) + + assert recorded == [(artifact, 3)] + + +def test_profile_auto_calibrates_once_and_uses_derived_iterations( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + request, target, _ = make_profile_case(tmp_path, force_run=True, show_summary=False) + recorded: list[int] = [] + + def fake_record(**kwargs: Any) -> dict[str, Any]: + recorded.append(kwargs["iterations"]) + return make_profile_data() + + mock_profile_resolution(monkeypatch, request, target) + monkeypatch.setattr( + bench, + "run_command", + lambda command, checkout_path, timeout_sec: bench.TimingResult("success", bench.TimingRow(wall_sec=2.0), None), + ) + monkeypatch.setattr(bench, "run_samply_record", fake_record) + monkeypatch.setattr(bench, "load_report", lambda destination: pytest.fail("profile mode should not load reports")) + monkeypatch.setattr( + bench, "append_rows", lambda destination, rows: pytest.fail("profile mode should not append rows") + ) + + bench.run_profile(argparse.Namespace(), bench.RunnerOutput(), ROOT, ROOT) + + assert recorded == [6] + + +def test_profile_auto_calibration_failure_stops_before_samply( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + request, target, _ = make_profile_case(tmp_path, force_run=True, show_summary=False) + mock_profile_resolution(monkeypatch, request, target) + monkeypatch.setattr( + bench, + "run_command", + lambda command, checkout_path, timeout_sec: bench.TimingResult( + "timed-out", bench.TimingRow(), bench.ErrorRow("timed out") + ), + ) + monkeypatch.setattr(bench, "run_samply_record", lambda **kwargs: pytest.fail("samply should not run")) + + with pytest.raises(ValueError, match="profile calibration failed"): + bench.run_profile(argparse.Namespace(), bench.RunnerOutput(), ROOT, ROOT) + + +def test_profile_auto_iteration_cap_prints_warning( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, capsys: Any +) -> None: + request, target, _ = make_profile_case(tmp_path, force_run=True, show_summary=False) + mock_profile_resolution(monkeypatch, request, target) + monkeypatch.setattr( + bench, + "run_command", + lambda command, checkout_path, timeout_sec: bench.TimingResult( + "success", bench.TimingRow(wall_sec=0.001), None + ), + ) + monkeypatch.setattr(bench, "calculate_profile_iterations", lambda elapsed_seconds, profile_seconds: (7, True)) + monkeypatch.setattr(bench, "run_samply_record", lambda **kwargs: make_profile_data()) + + bench.run_profile(argparse.Namespace(), bench.RunnerOutput(), ROOT, ROOT) + + assert "maximum profile iterations reached" in capsys.readouterr().err + + +def test_profile_rejects_cache_only_label_targets() -> None: + with pytest.raises(ValueError, match="cache-only label="): + bench.resolve_profile_target(bench.parse_target("cached="), "main", ROOT, ROOT, bench.RunnerOutput()) + + def test_run_command_records_peak_rss() -> None: result = bench.run_command([sys.executable, "-c", "print('ok')"], ROOT, 120)