perf: Tune ix shard counts and recalibrate for proving via --ram-budget

Ix/Aiur/Protocol.lean

@@ -173,54 +173,153 @@ def proveAddrWithEnv (system : @& AiurSystem)
   (proveAddrWithEnv' system funIdx envHandle addrBytes useBytecode).map
     fun r => (r.claimBytes, r.proof, .ofArrays r.ioData r.ioMap)
 
+/-- Result of a per-shard prove: the claim's wire bytes, the proof, the
+    projected prover RAM peak of the record that produced it
+    (`AiurSystem::peak_prove_bytes`), and the part count the peak model
+    projects will fit the budget
+    (`AiurSystem::suggested_split_parts`).
+
+    `proof` is `none` exactly when the peak exceeded the budget — a
+    RESULT rather than an error, since the caller's answer is to cut
+    the shard into `suggestedParts` parts and prove those. The count is
+    computed Rust-side because only there does the executed record
+    still exist to read per-circuit heights from; it is optimistic
+    (parts re-execute dependencies shared across the cut), so each part
+    must still be gated on its own record. `suggestedParts` is 1
+    whenever the prove ran. The claim bytes are filled either way (the
+    claim is known before proving starts).
+
+    The final IO buffer is not returned — it is the shard's whole
+    ingested byte scope and no caller reads it. -/
+structure ShardProveResult where
+  claimBytes : ByteArray
+  proof : Option Proof
+  peakBytes : Nat
+  suggestedParts : Nat
+
 @[extern "rs_aiur_system_shard_prove_with_env"]
 private opaque shardProveWithEnv' : @& AiurSystem →
-  @& Bytecode.FunIdx → @& EnvHandle → @& ByteArray →
-    Except String ProveEnvResult
-
-/-- Per-shard prove against a Rust-owned `EnvHandle`. -/
+  @& Bytecode.FunIdx → @& EnvHandle → @& ByteArray → @& Nat → Bool →
+    Except String ShardProveResult
+
+/-- Per-shard prove against a Rust-owned `EnvHandle`: ONE execution,
+    whose record is proven from directly.
+
+    `maxRamBytes` is a per-shard prover-RAM budget checked against that
+    record's projected peak before the witness phase begins; `0` means
+    detect (85% of `MemAvailable`, the policy the check batch's RAM gate
+    uses), and an unreadable `/proc/meminfo` disables the check rather
+    than guessing. Over budget, the record is dropped and `proof` is
+    `none` — learning that here costs one execution instead of an OOM
+    part-way through an FFT. The peak comes back either way, so a prove
+    run yields the same split/merge signal a check run does.
+
+    `execOnly` stops after execution + measurement (`proof` is `none`
+    either way; `suggestedParts` is 1 exactly when the peak fits): the
+    split loop runs on executions alone, never starting a STARK. -/
 def shardProveWithEnv (system : @& AiurSystem)
-  (funIdx : @& Bytecode.FunIdx) (envHandle : @& EnvHandle) (ownedBlob : ByteArray) :
-    Except String (ByteArray × Proof × IOBuffer) :=
-  (shardProveWithEnv' system funIdx envHandle ownedBlob).map
-    fun r => (r.claimBytes, r.proof, .ofArrays r.ioData r.ioMap)
+  (funIdx : @& Bytecode.FunIdx) (envHandle : @& EnvHandle)
+  (ownedBlob : ByteArray) (maxRamBytes : Nat := 0)
+  (execOnly : Bool := false) :
+    Except String ShardProveResult :=
+  shardProveWithEnv' system funIdx envHandle ownedBlob maxRamBytes execOnly
 
 @[extern "rs_aiur_system_verify"]
 opaque verify : @& AiurSystem →
   @& Array G → @& Proof → Except String Unit
 
 end AiurSystem
 
+/-- Write a `.ixes` manifest for an EXPLICIT partition — the block lists
+    a run actually produced (splits included) rather than a planner's
+    output. `shardsBlob`: per shard, a 4-byte LE block count followed by
+    that many 32-byte block addresses; every env block must appear in
+    exactly one shard. `peaksBlob`: one 8-byte LE measured prover peak
+    per shard in order, recorded on the manifest for schedulers. Own
+    sizes, foreign blocks, cross-ingress and assumption roots are
+    recomputed from the env's static profile; prints the manifest
+    summary to stderr. -/
+@[extern "rs_shard_manifest_from_partition"]
+opaque shardManifestFromPartition : @& EnvHandle →
+  @& ByteArray → @& ByteArray → @& String → IO Unit
+
 namespace Bytecode.Toplevel
 
+/-- One shard's result from `shardCheckBatchWithEnv`. `weights` is the
+    shard's per-constant virtual-gas table, packed as 48-byte rows —
+    32-byte address, then `vspan` and `mult` as little-endian `UInt64`s
+    (see `ShardResult.foldWeights`). Empty unless the batch ran with
+    `profile := true`; the record it is read from is reduced to these
+    rows and dropped inside the shard's own task, so a whole partition's
+    weights fit in RAM when its records could not. -/
+structure ShardResult where
+  error : String
+  peakBytes : Nat
+  weights : ByteArray
+  /-- 1 when `peakBytes` fits the batch's `maxRamBytes` (or no budget
+      was given); otherwise the part count the peak model projects will
+      fit (`AiurSystem::suggested_split_parts`, measured on the record
+      in-task). -/
+  suggestedParts : Nat
+  deriving Inhabited
+
+/-- Bytes per packed `weights` row: 32 address + 8 vspan + 8 mult. -/
+def shardWeightRow : Nat := 48
+
+/-- Fold `f` over the packed `(addrBytes, vspan, mult)` rows of `weights`.
+    Trailing bytes that do not complete a row are ignored. -/
+@[inline] def ShardResult.foldWeights {α : Type} (r : ShardResult) (init : α)
+    (f : α → ByteArray → UInt64 → UInt64 → α) : α :=
+  go (r.weights.size / shardWeightRow) 0 init
+where
+  go : Nat → Nat → α → α
+    | 0, _, acc => acc
+    | rows + 1, off, acc =>
+      go rows (off + shardWeightRow) <| f acc
+        (r.weights.extract off (off + 32))
+        (r.weights.extract (off + 32) (off + 40)).toUInt64LE!
+        (r.weights.extract (off + 40) (off + 48)).toUInt64LE!
+
 @[extern "rs_aiur_toplevel_shard_check_batch"]
 private opaque shardCheckBatchWithEnv' : @& Bytecode.Toplevel →
   @& Bytecode.FunIdx → @& EnvHandle → @& ByteArray → Bool → @& Nat →
-  @& CommitmentParameters → @& FriParameters →
-    Except String (Array (String × Nat))
+  @& CommitmentParameters → @& FriParameters → Bool → @& Nat → @& Nat →
+    Except String (Array ShardResult)
 
 /-- Check EVERY shard of a partition in one call: rayon over the shard
     list with true work-stealing (no chunk barriers), each shard
     through the exact single-shard machinery over its own private
     record and witness io. `shardsBlob` encodes, per shard, a 4-byte LE
     owned-constant count followed by that many 32-byte addresses.
-    Returns one `(error, peakBytes)` pair per shard in shard order:
-    empty error = clean, and `peakBytes` is the analytic prover RAM
-    peak ([`AiurSystem::peak_prove_bytes`] Rust-side) of the shard's
-    executed record — the split/merge input (0 on failure).
+    Returns one `ShardResult` per shard in shard order: empty error =
+    clean, and `peakBytes` is the analytic prover RAM peak
+    ([`AiurSystem::peak_prove_bytes`] Rust-side) of the shard's executed
+    record — the split/merge input (0 on failure). `profile` turns on the
+    virtual-gas meter and fills each result's `weights` with the shard's
+    per-constant cost rows, read off the record inside the shard's own
+    task; `checkConstIdx` names the `check_const` function whose queries
+    those rows come from (ignored when not profiling).
     `jobs = 0` uses rayon's default pool width (all cores): peak RSS
     is bounded by the Rust-side RAM gate (a byte-weighted admission
     semaphore over estimated per-shard execution RSS vs available
     system RAM), not by thread count — pass `jobs` only to narrow
-    CPU use. -/
+    CPU use.
+
+    `maxRamBytes > 0` is a per-shard prover-RAM budget: each result's
+    `suggestedParts` is 1 when its peak fits and the model's projected
+    part count otherwise, so a caller can cut over-budget shards and
+    re-batch the parts — the wave loop that audits a partition's split
+    behavior on executions alone. -/
 def shardCheckBatchWithEnv (toplevel : @& Bytecode.Toplevel)
   (funIdx : @& Bytecode.FunIdx) (envHandle : @& EnvHandle)
   (shardsBlob : ByteArray) (useBytecode : Bool := false) (jobs : Nat := 0)
   (commitmentParameters : CommitmentParameters := defaultCommitmentParameters)
   (friParameters : FriParameters := defaultFriParameters)
-  : Except String (Array (String × Nat)) :=
+  (profile : Bool := false) (checkConstIdx : Nat := 0)
+  (maxRamBytes : Nat := 0)
+  : Except String (Array ShardResult) :=
   shardCheckBatchWithEnv' toplevel funIdx envHandle shardsBlob useBytecode
-    jobs commitmentParameters friParameters
+    jobs commitmentParameters friParameters profile checkConstIdx maxRamBytes
 
 end Bytecode.Toplevel
 

Ix/Aiur/Semantics/BytecodeFfi.lean

@@ -75,14 +75,61 @@ structure QueryCount where
   totalHits : Nat
   deriving Inhabited
 
+-- ===========================================================================
+-- QueryRecordHandle: the execution's Rust-owned `QueryRecord` exposed to
+-- Lean as an opaque handle inside every `ExecuteResult`. Lean inspects
+-- queries with its higher-level knowledge (fn names via `nameMap`,
+-- layouts, addr→name maps) through the accessors below; the record —
+-- the dominant RAM consumer of an execution — frees when Lean drops it.
+-- ===========================================================================
+
+private opaque QueryRecordHandleNonempty : NonemptyType
+def QueryRecordHandle : Type := QueryRecordHandleNonempty.type
+instance : Nonempty QueryRecordHandle := QueryRecordHandleNonempty.property
+
+namespace QueryRecordHandle
+
+/-- Number of query entries registered for function `funIdx`. -/
+@[extern "rs_aiur_qr_fn_len"]
+opaque fnLen : @& QueryRecordHandle → (funIdx : @& Nat) → Nat
+
+/-- Input tuple of entry `i` of function `funIdx` (empty when out of
+    range). Entries are in insertion order; iterate within `fnLen`. -/
+@[extern "rs_aiur_qr_fn_key"]
+opaque fnKey : @& QueryRecordHandle → (funIdx i : @& Nat) → Array G
+
+/-- Multiplicity of entry `i` of function `funIdx`. `0` marks an
+    unconstrained hint row (or an out-of-range index). -/
+@[extern "rs_aiur_qr_fn_mult"]
+opaque fnMult : @& QueryRecordHandle → (funIdx i : @& Nat) → Nat
+
+/-- Recorded virtual-gas span of entry `i` of function `funIdx` — the
+    entry's standalone cost, with memo hits replaying the callee's
+    recorded span. `0` unless the execution ran with `profile := true`
+    (the `Full` execution wrappers' trailing flag). -/
+@[extern "rs_aiur_qr_fn_vspan"]
+opaque fnVspan : @& QueryRecordHandle → (funIdx i : @& Nat) → Nat
+
+/-- Stored values of entry `ptr` in the `memory[size]` map — the deref
+    of an Aiur `&[T; size]` pointer (empty when unbound). E.g. an
+    `Addr = &[U8; 32]` argument derefs to its 32 blake3 bytes via
+    `memKey 32 ptr`. -/
+@[extern "rs_aiur_qr_mem_key"]
+opaque memKey : @& QueryRecordHandle → (size ptr : @& Nat) → Array G
+
+end QueryRecordHandle
+
 /-- Result of an execution FFI call, built directly by Rust
 (`LeanAiurExecuteResult` in `crates/ffi/src/lean.rs`). `ioData`/`ioMap`
-are the flattened `IOBuffer`; see `IOBuffer.ofArrays`. -/
+are the flattened `IOBuffer`; see `IOBuffer.ofArrays`. `record` owns the
+execution's `QueryRecord` — drop the `ExecuteResult` (or project the
+other fields out) before the next heavy execution to release its RAM. -/
 structure ExecuteResult where
   output : Array G
   ioData : Array (G × Array G)
   ioMap : Array ((G × Array G) × IOKeyInfo)
   queryCounts : Array QueryCount
+  record : QueryRecordHandle
 
 -- ===========================================================================
 -- EnvHandle: Rust-owned `ixon::Env` exposed to Lean as an opaque handle.
@@ -122,6 +169,7 @@ private opaque execute' : @& Bytecode.Toplevel →
   @& Bytecode.FunIdx → @& Array G →
   (ioData : @& Array (G × Array G)) →
   (ioMap : @& Array ((G × Array G) × IOKeyInfo)) →
+  (profile : Bool) →
     Except String ExecuteResult
 
 /-- Executes the bytecode function `funIdx` with the given `args` and `ioBuffer`,
@@ -132,14 +180,24 @@ callers can recover instead of crashing. -/
 def execute (toplevel : @& Bytecode.Toplevel)
   (funIdx : @& Bytecode.FunIdx) (args : @& Array G) (ioBuffer : IOBuffer) :
     Except String (Array G × IOBuffer × Array QueryCount) :=
-  (execute' toplevel funIdx args ioBuffer.data.toArray ioBuffer.map.toArray).map
+  (execute' toplevel funIdx args ioBuffer.data.toArray ioBuffer.map.toArray false).map
     fun r => (r.output, .ofArrays r.ioData r.ioMap, r.queryCounts)
 
+/-- `execute` variant returning the full `ExecuteResult`, including the
+    `QueryRecordHandle` for post-run query inspection. `profile` enables
+    per-entry virtual-span recording (`QueryRecordHandle.fnVspan`). -/
+def executeFull (toplevel : @& Bytecode.Toplevel)
+  (funIdx : @& Bytecode.FunIdx) (args : @& Array G) (ioBuffer : IOBuffer)
+  (profile : Bool := false) :
+    Except String ExecuteResult :=
+  execute' toplevel funIdx args ioBuffer.data.toArray ioBuffer.map.toArray profile
+
 @[extern "rs_aiur_toplevel_execute_ixvm"]
 private opaque executeIxVM' : @& Bytecode.Toplevel →
   @& Bytecode.FunIdx → @& Array G →
   (ioData : @& Array (G × Array G)) →
   (ioMap : @& Array ((G × Array G) × IOKeyInfo)) →
+  (profile : Bool) →
     Except String ExecuteResult
 
 /-- IxVM-native execution: same shape as `execute`, but routes the
@@ -153,9 +211,19 @@ private opaque executeIxVM' : @& Bytecode.Toplevel →
 def executeIxVM (toplevel : @& Bytecode.Toplevel)
   (funIdx : @& Bytecode.FunIdx) (args : @& Array G) (ioBuffer : IOBuffer) :
     Except String (Array G × IOBuffer × Array QueryCount) :=
-  (executeIxVM' toplevel funIdx args ioBuffer.data.toArray ioBuffer.map.toArray).map
+  (executeIxVM' toplevel funIdx args ioBuffer.data.toArray ioBuffer.map.toArray false).map
     fun r => (r.output, .ofArrays r.ioData r.ioMap, r.queryCounts)
 
+/-- `executeIxVM` variant returning the full `ExecuteResult`, including
+    the `QueryRecordHandle` for post-run query inspection. `profile`
+    enables per-entry virtual-span recording
+    (`QueryRecordHandle.fnVspan`). -/
+def executeIxVMFull (toplevel : @& Bytecode.Toplevel)
+  (funIdx : @& Bytecode.FunIdx) (args : @& Array G) (ioBuffer : IOBuffer)
+  (profile : Bool := false) :
+    Except String ExecuteResult :=
+  executeIxVM' toplevel funIdx args ioBuffer.data.toArray ioBuffer.map.toArray profile
+
 /-- MultiStark-native execution of `verify_multi_stark_proof`: the IO
     advice buffer (channel 0 = proof, 1 = vk, 2 = claims, key `[0]`
     each) is built natively in Rust from the raw byte blobs — no
@@ -181,7 +249,8 @@ opaque executeMultiStark (toplevel : @& Bytecode.Toplevel)
 
 @[extern "rs_aiur_toplevel_check_addr_with_env"]
 private opaque checkAddrWithEnv' : @& Bytecode.Toplevel →
-  @& Bytecode.FunIdx → @& EnvHandle → @& ByteArray → Bool →
+  @& Bytecode.FunIdx → @& EnvHandle → @& ByteArray →
+  (useBytecode profile : Bool) →
     Except String ExecuteResult
 
 /-- Per-claim check against a Rust-owned `EnvHandle`. `useBytecode`
@@ -193,7 +262,7 @@ def checkAddrWithEnv (toplevel : @& Bytecode.Toplevel)
   (funIdx : @& Bytecode.FunIdx) (envHandle : @& EnvHandle)
   (addrBytes : ByteArray) (useBytecode : Bool := false)
   : Except String (Array G × IOBuffer × Array QueryCount) :=
-  (checkAddrWithEnv' toplevel funIdx envHandle addrBytes useBytecode).map
+  (checkAddrWithEnv' toplevel funIdx envHandle addrBytes useBytecode false).map
     fun r => (r.output, .ofArrays r.ioData r.ioMap, r.queryCounts)
 
 @[extern "rs_aiur_toplevel_check_addrs_with_env"]
@@ -207,7 +276,8 @@ private opaque checkAddrsWithEnv' : @& Bytecode.Toplevel →
     through the exact single-claim machinery over task-private data
     (own witness io, own query record), nothing shared between tasks
     but the read-only toplevel and env. Returns the FAILURES as
-    `(addrHex, error)` pairs; empty means all passed. Per-claim
+    `(batch index as a decimal string, error)` pairs, resolving back to
+    the caller's label order; empty means all passed. Per-claim
     outputs/records are not returned — use `checkAddrWithEnv` for a
     single claim's full result. -/
 def checkAddrsWithEnv (toplevel : @& Bytecode.Toplevel)
@@ -216,9 +286,21 @@ def checkAddrsWithEnv (toplevel : @& Bytecode.Toplevel)
   : Except String (Array (String × String)) :=
   checkAddrsWithEnv' toplevel funIdx envHandle addrsBlob useBytecode jobs
 
+/-- `checkAddrWithEnv` variant returning the full `ExecuteResult`,
+    including the `QueryRecordHandle` for post-run query inspection.
+    `profile` enables per-entry virtual-span recording
+    (`QueryRecordHandle.fnVspan`). -/
+def checkAddrWithEnvFull (toplevel : @& Bytecode.Toplevel)
+  (funIdx : @& Bytecode.FunIdx) (envHandle : @& EnvHandle)
+  (addrBytes : ByteArray) (useBytecode : Bool := false)
+  (profile : Bool := false)
+  : Except String ExecuteResult :=
+  checkAddrWithEnv' toplevel funIdx envHandle addrBytes useBytecode profile
+
 @[extern "rs_aiur_toplevel_shard_check_with_env"]
 private opaque shardCheckWithEnv' : @& Bytecode.Toplevel →
-  @& Bytecode.FunIdx → @& EnvHandle → @& ByteArray → Bool →
+  @& Bytecode.FunIdx → @& EnvHandle → @& ByteArray →
+  (useBytecode profile : Bool) →
     Except String ExecuteResult
 
 /-- Per-shard check with the witness shape (wrapper-augmented byte
@@ -228,9 +310,20 @@ def shardCheckWithEnv (toplevel : @& Bytecode.Toplevel)
   (funIdx : @& Bytecode.FunIdx) (envHandle : @& EnvHandle)
   (ownedBlob : ByteArray) (useBytecode : Bool := false)
   : Except String (Array G × IOBuffer × Array QueryCount) :=
-  (shardCheckWithEnv' toplevel funIdx envHandle ownedBlob useBytecode).map
+  (shardCheckWithEnv' toplevel funIdx envHandle ownedBlob useBytecode false).map
     fun r => (r.output, .ofArrays r.ioData r.ioMap, r.queryCounts)
 
+/-- `shardCheckWithEnv` variant returning the full `ExecuteResult`,
+    including the `QueryRecordHandle` for post-run query inspection.
+    `profile` enables per-entry virtual-span recording
+    (`QueryRecordHandle.fnVspan`). -/
+def shardCheckWithEnvFull (toplevel : @& Bytecode.Toplevel)
+  (funIdx : @& Bytecode.FunIdx) (envHandle : @& EnvHandle)
+  (ownedBlob : ByteArray) (useBytecode : Bool := false)
+  (profile : Bool := false)
+  : Except String ExecuteResult :=
+  shardCheckWithEnv' toplevel funIdx envHandle ownedBlob useBytecode profile
+
 end Bytecode.Toplevel
 
 end Aiur