Aggregate-first recursion: lift shard proofs, fold CheckEnv claims, and verify one root

Ix/Aiur/Protocol.lean

@@ -154,6 +154,21 @@ opaque proveMultiStark (system : @& AiurSystem)
   (proofBytes vkBytes claimBytes : @& ByteArray) (useBytecode : Bool := false) :
     Array G × Proof
 
+/-- Prove one flat or structural aggregate-first binary join over raw child
+proof/claim advice. The compact preimage/tree/path blobs are produced by
+`MultiStark.joinPreimagesBlob`, `MultiStark.joinTreesBlob`, and
+`MultiStark.joinPathsBlob`. Malformed
+framing is returned as an error; as with `prove`/`proveMultiStark`, callers
+must supply an accepting execution witness. The final native IO buffer is
+intentionally not marshalled back to Lean. -/
+@[extern "rs_aiur_multi_stark_join_prove"]
+opaque proveMultiStarkJoin (system : @& AiurSystem)
+  (funIdx : @& Bytecode.FunIdx) (pubInput : @& Array G)
+  (leftProofBytes rightProofBytes recursionVkBytes : @& ByteArray)
+  (leftClaimsBytes rightClaimsBytes outputClaimBytes allowedBytes : @& ByteArray)
+  (preimagesBlob treesBlob pathsBlob : @& ByteArray) (useBytecode : Bool := false) :
+    Except String (Array G × Proof)
+
 @[extern "rs_aiur_system_prove_addr_with_env"]
 private opaque proveAddrWithEnv' : @& AiurSystem →
   @& Bytecode.FunIdx → @& EnvHandle → @& ByteArray → Bool →

Ix/Aiur/Semantics/BytecodeFfi.lean

@@ -174,6 +174,22 @@ opaque executeMultiStark (toplevel : @& Bytecode.Toplevel)
   (proofBytes vkBytes claimBytes : @& ByteArray) (useBytecode : Bool := false) :
     Except String (Array G × Array QueryCount)
 
+/-- Native execution of either aggregate-first join entrypoint. The two
+child proofs/claims and the fixed vk/output/allowed blobs are passed without
+per-byte Lean field boxing. `preimagesBlob` and `treesBlob` use the compact
+count/key/length framing produced by `MultiStark.joinPreimagesBlob` and
+`MultiStark.joinTreesBlob`; `pathsBlob` carries structural discharge choices.
+Rust expands them directly into IO channels 4–6. As with
+`executeMultiStark`, `useBytecode` selects the generic interpreter over the
+generated production verifier. -/
+@[extern "rs_aiur_multi_stark_join_execute"]
+opaque executeMultiStarkJoin (toplevel : @& Bytecode.Toplevel)
+  (funIdx : @& Bytecode.FunIdx) (pubInput : @& Array G)
+  (leftProofBytes rightProofBytes recursionVkBytes : @& ByteArray)
+  (leftClaimsBytes rightClaimsBytes outputClaimBytes allowedBytes : @& ByteArray)
+  (preimagesBlob treesBlob pathsBlob : @& ByteArray) (useBytecode : Bool := false) :
+    Except String (Array G × Array QueryCount)
+
 -- (EnvHandle opaque type + constructors live above `namespace
 -- Bytecode.Toplevel`; see `Aiur.EnvHandle`. The with-env FFI
 -- declarations below reference `EnvHandle` and `Bytecode.Toplevel`

Ix/Cli/CheckCmd.lean

@@ -428,34 +428,131 @@ private def ixesAddr : IxesP Address := do
   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`
+/-- The binary bisection tree stored at the tail of a `.ixes` manifest. Its
+leaves are shard ids; internal nodes prescribe the aggregation order that
+minimizes how long cross-shard assumptions remain live. -/
+inductive AggregationTree where
+  | leaf (shard : Nat)
+  | node (left right : AggregationTree)
+  deriving BEq, Repr, Inhabited
+
+namespace AggregationTree
+
+partial def leaves : AggregationTree → Array Nat
+  | .leaf shard => #[shard]
+  | .node left right => left.leaves ++ right.leaves
+
+/-- A legacy fallback for manifests written before the optional bisection-tree
+tail existed. Leaves remain in ascending shard order. -/
+partial def balancedRange (start count : Nat) : AggregationTree :=
+  if count ≤ 1 then .leaf start
+  else
+    let leftCount := count / 2
+    .node (balancedRange start leftCount)
+      (balancedRange (start + leftCount) (count - leftCount))
+
+/-- One post-order binary aggregation operation. Join child indices always
+refer to earlier plan slots, and the final slot is the root. -/
+inductive FoldOp where
+  | leaf (shard : Nat)
+  | join (left right : Nat)
+  deriving BEq, Repr
+
+/-- Lower a bisection tree to the slot plan consumed by the binary join host
+driver. -/
+partial def foldPlan (tree : AggregationTree) : Array FoldOp :=
+  (go tree #[]).2
+where
+  go (node : AggregationTree) (ops : Array FoldOp) : Nat × Array FoldOp :=
+    match node with
+    | .leaf shard => (ops.size, ops.push (.leaf shard))
+    | .node left right =>
+      let (leftIdx, ops) := go left ops
+      let (rightIdx, ops) := go right ops
+      (ops.size, ops.push (.join leftIdx rightIdx))
+
+end AggregationTree
+
+structure IxesManifestView where
+  shards : Array (Array Address)
+  aggregationTree : AggregationTree
+  deriving BEq, Repr
+
+private partial def ixesAggregationTree : IxesP AggregationTree := do
+  match ← ixesU8 with
+  | 0 => pure (.leaf (← ixesU32).toNat)
+  | 1 => pure (.node (← ixesAggregationTree) (← ixesAggregationTree))
+  | tag => throw s!"ixes: invalid aggregation-tree tag {tag.toNat}"
+
+private def validateAggregationTree (tree : AggregationTree)
+    (numShards : Nat) : Except String Unit := do
+  let leaves := tree.leaves
+  if leaves.size != numShards then
+    throw s!"ixes: aggregation tree has {leaves.size} leaves for {numShards} shards"
+  let mut seen : Std.HashSet Nat := {}
+  for shard in leaves do
+    if shard ≥ numShards then
+      throw s!"ixes: aggregation tree leaf {shard} is out of range"
+    if seen.contains shard then
+      throw s!"ixes: aggregation tree repeats shard {shard}"
+    seen := seen.insert shard
+
+/-- Parse every shard's owned block addresses and aggregation tree 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
+    The optional trailing tree uses preorder `leaf(0,id:u32)` /
+    `node(1,left,right)` encoding. Legacy manifests with no tree (or an explicit
+    zero presence byte) receive a balanced ascending-id fallback. Bounds-checked:
+    a truncated/malformed file yields `.error`, never a panic. -/
+def parseIxesManifest (bytes : ByteArray) : Except String IxesManifestView :=
+  let go : IxesP IxesManifestView := 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
+    if n == 0 then throw "ixes: manifest contains no shards"
     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
+    for shardIdx in [0:n.toNat] do
+      let shardId ← ixesU32
+      if shardId.toNat != shardIdx then
+        throw s!"ixes: shard entry {shardIdx} has id {shardId.toNat}"
+      ixesSkip (8 + 8 + 8)  -- heartbeats + own_size + cross_ingress
+      match ← ixesU8 with
+      | 0 => pure ()
+      | 1 => ixesSkip 32  -- assumption_root present
+      | tag => throw s!"ixes: invalid assumption-root presence tag {tag.toNat}"
       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
+    let (buffer, position) ← get
+    let tree ← if position == buffer.size then
+        pure (AggregationTree.balancedRange 0 n.toNat)
+      else
+        match ← ixesU8 with
+        | 0 => pure (AggregationTree.balancedRange 0 n.toNat)
+        | 1 => ixesAggregationTree
+        | tag => throw s!"ixes: invalid aggregation-tree presence tag {tag.toNat}"
+    validateAggregationTree tree n.toNat
+    let (buffer, position) ← get
+    if position != buffer.size then
+      throw s!"ixes: {buffer.size - position} trailing bytes after aggregation tree"
+    pure { shards, aggregationTree := tree }
   go.run' (bytes, 0)
 
+/-- Backward-compatible shard-only view used by check/prove/verify paths that
+do not yet consume the aggregation order. -/
+def parseIxesAllShards (bytes : ByteArray) : Except String (Array (Array Address)) :=
+  (parseIxesManifest bytes).map (·.shards)
+
 /-- 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`). -/
@@ -482,6 +579,22 @@ def ownedConstsForBlocks (ixonEnv : Ixon.Env) (blocks : Array Address) : Array A
     if blockSet.contains (blockAddrOf addr c) then o := o.push addr
   return o
 
+/-- Count each shard's owned constants in one environment pass. Callers first
+run `shardsCover`, so every check-schedule block has exactly one owner. This is
+the subject-leaf count used by aggregate structural-threshold scheduling. -/
+def ownedConstCountsForShards (ixonEnv : Ixon.Env)
+    (shards : Array (Array Address)) : Array Nat := Id.run do
+  let mut blockToShard : Std.HashMap Address Nat := {}
+  for (blocks, shard) in shards.mapIdx fun shard blocks => (blocks, shard) do
+    for block in blocks do
+      blockToShard := blockToShard.insert block shard
+  let mut counts := Array.replicate shards.size 0
+  for (addr, lc) in ixonEnv.consts do
+    let some c := lc.get? | continue
+    let some shard := blockToShard.get? (blockAddrOf addr c) | continue
+    counts := counts.modify shard (· + 1)
+  return counts
+
 /-- 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. -/