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..12d8550b 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,100 @@ 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` + +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. +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 +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: + +```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 +120,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/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/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"])], 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..c71838a9 --- /dev/null +++ b/crates/compile/src/compile/admission.rs @@ -0,0 +1,573 @@ +//! 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, 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 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 = + 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 crate::compile::memory::MIB; + 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..6e9036f2 100644 --- a/crates/compile/src/compile/aux_gen.rs +++ b/crates/compile/src/compile/aux_gen.rs @@ -77,7 +77,12 @@ pub mod below; pub mod brecon; pub mod cases_on; +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/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..84b51fcd 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::{ @@ -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 // ========================================================================= @@ -521,71 +528,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 +557,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 +578,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 +656,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 +681,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 +701,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 +1332,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. @@ -2739,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, ) } } @@ -2817,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 @@ -2875,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 @@ -3323,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()) { ( @@ -3353,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), ) }, ( @@ -3362,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(), ), ( @@ -3371,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(), ), ( @@ -3380,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, ), ( @@ -3394,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(), @@ -3422,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); @@ -3440,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()) @@ -3460,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, _) => { @@ -3503,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()); @@ -3539,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/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/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/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..1c909bad --- /dev/null +++ b/crates/compile/src/compile/memory.rs @@ -0,0 +1,226 @@ +//! 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, +} + +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() + }) +} +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 reserve = now.reserve(budget); + let available = now.headroom(budget); + 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..88f13042 --- /dev/null +++ b/crates/compile/src/compile/validation.rs @@ -0,0 +1,944 @@ +//! 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, reclaiming, + 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 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); + 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 blocked { + pressured_at = Some(now); + } + if matches!(state, Pressure::Backoff | Pressure::Critical) { + 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::Hold { + limit = 1; + } 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)); + // 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() < admission_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 (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 { + 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 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); + 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, + || { + // 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() + }, + ) + }, + ) + }) + .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 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); + 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/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/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}"), } } } 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 —