Aiur hash circuits: value-based keccak-256, inlined calls, self-recursive rounds, and a generalized chain-rotate gadget

Benchmarks/Blake3.lean

@@ -1,54 +1,12 @@
 import Ix.IxVM.Core
 import Ix.IxVM.ByteStream
 import Ix.IxVM.Blake3
-import Ix.Aiur.Protocol
-import Ix.Aiur.Compiler
-import Ix.Benchmark.Bench
-
-open BgroupM
-
-abbrev dataSizes := #[64, 128, 256, 512, 1024, 2048]
-abbrev numHashesPerProof := #[1, 2, 4, 8, 16, 32]
-
-def commitmentParameters : Aiur.CommitmentParameters := {
-  logBlowup := 1
-  capHeight := 0
-}
-
-def friParameters : Aiur.FriParameters := {
-  logFinalPolyLen := 0
-  maxLogArity := 1
-  numQueries := 100
-  commitProofOfWorkBits := 20
-  queryProofOfWorkBits := 0
-}
+import Benchmarks.HashCommon
 
 def mergedToplevel : Except Aiur.Global Aiur.Source.Toplevel := do
   let tl ← IxVM.core.merge IxVM.byteStream
   tl.merge IxVM.blake3
 
-def blake3Bench : IO $ Array BenchReport := do
-  let .ok toplevel := mergedToplevel
-    | throw (IO.userError "Merging failed")
-  let .ok compiled := toplevel.compile
-    | throw (IO.userError "Compilation failed")
-  let some funIdx := compiled.getFuncIdx `blake3_bench
-    | throw (IO.userError "Aiur function not found")
-  let aiurSystem := Aiur.AiurSystem.build compiled.bytecode commitmentParameters friParameters
-  bgroup "prove blake3" { oneShot := true, avgThroughput := true, report := true } do
-    for dataSize in dataSizes do
-      for numHashes in numHashesPerProof do
-        let ioBuffer := Array.range numHashes |>.foldl
-          (init := default)
-          fun ioBuffer idx =>
-            let data := Array.range dataSize |>.map
-              -- Add `idx` so every preimage is different and avoids memoization.
-              fun i => Aiur.G.ofUInt8 (i + idx).toUInt8
-            ioBuffer.extend 0 #[.ofNat idx] data
-        throughput (.ElementsAndBytes numHashes.toUInt64 (dataSize * numHashes).toUInt64 "hashes")
-        bench s!"dataSize={dataSize} numHashes={numHashes}"
-          (aiurSystem.prove funIdx #[Aiur.G.ofNat numHashes]) ioBuffer
-
 def main : IO Unit := do
-  let _ ← blake3Bench
+  let _ ← HashBench.run "prove blake3" `blake3_bench mergedToplevel
   return

Benchmarks/HashCommon.lean

@@ -0,0 +1,80 @@
+import Ix.Aiur.Protocol
+import Ix.Aiur.Compiler
+import Ix.Aiur.Statistics
+import Ix.Benchmark.Bench
+
+/-!
+Shared harness for the hash-proving benchmarks (`bench-blake3`,
+`bench-sha256` and `bench-keccak`).
+
+Each benchmark case proves `numHashes` hashes of `dataSize` bytes in a single
+Aiur proof and reports proving throughput in bytes/s. Before the timed runs,
+the harness executes the largest workload once (execution only, no proving)
+and prints the circuit statistics — per-circuit width, height and FFT cost,
+plus the total circuit width — so width regressions show up in the benchmark
+output.
+-/
+
+open BgroupM
+
+namespace HashBench
+
+/-- Sizes in bytes of each hashed message. The largest case proves 4 hashes
+of 128 KiB — 512 KiB per proof — which sits on the sustained bytes/s plateau
+(reached at ~128 KiB per proof) while keeping a full grid run fast enough to
+iterate on. -/
+def dataSizes : Array Nat := #[4096, 32768, 131072]
+
+/-- Number of hashes proven in a single proof. -/
+def numHashesPerProof : Array Nat := #[1, 4]
+
+def commitmentParameters : Aiur.CommitmentParameters := {
+  logBlowup := 1
+  capHeight := 0
+}
+
+def friParameters : Aiur.FriParameters := {
+  logFinalPolyLen := 0
+  maxLogArity := 1
+  numQueries := 100
+  commitProofOfWorkBits := 20
+  queryProofOfWorkBits := 0
+}
+
+/-- Builds an `IOBuffer` with `numHashes` streams of `dataSize` bytes each.
+
+The bytes are non-periodic pseudorandom (LCG), seeded per hash: a
+`(i + idx) % 256` pattern repeats every 256 bytes, which memoizes repeated
+sub-computations across blocks and understates honest proving cost. -/
+def mkIOBuffer (dataSize numHashes : Nat) : Aiur.IOBuffer :=
+  Array.range numHashes |>.foldl (init := default) fun ioBuffer idx =>
+    let data := Array.range dataSize |>.map fun i =>
+      Aiur.G.ofUInt8 (((i + idx * dataSize) * 1103515245 + 12345) >>> 16).toUInt8
+    ioBuffer.extend 0 #[.ofNat idx] data
+
+def run (groupName : String) (funcName : Lean.Name)
+    (toplevel : Except Aiur.Global Aiur.Source.Toplevel) : IO (Array BenchReport) := do
+  let .ok toplevel := toplevel
+    | throw (IO.userError "Merging failed")
+  let .ok compiled := toplevel.compile
+    | throw (IO.userError "Compilation failed")
+  let some funIdx := compiled.getFuncIdx funcName
+    | throw (IO.userError "Aiur function not found")
+  let maxDataSize := dataSizes.foldl Nat.max 0
+  let maxNumHashes := numHashesPerProof.foldl Nat.max 0
+  match compiled.bytecode.execute funIdx #[.ofNat maxNumHashes]
+      (mkIOBuffer maxDataSize maxNumHashes) with
+  | .error e => throw (IO.userError s!"Execution failed: {e}")
+  | .ok (_, _, queryCounts) =>
+    IO.println s!"Circuit statistics at dataSize={maxDataSize} numHashes={maxNumHashes}:"
+    Aiur.printStats (Aiur.computeStats compiled queryCounts)
+  let aiurSystem := Aiur.AiurSystem.build compiled.bytecode commitmentParameters friParameters
+  bgroup groupName { oneShot := true, avgThroughput := true, report := true } do
+    for dataSize in dataSizes do
+      for numHashes in numHashesPerProof do
+        let ioBuffer := mkIOBuffer dataSize numHashes
+        throughput (.ElementsAndBytes numHashes.toUInt64 (dataSize * numHashes).toUInt64 "hashes")
+        bench s!"dataSize={dataSize} numHashes={numHashes}"
+          (aiurSystem.prove funIdx #[Aiur.G.ofNat numHashes]) ioBuffer
+
+end HashBench

Benchmarks/Keccak.lean

@@ -0,0 +1,12 @@
+import Ix.IxVM.Core
+import Ix.IxVM.ByteStream
+import Ix.IxVM.Keccak
+import Benchmarks.HashCommon
+
+def mergedToplevel : Except Aiur.Global Aiur.Source.Toplevel := do
+  let tl ← IxVM.core.merge IxVM.byteStream
+  tl.merge IxVM.keccak
+
+def main : IO Unit := do
+  let _ ← HashBench.run "prove keccak256" `keccak256_bench mergedToplevel
+  return

Benchmarks/Sha256.lean

@@ -1,54 +1,12 @@
 import Ix.IxVM.Core
 import Ix.IxVM.ByteStream
 import Ix.IxVM.Sha256
-import Ix.Aiur.Protocol
-import Ix.Aiur.Compiler
-import Ix.Benchmark.Bench
-
-open BgroupM
-
-abbrev dataSizes := #[64, 128, 256, 512, 1024, 2048]
-abbrev numHashesPerProof := #[1, 2, 4, 8, 16, 32]
-
-def commitmentParameters : Aiur.CommitmentParameters := {
-  logBlowup := 1
-  capHeight := 0
-}
-
-def friParameters : Aiur.FriParameters := {
-  logFinalPolyLen := 0
-  maxLogArity := 1
-  numQueries := 100
-  commitProofOfWorkBits := 20
-  queryProofOfWorkBits := 0
-}
+import Benchmarks.HashCommon
 
 def mergedToplevel : Except Aiur.Global Aiur.Source.Toplevel := do
   let tl ← IxVM.core.merge IxVM.byteStream
   tl.merge IxVM.sha256
 
-def sha256Bench : IO $ Array BenchReport := do
-  let .ok toplevel := mergedToplevel
-    | throw (IO.userError "Merging failed")
-  let .ok compiled := toplevel.compile
-    | throw (IO.userError "Compilation failed")
-  let some funIdx := compiled.getFuncIdx `sha256_bench
-    | throw (IO.userError "Aiur function not found")
-  let aiurSystem := Aiur.AiurSystem.build compiled.bytecode commitmentParameters friParameters
-  bgroup "prove sha256" { oneShot := true, avgThroughput := true, report := true } do
-    for dataSize in dataSizes do
-      for numHashes in numHashesPerProof do
-        let ioBuffer := Array.range numHashes |>.foldl
-          (init := default)
-          fun ioBuffer idx =>
-            let data := Array.range dataSize |>.map
-              -- Add `idx` so every preimage is different and avoids memoization.
-              fun i => Aiur.G.ofUInt8 (i + idx).toUInt8
-            ioBuffer.extend 0 #[.ofNat idx] data
-        throughput (.ElementsAndBytes numHashes.toUInt64 (dataSize * numHashes).toUInt64 "hashes")
-        bench s!"dataSize={dataSize} numHashes={numHashes}"
-          (aiurSystem.prove funIdx #[Aiur.G.ofNat numHashes]) ioBuffer
-
 def main : IO Unit := do
-  let _ ← sha256Bench
+  let _ ← HashBench.run "prove sha256" `sha256_bench mergedToplevel
   return

Ix/Aiur/Compiler.lean

@@ -112,6 +112,7 @@ def Bytecode.Toplevel.needsCircuit (t : Bytecode.Toplevel) : Array Bool := Id.ru
 
 /-- Full compilation pipeline. -/
 def Source.Toplevel.compile (t : Source.Toplevel) : Except String CompiledToplevel := do
+  let t ← t.inlineCalls
   let typedDecls ← t.checkAndSimplify.mapError toString
   let concDecls ← typedDecls.concretize.mapError toString
   let (bytecodeRaw, preNameMap) ← concDecls.toBytecode
@@ -129,13 +130,14 @@ def Source.Toplevel.compile (t : Source.Toplevel) : Except String CompiledToplev
 `needsCircuit`, the field-setter `mapIdx`, the name-map `fold`, and the
 terminating `pure` — are all total). -/
 theorem Source.Toplevel.compile_ok_of_stages
-    {t : Source.Toplevel} {typedDecls concDecls bytecodeRaw preNameMap}
-    (hts : t.checkAndSimplify = .ok typedDecls)
+    {t inlined : Source.Toplevel} {typedDecls concDecls bytecodeRaw preNameMap}
+    (hinline : t.inlineCalls = .ok inlined)
+    (hts : inlined.checkAndSimplify = .ok typedDecls)
     (hconc : typedDecls.concretize = .ok concDecls)
     (hbc : concDecls.toBytecode = .ok (bytecodeRaw, preNameMap)) :
     ∃ ct, t.compile = .ok ct := by
-  simp only [Source.Toplevel.compile, hts, hconc, hbc, Except.mapError, bind,
-             Except.bind, pure, Except.pure]
+  simp only [Source.Toplevel.compile, hinline, hts, hconc, hbc, Except.mapError,
+             bind, Except.bind, pure, Except.pure]
   exact ⟨_, rfl⟩
 
 /-- Inverse of `compile_ok_of_stages`: unpack a successful `compile` into
@@ -144,27 +146,28 @@ lemmas through the composition. -/
 theorem Source.Toplevel.compile_stages_of_ok
     {t : Source.Toplevel} {ct : CompiledToplevel}
     (_hct : t.compile = .ok ct) :
-    ∃ typedDecls concDecls bytecodeRaw preNameMap,
-      t.checkAndSimplify = .ok typedDecls ∧
+    ∃ inlined typedDecls concDecls bytecodeRaw preNameMap,
+      t.inlineCalls = .ok inlined ∧
+      inlined.checkAndSimplify = .ok typedDecls ∧
       typedDecls.concretize = .ok concDecls ∧
       concDecls.toBytecode = .ok (bytecodeRaw, preNameMap) := by
   -- Case on each stage result; `mapError` on `.ok` is definitionally `.ok`.
-  cases hts : t.checkAndSimplify with
-  | error e => simp [Source.Toplevel.compile, hts, bind, Except.bind, Except.mapError] at _hct
-  | ok typedDecls =>
-    cases hconc : typedDecls.concretize with
-    | error e =>
-      simp [Source.Toplevel.compile, hts, hconc, bind, Except.bind, Except.mapError] at _hct
-    | ok concDecls =>
-      cases hbc : concDecls.toBytecode with
+  cases hinline : t.inlineCalls with
+  | error e => simp [Source.Toplevel.compile, hinline, bind, Except.bind] at _hct
+  | ok inlined =>
+    cases hts : inlined.checkAndSimplify with
+    | error e => simp [Source.Toplevel.compile, hinline, hts, bind, Except.bind, Except.mapError] at _hct
+    | ok typedDecls =>
+      cases hconc : typedDecls.concretize with
       | error e =>
-        simp [Source.Toplevel.compile, hts, hconc, hbc, bind, Except.bind, Except.mapError] at _hct
-      | ok bc =>
-        obtain ⟨bytecodeRaw, preNameMap⟩ := bc
-        -- `checkAndSimplify`/`concretize` are unfolded; `toBytecode` is not
-        -- (private module), so the first two equalities reduce to `rfl` but
-        -- the third still mentions `concDecls.toBytecode` — supply `hbc`.
-        exact ⟨typedDecls, concDecls, bytecodeRaw, preNameMap, rfl, hconc, hbc⟩
+        simp [Source.Toplevel.compile, hinline, hts, hconc, bind, Except.bind, Except.mapError] at _hct
+      | ok concDecls =>
+        cases hbc : concDecls.toBytecode with
+        | error e =>
+          simp [Source.Toplevel.compile, hinline, hts, hconc, hbc, bind, Except.bind, Except.mapError] at _hct
+        | ok bc =>
+          obtain ⟨bytecodeRaw, preNameMap⟩ := bc
+          exact ⟨inlined, typedDecls, concDecls, bytecodeRaw, preNameMap, rfl, hts, hconc, hbc⟩
 
 -- Compile post-conditions moved to `Proofs/StructCompatible.lean`.
 

Ix/Aiur/Compiler/Check.lean

@@ -49,6 +49,7 @@ inductive CheckError
   | infiniteType : Nat → Typ → CheckError
   | unresolvedMVar : Nat → CheckError
   | u8LitOutOfRange : Nat → CheckError
+  | chainRotrOutOfRange : Nat → CheckError
   | entryHasPointer : Global → CheckError
   deriving Repr
 
@@ -622,7 +623,14 @@ def inferTerm (t : Term) : CheckM Typed.Term := match t with
     match typOpt with
     | some (typ, escapes) => pure (Typed.Term.match typ escapes term' branches')
     | none => throw .emptyMatch
-  | .app func args u => do
+  | .app func args mode => do
+    -- `.inlined` is eliminated by `Toplevel.inlineCalls` before checking;
+    -- reaching it here means the pass was skipped. Downstream stages carry
+    -- only the unconstrained flag.
+    let u : Bool ← match mode with
+      | .normal => pure false
+      | .unconstrained => pure true
+      | .inlined => throw (.cannotApply func)
     let ctx ← read
     -- Local function lookup (only for unqualified names); returns
     -- `some (.ok …)` on hit, `some (.error …)` on wrong-type local, `none` to
@@ -758,14 +766,12 @@ def inferTerm (t : Term) : CheckM Typed.Term := match t with
     let a' ← checkNoEscape a .u8
     let b' ← checkNoEscape b .u8
     pure (Typed.Term.u8Mul (.tuple #[.u8, .u8]) false a' b')
-  | .u8ChainRotr7 a b => do
+  | .u8ChainRotr k a b => do
+    if k == 0 || k > 7 then
+      throw (.chainRotrOutOfRange k)
     let a' ← checkNoEscape a .u8
     let b' ← checkNoEscape b .u8
-    pure (Typed.Term.u8ChainRotr7 (.tuple #[.u8, .u8, .u8]) false a' b')
-  | .u8ChainRotr4 a b => do
-    let a' ← checkNoEscape a .u8
-    let b' ← checkNoEscape b .u8
-    pure (Typed.Term.u8ChainRotr4 (.tuple #[.u8, .u8, .u8]) false a' b')
+    pure (Typed.Term.u8ChainRotr k (.tuple #[.u8, .u8, .u8]) false a' b')
   | .u8Sub a b => do
     -- Low byte and the 0/1 borrow are both `u8` (same range argument as add).
     let a' ← checkNoEscape a .u8
@@ -961,8 +967,7 @@ def zonkTypedTerm (t : Typed.Term) : CheckM Typed.Term := match t with
   | .u8Xor τ e a b => do pure (.u8Xor (← zonkTyp τ) e (← zonkTypedTerm a) (← zonkTypedTerm b))
   | .u8Add τ e a b => do pure (.u8Add (← zonkTyp τ) e (← zonkTypedTerm a) (← zonkTypedTerm b))
   | .u8Mul τ e a b => do pure (.u8Mul (← zonkTyp τ) e (← zonkTypedTerm a) (← zonkTypedTerm b))
-  | .u8ChainRotr7 τ e a b => do pure (.u8ChainRotr7 (← zonkTyp τ) e (← zonkTypedTerm a) (← zonkTypedTerm b))
-  | .u8ChainRotr4 τ e a b => do pure (.u8ChainRotr4 (← zonkTyp τ) e (← zonkTypedTerm a) (← zonkTypedTerm b))
+  | .u8ChainRotr k τ e a b => do pure (.u8ChainRotr k (← zonkTyp τ) e (← zonkTypedTerm a) (← zonkTypedTerm b))
   | .u8Sub τ e a b => do pure (.u8Sub (← zonkTyp τ) e (← zonkTypedTerm a) (← zonkTypedTerm b))
   | .u8And τ e a b => do pure (.u8And (← zonkTyp τ) e (← zonkTypedTerm a) (← zonkTypedTerm b))
   | .u8Or τ e a b => do pure (.u8Or (← zonkTyp τ) e (← zonkTypedTerm a) (← zonkTypedTerm b))

Ix/Aiur/Compiler/Concretize.lean

@@ -339,11 +339,8 @@ def termToConcrete
   | .u8Mul τ e a b => do
       pure (.u8Mul (← typToConcrete mono τ) e
                    (← termToConcrete mono a) (← termToConcrete mono b))
-  | .u8ChainRotr7 τ e a b => do
-      pure (.u8ChainRotr7 (← typToConcrete mono τ) e
-                   (← termToConcrete mono a) (← termToConcrete mono b))
-  | .u8ChainRotr4 τ e a b => do
-      pure (.u8ChainRotr4 (← typToConcrete mono τ) e
+  | .u8ChainRotr k τ e a b => do
+      pure (.u8ChainRotr k (← typToConcrete mono τ) e
                    (← termToConcrete mono a) (← termToConcrete mono b))
   | .u8Sub τ e a b => do
       pure (.u8Sub (← typToConcrete mono τ) e
@@ -557,9 +554,7 @@ def rewriteTypedTerm (decls : Typed.Decls)
       (rewriteTypedTerm decls subst mono a) (rewriteTypedTerm decls subst mono b)
   | .u8Mul τ e a b => .u8Mul (rewriteTyp subst mono τ) e
       (rewriteTypedTerm decls subst mono a) (rewriteTypedTerm decls subst mono b)
-  | .u8ChainRotr7 τ e a b => .u8ChainRotr7 (rewriteTyp subst mono τ) e
-      (rewriteTypedTerm decls subst mono a) (rewriteTypedTerm decls subst mono b)
-  | .u8ChainRotr4 τ e a b => .u8ChainRotr4 (rewriteTyp subst mono τ) e
+  | .u8ChainRotr k τ e a b => .u8ChainRotr k (rewriteTyp subst mono τ) e
       (rewriteTypedTerm decls subst mono a) (rewriteTypedTerm decls subst mono b)
   | .u8Sub τ e a b => .u8Sub (rewriteTyp subst mono τ) e
       (rewriteTypedTerm decls subst mono a) (rewriteTypedTerm decls subst mono b)
@@ -646,7 +641,7 @@ def collectInTypedTerm (seen : Std.HashSet (Global × Array Typ)) :
     args.attach.foldl (fun s ⟨a, _⟩ => collectInTypedTerm s a) seen
   | .add τ _ a b | .sub τ _ a b | .mul τ _ a b
   | .u8Xor τ _ a b | .u8Add τ _ a b | .u8Mul τ _ a b | .u8Sub τ _ a b
-  | .u8ChainRotr7 τ _ a b | .u8ChainRotr4 τ _ a b | .u8RangeCheck τ _ a b
+  | .u8ChainRotr _ τ _ a b | .u8RangeCheck τ _ a b
   | .unconstrainedBigUintDivMod τ _ a b
   | .u8And τ _ a b | .u8Or τ _ a b
   | .u8LessThan τ _ a b | .u32LessThan τ _ a b =>
@@ -717,7 +712,7 @@ def collectCalls (decls : Typed.Decls)
     bs.attach.foldl (fun s ⟨(_, b), _⟩ => collectCalls decls s b) seen
   | .add _ _ a b | .sub _ _ a b | .mul _ _ a b
   | .u8Xor _ _ a b | .u8Add _ _ a b | .u8Mul _ _ a b | .u8Sub _ _ a b
-  | .u8ChainRotr7 _ _ a b | .u8ChainRotr4 _ _ a b | .u8RangeCheck _ _ a b
+  | .u8ChainRotr _ _ _ a b | .u8RangeCheck _ _ a b
   | .unconstrainedBigUintDivMod _ _ a b
   | .u8And _ _ a b | .u8Or _ _ a b
   | .u8LessThan _ _ a b | .u32LessThan _ _ a b =>
@@ -821,9 +816,7 @@ def substInTypedTerm (subst : Global → Option Typ) : Typed.Term → Typed.Term
       (substInTypedTerm subst a) (substInTypedTerm subst b)
   | .u8Mul τ e a b => .u8Mul (Typ.instantiate subst τ) e
       (substInTypedTerm subst a) (substInTypedTerm subst b)
-  | .u8ChainRotr7 τ e a b => .u8ChainRotr7 (Typ.instantiate subst τ) e
-      (substInTypedTerm subst a) (substInTypedTerm subst b)
-  | .u8ChainRotr4 τ e a b => .u8ChainRotr4 (Typ.instantiate subst τ) e
+  | .u8ChainRotr k τ e a b => .u8ChainRotr k (Typ.instantiate subst τ) e
       (substInTypedTerm subst a) (substInTypedTerm subst b)
   | .u8Sub τ e a b => .u8Sub (Typ.instantiate subst τ) e
       (substInTypedTerm subst a) (substInTypedTerm subst b)

Ix/Aiur/Compiler/Layout.lean

@@ -198,7 +198,7 @@ def opLayout : Bytecode.Op → LayoutM Unit
     pushDegree ((aDegree.max bDegree).max 1)
     bumpAuxiliaries; bumpLookups
   | .u8Mul .. => do pushDegrees #[1, 1]; bumpAuxiliaries 2; bumpLookups
-  | .u8ChainRotr7 .. | .u8ChainRotr4 .. => do pushDegrees #[1, 1, 1]; bumpAuxiliaries 3; bumpLookups
+  | .u8ChainRotr .. => do pushDegrees #[1, 1, 1]; bumpAuxiliaries 3; bumpLookups
   | .u8LessThan .. => do pushDegree 1; bumpAuxiliaries; bumpLookups
   | .u32LessThan .. => do pushDegree 1; bumpAuxiliaries 12; bumpLookups 6
   -- Pure range-check lookup: no output columns/degrees, just one lookup.