Sharding: per-shard check / prove / verify on the IxVM kernel

Ix/Cli/CheckCmd.lean

@@ -242,6 +242,190 @@ def forEachClaim
     for n in failures do IO.eprintln s!"  {n}"
     pure 1
 
+-- Bounds-checked little-endian cursor over a `.ixes` byte buffer. Every read
+-- verifies the buffer has enough bytes and returns `.error` on underflow — no
+-- `b[p]!` panics on a truncated/malformed manifest.
+private abbrev IxesP := StateT (ByteArray × Nat) (Except String)
+
+private def ixesU8 : IxesP UInt8 := do
+  let (b, p) ← get
+  if h : p < b.size then
+    modify (fun _ => (b, p + 1))
+    pure b[p]
+  else throw "ixes: truncated (expected a byte)"
+
+private def ixesU32 : IxesP UInt32 := do
+  let a ← ixesU8; let b ← ixesU8; let c ← ixesU8; let d ← ixesU8
+  pure (a.toUInt32 ||| (b.toUInt32 <<< 8) ||| (c.toUInt32 <<< 16) ||| (d.toUInt32 <<< 24))
+
+private def ixesSkip (n : Nat) : IxesP Unit := do
+  let (b, p) ← get
+  if p + n ≤ b.size then modify (fun _ => (b, p + n)) else throw s!"ixes: truncated (skip {n})"
+
+private def ixesAddr : IxesP Address := do
+  let (b, p) ← get
+  if p + 32 ≤ b.size then modify (fun _ => (b, p + 32)); pure ⟨b.extract p (p + 32)⟩
+  else throw "ixes: truncated (expected a 32-byte address)"
+
+/-- Parse every shard's owned block addresses from a serialized `.ixes`
+    manifest (`ShardManifest::to_bytes`, `src/ix/shard.rs`):
+    magic(8) ‖ total_cross_ingress(u128) ‖ num_shards(u32) ‖ per shard
+    { id(u32) ‖ heartbeats(u64) ‖ own_size(u64) ‖ cross_ingress(u64) ‖
+      assumption_root(u8 tag + 32?) ‖ blocks(u32 len + 32·len) ‖
+      foreign_blocks(u32 len + 32·len) }.
+    Bounds-checked: a truncated/malformed file yields `.error`, never a panic. -/
+def parseIxesAllShards (bytes : ByteArray) : Except String (Array (Array Address)) :=
+  let go : IxesP (Array (Array Address)) := do
+    let m0 ← ixesU8; let m1 ← ixesU8; let m2 ← ixesU8; let m3 ← ixesU8
+    if !(m0 == 0x49 && m1 == 0x58 && m2 == 0x45 && m3 == 0x53) then
+      throw "not an .ixes file (bad magic)"
+    ixesSkip 4    -- rest of the 8-byte magic
+    ixesSkip 16   -- total_cross_ingress (u128)
+    let n ← ixesU32
+    let mut shards : Array (Array Address) := #[]
+    for _ in [0:n.toNat] do
+      ixesSkip (4 + 8 + 8 + 8)  -- id + heartbeats + own_size + cross_ingress
+      if (← ixesU8) == 1 then ixesSkip 32  -- assumption_root present
+      let blen ← ixesU32
+      let mut blocks : Array Address := #[]
+      for _ in [0:blen.toNat] do
+        blocks := blocks.push (← ixesAddr)
+      ixesSkip ((← ixesU32).toNat * 32)  -- skip foreign_blocks
+      shards := shards.push blocks
+    pure shards
+  go.run' (bytes, 0)
+
+/-- The check-schedule block address of a constant: a projection collapses
+    to its SCC/Muts wrapper (`p.block`); everything else is its own block.
+    Mirrors `check_schedule_block_addr` (`src/ffi/kernel.rs`). -/
+private def blockAddrOf (addr : Address) (c : Ixon.Constant) : Address :=
+  match c.info with
+  | .iPrj prj => prj.block
+  | .cPrj prj => prj.block
+  | .rPrj prj => prj.block
+  | .dPrj prj => prj.block
+  | _ => addr
+
+/-- Owned constants of a shard: every env constant whose check-schedule block
+    is in `blocks`. -/
+def ownedConstsForBlocks (ixonEnv : Ixon.Env) (blocks : Array Address) : Array Address := Id.run do
+  let blockSet : Std.HashSet Address := blocks.foldl (·.insert ·) {}
+  let mut o : Array Address := #[]
+  for (addr, c) in ixonEnv.consts do
+    if blockSet.contains (blockAddrOf addr c) then o := o.push addr
+  return o
+
+/-- The `CheckEnv` claim digest a shard's proof commits to — reconstructed
+    deterministically from the env + the shard's owned blocks. Matches the
+    digest `prove --shard K` produced, so a proof can be bound to its shard. -/
+def shardClaimDigest (ixonEnv : Ixon.Env) (blocks : Array Address) : Except String Address := do
+  let (claim, _, _) ← IxVM.ClaimHarness.shardCheckEnvClaim ixonEnv (ownedConstsForBlocks ixonEnv blocks)
+  pure (Address.blake3 (Ix.Claim.ser claim))
+
+/-- Load the `.ixe` env and the `.ixes` shard partition together (each file
+    read once). Shared by every manifest-driven shard path. -/
+def loadEnvAndShards (manifestPath ixePath : String) :
+    IO (Except String (Ixon.Env × Array (Array Address))) := do
+  match parseIxesAllShards (← IO.FS.readBinFile manifestPath) with
+  | .error e => return .error s!"manifest parse failed: {e}"
+  | .ok shards => match Ixon.rsDeEnv (← IO.FS.readBinFile ixePath) with
+    | .error e => return .error s!"deserialize {ixePath} failed: {e}"
+    | .ok env => return .ok (env, shards)
+
+/-- Run the shard operation for one shard, given the already-loaded env and the
+    shard's owned `blocks`: build the `CheckEnv` witness over the owned consts
+    (ingress their closure, skip the frontier) and dispatch `runOne`. -/
+def runShardOwned (ixonEnv : Ixon.Env) (blocks : Array Address) (shardK : Nat)
+    (runOne : Ix.Claim → IxVM.ClaimHarness.ClaimWitness → String → IO UInt32) : IO UInt32 := do
+  let owned := ownedConstsForBlocks ixonEnv blocks
+  IO.println s!"[shard] shard {shardK}: {blocks.size} owned blocks → \
+    {owned.size}/{ixonEnv.consts.size} owned consts"
+  match IxVM.ClaimHarness.buildShardCheckEnvWitness ixonEnv owned with
+  | .error e => IO.eprintln s!"shard witness build failed: {e}"; return 1
+  | .ok (claim, witness) => runOne claim witness s!"shard {shardK}"
+
+/-- Manifest-driven check/prove of one shard `shardK` of the partition. -/
+def runShardCheckManifest (manifestPath ixePath : String) (shardK : Nat)
+    (runOne : Ix.Claim → IxVM.ClaimHarness.ClaimWitness → String → IO UInt32) : IO UInt32 := do
+  match (← loadEnvAndShards manifestPath ixePath) with
+  | .error e => IO.eprintln e; return 1
+  | .ok (ixonEnv, shards) => match shards[shardK]? with
+    | none => IO.eprintln s!"shard {shardK} out of range ({shards.size} shards)"; return 1
+    | some blocks => runShardOwned ixonEnv blocks shardK runOne
+
+/-- Run the shard operation over EVERY shard — the whole-partition behavior of
+    `--ixes` with no `--shard` (used by `prove`). Loads the env once. Returns 1
+    if any shard fails, else 0. -/
+def runShardManifestAll (manifestPath ixePath : String)
+    (runOne : Ix.Claim → IxVM.ClaimHarness.ClaimWitness → String → IO UInt32) : IO UInt32 := do
+  match (← loadEnvAndShards manifestPath ixePath) with
+  | .error e => IO.eprintln e; return 1
+  | .ok (ixonEnv, shards) =>
+    let mut rc : UInt32 := 0
+    for (blocks, k) in shards.mapIdx (fun k b => (b, k)) do
+      if (← runShardOwned ixonEnv blocks k runOne) != 0 then rc := 1
+    pure rc
+
+/-- Coverage check over already-loaded env + shards: every constant's
+    check-schedule block is owned by **exactly one** shard. That is the whole
+    soundness condition for the check case — each constant is type-checked
+    once, and every shard's frontier (closure minus owned) is therefore owned
+    (checked) by some other shard. Prints the per-shard report; returns whether
+    the partition is a valid disjoint cover (no block owned by two shards, no
+    constant whose block no shard owns). -/
+def shardsCover (ixonEnv : Ixon.Env) (shards : Array (Array Address)) : IO Bool := do
+  -- block → shard, detecting blocks claimed by more than one shard.
+  let mut blockToShard : Std.HashMap Address Nat := {}
+  let mut dup : Nat := 0
+  for (blocks, k) in shards.mapIdx (fun k b => (b, k)) do
+    for blk in blocks do
+      match blockToShard.get? blk with
+      | some _ => dup := dup + 1
+      | none => blockToShard := blockToShard.insert blk k
+  -- assign every const to a shard via its block; count + detect unowned.
+  let mut counts : Array Nat := Array.replicate shards.size 0
+  let mut unowned : Nat := 0
+  for (addr, c) in ixonEnv.consts do
+    match blockToShard.get? (blockAddrOf addr c) with
+    | some k => counts := counts.modify k (· + 1)  -- total: no-op if out of range
+    | none => unowned := unowned + 1
+  IO.println s!"[shards] {shards.size} shards, {ixonEnv.consts.size} consts"
+  for (blocks, k) in shards.mapIdx (fun k b => (b, k)) do
+    IO.println s!"  shard {k}: {blocks.size} blocks, {(counts[k]?).getD 0} consts"
+  if dup != 0 then
+    IO.eprintln s!"[shards] FAIL: {dup} block(s) owned by >1 shard (not disjoint)"
+  if unowned != 0 then
+    IO.eprintln s!"[shards] FAIL: {unowned} const(s) with no owning shard (coverage gap)"
+  let ok := dup == 0 && unowned == 0
+  if ok then
+    IO.println s!"[shards] OK: partition covers all {ixonEnv.consts.size} consts, disjoint"
+  pure ok
+
+/-- Check EVERY shard of the partition concurrently (shards are independent
+    bytecode runs) after verifying coverage — the whole-partition behavior of
+    `check --ixes` with no `--shard`. At most `jobs` shards run at once
+    (`none` ⇒ all of them); cap it to bound peak RAM, since each in-flight
+    shard's IO buffer re-ingests its whole closure. Returns 1 on a coverage gap
+    or any shard failure. -/
+def runShardCheckAll (manifestPath ixePath : String) (jobs? : Option Nat)
+    (runOne : Ix.Claim → IxVM.ClaimHarness.ClaimWitness → String → IO UInt32) : IO UInt32 := do
+  let (ixonEnv, shards) ← match (← loadEnvAndShards manifestPath ixePath) with
+    | .error e => IO.eprintln e; return 1
+    | .ok r => pure r
+  if !(← shardsCover ixonEnv shards) then return 1
+  -- The env + compiled toplevel are read-only, so each shard runs on its own
+  -- dedicated task; chunk by `jobs` to cap the number in flight at once.
+  let maxJobs := max 1 (jobs?.getD shards.size)
+  let mut rc : UInt32 := 0
+  for chunk in (shards.mapIdx (fun k b => (b, k))).toList.toChunks maxJobs do
+    let tasks ← chunk.mapM fun (blocks, k) =>
+      IO.asTask (prio := .dedicated) (runShardOwned ixonEnv blocks k runOne)
+    for t in tasks do
+      match t.get with
+      | .ok r => if r != 0 then rc := 1
+      | .error e => IO.eprintln s!"shard check task failed: {e}"; rc := 1
+  pure rc
+
 def runCheckCmd (p : Cli.Parsed) : IO UInt32 := do
   -- Always silence the Rust-side `[compile_env]` progress logs. The
   -- per-name labels + stats are signal enough at this layer.
@@ -255,7 +439,11 @@ def runCheckCmd (p : Cli.Parsed) : IO UInt32 := do
   let claimHex : Option String :=
     (p.flag? "claim").map (·.as! String)
   let names := (p.variableArgsAs! String).toList
-  let printStats := names.length == 1 || claimHex.isSome
+  let ixesPath := (p.flag? "ixes").map (·.as! String)
+  let shardK := (p.flag? "shard").map (·.as! Nat)
+  -- a single targeted constant, a `--claim`, or a single shard each print
+  -- per-circuit stats; whole-env / whole-partition iteration suppresses them.
+  let printStats := names.length == 1 || claimHex.isSome || (ixesPath.isSome && shardK.isSome)
   let toplevel ← match IxVM.ixVM with
     | .error e => IO.eprintln s!"Toplevel merging failed: {e}"; return 1
     | .ok t => pure t
@@ -270,8 +458,15 @@ def runCheckCmd (p : Cli.Parsed) : IO UInt32 := do
         | .error e => IO.eprintln s!"Compilation failed: {e}"; return 1
         | .ok c => pure c
       pure (runCompiled compiled printStats statsOut)
-  forEachClaim ixePath claimHex names keepGoing "check"
-    (fun _ w l => runOne w l)
+  match ixePath, ixesPath, shardK with
+  | some ixe, some manifest, some k =>
+    return (← runShardCheckManifest manifest ixe k (fun _c w l => runOne w l))
+  | some ixe, some manifest, none   =>
+    return (← runShardCheckAll manifest ixe ((p.flag? "jobs").map (·.as! Nat))
+      (fun _c w l => runOne w l))
+  | _, _, _ =>
+    forEachClaim ixePath claimHex names keepGoing "check"
+      (fun _ w l => runOne w l)
 
 end Ix.Cli.CheckCmd
 
@@ -286,6 +481,9 @@ def checkCmd : Cli.Cmd := `[Cli|
     "ixe"       : String;   "Path to a serialized `.ixe` env. When set, the binary reads the env from disk instead of using the compiled-in Lean env."
     "claim"     : String;   "32-byte hex address of a persisted `Ix.Claim` in `~/.ix/store/`. When set, runs the `verify_claim` entrypoint once over the claim's witness against the `--ixe` env (single execution, skips per-const iteration)."
     "stats-out" : String;   "Redirect the per-circuit statistics dump to this file (only used when exactly one constant is targeted)."
+    "ixes"      : String;   "Path to a `.ixes` shard manifest (with --ixe). With --shard K: check the constants owned by shard K (ingress their closure, skip the frontier). Without --shard: check every shard of the partition concurrently, after a coverage check."
+    "shard"     : Nat;      "0-based shard index K (with --ixe + --ixes): check the constants owned by shard K of the manifest's partition."
+    "jobs"      : Nat;      "Max shards to check concurrently when checking a whole partition (--ixes without --shard). Default: all at once. Lower it to bound peak RAM — each in-flight shard re-ingests its closure into its own IO buffer."
 
   ARGS:
     ...names : String; "Fully-qualified Lean.Name(s) to check. With none, iterate every named constant in the env (sorted)."

Ix/Cli/ProveCmd.lean

@@ -94,8 +94,14 @@ def runProveCmd (p : Cli.Parsed) : IO UInt32 := do
     | .error e => IO.eprintln s!"compilation failed: {e}"; return 1
     | .ok c => pure c
   let aiurSystem := Aiur.AiurSystem.build compiled.bytecode commitmentParameters
-  Ix.Cli.CheckCmd.forEachClaim ixePath claimHex names keepGoing "prove"
-    (proveOne aiurSystem compiled)
+  let runOne := proveOne aiurSystem compiled
+  match ixePath, (p.flag? "ixes").map (·.as! String), (p.flag? "shard").map (·.as! Nat) with
+  | some ixe, some manifest, some k =>
+    Ix.Cli.CheckCmd.runShardCheckManifest manifest ixe k runOne
+  | some ixe, some manifest, none =>
+    Ix.Cli.CheckCmd.runShardManifestAll manifest ixe runOne
+  | _, _, _ =>
+    Ix.Cli.CheckCmd.forEachClaim ixePath claimHex names keepGoing "prove" runOne
 
 end Ix.Cli.ProveCmd
 
@@ -108,6 +114,8 @@ def proveCmd : Cli.Cmd := `[Cli|
     "keep-going";       "Continue past failures and report them at the end instead of halting on the first."
     "ixe"   : String;   "Path to a serialized `.ixe` env. When set, the binary reads the env from disk instead of using the compiled-in Lean env."
     "claim" : String;   "32-byte hex address of a persisted `Ix.Claim` in `~/.ix/store/`. When set, proves the persisted claim against the `--ixe` env (single proof, skips per-const iteration)."
+    "ixes"  : String;   "Path to a `.ixes` shard manifest (with --ixe). With --shard K: prove shard K. Without --shard: prove every shard in the partition."
+    "shard" : Nat;      "0-based shard index K (with --ixes and --ixe): prove that one shard's CheckEnv claim."
 
   ARGS:
     ...names : String; "Fully-qualified Lean.Name(s) to prove. With none, iterate every named constant in the env (sorted)."