From 57c28bea06f1a9438228776367602ea08136b29f Mon Sep 17 00:00:00 2001 From: "John C. Burnham" Date: Fri, 4 Sep 2026 15:12:07 -0400 Subject: [PATCH 1/4] bench: add Anthropic FLT formalization --- Benchmarks/Compile/CompileAnthropicFLT.lean | 1 + Benchmarks/Compile/README.md | 23 ++++++++++++++++++--- Benchmarks/Compile/lake-manifest.json | 10 +++++++++ Benchmarks/Compile/lakefile.toml | 8 +++++++ Ix/Cli/BenchCmd.lean | 8 ++++++- 5 files changed, 46 insertions(+), 4 deletions(-) create mode 100644 Benchmarks/Compile/CompileAnthropicFLT.lean diff --git a/Benchmarks/Compile/CompileAnthropicFLT.lean b/Benchmarks/Compile/CompileAnthropicFLT.lean new file mode 100644 index 00000000..12d777fa --- /dev/null +++ b/Benchmarks/Compile/CompileAnthropicFLT.lean @@ -0,0 +1 @@ +import FinalCheck diff --git a/Benchmarks/Compile/README.md b/Benchmarks/Compile/README.md index 31ee0c70..1cba8377 100644 --- a/Benchmarks/Compile/README.md +++ b/Benchmarks/Compile/README.md @@ -4,7 +4,8 @@ Test libraries for the Ix compiler - [Init, Std, and Lean libraries](https://github.com/leanprover/lean4) - [Mathlib](https://github.com/leanprover-community/mathlib4) -- [FLT project](https://github.com/ImperialCollegeLondon/FLT) +- [Imperial College London FLT project](https://github.com/ImperialCollegeLondon/FLT) +- [Anthropic FLT proof artifact](https://github.com/anthropics/fermats-last-theorem) - Every native TruthMines member, independently, through the generated `TruthMines/Members/.lean` fidelity drivers - [Palomar.ix](https://github.com/argumentcomputer/Palomar.ix) as one aggregate @@ -14,7 +15,19 @@ Test libraries for the Ix compiler First ensure the Lean version used to build Ix matches the `Benchmarks/Compile/lean-toolchain` version (check against `ix --version`). Then run -`ix compile /path/to/Compile.lean` # replace `` with `Init`, `InitStd`, `Lean`, `Mathlib`, or `FLT` +`ix compile /path/to/Compile.lean` # replace `` with `Init`, `InitStd`, `Lean`, `Mathlib`, `FLT`, or `AnthropicFLT` + +The Anthropic artifact is also registered as the on-demand `AnthropicFLT` +benchmark environment. After building its oleans, benchmark the Ix compiler +and Rust kernel from the repository root with: + +```sh +cd Benchmarks/Compile +lake build +CompileAnthropicFLT:olean +cd ../.. +ix bench run --backend compile --env AnthropicFLT +ix bench run --backend ooc --env AnthropicFLT --ixe AnthropicFLT.ixe +``` For a TruthMines constituent, use the nested fidelity workspace, for example: @@ -26,4 +39,8 @@ complete sweep with `lake exe truthmines validate`; use `--only Cli,Palomar` to select libraries. > [!NOTE] -> Compiling Mathlib and FLT currently requires a multi-core CPU and >64 GB RAM. +> Compiling Mathlib and the Imperial FLT project currently requires a +> multi-core CPU and >64 GB RAM. Anthropic reports that building its FLT +> artifact from scratch peaked at 153 GB RAM and used about 67 GB under +> `.lake/`, plus roughly 220 GB of generated C files. The `olean` facet above +> skips native object compilation, but Lean still emits those C files. diff --git a/Benchmarks/Compile/lake-manifest.json b/Benchmarks/Compile/lake-manifest.json index 74b77d83..1112621d 100644 --- a/Benchmarks/Compile/lake-manifest.json +++ b/Benchmarks/Compile/lake-manifest.json @@ -51,6 +51,16 @@ "inputRev": "75c0681bd37567af00e8f0bd13fd59f1423e4217", "inherited": false, "configFile": "lakefile.lean"}, + {"url": "https://github.com/anthropics/fermats-last-theorem", + "type": "git", + "subDir": null, + "scope": "", + "rev": "aa2d8b34692b16c70f699536de0d8e75b9a3e9ef", + "name": "flt_e2e", + "manifestFile": "lake-manifest.json", + "inputRev": "aa2d8b34692b16c70f699536de0d8e75b9a3e9ef", + "inherited": false, + "configFile": "lakefile.lean"}, {"url": "https://github.com/ImperialCollegeLondon/FLT", "type": "git", "subDir": null, diff --git a/Benchmarks/Compile/lakefile.toml b/Benchmarks/Compile/lakefile.toml index ad6971ab..a4cb6272 100644 --- a/Benchmarks/Compile/lakefile.toml +++ b/Benchmarks/Compile/lakefile.toml @@ -26,6 +26,9 @@ name = "CompileRedStep" [[lean_lib]] name = "CompileFLT" +[[lean_lib]] +name = "CompileAnthropicFLT" + [[lean_lib]] name = "CompileMutualFixtures" @@ -41,6 +44,11 @@ name = "flt" git = "https://github.com/ImperialCollegeLondon/FLT" rev = "v4.33.0" +[[require]] +name = "flt_e2e" +git = "https://github.com/anthropics/fermats-last-theorem" +rev = "aa2d8b34692b16c70f699536de0d8e75b9a3e9ef" + [[require]] name = "CompPoly" git = "https://github.com/Verified-zkEVM/CompPoly" diff --git a/Ix/Cli/BenchCmd.lean b/Ix/Cli/BenchCmd.lean index a6949bc6..13915437 100644 --- a/Ix/Cli/BenchCmd.lean +++ b/Ix/Cli/BenchCmd.lean @@ -103,7 +103,9 @@ def envSpecs : List EnvSpec := [ -- carries Init and Std). { name := "ISLB", module := "Benchmarks/Compile/CompileISLB.lean" }, { name := "Mathlib", module := "Benchmarks/Compile/CompileMathlib.lean" }, - { name := "FLT", module := "Benchmarks/Compile/CompileFLT.lean" } + { name := "FLT", module := "Benchmarks/Compile/CompileFLT.lean" }, + { name := "AnthropicFLT", + module := "Benchmarks/Compile/CompileAnthropicFLT.lean" } ] def findEnv (token : String) : Option EnvSpec := @@ -349,7 +351,11 @@ def backendSpecs : List BackendSpec := [ testbeds := [("execute", "lean4lean-check-x64-32x")], metrics := [("execute", ["check-time", "throughput", "peak-rss", "constants"])] }, + -- AnthropicFLT remains on-demand: its from-scratch upstream build needs + -- substantially more than the per-push workflow's one-hour budget. An + -- explicit `--env AnthropicFLT` or `BENCH_ENVS=AnthropicFLT` still runs it. { name := "compile", defaultMode := "execute", inputs := .perEnv, + envs := some ["InitStd", "Lean", "ISLB", "Mathlib", "FLT"], testbeds := [("execute", "ix-compile-x64-32x")], metrics := [("execute", ["compile-time", "throughput", "peak-rss", "file-size", "constants"])], From 0f1901cf9572d791128aefc001bfa00aaae4c5b8 Mon Sep 17 00:00:00 2001 From: "J. C. Burnham" Date: Sat, 5 Sep 2026 14:57:31 -0400 Subject: [PATCH 2/4] feat(compile): add memory-aware scheduling and cached-artifact retries Add --no-build for compiler retries, pin mimalloc 3.5.1, and log memory and active-block progress. Preserve expression DAG sharing during validation and main-stage auxiliary generation. Adapt validation and main-block admission to memory headroom, swap pressure, and recent growth; distinguish resource exhaustion from invalid proofs. Cover scheduler backoff, expression equivalence, worker-count determinism, and CLI/FFI behavior with tests. --- Benchmarks/Compile/README.md | 78 ++ Cargo.lock | 10 +- Cargo.toml | 3 +- Ix/Cli/CompileCmd.lean | 6 +- Ix/CompileM.lean | 4 +- Tests/Cli.lean | 35 + Tests/FFI/Ix.lean | 3 +- Tests/Gen/Ix.lean | 3 +- crates/compile/src/compile.rs | 86 +- crates/compile/src/compile/admission.rs | 577 ++++++++++++ crates/compile/src/compile/aux_gen.rs | 3 + .../src/compile/aux_gen/checked_expr.rs | 670 ++++++++++++++ .../src/compile/aux_gen/expr_reference.rs | 531 +++++++++++ .../compile/src/compile/aux_gen/expr_utils.rs | 416 +-------- crates/compile/src/compile/aux_gen/nested.rs | 373 +++++--- crates/compile/src/compile/env.rs | 75 +- crates/compile/src/compile/memory.rs | 216 +++++ crates/compile/src/compile/validation.rs | 845 ++++++++++++++++++ crates/compile/src/diag.rs | 155 +++- crates/ffi/src/compile.rs | 11 +- crates/ffi/src/lean.rs | 1 + crates/ixon/src/error.rs | 5 +- 22 files changed, 3577 insertions(+), 529 deletions(-) create mode 100644 crates/compile/src/compile/admission.rs create mode 100644 crates/compile/src/compile/aux_gen/checked_expr.rs create mode 100644 crates/compile/src/compile/aux_gen/expr_reference.rs create mode 100644 crates/compile/src/compile/memory.rs create mode 100644 crates/compile/src/compile/validation.rs diff --git a/Benchmarks/Compile/README.md b/Benchmarks/Compile/README.md index 1cba8377..3289687a 100644 --- a/Benchmarks/Compile/README.md +++ b/Benchmarks/Compile/README.md @@ -17,6 +17,84 @@ First ensure the Lean version used to build Ix matches the `Benchmarks/Compile/l `ix compile /path/to/Compile.lean` # replace `` with `Init`, `InitStd`, `Lean`, `Mathlib`, `FLT`, or `AnthropicFLT` +To retry only the Ix compiler using existing, up-to-date import artifacts: + +```sh +lake exe ix compile Benchmarks/Compile/CompileAnthropicFLT.lean --no-build --verbose +``` + +Run this from the repository root. `--no-build` skips the target's Lake +build/cache step, including checks for generated C files. The input source is +still elaborated; missing imports fail, and stale imports are not rebuilt. +The outer `lake exe` may still rebuild the Ix executable itself. + +For allocation-failure diagnostics, set `IX_MEMORY_DIAG=1` to log process +RSS, virtual memory, swap, system memory availability, and mapping counts +every five seconds during Ix compilation. `IX_LOG_IND_GROUPS=1` also logs +each inductive-validation group's entry/exit; unmatched `BEGIN` lines show +which groups were active if the process aborts. Both are opt-in and leave +the compilation algorithm unchanged. + +Ix pins the [mimalloc Rust fork](https://github.com/argumentcomputer/mimalloc_rust) +with mimalloc v3.5.1 in `Cargo.toml`. This fixes v3.3.x rejecting its own +metadata for 16 GiB arenas and eventually exhausting `vm.max_map_count` +despite available RAM and swap ([upstream issue #1309](https://github.com/microsoft/mimalloc/issues/1309)). + +Inductive-flag validation (setup stage 4) uses adaptive admission on Linux. +The Rayon pool keeps its configured size, but validation starts with two +active jobs and ramps up while memory is healthy. The controller samples +`MemAvailable`, visible cgroup-v2 ancestor limits, swap growth, and memory +stall pressure every 250 ms. It stops admissions and cancels excess attempts +under pressure; those attempts return normally, drop their scratch data, and +retry once alone after other work finishes. Completed validations are retained. +Expression walks preserve DAG sharing and check cancellation within the walk. + +This is a soft safety mechanism, not a hard allocation limit: an individual +allocation, lazy import fetch, or destruction cannot be interrupted. A lone +attempt that exhausts the safety reserve returns `resourceLimit`, not an +invalid-proof error. +Non-Linux hosts without telemetry retain ordinary parallel validation. + +`--verbose` reports `[validate_memory]` admission/pressure progress. +`IX_LOG_IND_GROUPS=1` distinguishes `END`, `CANCEL`, and `ERROR` with stable +group IDs. Optional controls: + +- `RAYON_NUM_THREADS`: the pool-size ceiling, not a fixed active-job count. +- `IX_VALIDATE_MEMORY_GIB`: an additional process RSS-plus-swap soft budget + during validation; system/cgroup headroom still applies. Swap is not + counted as extra available RAM. +- `IX_VALIDATE_ADAPTIVE=0`: disable admission control (cannot be combined + with an explicit memory budget). DAG-preserving walks remain enabled. + +The main dependency scheduler also uses adaptive admission on Linux. It +starts with up to two active blocks and samples the same memory signals +every 250 ms, plus a ten-second forecast of recent memory growth. It raises +concurrency only after completed work and sufficient headroom, stops new +admissions under pressure, and waits five seconds after recovery before +resuming. Stable retained output can continue growing with one active block +while the safety reserve remains available. Main-stage expression transforms +also preserve DAG sharing instead of repeatedly copying shared subexpressions. + +Unlike validation, main-stage blocks publish shared metadata during their +execution, so active blocks are **not cancelled or retried**: they finish and +release their scratch before the slot is reused. This is a soft safeguard, +not a hard memory cap; one large block can still exhaust memory. If no block +is active and admission cannot resume for 30 seconds, compilation returns +`resourceLimit`. The gate covers the main block scheduler, not graph setup, +final serialization, or other work outside that scheduler. Without Linux +telemetry, the scheduler retains fixed concurrency. + +`--verbose` reports `[compile_memory]` limits, active blocks, admissions, +memory, and growth forecasts. Compilation progress shows recent completions, +active blocks, and time since the last completion instead of a lifetime-average +ETA or a `STALLED` label. Optional main-stage controls: + +- `IX_COMPILE_WORKERS`: the worker-count ceiling (bounded by available CPUs). +- `IX_COMPILE_MEMORY_GIB`: an additional process RSS-plus-swap soft budget + during main compilation; system/cgroup headroom still applies. +- `IX_COMPILE_ADAPTIVE=0`: use fixed admission (cannot be combined with + `IX_COMPILE_MEMORY_GIB`). DAG-preserving transforms remain enabled. + The Anthropic artifact is also registered as the on-demand `AnthropicFLT` benchmark environment. After building its oleans, benchmark the Ix compiler and Rust kernel from the repository root with: diff --git a/Cargo.lock b/Cargo.lock index 9e385b38..fe8e1938 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1960,9 +1960,8 @@ checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" [[package]] name = "libmimalloc-sys" -version = "0.1.47" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2d1eacfa31c33ec25e873c136ba5669f00f9866d0688bea7be4d3f7e43067df6" +version = "0.1.49" +source = "git+https://github.com/argumentcomputer/mimalloc_rust.git?rev=616cc1f130d9ea2a1111a46b4e98c140e078f882#616cc1f130d9ea2a1111a46b4e98c140e078f882" dependencies = [ "cc", ] @@ -2060,9 +2059,8 @@ dependencies = [ [[package]] name = "mimalloc" -version = "0.1.50" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b3627c4272df786b9260cabaa46aec1d59c93ede723d4c3ef646c503816b0640" +version = "0.1.52" +source = "git+https://github.com/argumentcomputer/mimalloc_rust.git?rev=616cc1f130d9ea2a1111a46b4e98c140e078f882#616cc1f130d9ea2a1111a46b4e98c140e078f882" dependencies = [ "libmimalloc-sys", ] diff --git a/Cargo.toml b/Cargo.toml index 55b5135d..e728cb1d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -55,7 +55,8 @@ itertools = "0.14.0" libc = "0.2" log = "0.4" memmap2 = "0.9" -mimalloc = { version = "0.1", default-features = false } +# Bundle v3.5.1 to fix large-arena failures (microsoft/mimalloc#1309). +mimalloc = { git = "https://github.com/argumentcomputer/mimalloc_rust.git", rev = "616cc1f130d9ea2a1111a46b4e98c140e078f882", default-features = false } multi-stark = { git = "https://github.com/argumentcomputer/multi-stark.git", rev = "9a90612286625b6cec59516532d4501677cf468c" } nom = "7.1.3" num-bigint = "0.4.6" diff --git a/Ix/Cli/CompileCmd.lean b/Ix/Cli/CompileCmd.lean index d07076ca..a7998f59 100644 --- a/Ix/Cli/CompileCmd.lean +++ b/Ix/Cli/CompileCmd.lean @@ -33,7 +33,10 @@ def runCompileCmd (p : Cli.Parsed) : IO UInt32 := do let outPath : String := (p.flag? "out").map (·.as! String) |>.getD (defaultOutPathFor pathStr) - buildFile pathStr + -- Reuse existing imports when retrying the Ix compiler. The frontend still + -- elaborates the input file; only Lake's dependency build/cache step is skipped. + unless p.hasFlag "no-build" do + buildFile pathStr let fe ← getFileEnvCore pathStr let leanEnv := fe.env @@ -241,6 +244,7 @@ def compileCmd : Cli.Cmd := `[Cli| FLAGS: v, verbose; "Print compiler phase timings, scheduler progress, and serialization progress. Equivalent to IX_VERBOSE=1." + "no-build"; "Skip the Lake build/cache step and reuse existing import artifacts. The input file is still elaborated. The caller must ensure all imports are built and up to date." out : String; "Output path for serialized Ixon.Env bytes; defaults to the lowercased input file stem with `.ixe` (e.g. CompileMathlib.lean -> compilemathlib.ixe)" consts : String; "Comma-separated EXACT constant names to compile (transitive deps pulled in automatically) instead of the whole import env — e.g. `Nat.add_comm`. Same flag/shape as `ix check --consts`. Mutually exclusive with --module; --exclude does not apply." "consts-file" : String; "Additionally read seed constant names from a file (one per line; `#` comments and blank lines ignored). Unions with --consts." diff --git a/Ix/CompileM.lean b/Ix/CompileM.lean index 5e51fe61..d3a4d8c8 100644 --- a/Ix/CompileM.lean +++ b/Ix/CompileM.lean @@ -250,7 +250,7 @@ structure BlockEnv where /-! ## Compilation Error -/ -/-- Compilation error type. Variant order matches Rust CompileError (tags 0–5). -/ +/-- Compilation error type. Variant order matches Rust CompileError (tags 0–6). -/ inductive CompileError where | missingConstant (name : String) | missingAddress (addr : Address) @@ -258,6 +258,7 @@ inductive CompileError where | unsupportedExpr (desc : String) | unknownUnivParam (curr param : String) | serializeError (err : Ixon.SerializeError) + | resourceLimit (reason : String) deriving Repr, BEq instance : ToString CompileError where @@ -268,6 +269,7 @@ instance : ToString CompileError where | .unsupportedExpr desc => s!"unsupportedExpr: {desc}" | .unknownUnivParam curr param => s!"unknownUnivParam: compiling {curr}, param {param}" | .serializeError err => s!"serializeError: {err}" + | .resourceLimit reason => s!"resourceLimit: {reason}" abbrev CompileM := ReaderT (CompileEnv × BlockEnv) (ExceptT CompileError (StateT BlockState Id)) diff --git a/Tests/Cli.lean b/Tests/Cli.lean index 58c60e13..973ce475 100644 --- a/Tests/Cli.lean +++ b/Tests/Cli.lean @@ -14,8 +14,43 @@ def Tests.Cli.run (buildCmd: String) (buildArgs : Array String) (buildDir : Opti else IO.println out.stdout +private def Tests.Cli.testCompileNoBuild : IO Unit := do + let ix ← IO.FS.realPath ".lake/build/bin/ix" + let dir ← IO.FS.createTempDir + let source := dir / "NoBuild.lean" + let output := dir / "no-build.ixe" + try + -- No Lake project or compiled target exists here. Only the toolchain's + -- implicit Init imports are available; the CLI must elaborate this body. + IO.FS.writeFile source + "def noBuildMarker : Nat := 7\ntheorem noBuildProof : noBuildMarker = 7 := rfl\n" + let args := #["compile", source.toString, "--consts", "noBuildProof", + "--out", output.toString] + let built ← IO.Process.output { cmd := ix.toString, args := args.push "--no-build" } + unless built.exitCode == 0 do + throw <| IO.userError s!"compile --no-build failed:\n{built.stdout}\n{built.stderr}" + unless (← output.pathExists) && !(← IO.FS.readBinFile output).isEmpty do + throw <| IO.userError "compile --no-build did not write a nonempty .ixe" + for ext in ["olean", "c"] do + if ← (source.withExtension ext).pathExists then + throw <| IO.userError s!"compile --no-build unexpectedly wrote a .{ext} file" + -- The default path must still require a Lake project/build. + let defaultRun ← IO.Process.output { cmd := ix.toString, args } + if defaultRun.exitCode == 0 then + throw <| IO.userError "compile without --no-build unexpectedly skipped the Lake build" + -- Missing imports must fail instead of silently fetching/building them. + IO.FS.removeFile output + IO.FS.writeFile source "import IxCliNoBuildMissingImport\n" + let missing ← IO.Process.output { cmd := ix.toString, args := args.push "--no-build" } + if missing.exitCode == 0 || (← output.pathExists) then + throw <| IO.userError "compile --no-build accepted a missing import" + IO.println "compile --no-build: source elaboration, output, default build, and missing-import checks passed" + finally + IO.FS.removeDirAll dir + public def Tests.Cli.suite : IO UInt32 := do Tests.Cli.run "lake" (#["exe", "ix", "--help"]) none + Tests.Cli.testCompileNoBuild --Tests.Cli.run "ix" (#["store", "ix_test/IxTest.lean"]) none --Tests.Cli.run "ix" (#["prove", "ix_test/IxTest.lean", "one"]) none return 0 diff --git a/Tests/FFI/Ix.lean b/Tests/FFI/Ix.lean index 3a13de6a..22d895d9 100644 --- a/Tests/FFI/Ix.lean +++ b/Tests/FFI/Ix.lean @@ -250,7 +250,8 @@ def compileErrorTests : TestSeq := test "CompileError.invalidMutualBlock" (roundtripCompileError (.invalidMutualBlock "empty") == .invalidMutualBlock "empty") ++ test "CompileError.unsupportedExpr" (roundtripCompileError (.unsupportedExpr "mvar") == .unsupportedExpr "mvar") ++ test "CompileError.unknownUnivParam" (roundtripCompileError (.unknownUnivParam "Nat" "u") == .unknownUnivParam "Nat" "u") ++ - test "CompileError.serializeError" (roundtripCompileError (.serializeError se) == .serializeError se) + test "CompileError.serializeError" (roundtripCompileError (.serializeError se) == .serializeError se) ++ + test "CompileError.resourceLimit" (roundtripCompileError (.resourceLimit "memory reserve") == .resourceLimit "memory reserve") /-! ## Test Suite -/ diff --git a/Tests/Gen/Ix.lean b/Tests/Gen/Ix.lean index 013af41e..a0822088 100644 --- a/Tests/Gen/Ix.lean +++ b/Tests/Gen/Ix.lean @@ -665,7 +665,8 @@ def genCompileError : Gen Ix.CompileM.CompileError := do (1, pure (.invalidMutualBlock s)), (1, pure (.unsupportedExpr s)), (1, do let s2 ← genIxString; pure (.unknownUnivParam s s2)), - (1, pure (.serializeError se)) + (1, pure (.serializeError se)), + (1, pure (.resourceLimit s)) ] (pure default) instance : Shrinkable Ix.CompileM.CompileError where diff --git a/crates/compile/src/compile.rs b/crates/compile/src/compile.rs index cae62d13..facc0484 100644 --- a/crates/compile/src/compile.rs +++ b/crates/compile/src/compile.rs @@ -86,8 +86,8 @@ pub static IX_SURGERY_APPLY_DEBUG: std::sync::LazyLock = /// Options controlling whole-environment compilation. #[derive(Clone, Copy, Debug, Default)] pub struct CompileOptions { - /// Override scheduler worker count. `None` uses available parallelism or - /// the `IX_COMPILE_WORKERS` environment variable if set. + /// Override the scheduler worker ceiling. `None` uses available parallelism + /// or `IX_COMPILE_WORKERS`; adaptive admission may run fewer active blocks. pub max_workers: Option, } @@ -4894,11 +4894,14 @@ fn compile_mutual( .map(|r| r.clone()) } +mod admission; pub mod aux_gen; mod env; +mod memory; pub mod mutual; pub mod nat_conv; pub mod surgery; +pub(crate) mod validation; pub use env::{compile_env, compile_env_with_options}; #[cfg(test)] @@ -5660,6 +5663,85 @@ mod tests { assert_eq!(stt.env.const_count(), 2); } + #[test] + fn test_compile_env_worker_limits_preserve_serialized_dependency_graph() { + use ix_common::env::{ + AxiomVal, ConstantVal, DefinitionSafety, DefinitionVal, + }; + let base = Name::str(Name::anon(), "Base".into()); + let typ = LeanExpr::sort(Level::succ(Level::zero())); + let mut source = LeanEnv::default(); + source.insert( + base.clone(), + LeanConstantInfo::AxiomInfo(AxiomVal { + cnst: ConstantVal { + name: base.clone(), + level_params: vec![], + typ: typ.clone(), + }, + is_unsafe: false, + }), + ); + let mut previous = vec![base; 8]; + for layer in 0..5 { + for (column, prior) in previous.iter_mut().enumerate() { + let name = Name::str(Name::anon(), format!("alias_{layer}_{column}")); + source.insert( + name.clone(), + LeanConstantInfo::DefnInfo(DefinitionVal { + cnst: ConstantVal { + name: name.clone(), + level_params: vec![], + typ: typ.clone(), + }, + value: LeanExpr::cnst(prior.clone(), vec![]), + hints: ReducibilityHints::Abbrev, + safety: DefinitionSafety::Safe, + all: vec![name.clone()], + }), + ); + *prior = name; + } + } + // A final block depends on every branch, testing dependency publication + // as well as byte-identical metadata for alpha-equivalent aliases. + let join = Name::str(Name::anon(), "Join".into()); + let mut value = LeanExpr::sort(Level::zero()); + for name in previous { + value = LeanExpr::all( + Name::anon(), + LeanExpr::cnst(name, vec![]), + value, + BinderInfo::Default, + ); + } + source.insert( + join.clone(), + LeanConstantInfo::DefnInfo(DefinitionVal { + cnst: ConstantVal { name: join.clone(), level_params: vec![], typ }, + value, + hints: ReducibilityHints::Abbrev, + safety: DefinitionSafety::Safe, + all: vec![join], + }), + ); + let source = Arc::new(source); + let mut outputs = Vec::new(); + for max_workers in [1, 4] { + let compiled = compile_env_with_options( + &source, + CompileOptions { max_workers: Some(max_workers) }, + ) + .unwrap(); + assert!(compiled.ungrounded.is_empty()); + assert_eq!(compiled.name_to_addr.len(), 42); + let mut bytes = Vec::new(); + compiled.env.put(&mut bytes).unwrap(); + outputs.push(bytes); + } + assert_eq!(outputs[0], outputs[1]); + } + /// Test that alpha-equivalent mutual definitions produce correct projection /// indices. Two definitions with identical type/value structure (but different /// names) should form one equivalence class, and projections should resolve diff --git a/crates/compile/src/compile/admission.rs b/crates/compile/src/compile/admission.rs new file mode 100644 index 00000000..9039293b --- /dev/null +++ b/crates/compile/src/compile/admission.rs @@ -0,0 +1,577 @@ +//! Non-preemptive memory-aware admission for the main dependency scheduler. +//! +//! A main-compiler block publishes shared names and auxiliary metadata before +//! returning, unlike a pure inductive validation. It cannot safely be retried +//! after arbitrary cancellation. Permits are therefore acquired BEFORE +//! dequeuing work and released AFTER block-local scratch and publication. +//! Backoff stops admissions; already-active blocks keep running, never park +//! holding their scratch. This is not an allocator-enforced memory limit. + +use std::sync::{Arc, Condvar, Mutex}; +use std::thread::{self, JoinHandle}; +use std::time::{Duration, Instant}; + +use super::memory::{ + GIB, MIB, Memory, MemoryReader, Pressure, budget_from_env, pressure, + resource_error, +}; +use ixon::CompileError; + +#[derive(Clone, Copy)] +struct Options { + max_workers: usize, + process_budget: Option, + tick: Duration, + ramp_interval: Duration, + recovery: Duration, + idle_timeout: Duration, + lookahead: Duration, + verbose: bool, +} + +impl Options { + fn from_env(max_workers: usize) -> Result { + Ok(Self { + max_workers: max_workers.max(1), + process_budget: budget_from_env("IX_COMPILE_MEMORY_GIB")?, + tick: Duration::from_millis(250), + ramp_interval: Duration::from_secs(1), + recovery: Duration::from_secs(5), + idle_timeout: Duration::from_secs(30), + lookahead: Duration::from_secs(10), + verbose: *super::env::IX_VERBOSE, + }) + } +} + +struct Policy { + limit: usize, + open: bool, + pressure: Pressure, + previous: Memory, + sampled: Instant, + ramped: Instant, + pressured: Option, + growth_per_second: u64, + projected_growth: u64, + completed: u64, +} + +impl Policy { + fn new(memory: Memory, now: Instant, options: Options) -> Self { + let pressure = + pressure(memory, memory, options.tick, options.process_budget); + Self { + limit: options.max_workers.min(2), + open: pressure == Pressure::Healthy, + pressure, + previous: memory, + sampled: now, + ramped: now, + pressured: (pressure != Pressure::Healthy).then_some(now), + growth_per_second: 0, + projected_growth: 0, + completed: 0, + } + } + + fn sample( + &mut self, + memory: Memory, + now: Instant, + active: usize, + completed: u64, + options: Options, + ) { + let elapsed = now.duration_since(self.sampled); + let delta = memory + .process + .saturating_sub(self.previous.process) + .max(self.previous.available.saturating_sub(memory.available)); + // Fast rise, gradual decay: reserve space for another ten seconds of + // recent growth, including external pressure, before starting more work. + // This is process-level prediction, NOT an attributed per-block footprint. + let growth = (u128::from(delta) * 1_000_000 / elapsed.as_micros().max(1)) + .min(u128::from(u64::MAX)) as u64; + self.growth_per_second = + growth.max(self.growth_per_second.saturating_mul(3) / 4); + self.projected_growth = + self.growth_per_second.saturating_mul(options.lookahead.as_secs()); + let capacity = options + .process_budget + .map_or(memory.capacity, |b| b.min(memory.capacity)); + let reserve = (capacity / 8).max(64 * MIB).min(capacity / 2); + let headroom = options.process_budget.map_or(memory.available, |b| { + memory.available.min(b.saturating_sub(memory.process)) + }); + self.pressure = + pressure(memory, self.previous, elapsed, options.process_budget); + let forecast_tight = + headroom < reserve.saturating_add(self.projected_growth); + let reclaiming = super::memory::reclaiming(memory, self.previous, elapsed); + let blocked = + self.pressure == Pressure::Critical || reclaiming || forecast_tight; + let wide_headroom = self.pressure == Pressure::Healthy + && headroom + >= reserve.saturating_mul(3).saturating_add(self.projected_growth); + if blocked { + self.limit = self.limit.min((active / 2).max(1)); + self.pressured = Some(now); + } else if !wide_headroom { + // Retained, successfully compiled data cannot be reclaimed by reducing + // workers. With stable memory and a reserve, allow one job at a time + // instead of waiting forever for the accumulator to shrink. + self.limit = 1; + } + let recovered = + self.pressured.is_none_or(|t| now.duration_since(t) >= options.recovery); + self.open = !blocked && recovered; + // A long-running block with no completions is not evidence that it is + // safe to ramp up. Require actual completed work before raising the cap. + if self.open + && wide_headroom + && completed > self.completed + && now.duration_since(self.ramped) >= options.ramp_interval + { + self.limit = (self.limit + 1).min(options.max_workers); + self.ramped = now; + } + self.previous = memory; + self.sampled = now; + self.completed = completed; + } +} + +struct State { + policy: Option, + active: usize, + completed: u64, + stopped: bool, + error: Option, + idle_since: Option, +} +struct Shared { + options: Options, + state: Mutex, + changed: Condvar, + shutdown: Condvar, +} + +pub(super) struct Admission { + shared: Arc, + monitor: Option>, +} + +/// Borrowed permit: no per-block thread or Arc allocation. +pub(super) struct Permit<'a> { + shared: &'a Shared, + completed: bool, +} +impl Permit<'_> { + pub(super) fn complete(mut self) { + self.completed = true; + } +} +impl Drop for Permit<'_> { + fn drop(&mut self) { + let mut state = self.shared.state.lock().unwrap(); + state.active -= 1; + state.completed += u64::from(self.completed); + if state.active == 0 { + state.idle_since = Some(Instant::now()); + } + self.shared.changed.notify_one(); + } +} + +impl Admission { + pub(super) fn new(max_workers: usize) -> Result { + let options = Options::from_env(max_workers)?; + let disabled = std::env::var("IX_COMPILE_ADAPTIVE").as_deref() == Ok("0"); + if disabled && options.process_budget.is_some() { + return Err(resource_error( + "IX_COMPILE_MEMORY_GIB cannot be combined with IX_COMPILE_ADAPTIVE=0", + )); + } + let mut reader = MemoryReader::new(); + let initial = if disabled { None } else { reader.read() }; + if initial.is_none() && options.process_budget.is_some() { + return Err(resource_error( + "IX_COMPILE_MEMORY_GIB requires Linux memory telemetry", + )); + } + Self::with_reader(options, initial, move || reader.read()) + } + + fn with_reader( + options: Options, + initial: Option, + mut read: impl FnMut() -> Option + Send + 'static, + ) -> Result { + let now = Instant::now(); + let shared = Arc::new(Shared { + options, + state: Mutex::new(State { + policy: initial.map(|memory| Policy::new(memory, now, options)), + active: 0, + completed: 0, + stopped: false, + error: None, + idle_since: Some(now), + }), + changed: Condvar::new(), + shutdown: Condvar::new(), + }); + if options.verbose { + eprintln!( + "[compile_memory] {}: initial={} max={} budget_gib={:?}; active blocks finish before slots are reused", + if initial.is_some() { + "adaptive admission" + } else { + "fixed admission (disabled or telemetry unavailable)" + }, + if initial.is_some() { + options.max_workers.min(2) + } else { + options.max_workers + }, + options.max_workers, + options.process_budget.map(|v| v as f64 / GIB as f64) + ); + } + let monitor = if initial.is_some() { + let shared = Arc::clone(&shared); + Some(thread::Builder::new().name("ix-compile-memory".into()).spawn(move || { + let mut logged = now; + loop { + let mut state = shared.state.lock().unwrap(); + let deadline = state.policy.as_ref().unwrap().sampled + options.tick; + // Permit releases wake waiters often; they must not cause /proc + // sampling on every block (millions of blocks in a large env). + while !state.stopped && Instant::now() < deadline { + let wait = deadline.saturating_duration_since(Instant::now()); + state = shared.shutdown.wait_timeout(state, wait).unwrap().0; + } + if state.stopped { return; } + drop(state); + let sample = read(); + let mut state = shared.state.lock().unwrap(); + if state.stopped { return; } + let Some(memory) = sample else { + state.error = Some(resource_error("memory telemetry became unavailable during main compilation")); + state.stopped = true; + shared.changed.notify_all(); + return; + }; + let now = Instant::now(); + let active = state.active; + let completed = state.completed; + let policy = state.policy.as_mut().unwrap(); + let before = (policy.limit, policy.open, policy.pressure); + policy.sample(memory, now, active, completed, options); + if options.verbose && (before != (policy.limit, policy.open, policy.pressure) + || now.duration_since(logged) >= Duration::from_secs(5)) { + eprintln!("[compile_memory] {:?}: limit={}/{} active={active} admitting={} completed={completed} available={:.1} GiB resident+swap={:.1} GiB growth_forecast={:.1} GiB", + policy.pressure, policy.limit, options.max_workers, policy.open, + memory.available as f64 / GIB as f64, memory.process as f64 / GIB as f64, + policy.projected_growth as f64 / GIB as f64); + logged = now; + } + shared.changed.notify_all(); + } + }).map_err(|e| resource_error(format!("could not start memory admission monitor: {e}")))?) + } else { + None + }; + Ok(Self { shared, monitor }) + } + + pub(super) fn acquire(&self) -> Result>, CompileError> { + let mut state = self.shared.state.lock().unwrap(); + loop { + if let Some(error) = &state.error { + return Err(error.clone()); + } + if state.stopped { + return Ok(None); + } + let (open, limit) = state + .policy + .as_ref() + .map_or((true, self.shared.options.max_workers), |p| (p.open, p.limit)); + if open && state.active < limit { + state.active += 1; + state.idle_since = None; + return Ok(Some(Permit { shared: &self.shared, completed: false })); + } + if state.active == 0 { + let start = *state.idle_since.get_or_insert_with(Instant::now); + if start.elapsed() >= self.shared.options.idle_timeout { + let error = resource_error(format!( + "not enough memory headroom to admit a compiler block; no active work for {:.0}s. Increase IX_COMPILE_MEMORY_GIB if set, or make more RAM available", + self.shared.options.idle_timeout.as_secs_f64() + )); + state.error = Some(error.clone()); + state.stopped = true; + self.shared.changed.notify_all(); + self.shared.shutdown.notify_all(); + return Err(error); + } + } + state = self + .shared + .changed + .wait_timeout(state, self.shared.options.tick) + .unwrap() + .0; + } + } + + pub(super) fn stop(&self) { + self.shared.state.lock().unwrap().stopped = true; + self.shared.changed.notify_all(); + self.shared.shutdown.notify_all(); + } + + pub(super) fn check(&self) -> Result<(), CompileError> { + self.shared.state.lock().unwrap().error.clone().map_or(Ok(()), Err) + } +} +impl Drop for Admission { + fn drop(&mut self) { + self.stop(); + if let Some(monitor) = self.monitor.take() { + monitor.join().unwrap(); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + + fn memory() -> Memory { + Memory { + capacity: 100 * GIB, + available: 90 * GIB, + process: 10 * GIB, + swap_used: 0, + stall_us: 0, + } + } + fn options(max_workers: usize) -> Options { + Options { + max_workers, + process_budget: None, + tick: Duration::from_millis(2), + ramp_interval: Duration::from_millis(5), + recovery: Duration::from_millis(5), + idle_timeout: Duration::from_millis(100), + lookahead: Duration::from_secs(10), + verbose: false, + } + } + fn until(mut predicate: impl FnMut() -> bool) { + let start = Instant::now(); + while !predicate() { + assert!( + start.elapsed() < Duration::from_secs(5), + "test made no progress" + ); + thread::sleep(Duration::from_millis(1)); + } + } + + #[test] + fn healthy_memory_requires_completions_to_ramp_and_respects_the_ceiling() { + let opt = options(4); + let now = Instant::now(); + let mut policy = Policy::new(memory(), now, opt); + for second in 1..10 { + policy.sample(memory(), now + Duration::from_secs(second), 2, 0, opt); + assert_eq!(policy.limit, 2, "a long job should not trigger ramp-up"); + } + for second in 10..20 { + policy.sample( + memory(), + now + Duration::from_secs(second), + 2, + second, + opt, + ); + assert!(policy.limit <= 4); + assert!(policy.open); + } + assert_eq!(policy.limit, 4); + } + + #[test] + fn growth_forecast_stops_admissions_before_low_memory_threshold() { + let opt = options(64); + let now = Instant::now(); + let mut policy = Policy::new(memory(), now, opt); + let rising = Memory { process: 16 * GIB, available: 84 * GIB, ..memory() }; + policy.sample(rising, now + Duration::from_millis(250), 2, 0, opt); + assert_eq!(policy.pressure, Pressure::Healthy); + assert_eq!(policy.projected_growth, 240 * GIB); + assert!( + !policy.open, + "high current availability must not hide a growth burst" + ); + assert_eq!(policy.limit, 1); + } + + #[test] + fn swap_pressure_reduces_limit_and_recovers_with_hysteresis() { + let opt = options(8); + let now = Instant::now(); + let mut policy = Policy::new(memory(), now, opt); + policy.limit = 8; + let swapped = Memory { swap_used: 9 * MIB, ..memory() }; + policy.sample(swapped, now + Duration::from_millis(2), 8, 0, opt); + assert_eq!(policy.pressure, Pressure::Backoff); + assert_eq!(policy.limit, 4); + assert!(!policy.open); + policy.sample(swapped, now + Duration::from_millis(4), 4, 1, opt); + assert!(!policy.open, "one healthy tick is not enough to recover"); + policy.sample(swapped, now + Duration::from_millis(20), 2, 2, opt); + assert!(policy.open); + assert_eq!(policy.limit, 5); + } + + #[test] + fn process_budget_uses_resident_plus_swap_headroom() { + let opt = Options { process_budget: Some(20 * GIB), ..options(8) }; + let now = Instant::now(); + let memory = Memory { process: 19 * GIB, ..memory() }; + let policy = Policy::new(memory, now, opt); + assert!(!policy.open); + assert_eq!(policy.pressure, Pressure::Critical); + } + + #[test] + fn stable_retained_data_allows_serial_forward_progress() { + let opt = options(8); + let now = Instant::now(); + let memory = Memory { process: 70 * GIB, available: 30 * GIB, ..memory() }; + let mut policy = Policy::new(memory, now, opt); + policy.sample(memory, now + Duration::from_secs(1), 0, 0, opt); + assert_eq!(policy.pressure, Pressure::Hold); + assert!(policy.open); + assert_eq!(policy.limit, 1); + } + + #[test] + fn pressure_drains_existing_work_before_reusing_its_slot() { + let pressured = Arc::new(AtomicBool::new(false)); + let reader_pressure = Arc::clone(&pressured); + let gate = Admission::with_reader(options(2), Some(memory()), move || { + Some(if reader_pressure.load(Ordering::SeqCst) { + Memory { available: 20 * GIB, ..memory() } + } else { + memory() + }) + }) + .unwrap(); + let first = gate.acquire().unwrap().unwrap(); + let second = gate.acquire().unwrap().unwrap(); + pressured.store(true, Ordering::SeqCst); + until(|| !gate.shared.state.lock().unwrap().policy.as_ref().unwrap().open); + let started = AtomicBool::new(false); + let published = AtomicBool::new(false); + thread::scope(|scope| { + let waiting = scope.spawn(|| { + let permit = gate.acquire().unwrap().unwrap(); + started.store(true, Ordering::SeqCst); + assert!( + published.load(Ordering::SeqCst), + "slot reused before prior block published" + ); + permit.complete(); + }); + // Neither active task is cancelled or parked. A completed block's + // slot remains unavailable while pressure persists. + drop(first); + assert_eq!(gate.shared.state.lock().unwrap().active, 1); + assert!(!started.load(Ordering::SeqCst)); + pressured.store(false, Ordering::SeqCst); + until(|| gate.shared.state.lock().unwrap().policy.as_ref().unwrap().open); + assert!( + !started.load(Ordering::SeqCst), + "backoff must respect the reduced active cap" + ); + published.store(true, Ordering::SeqCst); + second.complete(); + waiting.join().unwrap(); + }); + let state = gate.shared.state.lock().unwrap(); + assert_eq!(state.active, 0); + assert_eq!(state.completed, 2); + } + + #[test] + fn idle_permits_do_not_count_as_completed_blocks() { + let gate = Admission::with_reader(options(2), None, || { + panic!("fixed mode sampled memory") + }) + .unwrap(); + drop(gate.acquire().unwrap().unwrap()); + assert_eq!(gate.shared.state.lock().unwrap().completed, 0); + gate.acquire().unwrap().unwrap().complete(); + assert_eq!(gate.shared.state.lock().unwrap().completed, 1); + } + + #[test] + fn stop_wakes_waiting_workers_even_while_other_permits_are_live() { + let gate = Admission::with_reader(options(1), None, || None).unwrap(); + let first = gate.acquire().unwrap().unwrap(); + thread::scope(|scope| { + let waiting = scope.spawn(|| assert!(gate.acquire().unwrap().is_none())); + gate.stop(); + waiting.join().unwrap(); + }); + drop(first); + assert_eq!(gate.shared.state.lock().unwrap().active, 0); + } + + #[test] + fn no_headroom_returns_resource_error_without_starting_a_block() { + let low = Memory { available: GIB, ..memory() }; + let gate = + Admission::with_reader(options(4), Some(low), move || Some(low)).unwrap(); + assert!(matches!(gate.acquire(), Err(CompileError::ResourceLimit { .. }))); + assert_eq!(gate.shared.state.lock().unwrap().active, 0); + assert!(matches!(gate.check(), Err(CompileError::ResourceLimit { .. }))); + } + + #[test] + fn telemetry_failure_is_visible_even_if_no_worker_acquires_again() { + let gate = + Admission::with_reader(options(1), Some(memory()), || None).unwrap(); + let last = gate.acquire().unwrap().unwrap(); + until(|| gate.check().is_err()); + last.complete(); + assert!(matches!(gate.check(), Err(CompileError::ResourceLimit { .. }))); + assert!(matches!(gate.acquire(), Err(CompileError::ResourceLimit { .. }))); + } + + #[test] + fn monitor_does_not_resample_for_each_block_and_shutdown_is_prompt() { + let reads = Arc::new(AtomicUsize::new(0)); + let reader_count = Arc::clone(&reads); + let opt = Options { tick: Duration::from_secs(10), ..options(2) }; + let gate = Admission::with_reader(opt, Some(memory()), move || { + reader_count.fetch_add(1, Ordering::SeqCst); + Some(memory()) + }) + .unwrap(); + for _ in 0..1000 { + gate.acquire().unwrap().unwrap().complete(); + } + assert_eq!(reads.load(Ordering::SeqCst), 0); + let start = Instant::now(); + drop(gate); + assert!(start.elapsed() < Duration::from_secs(1)); + } +} diff --git a/crates/compile/src/compile/aux_gen.rs b/crates/compile/src/compile/aux_gen.rs index 0f1b207b..c687174d 100644 --- a/crates/compile/src/compile/aux_gen.rs +++ b/crates/compile/src/compile/aux_gen.rs @@ -77,6 +77,9 @@ pub mod below; pub mod brecon; pub mod cases_on; +mod checked_expr; +#[cfg(test)] +mod expr_reference; pub mod expr_utils; pub mod nested; pub mod rec_on; diff --git a/crates/compile/src/compile/aux_gen/checked_expr.rs b/crates/compile/src/compile/aux_gen/checked_expr.rs new file mode 100644 index 00000000..46246641 --- /dev/null +++ b/crates/compile/src/compile/aux_gen/checked_expr.rs @@ -0,0 +1,670 @@ +//! Cancellable, DAG-preserving versions of the transforms used by validation. +//! Caches are local to one substitution context, and include binder depth: +//! the same shared BVar can mean different things beneath different binders. + +use bignat::Nat; +use blake3::Hash; +use ix_common::env::{Expr as LeanExpr, ExprData, Level, LevelData, Name}; +use rustc_hash::{FxHashMap, FxHashSet}; + +use super::expr_utils::{LocalDecl, fresh_fvar}; +use crate::compile::nat_conv::nat_to_u64; +use crate::compile::validation::{Cancelled, Checkpoint}; + +type Key = (Hash, u64); +type Result = std::result::Result; + +fn rewrite( + e: &LeanExpr, + depth: u64, + control: &Checkpoint, + cache: &mut FxHashMap, + leaf: &mut impl FnMut(&LeanExpr, u64) -> Result>, +) -> Result { + control.visit()?; + let key = (*e.get_hash(), depth); + if let Some(value) = cache.get(&key) { + return Ok(value.clone()); + } + let result = if let Some(value) = leaf(e, depth)? { + value + } else { + match e.as_data() { + ExprData::App(f, a, _) => LeanExpr::app( + rewrite(f, depth, control, cache, leaf)?, + rewrite(a, depth, control, cache, leaf)?, + ), + ExprData::Lam(n, t, b, bi, _) => LeanExpr::lam( + n.clone(), + rewrite(t, depth, control, cache, leaf)?, + rewrite(b, depth + 1, control, cache, leaf)?, + bi.clone(), + ), + ExprData::ForallE(n, t, b, bi, _) => LeanExpr::all( + n.clone(), + rewrite(t, depth, control, cache, leaf)?, + rewrite(b, depth + 1, control, cache, leaf)?, + bi.clone(), + ), + ExprData::LetE(n, t, v, b, nd, _) => LeanExpr::letE( + n.clone(), + rewrite(t, depth, control, cache, leaf)?, + rewrite(v, depth, control, cache, leaf)?, + rewrite(b, depth + 1, control, cache, leaf)?, + *nd, + ), + ExprData::Proj(n, i, v, _) => LeanExpr::proj( + n.clone(), + i.clone(), + rewrite(v, depth, control, cache, leaf)?, + ), + ExprData::Mdata(md, v, _) => { + LeanExpr::mdata(md.clone(), rewrite(v, depth, control, cache, leaf)?) + }, + _ => e.clone(), + } + }; + // Avoid retaining a second copy of unchanged nodes and their children. + let result = + if result.get_hash() == e.get_hash() { e.clone() } else { result }; + cache.insert(key, result.clone()); + control.scratch( + cache.capacity() * (size_of::<(Key, LeanExpr)>() + size_of::()), + ); + Ok(result) +} + +fn transform( + e: &LeanExpr, + control: &Checkpoint, + mut leaf: impl FnMut(&LeanExpr, u64) -> Result>, +) -> Result { + rewrite(e, 0, control, &mut FxHashMap::default(), &mut leaf) +} + +pub(super) fn instantiate1( + e: &LeanExpr, + replacement: &LeanExpr, + c: &Checkpoint, +) -> Result { + instantiate1_at(e, replacement, 0, c) +} + +pub(super) fn instantiate1_at( + e: &LeanExpr, + replacement: &LeanExpr, + depth: u64, + c: &Checkpoint, +) -> Result { + rewrite(e, depth, c, &mut FxHashMap::default(), &mut |e, depth| { + Ok(match e.as_data() { + ExprData::Bvar(i, _) if nat_to_u64(i) == depth => { + Some(replacement.clone()) + }, + ExprData::Bvar(i, _) if nat_to_u64(i) > depth => { + Some(LeanExpr::bvar(Nat::from(nat_to_u64(i) - 1))) + }, + _ => None, + }) + }) +} + +pub(super) fn instantiate_rev( + e: &LeanExpr, + args: &[LeanExpr], + c: &Checkpoint, +) -> Result { + if args.is_empty() { + c.visit()?; + return Ok(e.clone()); + } + let mut shifted = FxHashMap::default(); + transform(e, c, |e, depth| { + let ExprData::Bvar(i, _) = e.as_data() else { + return Ok(None); + }; + let i = nat_to_u64(i); + if i < depth { + return Ok(None); + } + let index = i - depth; + if index >= args.len() as u64 { + return Ok(Some(LeanExpr::bvar(Nat::from(i - args.len() as u64)))); + } + if let Some(value) = shifted.get(&(index, depth)) { + return Ok(Some(LeanExpr::clone(value))); + } + let value = if depth == 0 { + args[index as usize].clone() + } else { + transform(&args[index as usize], c, |e, inner| { + Ok(match e.as_data() { + ExprData::Bvar(i, _) if nat_to_u64(i) >= inner => { + Some(LeanExpr::bvar(Nat::from(nat_to_u64(i) + depth))) + }, + _ => None, + }) + })? + }; + shifted.insert((index, depth), value.clone()); + Ok(Some(value)) + }) +} + +pub(super) fn instantiate_pi_params( + e: &LeanExpr, + n: usize, + args: &[LeanExpr], + c: &Checkpoint, +) -> Result { + debug_assert!(args.len() >= n); + let mut cur = e.clone(); + for arg in args.iter().take(n) { + c.visit()?; + if let ExprData::ForallE(_, _, body, _, _) = cur.as_data() { + cur = instantiate_rev(body, std::slice::from_ref(arg), c)?; + } else { + break; + } + } + Ok(cur) +} + +pub(super) fn forall_telescope( + e: &LeanExpr, + n: usize, + prefix: &str, + start: usize, + c: &Checkpoint, +) -> Result<(Vec, Vec, LeanExpr)> { + let mut fvars = Vec::new(); + let mut decls = Vec::new(); + let mut cur = e.clone(); + for i in 0..n { + c.visit()?; + while let ExprData::Mdata(_, inner, _) = cur.as_data() { + c.visit()?; + cur = inner.clone(); + } + let ExprData::ForallE(name, dom, body, bi, _) = cur.as_data() else { + break; + }; + let (fvar_name, fv) = fresh_fvar(prefix, start + i); + decls.push(LocalDecl { + fvar_name, + binder_name: name.clone(), + domain: dom.clone(), + info: bi.clone(), + }); + fvars.push(fv.clone()); + cur = instantiate1(body, &fv, c)?; + } + Ok((fvars, decls, cur)) +} + +fn batch_abstract( + e: &LeanExpr, + vars: &FxHashMap, + scope: usize, + c: &Checkpoint, +) -> Result { + batch_abstract_at(e, vars, scope, 0, c) +} + +pub(super) fn batch_abstract_at( + e: &LeanExpr, + vars: &FxHashMap, + scope: usize, + depth: u64, + c: &Checkpoint, +) -> Result { + if scope == 0 { + c.visit()?; + return Ok(e.clone()); + } + rewrite(e, depth, c, &mut FxHashMap::default(), &mut |e, depth| { + Ok(match e.as_data() { + ExprData::Fvar(n, _) => { + vars.get(n).filter(|&&pos| pos < scope).map(|&pos| { + LeanExpr::bvar(Nat::from((scope - 1 - pos) as u64 + depth)) + }) + }, + ExprData::Bvar(i, _) if nat_to_u64(i) >= depth => { + Some(LeanExpr::bvar(Nat::from(nat_to_u64(i) + scope as u64))) + }, + _ => None, + }) + }) +} + +pub(super) fn mk_forall( + body: LeanExpr, + binders: &[LocalDecl], + c: &Checkpoint, +) -> Result { + c.visit()?; + if binders.is_empty() { + return Ok(body); + } + let vars = + binders.iter().enumerate().map(|(i, d)| (d.fvar_name.clone(), i)).collect(); + let mut result = batch_abstract(&body, &vars, binders.len(), c)?; + for (j, decl) in binders.iter().enumerate().rev() { + c.visit()?; + let domain = batch_abstract(&decl.domain, &vars, j, c)?; + result = LeanExpr::all( + decl.binder_name.clone(), + domain, + result, + decl.info.clone(), + ); + } + Ok(result) +} + +pub(super) fn replace_params( + e: &LeanExpr, + from: &[LeanExpr], + to: &[LeanExpr], + c: &Checkpoint, +) -> Result { + if from.is_empty() { + c.visit()?; + return Ok(e.clone()); + } + let vars: FxHashMap<_, _> = from + .iter() + .zip(to) + .filter_map(|(a, b)| match a.as_data() { + ExprData::Fvar(n, _) => Some((n.clone(), b.clone())), + _ => None, + }) + .collect(); + transform(e, c, |e, _| { + Ok(match e.as_data() { + ExprData::Fvar(n, _) => vars.get(n).cloned(), + _ => None, + }) + }) +} + +pub(super) fn subst_fvar( + e: &LeanExpr, + name: &Name, + value: &LeanExpr, + c: &Checkpoint, +) -> Result { + transform(e, c, |e, _| { + Ok(match e.as_data() { + ExprData::Fvar(n, _) if n == name => Some(value.clone()), + _ => None, + }) + }) +} + +pub(super) fn shift_vars( + e: &LeanExpr, + amount: usize, + cutoff: usize, + lower: bool, + c: &Checkpoint, +) -> Result { + c.visit()?; + if amount == 0 { + return Ok(e.clone()); + } + rewrite(e, cutoff as u64, c, &mut FxHashMap::default(), &mut |e, depth| { + Ok(match e.as_data() { + ExprData::Bvar(i, _) + if nat_to_u64(i) >= depth + if lower { amount as u64 } else { 0 } => + { + let i = nat_to_u64(i); + Some(LeanExpr::bvar(Nat::from(if lower { + i - amount as u64 + } else { + i + amount as u64 + }))) + }, + _ => None, + }) + }) +} + +pub(super) fn subst_levels( + e: &LeanExpr, + params: &[Name], + univs: &[Level], + c: &Checkpoint, +) -> Result { + if params.is_empty() || univs.is_empty() { + c.visit()?; + return Ok(e.clone()); + } + fn level( + l: &Level, + params: &[Name], + univs: &[Level], + c: &Checkpoint, + cache: &mut FxHashMap, + ) -> Result { + c.visit()?; + let key = *l.get_hash(); + if let Some(value) = cache.get(&key) { + return Ok(value.clone()); + } + let value = match l.as_data() { + LevelData::Param(n, _) => params + .iter() + .position(|p| p == n) + .and_then(|i| univs.get(i)) + .unwrap_or(l) + .clone(), + LevelData::Succ(v, _) => Level::succ(level(v, params, univs, c, cache)?), + LevelData::Max(a, b, _) => Level::max_smart( + level(a, params, univs, c, cache)?, + level(b, params, univs, c, cache)?, + ), + LevelData::Imax(a, b, _) => Level::imax_smart( + level(a, params, univs, c, cache)?, + level(b, params, univs, c, cache)?, + ), + _ => l.clone(), + }; + cache.insert(key, value.clone()); + Ok(value) + } + let mut levels = FxHashMap::default(); + transform(e, c, |e, _| { + Ok(match e.as_data() { + ExprData::Sort(l, _) => { + Some(LeanExpr::sort(level(l, params, univs, c, &mut levels)?)) + }, + ExprData::Const(n, ls, _) => Some(LeanExpr::cnst( + n.clone(), + ls.iter() + .map(|l| level(l, params, univs, c, &mut levels)) + .collect::>()?, + )), + _ => None, + }) + }) +} + +/// DAG-memoized predicate traversal, including binder depth for free-BVar tests. +pub(super) fn any( + e: &LeanExpr, + c: &Checkpoint, + mut predicate: impl FnMut(&LeanExpr, u64) -> bool, +) -> Result { + let mut seen = FxHashSet::default(); + let mut stack = vec![(e, 0)]; + while let Some((e, depth)) = stack.pop() { + c.visit()?; + if !seen.insert((*e.get_hash(), depth)) { + continue; + } + if predicate(e, depth) { + return Ok(true); + } + match e.as_data() { + ExprData::App(f, a, _) => { + stack.push((f, depth)); + stack.push((a, depth)); + }, + ExprData::Lam(_, t, b, _, _) | ExprData::ForallE(_, t, b, _, _) => { + stack.push((t, depth)); + stack.push((b, depth + 1)); + }, + ExprData::LetE(_, t, v, b, _, _) => { + stack.push((t, depth)); + stack.push((v, depth)); + stack.push((b, depth + 1)); + }, + ExprData::Proj(_, _, v, _) | ExprData::Mdata(_, v, _) => { + stack.push((v, depth)) + }, + _ => {}, + } + c.scratch( + seen.capacity() * size_of::() + + stack.capacity() * size_of::<(&LeanExpr, u64)>(), + ); + } + Ok(false) +} + +#[cfg(test)] +mod tests { + use super::super::expr_reference as old; + use super::*; + use ix_common::env::BinderInfo; + + fn name(s: &str) -> Name { + Name::str(Name::anon(), s.into()) + } + fn b(i: u64) -> LeanExpr { + LeanExpr::bvar(Nat::from(i)) + } + fn pi(t: LeanExpr, body: LeanExpr) -> LeanExpr { + LeanExpr::all(name("x"), t, body, BinderInfo::Default) + } + + #[test] + fn matches_reference_transforms_with_shared_terms_at_different_depths() { + let c = Checkpoint::default(); + let u = name("u"); + let f = LeanExpr::fvar(name("f")); + let leaf = LeanExpr::app(b(1), f.clone()); + let terms = vec![ + b(0), + b(3), + f.clone(), + LeanExpr::sort(Level::param(u.clone())), + LeanExpr::cnst( + name("C"), + vec![Level::imax(Level::param(u.clone()), Level::zero())], + ), + pi(leaf.clone(), pi(leaf.clone(), LeanExpr::app(leaf.clone(), b(2)))), + LeanExpr::lam( + name("x"), + leaf.clone(), + leaf.clone(), + BinderInfo::Implicit, + ), + LeanExpr::letE(name("v"), leaf.clone(), b(0), leaf.clone(), false), + LeanExpr::proj(name("P"), Nat::from(0u64), leaf), + ]; + let vars = [(name("f"), 0), (name("g"), 1)].into_iter().collect(); + for e in &terms { + assert_eq!(instantiate1(e, &f, &c).unwrap(), old::instantiate1(e, &f)); + assert_eq!( + instantiate_rev(e, &[b(2), f.clone()], &c).unwrap(), + old::instantiate_rev(e, &[b(2), f.clone()]) + ); + for scope in 0..=2 { + assert_eq!( + batch_abstract(e, &vars, scope, &c).unwrap(), + old::batch_abstract(e, &vars, scope, 0) + ); + } + assert_eq!( + subst_levels(e, std::slice::from_ref(&u), &[Level::zero()], &c) + .unwrap(), + old::subst_levels(e, std::slice::from_ref(&u), &[Level::zero()]) + ); + let ty = pi(e.clone(), pi(e.clone(), b(1))); + assert_eq!( + instantiate_pi_params(&ty, 1, &[b(2)], &c).unwrap(), + old::instantiate_pi_params(&ty, 1, &[b(2)]) + ); + let (fvars, decls, body) = + forall_telescope(&ty, 2, "test", 3, &c).unwrap(); + let (ref_fvars, ref_decls, ref_body) = + old::forall_telescope(&ty, 2, "test", 3); + assert_eq!(fvars, ref_fvars); + assert_eq!(body, ref_body); + assert_eq!( + mk_forall(body, &decls, &c).unwrap(), + old::mk_forall(ref_body, &ref_decls) + ); + } + } + + #[test] + fn diamond_dag_is_visited_linearly() { + let mut e = b(0); + for _ in 0..40 { + e = LeanExpr::app(e.clone(), e); + } + let mut visits = 0; + let result = transform(&e, &Checkpoint::default(), |_, _| { + visits += 1; + Ok(None) + }) + .unwrap(); + assert_eq!(result, e); + assert_eq!(visits, 41); // Not 2^40 visits/copies. + } + + #[test] + fn cancellation_unwinds_an_in_progress_walk_without_panicking() { + let mut e = b(0); + for _ in 0..40 { + e = LeanExpr::app(e.clone(), e); + } + let mut visits = 0; + let result = transform(&e, &Checkpoint::default(), |_, _| { + visits += 1; + if visits == 20 { Err(Cancelled) } else { Ok(None) } + }); + assert!(matches!(result, Err(Cancelled))); + assert_eq!(visits, 20); + } + + #[test] + fn predicate_cache_distinguishes_binder_depth() { + // BVar 0 is bound in the body (visited first), but free in the domain. + let e = pi(b(0), b(0)); + assert!(any(&e, &Checkpoint::default(), |e, depth| + matches!(e.as_data(), ExprData::Bvar(i, _) if nat_to_u64(i) >= depth)).unwrap()); + } + + #[test] + fn main_compiler_helpers_match_frozen_reference_on_mixed_shared_terms() { + use super::super::expr_utils as main; + let f = name("f"); + let g = name("g"); + let u = name("u"); + let fv = LeanExpr::fvar(f.clone()); + let mut terms = vec![ + b(0), + b(1), + b(3), + fv.clone(), + LeanExpr::fvar(g.clone()), + LeanExpr::sort(Level::param(u.clone())), + LeanExpr::cnst( + name("C"), + vec![Level::max(Level::param(u.clone()), Level::zero())], + ), + ]; + for round in 0..5 { + let previous = terms.clone(); + for (i, e) in previous.iter().take(12).enumerate() { + let other = &previous[(i * 7 + round) % previous.len()]; + terms.push(match (i + round) % 6 { + 0 => LeanExpr::app(e.clone(), e.clone()), + 1 => pi(e.clone(), other.clone()), + 2 => LeanExpr::lam( + name("x"), + e.clone(), + other.clone(), + BinderInfo::Implicit, + ), + 3 => { + LeanExpr::letE(name("x"), e.clone(), other.clone(), e.clone(), true) + }, + 4 => LeanExpr::proj(name("P"), Nat::from(0u64), e.clone()), + _ => LeanExpr::mdata(vec![], e.clone()), + }); + } + } + let vars = [(f.clone(), 0), (g.clone(), 1)].into_iter().collect(); + for e in &terms { + for depth in 0..3 { + assert_eq!( + main::instantiate1_at(e, &fv, depth), + old::instantiate1_at(e, &fv, depth) + ); + assert_eq!( + main::batch_abstract(e, &vars, 2, depth), + old::batch_abstract(e, &vars, 2, depth) + ); + } + assert_eq!( + main::instantiate_rev(e, &[b(1), fv.clone()]), + old::instantiate_rev(e, &[b(1), fv.clone()]) + ); + for amount in 0..3 { + for cutoff in 0..3 { + assert_eq!( + main::shift_vars(e, amount, cutoff), + old::shift_vars(e, amount, cutoff) + ); + assert_eq!( + main::lower_vars(e, amount, cutoff), + old::lower_vars(e, amount, cutoff) + ); + } + } + assert_eq!(main::subst_fvar(e, &f, &b(2)), old::subst_fvar(e, &f, &b(2))); + assert_eq!( + main::subst_levels(e, std::slice::from_ref(&u), &[Level::zero()]), + old::subst_levels(e, std::slice::from_ref(&u), &[Level::zero()]) + ); + } + } + + #[test] + fn main_compiler_transforms_preserve_large_diamond_sharing() { + use super::super::expr_utils as main; + fn diamond(mut e: LeanExpr) -> LeanExpr { + for _ in 0..40 { + e = LeanExpr::app(e.clone(), e); + } + e + } + fn unique_nodes(e: &LeanExpr) -> usize { + let mut count = 0; + any(e, &Checkpoint::default(), |_, _| { + count += 1; + false + }) + .unwrap(); + count + } + let fv = LeanExpr::fvar(name("f")); + let e = diamond(b(0)); + let instantiated = main::instantiate1(&e, &fv); + assert_eq!(instantiated.get_hash(), diamond(fv.clone()).get_hash()); + assert_eq!(unique_nodes(&instantiated), 41); + let vars = [(name("f"), 0)].into_iter().collect(); + let abstracted = main::batch_abstract(&instantiated, &vars, 1, 0); + assert_eq!(abstracted.get_hash(), e.get_hash()); + assert_eq!(unique_nodes(&abstracted), 41); + let shifted = main::shift_vars(&e, 2, 0); + assert_eq!(shifted.get_hash(), diamond(b(2)).get_hash()); + assert_eq!(unique_nodes(&shifted), 41); + let replaced = main::subst_fvar(&instantiated, &name("f"), &b(1)); + assert_eq!(replaced.get_hash(), diamond(b(1)).get_hash()); + assert_eq!(unique_nodes(&replaced), 41); + let universe = diamond(LeanExpr::sort(Level::param(name("u")))); + let substituted = + main::subst_levels(&universe, &[name("u")], &[Level::zero()]); + assert_eq!( + substituted.get_hash(), + diamond(LeanExpr::sort(Level::zero())).get_hash() + ); + assert_eq!(unique_nodes(&substituted), 41); + } +} diff --git a/crates/compile/src/compile/aux_gen/expr_reference.rs b/crates/compile/src/compile/aux_gen/expr_reference.rs new file mode 100644 index 00000000..ae660cf0 --- /dev/null +++ b/crates/compile/src/compile/aux_gen/expr_reference.rs @@ -0,0 +1,531 @@ +//! Frozen pre-memoization implementations for differential tests only. +#![allow(dead_code)] + +use super::expr_utils::{LocalDecl, fresh_fvar}; +use crate::compile::nat_conv::{nat_to_u64, nat_to_usize}; +use bignat::Nat; +use ix_common::env::{Expr as LeanExpr, ExprData, Level, LevelData, Name}; +use rustc_hash::FxHashMap; + +#[derive(Clone, Copy)] +enum BinderKind { + Forall, + Lambda, +} + +pub(super) fn forall_telescope( + expr: &LeanExpr, + n: usize, + prefix: &str, + start_idx: usize, +) -> (Vec, Vec, LeanExpr) { + let mut fvars = Vec::with_capacity(n); + let mut decls = Vec::with_capacity(n); + let mut cur = expr.clone(); + for i in 0..n { + // Peel any Mdata wrappers before matching — they're structural no-ops. + while let ExprData::Mdata(_, inner, _) = cur.as_data() { + cur = inner.clone(); + } + match cur.as_data() { + ExprData::ForallE(name, dom, body, bi, _) => { + let (fv_name, fv) = fresh_fvar(prefix, start_idx + i); + decls.push(LocalDecl { + fvar_name: fv_name, + binder_name: name.clone(), + domain: dom.clone(), + info: bi.clone(), + }); + fvars.push(fv.clone()); + cur = instantiate1(body, &fv); + }, + _ => break, + } + } + (fvars, decls, cur) +} + +pub(super) fn mk_forall(body: LeanExpr, binders: &[LocalDecl]) -> LeanExpr { + mk_binder_chain(body, binders, BinderKind::Forall) +} + +pub(super) fn mk_lambda(body: LeanExpr, binders: &[LocalDecl]) -> LeanExpr { + mk_binder_chain(body, binders, BinderKind::Lambda) +} + +fn mk_binder_chain( + body: LeanExpr, + binders: &[LocalDecl], + kind: BinderKind, +) -> LeanExpr { + let k = binders.len(); + if k == 0 { + return body; + } + + // Build FVar name → binder position map (0 = outermost). + let fvar_map: FxHashMap = + binders.iter().enumerate().map(|(i, d)| (d.fvar_name.clone(), i)).collect(); + + // Abstract body: all k binders in scope. + let mut result = batch_abstract(&body, &fvar_map, k, 0); + + // Build binder chain from innermost to outermost. + for j in (0..k).rev() { + let decl = &binders[j]; + // Domain D_j: only binders 0..j-1 are in scope (scope_depth = j). + // Binder j's domain is NOT under binder j itself — only the body is. + let domain = batch_abstract(&decl.domain, &fvar_map, j, 0); + result = match kind { + BinderKind::Forall => LeanExpr::all( + decl.binder_name.clone(), + domain, + result, + decl.info.clone(), + ), + BinderKind::Lambda => LeanExpr::lam( + decl.binder_name.clone(), + domain, + result, + decl.info.clone(), + ), + }; + } + result +} + +pub(super) fn batch_abstract( + expr: &LeanExpr, + fvar_map: &FxHashMap, + scope_depth: usize, + internal_depth: u64, +) -> LeanExpr { + // Fast path: no binders to abstract. + if scope_depth == 0 { + return expr.clone(); + } + match expr.as_data() { + ExprData::Fvar(name, _) => { + if let Some(&pos) = fvar_map.get(name) { + if pos < scope_depth { + let idx = (scope_depth - 1 - pos) as u64 + internal_depth; + LeanExpr::bvar(Nat::from(idx)) + } else { + // FVar not yet in scope (e.g., a forward reference in a domain + // to a binder declared later). Leave as-is. + expr.clone() + } + } else { + // FVar not in our telescope — leave as-is. + expr.clone() + } + }, + ExprData::Bvar(idx, _) => { + let i = nat_to_u64(idx); + if i >= internal_depth { + // Free BVar: shift up by scope_depth to make room for our binders. + LeanExpr::bvar(Nat::from(i + scope_depth as u64)) + } else { + // Bound by an expression-internal binder — unchanged. + expr.clone() + } + }, + ExprData::App(f, a, _) => LeanExpr::app( + batch_abstract(f, fvar_map, scope_depth, internal_depth), + batch_abstract(a, fvar_map, scope_depth, internal_depth), + ), + ExprData::Lam(n, t, b, bi, _) => LeanExpr::lam( + n.clone(), + batch_abstract(t, fvar_map, scope_depth, internal_depth), + batch_abstract(b, fvar_map, scope_depth, internal_depth + 1), + bi.clone(), + ), + ExprData::ForallE(n, t, b, bi, _) => LeanExpr::all( + n.clone(), + batch_abstract(t, fvar_map, scope_depth, internal_depth), + batch_abstract(b, fvar_map, scope_depth, internal_depth + 1), + bi.clone(), + ), + ExprData::LetE(n, t, v, b, nd, _) => LeanExpr::letE( + n.clone(), + batch_abstract(t, fvar_map, scope_depth, internal_depth), + batch_abstract(v, fvar_map, scope_depth, internal_depth), + batch_abstract(b, fvar_map, scope_depth, internal_depth + 1), + *nd, + ), + ExprData::Proj(n, i, e, _) => LeanExpr::proj( + n.clone(), + i.clone(), + batch_abstract(e, fvar_map, scope_depth, internal_depth), + ), + ExprData::Mdata(kvs, e, _) => LeanExpr::mdata( + kvs.clone(), + batch_abstract(e, fvar_map, scope_depth, internal_depth), + ), + // Sort, Const, MVar, Lit — no FVars or BVars to process. + _ => expr.clone(), + } +} + +pub(super) fn instantiate1( + body: &LeanExpr, + replacement: &LeanExpr, +) -> LeanExpr { + instantiate1_at(body, replacement, 0) +} + +pub(super) fn instantiate1_at( + body: &LeanExpr, + replacement: &LeanExpr, + depth: u64, +) -> LeanExpr { + match body.as_data() { + ExprData::Bvar(idx, _) => { + let i = nat_to_u64(idx); + if i == depth { + replacement.clone() + } else if i > depth { + LeanExpr::bvar(Nat::from(i - 1)) + } else { + body.clone() + } + }, + ExprData::App(f, a, _) => LeanExpr::app( + instantiate1_at(f, replacement, depth), + instantiate1_at(a, replacement, depth), + ), + ExprData::Lam(n, t, b, bi, _) => LeanExpr::lam( + n.clone(), + instantiate1_at(t, replacement, depth), + instantiate1_at(b, replacement, depth + 1), + bi.clone(), + ), + ExprData::ForallE(n, t, b, bi, _) => LeanExpr::all( + n.clone(), + instantiate1_at(t, replacement, depth), + instantiate1_at(b, replacement, depth + 1), + bi.clone(), + ), + ExprData::LetE(n, t, v, b, nd, _) => LeanExpr::letE( + n.clone(), + instantiate1_at(t, replacement, depth), + instantiate1_at(v, replacement, depth), + instantiate1_at(b, replacement, depth + 1), + *nd, + ), + ExprData::Proj(n, i, e, _) => LeanExpr::proj( + n.clone(), + i.clone(), + instantiate1_at(e, replacement, depth), + ), + ExprData::Mdata(kvs, e, _) => { + LeanExpr::mdata(kvs.clone(), instantiate1_at(e, replacement, depth)) + }, + _ => body.clone(), + } +} + +pub(super) fn instantiate_rev(body: &LeanExpr, args: &[LeanExpr]) -> LeanExpr { + if args.is_empty() { + return body.clone(); + } + instantiate_rev_at(body, args, 0) +} + +fn instantiate_rev_at( + body: &LeanExpr, + args: &[LeanExpr], + depth: u64, +) -> LeanExpr { + let n = args.len() as u64; + match body.as_data() { + ExprData::Bvar(idx, _) => { + let i = nat_to_u64(idx); + if i >= depth { + let ridx = i - depth; + if ridx < n { + // Replace with args[ridx], shifted up by depth for the binders we're under. + shift_vars(&args[ridx as usize], depth as usize, 0) + } else { + // Free BVar past our substitution range: decrement by n. + LeanExpr::bvar(Nat::from(i - n)) + } + } else { + // Bound by an expression-internal binder — unchanged. + body.clone() + } + }, + ExprData::App(f, a, _) => LeanExpr::app( + instantiate_rev_at(f, args, depth), + instantiate_rev_at(a, args, depth), + ), + ExprData::Lam(name, t, b, bi, _) => LeanExpr::lam( + name.clone(), + instantiate_rev_at(t, args, depth), + instantiate_rev_at(b, args, depth + 1), + bi.clone(), + ), + ExprData::ForallE(name, t, b, bi, _) => LeanExpr::all( + name.clone(), + instantiate_rev_at(t, args, depth), + instantiate_rev_at(b, args, depth + 1), + bi.clone(), + ), + ExprData::LetE(name, t, v, b, nd, _) => LeanExpr::letE( + name.clone(), + instantiate_rev_at(t, args, depth), + instantiate_rev_at(v, args, depth), + instantiate_rev_at(b, args, depth + 1), + *nd, + ), + ExprData::Proj(name, i, e, _) => LeanExpr::proj( + name.clone(), + i.clone(), + instantiate_rev_at(e, args, depth), + ), + ExprData::Mdata(kvs, e, _) => { + LeanExpr::mdata(kvs.clone(), instantiate_rev_at(e, args, depth)) + }, + // Sort, Const, Lit, FVar, MVar — no BVars to substitute. + _ => body.clone(), + } +} + +pub(crate) fn instantiate_pi_params( + typ: &LeanExpr, + n: usize, + args: &[LeanExpr], +) -> LeanExpr { + debug_assert!( + args.len() >= n, + "instantiate_pi_params: args.len()={} < n={}", + args.len(), + n + ); + let mut cur = typ.clone(); + for arg in args.iter().take(n) { + match cur.as_data() { + ExprData::ForallE(_, _, body, _, _) => { + cur = instantiate_rev(body, std::slice::from_ref(arg)); + }, + _ => break, + } + } + cur +} + +pub(crate) fn shift_vars( + expr: &LeanExpr, + amount: usize, + cutoff: usize, +) -> LeanExpr { + if amount == 0 { + return expr.clone(); + } + match expr.as_data() { + ExprData::Bvar(idx, _) => { + let i = nat_to_usize(idx); + if i >= cutoff { + LeanExpr::bvar(Nat::from((i + amount) as u64)) + } else { + expr.clone() + } + }, + ExprData::App(f, a, _) => LeanExpr::app( + shift_vars(f, amount, cutoff), + shift_vars(a, amount, cutoff), + ), + ExprData::Lam(n, t, b, bi, _) => LeanExpr::lam( + n.clone(), + shift_vars(t, amount, cutoff), + shift_vars(b, amount, cutoff + 1), + bi.clone(), + ), + ExprData::ForallE(n, t, b, bi, _) => LeanExpr::all( + n.clone(), + shift_vars(t, amount, cutoff), + shift_vars(b, amount, cutoff + 1), + bi.clone(), + ), + ExprData::LetE(n, t, v, b, nd, _) => LeanExpr::letE( + n.clone(), + shift_vars(t, amount, cutoff), + shift_vars(v, amount, cutoff), + shift_vars(b, amount, cutoff + 1), + *nd, + ), + ExprData::Proj(n, i, e, _) => { + LeanExpr::proj(n.clone(), i.clone(), shift_vars(e, amount, cutoff)) + }, + ExprData::Mdata(kvs, e, _) => { + LeanExpr::mdata(kvs.clone(), shift_vars(e, amount, cutoff)) + }, + _ => expr.clone(), + } +} + +pub(crate) fn lower_vars( + expr: &LeanExpr, + amount: usize, + cutoff: usize, +) -> LeanExpr { + if amount == 0 { + return expr.clone(); + } + match expr.as_data() { + ExprData::Bvar(idx, _) => { + let i = nat_to_usize(idx); + if i >= cutoff + amount { + LeanExpr::bvar(Nat::from((i - amount) as u64)) + } else { + expr.clone() + } + }, + ExprData::App(f, a, _) => LeanExpr::app( + lower_vars(f, amount, cutoff), + lower_vars(a, amount, cutoff), + ), + ExprData::Lam(n, t, b, bi, _) => LeanExpr::lam( + n.clone(), + lower_vars(t, amount, cutoff), + lower_vars(b, amount, cutoff + 1), + bi.clone(), + ), + ExprData::ForallE(n, t, b, bi, _) => LeanExpr::all( + n.clone(), + lower_vars(t, amount, cutoff), + lower_vars(b, amount, cutoff + 1), + bi.clone(), + ), + ExprData::LetE(n, t, v, b, nd, _) => LeanExpr::letE( + n.clone(), + lower_vars(t, amount, cutoff), + lower_vars(v, amount, cutoff), + lower_vars(b, amount, cutoff + 1), + *nd, + ), + ExprData::Proj(n, i, e, _) => { + LeanExpr::proj(n.clone(), i.clone(), lower_vars(e, amount, cutoff)) + }, + ExprData::Mdata(kvs, e, _) => { + LeanExpr::mdata(kvs.clone(), lower_vars(e, amount, cutoff)) + }, + _ => expr.clone(), + } +} + +pub(super) fn subst_levels( + expr: &LeanExpr, + params: &[Name], + univs: &[Level], +) -> LeanExpr { + if params.is_empty() || univs.is_empty() { + return expr.clone(); + } + match expr.as_data() { + ExprData::Sort(lvl, _) => LeanExpr::sort(subst_level(lvl, params, univs)), + ExprData::Const(name, us, _) => LeanExpr::cnst( + name.clone(), + us.iter().map(|u| subst_level(u, params, univs)).collect(), + ), + ExprData::App(f, a, _) => LeanExpr::app( + subst_levels(f, params, univs), + subst_levels(a, params, univs), + ), + ExprData::Lam(n, t, b, bi, _) => LeanExpr::lam( + n.clone(), + subst_levels(t, params, univs), + subst_levels(b, params, univs), + bi.clone(), + ), + ExprData::ForallE(n, t, b, bi, _) => LeanExpr::all( + n.clone(), + subst_levels(t, params, univs), + subst_levels(b, params, univs), + bi.clone(), + ), + ExprData::LetE(n, t, v, b, nd, _) => LeanExpr::letE( + n.clone(), + subst_levels(t, params, univs), + subst_levels(v, params, univs), + subst_levels(b, params, univs), + *nd, + ), + ExprData::Proj(n, i, e, _) => { + LeanExpr::proj(n.clone(), i.clone(), subst_levels(e, params, univs)) + }, + ExprData::Mdata(md, e, _) => { + LeanExpr::mdata(md.clone(), subst_levels(e, params, univs)) + }, + _ => expr.clone(), + } +} + +pub(super) fn subst_level( + lvl: &Level, + params: &[Name], + univs: &[Level], +) -> Level { + match lvl.as_data() { + LevelData::Zero(_) | LevelData::Mvar(_, _) => lvl.clone(), + LevelData::Succ(l, _) => Level::succ(subst_level(l, params, univs)), + LevelData::Max(a, b, _) => Level::max_smart( + subst_level(a, params, univs), + subst_level(b, params, univs), + ), + LevelData::Imax(a, b, _) => Level::imax_smart( + subst_level(a, params, univs), + subst_level(b, params, univs), + ), + LevelData::Param(name, _) => { + for (i, p) in params.iter().enumerate() { + if p == name && i < univs.len() { + return univs[i].clone(); + } + } + lvl.clone() + }, + } +} + +pub(super) fn subst_fvar( + expr: &LeanExpr, + fvar_name: &Name, + replacement: &LeanExpr, +) -> LeanExpr { + match expr.as_data() { + ExprData::Fvar(n, _) if n == fvar_name => replacement.clone(), + ExprData::App(f, a, _) => LeanExpr::app( + subst_fvar(f, fvar_name, replacement), + subst_fvar(a, fvar_name, replacement), + ), + ExprData::Lam(n, t, b, bi, _) => LeanExpr::lam( + n.clone(), + subst_fvar(t, fvar_name, replacement), + subst_fvar(b, fvar_name, replacement), + bi.clone(), + ), + ExprData::ForallE(n, t, b, bi, _) => LeanExpr::all( + n.clone(), + subst_fvar(t, fvar_name, replacement), + subst_fvar(b, fvar_name, replacement), + bi.clone(), + ), + ExprData::LetE(n, t, v, b, nd, _) => LeanExpr::letE( + n.clone(), + subst_fvar(t, fvar_name, replacement), + subst_fvar(v, fvar_name, replacement), + subst_fvar(b, fvar_name, replacement), + *nd, + ), + ExprData::Proj(n, i, e, _) => LeanExpr::proj( + n.clone(), + i.clone(), + subst_fvar(e, fvar_name, replacement), + ), + ExprData::Mdata(kvs, e, _) => { + LeanExpr::mdata(kvs.clone(), subst_fvar(e, fvar_name, replacement)) + }, + _ => expr.clone(), + } +} diff --git a/crates/compile/src/compile/aux_gen/expr_utils.rs b/crates/compile/src/compile/aux_gen/expr_utils.rs index 032d3cf8..6cfef80b 100644 --- a/crates/compile/src/compile/aux_gen/expr_utils.rs +++ b/crates/compile/src/compile/aux_gen/expr_utils.rs @@ -9,7 +9,7 @@ use rustc_hash::{FxHashMap, FxHashSet}; -use crate::compile::nat_conv::{nat_to_u64, nat_to_usize}; +use crate::compile::nat_conv::nat_to_u64; use bignat::Nat; use ix_common::address::Address; use ix_common::env::{ @@ -521,71 +521,14 @@ pub(super) fn batch_abstract( scope_depth: usize, internal_depth: u64, ) -> LeanExpr { - // Fast path: no binders to abstract. - if scope_depth == 0 { - return expr.clone(); - } - match expr.as_data() { - ExprData::Fvar(name, _) => { - if let Some(&pos) = fvar_map.get(name) { - if pos < scope_depth { - let idx = (scope_depth - 1 - pos) as u64 + internal_depth; - LeanExpr::bvar(Nat::from(idx)) - } else { - // FVar not yet in scope (e.g., a forward reference in a domain - // to a binder declared later). Leave as-is. - expr.clone() - } - } else { - // FVar not in our telescope — leave as-is. - expr.clone() - } - }, - ExprData::Bvar(idx, _) => { - let i = nat_to_u64(idx); - if i >= internal_depth { - // Free BVar: shift up by scope_depth to make room for our binders. - LeanExpr::bvar(Nat::from(i + scope_depth as u64)) - } else { - // Bound by an expression-internal binder — unchanged. - expr.clone() - } - }, - ExprData::App(f, a, _) => LeanExpr::app( - batch_abstract(f, fvar_map, scope_depth, internal_depth), - batch_abstract(a, fvar_map, scope_depth, internal_depth), - ), - ExprData::Lam(n, t, b, bi, _) => LeanExpr::lam( - n.clone(), - batch_abstract(t, fvar_map, scope_depth, internal_depth), - batch_abstract(b, fvar_map, scope_depth, internal_depth + 1), - bi.clone(), - ), - ExprData::ForallE(n, t, b, bi, _) => LeanExpr::all( - n.clone(), - batch_abstract(t, fvar_map, scope_depth, internal_depth), - batch_abstract(b, fvar_map, scope_depth, internal_depth + 1), - bi.clone(), - ), - ExprData::LetE(n, t, v, b, nd, _) => LeanExpr::letE( - n.clone(), - batch_abstract(t, fvar_map, scope_depth, internal_depth), - batch_abstract(v, fvar_map, scope_depth, internal_depth), - batch_abstract(b, fvar_map, scope_depth, internal_depth + 1), - *nd, - ), - ExprData::Proj(n, i, e, _) => LeanExpr::proj( - n.clone(), - i.clone(), - batch_abstract(e, fvar_map, scope_depth, internal_depth), - ), - ExprData::Mdata(kvs, e, _) => LeanExpr::mdata( - kvs.clone(), - batch_abstract(e, fvar_map, scope_depth, internal_depth), - ), - // Sort, Const, MVar, Lit — no FVars or BVars to process. - _ => expr.clone(), - } + super::checked_expr::batch_abstract_at( + expr, + fvar_map, + scope_depth, + internal_depth, + &Default::default(), + ) + .expect("disabled cancellation checkpoint") } // ========================================================================= @@ -607,50 +550,13 @@ pub(super) fn instantiate1_at( replacement: &LeanExpr, depth: u64, ) -> LeanExpr { - match body.as_data() { - ExprData::Bvar(idx, _) => { - let i = nat_to_u64(idx); - if i == depth { - replacement.clone() - } else if i > depth { - LeanExpr::bvar(Nat::from(i - 1)) - } else { - body.clone() - } - }, - ExprData::App(f, a, _) => LeanExpr::app( - instantiate1_at(f, replacement, depth), - instantiate1_at(a, replacement, depth), - ), - ExprData::Lam(n, t, b, bi, _) => LeanExpr::lam( - n.clone(), - instantiate1_at(t, replacement, depth), - instantiate1_at(b, replacement, depth + 1), - bi.clone(), - ), - ExprData::ForallE(n, t, b, bi, _) => LeanExpr::all( - n.clone(), - instantiate1_at(t, replacement, depth), - instantiate1_at(b, replacement, depth + 1), - bi.clone(), - ), - ExprData::LetE(n, t, v, b, nd, _) => LeanExpr::letE( - n.clone(), - instantiate1_at(t, replacement, depth), - instantiate1_at(v, replacement, depth), - instantiate1_at(b, replacement, depth + 1), - *nd, - ), - ExprData::Proj(n, i, e, _) => LeanExpr::proj( - n.clone(), - i.clone(), - instantiate1_at(e, replacement, depth), - ), - ExprData::Mdata(kvs, e, _) => { - LeanExpr::mdata(kvs.clone(), instantiate1_at(e, replacement, depth)) - }, - _ => body.clone(), - } + super::checked_expr::instantiate1_at( + body, + replacement, + depth, + &Default::default(), + ) + .expect("disabled cancellation checkpoint") } /// Multi-argument reverse instantiation: replace BVar(0)..BVar(n-1) with @@ -665,69 +571,8 @@ pub(super) fn instantiate1_at( /// argument may reference the caller's telescope (e.g. call-site surgery /// on an application under binders, as in `.brecOn_N.go` bodies). pub fn instantiate_rev(body: &LeanExpr, args: &[LeanExpr]) -> LeanExpr { - if args.is_empty() { - return body.clone(); - } - instantiate_rev_at(body, args, 0) -} - -fn instantiate_rev_at( - body: &LeanExpr, - args: &[LeanExpr], - depth: u64, -) -> LeanExpr { - let n = args.len() as u64; - match body.as_data() { - ExprData::Bvar(idx, _) => { - let i = nat_to_u64(idx); - if i >= depth { - let ridx = i - depth; - if ridx < n { - // Replace with args[ridx], shifted up by depth for the binders we're under. - shift_vars(&args[ridx as usize], depth as usize, 0) - } else { - // Free BVar past our substitution range: decrement by n. - LeanExpr::bvar(Nat::from(i - n)) - } - } else { - // Bound by an expression-internal binder — unchanged. - body.clone() - } - }, - ExprData::App(f, a, _) => LeanExpr::app( - instantiate_rev_at(f, args, depth), - instantiate_rev_at(a, args, depth), - ), - ExprData::Lam(name, t, b, bi, _) => LeanExpr::lam( - name.clone(), - instantiate_rev_at(t, args, depth), - instantiate_rev_at(b, args, depth + 1), - bi.clone(), - ), - ExprData::ForallE(name, t, b, bi, _) => LeanExpr::all( - name.clone(), - instantiate_rev_at(t, args, depth), - instantiate_rev_at(b, args, depth + 1), - bi.clone(), - ), - ExprData::LetE(name, t, v, b, nd, _) => LeanExpr::letE( - name.clone(), - instantiate_rev_at(t, args, depth), - instantiate_rev_at(v, args, depth), - instantiate_rev_at(b, args, depth + 1), - *nd, - ), - ExprData::Proj(name, i, e, _) => LeanExpr::proj( - name.clone(), - i.clone(), - instantiate_rev_at(e, args, depth), - ), - ExprData::Mdata(kvs, e, _) => { - LeanExpr::mdata(kvs.clone(), instantiate_rev_at(e, args, depth)) - }, - // Sort, Const, Lit, FVar, MVar — no BVars to substitute. - _ => body.clone(), - } + super::checked_expr::instantiate_rev(body, args, &Default::default()) + .expect("disabled cancellation checkpoint") } /// Peel `n` forall binders and substitute their variables with `args`. @@ -804,56 +649,21 @@ pub(super) fn instantiate_spec_with_fvars( /// Shift BVars UP by `amount` for BVars >= cutoff. /// -/// Used internally by `instantiate_rev_at` when substituting args under +/// Used when substituting args under /// inner binders (each args element is re-shifted by the current depth). pub(crate) fn shift_vars( expr: &LeanExpr, amount: usize, cutoff: usize, ) -> LeanExpr { - if amount == 0 { - return expr.clone(); - } - match expr.as_data() { - ExprData::Bvar(idx, _) => { - let i = nat_to_usize(idx); - if i >= cutoff { - LeanExpr::bvar(Nat::from((i + amount) as u64)) - } else { - expr.clone() - } - }, - ExprData::App(f, a, _) => LeanExpr::app( - shift_vars(f, amount, cutoff), - shift_vars(a, amount, cutoff), - ), - ExprData::Lam(n, t, b, bi, _) => LeanExpr::lam( - n.clone(), - shift_vars(t, amount, cutoff), - shift_vars(b, amount, cutoff + 1), - bi.clone(), - ), - ExprData::ForallE(n, t, b, bi, _) => LeanExpr::all( - n.clone(), - shift_vars(t, amount, cutoff), - shift_vars(b, amount, cutoff + 1), - bi.clone(), - ), - ExprData::LetE(n, t, v, b, nd, _) => LeanExpr::letE( - n.clone(), - shift_vars(t, amount, cutoff), - shift_vars(v, amount, cutoff), - shift_vars(b, amount, cutoff + 1), - *nd, - ), - ExprData::Proj(n, i, e, _) => { - LeanExpr::proj(n.clone(), i.clone(), shift_vars(e, amount, cutoff)) - }, - ExprData::Mdata(kvs, e, _) => { - LeanExpr::mdata(kvs.clone(), shift_vars(e, amount, cutoff)) - }, - _ => expr.clone(), - } + super::checked_expr::shift_vars( + expr, + amount, + cutoff, + false, + &Default::default(), + ) + .expect("disabled cancellation checkpoint") } /// Inverse of [`shift_vars`] for an expression known to have been lifted by @@ -864,49 +674,14 @@ pub(crate) fn lower_vars( amount: usize, cutoff: usize, ) -> LeanExpr { - if amount == 0 { - return expr.clone(); - } - match expr.as_data() { - ExprData::Bvar(idx, _) => { - let i = nat_to_usize(idx); - if i >= cutoff + amount { - LeanExpr::bvar(Nat::from((i - amount) as u64)) - } else { - expr.clone() - } - }, - ExprData::App(f, a, _) => LeanExpr::app( - lower_vars(f, amount, cutoff), - lower_vars(a, amount, cutoff), - ), - ExprData::Lam(n, t, b, bi, _) => LeanExpr::lam( - n.clone(), - lower_vars(t, amount, cutoff), - lower_vars(b, amount, cutoff + 1), - bi.clone(), - ), - ExprData::ForallE(n, t, b, bi, _) => LeanExpr::all( - n.clone(), - lower_vars(t, amount, cutoff), - lower_vars(b, amount, cutoff + 1), - bi.clone(), - ), - ExprData::LetE(n, t, v, b, nd, _) => LeanExpr::letE( - n.clone(), - lower_vars(t, amount, cutoff), - lower_vars(v, amount, cutoff), - lower_vars(b, amount, cutoff + 1), - *nd, - ), - ExprData::Proj(n, i, e, _) => { - LeanExpr::proj(n.clone(), i.clone(), lower_vars(e, amount, cutoff)) - }, - ExprData::Mdata(kvs, e, _) => { - LeanExpr::mdata(kvs.clone(), lower_vars(e, amount, cutoff)) - }, - _ => expr.clone(), - } + super::checked_expr::shift_vars( + expr, + amount, + cutoff, + true, + &Default::default(), + ) + .expect("disabled cancellation checkpoint") } // ========================================================================= @@ -919,84 +694,8 @@ pub fn subst_levels( params: &[Name], univs: &[Level], ) -> LeanExpr { - if params.is_empty() || univs.is_empty() { - return expr.clone(); - } - match expr.as_data() { - ExprData::Sort(lvl, _) => LeanExpr::sort(subst_level(lvl, params, univs)), - ExprData::Const(name, us, _) => LeanExpr::cnst( - name.clone(), - us.iter().map(|u| subst_level(u, params, univs)).collect(), - ), - ExprData::App(f, a, _) => LeanExpr::app( - subst_levels(f, params, univs), - subst_levels(a, params, univs), - ), - ExprData::Lam(n, t, b, bi, _) => LeanExpr::lam( - n.clone(), - subst_levels(t, params, univs), - subst_levels(b, params, univs), - bi.clone(), - ), - ExprData::ForallE(n, t, b, bi, _) => LeanExpr::all( - n.clone(), - subst_levels(t, params, univs), - subst_levels(b, params, univs), - bi.clone(), - ), - ExprData::LetE(n, t, v, b, nd, _) => LeanExpr::letE( - n.clone(), - subst_levels(t, params, univs), - subst_levels(v, params, univs), - subst_levels(b, params, univs), - *nd, - ), - ExprData::Proj(n, i, e, _) => { - LeanExpr::proj(n.clone(), i.clone(), subst_levels(e, params, univs)) - }, - ExprData::Mdata(md, e, _) => { - LeanExpr::mdata(md.clone(), subst_levels(e, params, univs)) - }, - _ => expr.clone(), - } -} - -/// Substitute universe parameters in a level. -/// -/// Uses the smart constructors `Level::max_smart` and `Level::imax_smart` so -/// that substituting away parameters produces the same canonical form the -/// kernel sees post-ingress (`KUniv::max` does the same simplifications at -/// kernel-side construction time). Without this normalization, `Max(Succ Param u, -/// Succ Param v)` substituted to `Max(Succ Zero, Succ Zero)` stays as a `Max` -/// node compile-side while the kernel collapses it to `Succ Zero` — -/// `sort_aux_by_partition_refinement` would then disagree with the kernel's -/// `canonical_aux_order` on whether two structurally-different aux types -/// (e.g. `Sort 1` vs `Sort (max 1 1)`) are equivalent. -pub(super) fn subst_level( - lvl: &Level, - params: &[Name], - univs: &[Level], -) -> Level { - match lvl.as_data() { - LevelData::Zero(_) | LevelData::Mvar(_, _) => lvl.clone(), - LevelData::Succ(l, _) => Level::succ(subst_level(l, params, univs)), - LevelData::Max(a, b, _) => Level::max_smart( - subst_level(a, params, univs), - subst_level(b, params, univs), - ), - LevelData::Imax(a, b, _) => Level::imax_smart( - subst_level(a, params, univs), - subst_level(b, params, univs), - ), - LevelData::Param(name, _) => { - for (i, p) in params.iter().enumerate() { - if p == name && i < univs.len() { - return univs[i].clone(); - } - } - lvl.clone() - }, - } + super::checked_expr::subst_levels(expr, params, univs, &Default::default()) + .expect("disabled cancellation checkpoint") } // ========================================================================= @@ -1626,44 +1325,11 @@ pub(super) fn mk_app_n(f: LeanExpr, args: &[LeanExpr]) -> LeanExpr { /// that shouldn't appear in the final output. pub(super) fn subst_fvar( expr: &LeanExpr, - fvar_name: &Name, - replacement: &LeanExpr, + name: &Name, + value: &LeanExpr, ) -> LeanExpr { - match expr.as_data() { - ExprData::Fvar(n, _) if n == fvar_name => replacement.clone(), - ExprData::App(f, a, _) => LeanExpr::app( - subst_fvar(f, fvar_name, replacement), - subst_fvar(a, fvar_name, replacement), - ), - ExprData::Lam(n, t, b, bi, _) => LeanExpr::lam( - n.clone(), - subst_fvar(t, fvar_name, replacement), - subst_fvar(b, fvar_name, replacement), - bi.clone(), - ), - ExprData::ForallE(n, t, b, bi, _) => LeanExpr::all( - n.clone(), - subst_fvar(t, fvar_name, replacement), - subst_fvar(b, fvar_name, replacement), - bi.clone(), - ), - ExprData::LetE(n, t, v, b, nd, _) => LeanExpr::letE( - n.clone(), - subst_fvar(t, fvar_name, replacement), - subst_fvar(v, fvar_name, replacement), - subst_fvar(b, fvar_name, replacement), - *nd, - ), - ExprData::Proj(n, i, e, _) => LeanExpr::proj( - n.clone(), - i.clone(), - subst_fvar(e, fvar_name, replacement), - ), - ExprData::Mdata(kvs, e, _) => { - LeanExpr::mdata(kvs.clone(), subst_fvar(e, fvar_name, replacement)) - }, - _ => expr.clone(), - } + super::checked_expr::subst_fvar(expr, name, value, &Default::default()) + .expect("disabled cancellation checkpoint") } /// Replace constant names throughout an expression according to a name map. diff --git a/crates/compile/src/compile/aux_gen/nested.rs b/crates/compile/src/compile/aux_gen/nested.rs index 101217a5..2051a8a9 100644 --- a/crates/compile/src/compile/aux_gen/nested.rs +++ b/crates/compile/src/compile/aux_gen/nested.rs @@ -16,14 +16,17 @@ use bignat::Nat; use blake3::Hash; -use rayon::iter::{IntoParallelRefIterator, ParallelIterator}; use rustc_hash::{FxHashMap, FxHashSet}; +use std::sync::LazyLock; +use std::time::Instant; +use super::checked_expr; use super::expr_utils::{ - LocalDecl, batch_abstract, decompose_apps, forall_telescope, - instantiate_pi_params, instantiate1, mk_forall, subst_levels, + LocalDecl, batch_abstract, decompose_apps, forall_telescope, instantiate1, + subst_levels, }; use crate::compile::nat_conv::{nat_to_u64, nat_to_usize}; +use crate::compile::validation::{self, AttemptError, Cancelled, Checkpoint}; use ix_common::env::{ ConstantInfo, Env as LeanEnv, Expr as LeanExpr, ExprData, Level, Name, }; @@ -138,6 +141,7 @@ struct ExpandCtx<'a> { block_param_decls: Vec, block_param_fvar_names: Vec, lean_env: &'a LeanEnv, + control: &'a Checkpoint, n_params: usize, } @@ -171,55 +175,62 @@ impl<'a> ExpandCtx<'a> { as_fvars: &[LeanExpr], source_owner: &Name, cache: &mut FxHashMap, - ) -> LeanExpr { + ) -> Result { + self.control.visit()?; let key = *e.get_hash(); if let Some(cached) = cache.get(&key) { - return cached.clone(); + return Ok(cached.clone()); } // Try top-level replacement first. - if let Some(replaced) = self.replace_if_nested(e, as_fvars, source_owner) { + if let Some(replaced) = self.replace_if_nested(e, as_fvars, source_owner)? { cache.insert(key, replaced.clone()); - return replaced; + return Ok(replaced); } // No match — recurse into sub-expressions. let result = match e.as_data() { ExprData::App(f, a, _) => LeanExpr::app( - self.replace_all_nested(f, as_fvars, source_owner, cache), - self.replace_all_nested(a, as_fvars, source_owner, cache), + self.replace_all_nested(f, as_fvars, source_owner, cache)?, + self.replace_all_nested(a, as_fvars, source_owner, cache)?, ), ExprData::Lam(n, t, b, bi, _) => LeanExpr::lam( n.clone(), - self.replace_all_nested(t, as_fvars, source_owner, cache), - self.replace_all_nested(b, as_fvars, source_owner, cache), + self.replace_all_nested(t, as_fvars, source_owner, cache)?, + self.replace_all_nested(b, as_fvars, source_owner, cache)?, bi.clone(), ), ExprData::ForallE(n, t, b, bi, _) => LeanExpr::all( n.clone(), - self.replace_all_nested(t, as_fvars, source_owner, cache), - self.replace_all_nested(b, as_fvars, source_owner, cache), + self.replace_all_nested(t, as_fvars, source_owner, cache)?, + self.replace_all_nested(b, as_fvars, source_owner, cache)?, bi.clone(), ), ExprData::LetE(n, t, v, b, nd, _) => LeanExpr::letE( n.clone(), - self.replace_all_nested(t, as_fvars, source_owner, cache), - self.replace_all_nested(v, as_fvars, source_owner, cache), - self.replace_all_nested(b, as_fvars, source_owner, cache), + self.replace_all_nested(t, as_fvars, source_owner, cache)?, + self.replace_all_nested(v, as_fvars, source_owner, cache)?, + self.replace_all_nested(b, as_fvars, source_owner, cache)?, *nd, ), ExprData::Proj(n, i, val, _) => LeanExpr::proj( n.clone(), i.clone(), - self.replace_all_nested(val, as_fvars, source_owner, cache), + self.replace_all_nested(val, as_fvars, source_owner, cache)?, ), ExprData::Mdata(md, inner, _) => LeanExpr::mdata( md.clone(), - self.replace_all_nested(inner, as_fvars, source_owner, cache), + self.replace_all_nested(inner, as_fvars, source_owner, cache)?, ), _ => e.clone(), }; + let result = + if result.get_hash() == e.get_hash() { e.clone() } else { result }; cache.insert(key, result.clone()); - result + self.control.scratch( + cache.capacity() + * (size_of::<(Hash, LeanExpr)>() + size_of::()), + ); + Ok(result) } /// Check if `e` is a nested inductive application and, if so, create @@ -231,41 +242,50 @@ impl<'a> ExpandCtx<'a> { e: &LeanExpr, as_fvars: &[LeanExpr], source_owner: &Name, - ) -> Option { + ) -> Result, Cancelled> { + self.control.visit()?; let (head, args) = decompose_apps(e); let (head_name, head_levels) = match head.as_data() { ExprData::Const(name, levels, _) => (name.clone(), levels.clone()), - _ => return None, + _ => return Ok(None), }; // Skip if head is in the block (direct recursive, not nested). The // `type_name_set` mirrors `self.types` names and is maintained // incrementally by `push_type`, so this is O(1) rather than O(n_types). if self.type_name_set.contains(&head_name) { - return None; + return Ok(None); } // Verify head is an external inductive. let ext_ind_ref = self.lean_env.get(&head_name); let ext_ind = match ext_ind_ref.as_deref() { Some(ConstantInfo::InductInfo(v)) => v, - _ => return None, + _ => return Ok(None), }; let ext_n_params = nat_to_usize(&ext_ind.num_params); if args.len() < ext_n_params { - return None; + return Ok(None); } // Check if any parameter arg mentions a block/flat-block member. // `expr_mentions_any_name` takes the incremental set directly so each // Const check is O(1) instead of a linear Vec scan. - if !args - .iter() - .take(ext_n_params) - .any(|a| expr_mentions_any_name(a, &self.type_name_set)) - { - return None; + let mut mentions = false; + for arg in args.iter().take(ext_n_params) { + if checked_expr::any(arg, self.control, |e, _| match e.as_data() { + ExprData::Const(n, _, _) | ExprData::Proj(n, _, _, _) => { + self.type_name_set.contains(n) + }, + _ => false, + })? { + mentions = true; + break; + } + } + if !mentions { + return Ok(None); } // Extract spec_params, normalizing constructor-local parameter FVars to @@ -275,11 +295,22 @@ impl<'a> ExpandCtx<'a> { // auxiliary identity must be expressed in the shared block-param space. let spec_params: Vec = args[..ext_n_params] .iter() - .map(|sp| replace_params_expr(sp, as_fvars, &self.block_param_fvars)) - .collect(); + .map(|sp| { + checked_expr::replace_params( + sp, + as_fvars, + &self.block_param_fvars, + self.control, + ) + }) + .collect::>()?; for sp in &spec_params { - if has_invalid_spec_ref(sp, &self.block_param_fvar_names) { - return None; + if checked_expr::any(sp, self.control, |e, depth| match e.as_data() { + ExprData::Bvar(i, _) => nat_to_u64(i) >= depth, + ExprData::Fvar(n, _) => !self.block_param_fvar_names.contains(n), + _ => false, + })? { + return Ok(None); } } @@ -303,7 +334,7 @@ impl<'a> ExpandCtx<'a> { for idx_arg in args.iter().skip(ext_n_params) { result = LeanExpr::app(result, idx_arg.clone()); } - return Some(result); + return Ok(Some(result)); } // New nested occurrence — create auxiliary types for all members of @@ -312,6 +343,7 @@ impl<'a> ExpandCtx<'a> { let mut result: Option = None; for j_name in &ext_all { + self.control.visit()?; let j_info_ref = self.lean_env.get(j_name); let j_info = match j_info_ref.as_deref() { Some(ConstantInfo::InductInfo(v)) => v, @@ -344,35 +376,58 @@ impl<'a> ExpandCtx<'a> { // 1. subst_levels(J.type, J.level_params, I_lvls) // 2. instantiate_pi_params(result, ext_n_params, spec_params) // 3. mk_forall(block_params, result) - let j_type_inst = - subst_levels(&j_info.cnst.typ, &j_info.cnst.level_params, &head_levels); - let j_type_peeled = - instantiate_pi_params(&j_type_inst, ext_n_params, &spec_params); - let j_type_block = - replace_params_expr(&j_type_peeled, as_fvars, &self.block_param_fvars); - let aux_type = mk_forall(j_type_block, &self.block_param_decls); + let j_type_inst = checked_expr::subst_levels( + &j_info.cnst.typ, + &j_info.cnst.level_params, + &head_levels, + self.control, + )?; + let j_type_peeled = checked_expr::instantiate_pi_params( + &j_type_inst, + ext_n_params, + &spec_params, + self.control, + )?; + let j_type_block = checked_expr::replace_params( + &j_type_peeled, + as_fvars, + &self.block_param_fvars, + self.control, + )?; + let aux_type = checked_expr::mk_forall( + j_type_block, + &self.block_param_decls, + self.control, + )?; // Build auxiliary constructors. let mut aux_ctors: Vec = Vec::new(); for j_ctor_name in &j_info.ctors { + self.control.visit()?; let j_ctor_ref = self.lean_env.get(j_ctor_name); let j_ctor = match j_ctor_ref.as_deref() { Some(ConstantInfo::CtorInfo(c)) => c, _ => continue, }; let aux_ctor_name = name_replace_prefix(j_ctor_name, j_name, &aux_name); - let ctor_type_inst = subst_levels( + let ctor_type_inst = checked_expr::subst_levels( &j_ctor.cnst.typ, &j_info.cnst.level_params, &head_levels, - ); - let ctor_type_peeled = - instantiate_pi_params(&ctor_type_inst, ext_n_params, &spec_params); - let ctor_type_block = replace_params_expr( + self.control, + )?; + let ctor_type_peeled = checked_expr::instantiate_pi_params( + &ctor_type_inst, + ext_n_params, + &spec_params, + self.control, + )?; + let ctor_type_block = checked_expr::replace_params( &ctor_type_peeled, as_fvars, &self.block_param_fvars, - ); + self.control, + )?; let ctor_type_block = replace_ctor_result_head_with_aux( &ctor_type_block, j_name, @@ -380,8 +435,13 @@ impl<'a> ExpandCtx<'a> { ext_n_params, &self.block_levels, &self.block_param_fvars, - ); - let aux_ctor_type = mk_forall(ctor_type_block, &self.block_param_decls); + self.control, + )?; + let aux_ctor_type = checked_expr::mk_forall( + ctor_type_block, + &self.block_param_decls, + self.control, + )?; self.aux_ctor_map.insert( aux_ctor_name.clone(), @@ -416,7 +476,7 @@ impl<'a> ExpandCtx<'a> { }); } - result + Ok(result) } } @@ -430,6 +490,22 @@ pub fn expand_nested_block( lean_env: &LeanEnv, alias_to_rep: &FxHashMap, ) -> Result { + expand_nested_block_checked( + ordered_originals, + lean_env, + alias_to_rep, + &Checkpoint::default(), + ) + .map_err(AttemptError::into_compile) +} + +fn expand_nested_block_checked( + ordered_originals: &[Name], + lean_env: &LeanEnv, + alias_to_rep: &FxHashMap, + control: &Checkpoint, +) -> Result { + control.visit()?; let first_name = ordered_originals.first().ok_or_else(|| { CompileError::InvalidMutualBlock { reason: "expand_nested_block: empty ordered_originals".into(), @@ -439,10 +515,13 @@ pub fn expand_nested_block( let first_ind = match first_ind_ref.as_deref() { Some(ConstantInfo::InductInfo(v)) => v, _ => { - return Err(CompileError::MissingConstant { - name: first_name.pretty(), - caller: "expand_nested_block: first original not an inductive".into(), - }); + return Err( + CompileError::MissingConstant { + name: first_name.pretty(), + caller: "expand_nested_block: first original not an inductive".into(), + } + .into(), + ); }, }; @@ -452,7 +531,13 @@ pub fn expand_nested_block( level_params.iter().map(|lp| Level::param(lp.clone())).collect(); let (block_param_fvars, block_param_decls, _) = - forall_telescope(&first_ind.cnst.typ, n_params, "bp", 0); + checked_expr::forall_telescope( + &first_ind.cnst.typ, + n_params, + "bp", + 0, + control, + )?; let block_param_fvar_names: Vec = block_param_decls.iter().map(|d| d.fvar_name.clone()).collect(); @@ -475,19 +560,24 @@ pub fn expand_nested_block( block_param_decls: block_param_decls.clone(), block_param_fvar_names, lean_env, + control, n_params, }; // Seed with original inductives. for name in ordered_originals { + control.visit()?; let ind_ref = lean_env.get(name); let ind = match ind_ref.as_deref() { Some(ConstantInfo::InductInfo(v)) => v, _ => { - return Err(CompileError::MissingConstant { - name: name.pretty(), - caller: "expand_nested_block: original not an inductive".into(), - }); + return Err( + CompileError::MissingConstant { + name: name.pretty(), + caller: "expand_nested_block: original not an inductive".into(), + } + .into(), + ); }, }; let ctors: Vec = ind @@ -543,14 +633,20 @@ pub fn expand_nested_block( // the cache turns DAG traversal from O(shared × nodes) into O(nodes). let mut qi = 0; while qi < ctx.types.len() { + control.visit()?; let n_ctors = ctx.types[qi].ctors.len(); let source_owner = ctx.types[qi].source_owner.clone(); for ci in 0..n_ctors { let ctor_type = ctx.types[qi].ctors[ci].typ.clone(); // Peel params, re-creating FVars per constructor for binding info. - let (as_fvars, as_decls, peeled) = - forall_telescope(&ctor_type, n_params, "cp", qi * 100 + ci); + let (as_fvars, as_decls, peeled) = checked_expr::forall_telescope( + &ctor_type, + n_params, + "cp", + qi * 100 + ci, + control, + )?; // Replace all nested occurrences in the peeled body. let mut walk_cache: FxHashMap = FxHashMap::default(); @@ -559,10 +655,11 @@ pub fn expand_nested_block( &as_fvars, &source_owner, &mut walk_cache, - ); + )?; // Re-wrap with constructor-local params. - let new_ctor_type = mk_forall(replaced, &as_decls); + let new_ctor_type = + checked_expr::mk_forall(replaced, &as_decls, control)?; ctx.types[qi].ctors[ci].typ = new_ctor_type; } qi += 1; @@ -1937,8 +2034,10 @@ fn replace_ctor_result_head_with_aux( original_n_params: usize, block_levels: &[Level], block_param_fvars: &[LeanExpr], -) -> LeanExpr { - match e.as_data() { + control: &Checkpoint, +) -> Result { + control.visit()?; + Ok(match e.as_data() { ExprData::ForallE(n, t, b, bi, _) => LeanExpr::all( n.clone(), t.clone(), @@ -1949,7 +2048,8 @@ fn replace_ctor_result_head_with_aux( original_n_params, block_levels, block_param_fvars, - ), + control, + )?, bi.clone(), ), ExprData::Mdata(md, inner, _) => LeanExpr::mdata( @@ -1961,15 +2061,16 @@ fn replace_ctor_result_head_with_aux( original_n_params, block_levels, block_param_fvars, - ), + control, + )?, ), _ => { let (head, args) = decompose_apps(e); let ExprData::Const(head_name, _, _) = head.as_data() else { - return e.clone(); + return Ok(e.clone()); }; if head_name != original_ind || args.len() < original_n_params { - return e.clone(); + return Ok(e.clone()); } let mut result = LeanExpr::cnst(aux_name.clone(), block_levels.to_vec()); @@ -1981,7 +2082,7 @@ fn replace_ctor_result_head_with_aux( } result }, - } + }) } // ========================================================================= @@ -2464,7 +2565,17 @@ pub fn compute_lean_ind_flags( all: &[Name], lean_env: &LeanEnv, ) -> Result { - let expanded = expand_nested_block(all, lean_env, &FxHashMap::default())?; + compute_lean_ind_flags_checked(all, lean_env, &Checkpoint::default()) + .map_err(AttemptError::into_compile) +} + +fn compute_lean_ind_flags_checked( + all: &[Name], + lean_env: &LeanEnv, + control: &Checkpoint, +) -> Result { + let expanded = + expand_nested_block_checked(all, lean_env, &FxHashMap::default(), control)?; let num_nested = (expanded.types.len() - expanded.n_originals) as u64; let block_names: FxHashSet<&Name> = expanded.types.iter().map(|m| &m.name).collect(); @@ -2476,7 +2587,8 @@ pub fn compute_lean_ind_flags( for ctor in &member.ctors { let mut ty = &ctor.typ; while let ExprData::ForallE(_, dom, body, _, _) = ty.as_data() { - if has_ind_occ(dom, &block_names, &mut occ_cache) { + control.visit()?; + if has_ind_occ(dom, &block_names, &mut occ_cache, control)? { is_rec = true; if matches!(dom.as_data(), ExprData::ForallE(..)) { is_reflexive = true; @@ -2503,38 +2615,75 @@ pub fn validate_lean_ind_flags(lean_env: &LeanEnv) -> Result<(), CompileError> { validate_ind_groups(&groups, lean_env) } +static LOG_IND_GROUPS: LazyLock = + LazyLock::new(|| std::env::var_os("IX_LOG_IND_GROUPS").is_some()); + /// Per-group half of [`validate_lean_ind_flags`], for callers that /// already hold the inductive groups from a wider env pass. +/// `IX_LOG_IND_GROUPS=1` logs group entry/exit for allocation-failure diagnosis. pub fn validate_ind_groups( groups: &FxHashMap>, lean_env: &LeanEnv, ) -> Result<(), CompileError> { - groups.par_iter().try_for_each(|(_, all)| { - for member in all.iter() { - let entry = lean_env.get(member); - let Some(ConstantInfo::InductInfo(v)) = entry.as_deref() else { - return Ok(()); - }; - for cn in &v.ctors { - if !matches!( - lean_env.get(cn).as_deref(), - Some(ConstantInfo::CtorInfo(_)) - ) { - return Ok(()); - } - } + // Stable IDs make cancellations/retries identifiable in progress logs. + let mut entries: Vec<_> = groups.iter().collect(); + entries.sort_unstable_by(|(a, _), (b, _)| { + a.get_hash().as_bytes().cmp(b.get_hash().as_bytes()) + }); + let entries: Vec<_> = entries.into_iter().enumerate().collect(); + validation::run(&entries, |(id, (leader, all)), control| { + let start = Instant::now(); + if *LOG_IND_GROUPS { + eprintln!("[validate_ind_flags] BEGIN {} (group #{id})", leader.pretty()); } - let flags = compute_lean_ind_flags(all, lean_env)?; - for member in all.iter() { - let entry = lean_env.get(member); - let Some(ConstantInfo::InductInfo(v)) = entry.as_deref() else { - continue; // unreachable + let result = validate_one_group(all, lean_env, control); + if *LOG_IND_GROUPS { + let status = match &result { + Ok(()) => "END", + Err(AttemptError::Cancelled) => "CANCEL", + Err(AttemptError::Compile(_)) => "ERROR", }; - if v.is_rec != flags.is_rec - || v.is_reflexive != flags.is_reflexive - || v.num_nested != Nat::from(flags.num_nested) + eprintln!( + "[validate_ind_flags] {status} {} (group #{id}, {:.2}s)", + leader.pretty(), + start.elapsed().as_secs_f32() + ); + } + result + }) +} + +fn validate_one_group( + all: &[Name], + lean_env: &LeanEnv, + control: &Checkpoint, +) -> Result<(), AttemptError> { + for member in all.iter() { + control.visit()?; + let entry = lean_env.get(member); + let Some(ConstantInfo::InductInfo(v)) = entry.as_deref() else { + return Ok(()); + }; + for cn in &v.ctors { + control.visit()?; + if !matches!(lean_env.get(cn).as_deref(), Some(ConstantInfo::CtorInfo(_))) { - return Err(CompileError::InvalidMutualBlock { + return Ok(()); + } + } + } + let flags = compute_lean_ind_flags_checked(all, lean_env, control)?; + for member in all.iter() { + let entry = lean_env.get(member); + let Some(ConstantInfo::InductInfo(v)) = entry.as_deref() else { + continue; // unreachable + }; + if v.is_rec != flags.is_rec + || v.is_reflexive != flags.is_reflexive + || v.num_nested != Nat::from(flags.num_nested) + { + return Err( + CompileError::InvalidMutualBlock { reason: format!( "non-canonical inductive flags for '{}': \ stored isRec={}isReflexive={} numNested={}, \ @@ -2547,11 +2696,12 @@ pub fn validate_ind_groups( flags.is_reflexive, flags.num_nested, ), - }); - } + } + .into(), + ); } - Ok(()) - }) + } + Ok(()) } // does `expr` mention any block type by name anywhere? @@ -2560,31 +2710,36 @@ fn has_ind_occ( expr: &LeanExpr, names: &FxHashSet<&Name>, cache: &mut FxHashMap, -) -> bool { + control: &Checkpoint, +) -> Result { + control.visit()?; let key = *expr.get_hash(); if let Some(&cached) = cache.get(&key) { - return cached; + return Ok(cached); } let result = match expr.as_data() { ExprData::Const(name, _, _) => names.contains(name), ExprData::App(f, a, _) => { - has_ind_occ(f, names, cache) || has_ind_occ(a, names, cache) + has_ind_occ(f, names, cache, control)? + || has_ind_occ(a, names, cache, control)? }, ExprData::Lam(_, t, b, _, _) | ExprData::ForallE(_, t, b, _, _) => { - has_ind_occ(t, names, cache) || has_ind_occ(b, names, cache) + has_ind_occ(t, names, cache, control)? + || has_ind_occ(b, names, cache, control)? }, ExprData::LetE(_, t, v, b, _, _) => { - has_ind_occ(t, names, cache) - || has_ind_occ(v, names, cache) - || has_ind_occ(b, names, cache) + has_ind_occ(t, names, cache, control)? + || has_ind_occ(v, names, cache, control)? + || has_ind_occ(b, names, cache, control)? }, ExprData::Proj(_, _, e, _) | ExprData::Mdata(_, e, _) => { - has_ind_occ(e, names, cache) + has_ind_occ(e, names, cache, control)? }, _ => false, }; cache.insert(key, result); - result + control.scratch(cache.capacity() * size_of::<(Hash, bool)>()); + Ok(result) } #[cfg(test)] diff --git a/crates/compile/src/compile/env.rs b/crates/compile/src/compile/env.rs index b30545b8..f0f7dd20 100644 --- a/crates/compile/src/compile/env.rs +++ b/crates/compile/src/compile/env.rs @@ -18,8 +18,9 @@ //! `.ixe`. //! //! Three knobs: -//! - `IX_COMPILE_WORKERS=N` — scheduler worker count (default: all -//! cores). Scales the per-worker transients. +//! - `IX_COMPILE_WORKERS=N` — scheduler worker ceiling (default: all +//! cores). Adaptive admission controls the active count under memory +//! pressure; set `IX_COMPILE_ADAPTIVE=0` for fixed parallelism. //! - `IX_COMPILE_DEMOTE=0` — keep materialized caches next to the //! accumulator's bytes instead of demoting to bytes-only. Spends //! RAM to make post-compile structural reads free, which only pays @@ -43,6 +44,7 @@ use dashmap::DashMap; use rayon::prelude::*; use rustc_hash::FxHashSet; +use super::admission::Admission; use crate::compile::{ BlockCache, CompileOptions, CompileState, aux_gen::nested::validate_ind_groups, compile_const, compile_const_no_aux, @@ -62,7 +64,7 @@ use ixon::CompileError; /// per-phase timings and scheduler progress. `IX_COMPILE_DBG=1` (the /// Lean driver's phase-attribution knob) is accepted as an alias so one /// flag lights up both sides of the pipeline. -static IX_VERBOSE: LazyLock = LazyLock::new(|| { +pub(super) static IX_VERBOSE: LazyLock = LazyLock::new(|| { std::env::var("IX_VERBOSE").is_ok() || std::env::var("IX_COMPILE_DBG").is_ok() }); @@ -144,6 +146,7 @@ pub fn compile_env_with_options( lean_env: &Arc, options: CompileOptions, ) -> Result { + let _memory_sampler = crate::diag::memory_sampler("compile_env"); let setup_start = Instant::now(); // Whole-env scan: ref graph + immediate groundedness + inductive // groups in one decode per constant — the env decodes lazily, so @@ -240,11 +243,20 @@ pub fn compile_env_with_options( // flags. Groups come from the fused scan; only inductive families are // re-read here. let phase_start = Instant::now(); + if *IX_VERBOSE { + eprintln!( + "[compile_env] setup 4/7 validate_ind_flags BEGIN ({} groups, {} Rayon threads){}", + scan.ind_groups.len(), + rayon::current_num_threads(), + crate::diag::rss_log_suffix(), + ); + } validate_ind_groups(&scan.ind_groups, lean_env.as_ref())?; if *IX_VERBOSE { eprintln!( - "[compile_env] setup 4/7 validate_ind_flags: {:.2}s", - phase_start.elapsed().as_secs_f32() + "[compile_env] setup 4/7 validate_ind_flags: {:.2}s{}", + phase_start.elapsed().as_secs_f32(), + crate::diag::rss_log_suffix(), ); } @@ -408,6 +420,7 @@ pub fn compile_env_with_options( .unwrap_or(available_threads) .min(available_threads) .max(1); + let admission = Admission::new(num_threads)?; // Progress tracking. `active` holds currently-compiling blocks per worker // so the reporter thread can show blocks that are still in-flight (useful @@ -435,6 +448,7 @@ pub fn compile_env_with_options( let condvar_ref = &work_available; let active_ref = &active; let stop_progress_ref = &stop_progress; + let admission_ref = &admission; thread::scope(|s| { // Periodic progress reporter. Wakes every IX_PROGRESS_MS to print @@ -459,6 +473,7 @@ pub fn compile_env_with_options( s.spawn(move || { let mut last_completed = 0usize; let mut last_print = Instant::now(); + let mut last_progress = Instant::now(); while !stop_p.load(AtomicOrdering::Relaxed) { thread::sleep(check_interval); if stop_p.load(AtomicOrdering::Relaxed) { @@ -469,11 +484,14 @@ pub fn compile_env_with_options( if last_print.elapsed() < interval { continue; } + let print_interval = last_print.elapsed().as_secs_f64(); last_print = Instant::now(); let done = completed_p.load(AtomicOrdering::SeqCst); - // Skip if no change and we're not in the first tick — reduces - // noise when the scheduler is blocked on a single slow block. + // No new completion is not a deadlock signal: report how long + // current blocks have been running, with their active count. let changed = done != last_completed; + let delta = done.saturating_sub(last_completed); + if changed { last_progress = Instant::now(); } last_completed = done; let pct = if total == 0 { 100.0 @@ -481,28 +499,21 @@ pub fn compile_env_with_options( (done as f64 / total as f64) * 100.0 }; let elapsed = start.elapsed().as_secs_f64(); - let rate = - if elapsed > 0.0 { done as f64 / elapsed } else { 0.0 }; - let eta = if rate > 0.0 && done < total { - let remaining = (total - done) as f64 / rate; - format!(" eta {:.0}s", remaining) - } else { - String::new() - }; // Oldest in-flight blocks (up to 3) for visibility into // slow/stuck compilations. Sort by start time ascending. - let in_flight: Vec = { + let (active_count, in_flight): (usize, Vec) = { let mut entries: Vec<(Name, Instant)> = active_p.lock().unwrap().clone(); entries.sort_by_key(|(_, t)| *t); - entries + let names = entries .iter() .take(3) .map(|(n, t)| { format!("{} ({:.0}s)", n.pretty(), t.elapsed().as_secs_f64()) }) - .collect() + .collect(); + (entries.len(), names) }; let suffix = if in_flight.is_empty() { String::new() @@ -510,16 +521,17 @@ pub fn compile_env_with_options( format!(" · in-flight: {}", in_flight.join(", ")) }; - // Always print the first tick and any tick with progress; - // print "stalled" ticks less often so the log doesn't churn. + // A lifetime-average ETA is misleading when a fast prefix is + // followed by a handful of expensive blocks. Show recent work. if changed || done == 0 { eprintln!( - "[compile_env] {done}/{total} ({pct:.1}%) · {elapsed:.0}s{eta}{suffix}{}", + "[compile_env] {done}/{total} ({pct:.1}%) · {elapsed:.0}s · +{delta} blocks/{print_interval:.1}s · active={active_count}{suffix}{}", crate::diag::rss_log_suffix(), ); } else { eprintln!( - "[compile_env] {done}/{total} ({pct:.1}%) · STALLED{suffix}" + "[compile_env] {done}/{total} ({pct:.1}%) · waiting: no completed blocks for {:.0}s · active={active_count}{suffix}", + last_progress.elapsed().as_secs_f64(), ); } } @@ -542,6 +554,17 @@ pub fn compile_env_with_options( for _ in 0..num_threads { s.spawn(move || { loop { + // Wait without holding a ready-queue lock or any block scratch. + // A permit covers compilation, publication, and dependency release. + let permit = match admission_ref.acquire() { + Ok(Some(permit)) => permit, + Ok(None) => return, + Err(failure) => { + error_ref.lock().unwrap().get_or_insert(failure); + condvar_ref.notify_all(); + return; + }, + }; // Try to get work from the ready queue let work = { let mut queue = ready_queue_ref.lock().unwrap(); @@ -899,9 +922,11 @@ pub fn compile_env_with_options( } else { condvar_ref.notify_one(); } + permit.complete(); }, None => { + drop(permit); // No work available - check if we're done if completed_ref.load(AtomicOrdering::SeqCst) == total_blocks { return; @@ -937,10 +962,16 @@ pub fn compile_env_with_options( while completed_ref.load(AtomicOrdering::SeqCst) < total_blocks && error_ref.lock().unwrap().is_none() { + if let Err(failure) = admission_ref.check() { + error_ref.lock().unwrap().get_or_insert(failure); + break; + } thread::sleep(Duration::from_millis(25)); } + admission_ref.stop(); stop_progress_ref.store(true, AtomicOrdering::Relaxed); }); + admission.check()?; if *IX_VERBOSE { let scheduler_elapsed = compile_start.elapsed().as_secs_f64(); diff --git a/crates/compile/src/compile/memory.rs b/crates/compile/src/compile/memory.rs new file mode 100644 index 00000000..3152aa1d --- /dev/null +++ b/crates/compile/src/compile/memory.rs @@ -0,0 +1,216 @@ +//! Linux memory telemetry and pressure policy shared by both compilation stages. + +use ixon::CompileError; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +pub(super) const MIB: u64 = 1024 * 1024; +pub(super) const GIB: u64 = 1024 * MIB; + +pub(super) fn budget_from_env(key: &str) -> Result, CompileError> { + std::env::var(key).ok().map(|value| parse_budget(key, &value)).transpose() +} + +fn parse_budget(key: &str, value: &str) -> Result { + value + .parse::() + .ok() + .filter(|v| v.is_finite() && *v > 0.0 && *v < (u64::MAX / GIB) as f64) + .map(|v| { + // The filter proves this is finite and positive. + #[allow(clippy::cast_sign_loss)] + let bytes = (v * GIB as f64) as u64; + bytes + }) + .filter(|&bytes| bytes > 0) + .ok_or_else(|| { + resource_error(format!("{key} must be a positive finite number")) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn memory_budget_rejects_invalid_and_unrepresentable_values() { + for value in ["", "0", "-1", "nan", "inf", "1e50", "1e-100", "no"] { + assert!(matches!( + parse_budget("TEST", value), + Err(CompileError::ResourceLimit { .. }) + )); + } + assert_eq!(parse_budget("TEST", "0.5").unwrap(), GIB / 2); + assert_eq!(parse_budget("TEST", "100").unwrap(), 100 * GIB); + } +} + +#[derive(Clone, Copy, Debug)] +pub(super) struct Memory { + pub(super) capacity: u64, + pub(super) available: u64, + pub(super) process: u64, + pub(super) swap_used: u64, + pub(super) stall_us: u64, +} + +pub(super) fn field(text: &str, key: &str) -> Option { + text.lines().find_map(|line| { + line.strip_prefix(key)?.split_whitespace().next()?.parse().ok() + }) +} +pub(super) fn psi_total(text: &str) -> u64 { + text + .lines() + .find(|line| line.starts_with("full ")) + .and_then(|line| { + line + .split_whitespace() + .find_map(|v| v.strip_prefix("total=")?.parse().ok()) + }) + .unwrap_or(0) +} +fn proc_path(text: &str) -> PathBuf { + PathBuf::from( + text + .replace("\\040", " ") + .replace("\\011", "\t") + .replace("\\012", "\n") + .replace("\\134", "\\"), + ) +} + +/// Resolve the unified cgroup relative to its visible mount, including +/// namespace roots and every visible ancestor's memory.max/memory.high. +pub(super) fn cgroup_dirs(cgroups: &str, mounts: &str) -> Vec { + let Some(group) = cgroups.lines().find_map(|line| line.strip_prefix("0::")) + else { + return Vec::new(); + }; + let group = proc_path(group); + if group.components().any(|c| matches!(c, std::path::Component::ParentDir)) { + return Vec::new(); + } + for line in mounts.lines() { + let Some((left, right)) = line.split_once(" - ") else { + continue; + }; + if right.split_whitespace().next() != Some("cgroup2") { + continue; + } + let fields: Vec<_> = left.split_whitespace().collect(); + if fields.len() < 5 { + continue; + } + let root = proc_path(fields[3]); + let mount = proc_path(fields[4]); + let relative = if group == Path::new("/") { + Path::new("") + } else if let Ok(relative) = group.strip_prefix(&root) { + relative + } else { + continue; + }; + let leaf = mount.join(relative); + return leaf + .ancestors() + .take_while(|path| path.starts_with(&mount)) + .map(Path::to_path_buf) + .collect(); + } + Vec::new() +} + +pub(super) struct MemoryReader { + cgroups: Vec, +} +impl MemoryReader { + pub(super) fn new() -> Self { + Self { + cgroups: cgroup_dirs( + &std::fs::read_to_string("/proc/self/cgroup").unwrap_or_default(), + &std::fs::read_to_string("/proc/self/mountinfo").unwrap_or_default(), + ), + } + } + pub(super) fn read(&mut self) -> Option { + let mem = std::fs::read_to_string("/proc/meminfo").ok()?; + let status = std::fs::read_to_string("/proc/self/status").ok()?; + let mut capacity = field(&mem, "MemTotal:")?.saturating_mul(1024); + let mut available = field(&mem, "MemAvailable:")?.saturating_mul(1024); + for dir in &self.cgroups { + let current = std::fs::read_to_string(dir.join("memory.current")) + .ok() + .and_then(|v| v.trim().parse::().ok()); + for file in ["memory.max", "memory.high"] { + if let (Some(current), Some(limit)) = ( + current, + std::fs::read_to_string(dir.join(file)) + .ok() + .and_then(|v| v.trim().parse::().ok()), + ) { + capacity = capacity.min(limit); + available = available.min(limit.saturating_sub(current)); + } + } + } + Some(Memory { + capacity, + available, + process: field(&status, "VmRSS:")? + .saturating_add(field(&status, "VmSwap:").unwrap_or(0)) + .saturating_mul(1024), + swap_used: field(&mem, "SwapTotal:")? + .saturating_sub(field(&mem, "SwapFree:")?) + .saturating_mul(1024), + stall_us: psi_total( + &std::fs::read_to_string("/proc/pressure/memory").unwrap_or_default(), + ), + }) + } +} + +pub(super) fn resource_error(reason: impl Into) -> CompileError { + CompileError::ResourceLimit { reason: reason.into() } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum Pressure { + Healthy, + Hold, + Backoff, + Critical, +} + +pub(super) fn reclaiming( + now: Memory, + before: Memory, + elapsed: Duration, +) -> bool { + let stalls = now.stall_us.saturating_sub(before.stall_us); + stalls > (elapsed.as_micros() as u64 / 10).max(1) + || now.swap_used.saturating_sub(before.swap_used) > 8 * MIB +} + +pub(super) fn pressure( + now: Memory, + before: Memory, + elapsed: Duration, + budget: Option, +) -> Pressure { + let capacity = budget.map_or(now.capacity, |b| b.min(now.capacity)); + let reserve = (capacity / 8).max(64 * MIB).min(capacity / 2); + let available = budget.map_or(now.available, |b| { + now.available.min(b.saturating_sub(now.process)) + }); + if available < reserve / 2 { + return Pressure::Critical; + } + if available < reserve * 2 || reclaiming(now, before, elapsed) { + Pressure::Backoff + } else if available < reserve * 3 { + Pressure::Hold + } else { + Pressure::Healthy + } +} diff --git a/crates/compile/src/compile/validation.rs b/crates/compile/src/compile/validation.rs new file mode 100644 index 00000000..7ec1b2b3 --- /dev/null +++ b/crates/compile/src/compile/validation.rs @@ -0,0 +1,845 @@ +//! Memory-aware admission and cooperative retry for inductive validation. +//! +//! The coordinator never occupies a Rayon worker. Cancelled attempts return +//! normally (the host uses panic=abort), dropping scratch data before retry. +//! This is a soft budget, not an allocator limit: checkpoints cannot interrupt +//! an allocation, a lazy-environment fetch, or destruction already in progress. + +use std::cell::Cell; +use std::collections::VecDeque; +use std::sync::{ + Arc, + atomic::{AtomicBool, AtomicU64, Ordering}, + mpsc, +}; +use std::time::{Duration, Instant}; + +use ixon::CompileError; + +use super::memory::{ + GIB, MIB, Memory, MemoryReader, Pressure, pressure, resource_error, +}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct Cancelled; + +#[derive(Debug)] +pub(crate) enum AttemptError { + Cancelled, + Compile(CompileError), +} + +impl From for AttemptError { + fn from(_: Cancelled) -> Self { + Self::Cancelled + } +} +impl From for AttemptError { + fn from(error: CompileError) -> Self { + Self::Compile(error) + } +} +impl AttemptError { + pub(crate) fn into_compile(self) -> CompileError { + match self { + Self::Compile(error) => error, + Self::Cancelled => resource_error("inductive validation interrupted"), + } + } +} + +#[derive(Default)] +struct Signals { + cancel: AtomicBool, + visits: AtomicU64, + scratch_peak: AtomicU64, +} + +/// Worker-local checkpoint state. Scratch telemetry estimates the largest +/// observed walk cache plus node payload, not total retained scratch or +/// per-job RSS (which cannot be attributed by subtracting process RSS +/// samples while multiple jobs execute). Admission uses OS telemetry. +#[derive(Default)] +pub(crate) struct Checkpoint { + signals: Option>, + visits: Cell, +} + +impl Checkpoint { + pub(crate) fn visit(&self) -> Result<(), Cancelled> { + if let Some(signals) = &self.signals { + if signals.cancel.load(Ordering::Relaxed) { + return Err(Cancelled); + } + let n = self.visits.get().saturating_add(1); + self.visits.set(n); + if n.is_multiple_of(1024) { + signals.visits.store(n, Ordering::Relaxed); + } + } + Ok(()) + } + + pub(crate) fn scratch(&self, estimate: usize) { + if let Some(signals) = &self.signals { + signals.scratch_peak.fetch_max(estimate as u64, Ordering::Relaxed); + } + } +} +impl Drop for Checkpoint { + fn drop(&mut self) { + if let Some(signals) = &self.signals { + signals.visits.store(self.visits.get(), Ordering::Relaxed); + } + } +} + +#[derive(Clone, Copy, Debug)] +struct Options { + max_workers: usize, + initial_workers: usize, + process_budget: Option, + tick: Duration, + ramp_interval: Duration, + recovery: Duration, + idle_timeout: Duration, + verbose: bool, +} +impl Options { + fn from_env() -> Result { + let budget = super::memory::budget_from_env("IX_VALIDATE_MEMORY_GIB")?; + let max_workers = rayon::current_num_threads().max(1); + Ok(Self { + max_workers, + initial_workers: max_workers.min(2), + process_budget: budget, + tick: Duration::from_millis(250), + ramp_interval: Duration::from_secs(1), + recovery: Duration::from_secs(5), + idle_timeout: Duration::from_secs(30), + verbose: *super::env::IX_VERBOSE, + }) + } +} + +struct Active { + id: usize, + solo: bool, + signals: Arc, +} + +// Production builds abort on panic. In unwind-enabled test/library builds, +// still notify the coordinator if a task exits abnormally, so the enclosing +// Rayon scope can propagate the panic instead of waiting for a lost result. +struct Completion { + id: usize, + sender: Option)>>, +} +impl Completion { + fn finish(mut self, result: Result<(), AttemptError>) { + let _ = self.sender.take().unwrap().send((self.id, result)); + } +} +impl Drop for Completion { + fn drop(&mut self) { + if let Some(sender) = self.sender.take() { + let _ = sender.send(( + self.id, + Err(resource_error("validation task exited without a result").into()), + )); + } + } +} + +/// Cancel the newest attempts, preserving the oldest work. Cancelled work +/// stays active until its result arrives *after* scratch has been dropped. +fn cancel_to(active: &[Active], keep: usize) { + for job in active + .iter() + .filter(|j| !j.signals.cancel.load(Ordering::Relaxed)) + .skip(keep) + { + job.signals.cancel.store(true, Ordering::Relaxed); + } +} + +fn coordinate<'scope, T: Sync, F>( + scope: &rayon::Scope<'scope>, + jobs: &'scope [T], + work: &'scope F, + options: &Options, + mut read: impl FnMut() -> Option, +) -> Result<(), CompileError> +where + F: Fn(&T, &Checkpoint) -> Result<(), AttemptError> + Sync, +{ + let (sender, receiver) = mpsc::channel(); + let mut ready: VecDeque = (0..jobs.len()).collect(); + let mut deferred = VecDeque::new(); + let mut active: Vec = Vec::new(); + let mut error = None; + let mut limit = options.initial_workers.max(1).min(options.max_workers); + let mut previous = + read().ok_or_else(|| resource_error("memory telemetry unavailable"))?; + let mut sampled = Instant::now(); + let mut ramped = sampled; + let mut pressured_at = None; + let mut idle_since = None; + let mut logged = sampled; + let mut completed = 0; + let mut retries = 0; + // Apply limits before the first allocation-heavy job, not only on tick 1. + let mut state = + pressure(previous, previous, options.tick, options.process_budget); + + loop { + let now = Instant::now(); + if now.duration_since(sampled) >= options.tick { + let previous_state = state; + if let Some(memory) = read() { + state = pressure( + memory, + previous, + now.duration_since(sampled), + options.process_budget, + ); + if matches!(state, Pressure::Backoff | Pressure::Critical) { + pressured_at = Some(now); + let running = active + .iter() + .filter(|j| !j.signals.cancel.load(Ordering::Relaxed)) + .count(); + let next = + if state == Pressure::Critical { 1 } else { (running / 2).max(1) }; + limit = limit.min(next); + cancel_to(&active, limit); + // Don't mistake memory still being released by other attempts for + // the survivor's footprint. Only fail a genuinely isolated attempt. + if state == Pressure::Critical && active.len() == 1 && error.is_none() + { + active[0].signals.cancel.store(true, Ordering::Relaxed); + error = Some(resource_error(format!( + "inductive validation group #{} reached the memory safety reserve while running alone (resident+swap {:.1} GiB); increase the budget or reduce this group's working set", + active[0].id, + memory.process as f64 / GIB as f64 + ))); + } + } else if state == Pressure::Healthy + && pressured_at + .is_none_or(|t| now.duration_since(t) >= options.recovery) + && now.duration_since(ramped) >= options.ramp_interval + { + limit = (limit + 1).min(options.max_workers); + ramped = now; + } + if options.verbose + && (state != previous_state + || now.duration_since(logged) >= Duration::from_secs(5)) + { + eprintln!( + "[validate_memory] {state:?}: limit={limit}/{} active={} ready={} deferred={} completed={completed}/{} retries={retries} available={:.1} GiB resident+swap={:.1} GiB", + options.max_workers, + active.len(), + ready.len(), + deferred.len(), + jobs.len(), + memory.available as f64 / GIB as f64, + memory.process as f64 / GIB as f64 + ); + for job in &active { + let visits = job.signals.visits.load(Ordering::Relaxed); + if visits >= 1_000_000 { + eprintln!( + "[validate_memory] active group #{}: visits={visits} scratch_estimate_peak={:.1} MiB cancelling={}", + job.id, + job.signals.scratch_peak.load(Ordering::Relaxed) as f64 + / MIB as f64, + job.signals.cancel.load(Ordering::Relaxed) + ); + } + } + logged = now; + } + previous = memory; + } else if error.is_none() { + error = Some(resource_error( + "memory telemetry became unavailable during validation", + )); + } + sampled = now; + } + + if error.is_some() { + cancel_to(&active, 0); + } + let recovering = + pressured_at.is_some_and(|t| now.duration_since(t) < options.recovery); + let draining = + active.iter().any(|j| j.signals.cancel.load(Ordering::Relaxed)); + let can_admit = + error.is_none() && !recovering && !draining && state == Pressure::Healthy; + if can_admit && !active.iter().any(|j| j.solo) { + while active.len() < limit { + let next = if !ready.is_empty() { + ready.pop_front().map(|id| (id, false)) + } else if active.is_empty() { + deferred.pop_front().map(|id| (id, true)) + } else { + None + }; + let Some((id, solo)) = next else { + break; + }; + let signals = Arc::new(Signals::default()); + let token = Arc::clone(&signals); + let result_sender = sender.clone(); + scope.spawn(move |_| { + let completion = Completion { id, sender: Some(result_sender) }; + let checkpoint = + Checkpoint { signals: Some(token), visits: Cell::new(0) }; + let result = work(&jobs[id], &checkpoint); + drop(checkpoint); + completion.finish(result); + }); + active.push(Active { id, solo, signals }); + if solo { + break; + } + } + } + + if active.is_empty() { + if let Some(error) = error { + return Err(error); + } + if ready.is_empty() && deferred.is_empty() { + return Ok(()); + } + let start = *idle_since.get_or_insert(now); + if now.duration_since(start) >= options.idle_timeout { + return Err(resource_error(format!( + "not enough available memory to admit an inductive validation job; no progress for {:.1} seconds", + options.idle_timeout.as_secs_f64() + ))); + } + } else { + idle_since = None; + } + + if let Ok((id, result)) = + receiver.recv_timeout(options.tick.min(Duration::from_millis(50))) + { + let position = active + .iter() + .position(|j| j.id == id) + .expect("result for active validation job"); + let job = active.remove(position); + if options.verbose + && (job.solo || job.signals.visits.load(Ordering::Relaxed) >= 1_000_000) + { + eprintln!( + "[validate_memory] group #{id}: visits={} scratch_estimate_peak={:.1} MiB solo={}", + job.signals.visits.load(Ordering::Relaxed), + job.signals.scratch_peak.load(Ordering::Relaxed) as f64 / MIB as f64, + job.solo + ); + } + match result { + Ok(()) => completed += 1, + Err(AttemptError::Compile(failure)) => { + if error.is_none() { + error = Some(failure); + } + }, + Err(AttemptError::Cancelled) => { + if error.is_none() { + if job.solo { + error = Some(resource_error(format!( + "inductive validation group #{id} was cancelled while already running alone" + ))); + } else { + retries += 1; + deferred.push_back(id); + } + } + }, + } + } + } +} + +fn run_with( + jobs: &[T], + work: &F, + options: Options, + read: R, +) -> Result<(), CompileError> +where + F: Fn(&T, &Checkpoint) -> Result<(), AttemptError> + Sync, + R: FnMut() -> Option + Send, +{ + rayon::in_place_scope(|scope| { + std::thread::scope(|threads| { + let controller = + threads.spawn(|| coordinate(scope, jobs, work, &options, read)); + // A caller may itself occupy the only thread of a custom Rayon pool. + // Let it execute scoped work while the independent controller samples + // memory and sends cancellations; never block that last Rayon worker. + while !controller.is_finished() { + rayon::yield_now(); + std::thread::sleep(Duration::from_millis(1)); + } + controller.join().unwrap() + }) + }) +} + +pub(crate) fn run(jobs: &[T], work: F) -> Result<(), CompileError> +where + F: Fn(&T, &Checkpoint) -> Result<(), AttemptError> + Sync, +{ + if jobs.is_empty() { + return Ok(()); + } + let options = Options::from_env()?; + let mut reader = MemoryReader::new(); + let disabled = std::env::var("IX_VALIDATE_ADAPTIVE").as_deref() == Ok("0"); + if disabled && options.process_budget.is_some() { + return Err(resource_error( + "IX_VALIDATE_MEMORY_GIB cannot be combined with IX_VALIDATE_ADAPTIVE=0", + )); + } + if disabled || reader.read().is_none() { + if !disabled && options.process_budget.is_some() { + return Err(resource_error( + "IX_VALIDATE_MEMORY_GIB requires Linux memory telemetry", + )); + } + if options.verbose { + eprintln!( + "[validate_memory] adaptive admission disabled or Linux memory telemetry unavailable" + ); + } + use rayon::prelude::*; + return jobs.par_iter().try_for_each(|job| { + work(job, &Checkpoint::default()).map_err(AttemptError::into_compile) + }); + } + if options.verbose { + eprintln!( + "[validate_memory] adaptive admission: initial={} max={} budget_gib={:?}", + options.initial_workers, + options.max_workers, + options.process_budget.map(|b| b as f64 / GIB as f64) + ); + } + run_with(jobs, &work, options, move || reader.read()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::compile::memory::{cgroup_dirs, field, psi_total}; + use std::path::PathBuf; + use std::sync::atomic::AtomicUsize; + + fn healthy() -> Memory { + Memory { + capacity: 100 * GIB, + available: 90 * GIB, + process: 10 * GIB, + swap_used: 0, + stall_us: 0, + } + } + fn options(max_workers: usize) -> Options { + Options { + max_workers, + initial_workers: 1, + process_budget: None, + tick: Duration::from_millis(2), + ramp_interval: Duration::from_millis(5), + recovery: Duration::from_millis(5), + idle_timeout: Duration::from_millis(100), + verbose: false, + } + } + fn pool(n: usize) -> rayon::ThreadPool { + rayon::ThreadPoolBuilder::new().num_threads(n).build().unwrap() + } + fn until(mut ready: impl FnMut() -> bool) { + let start = Instant::now(); + while !ready() { + assert!( + start.elapsed() < Duration::from_secs(5), + "test made no progress" + ); + std::thread::sleep(Duration::from_millis(1)); + } + } + struct Scratch<'a> { + live: &'a AtomicUsize, + dropping: Option<&'a AtomicUsize>, + } + impl Drop for Scratch<'_> { + fn drop(&mut self) { + if let Some(dropping) = self.dropping { + dropping.fetch_add(1, Ordering::SeqCst); + std::thread::sleep(Duration::from_millis(20)); + dropping.fetch_sub(1, Ordering::SeqCst); + } + self.live.fetch_sub(1, Ordering::SeqCst); + } + } + + #[test] + fn pressure_accounts_for_resident_plus_swap_budget_and_stalls() { + let mem = healthy(); + assert_eq!( + pressure(mem, mem, Duration::from_secs(1), None), + Pressure::Healthy + ); + assert_eq!( + pressure( + Memory { available: 30 * GIB, ..mem }, + mem, + Duration::from_secs(1), + None + ), + Pressure::Hold + ); + assert_eq!( + pressure( + Memory { available: 20 * GIB, ..mem }, + mem, + Duration::from_secs(1), + None + ), + Pressure::Backoff + ); + assert_eq!( + pressure( + Memory { available: 4 * GIB, ..mem }, + mem, + Duration::from_secs(1), + None + ), + Pressure::Critical + ); + assert_eq!( + pressure( + Memory { stall_us: 110_000, ..mem }, + mem, + Duration::from_secs(1), + None + ), + Pressure::Backoff + ); + assert_eq!( + pressure( + Memory { swap_used: 9 * MIB, ..mem }, + mem, + Duration::from_secs(1), + None + ), + Pressure::Backoff + ); + assert_eq!( + pressure( + Memory { process: 96 * GIB, ..mem }, + mem, + Duration::from_secs(1), + Some(100 * GIB) + ), + Pressure::Critical + ); + // Old swap occupancy is not new pressure. + let swapped = Memory { swap_used: 90 * GIB, ..mem }; + assert_eq!( + pressure(swapped, swapped, Duration::from_secs(1), None), + Pressure::Healthy + ); + } + + #[test] + fn parses_proc_and_cgroup_namespace_ancestors() { + assert_eq!( + field("MemTotal: 123 kB\nMemAvailable: 45 kB\n", "MemAvailable:"), + Some(45) + ); + assert_eq!( + psi_total("some avg10=0 total=999\nfull avg10=0 total=123\n"), + 123 + ); + let mount = "30 20 0:25 / /sys/fs/cgroup rw - cgroup2 cgroup rw"; + assert_eq!( + cgroup_dirs("0::/a/b\n", mount), + vec![ + PathBuf::from("/sys/fs/cgroup/a/b"), + PathBuf::from("/sys/fs/cgroup/a"), + PathBuf::from("/sys/fs/cgroup"), + ] + ); + let namespaced = + "30 20 0:25 /host/container /cg\\040mount rw - cgroup2 cgroup rw"; + assert_eq!( + cgroup_dirs("0::/\n", namespaced), + vec![PathBuf::from("/cg mount")] + ); + assert_eq!( + cgroup_dirs("0::/host/container/child\n", namespaced), + vec![PathBuf::from("/cg mount/child"), PathBuf::from("/cg mount"),] + ); + assert!(cgroup_dirs("0::/../../escape\n", mount).is_empty()); + assert!(cgroup_dirs("1:memory:/old\n", mount).is_empty()); + } + + #[test] + fn healthy_admission_ramps_up_without_exceeding_pool_limit() { + let live = AtomicUsize::new(0); + let peak = AtomicUsize::new(0); + let calls: Vec<_> = (0..12).map(|_| AtomicUsize::new(0)).collect(); + pool(4) + .install(|| { + run_with( + &calls, + &|calls: &AtomicUsize, c: &Checkpoint| { + calls.fetch_add(1, Ordering::SeqCst); + let n = live.fetch_add(1, Ordering::SeqCst) + 1; + let _scratch = Scratch { live: &live, dropping: None }; + peak.fetch_max(n, Ordering::SeqCst); + std::thread::sleep(Duration::from_millis(40)); + c.visit()?; + Ok(()) + }, + options(4), + || Some(healthy()), + ) + }) + .unwrap(); + assert_eq!(peak.load(Ordering::SeqCst), 4); + assert!(calls.iter().all(|n| n.load(Ordering::SeqCst) == 1)); + assert_eq!(live.load(Ordering::SeqCst), 0); + } + + #[test] + fn cancellation_drops_scratch_before_admitting_and_retries_once_alone() { + let live = AtomicUsize::new(0); + let dropping = AtomicUsize::new(0); + let cancelled = AtomicBool::new(false); + let calls: Vec<_> = (0..5).map(|_| AtomicUsize::new(0)).collect(); + let jobs: Vec<_> = (0..5).collect(); + let mut opt = options(2); + opt.initial_workers = 2; + pool(2) + .install(|| { + run_with( + &jobs, + &|id: &usize, c: &Checkpoint| { + let attempt = calls[*id].fetch_add(1, Ordering::SeqCst); + assert_eq!( + dropping.load(Ordering::SeqCst), + 0, + "admitted while cancelled scratch was dropping" + ); + let n = live.fetch_add(1, Ordering::SeqCst); + let _scratch = Scratch { + live: &live, + dropping: (*id == 1 && attempt == 0).then_some(&dropping), + }; + if *id == 1 && attempt == 0 { + until(|| { + c.signals.as_ref().unwrap().cancel.load(Ordering::Relaxed) + }); + cancelled.store(true, Ordering::SeqCst); + c.visit()?; + unreachable!(); + } + if *id == 0 { + until(|| cancelled.load(Ordering::SeqCst)); + } + if attempt > 0 { + assert_eq!(n, 0, "retry was not alone"); + } + Ok(()) + }, + opt, + || { + Some(if live.load(Ordering::SeqCst) >= 2 { + Memory { available: 20 * GIB, ..healthy() } + } else { + healthy() + }) + }, + ) + }) + .unwrap(); + assert_eq!(calls[1].load(Ordering::SeqCst), 2); + for id in [0, 2, 3, 4] { + assert_eq!(calls[id].load(Ordering::SeqCst), 1); + } + assert_eq!(live.load(Ordering::SeqCst), 0); + } + + #[test] + fn completed_work_is_not_retried_even_if_cancellation_was_requested() { + let live = AtomicUsize::new(0); + let finished = AtomicBool::new(false); + let calls = [AtomicUsize::new(0), AtomicUsize::new(0)]; + let mut opt = options(2); + opt.initial_workers = 2; + pool(2) + .install(|| { + run_with( + &[0, 1], + &|id: &usize, c: &Checkpoint| { + calls[*id].fetch_add(1, Ordering::SeqCst); + live.fetch_add(1, Ordering::SeqCst); + let _scratch = Scratch { live: &live, dropping: None }; + if *id == 1 { + until(|| { + c.signals.as_ref().unwrap().cancel.load(Ordering::Relaxed) + }); + finished.store(true, Ordering::SeqCst); + } else { + until(|| finished.load(Ordering::SeqCst)); + } + Ok(()) + }, + opt, + || { + Some(if live.load(Ordering::SeqCst) == 2 { + Memory { available: 20 * GIB, ..healthy() } + } else { + healthy() + }) + }, + ) + }) + .unwrap(); + assert!(calls.iter().all(|n| n.load(Ordering::SeqCst) == 1)); + } + + #[test] + fn oversized_solo_job_returns_resource_error_and_releases_memory() { + let live = AtomicUsize::new(0); + let calls = AtomicUsize::new(0); + let result = pool(1).install(|| { + run_with( + &[()], + &|_: &(), c: &Checkpoint| { + calls.fetch_add(1, Ordering::SeqCst); + live.fetch_add(1, Ordering::SeqCst); + let _scratch = Scratch { live: &live, dropping: Some(&calls) }; + until(|| c.signals.as_ref().unwrap().cancel.load(Ordering::Relaxed)); + c.visit()?; + unreachable!() + }, + options(1), + || { + Some(if live.load(Ordering::SeqCst) == 1 { + Memory { available: GIB, process: 99 * GIB, ..healthy() } + } else { + healthy() + }) + }, + ) + }); + assert!(matches!(result, Err(CompileError::ResourceLimit { .. }))); + assert_eq!(live.load(Ordering::SeqCst), 0); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + + #[test] + fn refuses_to_launch_without_headroom_and_does_not_wait_forever() { + let calls = AtomicUsize::new(0); + let result = pool(1).install(|| { + run_with( + &[()], + &|_: &(), _: &Checkpoint| { + calls.fetch_add(1, Ordering::SeqCst); + Ok(()) + }, + options(1), + || Some(Memory { available: GIB, ..healthy() }), + ) + }); + assert!(matches!(result, Err(CompileError::ResourceLimit { .. }))); + assert_eq!(calls.load(Ordering::SeqCst), 0); + } + + #[test] + fn unavailable_telemetry_cancels_and_drains_active_work() { + let live = AtomicUsize::new(0); + let result = pool(1).install(|| { + run_with( + &[()], + &|_: &(), c: &Checkpoint| { + live.fetch_add(1, Ordering::SeqCst); + let _scratch = Scratch { live: &live, dropping: None }; + until(|| c.signals.as_ref().unwrap().cancel.load(Ordering::Relaxed)); + c.visit()?; + unreachable!() + }, + options(1), + || { + if live.load(Ordering::SeqCst) == 0 { Some(healthy()) } else { None } + }, + ) + }); + assert!(matches!(result, Err(CompileError::ResourceLimit { .. }))); + assert_eq!(live.load(Ordering::SeqCst), 0); + } + + #[test] + fn compile_errors_are_preserved_and_other_jobs_are_drained() { + let live = AtomicUsize::new(0); + let failure = + CompileError::InvalidMutualBlock { reason: "test mismatch".into() }; + let mut opt = options(2); + opt.initial_workers = 2; + let result = pool(2).install(|| { + run_with( + &[0, 1], + &|id: &usize, c: &Checkpoint| { + live.fetch_add(1, Ordering::SeqCst); + let _scratch = Scratch { live: &live, dropping: None }; + if *id == 0 { + until(|| live.load(Ordering::SeqCst) == 2); + Err(failure.clone().into()) + } else { + until(|| { + c.signals.as_ref().unwrap().cancel.load(Ordering::Relaxed) + }); + c.visit()?; + unreachable!() + } + }, + opt, + || Some(healthy()), + ) + }); + assert_eq!(result, Err(failure)); + assert_eq!(live.load(Ordering::SeqCst), 0); + } + + #[test] + fn custom_single_worker_pool_can_make_forward_progress() { + let calls = AtomicUsize::new(0); + pool(1) + .install(|| { + run_with( + &[0, 1, 2], + &|_: &usize, c: &Checkpoint| { + c.visit()?; + calls.fetch_add(1, Ordering::SeqCst); + Ok(()) + }, + options(1), + || Some(healthy()), + ) + }) + .unwrap(); + assert_eq!(calls.load(Ordering::SeqCst), 3); + } +} diff --git a/crates/compile/src/diag.rs b/crates/compile/src/diag.rs index b7b12173..9a84d395 100644 --- a/crates/compile/src/diag.rs +++ b/crates/compile/src/diag.rs @@ -1,5 +1,19 @@ //! Shared diagnostics helpers for phase logging. +use std::io::Read; +use std::sync::{LazyLock, mpsc}; +use std::thread::{self, JoinHandle}; +use std::time::Duration; + +static MEMORY_DIAG: LazyLock = + LazyLock::new(|| std::env::var_os("IX_MEMORY_DIAG").is_some()); + +fn field_kb(text: &str, name: &str) -> Option { + text.lines().find_map(|line| { + line.strip_prefix(name)?.split_whitespace().next()?.parse().ok() + }) +} + /// ` · rss X.X GiB (anon Y.Y, file Z.Z)` sampled from /// `/proc/self/status`, appended to phase logs. Anon can only leave RAM /// via swap; file RSS is reclaimable page cache — the split shows which @@ -8,18 +22,15 @@ pub fn rss_log_suffix() -> String { let Ok(status) = std::fs::read_to_string("/proc/self/status") else { return String::new(); }; - let field_kb = |name: &str| -> Option { - status - .lines() - .find(|l| l.starts_with(name)) - .and_then(|l| l.split_whitespace().nth(1)) - .and_then(|v| v.parse().ok()) - }; let gib_tenths = |kb: u64| -> (u64, u64) { let tenths = kb * 10 / (1024 * 1024); (tenths / 10, tenths % 10) }; - match (field_kb("VmRSS:"), field_kb("RssAnon:"), field_kb("RssFile:")) { + match ( + field_kb(&status, "VmRSS:"), + field_kb(&status, "RssAnon:"), + field_kb(&status, "RssFile:"), + ) { (Some(rss), Some(anon), Some(file)) => { let (r, rt) = gib_tenths(rss); let (a, at) = gib_tenths(anon); @@ -29,3 +40,131 @@ pub fn rss_log_suffix() -> String { _ => String::new(), } } + +// Stream maps instead of allocating a String proportional to the mapping +// count: this diagnostic must remain small even near vm.max_map_count. +fn count_lines(mut input: impl Read) -> std::io::Result { + let mut buffer = [0u8; 8192]; + let mut count = 0; + loop { + let n = input.read(&mut buffer)?; + if n == 0 { + return Ok(count); + } + count += buffer[..n].iter().filter(|&&b| b == b'\n').count() as u64; + } +} + +// Diagnostic output rounds to tenths of a GiB, not an exact byte count. +#[allow(clippy::cast_precision_loss)] +fn gib_field(text: &str, name: &str) -> String { + field_kb(text, name).map_or_else( + || "?".into(), + |kb| format!("{:.1}", kb as f64 / (1024.0 * 1024.0)), + ) +} + +/// Detailed Linux memory snapshot for diagnosing allocation failures that +/// need not be physical OOMs (e.g. address-space or mapping-count limits). +fn memory_snapshot() -> String { + let Ok(status) = std::fs::read_to_string("/proc/self/status") else { + return "memory diagnostics unavailable (no /proc/self/status)".into(); + }; + let meminfo = std::fs::read_to_string("/proc/meminfo").unwrap_or_default(); + let maps = std::fs::File::open("/proc/self/maps") + .ok() + .and_then(|file| count_lines(file).ok()) + .map_or_else(|| "?".into(), |n| n.to_string()); + let map_limit = std::fs::read_to_string("/proc/sys/vm/max_map_count") + .unwrap_or_else(|_| "?".into()); + let overcommit = std::fs::read_to_string("/proc/sys/vm/overcommit_memory") + .unwrap_or_else(|_| "?".into()); + format!( + "rss={} GiB (anon {}, file {}), vmsize={} GiB, peak_rss={} GiB, \ + process_swap={} GiB, available={} GiB, swap_free={}/{} GiB, \ + committed={}/{} GiB (overcommit={}), maps={}/{}", + gib_field(&status, "VmRSS:"), + gib_field(&status, "RssAnon:"), + gib_field(&status, "RssFile:"), + gib_field(&status, "VmSize:"), + gib_field(&status, "VmHWM:"), + gib_field(&status, "VmSwap:"), + gib_field(&meminfo, "MemAvailable:"), + gib_field(&meminfo, "SwapFree:"), + gib_field(&meminfo, "SwapTotal:"), + gib_field(&meminfo, "Committed_AS:"), + gib_field(&meminfo, "CommitLimit:"), + overcommit.trim(), + maps, + map_limit.trim(), + ) +} + +/// Scope guard for an opt-in periodic memory sampler. Dropping it wakes and +/// joins the thread immediately, including on early compilation errors. +pub struct MemorySampler { + stop: mpsc::Sender<()>, + thread: Option>, +} + +impl Drop for MemorySampler { + fn drop(&mut self) { + let _ = self.stop.send(()); + if let Some(thread) = self.thread.take() { + let _ = thread.join(); + } + } +} + +/// With `IX_MEMORY_DIAG=1`, log process/system memory every five seconds, +/// including inside phases that never reach their completion log on OOM. +pub fn memory_sampler(label: &'static str) -> Option { + if !*MEMORY_DIAG { + return None; + } + let (stop, receiver) = mpsc::channel(); + match thread::Builder::new().name("ix-memory-diag".into()).spawn(move || { + loop { + eprintln!("[{label}] memory: {}", memory_snapshot()); + match receiver.recv_timeout(Duration::from_secs(5)) { + Err(mpsc::RecvTimeoutError::Timeout) => {}, + _ => break, + } + } + }) { + Ok(thread) => Some(MemorySampler { stop, thread: Some(thread) }), + Err(error) => { + eprintln!("[{label}] could not start memory diagnostics: {error}"); + None + }, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_proc_fields_without_confusing_similar_names() { + let status = "VmRSS:\t1048576 kB\nRssAnon:\t524288 kB\n"; + assert_eq!(field_kb(status, "VmRSS:"), Some(1048576)); + assert_eq!(gib_field(status, "RssAnon:"), "0.5"); + assert_eq!(field_kb(status, "VmSwap:"), None); + assert_eq!(gib_field(status, "VmSwap:"), "?"); + assert_eq!(field_kb("VmRSS: invalid kB\n", "VmRSS:"), None); + } + + #[test] + fn counts_maps_across_buffer_boundaries() { + let maps = "some mapping\n".repeat(2000); + assert_eq!(count_lines(maps.as_bytes()).unwrap(), 2000); + assert_eq!(count_lines(&b""[..]).unwrap(), 0); + } + + #[test] + fn sampler_guard_wakes_and_joins_its_thread() { + let (stop, receiver) = mpsc::channel(); + let thread = thread::spawn(move || receiver.recv().unwrap()); + drop(MemorySampler { stop, thread: Some(thread) }); + } +} diff --git a/crates/ffi/src/compile.rs b/crates/ffi/src/compile.rs index bc127618..3b4021a3 100644 --- a/crates/ffi/src/compile.rs +++ b/crates/ffi/src/compile.rs @@ -1491,13 +1491,14 @@ impl LeanIxDecompileError { impl LeanIxCompileError { /// Build a Lean CompileError from a Rust CompileError. /// - /// Tags 0–5: + /// Tags 0–6: /// 0: missingConstant (name : String) → 1 obj /// 1: missingAddress (addr : Address) → 1 obj /// 2: invalidMutualBlock (reason : String) → 1 obj /// 3: unsupportedExpr (desc : String) → 1 obj /// 4: unknownUnivParam (curr param : String) → 2 obj /// 5: serializeError (msg : String) → 1 obj + /// 6: resourceLimit (reason : String) → 1 obj pub fn build(err: &CompileError) -> Self { match err { CompileError::MissingConstant { name, .. } => { @@ -1531,6 +1532,11 @@ impl LeanIxCompileError { ctor.set_obj(0, LeanIxSerializeError::build(se)); ctor }, + CompileError::ResourceLimit { reason } => { + let ctor = LeanIxCompileError::alloc(6); + ctor.set_obj(0, build_lean_string(reason)); + ctor + }, } } } @@ -1565,6 +1571,9 @@ impl LeanIxCompileError { 5 => { CompileError::Serialize(LeanIxSerializeError(self.get_obj(0)).decode()) }, + 6 => CompileError::ResourceLimit { + reason: self.get_obj(0).as_string().to_string(), + }, tag => unreachable!("Invalid CompileError tag: {tag}"), } } diff --git a/crates/ffi/src/lean.rs b/crates/ffi/src/lean.rs index 7b7428fc..4c57a4b2 100644 --- a/crates/ffi/src/lean.rs +++ b/crates/ffi/src/lean.rs @@ -345,6 +345,7 @@ lean_ffi::lean_inductive! { { num_obj: 1 }, // tag 3: unsupportedExpr { num_obj: 2 }, // tag 4: unknownUnivParam { num_obj: 1 }, // tag 5: serialize + { num_obj: 1 }, // tag 6: resourceLimit ]; // Defined in `Ix/KernelCheck.lean`. diff --git a/crates/ixon/src/error.rs b/crates/ixon/src/error.rs index 0e2c8462..04b65e74 100644 --- a/crates/ixon/src/error.rs +++ b/crates/ixon/src/error.rs @@ -49,7 +49,7 @@ impl std::error::Error for SerializeError {} /// Errors during compilation (Lean → Ixon). /// -/// Variant order matches Lean constructor tags (0–5). +/// Variant order matches Lean constructor tags (0–6). #[derive(Debug, Clone, PartialEq, Eq)] pub enum CompileError { /// Referenced constant not found (tag 0). @@ -65,6 +65,8 @@ pub enum CompileError { UnknownUnivParam { curr: String, param: String }, /// Serialization error during compilation (tag 5) Serialize(SerializeError), + /// Compilation could not stay within its resource budget (tag 6). + ResourceLimit { reason: String }, } impl std::fmt::Display for CompileError { @@ -84,6 +86,7 @@ impl std::fmt::Display for CompileError { write!(f, "unknown universe parameter: compiling {curr}, param {param}") }, Self::Serialize(e) => write!(f, "serialization error: {e}"), + Self::ResourceLimit { reason } => write!(f, "resource limit: {reason}"), } } } From ac6a477c6bba33087ce0a5ad0532d204f654ca15 Mon Sep 17 00:00:00 2001 From: "John C. Burnham" Date: Sat, 5 Sep 2026 18:50:20 -0400 Subject: [PATCH 3/4] perf(compile): memoize shared expression DAG traversals Fold graph setup into chunk-local accumulators and preserve DAG sharing across Lean/kernel conversion, expression interning, and source-name restoration. Keep memoization pass-local and use exact structural hint keys, with differential, collision, and output-parity regression tests. --- crates/compile/src/compile/aux_gen.rs | 2 + .../compile/src/compile/aux_gen/expr_utils.rs | 665 ++++++------------ .../aux_gen/kernel_bridge_reference.rs | 304 ++++++++ .../compile/aux_gen/kernel_bridge_tests.rs | 336 +++++++++ .../src/compile/aux_gen/source_name_hints.rs | 406 +++++++++++ .../aux_gen/source_name_hints_reference.rs | 393 +++++++++++ .../aux_gen/source_name_hints_tests.rs | 353 ++++++++++ crates/compile/src/graph.rs | 318 ++++++++- crates/kernel/src/env.rs | 197 ++++-- crates/kernel/src/intern_reference.rs | 156 ++++ crates/kernel/src/intern_tests.rs | 436 ++++++++++++ docs/ix_canonicity.md | 26 +- 12 files changed, 3066 insertions(+), 526 deletions(-) create mode 100644 crates/compile/src/compile/aux_gen/kernel_bridge_reference.rs create mode 100644 crates/compile/src/compile/aux_gen/kernel_bridge_tests.rs create mode 100644 crates/compile/src/compile/aux_gen/source_name_hints.rs create mode 100644 crates/compile/src/compile/aux_gen/source_name_hints_reference.rs create mode 100644 crates/compile/src/compile/aux_gen/source_name_hints_tests.rs create mode 100644 crates/kernel/src/intern_reference.rs create mode 100644 crates/kernel/src/intern_tests.rs diff --git a/crates/compile/src/compile/aux_gen.rs b/crates/compile/src/compile/aux_gen.rs index c687174d..6e9036f2 100644 --- a/crates/compile/src/compile/aux_gen.rs +++ b/crates/compile/src/compile/aux_gen.rs @@ -81,6 +81,8 @@ mod checked_expr; #[cfg(test)] mod expr_reference; pub mod expr_utils; +#[cfg(test)] +mod kernel_bridge_reference; pub mod nested; pub mod rec_on; pub mod recursor; diff --git a/crates/compile/src/compile/aux_gen/expr_utils.rs b/crates/compile/src/compile/aux_gen/expr_utils.rs index 6cfef80b..84b51fcd 100644 --- a/crates/compile/src/compile/aux_gen/expr_utils.rs +++ b/crates/compile/src/compile/aux_gen/expr_utils.rs @@ -18,6 +18,13 @@ use ix_common::env::{ use ix_kernel::ingress::{lean_level_to_kuniv, resolve_lean_name_addr}; use ix_kernel::mode::Meta; +#[path = "source_name_hints.rs"] +mod source_name_hints; + +#[cfg(test)] +#[path = "source_name_hints_reference.rs"] +mod source_name_hints_reference; + // ========================================================================= // FVar infrastructure // ========================================================================= @@ -2405,22 +2412,13 @@ impl<'a> TcScope<'a> { if whnfed == kexpr { restore_source_names_same_content(&out, ty, self.stt) } else { - let mut source_name_hints = FxHashMap::default(); - collect_lean_source_name_hints( - ty, - &self.fvar_levels, - depth, - self.param_names, - self.stt, - &mut source_name_hints, - ); - restore_lean_source_name_hints( + source_name_hints::restore( &out, + ty, &self.fvar_levels, depth, self.param_names, self.stt, - &source_name_hints, ) } } @@ -2483,15 +2481,42 @@ impl<'a> TcScope<'a> { /// /// `Mdata` layers carried by the kernel expression are re-wrapped around /// the result in original order — matching `egress_expr`. +/// Memoization preserves shared subexpressions at each binder depth instead +/// of expanding the input DAG into a tree. The cache lives for one call only. pub(super) fn kexpr_to_lean( expr: &ix_kernel::expr::KExpr, outer_depth: usize, fvar_levels: &FxHashMap, local_depth: usize, param_names: &[Name], +) -> LeanExpr { + kexpr_to_lean_cached( + expr, + outer_depth, + fvar_levels, + local_depth, + param_names, + &mut FxHashMap::default(), + ) +} + +fn kexpr_to_lean_cached( + expr: &ix_kernel::expr::KExpr, + outer_depth: usize, + fvar_levels: &FxHashMap, + local_depth: usize, + param_names: &[Name], + cache: &mut FxHashMap<(usize, usize), LeanExpr>, ) -> LeanExpr { use ix_kernel::expr::ExprData as KED; + // Input nodes stay alive for this call. Intern uids are NOT sufficient: + // distinct metadata-bearing nodes can have the same name-erased identity. + let key = (std::ptr::from_ref(expr.data()).addr(), local_depth); + if let Some(result) = cache.get(&key) { + return result.clone(); + } + // Reverse `fvar_levels` lazily via linear search — the FVar context is // small in practice (a handful of param/motive/minor/index binders), // so an O(n) scan per Var hit is cheaper than maintaining an inverse @@ -2541,443 +2566,121 @@ pub(super) fn kexpr_to_lean( LeanExpr::cnst(kid.name.clone(), levels) }, KED::App(f, a, _) => LeanExpr::app( - kexpr_to_lean(f, outer_depth, fvar_levels, local_depth, param_names), - kexpr_to_lean(a, outer_depth, fvar_levels, local_depth, param_names), - ), - KED::All(name, bi, d, b, _) => LeanExpr::all( - name.clone(), - kexpr_to_lean(d, outer_depth, fvar_levels, local_depth, param_names), - kexpr_to_lean(b, outer_depth, fvar_levels, local_depth + 1, param_names), - bi.clone(), - ), - KED::Lam(name, bi, d, b, _) => LeanExpr::lam( - name.clone(), - kexpr_to_lean(d, outer_depth, fvar_levels, local_depth, param_names), - kexpr_to_lean(b, outer_depth, fvar_levels, local_depth + 1, param_names), - bi.clone(), - ), - KED::Let(name, ty, val, body, nd, _) => LeanExpr::letE( - name.clone(), - kexpr_to_lean(ty, outer_depth, fvar_levels, local_depth, param_names), - kexpr_to_lean(val, outer_depth, fvar_levels, local_depth, param_names), - kexpr_to_lean( - body, - outer_depth, - fvar_levels, - local_depth + 1, - param_names, - ), - *nd, - ), - KED::Prj(kid, field, val, _) => LeanExpr::proj( - kid.name.clone(), - Nat::from(*field), - kexpr_to_lean(val, outer_depth, fvar_levels, local_depth, param_names), - ), - KED::Nat(n, _, _) => { - use ix_common::env::Literal; - LeanExpr::lit(Literal::NatVal(n.clone())) - }, - KED::Str(s, _, _) => { - use ix_common::env::Literal; - LeanExpr::lit(Literal::StrVal(s.clone())) - }, - }; - - // Re-wrap mdata layers, outermost first (matching egress_expr's order). - expr - .mdata() - .iter() - .rev() - .fold(inner, |acc, kvs| LeanExpr::mdata(kvs.clone(), acc)) -} - -fn source_name_hint_candidate(expr: &LeanExpr) -> bool { - matches!(expr.as_data(), ExprData::App(..) | ExprData::Proj(..)) -} - -/// Name-erased structural content key for the source-name hint map. -/// -/// Mirrors the equivalence of the Lean pipeline's `Ix.Tc.KExpr` content -/// addresses (`toKexprStatic ... |>.addr` in `Ix/AuxGen/Kernel.lean`) and -/// of the kernel's `ExprKey`/`structural_eq`: display names, binder -/// names, binder infos, and mdata are excluded; `Const`/`Prj` contribute -/// their resolved content address, universes their index structure. Two -/// spellings of one alias pair (`Paths V` / `Symmetrify V`) therefore -/// agree on this key — which is the whole point of the hint map. -/// -/// `KExpr::hash_key()` is NOT usable here: it is the intern-assigned uid, -/// fresh for every un-interned construction, and `to_kexpr_static` does -/// not intern — so the collect-time and restore-time keys of two -/// content-equal subterms never matched, and the hint map restored -/// nothing (the Mathlib `Quiver.FreeGroupoid.redStep` metadata -/// divergence, canonicity §10.5). -fn kexpr_content_key(e: &ix_kernel::expr::KExpr) -> u64 { - use std::hash::Hasher; - let mut h = rustc_hash::FxHasher::default(); - kexpr_content_hash(e, &mut h); - h.finish() -} - -fn kuniv_content_hash( - u: &ix_kernel::level::KUniv, - h: &mut rustc_hash::FxHasher, -) { - use ix_kernel::level::UnivData as UD; - use std::hash::Hasher; - match u.data() { - UD::Zero(_) => h.write_u8(0), - UD::Succ(a, _) => { - h.write_u8(1); - kuniv_content_hash(a, h); - }, - UD::Max(a, b, _) => { - h.write_u8(2); - kuniv_content_hash(a, h); - kuniv_content_hash(b, h); - }, - UD::IMax(a, b, _) => { - h.write_u8(3); - kuniv_content_hash(a, h); - kuniv_content_hash(b, h); - }, - UD::Param(idx, _, _) => { - h.write_u8(4); - h.write_u64(*idx); - }, - } -} - -fn kexpr_content_hash( - e: &ix_kernel::expr::KExpr, - h: &mut rustc_hash::FxHasher, -) { - use ix_kernel::expr::ExprData as KED; - use std::hash::Hasher; - match e.data() { - KED::Var(i, _, _) => { - h.write_u8(0); - h.write_u64(*i); - }, - KED::FVar(id, _, _) => { - h.write_u8(1); - h.write_u64(id.0); - }, - KED::Sort(u, _) => { - h.write_u8(2); - kuniv_content_hash(u, h); - }, - KED::Const(id, us, _) => { - h.write_u8(3); - h.write(id.addr.as_bytes()); - h.write_u64(us.len() as u64); - for u in us.iter() { - kuniv_content_hash(u, h); - } - }, - KED::App(f, a, _) => { - h.write_u8(4); - kexpr_content_hash(f, h); - kexpr_content_hash(a, h); - }, - KED::Lam(_, _, t, b, _) => { - h.write_u8(5); - kexpr_content_hash(t, h); - kexpr_content_hash(b, h); - }, - KED::All(_, _, t, b, _) => { - h.write_u8(6); - kexpr_content_hash(t, h); - kexpr_content_hash(b, h); - }, - KED::Let(_, t, v, b, nd, _) => { - h.write_u8(7); - h.write_u8(u8::from(*nd)); - kexpr_content_hash(t, h); - kexpr_content_hash(v, h); - kexpr_content_hash(b, h); - }, - KED::Prj(id, f, v, _) => { - h.write_u8(8); - h.write(id.addr.as_bytes()); - h.write_u64(*f); - kexpr_content_hash(v, h); - }, - KED::Nat(_, ba, _) => { - h.write_u8(9); - h.write(ba.as_bytes()); - }, - KED::Str(_, ba, _) => { - h.write_u8(10); - h.write(ba.as_bytes()); - }, - } -} - -/// Collect source-shaped subterms that WHNF may copy into a reduct. -/// -/// Keys use the kernel content hash so alpha-collapsed aliases like -/// `CategoryTheory.Paths V` and `Quiver.Symmetrify V` line up, while values -/// keep the Lean display names from the caller. We skip BVar-containing terms: -/// WHNF may lift copied arguments under freshly-exposed binders, so matching -/// those by raw de Bruijn indices would be unstable. -fn collect_lean_source_name_hints( - source: &LeanExpr, - fvar_levels: &FxHashMap, - depth: usize, - param_names: &[Name], - stt: &crate::compile::CompileState, - out: &mut FxHashMap, -) { - if source_name_hint_candidate(source) && !expr_has_bvar(source) { - let key = kexpr_content_key(&to_kexpr_static( - source, - fvar_levels, - depth, - param_names, - stt, - )); - out.entry(key).or_insert_with(|| source.clone()); - } - - match source.as_data() { - ExprData::Mdata(_, inner, _) => collect_lean_source_name_hints( - inner, - fvar_levels, - depth, - param_names, - stt, - out, - ), - ExprData::App(f, a, _) => { - collect_lean_source_name_hints( - f, - fvar_levels, - depth, - param_names, - stt, - out, - ); - collect_lean_source_name_hints( - a, - fvar_levels, - depth, - param_names, - stt, - out, - ); - }, - ExprData::ForallE(_, d, b, _, _) | ExprData::Lam(_, d, b, _, _) => { - collect_lean_source_name_hints( - d, - fvar_levels, - depth, - param_names, - stt, - out, - ); - collect_lean_source_name_hints( - b, - fvar_levels, - depth, - param_names, - stt, - out, - ); - }, - ExprData::LetE(_, t, v, b, _, _) => { - collect_lean_source_name_hints( - t, - fvar_levels, - depth, - param_names, - stt, - out, - ); - collect_lean_source_name_hints( - v, - fvar_levels, - depth, - param_names, - stt, - out, - ); - collect_lean_source_name_hints( - b, - fvar_levels, - depth, - param_names, - stt, - out, - ); - }, - ExprData::Proj(_, _, v, _) => collect_lean_source_name_hints( - v, - fvar_levels, - depth, - param_names, - stt, - out, - ), - _ => {}, - } -} - -/// Restore source spellings for copied subterms after a real WHNF reduction. -/// -/// This is intentionally subterm-based rather than whole-expression based: -/// unfolding a reducible alias such as `HomRel (Paths (Symmetrify V))` should -/// keep the expanded `∀` telescope, but the repeated argument subterms inside -/// that telescope should retain the caller's `Symmetrify` spelling instead of -/// whichever same-address alias the kernel cache/intern table already held. -fn restore_lean_source_name_hints( - generated: &LeanExpr, - fvar_levels: &FxHashMap, - depth: usize, - param_names: &[Name], - stt: &crate::compile::CompileState, - hints: &FxHashMap, -) -> LeanExpr { - if source_name_hint_candidate(generated) && !expr_has_bvar(generated) { - let key = kexpr_content_key(&to_kexpr_static( - generated, - fvar_levels, - depth, - param_names, - stt, - )); - if let Some(source) = hints.get(&key) { - return source.clone(); - } - } - - match generated.as_data() { - ExprData::App(f, a, _) => LeanExpr::app( - restore_lean_source_name_hints( + kexpr_to_lean_cached( f, + outer_depth, fvar_levels, - depth, + local_depth, param_names, - stt, - hints, + cache, ), - restore_lean_source_name_hints( + kexpr_to_lean_cached( a, + outer_depth, fvar_levels, - depth, + local_depth, param_names, - stt, - hints, + cache, ), ), - ExprData::ForallE(n, d, b, bi, _) => LeanExpr::all( - n.clone(), - restore_lean_source_name_hints( + KED::All(name, bi, d, b, _) => LeanExpr::all( + name.clone(), + kexpr_to_lean_cached( d, + outer_depth, fvar_levels, - depth, + local_depth, param_names, - stt, - hints, + cache, ), - restore_lean_source_name_hints( + kexpr_to_lean_cached( b, + outer_depth, fvar_levels, - depth, + local_depth + 1, param_names, - stt, - hints, + cache, ), bi.clone(), ), - ExprData::Lam(n, d, b, bi, _) => LeanExpr::lam( - n.clone(), - restore_lean_source_name_hints( + KED::Lam(name, bi, d, b, _) => LeanExpr::lam( + name.clone(), + kexpr_to_lean_cached( d, + outer_depth, fvar_levels, - depth, + local_depth, param_names, - stt, - hints, + cache, ), - restore_lean_source_name_hints( + kexpr_to_lean_cached( b, + outer_depth, fvar_levels, - depth, + local_depth + 1, param_names, - stt, - hints, + cache, ), bi.clone(), ), - ExprData::LetE(n, t, v, b, nd, _) => LeanExpr::letE( - n.clone(), - restore_lean_source_name_hints( - t, + KED::Let(name, ty, val, body, nd, _) => LeanExpr::letE( + name.clone(), + kexpr_to_lean_cached( + ty, + outer_depth, fvar_levels, - depth, + local_depth, param_names, - stt, - hints, + cache, ), - restore_lean_source_name_hints( - v, + kexpr_to_lean_cached( + val, + outer_depth, fvar_levels, - depth, + local_depth, param_names, - stt, - hints, + cache, ), - restore_lean_source_name_hints( - b, + kexpr_to_lean_cached( + body, + outer_depth, fvar_levels, - depth, + local_depth + 1, param_names, - stt, - hints, + cache, ), *nd, ), - ExprData::Proj(n, i, v, _) => LeanExpr::proj( - n.clone(), - i.clone(), - restore_lean_source_name_hints( - v, - fvar_levels, - depth, - param_names, - stt, - hints, - ), - ), - ExprData::Mdata(kvs, v, _) => LeanExpr::mdata( - kvs.clone(), - restore_lean_source_name_hints( - v, + KED::Prj(kid, field, val, _) => LeanExpr::proj( + kid.name.clone(), + Nat::from(*field), + kexpr_to_lean_cached( + val, + outer_depth, fvar_levels, - depth, + local_depth, param_names, - stt, - hints, + cache, ), ), - _ => generated.clone(), - } -} - -fn expr_has_bvar(expr: &LeanExpr) -> bool { - match expr.as_data() { - ExprData::Bvar(..) => true, - ExprData::App(f, a, _) => expr_has_bvar(f) || expr_has_bvar(a), - ExprData::ForallE(_, d, b, _, _) | ExprData::Lam(_, d, b, _, _) => { - expr_has_bvar(d) || expr_has_bvar(b) + KED::Nat(n, _, _) => { + use ix_common::env::Literal; + LeanExpr::lit(Literal::NatVal(n.clone())) }, - ExprData::LetE(_, t, v, b, _, _) => { - expr_has_bvar(t) || expr_has_bvar(v) || expr_has_bvar(b) + KED::Str(s, _, _) => { + use ix_common::env::Literal; + LeanExpr::lit(Literal::StrVal(s.clone())) }, - ExprData::Proj(_, _, v, _) | ExprData::Mdata(_, v, _) => expr_has_bvar(v), - _ => false, - } + }; + + // Re-wrap mdata layers, outermost first (matching egress_expr's order). + let result = expr + .mdata() + .iter() + .rev() + .fold(inner, |acc, kvs| LeanExpr::mdata(kvs.clone(), acc)); + cache.insert(key, result.clone()); + result } /// Restore source-side display names after a WHNF roundtrip that did not @@ -2989,26 +2692,53 @@ fn expr_has_bvar(expr: &LeanExpr) -> bool { /// output are equal as kernel content we prefer the caller's Lean names while /// keeping the output's reduced levels/subterms. Real reductions are filtered /// by the caller's top-level content-hash check before this function is used. +/// Cache by both full Lean identities so sharing is preserved without merging +/// distinct source spellings of the same name-erased kernel expression. fn restore_source_names_same_content( generated: &LeanExpr, source: &LeanExpr, stt: &crate::compile::CompileState, ) -> LeanExpr { - let source = strip_mdata_ref(source); + restore_source_names_cached(generated, source, stt, &mut FxHashMap::default()) +} - match generated.as_data() { +fn restore_source_names_cached( + generated: &LeanExpr, + source: &LeanExpr, + stt: &crate::compile::CompileState, + cache: &mut FxHashMap<(blake3::Hash, blake3::Hash), LeanExpr>, +) -> LeanExpr { + let source = strip_mdata_ref(source); + if generated.get_hash() == source.get_hash() { + return generated.clone(); + } + // One generated node can inherit different aliases at different source + // occurrences, so both source and generated identity belong in the key. + let key = (*generated.get_hash(), *source.get_hash()); + if let Some(result) = cache.get(&key) { + return result.clone(); + } + let result = match generated.as_data() { ExprData::Mdata(kvs, inner, _) => LeanExpr::mdata( kvs.clone(), - restore_source_names_same_content(inner, source, stt), + restore_source_names_cached(inner, source, stt, cache), ), - _ => restore_source_names_same_content_inner(generated, source, stt), - } + _ => restore_source_names_same_content_inner(generated, source, stt, cache), + }; + let result = if result.get_hash() == generated.get_hash() { + generated.clone() + } else { + result + }; + cache.insert(key, result.clone()); + result } fn restore_source_names_same_content_inner( generated: &LeanExpr, source: &LeanExpr, stt: &crate::compile::CompileState, + cache: &mut FxHashMap<(blake3::Hash, blake3::Hash), LeanExpr>, ) -> LeanExpr { match (generated.as_data(), source.as_data()) { ( @@ -3019,8 +2749,8 @@ fn restore_source_names_same_content_inner( }, (ExprData::App(gen_f, gen_a, _), ExprData::App(source_f, source_a, _)) => { LeanExpr::app( - restore_source_names_same_content(gen_f, source_f, stt), - restore_source_names_same_content(gen_a, source_a, stt), + restore_source_names_cached(gen_f, source_f, stt, cache), + restore_source_names_cached(gen_a, source_a, stt, cache), ) }, ( @@ -3028,8 +2758,8 @@ fn restore_source_names_same_content_inner( ExprData::ForallE(source_name, source_dom, source_body, _, _), ) => LeanExpr::all( source_name.clone(), - restore_source_names_same_content(gen_dom, source_dom, stt), - restore_source_names_same_content(gen_body, source_body, stt), + restore_source_names_cached(gen_dom, source_dom, stt, cache), + restore_source_names_cached(gen_body, source_body, stt, cache), gen_bi.clone(), ), ( @@ -3037,8 +2767,8 @@ fn restore_source_names_same_content_inner( ExprData::Lam(source_name, source_dom, source_body, _, _), ) => LeanExpr::lam( source_name.clone(), - restore_source_names_same_content(gen_dom, source_dom, stt), - restore_source_names_same_content(gen_body, source_body, stt), + restore_source_names_cached(gen_dom, source_dom, stt, cache), + restore_source_names_cached(gen_body, source_body, stt, cache), gen_bi.clone(), ), ( @@ -3046,9 +2776,9 @@ fn restore_source_names_same_content_inner( ExprData::LetE(source_name, source_ty, source_val, source_body, _, _), ) => LeanExpr::letE( source_name.clone(), - restore_source_names_same_content(gen_ty, source_ty, stt), - restore_source_names_same_content(gen_val, source_val, stt), - restore_source_names_same_content(gen_body, source_body, stt), + restore_source_names_cached(gen_ty, source_ty, stt, cache), + restore_source_names_cached(gen_val, source_val, stt, cache), + restore_source_names_cached(gen_body, source_body, stt, cache), *gen_nd, ), ( @@ -3060,7 +2790,7 @@ fn restore_source_names_same_content_inner( LeanExpr::proj( source_name.clone(), gen_field.clone(), - restore_source_names_same_content(gen_val, source_val, stt), + restore_source_names_cached(gen_val, source_val, stt, cache), ) }, _ => generated.clone(), @@ -3088,17 +2818,34 @@ fn same_resolved_name_addr( == resolve_lean_name_addr(b, n2a, aux_n2a) } -/// Static version of `to_kexpr` that takes borrowed references. -/// -/// Identical to the closure-based `to_kexpr` in `get_level`, but as a -/// standalone function so it can be called from both `PreparedTC::new` -/// and `get_level_with_tc`. +/// Convert a Lean expression in the current FVar context into kernel syntax. +/// Memoize by full Lean identity and binder depth to preserve the input DAG. +/// The cache is scoped to this call: FVar/universe bindings and resolved +/// constant addresses can change between calls, even within one block. fn to_kexpr_static( expr: &LeanExpr, fvar_levels: &FxHashMap, ctx_depth: usize, param_names: &[Name], stt: &crate::compile::CompileState, +) -> ix_kernel::expr::KExpr { + to_kexpr_cached( + expr, + fvar_levels, + ctx_depth, + param_names, + stt, + &mut FxHashMap::default(), + ) +} + +fn to_kexpr_cached( + expr: &LeanExpr, + fvar_levels: &FxHashMap, + ctx_depth: usize, + param_names: &[Name], + stt: &crate::compile::CompileState, + cache: &mut FxHashMap<(blake3::Hash, usize), ix_kernel::expr::KExpr>, ) -> ix_kernel::expr::KExpr { let n2a = Some(&stt.name_to_addr); let aux_n2a = Some(&stt.aux_name_to_addr); @@ -3106,7 +2853,14 @@ fn to_kexpr_static( use ix_kernel::id::KId; use ix_kernel::level::KUniv; - match expr.as_data() { + // The FVar/universe/address context is fixed for this conversion call. + // Lean's digest includes source names and metadata; binder depth is + // additionally needed because FVars map to different de Bruijn indices. + let key = (*expr.get_hash(), ctx_depth); + if let Some(result) = cache.get(&key) { + return result.clone(); + } + let result = match expr.as_data() { ExprData::Fvar(fname, _) => { if let Some(&level) = fvar_levels.get(fname) { KExpr::var((ctx_depth - level - 1) as u64, Name::anon()) @@ -3126,33 +2880,58 @@ fn to_kexpr_static( KExpr::cnst(zid, zus) }, ExprData::App(f, a, _) => { - let kf = to_kexpr_static(f, fvar_levels, ctx_depth, param_names, stt); - let ka = to_kexpr_static(a, fvar_levels, ctx_depth, param_names, stt); + let kf = + to_kexpr_cached(f, fvar_levels, ctx_depth, param_names, stt, cache); + let ka = + to_kexpr_cached(a, fvar_levels, ctx_depth, param_names, stt, cache); KExpr::app(kf, ka) }, ExprData::ForallE(binder_name, dom, body, bi, _) => { - let kd = to_kexpr_static(dom, fvar_levels, ctx_depth, param_names, stt); - let kb = - to_kexpr_static(body, fvar_levels, ctx_depth + 1, param_names, stt); + let kd = + to_kexpr_cached(dom, fvar_levels, ctx_depth, param_names, stt, cache); + let kb = to_kexpr_cached( + body, + fvar_levels, + ctx_depth + 1, + param_names, + stt, + cache, + ); KExpr::all(binder_name.clone(), bi.clone(), kd, kb) }, ExprData::Lam(binder_name, dom, body, bi, _) => { - let kd = to_kexpr_static(dom, fvar_levels, ctx_depth, param_names, stt); - let kb = - to_kexpr_static(body, fvar_levels, ctx_depth + 1, param_names, stt); + let kd = + to_kexpr_cached(dom, fvar_levels, ctx_depth, param_names, stt, cache); + let kb = to_kexpr_cached( + body, + fvar_levels, + ctx_depth + 1, + param_names, + stt, + cache, + ); KExpr::lam(binder_name.clone(), bi.clone(), kd, kb) }, ExprData::LetE(binder_name, ty, val, body, nd, _) => { - let kt = to_kexpr_static(ty, fvar_levels, ctx_depth, param_names, stt); - let kv = to_kexpr_static(val, fvar_levels, ctx_depth, param_names, stt); - let kb = - to_kexpr_static(body, fvar_levels, ctx_depth + 1, param_names, stt); + let kt = + to_kexpr_cached(ty, fvar_levels, ctx_depth, param_names, stt, cache); + let kv = + to_kexpr_cached(val, fvar_levels, ctx_depth, param_names, stt, cache); + let kb = to_kexpr_cached( + body, + fvar_levels, + ctx_depth + 1, + param_names, + stt, + cache, + ); KExpr::let_(binder_name.clone(), kt, kv, kb, *nd) }, ExprData::Proj(pname, idx, e, _) => { let addr = resolve_lean_name_addr(pname, n2a, aux_n2a); let zid = KId::new(addr, pname.clone()); - let ke = to_kexpr_static(e, fvar_levels, ctx_depth, param_names, stt); + let ke = + to_kexpr_cached(e, fvar_levels, ctx_depth, param_names, stt, cache); KExpr::prj(zid, nat_to_u64(idx), ke) }, ExprData::Lit(lit, _) => { @@ -3169,15 +2948,21 @@ fn to_kexpr_static( } }, ExprData::Mdata(_, inner, _) => { - to_kexpr_static(inner, fvar_levels, ctx_depth, param_names, stt) + to_kexpr_cached(inner, fvar_levels, ctx_depth, param_names, stt, cache) }, _ => KExpr::sort(KUniv::zero()), - } + }; + cache.insert(key, result.clone()); + result } fn collect_lean_const_refs(expr: &LeanExpr, out: &mut FxHashSet) { + let mut visited = FxHashSet::default(); let mut stack = vec![expr]; while let Some(expr) = stack.pop() { + if !visited.insert(*expr.get_hash()) { + continue; + } match expr.as_data() { ExprData::Const(name, _, _) => { out.insert(name.clone()); @@ -3205,6 +2990,10 @@ fn collect_lean_const_refs(expr: &LeanExpr, out: &mut FxHashSet) { } } +#[cfg(test)] +#[path = "kernel_bridge_tests.rs"] +mod kernel_bridge_tests; + #[cfg(test)] mod tests { use super::*; diff --git a/crates/compile/src/compile/aux_gen/kernel_bridge_reference.rs b/crates/compile/src/compile/aux_gen/kernel_bridge_reference.rs new file mode 100644 index 00000000..086fc6b1 --- /dev/null +++ b/crates/compile/src/compile/aux_gen/kernel_bridge_reference.rs @@ -0,0 +1,304 @@ +//! Frozen pre-memoization bridge transforms for small differential fixtures. + +use crate::compile::nat_conv::nat_to_u64; +use bignat::Nat; +use ix_common::address::Address; +use ix_common::env::{Expr as LeanExpr, ExprData, Level, Name}; +use ix_kernel::ingress::{lean_level_to_kuniv, resolve_lean_name_addr}; +use ix_kernel::mode::Meta; +use rustc_hash::FxHashMap; + +pub(super) fn kexpr_to_lean( + expr: &ix_kernel::expr::KExpr, + outer_depth: usize, + fvar_levels: &FxHashMap, + local_depth: usize, + param_names: &[Name], +) -> LeanExpr { + use ix_kernel::expr::ExprData as KED; + + // Reverse `fvar_levels` lazily via linear search — the FVar context is + // small in practice (a handful of param/motive/minor/index binders), + // so an O(n) scan per Var hit is cheaper than maintaining an inverse + // map alongside `TcScope`. + let lookup_fvar = |level: usize| -> Option { + fvar_levels.iter().find_map(|(name, &lvl)| { + if lvl == level { Some(name.clone()) } else { None } + }) + }; + + let inner = match expr.data() { + KED::Var(i, _, _) => { + let i = *i as usize; + if i < local_depth { + LeanExpr::bvar(Nat::from(i as u64)) + } else { + let fvar_idx_from_top = i - local_depth; + let level = outer_depth + .checked_sub(fvar_idx_from_top + 1) + .expect("kexpr_to_lean: Var index out of range of outer context"); + let name = lookup_fvar(level).unwrap_or_else(|| { + // Unregistered FVar — indicates mismatched `fvar_levels` vs. + // the expression's Var indices. Use a synthetic placeholder + // rather than panic so diagnostics can surface the issue. + Name::str(Name::anon(), format!("_dangling_fvar_{level}")) + }); + LeanExpr::fvar(name) + } + }, + // Kernel-side FVar nodes (introduced by binder opening during type + // checking) should never appear in the inputs of `kexpr_to_lean`, + // which converts ingressed/compile-time expressions back to Lean + // syntax. If one does appear, it indicates a path leaked an open + // expression past its abstraction step — treat it as a synthetic + // free variable named after its id so diagnostics can surface it. + KED::FVar(id, _, _) => { + LeanExpr::fvar(Name::str(Name::anon(), format!("_kernel_fvar_{}", id.0))) + }, + KED::Sort(u, _) => { + LeanExpr::sort(super::below::kuniv_to_level(u, param_names)) + }, + KED::Const(kid, us, _) => { + let levels: Vec = us + .iter() + .map(|u| super::below::kuniv_to_level(u, param_names)) + .collect(); + LeanExpr::cnst(kid.name.clone(), levels) + }, + KED::App(f, a, _) => LeanExpr::app( + kexpr_to_lean(f, outer_depth, fvar_levels, local_depth, param_names), + kexpr_to_lean(a, outer_depth, fvar_levels, local_depth, param_names), + ), + KED::All(name, bi, d, b, _) => LeanExpr::all( + name.clone(), + kexpr_to_lean(d, outer_depth, fvar_levels, local_depth, param_names), + kexpr_to_lean(b, outer_depth, fvar_levels, local_depth + 1, param_names), + bi.clone(), + ), + KED::Lam(name, bi, d, b, _) => LeanExpr::lam( + name.clone(), + kexpr_to_lean(d, outer_depth, fvar_levels, local_depth, param_names), + kexpr_to_lean(b, outer_depth, fvar_levels, local_depth + 1, param_names), + bi.clone(), + ), + KED::Let(name, ty, val, body, nd, _) => LeanExpr::letE( + name.clone(), + kexpr_to_lean(ty, outer_depth, fvar_levels, local_depth, param_names), + kexpr_to_lean(val, outer_depth, fvar_levels, local_depth, param_names), + kexpr_to_lean( + body, + outer_depth, + fvar_levels, + local_depth + 1, + param_names, + ), + *nd, + ), + KED::Prj(kid, field, val, _) => LeanExpr::proj( + kid.name.clone(), + Nat::from(*field), + kexpr_to_lean(val, outer_depth, fvar_levels, local_depth, param_names), + ), + KED::Nat(n, _, _) => { + use ix_common::env::Literal; + LeanExpr::lit(Literal::NatVal(n.clone())) + }, + KED::Str(s, _, _) => { + use ix_common::env::Literal; + LeanExpr::lit(Literal::StrVal(s.clone())) + }, + }; + + // Re-wrap mdata layers, outermost first (matching egress_expr's order). + expr + .mdata() + .iter() + .rev() + .fold(inner, |acc, kvs| LeanExpr::mdata(kvs.clone(), acc)) +} + +pub(super) fn restore_source_names_same_content( + generated: &LeanExpr, + source: &LeanExpr, + stt: &crate::compile::CompileState, +) -> LeanExpr { + let source = strip_mdata_ref(source); + + match generated.as_data() { + ExprData::Mdata(kvs, inner, _) => LeanExpr::mdata( + kvs.clone(), + restore_source_names_same_content(inner, source, stt), + ), + _ => restore_source_names_same_content_inner(generated, source, stt), + } +} + +fn restore_source_names_same_content_inner( + generated: &LeanExpr, + source: &LeanExpr, + stt: &crate::compile::CompileState, +) -> LeanExpr { + match (generated.as_data(), source.as_data()) { + ( + ExprData::Const(gen_name, gen_lvls, _), + ExprData::Const(source_name, _, _), + ) if same_resolved_name_addr(gen_name, source_name, stt) => { + LeanExpr::cnst(source_name.clone(), gen_lvls.clone()) + }, + (ExprData::App(gen_f, gen_a, _), ExprData::App(source_f, source_a, _)) => { + LeanExpr::app( + restore_source_names_same_content(gen_f, source_f, stt), + restore_source_names_same_content(gen_a, source_a, stt), + ) + }, + ( + ExprData::ForallE(_, gen_dom, gen_body, gen_bi, _), + ExprData::ForallE(source_name, source_dom, source_body, _, _), + ) => LeanExpr::all( + source_name.clone(), + restore_source_names_same_content(gen_dom, source_dom, stt), + restore_source_names_same_content(gen_body, source_body, stt), + gen_bi.clone(), + ), + ( + ExprData::Lam(_, gen_dom, gen_body, gen_bi, _), + ExprData::Lam(source_name, source_dom, source_body, _, _), + ) => LeanExpr::lam( + source_name.clone(), + restore_source_names_same_content(gen_dom, source_dom, stt), + restore_source_names_same_content(gen_body, source_body, stt), + gen_bi.clone(), + ), + ( + ExprData::LetE(_, gen_ty, gen_val, gen_body, gen_nd, _), + ExprData::LetE(source_name, source_ty, source_val, source_body, _, _), + ) => LeanExpr::letE( + source_name.clone(), + restore_source_names_same_content(gen_ty, source_ty, stt), + restore_source_names_same_content(gen_val, source_val, stt), + restore_source_names_same_content(gen_body, source_body, stt), + *gen_nd, + ), + ( + ExprData::Proj(gen_name, gen_field, gen_val, _), + ExprData::Proj(source_name, source_field, source_val, _), + ) if gen_field == source_field + && same_resolved_name_addr(gen_name, source_name, stt) => + { + LeanExpr::proj( + source_name.clone(), + gen_field.clone(), + restore_source_names_same_content(gen_val, source_val, stt), + ) + }, + _ => generated.clone(), + } +} + +fn strip_mdata_ref(mut expr: &LeanExpr) -> &LeanExpr { + while let ExprData::Mdata(_, inner, _) = expr.as_data() { + expr = inner; + } + expr +} + +fn same_resolved_name_addr( + a: &Name, + b: &Name, + stt: &crate::compile::CompileState, +) -> bool { + if a == b { + return true; + } + let n2a = Some(&stt.name_to_addr); + let aux_n2a = Some(&stt.aux_name_to_addr); + resolve_lean_name_addr(a, n2a, aux_n2a) + == resolve_lean_name_addr(b, n2a, aux_n2a) +} + +/// Static version of `to_kexpr` that takes borrowed references. +/// +/// Identical to the closure-based `to_kexpr` in `get_level`, but as a +/// standalone function so it can be called from both `PreparedTC::new` +/// and `get_level_with_tc`. +pub(super) fn to_kexpr_static( + expr: &LeanExpr, + fvar_levels: &FxHashMap, + ctx_depth: usize, + param_names: &[Name], + stt: &crate::compile::CompileState, +) -> ix_kernel::expr::KExpr { + let n2a = Some(&stt.name_to_addr); + let aux_n2a = Some(&stt.aux_name_to_addr); + use ix_kernel::expr::KExpr; + use ix_kernel::id::KId; + use ix_kernel::level::KUniv; + + match expr.as_data() { + ExprData::Fvar(fname, _) => { + if let Some(&level) = fvar_levels.get(fname) { + KExpr::var((ctx_depth - level - 1) as u64, Name::anon()) + } else { + KExpr::sort(KUniv::zero()) + } + }, + ExprData::Bvar(idx, _) => KExpr::var(nat_to_u64(idx), Name::anon()), + ExprData::Sort(lvl, _) => { + KExpr::sort(lean_level_to_kuniv(lvl, param_names)) + }, + ExprData::Const(cname, us, _) => { + let addr = resolve_lean_name_addr(cname, n2a, aux_n2a); + let zid = KId::new(addr, cname.clone()); + let zus: Box<[KUniv]> = + us.iter().map(|u| lean_level_to_kuniv(u, param_names)).collect(); + KExpr::cnst(zid, zus) + }, + ExprData::App(f, a, _) => { + let kf = to_kexpr_static(f, fvar_levels, ctx_depth, param_names, stt); + let ka = to_kexpr_static(a, fvar_levels, ctx_depth, param_names, stt); + KExpr::app(kf, ka) + }, + ExprData::ForallE(binder_name, dom, body, bi, _) => { + let kd = to_kexpr_static(dom, fvar_levels, ctx_depth, param_names, stt); + let kb = + to_kexpr_static(body, fvar_levels, ctx_depth + 1, param_names, stt); + KExpr::all(binder_name.clone(), bi.clone(), kd, kb) + }, + ExprData::Lam(binder_name, dom, body, bi, _) => { + let kd = to_kexpr_static(dom, fvar_levels, ctx_depth, param_names, stt); + let kb = + to_kexpr_static(body, fvar_levels, ctx_depth + 1, param_names, stt); + KExpr::lam(binder_name.clone(), bi.clone(), kd, kb) + }, + ExprData::LetE(binder_name, ty, val, body, nd, _) => { + let kt = to_kexpr_static(ty, fvar_levels, ctx_depth, param_names, stt); + let kv = to_kexpr_static(val, fvar_levels, ctx_depth, param_names, stt); + let kb = + to_kexpr_static(body, fvar_levels, ctx_depth + 1, param_names, stt); + KExpr::let_(binder_name.clone(), kt, kv, kb, *nd) + }, + ExprData::Proj(pname, idx, e, _) => { + let addr = resolve_lean_name_addr(pname, n2a, aux_n2a); + let zid = KId::new(addr, pname.clone()); + let ke = to_kexpr_static(e, fvar_levels, ctx_depth, param_names, stt); + KExpr::prj(zid, nat_to_u64(idx), ke) + }, + ExprData::Lit(lit, _) => { + use ix_common::env::Literal; + match lit { + Literal::NatVal(n) => { + let addr = Address::hash(&nat_to_u64(n).to_le_bytes()); + KExpr::nat(n.clone(), addr) + }, + Literal::StrVal(s) => { + let addr = Address::hash(s.as_bytes()); + KExpr::str(s.clone(), addr) + }, + } + }, + ExprData::Mdata(_, inner, _) => { + to_kexpr_static(inner, fvar_levels, ctx_depth, param_names, stt) + }, + _ => KExpr::sort(KUniv::zero()), + } +} diff --git a/crates/compile/src/compile/aux_gen/kernel_bridge_tests.rs b/crates/compile/src/compile/aux_gen/kernel_bridge_tests.rs new file mode 100644 index 00000000..381fe6b2 --- /dev/null +++ b/crates/compile/src/compile/aux_gen/kernel_bridge_tests.rs @@ -0,0 +1,336 @@ +//! Differential and sharing regressions for the aux-generation kernel bridge. + +use super::*; +use crate::compile::CompileState; +use crate::compile::aux_gen::kernel_bridge_reference as reference; +use ix_common::env::{BinderInfo, DataValue, Literal}; +use ix_kernel::expr::{ExprData as KED, KExpr}; +use ix_kernel::id::KId; +use ix_kernel::level::KUniv; + +fn name(s: &str) -> Name { + Name::str(Name::anon(), s.to_owned()) +} + +fn diamond(mut leaf: LeanExpr, depth: usize) -> LeanExpr { + for _ in 0..depth { + leaf = LeanExpr::app(leaf.clone(), leaf); + } + leaf +} + +fn metadata(label: &str) -> Vec<(Name, DataValue)> { + vec![(name("tag"), DataValue::OfString(label.to_owned()))] +} + +fn assert_same_lean(actual: &LeanExpr, expected: &LeanExpr) { + // Full Lean digests include display names, binder info and mdata. Avoid + // formatting an exponentially large tree when a sharing regression fails. + assert_eq!(actual.get_hash(), expected.get_hash()); +} + +#[test] +fn kernel_bridge_matches_tree_reference_across_scopes() { + let stt = CompileState::new_empty(); + let params = [name("u"), name("v")]; + let fvars = FxHashMap::from_iter([(name("x"), 0), (name("y"), 1)]); + let u = Level::param(params[0].clone()); + let v = Level::param(params[1].clone()); + let shared = LeanExpr::app( + LeanExpr::cnst( + name("C"), + vec![Level::max(u.clone(), v.clone()), Level::imax(v, u)], + ), + LeanExpr::fvar(name("x")), + ); + let sort = LeanExpr::sort(Level::succ(Level::zero())); + let body = LeanExpr::app( + LeanExpr::proj(name("P"), Nat::from(2u64), shared.clone()), + LeanExpr::app(LeanExpr::bvar(Nat::from(0u64)), shared.clone()), + ); + let mut fixtures = vec![ + shared.clone(), + diamond(shared.clone(), 5), + LeanExpr::lit(Literal::NatVal(Nat::from(123u64))), + LeanExpr::lit(Literal::StrVal("bridge".to_owned())), + LeanExpr::fvar(name("unregistered")), + LeanExpr::mvar(name("unresolved")), + LeanExpr::mdata(metadata("outer"), shared.clone()), + ]; + for bi in [ + BinderInfo::Default, + BinderInfo::Implicit, + BinderInfo::StrictImplicit, + BinderInfo::InstImplicit, + ] { + fixtures.push(LeanExpr::all( + name("a"), + shared.clone(), + body.clone(), + bi.clone(), + )); + fixtures.push(LeanExpr::lam(name("b"), sort.clone(), body.clone(), bi)); + } + for non_dep in [false, true] { + fixtures.push(LeanExpr::letE( + name("c"), + sort.clone(), + shared.clone(), + body.clone(), + non_dep, + )); + } + for depth in [2, 3] { + for source in &fixtures { + let old = + reference::to_kexpr_static(source, &fvars, depth, ¶ms, &stt); + let new = to_kexpr_static(source, &fvars, depth, ¶ms, &stt); + assert_eq!(new, old); + let expected = reference::kexpr_to_lean(&old, depth, &fvars, 0, ¶ms); + assert_same_lean( + &reference::kexpr_to_lean(&new, depth, &fvars, 0, ¶ms), + &expected, + ); + assert_same_lean( + &kexpr_to_lean(&old, depth, &fvars, 0, ¶ms), + &expected, + ); + assert_same_lean( + &kexpr_to_lean(&new, depth, &fvars, 0, ¶ms), + &expected, + ); + } + } +} + +#[test] +fn kernel_bridge_cache_keys_include_binder_depth() { + let stt = CompileState::new_empty(); + let fvars = FxHashMap::from_iter([(name("x"), 0)]); + let x = LeanExpr::fvar(name("x")); + let source = LeanExpr::app( + x.clone(), + LeanExpr::lam( + name("bound"), + LeanExpr::sort(Level::zero()), + x, + BinderInfo::Default, + ), + ); + let ingressed = to_kexpr_static(&source, &fvars, 1, &[], &stt); + let KED::App(outer, lam, _) = ingressed.data() else { panic!("app") }; + let KED::Lam(_, _, _, inner, _) = lam.data() else { panic!("lam") }; + assert!(matches!(outer.data(), KED::Var(0, ..))); + assert!(matches!(inner.data(), KED::Var(1, ..))); + assert_same_lean(&kexpr_to_lean(&ingressed, 1, &fvars, 0, &[]), &source); + + // One kernel node is free at the root but bound under the lambda. + let var = KExpr::var(0, Name::anon()); + let kernel = KExpr::app( + var.clone(), + KExpr::lam( + name("bound"), + BinderInfo::Default, + KExpr::sort(KUniv::zero()), + var, + ), + ); + let expected = reference::kexpr_to_lean(&kernel, 1, &fvars, 0, &[]); + assert_same_lean(&kexpr_to_lean(&kernel, 1, &fvars, 0, &[]), &expected); + let ExprData::App(outer, lam, _) = expected.as_data() else { panic!("app") }; + let ExprData::Lam(_, _, inner, _, _) = lam.as_data() else { panic!("lam") }; + assert!(matches!(outer.as_data(), ExprData::Fvar(..))); + assert!(matches!(inner.as_data(), ExprData::Bvar(..))); +} + +#[test] +fn kernel_egress_distinguishes_metadata_nodes_with_the_same_uid() { + let address = Address::hash(b"alias class"); + let a = KExpr::::cnst_mdata( + KId::new(address.clone(), name("A")), + Box::new([]), + vec![metadata("outer-A"), metadata("inner-A")], + ); + let mut b_info = a.info().clone(); + b_info.mdata = vec![metadata("B")]; + let b = + KExpr::new(KED::Const(KId::new(address, name("B")), Box::new([]), b_info)); + assert_eq!(a.hash_key(), b.hash_key()); + let kernel = KExpr::app(a, b); + let fvars = FxHashMap::default(); + let expected = reference::kexpr_to_lean(&kernel, 0, &fvars, 0, &[]); + assert_same_lean(&kexpr_to_lean(&kernel, 0, &fvars, 0, &[]), &expected); + let ExprData::App(a, b, _) = expected.as_data() else { panic!("app") }; + assert_ne!(a.get_hash(), b.get_hash()); + assert_same_lean( + a, + &LeanExpr::mdata( + metadata("outer-A"), + LeanExpr::mdata(metadata("inner-A"), LeanExpr::cnst(name("A"), vec![])), + ), + ); +} + +#[test] +fn kernel_bridge_caches_do_not_outlive_their_context() { + let stt = CompileState::new_empty(); + let source = LeanExpr::app( + LeanExpr::cnst(name("C"), vec![Level::param(name("u"))]), + LeanExpr::fvar(name("x")), + ); + for (i, params) in + [[name("u"), name("v")], [name("v"), name("u")]].iter().enumerate() + { + let address = Address::hash(&[i as u8]); + stt.name_to_addr.insert(name("C"), address.clone()); + let fvars = FxHashMap::from_iter([(name("x"), i)]); + let kernel = to_kexpr_static(&source, &fvars, 2, params, &stt); + let KED::App(c, x, _) = kernel.data() else { panic!("app") }; + let KED::Const(id, levels, _) = c.data() else { panic!("const") }; + assert_eq!(id.addr, address); + assert_eq!( + levels[0], + lean_level_to_kuniv(&Level::param(name("u")), params) + ); + assert!(matches!(x.data(), KED::Var(idx, ..) if *idx == (1 - i) as u64)); + assert_same_lean(&kexpr_to_lean(&kernel, 2, &fvars, 0, params), &source); + } +} + +#[test] +fn source_restoration_matches_reference_and_preserves_occurrence_names() { + let stt = CompileState::new_empty(); + for alias in ["G", "A", "B"] { + stt.name_to_addr.insert(name(alias), Address::hash(b"same declaration")); + } + let g = LeanExpr::cnst(name("G"), vec![Level::zero()]); + let a = LeanExpr::cnst(name("A"), vec![Level::succ(Level::zero())]); + let b = LeanExpr::cnst(name("B"), vec![]); + let generated = LeanExpr::app(g.clone(), g.clone()); + let source = LeanExpr::app(a.clone(), b.clone()); + let restored = restore_source_names_same_content(&generated, &source, &stt); + assert_same_lean( + &restored, + &reference::restore_source_names_same_content(&generated, &source, &stt), + ); + assert_same_lean( + &restored, + &LeanExpr::app( + LeanExpr::cnst(name("A"), vec![Level::zero()]), + LeanExpr::cnst(name("B"), vec![Level::zero()]), + ), + ); + // Also cover binders, non-dependency flags, projection indices, mismatches + // and independent source/generated metadata layers. + let pairs = [ + ( + LeanExpr::all(name("g"), g.clone(), g.clone(), BinderInfo::Implicit), + LeanExpr::all(name("s"), a.clone(), b.clone(), BinderInfo::Default), + ), + ( + LeanExpr::lam(name("g"), g.clone(), g.clone(), BinderInfo::InstImplicit), + LeanExpr::lam(name("s"), a.clone(), b.clone(), BinderInfo::Default), + ), + ( + LeanExpr::letE(name("g"), g.clone(), g.clone(), g.clone(), true), + LeanExpr::letE(name("s"), a.clone(), b.clone(), a.clone(), false), + ), + ( + LeanExpr::proj(name("G"), Nat::from(0u64), g.clone()), + LeanExpr::proj(name("A"), Nat::from(0u64), b.clone()), + ), + ( + LeanExpr::proj(name("G"), Nat::from(0u64), g.clone()), + LeanExpr::proj(name("A"), Nat::from(1u64), b.clone()), + ), + ( + LeanExpr::mdata(metadata("generated"), generated), + LeanExpr::mdata(metadata("source"), source), + ), + (g.clone(), LeanExpr::cnst(name("unrelated"), vec![])), + (g.clone(), LeanExpr::sort(Level::zero())), + (g.clone(), LeanExpr::mdata(metadata("source"), g)), + ]; + for (generated, source) in pairs { + assert_same_lean( + &restore_source_names_same_content(&generated, &source, &stt), + &reference::restore_source_names_same_content(&generated, &source, &stt), + ); + } +} + +#[test] +fn kernel_bridge_preserves_a_trillion_path_dag() { + let stt = CompileState::new_empty(); + for alias in ["A", "B"] { + stt.name_to_addr.insert(name(alias), Address::hash(b"shared leaf")); + } + // 41 unique nodes, but 2^40 leaf occurrences if expanded as a tree. + let depth = 40; + let source = diamond(LeanExpr::cnst(name("A"), vec![]), depth); + let aliases = diamond(LeanExpr::cnst(name("B"), vec![]), depth); + let fvars = FxHashMap::default(); + let mut ingress_cache = FxHashMap::default(); + let kernel = + to_kexpr_cached(&source, &fvars, 0, &[], &stt, &mut ingress_cache); + assert_eq!(ingress_cache.len(), depth + 1); + let mut cursor = &kernel; + for _ in 0..depth { + let KED::App(f, a, _) = cursor.data() else { panic!("app") }; + assert!(std::ptr::eq(f.data(), a.data())); + cursor = f; + } + let mut egress_cache = FxHashMap::default(); + let generated = + kexpr_to_lean_cached(&kernel, 0, &fvars, 0, &[], &mut egress_cache); + assert_eq!(egress_cache.len(), depth + 1); + assert_same_lean(&generated, &source); + let mut restore_cache = FxHashMap::default(); + let restored = + restore_source_names_cached(&generated, &aliases, &stt, &mut restore_cache); + assert_eq!(restore_cache.len(), depth + 1); + assert_same_lean(&restored, &aliases); + for root in [&generated, &restored] { + let mut cursor = root; + for _ in 0..depth { + let ExprData::App(f, a, _) = cursor.as_data() else { panic!("app") }; + assert!(std::ptr::eq(f.as_data(), a.as_data())); + cursor = f; + } + } + let mut refs = FxHashSet::default(); + collect_lean_const_refs(&source, &mut refs); + assert_eq!(refs, FxHashSet::from_iter([name("A")])); +} + +#[test] +#[ignore = "manual release-mode comparison against the pre-memoization bridge"] +fn kernel_bridge_shared_dag_benchmark() { + use std::time::Instant; + let stt = CompileState::new_empty(); + for alias in ["A", "B"] { + stt.name_to_addr.insert(name(alias), Address::hash(b"shared leaf")); + } + let fvars = FxHashMap::default(); + for depth in [12, 16, 18] { + let source = diamond(LeanExpr::cnst(name("A"), vec![]), depth); + let aliases = diamond(LeanExpr::cnst(name("B"), vec![]), depth); + let start = Instant::now(); + let old_k = reference::to_kexpr_static(&source, &fvars, 0, &[], &stt); + let old_l = reference::kexpr_to_lean(&old_k, 0, &fvars, 0, &[]); + let old_r = + reference::restore_source_names_same_content(&old_l, &aliases, &stt); + let old_time = start.elapsed(); + let start = Instant::now(); + let new_k = to_kexpr_static(&source, &fvars, 0, &[], &stt); + let new_l = kexpr_to_lean(&new_k, 0, &fvars, 0, &[]); + let new_r = restore_source_names_same_content(&new_l, &aliases, &stt); + let new_time = start.elapsed(); + assert_same_lean(&new_r, &old_r); + eprintln!( + "bridge DAG: {} unique nodes / {} leaf paths: old={old_time:?} new={new_time:?}", + depth + 1, + 1usize << depth + ); + } +} diff --git a/crates/compile/src/compile/aux_gen/source_name_hints.rs b/crates/compile/src/compile/aux_gen/source_name_hints.rs new file mode 100644 index 00000000..2f25752f --- /dev/null +++ b/crates/compile/src/compile/aux_gen/source_name_hints.rs @@ -0,0 +1,406 @@ +//! DAG-aware source-name restoration after a real WHNF reduction. +//! +//! One pass owns all memoization and a fixed view of referenced addresses. +//! Collection remains preorder/left-to-right/first-wins; restoration starts +//! only after the hint table is complete. Eligibility is still App/Proj with +//! NO Bvar node anywhere, including under binders (canonicity §10.5). + +use super::*; +use ix_common::env::Literal; +use ix_kernel::env::InternTable; +use ix_kernel::level::KUniv; +use ix_kernel::mode::Anon; +use std::collections::HashMap; +use std::hash::BuildHasher; +use std::sync::Arc; + +fn ptr(expr: &LeanExpr) -> usize { + std::ptr::from_ref(expr.as_data()).addr() +} + +fn same_node(a: &LeanExpr, b: &LeanExpr) -> bool { + Arc::ptr_eq(&a.0, &b.0) +} + +/// Capture each referenced name once, preserving primary/aux/fallback +/// precedence. No live DashMap reads occur during collection or restoration. +/// This is a pass-local resolution view, not an atomic snapshot of all global +/// compiler state; later passes capture a fresh view after further publication. +fn capture_addresses( + source: &LeanExpr, + generated: &LeanExpr, + stt: &crate::compile::CompileState, +) -> FxHashMap { + let mut addresses = FxHashMap::default(); + let mut seen = FxHashSet::default(); + let mut stack = vec![generated, source]; + while let Some(expr) = stack.pop() { + if !seen.insert(ptr(expr)) { + continue; + } + match expr.as_data() { + ExprData::Const(name, _, _) | ExprData::Proj(name, _, _, _) => { + addresses.entry(name.clone()).or_insert_with(|| { + resolve_lean_name_addr( + name, + Some(&stt.name_to_addr), + Some(&stt.aux_name_to_addr), + ) + }); + }, + _ => {}, + } + match expr.as_data() { + ExprData::App(f, a, _) => stack.extend([a, f]), + ExprData::ForallE(_, d, b, _, _) | ExprData::Lam(_, d, b, _, _) => { + stack.extend([b, d]); + }, + ExprData::LetE(_, t, v, b, _, _) => stack.extend([b, v, t]), + ExprData::Proj(_, _, v, _) | ExprData::Mdata(_, v, _) => stack.push(v), + _ => {}, + } + } + addresses +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +struct ContentId(usize); + +/// Exact name-erased structure of to_kexpr_static's result. Children and +/// normalized universes have exact interned identities; hashing only chooses +/// a bucket, and Eq confirms the entire shallow key. No lossy digest is used +/// as evidence of equality. Literal VALUES are included, not just blob hashes. +/// These ids never enter a kernel cache or serialized output. +#[derive(Debug, PartialEq, Eq, Hash)] +enum ContentKey { + Var(u64), + Sort(u64), + Const(Address, Box<[u64]>), + App(ContentId, ContentId), + Lam(ContentId, ContentId), + All(ContentId, ContentId), + Let(ContentId, ContentId, ContentId, bool), + Prj(Address, u64, ContentId), + Nat(Nat), + Str(String), +} + +struct ContentTable<'a, S = rustc_hash::FxBuildHasher> { + addresses: &'a FxHashMap, + fvars: &'a FxHashMap, + params: &'a [Name], + shapes: HashMap, + converted: FxHashMap<(usize, usize), ContentId>, + levels: FxHashMap>, + univs: InternTable, +} + +impl<'a, S: BuildHasher + Default> ContentTable<'a, S> { + fn new( + addresses: &'a FxHashMap, + fvars: &'a FxHashMap, + params: &'a [Name], + ) -> Self { + Self { + addresses, + fvars, + params, + shapes: HashMap::default(), + converted: FxHashMap::default(), + levels: FxHashMap::default(), + univs: InternTable::new(), + } + } + + fn intern(&mut self, key: ContentKey) -> ContentId { + let next = ContentId(self.shapes.len()); + *self.shapes.entry(key).or_insert(next) + } + + /// Use the SAME mk_max/mk_imax reductions as lean_level_to_kuniv, then + /// intern normalized levels structurally. Parameter display names are + /// irrelevant to hint equality. These private universe nodes never egress. + fn level(&mut self, level: &Level) -> KUniv { + let ptr = std::ptr::from_ref(level.as_data()).addr(); + if let Some(result) = self.levels.get(&ptr) { + return result.clone(); + } + let result = match level.as_data() { + LevelData::Zero(_) => KUniv::zero(), + LevelData::Succ(a, _) => KUniv::succ(self.level(a)), + LevelData::Max(a, b, _) => { + let a = self.level(a); + let b = self.level(b); + KUniv::max(a, b) + }, + LevelData::Imax(a, b, _) => { + let a = self.level(a); + let b = self.level(b); + KUniv::imax(a, b) + }, + LevelData::Param(name, _) => { + let idx = + self.params.iter().position(|n| n == name).unwrap_or_else(|| { + panic!( + "unknown level param `{}` in source-name hints", + name.pretty() + ) + }); + KUniv::param(idx as u64, ()) + }, + LevelData::Mvar(name, _) => { + panic!( + "unexpected level metavariable `{}` in source-name hints", + name.pretty() + ); + }, + }; + let result = self.univs.intern_univ(result); + self.levels.insert(ptr, result.clone()); + result + } + + fn content(&mut self, expr: &LeanExpr, depth: usize) -> ContentId { + let memo_key = (ptr(expr), depth); + if let Some(id) = self.converted.get(&memo_key) { + return *id; + } + let key = match expr.as_data() { + ExprData::Bvar(idx, _) => ContentKey::Var(nat_to_u64(idx)), + ExprData::Fvar(name, _) => match self.fvars.get(name) { + Some(level) => ContentKey::Var((depth - level - 1) as u64), + None => { + let zero = self.univs.intern_univ(KUniv::zero()); + ContentKey::Sort(*zero.addr()) + }, + }, + ExprData::Mvar(..) => { + let zero = self.univs.intern_univ(KUniv::zero()); + ContentKey::Sort(*zero.addr()) + }, + ExprData::Sort(level, _) => ContentKey::Sort(*self.level(level).addr()), + ExprData::Const(name, levels, _) => { + let address = self.addresses[name].clone(); + let levels = levels.iter().map(|u| *self.level(u).addr()).collect(); + ContentKey::Const(address, levels) + }, + ExprData::App(f, a, _) => { + ContentKey::App(self.content(f, depth), self.content(a, depth)) + }, + ExprData::Lam(_, ty, body, _, _) => { + ContentKey::Lam(self.content(ty, depth), self.content(body, depth + 1)) + }, + ExprData::ForallE(_, ty, body, _, _) => { + ContentKey::All(self.content(ty, depth), self.content(body, depth + 1)) + }, + ExprData::LetE(_, ty, val, body, nd, _) => ContentKey::Let( + self.content(ty, depth), + self.content(val, depth), + self.content(body, depth + 1), + *nd, + ), + ExprData::Proj(name, field, val, _) => ContentKey::Prj( + self.addresses[name].clone(), + nat_to_u64(field), + self.content(val, depth), + ), + ExprData::Lit(Literal::NatVal(n), _) => ContentKey::Nat(n.clone()), + ExprData::Lit(Literal::StrVal(s), _) => ContentKey::Str(s.clone()), + ExprData::Mdata(_, inner, _) => { + let id = self.content(inner, depth); + self.converted.insert(memo_key, id); + return id; + }, + }; + let id = self.intern(key); + self.converted.insert(memo_key, id); + id + } +} + +struct Pass<'a> { + content: ContentTable<'a>, + depth: usize, + has_bvar: FxHashMap, + collected: FxHashSet, + restored: FxHashMap, +} + +impl<'a> Pass<'a> { + fn new( + addresses: &'a FxHashMap, + fvars: &'a FxHashMap, + depth: usize, + params: &'a [Name], + ) -> Self { + Self { + content: ContentTable::new(addresses, fvars, params), + depth, + has_bvar: FxHashMap::default(), + collected: FxHashSet::default(), + restored: FxHashMap::default(), + } + } + + /// This is "contains ANY Bvar", not a loose/free-variable range test. + fn has_bvar(&mut self, expr: &LeanExpr) -> bool { + let key = ptr(expr); + if let Some(result) = self.has_bvar.get(&key) { + return *result; + } + let result = match expr.as_data() { + ExprData::Bvar(..) => true, + ExprData::App(f, a, _) => self.has_bvar(f) || self.has_bvar(a), + ExprData::ForallE(_, d, b, _, _) | ExprData::Lam(_, d, b, _, _) => { + self.has_bvar(d) || self.has_bvar(b) + }, + ExprData::LetE(_, t, v, b, _, _) => { + self.has_bvar(t) || self.has_bvar(v) || self.has_bvar(b) + }, + ExprData::Proj(_, _, v, _) | ExprData::Mdata(_, v, _) => self.has_bvar(v), + _ => false, + }; + self.has_bvar.insert(key, result); + result + } + + fn candidate(&mut self, expr: &LeanExpr) -> bool { + matches!(expr.as_data(), ExprData::App(..) | ExprData::Proj(..)) + && !self.has_bvar(expr) + } + + fn collect( + &mut self, + source: &LeanExpr, + hints: &mut FxHashMap, + ) { + let key = ptr(source); + if self.collected.contains(&key) { + return; + } + if self.candidate(source) { + let id = self.content.content(source, self.depth); + hints.entry(id).or_insert_with(|| source.clone()); + } + // Do not skip children merely because this node's hint slot was occupied. + // Mark the exact input node done only after its entire subtree was offered. + match source.as_data() { + ExprData::App(f, a, _) => { + self.collect(f, hints); + self.collect(a, hints); + }, + ExprData::ForallE(_, d, b, _, _) | ExprData::Lam(_, d, b, _, _) => { + self.collect(d, hints); + self.collect(b, hints); + }, + ExprData::LetE(_, t, v, b, _, _) => { + self.collect(t, hints); + self.collect(v, hints); + self.collect(b, hints); + }, + ExprData::Proj(_, _, v, _) | ExprData::Mdata(_, v, _) => { + self.collect(v, hints) + }, + _ => {}, + } + self.collected.insert(key); + } + + fn restore( + &mut self, + generated: &LeanExpr, + hints: &FxHashMap, + ) -> LeanExpr { + let key = ptr(generated); + if let Some(result) = self.restored.get(&key) { + return result.clone(); + } + if self.candidate(generated) { + let id = self.content.content(generated, self.depth); + if let Some(source) = hints.get(&id) { + self.restored.insert(key, source.clone()); + return source.clone(); + } + } + let result = match generated.as_data() { + ExprData::App(f, a, _) => { + let cf = self.restore(f, hints); + let ca = self.restore(a, hints); + if same_node(f, &cf) && same_node(a, &ca) { + generated.clone() + } else { + LeanExpr::app(cf, ca) + } + }, + ExprData::ForallE(n, d, b, bi, _) => { + let cd = self.restore(d, hints); + let cb = self.restore(b, hints); + if same_node(d, &cd) && same_node(b, &cb) { + generated.clone() + } else { + LeanExpr::all(n.clone(), cd, cb, bi.clone()) + } + }, + ExprData::Lam(n, d, b, bi, _) => { + let cd = self.restore(d, hints); + let cb = self.restore(b, hints); + if same_node(d, &cd) && same_node(b, &cb) { + generated.clone() + } else { + LeanExpr::lam(n.clone(), cd, cb, bi.clone()) + } + }, + ExprData::LetE(n, t, v, b, nd, _) => { + let ct = self.restore(t, hints); + let cv = self.restore(v, hints); + let cb = self.restore(b, hints); + if same_node(t, &ct) && same_node(v, &cv) && same_node(b, &cb) { + generated.clone() + } else { + LeanExpr::letE(n.clone(), ct, cv, cb, *nd) + } + }, + ExprData::Proj(n, i, v, _) => { + let cv = self.restore(v, hints); + if same_node(v, &cv) { + generated.clone() + } else { + LeanExpr::proj(n.clone(), i.clone(), cv) + } + }, + ExprData::Mdata(kvs, v, _) => { + let cv = self.restore(v, hints); + if same_node(v, &cv) { + generated.clone() + } else { + LeanExpr::mdata(kvs.clone(), cv) + } + }, + _ => generated.clone(), + }; + self.restored.insert(key, result.clone()); + result + } +} + +/// Roots stay borrowed until all pointer-keyed caches are dropped. Only +/// original source/generated descendants become memo keys, never temporaries. +/// The fixed context and finalized hint map cannot change during restoration. +pub(super) fn restore( + generated: &LeanExpr, + source: &LeanExpr, + fvars: &FxHashMap, + depth: usize, + params: &[Name], + stt: &crate::compile::CompileState, +) -> LeanExpr { + let addresses = capture_addresses(source, generated, stt); + let mut pass = Pass::new(&addresses, fvars, depth, params); + let mut hints = FxHashMap::default(); + pass.collect(source, &mut hints); + pass.restore(generated, &hints) +} + +#[cfg(test)] +#[path = "source_name_hints_tests.rs"] +mod tests; diff --git a/crates/compile/src/compile/aux_gen/source_name_hints_reference.rs b/crates/compile/src/compile/aux_gen/source_name_hints_reference.rs new file mode 100644 index 00000000..77ab0bed --- /dev/null +++ b/crates/compile/src/compile/aux_gen/source_name_hints_reference.rs @@ -0,0 +1,393 @@ +//! Frozen pre-cache source-name hint traversal for differential tests. +//! The legacy u64 key is intentionally retained here, including its collision +//! behavior; collision regressions test the new exact matcher separately. + +use super::*; + +pub(super) fn source_name_hint_candidate(expr: &LeanExpr) -> bool { + matches!(expr.as_data(), ExprData::App(..) | ExprData::Proj(..)) +} + +/// Name-erased structural content key for the source-name hint map. +/// +/// Mirrors the equivalence of the Lean pipeline's `Ix.Tc.KExpr` content +/// addresses (`toKexprStatic ... |>.addr` in `Ix/AuxGen/Kernel.lean`) and +/// of the kernel's `ExprKey`/`structural_eq`: display names, binder +/// names, binder infos, and mdata are excluded; `Const`/`Prj` contribute +/// their resolved content address, universes their index structure. Two +/// spellings of one alias pair (`Paths V` / `Symmetrify V`) therefore +/// agree on this key — which is the whole point of the hint map. +/// +/// `KExpr::hash_key()` is NOT usable here: it is the intern-assigned uid, +/// fresh for every un-interned construction, and `to_kexpr_static` does +/// not intern — so the collect-time and restore-time keys of two +/// content-equal subterms never matched, and the hint map restored +/// nothing (the Mathlib `Quiver.FreeGroupoid.redStep` metadata +/// divergence, canonicity §10.5). +pub(super) fn kexpr_content_key(e: &ix_kernel::expr::KExpr) -> u64 { + use std::hash::Hasher; + let mut h = rustc_hash::FxHasher::default(); + kexpr_content_hash(e, &mut h); + h.finish() +} + +pub(super) fn kuniv_content_hash( + u: &ix_kernel::level::KUniv, + h: &mut rustc_hash::FxHasher, +) { + use ix_kernel::level::UnivData as UD; + use std::hash::Hasher; + match u.data() { + UD::Zero(_) => h.write_u8(0), + UD::Succ(a, _) => { + h.write_u8(1); + kuniv_content_hash(a, h); + }, + UD::Max(a, b, _) => { + h.write_u8(2); + kuniv_content_hash(a, h); + kuniv_content_hash(b, h); + }, + UD::IMax(a, b, _) => { + h.write_u8(3); + kuniv_content_hash(a, h); + kuniv_content_hash(b, h); + }, + UD::Param(idx, _, _) => { + h.write_u8(4); + h.write_u64(*idx); + }, + } +} + +pub(super) fn kexpr_content_hash( + e: &ix_kernel::expr::KExpr, + h: &mut rustc_hash::FxHasher, +) { + use ix_kernel::expr::ExprData as KED; + use std::hash::Hasher; + match e.data() { + KED::Var(i, _, _) => { + h.write_u8(0); + h.write_u64(*i); + }, + KED::FVar(id, _, _) => { + h.write_u8(1); + h.write_u64(id.0); + }, + KED::Sort(u, _) => { + h.write_u8(2); + kuniv_content_hash(u, h); + }, + KED::Const(id, us, _) => { + h.write_u8(3); + h.write(id.addr.as_bytes()); + h.write_u64(us.len() as u64); + for u in us.iter() { + kuniv_content_hash(u, h); + } + }, + KED::App(f, a, _) => { + h.write_u8(4); + kexpr_content_hash(f, h); + kexpr_content_hash(a, h); + }, + KED::Lam(_, _, t, b, _) => { + h.write_u8(5); + kexpr_content_hash(t, h); + kexpr_content_hash(b, h); + }, + KED::All(_, _, t, b, _) => { + h.write_u8(6); + kexpr_content_hash(t, h); + kexpr_content_hash(b, h); + }, + KED::Let(_, t, v, b, nd, _) => { + h.write_u8(7); + h.write_u8(u8::from(*nd)); + kexpr_content_hash(t, h); + kexpr_content_hash(v, h); + kexpr_content_hash(b, h); + }, + KED::Prj(id, f, v, _) => { + h.write_u8(8); + h.write(id.addr.as_bytes()); + h.write_u64(*f); + kexpr_content_hash(v, h); + }, + KED::Nat(_, ba, _) => { + h.write_u8(9); + h.write(ba.as_bytes()); + }, + KED::Str(_, ba, _) => { + h.write_u8(10); + h.write(ba.as_bytes()); + }, + } +} + +/// Collect source-shaped subterms that WHNF may copy into a reduct. +/// +/// Keys use the kernel content hash so alpha-collapsed aliases like +/// `CategoryTheory.Paths V` and `Quiver.Symmetrify V` line up, while values +/// keep the Lean display names from the caller. We skip BVar-containing terms: +/// WHNF may lift copied arguments under freshly-exposed binders, so matching +/// those by raw de Bruijn indices would be unstable. +pub(super) fn collect_lean_source_name_hints( + source: &LeanExpr, + fvar_levels: &FxHashMap, + depth: usize, + param_names: &[Name], + stt: &crate::compile::CompileState, + out: &mut FxHashMap, +) { + if source_name_hint_candidate(source) && !expr_has_bvar(source) { + let key = kexpr_content_key(&to_kexpr_static( + source, + fvar_levels, + depth, + param_names, + stt, + )); + out.entry(key).or_insert_with(|| source.clone()); + } + + match source.as_data() { + ExprData::Mdata(_, inner, _) => collect_lean_source_name_hints( + inner, + fvar_levels, + depth, + param_names, + stt, + out, + ), + ExprData::App(f, a, _) => { + collect_lean_source_name_hints( + f, + fvar_levels, + depth, + param_names, + stt, + out, + ); + collect_lean_source_name_hints( + a, + fvar_levels, + depth, + param_names, + stt, + out, + ); + }, + ExprData::ForallE(_, d, b, _, _) | ExprData::Lam(_, d, b, _, _) => { + collect_lean_source_name_hints( + d, + fvar_levels, + depth, + param_names, + stt, + out, + ); + collect_lean_source_name_hints( + b, + fvar_levels, + depth, + param_names, + stt, + out, + ); + }, + ExprData::LetE(_, t, v, b, _, _) => { + collect_lean_source_name_hints( + t, + fvar_levels, + depth, + param_names, + stt, + out, + ); + collect_lean_source_name_hints( + v, + fvar_levels, + depth, + param_names, + stt, + out, + ); + collect_lean_source_name_hints( + b, + fvar_levels, + depth, + param_names, + stt, + out, + ); + }, + ExprData::Proj(_, _, v, _) => collect_lean_source_name_hints( + v, + fvar_levels, + depth, + param_names, + stt, + out, + ), + _ => {}, + } +} + +/// Restore source spellings for copied subterms after a real WHNF reduction. +/// +/// This is intentionally subterm-based rather than whole-expression based: +/// unfolding a reducible alias such as `HomRel (Paths (Symmetrify V))` should +/// keep the expanded `∀` telescope, but the repeated argument subterms inside +/// that telescope should retain the caller's `Symmetrify` spelling instead of +/// whichever same-address alias the kernel cache/intern table already held. +pub(super) fn restore_lean_source_name_hints( + generated: &LeanExpr, + fvar_levels: &FxHashMap, + depth: usize, + param_names: &[Name], + stt: &crate::compile::CompileState, + hints: &FxHashMap, +) -> LeanExpr { + if source_name_hint_candidate(generated) && !expr_has_bvar(generated) { + let key = kexpr_content_key(&to_kexpr_static( + generated, + fvar_levels, + depth, + param_names, + stt, + )); + if let Some(source) = hints.get(&key) { + return source.clone(); + } + } + + match generated.as_data() { + ExprData::App(f, a, _) => LeanExpr::app( + restore_lean_source_name_hints( + f, + fvar_levels, + depth, + param_names, + stt, + hints, + ), + restore_lean_source_name_hints( + a, + fvar_levels, + depth, + param_names, + stt, + hints, + ), + ), + ExprData::ForallE(n, d, b, bi, _) => LeanExpr::all( + n.clone(), + restore_lean_source_name_hints( + d, + fvar_levels, + depth, + param_names, + stt, + hints, + ), + restore_lean_source_name_hints( + b, + fvar_levels, + depth, + param_names, + stt, + hints, + ), + bi.clone(), + ), + ExprData::Lam(n, d, b, bi, _) => LeanExpr::lam( + n.clone(), + restore_lean_source_name_hints( + d, + fvar_levels, + depth, + param_names, + stt, + hints, + ), + restore_lean_source_name_hints( + b, + fvar_levels, + depth, + param_names, + stt, + hints, + ), + bi.clone(), + ), + ExprData::LetE(n, t, v, b, nd, _) => LeanExpr::letE( + n.clone(), + restore_lean_source_name_hints( + t, + fvar_levels, + depth, + param_names, + stt, + hints, + ), + restore_lean_source_name_hints( + v, + fvar_levels, + depth, + param_names, + stt, + hints, + ), + restore_lean_source_name_hints( + b, + fvar_levels, + depth, + param_names, + stt, + hints, + ), + *nd, + ), + ExprData::Proj(n, i, v, _) => LeanExpr::proj( + n.clone(), + i.clone(), + restore_lean_source_name_hints( + v, + fvar_levels, + depth, + param_names, + stt, + hints, + ), + ), + ExprData::Mdata(kvs, v, _) => LeanExpr::mdata( + kvs.clone(), + restore_lean_source_name_hints( + v, + fvar_levels, + depth, + param_names, + stt, + hints, + ), + ), + _ => generated.clone(), + } +} + +pub(super) fn expr_has_bvar(expr: &LeanExpr) -> bool { + match expr.as_data() { + ExprData::Bvar(..) => true, + ExprData::App(f, a, _) => expr_has_bvar(f) || expr_has_bvar(a), + ExprData::ForallE(_, d, b, _, _) | ExprData::Lam(_, d, b, _, _) => { + expr_has_bvar(d) || expr_has_bvar(b) + }, + ExprData::LetE(_, t, v, b, _, _) => { + expr_has_bvar(t) || expr_has_bvar(v) || expr_has_bvar(b) + }, + ExprData::Proj(_, _, v, _) | ExprData::Mdata(_, v, _) => expr_has_bvar(v), + _ => false, + } +} diff --git a/crates/compile/src/compile/aux_gen/source_name_hints_tests.rs b/crates/compile/src/compile/aux_gen/source_name_hints_tests.rs new file mode 100644 index 00000000..6e3ec22c --- /dev/null +++ b/crates/compile/src/compile/aux_gen/source_name_hints_tests.rs @@ -0,0 +1,353 @@ +//! Differential semantics, collision handling, context, and DAG regressions. + +use super::super::source_name_hints_reference as reference; +use super::*; +use crate::compile::CompileState; +use ix_common::env::DataValue; +use std::hash::{BuildHasherDefault, Hasher}; + +fn name(s: &str) -> Name { + Name::str(Name::anon(), s.to_owned()) +} + +fn state() -> CompileState { + let stt = CompileState::new_empty(); + for alias in ["A", "B"] { + stt.name_to_addr.insert(name(alias), Address::hash(b"alias class")); + } + stt.aux_name_to_addr.insert(name("P"), Address::hash(b"projection class")); + stt.name_to_addr.insert(name("Q"), Address::hash(b"projection class")); + stt +} + +fn tag(s: &str) -> Vec<(Name, DataValue)> { + vec![(name("tag"), DataValue::OfString(s.to_owned()))] +} + +fn application(alias: &str) -> LeanExpr { + LeanExpr::app(LeanExpr::cnst(name(alias), vec![]), LeanExpr::fvar(name("x"))) +} + +fn diamond(mut leaf: LeanExpr, depth: usize) -> LeanExpr { + for _ in 0..depth { + leaf = LeanExpr::app(leaf.clone(), leaf); + } + leaf +} + +fn fixtures(alias: &str) -> Vec { + let u = Level::param(name("u")); + let v = Level::param(name("v")); + let shared = application(alias); + let mut roots = vec![ + LeanExpr::bvar(Nat::from(0u64)), + LeanExpr::fvar(name("x")), + LeanExpr::fvar(name("unregistered")), + LeanExpr::mvar(name("unresolved")), + LeanExpr::sort(Level::zero()), + LeanExpr::cnst( + name(alias), + vec![u.clone(), Level::imax(v.clone(), u.clone())], + ), + shared.clone(), + diamond(shared.clone(), 4), + LeanExpr::proj( + name(if alias == "A" { "P" } else { "Q" }), + Nat::from(1u64), + shared.clone(), + ), + LeanExpr::lit(Literal::NatVal(Nat::from(42u64))), + LeanExpr::lit(Literal::StrVal("a string".to_owned())), + LeanExpr::mdata(tag(alias), shared.clone()), + LeanExpr::mdata(tag(alias), LeanExpr::mdata(tag("inner"), shared.clone())), + ]; + for level in [ + u.clone(), + Level::succ(u.clone()), + Level::max(u.clone(), Level::zero()), + Level::max(u.clone(), v.clone()), + Level::max(v.clone(), u.clone()), + Level::imax(u, Level::zero()), + ] { + roots.push(LeanExpr::sort(level)); + } + for bi in [ + BinderInfo::Default, + BinderInfo::Implicit, + BinderInfo::StrictImplicit, + BinderInfo::InstImplicit, + ] { + roots.push(LeanExpr::lam( + name(alias), + shared.clone(), + LeanExpr::bvar(Nat::from(0u64)), + bi.clone(), + )); + roots.push(LeanExpr::all(name(alias), shared.clone(), shared.clone(), bi)); + } + for nd in [false, true] { + roots.push(LeanExpr::letE( + name(alias), + shared.clone(), + shared.clone(), + shared.clone(), + nd, + )); + } + roots +} + +fn old_restore( + generated: &LeanExpr, + source: &LeanExpr, + fvars: &FxHashMap, + depth: usize, + params: &[Name], + stt: &CompileState, +) -> LeanExpr { + let mut hints = FxHashMap::default(); + reference::collect_lean_source_name_hints( + source, fvars, depth, params, stt, &mut hints, + ); + reference::restore_lean_source_name_hints( + generated, fvars, depth, params, stt, &hints, + ) +} + +#[test] +fn source_hints_match_reference_including_all_metadata() { + let stt = state(); + let fvars = FxHashMap::from_iter([(name("x"), 0)]); + let params = [name("u"), name("v")]; + for depth in [1, 3] { + for (source, generated) in fixtures("A").iter().zip(fixtures("B")) { + // A real reduction can introduce a different surrounding telescope. + let generated = LeanExpr::lam( + name("generated binder"), + LeanExpr::sort(Level::zero()), + LeanExpr::app(generated, LeanExpr::bvar(Nat::from(0u64))), + BinderInfo::InstImplicit, + ); + let expected = + old_restore(&generated, source, &fvars, depth, ¶ms, &stt); + let actual = restore(&generated, source, &fvars, depth, ¶ms, &stt); + assert_eq!(actual, expected); + assert_eq!(actual.get_hash(), expected.get_hash()); + } + } +} + +#[test] +fn structural_ids_match_kernel_equality_and_include_conversion_depth() { + let stt = state(); + let fvars = FxHashMap::from_iter([(name("x"), 0)]); + let params = [name("u"), name("v")]; + let roots: Vec<_> = fixtures("A").into_iter().chain(fixtures("B")).collect(); + let mut addresses = FxHashMap::default(); + for root in &roots { + addresses.extend(capture_addresses(root, root, &stt)); + } + let mut table = + ContentTable::::new(&addresses, &fvars, ¶ms); + let mut ids = Vec::new(); + let mut kernel = Vec::new(); + for depth in [1, 3] { + for root in &roots { + ids.push(table.content(root, depth)); + kernel.push(to_kexpr_static(root, &fvars, depth, ¶ms, &stt)); + } + } + for i in 0..ids.len() { + for j in 0..ids.len() { + assert_eq!(ids[i] == ids[j], kernel[i] == kernel[j], "pair {i}/{j}"); + } + } +} + +#[test] +fn first_hint_wins_in_original_preorder() { + let stt = state(); + let fvars = FxHashMap::from_iter([(name("x"), 0)]); + let a = application("A"); + let b = application("B"); + for (first, second) in [(&a, &b), (&b, &a)] { + let source = LeanExpr::app( + first.clone(), + LeanExpr::app(second.clone(), first.clone()), + ); + let actual = restore(second, &source, &fvars, 1, &[], &stt); + assert!(same_node(&actual, first)); + assert_eq!(actual, old_restore(second, &source, &fvars, 1, &[], &stt)); + } +} + +#[test] +fn eligibility_still_excludes_any_bvar_even_when_bound() { + let stt = state(); + let fvars = FxHashMap::from_iter([(name("x"), 0)]); + let params = [name("u"), name("v")]; + let roots = fixtures("A"); + for root in &roots { + let addresses = capture_addresses(root, root, &stt); + let mut pass = Pass::new(&addresses, &fvars, 1, ¶ms); + assert_eq!(pass.has_bvar(root), reference::expr_has_bvar(root)); + } + let bound = LeanExpr::app( + LeanExpr::cnst(name("A"), vec![]), + LeanExpr::lam( + name("bound"), + LeanExpr::sort(Level::zero()), + LeanExpr::bvar(Nat::from(0u64)), + BinderInfo::Default, + ), + ); + let addresses = capture_addresses(&bound, &bound, &stt); + let mut pass = Pass::new(&addresses, &fvars, 1, ¶ms); + assert!(pass.has_bvar(&bound)); + assert!(!pass.candidate(&bound)); + let mut hints = FxHashMap::default(); + pass.collect(&bound, &mut hints); + assert!(hints.is_empty()); +} + +#[test] +fn restored_memo_distinguishes_alias_metadata_when_no_hint_matches() { + let stt = state(); + let fvars = FxHashMap::from_iter([(name("x"), 0)]); + let a = LeanExpr::mdata(tag("A"), application("A")); + let b = LeanExpr::mdata(tag("B"), application("B")); + let root = LeanExpr::app(a.clone(), b.clone()); + let addresses = capture_addresses(&root, &root, &stt); + let mut pass = Pass::new(&addresses, &fvars, 1, &[]); + assert_eq!(pass.content.content(&a, 1), pass.content.content(&b, 1)); + let result = pass.restore(&root, &FxHashMap::default()); + assert!(same_node(&result, &root)); + let ExprData::App(lhs, rhs, _) = result.as_data() else { panic!("app") }; + assert!(same_node(lhs, &a)); + assert!(same_node(rhs, &b)); + assert_ne!(lhs.get_hash(), rhs.get_hash()); +} + +#[test] +fn resolution_is_frozen_within_pass_but_fresh_on_the_next_call() { + let stt = CompileState::new_empty(); + let fvars = FxHashMap::from_iter([(name("x"), 0)]); + let source = application("A"); + let generated = application("B"); + stt.name_to_addr.insert(name("A"), Address::hash(b"a")); + let before_publication = capture_addresses(&source, &generated, &stt); + stt.aux_name_to_addr.insert(name("B"), Address::hash(b"a")); + let mut pass = Pass::new(&before_publication, &fvars, 1, &[]); + let mut hints = FxHashMap::default(); + pass.collect(&source, &mut hints); + assert!(same_node(&pass.restore(&generated, &hints), &generated)); + assert!(same_node( + &restore(&generated, &source, &fvars, 1, &[], &stt), + &source + )); + + // A new primary entry takes precedence over the already-captured aux entry. + let before_override = capture_addresses(&source, &generated, &stt); + stt.name_to_addr.insert(name("B"), Address::hash(b"b")); + let mut pass = Pass::new(&before_override, &fvars, 1, &[]); + let mut hints = FxHashMap::default(); + pass.collect(&source, &mut hints); + assert!(same_node(&pass.restore(&generated, &hints), &source)); + assert!(same_node( + &restore(&generated, &source, &fvars, 1, &[], &stt), + &generated + )); +} + +#[test] +fn shared_dag_visits_unique_nodes_and_preserves_sharing() { + let stt = state(); + let fvars = FxHashMap::from_iter([(name("x"), 0)]); + let source = application("A"); + let generated = diamond(application("B"), 60); + let addresses = capture_addresses(&source, &generated, &stt); + let mut pass = Pass::new(&addresses, &fvars, 1, &[]); + let mut hints = FxHashMap::default(); + pass.collect(&source, &mut hints); + let result = pass.restore(&generated, &hints); + assert!(pass.restored.len() < 70); + assert!(pass.content.converted.len() < 70); + assert!(pass.has_bvar.len() < 70); + let mut cursor = &result; + for _ in 0..60 { + let ExprData::App(a, b, _) = cursor.as_data() else { panic!("app") }; + assert!(same_node(a, b)); + cursor = a; + } + assert!(same_node(cursor, &source)); + + let source = diamond(source, 60); + let unmatched = diamond(application("unmatched"), 60); + let addresses = capture_addresses(&source, &unmatched, &stt); + let mut pass = Pass::new(&addresses, &fvars, 1, &[]); + let mut hints = FxHashMap::default(); + pass.collect(&source, &mut hints); + assert!(pass.collected.len() < 70); + assert!(hints.len() < 70); + let result = pass.restore(&unmatched, &hints); + assert!(same_node(&result, &unmatched)); + assert!(pass.content.converted.len() < 140); +} + +#[derive(Default)] +struct AlwaysZero; +impl Hasher for AlwaysZero { + fn write(&mut self, _: &[u8]) {} + fn finish(&self) -> u64 { + 0 + } +} + +#[test] +fn structural_hash_collisions_cannot_merge_distinct_terms() { + let stt = state(); + let fvars = FxHashMap::from_iter([(name("x"), 0)]); + let params = [name("u"), name("v")]; + let roots: Vec<_> = fixtures("A").into_iter().chain(fixtures("B")).collect(); + let mut addresses = FxHashMap::default(); + for root in &roots { + addresses.extend(capture_addresses(root, root, &stt)); + } + let mut table = ContentTable::>::new( + &addresses, &fvars, ¶ms, + ); + let ids: Vec<_> = roots.iter().map(|e| table.content(e, 1)).collect(); + let kernel: Vec<_> = roots + .iter() + .map(|e| to_kexpr_static(e, &fvars, 1, ¶ms, &stt)) + .collect(); + for i in 0..ids.len() { + for j in 0..ids.len() { + assert_eq!( + ids[i] == ids[j], + kernel[i] == kernel[j], + "collision pair {i}/{j}" + ); + } + } +} + +#[test] +#[ignore = "manual release benchmark against the pre-cache restoration"] +fn source_name_hints_shared_dag_benchmark() { + use std::time::Instant; + let stt = state(); + let fvars = FxHashMap::from_iter([(name("x"), 0)]); + for depth in [8, 12, 16] { + let source = diamond(application("A"), depth); + let generated = diamond(application("unmatched"), depth); + let start = Instant::now(); + let expected = old_restore(&generated, &source, &fvars, 1, &[], &stt); + let old_time = start.elapsed(); + let start = Instant::now(); + let result = restore(&generated, &source, &fvars, 1, &[], &stt); + let new_time = start.elapsed(); + assert_eq!(result.get_hash(), expected.get_hash()); + eprintln!("depth={depth}: old={old_time:?}, cached={new_time:?}"); + } +} diff --git a/crates/compile/src/graph.rs b/crates/compile/src/graph.rs index 015010e5..4800f5ea 100644 --- a/crates/compile/src/graph.rs +++ b/crates/compile/src/graph.rs @@ -59,11 +59,15 @@ pub struct SetupScan { /// Fused whole-env setup pass: one decode per constant feeding the ref /// graph, the groundedness check, and inductive-group collection. -/// One pass instead of separate `build_ref_graph` / -/// `ground_consts`' scan, `validate_lean_ind_flags`' scan) — under -/// (the compile path'''s default) decodes a constant per access, so -/// visiting the whole env once instead of three times cuts the setup -/// decode count to a third. Outputs are identical to the separate passes. +/// The compile path's lazy environment can decode a constant on each +/// access, so one scan avoids the repeat decodes of separate graph, +/// grounding, and inductive-group scans. Outputs match those passes. +/// +/// Fold directly into chunk-local maps instead of allocating and merging +/// graph fragments for every constant. Reverse edges to a common dependency +/// share one set within the chunk; only completed chunk maps are reduced. +/// Rayon chooses the chunks dynamically, retaining work-stealing balance +/// when some constants take much longer to decode and traverse than others. pub fn setup_scan(env: &Env) -> SetupScan { #[derive(Default)] struct Acc { @@ -76,14 +80,19 @@ pub fn setup_scan(env: &Env) -> SetupScan { let names: Vec<&Name> = env.keys().collect(); let acc = names .into_par_iter() - .filter_map(|name| { - let constant = env.get(name)?; - let deps = get_constant_info_references(&constant); - let mut acc = Acc { - in_refs: mk_in_refs(name, &deps), - out_refs: RefMap::from_iter([(name.clone(), deps)]), - ..Acc::default() + .fold(Acc::default, |mut acc, name| { + let Some(constant) = env.get(name) else { + return acc; }; + let deps = get_constant_info_references(&constant); + // Keep an empty reverse-edge entry even for an isolated constant. + // Missing referenced names also retain their incoming edges, as in + // the reference graph used by ungroundedness propagation. + acc.in_refs.entry(name.clone()).or_default(); + for dep in &deps { + acc.in_refs.entry(dep.clone()).or_default().insert(name.clone()); + } + acc.out_refs.insert(name.clone(), deps); if let Err(err) = crate::ground::ground_const_check(&constant, env) { acc.ungrounded.insert(name.clone(), err); } @@ -95,7 +104,7 @@ pub fn setup_scan(env: &Env) -> SetupScan { { acc.ind_groups.entry(first.clone()).or_insert_with(|| v.all.clone()); } - Some(acc) + acc }) .reduce(Acc::default, |mut l, r| { l.out_refs = merge_ref_maps(l.out_refs, r.out_refs); @@ -114,23 +123,6 @@ pub fn setup_scan(env: &Env) -> SetupScan { } } -/// `name → {name} ∪ deps`-shaped reverse-edge fragment for one -/// constant, merged across the parallel scan. -fn mk_in_refs(name: &Name, deps: &NameSet) -> RefMap { - let mut in_refs = RefMap::from_iter([(name.clone(), NameSet::default())]); - for dep in deps { - match in_refs.entry(dep.clone()) { - Entry::Vacant(entry) => { - entry.insert(NameSet::from_iter([name.clone()])); - }, - Entry::Occupied(mut entry) => { - entry.get_mut().insert(name.clone()); - }, - } - } - in_refs -} - /// Size-aware map union (drain the smaller side into the bigger). fn merge_ref_maps(l: RefMap, r: RefMap) -> RefMap { let (smaller, mut bigger) = if l.len() < r.len() { (l, r) } else { (r, l) }; @@ -305,6 +297,272 @@ mod tests { ConstantVal { name: n(name), level_params: vec![], typ: sort0() } } + /// Frozen per-constant map/reduce implementation for differential tests + /// and the opt-in performance comparison below. + fn setup_scan_reference(env: &Env) -> SetupScan { + let empty = || SetupScan { + graph: RefGraph::default(), + immediate_ungrounded: FxHashMap::default(), + ind_groups: FxHashMap::default(), + }; + let names: Vec<_> = env.keys().collect(); + names + .into_par_iter() + .filter_map(|name| { + let constant = env.get(name)?; + let deps = get_constant_info_references(&constant); + let mut scan = empty(); + scan.graph.in_refs = + RefMap::from_iter([(name.clone(), NameSet::default())]); + for dep in &deps { + match scan.graph.in_refs.entry(dep.clone()) { + Entry::Vacant(entry) => { + entry.insert(NameSet::from_iter([name.clone()])); + }, + Entry::Occupied(mut entry) => { + entry.get_mut().insert(name.clone()); + }, + } + } + scan.graph.out_refs = RefMap::from_iter([(name.clone(), deps)]); + if let Err(err) = crate::ground::ground_const_check(&constant, env) { + scan.immediate_ungrounded.insert(name.clone(), err); + } + if let ConstantInfo::InductInfo(v) = &*constant + && let Some(first) = v.all.first() + { + scan.ind_groups.insert(first.clone(), v.all.clone()); + } + Some(scan) + }) + .reduce(empty, |mut l, r| { + l.graph.out_refs = merge_ref_maps(l.graph.out_refs, r.graph.out_refs); + l.graph.in_refs = merge_ref_maps(l.graph.in_refs, r.graph.in_refs); + l.immediate_ungrounded.extend(r.immediate_ungrounded); + for (k, v) in r.ind_groups { + l.ind_groups.entry(k).or_insert(v); + } + l + }) + } + + fn assert_same_scan(actual: &SetupScan, expected: &SetupScan) { + assert_eq!(actual.graph.out_refs, expected.graph.out_refs); + assert_eq!(actual.graph.in_refs, expected.graph.in_refs); + assert_eq!(actual.ind_groups, expected.ind_groups); + assert_eq!( + actual.immediate_ungrounded.len(), + expected.immediate_ungrounded.len() + ); + for (name, expected) in &expected.immediate_ungrounded { + assert_eq!( + format!("{:?}", actual.immediate_ungrounded.get(name).unwrap()), + format!("{expected:?}"), + "grounding error for {name}", + ); + } + } + + /// Well-typed aliases with both shared and local dependencies. The input + /// graph has no cycles, while many reverse edges meet at the same names. + fn setup_fixture(count: usize) -> Env { + let mut env = Env::default(); + let typ = Expr::sort(Level::succ(Level::zero())); + for name in ["Base", "Isolated"] { + env.insert( + n(name), + ConstantInfo::AxiomInfo(AxiomVal { + cnst: ConstantVal { + name: n(name), + level_params: vec![], + typ: typ.clone(), + }, + is_unsafe: false, + }), + ); + } + let names: Vec<_> = (0..count).map(|i| n(&format!("alias_{i}"))).collect(); + let base = Expr::cnst(n("Base"), vec![]); + for (i, name) in names.iter().enumerate() { + let mut value = base.clone(); + for offset in [1, 3, 17, 127] { + let dep = + i.checked_sub(offset).map_or_else(|| n("Base"), |j| names[j].clone()); + value = Expr::all( + Name::anon(), + Expr::cnst(dep, vec![]), + value, + BinderInfo::Default, + ); + } + env.insert( + name.clone(), + ConstantInfo::DefnInfo(DefinitionVal { + cnst: ConstantVal { + name: name.clone(), + level_params: vec![], + typ: typ.clone(), + }, + value, + hints: ReducibilityHints::Abbrev, + safety: DefinitionSafety::Safe, + all: vec![name.clone()], + }), + ); + } + env + } + + #[test] + fn setup_scan_matches_reference_across_worker_counts() { + let mut env = setup_fixture(1024); + // Cover isolated nodes, self/cyclic edges, missing names, each family of + // grounding error, and mutual-group deduplication across fold chunks. + for (name, typ) in [ + ("Self", Expr::cnst(n("Self"), vec![])), + ("CycleA", Expr::cnst(n("CycleB"), vec![])), + ("CycleB", Expr::cnst(n("CycleA"), vec![])), + ("Missing", Expr::cnst(n("Absent"), vec![])), + ("Bvar", Expr::bvar(Nat::from(0u64))), + ("Fvar", Expr::fvar(n("local"))), + ("Mvar", Expr::mvar(n("hole"))), + ("Universe", Expr::sort(Level::param(n("u")))), + ] { + env.insert( + n(name), + ConstantInfo::AxiomInfo(AxiomVal { + cnst: ConstantVal { name: n(name), level_params: vec![], typ }, + is_unsafe: false, + }), + ); + } + for (name, all, ctors) in [ + ("MutA", vec![n("MutA"), n("MutB")], vec![]), + ("MutB", vec![n("MutA"), n("MutB")], vec![]), + ("EmptyGroup", vec![], vec![]), + ("BadCtor", vec![n("BadCtor")], vec![n("AbsentCtor")]), + ] { + env.insert( + n(name), + ConstantInfo::InductInfo(InductiveVal { + cnst: mk_cv(name), + num_params: Nat::from(0u64), + num_indices: Nat::from(0u64), + all, + ctors, + num_nested: Nat::from(0u64), + is_rec: false, + is_unsafe: false, + is_reflexive: false, + }), + ); + } + let expected = setup_scan_reference(&env); + assert_eq!(expected.immediate_ungrounded.len(), 6); + assert!(expected.graph.in_refs[&n("Isolated")].is_empty()); + assert_eq!(expected.ind_groups.len(), 2); + let graph = build_ref_graph(&env); + assert_eq!(expected.graph.out_refs, graph.out_refs); + assert_eq!(expected.graph.in_refs, graph.in_refs); + for workers in [1, 2, 4, 8] { + let pool = + rayon::ThreadPoolBuilder::new().num_threads(workers).build().unwrap(); + pool.install(|| { + assert_same_scan(&setup_scan(&env), &expected); + assert_same_scan( + &setup_scan(&Env::default()), + &setup_scan_reference(&Env::default()), + ); + }); + } + } + + #[test] + fn setup_scan_fetches_each_lazy_entry_once_and_skips_unavailable_entries() { + use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }; + let backing = setup_fixture(256); + let expected = setup_scan_reference(&backing); + let mut names: Vec<_> = backing.keys().cloned().collect(); + names.push(n("Unavailable")); + let calls = Arc::new(AtomicUsize::new(0)); + let counted = Arc::clone(&calls); + let count = names.len(); + let env = Env::new_lazy( + names, + Box::new(move |name| { + counted.fetch_add(1, Ordering::Relaxed); + backing.get(name).map(|entry| entry.cloned()) + }), + 16, + ); + assert_same_scan(&setup_scan(&env), &expected); + assert_eq!(calls.load(Ordering::Relaxed), count); + } + + #[test] + fn setup_scan_preserves_compiled_fixture_bytes() { + use crate::compile::{CompileOptions, compile_env_with_options}; + let source = std::sync::Arc::new(setup_fixture(32)); + for workers in [1, 4] { + let compiled = compile_env_with_options( + &source, + CompileOptions { max_workers: Some(workers) }, + ) + .unwrap(); + assert!(compiled.ungrounded.is_empty()); + let mut bytes = Vec::new(); + compiled.env.put(&mut bytes).unwrap(); + // Captured from the per-constant scanner before the fold change. + assert_eq!( + blake3::hash(&bytes).to_hex().as_str(), + "d0a1c1490fad7e9e9236792c62ac3697dc7d3b541e3b3db12e00b9dd29ba36ed" + ); + } + } + + /// Run explicitly in release mode: cargo test -p ix-compile --release + /// setup_scan_benchmark -- --ignored --nocapture --test-threads=1 + #[test] + #[ignore = "synthetic setup-scan timing comparison; run in release mode"] + fn setup_scan_benchmark() { + use std::time::{Duration, Instant}; + for count in [50_000, 250_000] { + let env = setup_fixture(count); + for workers in [1, 8, 32] { + let pool = + rayon::ThreadPoolBuilder::new().num_threads(workers).build().unwrap(); + pool.install(|| { + assert_same_scan(&setup_scan(&env), &setup_scan_reference(&env)); + let mut elapsed = [Duration::ZERO; 2]; + for trial in 0..6 { + // Alternate order to reduce allocator/cache warmup bias. Exclude + // destruction, which the compiler likewise performs after setup. + for index in [trial % 2, 1 - trial % 2] { + let start = Instant::now(); + let result = if index == 0 { + setup_scan_reference(&env) + } else { + setup_scan(&env) + }; + elapsed[index] += start.elapsed(); + std::hint::black_box(result); + } + } + eprintln!( + "setup_scan: {} constants, {workers} workers: per-constant {:.3}s, chunked {:.3}s ({:.2}x)", + env.len(), + elapsed[0].as_secs_f64() / 6.0, + elapsed[1].as_secs_f64() / 6.0, + elapsed[0].as_secs_f64() / elapsed[1].as_secs_f64() + ); + }); + } + } + } + #[test] fn empty_env() { let env = Env::default(); diff --git a/crates/kernel/src/env.rs b/crates/kernel/src/env.rs index e2531aaf..0f5c9134 100644 --- a/crates/kernel/src/env.rs +++ b/crates/kernel/src/env.rs @@ -128,6 +128,25 @@ pub struct InternTable { pub(crate) clo_scratch_pool: Vec>>, } +/// Original input allocation -> completed canonical result, for ONE public +/// interning call. The caller retains its input root, and recursive helpers +/// borrow only descendants of that root, so these pointer keys cannot be reused +/// while the memo is alive. Rebuilt temporaries must never become input keys. +/// +/// Exact pointers (not semantic uids) distinguish metadata-bearing occurrences. +/// Only changed nodes need entries: unchanged nodes hit the canonical-uid sets. +/// Dropping the memo on return also avoids retaining a second DAG across calls. +struct InternMemo { + exprs: FxHashMap>, + univs: FxHashMap>, +} + +impl Default for InternMemo { + fn default() -> Self { + Self { exprs: FxHashMap::default(), univs: FxHashMap::default() } + } +} + impl Default for InternTable { fn default() -> Self { Self::new() @@ -203,48 +222,79 @@ impl InternTable { /// identity, recursively canonicalizing children as needed so the /// shallow key is meaningful. pub fn intern_univ(&mut self, u: KUniv) -> KUniv { - use super::level::UnivData; crate::profile::bump_intern_nodes(); if self.canon_univs.contains(u.addr()) { return u; } + // The root cannot recur below itself. Only descendants need memo entries. + self.intern_univ_node(&u, &mut InternMemo::default()) + } + + #[inline] + fn intern_univ_cached( + &mut self, + input: &KUniv, + memo: &mut InternMemo, + ) -> KUniv { + crate::profile::bump_intern_nodes(); + if self.canon_univs.contains(input.addr()) { + return input.clone(); + } + let ptr = std::ptr::from_ref(input.data()).addr(); + if let Some(canonical) = memo.univs.get(&ptr) { + return canonical.clone(); + } + let canonical = self.intern_univ_node(input, memo); + if !canonical.ptr_eq(input) { + memo.univs.insert(ptr, canonical.clone()); + } + canonical + } + + fn intern_univ_node( + &mut self, + input: &KUniv, + memo: &mut InternMemo, + ) -> KUniv { + use super::level::UnivData; // Canonicalize children first; rebuild only if any child changed. - let u = match u.data() { + let u = match input.data() { UnivData::Succ(inner, _) => { - let ci = self.intern_univ(inner.clone()); + let ci = self.intern_univ_cached(inner, memo); if ci.ptr_eq(inner) { - u + input.clone() } else { KUniv::new(UnivData::Succ(ci, super::expr::fresh_uid())) } }, UnivData::Max(a, b, _) => { - let ca = self.intern_univ(a.clone()); - let cb = self.intern_univ(b.clone()); + let ca = self.intern_univ_cached(a, memo); + let cb = self.intern_univ_cached(b, memo); if ca.ptr_eq(a) && cb.ptr_eq(b) { - u + input.clone() } else { KUniv::new(UnivData::Max(ca, cb, super::expr::fresh_uid())) } }, UnivData::IMax(a, b, _) => { - let ca = self.intern_univ(a.clone()); - let cb = self.intern_univ(b.clone()); + let ca = self.intern_univ_cached(a, memo); + let cb = self.intern_univ_cached(b, memo); if ca.ptr_eq(a) && cb.ptr_eq(b) { - u + input.clone() } else { KUniv::new(UnivData::IMax(ca, cb, super::expr::fresh_uid())) } }, - UnivData::Zero(_) | UnivData::Param(..) => u, + UnivData::Zero(_) | UnivData::Param(..) => input.clone(), }; let key = univ_key(&u); if let Some(existing) = self.univs.get(&key) { - return existing.clone(); + existing.clone() + } else { + self.canon_univs.insert(*u.addr()); + self.univs.insert(key, u.clone()); + u } - self.canon_univs.insert(*u.addr()); - self.univs.insert(key, u.clone()); - u } /// Intern an expression: returns the canonical value for its structural @@ -253,85 +303,125 @@ impl InternTable { /// make the shallow key meaningless), preserving the historical /// content-hash interning semantics. pub fn intern_expr(&mut self, e: KExpr) -> KExpr { - use super::expr::ExprData; crate::profile::bump_intern_nodes(); if self.canon_exprs.contains(e.addr()) { return e; } - let e = match e.data() { + // Retain the entire input DAG until this call-local pointer memo is gone. + // Do not memoize the root: this avoids allocating for the common case of + // a new node whose children are already canonical. + self.intern_expr_node(&e, &mut InternMemo::default()) + } + + #[inline] + fn intern_expr_cached( + &mut self, + input: &KExpr, + memo: &mut InternMemo, + ) -> KExpr { + crate::profile::bump_intern_nodes(); + // Keep this before the pointer memo: already-canonical uids retain the + // historical fast path, including their occurrence metadata. + if self.canon_exprs.contains(input.addr()) { + return input.clone(); + } + let ptr = std::ptr::from_ref(input.data()).addr(); + if let Some(canonical) = memo.exprs.get(&ptr) { + return canonical.clone(); + } + let canonical = self.intern_expr_node(input, memo); + // Do not mark the original uid canonical when it was replaced. Future + // edges to this input must return the actual canonical OUTPUT, not input. + if !canonical.ptr_eq(input) { + memo.exprs.insert(ptr, canonical.clone()); + } + canonical + } + + fn intern_expr_node( + &mut self, + input: &KExpr, + memo: &mut InternMemo, + ) -> KExpr { + use super::expr::ExprData; + let e = match input.data() { ExprData::Sort(un, _) => { - let cu = self.intern_univ(un.clone()); + let cu = self.intern_univ_cached(un, memo); if cu.ptr_eq(un) { - e + input.clone() } else { // Child canonicalization only — same semantic level, same // occurrence: the spelling decoration rides along. - KExpr::sort_full(cu, e.mdata().clone(), e.univ_decor().clone()) + KExpr::sort_full( + cu, + input.mdata().clone(), + input.univ_decor().clone(), + ) } }, ExprData::Const(id, us, _) => { let cus: Box<[KUniv]> = - us.iter().map(|un| self.intern_univ(un.clone())).collect(); + us.iter().map(|un| self.intern_univ_cached(un, memo)).collect(); if cus.iter().zip(us.iter()).all(|(a, b)| a.ptr_eq(b)) { - e + input.clone() } else { KExpr::cnst_full( id.clone(), cus, - e.mdata().clone(), - e.univ_decor().clone(), + input.mdata().clone(), + input.univ_decor().clone(), ) } }, ExprData::App(f, a, _) => { - let cf = self.intern_expr(f.clone()); - let ca = self.intern_expr(a.clone()); + let cf = self.intern_expr_cached(f, memo); + let ca = self.intern_expr_cached(a, memo); if cf.ptr_eq(f) && ca.ptr_eq(a) { - e + input.clone() } else { - KExpr::app_mdata(cf, ca, e.mdata().clone()) + KExpr::app_mdata(cf, ca, input.mdata().clone()) } }, ExprData::Lam(n, bi, t, b, _) => { - let ct = self.intern_expr(t.clone()); - let cb = self.intern_expr(b.clone()); + let ct = self.intern_expr_cached(t, memo); + let cb = self.intern_expr_cached(b, memo); if ct.ptr_eq(t) && cb.ptr_eq(b) { - e + input.clone() } else { - KExpr::lam_mdata(n.clone(), bi.clone(), ct, cb, e.mdata().clone()) + KExpr::lam_mdata(n.clone(), bi.clone(), ct, cb, input.mdata().clone()) } }, ExprData::All(n, bi, t, b, _) => { - let ct = self.intern_expr(t.clone()); - let cb = self.intern_expr(b.clone()); + let ct = self.intern_expr_cached(t, memo); + let cb = self.intern_expr_cached(b, memo); if ct.ptr_eq(t) && cb.ptr_eq(b) { - e + input.clone() } else { - KExpr::all_mdata(n.clone(), bi.clone(), ct, cb, e.mdata().clone()) + KExpr::all_mdata(n.clone(), bi.clone(), ct, cb, input.mdata().clone()) } }, ExprData::Let(n, t, v, b, nd, _) => { - let ct = self.intern_expr(t.clone()); - let cv = self.intern_expr(v.clone()); - let cb = self.intern_expr(b.clone()); + let ct = self.intern_expr_cached(t, memo); + let cv = self.intern_expr_cached(v, memo); + let cb = self.intern_expr_cached(b, memo); if ct.ptr_eq(t) && cv.ptr_eq(v) && cb.ptr_eq(b) { - e + input.clone() } else { - KExpr::let_mdata(n.clone(), ct, cv, cb, *nd, e.mdata().clone()) + KExpr::let_mdata(n.clone(), ct, cv, cb, *nd, input.mdata().clone()) } }, ExprData::Prj(id, f, v, _) => { - let cv = self.intern_expr(v.clone()); + let cv = self.intern_expr_cached(v, memo); if cv.ptr_eq(v) { - e + input.clone() } else { - KExpr::prj_mdata(id.clone(), *f, cv, e.mdata().clone()) + KExpr::prj_mdata(id.clone(), *f, cv, input.mdata().clone()) } }, ExprData::Var(..) | ExprData::FVar(..) | ExprData::Nat(..) - | ExprData::Str(..) => e, + | ExprData::Str(..) => input.clone(), }; let key = expr_key(&e); if let Some(existing) = self.exprs.get(&key) { @@ -341,14 +431,23 @@ impl InternTable { // debug builds; a violation here would be an interning bug, not // an input an adversary can craft (uids are assigned, not hashed). debug_assert!(existing == &e, "intern hit is not structurally equal"); - return existing.clone(); + existing.clone() + } else { + self.canon_exprs.insert(*e.addr()); + self.exprs.insert(key, e.clone()); + e } - self.canon_exprs.insert(*e.addr()); - self.exprs.insert(key, e.clone()); - e } } +#[cfg(test)] +#[path = "intern_reference.rs"] +mod intern_reference; + +#[cfg(test)] +#[path = "intern_tests.rs"] +mod intern_tests; + /// Generated recursor, cached after inductive validation. #[derive(Clone, Debug)] pub struct GeneratedRecursor { diff --git a/crates/kernel/src/intern_reference.rs b/crates/kernel/src/intern_reference.rs new file mode 100644 index 00000000..b14fc6cb --- /dev/null +++ b/crates/kernel/src/intern_reference.rs @@ -0,0 +1,156 @@ +//! Frozen pre-memoization interner for differential tests and benchmarks. +//! Keep the traversal, rebuilding, and first-insert-wins behavior unchanged. + +use super::{InternTable, expr_key, univ_key}; +use crate::{expr::KExpr, level::KUniv, mode::KernelMode}; + +impl InternTable { + /// Intern a universe: returns the canonical value for its structural + /// identity, recursively canonicalizing children as needed so the + /// shallow key is meaningful. + pub(super) fn intern_univ_reference(&mut self, u: KUniv) -> KUniv { + use crate::level::UnivData; + crate::profile::bump_intern_nodes(); + if self.canon_univs.contains(u.addr()) { + return u; + } + // Canonicalize children first; rebuild only if any child changed. + let u = match u.data() { + UnivData::Succ(inner, _) => { + let ci = self.intern_univ_reference(inner.clone()); + if ci.ptr_eq(inner) { + u + } else { + KUniv::new(UnivData::Succ(ci, crate::expr::fresh_uid())) + } + }, + UnivData::Max(a, b, _) => { + let ca = self.intern_univ_reference(a.clone()); + let cb = self.intern_univ_reference(b.clone()); + if ca.ptr_eq(a) && cb.ptr_eq(b) { + u + } else { + KUniv::new(UnivData::Max(ca, cb, crate::expr::fresh_uid())) + } + }, + UnivData::IMax(a, b, _) => { + let ca = self.intern_univ_reference(a.clone()); + let cb = self.intern_univ_reference(b.clone()); + if ca.ptr_eq(a) && cb.ptr_eq(b) { + u + } else { + KUniv::new(UnivData::IMax(ca, cb, crate::expr::fresh_uid())) + } + }, + UnivData::Zero(_) | UnivData::Param(..) => u, + }; + let key = univ_key(&u); + if let Some(existing) = self.univs.get(&key) { + return existing.clone(); + } + self.canon_univs.insert(*u.addr()); + self.univs.insert(key, u.clone()); + u + } + + /// Intern an expression: returns the canonical value for its structural + /// identity. Children are canonicalized recursively when needed (a node + /// built outside the table has non-canonical children whose uids would + /// make the shallow key meaningless), preserving the historical + /// content-hash interning semantics. + pub(super) fn intern_expr_reference(&mut self, e: KExpr) -> KExpr { + use crate::expr::ExprData; + crate::profile::bump_intern_nodes(); + if self.canon_exprs.contains(e.addr()) { + return e; + } + let e = match e.data() { + ExprData::Sort(un, _) => { + let cu = self.intern_univ_reference(un.clone()); + if cu.ptr_eq(un) { + e + } else { + // Child canonicalization only — same semantic level, same + // occurrence: the spelling decoration rides along. + KExpr::sort_full(cu, e.mdata().clone(), e.univ_decor().clone()) + } + }, + ExprData::Const(id, us, _) => { + let cus: Box<[KUniv]> = + us.iter().map(|un| self.intern_univ_reference(un.clone())).collect(); + if cus.iter().zip(us.iter()).all(|(a, b)| a.ptr_eq(b)) { + e + } else { + KExpr::cnst_full( + id.clone(), + cus, + e.mdata().clone(), + e.univ_decor().clone(), + ) + } + }, + ExprData::App(f, a, _) => { + let cf = self.intern_expr_reference(f.clone()); + let ca = self.intern_expr_reference(a.clone()); + if cf.ptr_eq(f) && ca.ptr_eq(a) { + e + } else { + KExpr::app_mdata(cf, ca, e.mdata().clone()) + } + }, + ExprData::Lam(n, bi, t, b, _) => { + let ct = self.intern_expr_reference(t.clone()); + let cb = self.intern_expr_reference(b.clone()); + if ct.ptr_eq(t) && cb.ptr_eq(b) { + e + } else { + KExpr::lam_mdata(n.clone(), bi.clone(), ct, cb, e.mdata().clone()) + } + }, + ExprData::All(n, bi, t, b, _) => { + let ct = self.intern_expr_reference(t.clone()); + let cb = self.intern_expr_reference(b.clone()); + if ct.ptr_eq(t) && cb.ptr_eq(b) { + e + } else { + KExpr::all_mdata(n.clone(), bi.clone(), ct, cb, e.mdata().clone()) + } + }, + ExprData::Let(n, t, v, b, nd, _) => { + let ct = self.intern_expr_reference(t.clone()); + let cv = self.intern_expr_reference(v.clone()); + let cb = self.intern_expr_reference(b.clone()); + if ct.ptr_eq(t) && cv.ptr_eq(v) && cb.ptr_eq(b) { + e + } else { + KExpr::let_mdata(n.clone(), ct, cv, cb, *nd, e.mdata().clone()) + } + }, + ExprData::Prj(id, f, v, _) => { + let cv = self.intern_expr_reference(v.clone()); + if cv.ptr_eq(v) { + e + } else { + KExpr::prj_mdata(id.clone(), *f, cv, e.mdata().clone()) + } + }, + ExprData::Var(..) + | ExprData::FVar(..) + | ExprData::Nat(..) + | ExprData::Str(..) => e, + }; + let key = expr_key(&e); + if let Some(existing) = self.exprs.get(&key) { + // The shallow key (exact structural Eq over variant tag + child + // uids + payload — never a truncated or content-hashed key) plus + // canonical children make this hit structurally exact. Checked in + // debug builds; a violation here would be an interning bug, not + // an input an adversary can craft (uids are assigned, not hashed). + debug_assert!(existing == &e, "intern hit is not structurally equal"); + return existing.clone(); + } + self.canon_exprs.insert(*e.addr()); + self.exprs.insert(key, e.clone()); + e + } +} diff --git a/crates/kernel/src/intern_tests.rs b/crates/kernel/src/intern_tests.rs new file mode 100644 index 00000000..28858113 --- /dev/null +++ b/crates/kernel/src/intern_tests.rs @@ -0,0 +1,436 @@ +//! Differential semantics and DAG-work bounds for the kernel interner. + +use super::*; +use crate::expr::{ExprData, FVarId, MData, fresh_uid}; +use crate::level::UnivData; +use crate::mode::{Anon, Meta}; +use crate::profile::take_op_counts; +use bignat::Nat; +use ix_common::env::{BinderInfo, DataValue, Name}; +use ixon::univ::Univ; + +fn name(s: &str) -> Name { + Name::str(Name::anon(), s.to_owned()) +} + +fn metadata(s: &str) -> Vec { + vec![vec![(name("tag"), DataValue::OfString(s.to_owned()))]] +} + +// Unlike KExpr/KUniv equality, these comparisons include ALL metadata and +// cached annotations, ignoring only fresh ephemeral uids. Compare DAGs without +// expanding shared subtrees or formatting an enormous expression on failure. +fn assert_same_univ(a: &KUniv, b: &KUniv) { + let mut pending = vec![(a, b)]; + let mut seen = FxHashSet::default(); + while let Some((a, b)) = pending.pop() { + if !seen.insert(( + std::ptr::from_ref(a.data()).addr(), + std::ptr::from_ref(b.data()).addr(), + )) { + continue; + } + match (a.data(), b.data()) { + (UnivData::Zero(_), UnivData::Zero(_)) => {}, + (UnivData::Param(i, n, _), UnivData::Param(j, m, _)) => { + assert_eq!(i, j); + assert_eq!(n, m); + }, + (UnivData::Succ(a, _), UnivData::Succ(b, _)) => pending.push((a, b)), + (UnivData::Max(a, b, _), UnivData::Max(c, d, _)) + | (UnivData::IMax(a, b, _), UnivData::IMax(c, d, _)) => { + pending.extend([(a, c), (b, d)]); + }, + _ => panic!("universe variants differ"), + } + } +} + +fn assert_same_expr(a: &KExpr, b: &KExpr) { + let mut pending = vec![(a, b)]; + let mut seen = FxHashSet::default(); + while let Some((a, b)) = pending.pop() { + if !seen.insert(( + std::ptr::from_ref(a.data()).addr(), + std::ptr::from_ref(b.data()).addr(), + )) { + continue; + } + assert_eq!(a.lbr(), b.lbr()); + assert_eq!(a.count_0(), b.count_0()); + assert_eq!(a.info().has_fvars, b.info().has_fvars); + assert_eq!(a.mdata(), b.mdata()); + assert_eq!(a.univ_decor(), b.univ_decor()); + match (a.data(), b.data()) { + (ExprData::Var(i, n, _), ExprData::Var(j, m, _)) => { + assert_eq!(i, j); + assert_eq!(n, m); + }, + (ExprData::FVar(i, n, _), ExprData::FVar(j, m, _)) => { + assert_eq!(i, j); + assert_eq!(n, m); + }, + (ExprData::Sort(u, _), ExprData::Sort(v, _)) => assert_same_univ(u, v), + (ExprData::Const(i, us, _), ExprData::Const(j, vs, _)) => { + assert_eq!(i, j); + assert_eq!(us.len(), vs.len()); + for (u, v) in us.iter().zip(vs) { + assert_same_univ(u, v); + } + }, + (ExprData::App(f, a, _), ExprData::App(g, b, _)) => { + pending.extend([(f, g), (a, b)]); + }, + (ExprData::Lam(n, bi, t, b, _), ExprData::Lam(m, bj, u, c, _)) + | (ExprData::All(n, bi, t, b, _), ExprData::All(m, bj, u, c, _)) => { + assert_eq!(n, m); + assert_eq!(bi, bj); + pending.extend([(t, u), (b, c)]); + }, + (ExprData::Let(n, t, v, b, nd, _), ExprData::Let(m, u, w, c, md, _)) => { + assert_eq!(n, m); + assert_eq!(nd, md); + pending.extend([(t, u), (v, w), (b, c)]); + }, + (ExprData::Prj(i, f, v, _), ExprData::Prj(j, g, w, _)) => { + assert_eq!(i, j); + assert_eq!(f, g); + pending.push((v, w)); + }, + (ExprData::Nat(n, a, _), ExprData::Nat(m, b, _)) => { + assert_eq!(n, m); + assert_eq!(a, b); + }, + (ExprData::Str(s, a, _), ExprData::Str(t, b, _)) => { + assert_eq!(s, t); + assert_eq!(a, b); + }, + _ => panic!("expression variants differ"), + } + } +} + +fn expr_dag(mut leaf: KExpr, depth: usize) -> KExpr { + for _ in 0..depth { + leaf = KExpr::all( + M::meta_field(name("x")), + M::meta_field(BinderInfo::Default), + leaf.clone(), + leaf, + ); + } + leaf +} + +fn univ_dag(mut leaf: KUniv, depth: usize) -> KUniv { + for _ in 0..depth { + // Use raw constructors: interning must preserve even unsimplified Max. + leaf = KUniv::new(UnivData::Max(leaf.clone(), leaf, fresh_uid())); + } + leaf +} + +fn fixtures(label: &str) -> Vec> { + let n = M::meta_field(name(label)); + let md = M::meta_field(metadata(label)); + let id = KId::new(Address::hash(b"constant"), n.clone()); + let u = KUniv::param(0, n.clone()); + let v = KUniv::param(1, n.clone()); + let levels = vec![ + KUniv::zero(), + KUniv::succ(u.clone()), + KUniv::new(UnivData::Max(u.clone(), v.clone(), fresh_uid())), + KUniv::new(UnivData::IMax(v, u.clone(), fresh_uid())), + u, + ]; + let var = KExpr::var_mdata(0, n.clone(), md.clone()); + let fvar = KExpr::fvar_mdata(FVarId(7), n.clone(), md.clone()); + let cnst = KExpr::cnst_mdata(id.clone(), levels.clone().into(), md.clone()); + let shared = KExpr::app_mdata(cnst.clone(), fvar.clone(), md.clone()); + let mut roots = vec![ + var.clone(), + fvar, + cnst, + shared.clone(), + KExpr::prj_mdata(id, 2, shared.clone(), md.clone()), + KExpr::nat_mdata(Nat::from(17u64), Address::hash(b"17"), md.clone()), + KExpr::str_mdata("hello".to_owned(), Address::hash(b"hello"), md.clone()), + expr_dag(shared.clone(), 6), + ]; + for level in levels { + roots.push(KExpr::sort_mdata(level, md.clone())); + } + for bi in [ + BinderInfo::Default, + BinderInfo::Implicit, + BinderInfo::StrictImplicit, + BinderInfo::InstImplicit, + ] { + roots.push(KExpr::lam_mdata( + n.clone(), + M::meta_field(bi.clone()), + shared.clone(), + var.clone(), + md.clone(), + )); + roots.push(KExpr::all_mdata( + n.clone(), + M::meta_field(bi), + shared.clone(), + var.clone(), + md.clone(), + )); + } + for non_dep in [false, true] { + roots.push(KExpr::let_mdata( + n.clone(), + shared.clone(), + var.clone(), + shared.clone(), + non_dep, + md.clone(), + )); + } + roots +} + +fn differential() { + for seeded in [false, true] { + let mut old = InternTable::::new(); + let mut new = InternTable::::new(); + if seeded { + for seed in fixtures("first") { + old.intern_expr_reference(seed.clone()); + new.intern_expr(seed); + } + } + for root in fixtures("second") { + let expected = old.intern_expr_reference(root.clone()); + let actual = new.intern_expr(root.clone()); + assert_same_expr(&actual, &expected); + assert!(new.intern_expr(root).ptr_eq(&actual)); + take_op_counts(); + assert!(new.intern_expr(actual.clone()).ptr_eq(&actual)); + assert_eq!(take_op_counts().intern_nodes, 1); + } + assert_eq!(new.exprs.len(), old.exprs.len()); + assert_eq!(new.univs.len(), old.univs.len()); + assert_eq!(new.canon_exprs.len(), old.canon_exprs.len()); + assert_eq!(new.canon_univs.len(), old.canon_univs.len()); + } +} + +#[test] +fn interning_matches_reference_in_both_modes_and_seeded_tables() { + differential::(); + differential::(); +} + +#[test] +fn shared_expression_dag_visits_edges_not_expanded_tree() { + let depth = 60; + let mut intern = InternTable::::new(); + intern.intern_expr(KExpr::sort(KUniv::zero())); + let root = expr_dag(KExpr::sort(KUniv::zero()), depth); + let mut previous: Option> = None; + for _ in 0..2 { + take_op_counts(); + let result = intern.intern_expr(root.clone()); + assert_eq!(take_op_counts().intern_nodes, (2 * depth + 2) as u64); + if let Some(previous) = previous.replace(result.clone()) { + assert!(previous.ptr_eq(&result)); + } + let mut cursor = &result; + for _ in 0..depth { + let ExprData::All(_, _, ty, body, _) = cursor.data() else { + panic!("all") + }; + assert!(ty.ptr_eq(body)); + cursor = ty; + } + assert_eq!(intern.exprs.len(), depth + 1); + assert!(!intern.canon_exprs.contains(root.addr())); + } +} + +#[test] +fn shared_universe_dag_is_memoized_across_constant_arguments() { + let depth = 60; + let mut intern = InternTable::::new(); + intern.intern_univ(KUniv::zero()); + let level = univ_dag(KUniv::zero(), depth); + take_op_counts(); + let canonical = intern.intern_univ(level.clone()); + assert_eq!(take_op_counts().intern_nodes, (2 * depth + 1) as u64); + let root = KExpr::cnst( + KId::new(Address::hash(b"C"), ()), + vec![level.clone(); 64].into(), + ); + take_op_counts(); + let result = intern.intern_expr(root); + assert_eq!(take_op_counts().intern_nodes, (2 * depth + 65) as u64); + let ExprData::Const(_, levels, _) = result.data() else { panic!("const") }; + assert!(levels.iter().all(|u| u.ptr_eq(&canonical))); + let mut cursor = &canonical; + for _ in 0..depth { + let UnivData::Max(a, b, _) = cursor.data() else { panic!("max") }; + assert!(a.ptr_eq(b)); + cursor = a; + } + assert_eq!(intern.univs.len(), depth + 1); + assert!(!intern.canon_univs.contains(level.addr())); +} + +#[test] +fn unchanged_inputs_use_canonical_fast_path_without_memo_entries() { + let mut intern = InternTable::::new(); + let root = expr_dag(KExpr::sort(KUniv::zero()), 40); + let mut memo = InternMemo::default(); + let result = intern.intern_expr_cached(&root, &mut memo); + assert!(result.ptr_eq(&root)); + assert_eq!(memo.exprs.capacity(), 0); + assert_eq!(memo.univs.capacity(), 0); +} + +#[test] +fn input_pointer_keys_preserve_same_uid_spelling_twins() { + for is_const in [false, true] { + let level = KUniv::zero(); + let a = if is_const { + KExpr::::cnst_mdata( + KId::new(Address::hash(b"C"), name("A")), + vec![level.clone()].into(), + metadata("A"), + ) + } else { + KExpr::sort_mdata(level.clone(), metadata("A")) + }; + let mut info = a.info().clone(); + info.mdata = metadata("B"); + let spelling = Univ::imax(Univ::succ(Univ::zero()), Univ::zero()); + let b = if is_const { + info.univ_decor = Some(UnivDecor::Const(vec![spelling].into())); + KExpr::new(ExprData::Const( + KId::new(Address::hash(b"C"), name("B")), + vec![level].into(), + info, + )) + } else { + info.univ_decor = Some(UnivDecor::Sort(spelling)); + KExpr::new(ExprData::Sort(level, info)) + }; + assert_eq!(a.addr(), b.addr()); + assert!(!a.ptr_eq(&b)); + let root = KExpr::app(a.clone(), KExpr::app(b, a)); + let mut old = InternTable::new(); + let mut new = InternTable::new(); + // Force both spelling occurrences to be rebuilt, so their shared input + // uid never becomes canonical and a uid-keyed memo would conflate them. + old.intern_univ_reference(KUniv::zero()); + new.intern_univ(KUniv::zero()); + let expected = old.intern_expr_reference(root.clone()); + let result = new.intern_expr(root); + assert_same_expr(&result, &expected); + let ExprData::App(a, rest, _) = result.data() else { panic!("app") }; + let ExprData::App(b, again, _) = rest.data() else { panic!("app") }; + assert!(!a.ptr_eq(b)); + assert!(a.ptr_eq(again)); + assert_eq!(a.mdata(), &metadata("A")); + assert_eq!(b.mdata(), &metadata("B")); + assert_ne!(a.univ_decor(), b.univ_decor()); + } +} + +#[test] +fn already_canonical_uid_fast_path_preserves_occurrence_metadata() { + let a = KExpr::::var(0, name("A")); + let mut info = a.info().clone(); + info.mdata = metadata("B"); + let b = KExpr::new(ExprData::Var(0, name("B"), info)); + let root = KExpr::app(a, b.clone()); + let expected = InternTable::new().intern_expr_reference(root.clone()); + let result = InternTable::new().intern_expr(root); + assert_same_expr(&result, &expected); + let ExprData::App(_, rhs, _) = result.data() else { panic!("app") }; + assert!(rhs.ptr_eq(&b)); +} + +#[test] +fn memo_does_not_survive_calls_or_environment_clears() { + let mut env = KEnv::::new(); + let leaf = KExpr::var(0, name("input")); + let root = expr_dag(leaf, 6); + for label in ["first", "second", "third"] { + env.clear(); + let seed = KExpr::var_mdata(0, name(label), metadata(label)); + env.intern.intern_expr(seed.clone()); + let result = env.intern.intern_expr(root.clone()); + let mut old = InternTable::new(); + old.intern_expr_reference(seed); + assert_same_expr(&result, &old.intern_expr_reference(root.clone())); + assert!(env.intern.intern_expr(root.clone()).ptr_eq(&result)); + } +} + +#[test] +#[ignore = "manual release benchmark; no wall-clock assertions"] +fn benchmark_shared_dag_interning() { + use std::time::Instant; + for depth in [12, 16, 20] { + let root = expr_dag(KExpr::sort(KUniv::zero()), depth); + let mut old = InternTable::::new(); + let mut new = InternTable::::new(); + old.intern_expr_reference(KExpr::sort(KUniv::zero())); + new.intern_expr(KExpr::sort(KUniv::zero())); + take_op_counts(); + let start = Instant::now(); + let expected = old.intern_expr_reference(root.clone()); + let old_time = start.elapsed(); + let old_visits = take_op_counts().intern_nodes; + let start = Instant::now(); + let result = new.intern_expr(root.clone()); + let new_time = start.elapsed(); + let new_visits = take_op_counts().intern_nodes; + assert_same_expr(&result, &expected); + eprintln!( + "depth={depth}: tree={old_time:?} ({old_visits} visits), \ + dag={new_time:?} ({new_visits} visits)" + ); + } +} + +#[test] +#[ignore = "manual release benchmark; no wall-clock assertions"] +fn benchmark_small_expression_interning() { + use std::hint::black_box; + use std::time::Instant; + for workload in ["canonical", "duplicate leaf", "duplicate app", "fresh app"] + { + for reference in [true, false] { + let intern: fn(&mut InternTable, KExpr) -> KExpr = + if reference { + InternTable::intern_expr_reference + } else { + InternTable::intern_expr + }; + let mut table = InternTable::new(); + let f = intern(&mut table, KExpr::var(1, ())); + let a = intern(&mut table, KExpr::var(0, ())); + let start = Instant::now(); + for i in 0..100_000 { + let input = match workload { + "canonical" => a.clone(), + "duplicate leaf" => KExpr::var(0, ()), + "duplicate app" => KExpr::app(f.clone(), a.clone()), + "fresh app" => { + let arg = intern(&mut table, KExpr::var(i + 2, ())); + KExpr::app(f.clone(), arg) + }, + _ => unreachable!(), + }; + black_box(intern(&mut table, input)); + } + eprintln!("{workload}, reference={reference}: {:?}", start.elapsed()); + } + } +} diff --git a/docs/ix_canonicity.md b/docs/ix_canonicity.md index e39fbbdb..0616e219 100644 --- a/docs/ix_canonicity.md +++ b/docs/ix_canonicity.md @@ -1384,13 +1384,19 @@ scheduler, on the work-stealing schedule itself. Both compilers must run aux synthesis in kernel contexts scoped to the block, and on egress restore source spellings structurally (`restore_source_names_same_content` and the source-name hint maps in -`expr_utils.rs`; `restoreSourceNamesSameContent` and mirrors in -`Ix/AuxGen/Kernel.lean`). The restoration heuristics — hint candidacy -restricted to `App`/`Proj` subterms, first-insert-wins hint slots — +`expr_utils.rs` / `source_name_hints.rs`; `restoreSourceNamesSameContent` +and mirrors in `Ix/AuxGen/Kernel.lean`). The restoration heuristics — +hint candidacy restricted to `App`/`Proj` subterms, first-insert-wins hint slots — are part of the specified behavior: they must remain mirrored hole-for-hole, since a restoration difference is a parity break even with block-scoped contexts (§17.8). +The Rust hint pass uses exact, shallow structural keys over canonical child +identities, not a lossy digest as evidence of equality. Its conversion and +restoration memos are scoped to one pass with a fixed view of referenced +name resolutions; first-wins collection precedes restoration against the +completed hint map. These temporary identities never enter serialized output. + **Positions in scope.** The rule governs the metadata positions that record a *reference to another constant*: @@ -2295,7 +2301,7 @@ Phase 2 paths already were). Remaining audit items: - **Restoration-heuristic mirror parity.** The kernel-egress restoration passes are heuristic: hint candidacy is `App`/`Proj` - only (`source_name_hint_candidate` in `expr_utils.rs` — a bare + only (`Pass::candidate` in `source_name_hints.rs` — a bare aliased `Const` surviving a genuinely-reducing WHNF is not restored), and hint slots are first-insert-wins (two same-address source subterms take the traversal-first spelling). The Lean @@ -2308,11 +2314,13 @@ Phase 2 paths already were). Remaining audit items: is an intern-assigned uid, fresh for every un-interned `to_kexpr_static` construction — so no restore-time key ever matched a collect-time key and the Rust hint pass restored nothing. - Fixed by `kexpr_content_key` (a pure name-erased structural digest - mirroring the `ExprKey`/`Ix.Tc` equivalence) and by making the - WHNF no-op test structural (`==`) rather than uid equality. Any - future keying change must preserve that both sides induce the - identical equivalence. + Originally fixed by `kexpr_content_key` (a pure name-erased structural + digest mirroring the `ExprKey`/`Ix.Tc` equivalence) and by making the + WHNF no-op test structural (`==`) rather than uid equality. The hint + pass now uses exact, pass-local structural identities instead of that + digest, with full shallow-key equality on hash collisions. Any future + keying change must preserve that both sides induce the identical + equivalence. - **Lean `nameForAddr` fallback.** `TcScopeSt.nameForAddr` (`Ix/AuxGen/Kernel.lean`) resolves a provisional kernel address by a linear scan of `cenv.nameToNamed` in `HashMap` iteration order — From 3f21465f791dc46583d681b697612c302ea2bfb0 Mon Sep 17 00:00:00 2001 From: "J. C. Burnham" Date: Sun, 6 Sep 2026 17:18:42 -0400 Subject: [PATCH 4/4] fix(compile): admit serial validation with limited headroom Allow validation to continue one job at a time when retained data leaves less than healthy headroom but the safety reserve remains available. Keep admissions blocked during reclaim pressure or insufficient headroom, and cover serial progress, retries, and reserve enforcement with regression tests. --- Benchmarks/Compile/README.md | 3 + crates/compile/src/compile/admission.rs | 12 +-- crates/compile/src/compile/memory.rs | 20 +++- crates/compile/src/compile/validation.rs | 125 ++++++++++++++++++++--- 4 files changed, 134 insertions(+), 26 deletions(-) diff --git a/Benchmarks/Compile/README.md b/Benchmarks/Compile/README.md index 3289687a..12d8550b 100644 --- a/Benchmarks/Compile/README.md +++ b/Benchmarks/Compile/README.md @@ -47,6 +47,9 @@ active jobs and ramps up while memory is healthy. The controller samples stall pressure every 250 ms. It stops admissions and cancels excess attempts under pressure; those attempts return normally, drop their scratch data, and retry once alone after other work finishes. Completed validations are retained. +When retained data leaves limited headroom, validation continues one job at a +time while the safety reserve is available and the system is not reclaiming +memory. Low headroom alone does not keep resetting the recovery delay. Expression walks preserve DAG sharing and check cancellation within the walk. This is a soft safety mechanism, not a hard allocation limit: an individual diff --git a/crates/compile/src/compile/admission.rs b/crates/compile/src/compile/admission.rs index 9039293b..c71838a9 100644 --- a/crates/compile/src/compile/admission.rs +++ b/crates/compile/src/compile/admission.rs @@ -12,7 +12,7 @@ use std::thread::{self, JoinHandle}; use std::time::{Duration, Instant}; use super::memory::{ - GIB, MIB, Memory, MemoryReader, Pressure, budget_from_env, pressure, + GIB, Memory, MemoryReader, Pressure, budget_from_env, pressure, resource_error, }; use ixon::CompileError; @@ -97,13 +97,8 @@ impl Policy { growth.max(self.growth_per_second.saturating_mul(3) / 4); self.projected_growth = self.growth_per_second.saturating_mul(options.lookahead.as_secs()); - let capacity = options - .process_budget - .map_or(memory.capacity, |b| b.min(memory.capacity)); - let reserve = (capacity / 8).max(64 * MIB).min(capacity / 2); - let headroom = options.process_budget.map_or(memory.available, |b| { - memory.available.min(b.saturating_sub(memory.process)) - }); + let reserve = memory.reserve(options.process_budget); + let headroom = memory.headroom(options.process_budget); self.pressure = pressure(memory, self.previous, elapsed, options.process_budget); let forecast_tight = @@ -349,6 +344,7 @@ impl Drop for Admission { #[cfg(test)] mod tests { use super::*; + use crate::compile::memory::MIB; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; fn memory() -> Memory { diff --git a/crates/compile/src/compile/memory.rs b/crates/compile/src/compile/memory.rs index 3152aa1d..1c909bad 100644 --- a/crates/compile/src/compile/memory.rs +++ b/crates/compile/src/compile/memory.rs @@ -54,6 +54,19 @@ pub(super) struct Memory { pub(super) stall_us: u64, } +impl Memory { + pub(super) fn headroom(self, budget: Option) -> u64 { + budget.map_or(self.available, |b| { + self.available.min(b.saturating_sub(self.process)) + }) + } + + pub(super) fn reserve(self, budget: Option) -> u64 { + let capacity = budget.map_or(self.capacity, |b| b.min(self.capacity)); + (capacity / 8).max(64 * MIB).min(capacity / 2) + } +} + pub(super) fn field(text: &str, key: &str) -> Option { text.lines().find_map(|line| { line.strip_prefix(key)?.split_whitespace().next()?.parse().ok() @@ -198,11 +211,8 @@ pub(super) fn pressure( elapsed: Duration, budget: Option, ) -> Pressure { - let capacity = budget.map_or(now.capacity, |b| b.min(now.capacity)); - let reserve = (capacity / 8).max(64 * MIB).min(capacity / 2); - let available = budget.map_or(now.available, |b| { - now.available.min(b.saturating_sub(now.process)) - }); + let reserve = now.reserve(budget); + let available = now.headroom(budget); if available < reserve / 2 { return Pressure::Critical; } diff --git a/crates/compile/src/compile/validation.rs b/crates/compile/src/compile/validation.rs index 7ec1b2b3..88f13042 100644 --- a/crates/compile/src/compile/validation.rs +++ b/crates/compile/src/compile/validation.rs @@ -17,7 +17,8 @@ use std::time::{Duration, Instant}; use ixon::CompileError; use super::memory::{ - GIB, MIB, Memory, MemoryReader, Pressure, pressure, resource_error, + GIB, MIB, Memory, MemoryReader, Pressure, pressure, reclaiming, + resource_error, }; #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -183,7 +184,6 @@ where read().ok_or_else(|| resource_error("memory telemetry unavailable"))?; let mut sampled = Instant::now(); let mut ramped = sampled; - let mut pressured_at = None; let mut idle_since = None; let mut logged = sampled; let mut completed = 0; @@ -191,20 +191,31 @@ where // Apply limits before the first allocation-heavy job, not only on tick 1. let mut state = pressure(previous, previous, options.tick, options.process_budget); + let mut blocked = previous.headroom(options.process_budget) + < previous.reserve(options.process_budget); + let mut pressured_at = blocked.then_some(sampled); + if state != Pressure::Healthy { + limit = 1; + } loop { let now = Instant::now(); if now.duration_since(sampled) >= options.tick { let previous_state = state; if let Some(memory) = read() { + blocked = memory.headroom(options.process_budget) + < memory.reserve(options.process_budget) + || reclaiming(memory, previous, now.duration_since(sampled)); state = pressure( memory, previous, now.duration_since(sampled), options.process_budget, ); - if matches!(state, Pressure::Backoff | Pressure::Critical) { + if blocked { pressured_at = Some(now); + } + if matches!(state, Pressure::Backoff | Pressure::Critical) { let running = active .iter() .filter(|j| !j.signals.cancel.load(Ordering::Relaxed)) @@ -224,6 +235,8 @@ where memory.process as f64 / GIB as f64 ))); } + } else if state == Pressure::Hold { + limit = 1; } else if state == Pressure::Healthy && pressured_at .is_none_or(|t| now.duration_since(t) >= options.recovery) @@ -276,10 +289,13 @@ where pressured_at.is_some_and(|t| now.duration_since(t) < options.recovery); let draining = active.iter().any(|j| j.signals.cancel.load(Ordering::Relaxed)); - let can_admit = - error.is_none() && !recovering && !draining && state == Pressure::Healthy; + // Retained environments can keep headroom below Healthy after every + // attempt has drained. With a reserve and no active reclaim pressure, + // continue serially rather than waiting for retained data to disappear. + let can_admit = error.is_none() && !recovering && !draining && !blocked; + let admission_limit = if state == Pressure::Healthy { limit } else { 1 }; if can_admit && !active.iter().any(|j| j.solo) { - while active.len() < limit { + while active.len() < admission_limit { let next = if !ready.is_empty() { ready.pop_front().map(|id| (id, false)) } else if active.is_empty() { @@ -318,8 +334,10 @@ where let start = *idle_since.get_or_insert(now); if now.duration_since(start) >= options.idle_timeout { return Err(resource_error(format!( - "not enough available memory to admit an inductive validation job; no progress for {:.1} seconds", - options.idle_timeout.as_secs_f64() + "not enough available memory to admit an inductive validation job; no progress for {:.1} seconds (headroom {:.1} GiB, reserve {:.1} GiB, pressure {state:?})", + options.idle_timeout.as_secs_f64(), + previous.headroom(options.process_budget) as f64 / GIB as f64, + previous.reserve(options.process_budget) as f64 / GIB as f64, ))); } } else { @@ -622,6 +640,46 @@ mod tests { assert_eq!(live.load(Ordering::SeqCst), 0); } + #[test] + fn retained_memory_allows_serial_progress_below_healthy_headroom() { + for headroom in [30 * GIB, 20 * GIB] { + for budget in [None, Some(100 * GIB)] { + let memory = Memory { + capacity: if budget.is_some() { 200 * GIB } else { 100 * GIB }, + available: if budget.is_some() { 100 * GIB } else { headroom }, + process: 100 * GIB - headroom, + ..healthy() + }; + let live = AtomicUsize::new(0); + let peak = AtomicUsize::new(0); + let calls: Vec<_> = (0..8).map(|_| AtomicUsize::new(0)).collect(); + let opt = + Options { initial_workers: 4, process_budget: budget, ..options(4) }; + pool(4) + .install(|| { + run_with( + &calls, + &|calls: &AtomicUsize, c: &Checkpoint| { + calls.fetch_add(1, Ordering::SeqCst); + let n = live.fetch_add(1, Ordering::SeqCst) + 1; + let _scratch = Scratch { live: &live, dropping: None }; + peak.fetch_max(n, Ordering::SeqCst); + std::thread::sleep(Duration::from_millis(10)); + c.visit()?; + Ok(()) + }, + opt, + || Some(memory), + ) + }) + .unwrap(); + assert_eq!(peak.load(Ordering::SeqCst), 1); + assert!(calls.iter().all(|n| n.load(Ordering::SeqCst) == 1)); + assert_eq!(live.load(Ordering::SeqCst), 0); + } + } + } + #[test] fn cancellation_drops_scratch_before_admitting_and_retries_once_alone() { let live = AtomicUsize::new(0); @@ -665,11 +723,17 @@ mod tests { }, opt, || { - Some(if live.load(Ordering::SeqCst) >= 2 { - Memory { available: 20 * GIB, ..healthy() } - } else { - healthy() - }) + // Retained data keeps headroom below Healthy even after the + // cancelled attempt releases its scratch. Retries must still run. + Some( + if live.load(Ordering::SeqCst) >= 2 + || cancelled.load(Ordering::SeqCst) + { + Memory { available: 20 * GIB, ..healthy() } + } else { + healthy() + }, + ) }, ) }) @@ -768,6 +832,41 @@ mod tests { assert_eq!(calls.load(Ordering::SeqCst), 0); } + #[test] + fn serial_admission_preserves_the_reserve_and_waits_out_reclaim_pressure() { + for reclaiming in [false, true] { + let calls = AtomicUsize::new(0); + let mut reads = 0; + let result = pool(1).install(|| { + run_with( + &[()], + &|_: &(), _: &Checkpoint| { + calls.fetch_add(1, Ordering::SeqCst); + Ok(()) + }, + options(1), + || { + reads += 1; + Some(Memory { + // Start below the admission reserve but above Critical. + // Later samples can have room for serial work but must not + // admit it while the system continues swapping. + available: if reclaiming && reads > 1 { + 20 * GIB + } else { + 8 * GIB + }, + swap_used: if reclaiming { reads * 9 * MIB } else { 0 }, + ..healthy() + }) + }, + ) + }); + assert!(matches!(result, Err(CompileError::ResourceLimit { .. }))); + assert_eq!(calls.load(Ordering::SeqCst), 0); + } + } + #[test] fn unavailable_telemetry_cancels_and_drains_active_work() { let live = AtomicUsize::new(0);