diff --git a/Benchmarks/RecursionDebug.lean b/Benchmarks/RecursionDebug.lean index 8bc20302..847c91ab 100644 --- a/Benchmarks/RecursionDebug.lean +++ b/Benchmarks/RecursionDebug.lean @@ -70,7 +70,7 @@ def proveConst (ixePath constName : String) (skipDeps : Bool) -- production toplevel no longer carries. let .ok toplevel := (if skipDeps then IxVM.ixVMFull else IxVM.ixVM) | IO.eprintln "IxVM toplevel merge failed"; return none - let .ok compiled := toplevel.compile + let .ok compiled := toplevel.compileWithGroups IxVM.coldGroups | IO.eprintln "IxVM compile failed"; return none let entrypoint := if skipDeps then `verify_const else `verify_claim let some funIdx := compiled.getFuncIdx entrypoint @@ -153,7 +153,7 @@ def main (args : List String) : IO UInt32 := do -- `--list-funcs`: dump the compiled verifier's funIdx → name table (for -- decoding fun_idx stacks printed by the Rust bytecode interpreter). if args.contains "--list-funcs" then - let .ok vCompiled := vTop.compile + let .ok vCompiled := vTop.compileWithGroups MultiStark.verifierColdGroups | IO.eprintln "multi-stark verifier compile failed"; return 1 let entries := vCompiled.nameMap.toArray.qsort (·.2 < ·.2) for (g, i) in entries do @@ -187,7 +187,7 @@ def main (args : List String) : IO UInt32 := do IO.println s!"ACCEPTED in {secs t0 t1} s: {Aiur.Value.ppDeref s.store depth v}" return 0 else - let .ok vCompiled := vTop.compile + let .ok vCompiled := vTop.compileWithGroups MultiStark.verifierColdGroups | IO.eprintln "multi-stark verifier compile failed"; return 1 let some vIdx := vCompiled.getFuncIdx `verify_multi_stark_proof | IO.eprintln "verify_multi_stark_proof entrypoint missing"; return 1 diff --git a/Benchmarks/RecursiveVerifier.lean b/Benchmarks/RecursiveVerifier.lean index e99d1f61..ed4ab5e3 100644 --- a/Benchmarks/RecursiveVerifier.lean +++ b/Benchmarks/RecursiveVerifier.lean @@ -147,7 +147,7 @@ def main (args : List String) : IO UInt32 := do let vTop ← match MultiStark.multiStark with | .ok t => pure t | .error e => IO.eprintln s!"verifier merge failed: {e}"; return 1 - let vCompiled ← match vTop.compile with + let vCompiled ← match vTop.compileWithGroups MultiStark.verifierColdGroups with | .ok c => pure c | .error e => IO.eprintln s!"verifier compile failed: {e}"; return 1 let vIdx := vCompiled.getFuncIdx `verify_multi_stark_proof |>.get! diff --git a/Benchmarks/Typecheck.lean b/Benchmarks/Typecheck.lean index 5d3e6ede..bebde57a 100644 --- a/Benchmarks/Typecheck.lean +++ b/Benchmarks/Typecheck.lean @@ -329,7 +329,7 @@ def runTypecheckCmd (p : Cli.Parsed) : IO UInt32 := do -- claim run, which is the honest reading of those numbers. let .ok toplevel := (if skipDeps then IxVM.ixVMFull else IxVM.ixVM) | throw (IO.userError "Merging IxVM kernel failed") - let .ok compiled := toplevel.compile + let .ok compiled := toplevel.compileWithGroups IxVM.coldGroups | throw (IO.userError "Compilation of IxVM kernel failed") let entrypoint := if skipDeps then `verify_const else `verify_claim let some funIdx := compiled.getFuncIdx entrypoint @@ -358,7 +358,7 @@ def runTypecheckCmd (p : Cli.Parsed) : IO UInt32 := do if !recursive then pure none else do let .ok vTop := MultiStark.multiStark | throw (IO.userError "Merging multi-stark verifier failed") - let .ok vCompiled := vTop.compile + let .ok vCompiled := vTop.compileWithGroups MultiStark.verifierColdGroups | throw (IO.userError "Compilation of multi-stark verifier failed") let some vIdx := vCompiled.getFuncIdx `verify_multi_stark_proof | throw (IO.userError "verify_multi_stark_proof entrypoint missing") diff --git a/Ix/Aiur/Compiler.lean b/Ix/Aiur/Compiler.lean index d8526663..0fe32a89 100644 --- a/Ix/Aiur/Compiler.lean +++ b/Ix/Aiur/Compiler.lean @@ -42,6 +42,73 @@ def CompiledToplevel.getFuncIdx (ct : CompiledToplevel) (name : Lean.Name) : Option Bytecode.FunIdx := ct.nameMap[Global.mk name]? +/-- Regroup the circuit partition: each `(name, members)` in `groups` becomes +ONE circuit proving all the listed functions (branching on the member; see +`Bytecode.Circuit`), positioned where its first member's singleton circuit +was; every other constrained function keeps its singleton circuit. Grouping +is a circuit-level choice — the function "library", its bytecode, execution, +and the query record are untouched, and callers still target function +indices on the function channel. + +Grouping is favorable for RARELY-CALLED functions of similar shape: the +merged circuit sums the members' selector columns but takes the max of their +auxiliary columns, so each (rare) row pays the group's selector count, while +the system sheds one circuit (vk entry, commitment matrix, verifier work) +per absorbed member. + +Errors if a name is unknown, unconstrained (it has no circuit to group), an +entry function, listed twice, or if a group is empty. -/ +def CompiledToplevel.groupFunctions (ct : CompiledToplevel) + (groups : Array (String × Array String)) : + Except String CompiledToplevel := do + let t := ct.bytecode + -- Function names as printed (`toString` of the `Global`), the exact + -- inverse of what statistics reports — so measured groupings can be fed + -- back verbatim. + let byName : Std.HashMap String Bytecode.FunIdx := + ct.nameMap.fold (init := {}) fun acc g i => acc.insert (toString g) i + -- Resolve and validate the groups into member-index arrays. + let mut grouped : Std.HashMap Bytecode.FunIdx Nat := {} + let mut resolved : Array (String × Array Bytecode.FunIdx) := #[] + for (gname, names) in groups do + if names.isEmpty then + throw s!"group {gname} is empty" + let mut members := #[] + for name in names do + let some i := byName[name]? + | throw s!"group {gname}: unknown function {name}" + let f := t.functions[i]! + unless f.constrained do + throw s!"group {gname}: {name} is unconstrained (it has no circuit)" + if f.entry then + throw s!"group {gname}: {name} is an entry function" + if grouped.contains i then + throw s!"group {gname}: {name} is already grouped" + grouped := grouped.insert i resolved.size + members := members.push i + resolved := resolved.push (gname, members) + -- Rebuild the partition in first-occurrence order over the existing + -- (singleton-ordered) circuits. + let mut circuits : Array Bytecode.Circuit := #[] + let mut placed : Array Bool := .replicate resolved.size false + for c in t.circuits do + let members := c.members + if h : members.size = 1 then + let i := members[0] + match grouped[i]? with + | none => circuits := circuits.push c + | some g => + unless placed[g]! do + placed := placed.set! g true + let (gname, ms) := resolved[g]! + let layout := ms.foldl (init := t.functions[ms[0]!]!.layout) + fun acc m => if m == ms[0]! then acc + else acc.merge t.functions[m]!.layout + circuits := circuits.push { name := gname, members := ms, layout } + else + throw "groupFunctions: partition already grouped; group from a freshly compiled toplevel" + pure { ct with bytecode := { t with circuits } } + /-- Termination helper for the `Block`/`Ctrl` traversal below. -/ private theorem Bytecode.Block.sizeOf_ctrl_lt'' (b : Bytecode.Block) : sizeOf b.ctrl < sizeOf b := by @@ -90,6 +157,18 @@ decreasing_by | (apply Prod.Lex.left; exact Bytecode.Block.sizeOf_ctrl_lt'' _) end +/-- The default circuit partition: one singleton circuit per constrained +function, in function-index order, named by `nameOf`. -/ +def Bytecode.Toplevel.singletonCircuits (t : Bytecode.Toplevel) + (nameOf : Bytecode.FunIdx → String) : Array Bytecode.Circuit := Id.run do + let mut circuits : Array Bytecode.Circuit := #[] + for h : i in [:t.functions.size] do + let f := t.functions[i] + if f.constrained then + circuits := circuits.push + { name := nameOf i, members := #[i], layout := f.layout } + pure circuits + /-- Compute which functions need a circuit. A function needs a circuit iff it is reachable from an entry point through a chain of constrained call edges. -/ def Bytecode.Toplevel.needsCircuit (t : Bytecode.Toplevel) : Array Bool := Id.run do @@ -118,13 +197,24 @@ def Source.Toplevel.compile (t : Source.Toplevel) : Except String CompiledToplev let (bytecodeRaw, preNameMap) ← concDecls.toBytecode let (bytecodeDedup, remap) := bytecodeRaw.deduplicate let needs := bytecodeDedup.needsCircuit - let bytecode := { bytecodeDedup with + let bytecode : Bytecode.Toplevel := { bytecodeDedup with functions := bytecodeDedup.functions.mapIdx fun i f => { f with constrained := needs[i]! } } let nameMap := preNameMap.fold (init := (∅ : Std.HashMap Global Bytecode.FunIdx)) fun acc name idx => acc.insert name (remap idx) + -- Singleton circuits are labeled with (one of) the function's source names. + let reverseMap := nameMap.fold (init := (∅ : Std.HashMap Bytecode.FunIdx String)) + fun acc global idx => if acc.contains idx then acc else acc.insert idx (toString global) + let bytecode := { bytecode with + circuits := bytecode.singletonCircuits fun i => reverseMap[i]?.getD s!"" } pure (CompiledToplevel.mk t bytecode nameMap) +/-- `compile`, then apply a function grouping (see +`CompiledToplevel.groupFunctions`). -/ +def Source.Toplevel.compileWithGroups (t : Source.Toplevel) + (groups : Array (String × Array String)) : Except String CompiledToplevel := do + (← t.compile).groupFunctions groups + /-- Progress helper: given success of the three `Except`-returning stages, `compile` as a whole returns `.ok` (the remaining stages — `deduplicate`, `needsCircuit`, the field-setter `mapIdx`, the name-map `fold`, and the diff --git a/Ix/Aiur/Compiler/Lower.lean b/Ix/Aiur/Compiler/Lower.lean index 4cc3dc33..136df82e 100644 --- a/Ix/Aiur/Compiler/Lower.lean +++ b/Ix/Aiur/Compiler/Lower.lean @@ -631,7 +631,7 @@ def Concrete.Decls.toBytecode (decls : Concrete.Decls) : let memSizes := layoutMState.memSizes.fold (·.insert ·) memSizes pure (functions.push function, memSizes, nameMap) | _ => pure acc - pure (⟨functions, memSizes.toArray⟩, nameMap) + pure (⟨functions, memSizes.toArray, #[]⟩, nameMap) end Aiur diff --git a/Ix/Aiur/Stages/Bytecode.lean b/Ix/Aiur/Stages/Bytecode.lean index acd8e257..21c52fd6 100644 --- a/Ix/Aiur/Stages/Bytecode.lean +++ b/Ix/Aiur/Stages/Bytecode.lean @@ -122,9 +122,33 @@ structure Function where constrained : Bool deriving Inhabited, Repr +/-- A circuit of the proving system, backing one or more functions. By +default every constrained function gets a singleton circuit named after it; +`CompiledToplevel.groupFunctions` can regroup several functions into one +circuit whose branching selects the member function. `layout` is the merged +layout: max `inputSize`, sum of `selectors`, max `auxiliaries` (which +includes the single shared multiplicity column), max `lookups` (slot 0 is +the shared return lookup). -/ +structure Circuit where + name : String + members : Array FunIdx + layout : FunctionLayout + deriving Inhabited, Repr + +/-- Merged layout of a group of functions (see `Circuit`). -/ +def FunctionLayout.merge (a b : FunctionLayout) : FunctionLayout where + inputSize := a.inputSize.max b.inputSize + selectors := a.selectors + b.selectors + auxiliaries := a.auxiliaries.max b.auxiliaries + lookups := a.lookups.max b.lookups + structure Toplevel where functions : Array Function memorySizes : Array Nat + /-- Circuit partition of the constrained functions, in first-occurrence + order. Built by `Source.Toplevel.compile` (singletons by default; see + `CompiledToplevel.groupFunctions`); empty on a freshly lowered toplevel. -/ + circuits : Array Circuit := #[] deriving Repr end Bytecode diff --git a/Ix/Aiur/Statistics.lean b/Ix/Aiur/Statistics.lean index aa91de16..05603516 100644 --- a/Ix/Aiur/Statistics.lean +++ b/Ix/Aiur/Statistics.lean @@ -45,6 +45,13 @@ structure CircuitStats where /-- FFT cost at the uncached height `height + cacheHits` (fixed-height gadget circuits keep their normal cost). Feeds `totalUncachedFftCost`. -/ uncachedFftCost : Float + -- TEMP (grouping instrumentation, revert with this commit): the circuit + -- layout shape, for picking group members — merging is cheapest between + -- circuits whose auxiliaries and lookups are CLOSE (both merge by max; + -- selectors sum). Zero for memory/gadget circuits (no function layout). + selectors : Nat := 0 + auxiliaries : Nat := 0 + lookups : Nat := 0 structure ExecutionStats where circuits : Array CircuitStats @@ -81,44 +88,39 @@ def computeStats (compiled : CompiledToplevel) (queryCounts : Array QueryCount) (logBlowup : Nat := defaultCommitmentParameters.logBlowup) : ExecutionStats := let t := compiled.bytecode - -- Invert nameMap to get FunIdx → String - let reverseMap := compiled.nameMap.fold (init := (∅ : Std.HashMap Bytecode.FunIdx String)) - fun acc global idx => if !acc.contains idx then acc.insert idx (toString global) else acc let nAllFuns := t.functions.size - let nConstrained := t.functions.foldl (fun n f => if f.constrained then n + 1 else n) 0 - -- Shapes arrive in canonical system order: constrained functions - -- (ascending index), memories, `Bytes1`, `Bytes2`. A mismatch means the - -- shapes were built from a different toplevel; misindexing would silently - -- attribute costs to the wrong circuits. - if shapes.size != nConstrained + t.memorySizes.size + 2 then + -- Shapes arrive in canonical system order: function circuits (grouped; + -- singletons for ungrouped functions, in ascending member index), + -- memories, `Bytes1`, `Bytes2`. A mismatch means the shapes were built + -- from a different toplevel; misindexing would silently attribute costs + -- to the wrong circuits. + if shapes.size != t.circuits.size + t.memorySizes.size + 2 then panic! s!"computeStats: {shapes.size} circuit shapes for \ - {nConstrained} constrained functions + {t.memorySizes.size} memories + 2 gadgets" + {t.circuits.size} function circuits + {t.memorySizes.size} memories + 2 gadgets" else let mkStats (name : String) (shape : CircuitShape) (h hits : Nat) : CircuitStats := { name, width := shape.committedWidth, height := h, cacheHits := hits, fftCost := fftCost shape h logBlowup, uncachedFftCost := fftCost shape (h + hits) logBlowup } - let functionCircuits := Id.run do - let mut acc := #[] - let mut shapeIdx := 0 - for i in [:nAllFuns] do - if t.functions[i]!.constrained then - let shape := shapes[shapeIdx]! - shapeIdx := shapeIdx + 1 - let qc := queryCounts[i]! - let name := reverseMap[i]?.getD s!"" - acc := acc.push - (mkStats name shape qc.uniqueRows (qc.totalHits - qc.uniqueRows)) - acc + -- One row per function circuit: heights and cache hits are summed over + -- the circuit's member functions (singletons sum over one). + let functionCircuits := t.circuits.mapIdx fun cIdx c => + let shape := shapes[cIdx]! + let (h, hits) := c.members.foldl (init := (0, 0)) fun (h, hits) i => + let qc := queryCounts[i]! + (h + qc.uniqueRows, hits + (qc.totalHits - qc.uniqueRows)) + { mkStats c.name shape h hits with + selectors := c.layout.selectors, auxiliaries := c.layout.auxiliaries, + lookups := c.layout.lookups } let memoryCircuits := t.memorySizes.mapIdx fun i size => - let shape := shapes[nConstrained + i]! + let shape := shapes[t.circuits.size + i]! let qc := queryCounts[nAllFuns + i]! mkStats s!"memory[{size}]" shape qc.uniqueRows (qc.totalHits - qc.uniqueRows) -- The byte gadgets commit full-table traces in every proof: their height -- is the (fixed) preprocessed height, independent of the query set, so -- they carry no cache-hit counterfactual. let gadgetCircuits := #["Bytes1", "Bytes2"].mapIdx fun i name => - let shape := shapes[nConstrained + t.memorySizes.size + i]! + let shape := shapes[t.circuits.size + t.memorySizes.size + i]! mkStats name shape shape.preprocessedHeight 0 let circuits := (functionCircuits ++ memoryCircuits ++ gadgetCircuits).qsort (·.fftCost > ·.fftCost) @@ -169,9 +171,14 @@ def printStats (stats : ExecutionStats) : IO Unit := do let n := f.round.toUInt64.toNat toString n let wFftCost := stats.circuits.foldl (fun m cs => Nat.max m (formatSci cs.fftCost).length) 8 + -- TEMP (grouping instrumentation, revert with this commit) + let wSel := stats.circuits.foldl (fun m cs => Nat.max m (toString cs.selectors).length) 3 + let wAux := stats.circuits.foldl (fun m cs => Nat.max m (toString cs.auxiliaries).length) 3 + let wLkp := stats.circuits.foldl (fun m cs => Nat.max m (toString cs.lookups).length) 3 let wPct := 7 let wCum := 7 - let totalW := wName + 1 + wWidth + 1 + wHeight + 1 + wHits + 1 + wFftCost + 1 + wPct + 1 + wCum + let totalW := wName + 1 + wWidth + 1 + wSel + 1 + wAux + 1 + wLkp + 1 + + wHeight + 1 + wHits + 1 + wFftCost + 1 + wPct + 1 + wCum let totalWidth := stats.circuits.foldl (· + ·.width) 0 let savedPct := if stats.totalUncachedFftCost == 0.0 then "0.00%" @@ -184,14 +191,14 @@ def printStats (stats : ExecutionStats) : IO Unit := do IO.println s!"Total cache hits: {stats.totalCacheHits}" IO.println s!"Total saved cost: {savedPct}" IO.println sep - IO.println s!"{padRight "Name" wName} {padLeft "Width" wWidth} {padLeft "Height" wHeight} {padLeft "Hits" wHits} {padLeft "FFT cost" wFftCost} {padLeft "%" wPct} {padLeft "%++" wCum}" + IO.println s!"{padRight "Name" wName} {padLeft "Width" wWidth} {padLeft "Sel" wSel} {padLeft "Aux" wAux} {padLeft "Lkp" wLkp} {padLeft "Height" wHeight} {padLeft "Hits" wHits} {padLeft "FFT cost" wFftCost} {padLeft "%" wPct} {padLeft "%++" wCum}" IO.println sep let mut cumFftCost : Float := 0.0 for cs in stats.circuits do cumFftCost := cumFftCost + cs.fftCost let pct := formatPercent cs.fftCost stats.totalFftCost let cum := formatPercent cumFftCost stats.totalFftCost - IO.println s!"{padRight cs.name wName} {padLeft (toString cs.width) wWidth} {padLeft (toString cs.height) wHeight} {padLeft (toString cs.cacheHits) wHits} {padLeft (formatSci cs.fftCost) wFftCost} {padLeft pct wPct} {padLeft cum wCum}" + IO.println s!"{padRight cs.name wName} {padLeft (toString cs.width) wWidth} {padLeft (toString cs.selectors) wSel} {padLeft (toString cs.auxiliaries) wAux} {padLeft (toString cs.lookups) wLkp} {padLeft (toString cs.height) wHeight} {padLeft (toString cs.cacheHits) wHits} {padLeft (formatSci cs.fftCost) wFftCost} {padLeft pct wPct} {padLeft cum wCum}" end Aiur diff --git a/Ix/Cli/CheckCmd.lean b/Ix/Cli/CheckCmd.lean index 56553878..ef647116 100644 --- a/Ix/Cli/CheckCmd.lean +++ b/Ix/Cli/CheckCmd.lean @@ -656,7 +656,7 @@ def runCheckCmd (p : Cli.Parsed) : IO UInt32 := do pure 1 pure go else do - let compiled ← match toplevel.compile with + let compiled ← match toplevel.compileWithGroups IxVM.coldGroups with | .error e => IO.eprintln s!"Compilation failed: {e}"; return 1 | .ok c => pure c let go (_ : Ix.Claim) (envHandle? : Option Aiur.EnvHandle) (target : Target) @@ -669,7 +669,7 @@ def runCheckCmd (p : Cli.Parsed) : IO UInt32 := do return (← runShardCheckManifest manifest ixe k (fun c w l => runOne c none (.leanW w) l)) else do - let compiled ← match toplevel.compile with + let compiled ← match toplevel.compileWithGroups IxVM.coldGroups with | .error e => IO.eprintln s!"Compilation failed: {e}"; return 1 | .ok c => pure c return (← runShardCheckManifestNative manifest ixe k compiled printStats statsOut useBytecode) @@ -678,7 +678,7 @@ def runCheckCmd (p : Cli.Parsed) : IO UInt32 := do return (← runShardCheckAll manifest ixe ((p.flag? "jobs").map (·.as! Nat)) (fun c w l => runOne c none (.leanW w) l)) else do - let compiled ← match toplevel.compile with + let compiled ← match toplevel.compileWithGroups IxVM.coldGroups with | .error e => IO.eprintln s!"Compilation failed: {e}"; return 1 | .ok c => pure c return (← runShardManifestAllNative manifest ixe diff --git a/Ix/Cli/ProveCmd.lean b/Ix/Cli/ProveCmd.lean index a35609d9..787853b3 100644 --- a/Ix/Cli/ProveCmd.lean +++ b/Ix/Cli/ProveCmd.lean @@ -146,7 +146,7 @@ def runProveCmd (p : Cli.Parsed) : IO UInt32 := do let toplevel ← match IxVM.ixVM with | .error e => IO.eprintln s!"toplevel merging failed: {e}"; return 1 | .ok t => pure t - let compiled ← match toplevel.compile with + let compiled ← match toplevel.compileWithGroups IxVM.coldGroups with | .error e => IO.eprintln s!"compilation failed: {e}"; return 1 | .ok c => pure c let aiurSystem := Aiur.AiurSystem.build compiled.bytecode commitmentParameters friParameters diff --git a/Ix/Cli/VerifyCmd.lean b/Ix/Cli/VerifyCmd.lean index 38135257..c2262b0d 100644 --- a/Ix/Cli/VerifyCmd.lean +++ b/Ix/Cli/VerifyCmd.lean @@ -73,7 +73,7 @@ def verifyOneProof (aiurSystem : Aiur.AiurSystem) (compiled : Aiur.CompiledTople def buildBackend : IO (Except String (Aiur.AiurSystem × Aiur.CompiledToplevel)) := do match IxVM.ixVM with | .error e => return .error s!"toplevel merging failed: {e}" - | .ok toplevel => match toplevel.compile with + | .ok toplevel => match toplevel.compileWithGroups IxVM.coldGroups with | .error e => return .error s!"compilation failed: {e}" | .ok compiled => return .ok (Aiur.AiurSystem.build compiled.bytecode commitmentParameters friParameters, compiled) diff --git a/Ix/IxVM.lean b/Ix/IxVM.lean index 45c6b898..0add2290 100644 --- a/Ix/IxVM.lean +++ b/Ix/IxVM.lean @@ -1,6 +1,7 @@ module public import Ix.Aiur.Meta public import Ix.IxVM.Core +public import Ix.IxVM.ColdGroups public import Ix.IxVM.ByteStream public import Ix.IxVM.Blake3 public import Ix.IxVM.RBTreeMap diff --git a/Ix/IxVM/ColdGroups.lean b/Ix/IxVM/ColdGroups.lean new file mode 100644 index 00000000..32873147 --- /dev/null +++ b/Ix/IxVM/ColdGroups.lean @@ -0,0 +1,20 @@ +module + +/-! +Circuit-grouping data for the IxVM kernel toplevel, applied wherever the +kernel is compiled for proving or verifying (see +`CompiledToplevel.groupFunctions`). Empty = no grouping: every constrained +function keeps its singleton circuit. Fill from measured workload +statistics; a stale grouping stays sound (grouping never affects +semantics), only less efficient. +-/ + +public section + +namespace IxVM + +def coldGroups : Array (String × Array String) := #[] + +end IxVM + +end diff --git a/Ix/MultiStark.lean b/Ix/MultiStark.lean index c9911442..6881ad13 100644 --- a/Ix/MultiStark.lean +++ b/Ix/MultiStark.lean @@ -11,6 +11,7 @@ public import Ix.MultiStark.Keccak public import Ix.MultiStark.Pcs public import Ix.MultiStark.SystemDeserialize public import Ix.MultiStark.Verifier +public import Ix.MultiStark.VerifierColdGroups public import Ix.MultiStark.Tests /-! diff --git a/Ix/MultiStark/VerifierColdGroups.lean b/Ix/MultiStark/VerifierColdGroups.lean new file mode 100644 index 00000000..58359883 --- /dev/null +++ b/Ix/MultiStark/VerifierColdGroups.lean @@ -0,0 +1,20 @@ +module + +/-! +Circuit-grouping data for the recursive-verifier toplevel, applied wherever +it is compiled for proving or verifying (see +`CompiledToplevel.groupFunctions`). Empty = no grouping: every constrained +function keeps its singleton circuit. Fill from measured workload +statistics; a stale grouping stays sound (grouping never affects +semantics), only less efficient. +-/ + +public section + +namespace MultiStark + +def verifierColdGroups : Array (String × Array String) := #[] + +end MultiStark + +end diff --git a/Tests/Aiur/Aiur.lean b/Tests/Aiur/Aiur.lean index d78fa2f6..d62f0d04 100644 --- a/Tests/Aiur/Aiur.lean +++ b/Tests/Aiur/Aiur.lean @@ -719,6 +719,38 @@ def toplevel := ⟦ let s5 = c[4] + c[0]; -- 255 s1 + s2 + 10 * s3 + s4 + s5 -- 1309 } + + --------------------------------------------------------------------------- + -- Grouped circuits (`CompiledToplevel.groupFunctions`): the test runner + -- groups these three into one circuit whose branching selects the member. + -- Grouping is a circuit-level choice, so there is NO source annotation: + -- the same functions also run ungrouped in the plain suite. Members + -- differ in arity, output and branch count, call each other (through the + -- shared circuit) and recurse. + --------------------------------------------------------------------------- + fn grouped_double(x: G) -> G { + x + x + } + + -- Different arity, a match (two selectors), calls a fellow group member. + fn grouped_pick(t: G, a: G, b: G) -> G { + match t { + 0 => grouped_double(a), + _ => b, + } + } + + -- Recursive group member: self-calls route through the shared circuit. + fn grouped_sum_range(n: G) -> G { + match n { + 0 => 0, + _ => n + grouped_sum_range(n - 1), + } + } + + pub fn calls_grouped(t: G, a: G, b: G) -> G { + grouped_pick(t, a, b) + grouped_sum_range(a) + } ⟧ /-- The PROVING suite: every case runs the full prove+verify pipeline @@ -843,6 +875,53 @@ def aiurTestCases : List AiurTestCase := [ -- Unconstrained g_to_bytes / g_inverse hints: all cases in one proof .prove `hint_test #[] #[1309], + + -- Grouped-circuit member functions, run UNGROUPED here (the grouped + -- variant runs in the grouped env; see `testGroups`). + -- t=0 → grouped_double(5) + Σ1..5 = 10 + 15 = 25; t≠0 → 9 + Σ1..3 = 15. + .prove `calls_grouped #[0, 5, 9] #[25] + (label := "calls_grouped(0,5,9)"), + .prove `calls_grouped #[1, 3, 9] #[15] + (label := "calls_grouped(1,3,9)"), ] +/-- The grouping the `aiur` runner applies for the grouped environment. -/ +def testGroups : Array (String × Array String) := + #[("test_group", #["grouped_double", "grouped_pick", "grouped_sum_range"])] + +def groupedTestCases : List AiurTestCase := [ + .prove `calls_grouped #[0, 5, 9] #[25] + (label := "calls_grouped(0,5,9) [grouped]"), + .prove `calls_grouped #[1, 3, 9] #[15] + (label := "calls_grouped(1,3,9) [grouped]"), +] + +/-- Structural checks on the grouped partition: the grouped circuit exists, +holds exactly its members, its layout follows the merge rule (max inputs, +summed selectors, max auxiliaries, max lookups), and every constrained +function lands in exactly one circuit. -/ +def groupingStructureChecks (compiled : Aiur.CompiledToplevel) : TestSeq := + let t := compiled.bytecode + let memberOf := fun (name : Lean.Name) => compiled.getFuncIdx name |>.get! + let expectedMembers := + #[`grouped_double, `grouped_pick, `grouped_sum_range].map memberOf + match t.circuits.find? (·.name == "test_group") with + | none => test "test_group circuit exists" false + | some c => + let layouts := c.members.map (t.functions[·]!.layout) + let expected := layouts.foldl (init := (⟨0, 0, 0, 0⟩ : Aiur.Bytecode.FunctionLayout)) + Aiur.Bytecode.FunctionLayout.merge + let allCircuitMembers := t.circuits.flatMap (·.members) + let constrained := (Array.range t.functions.size).filter + (t.functions[·]!.constrained) + test "test_group circuit exists" true ++ + test "test_group members" (c.members == expectedMembers) ++ + test "test_group layout follows the merge rule" + (c.layout.inputSize == expected.inputSize && + c.layout.selectors == expected.selectors && + c.layout.auxiliaries == expected.auxiliaries && + c.layout.lookups == expected.lookups) ++ + test "every constrained function is in exactly one circuit" + (allCircuitMembers.qsort (· < ·) == constrained) + end diff --git a/Tests/Aiur/Common.lean b/Tests/Aiur/Common.lean index f90e64f6..39ffb930 100644 --- a/Tests/Aiur/Common.lean +++ b/Tests/Aiur/Common.lean @@ -72,10 +72,13 @@ structure AiurTestEnv where aiurSystem : Aiur.AiurSystem shapes : Array Aiur.CircuitShape -def AiurTestEnv.build (toplevelFn : Except Aiur.Global Aiur.Source.Toplevel) : +def AiurTestEnv.build (toplevelFn : Except Aiur.Global Aiur.Source.Toplevel) + (groups : Array (String × Array String) := #[]) : Except String AiurTestEnv := do let toplevel ← toplevelFn.mapError toString let compiled ← toplevel.compile + let compiled ← if groups.isEmpty then pure compiled + else compiled.groupFunctions groups let decls ← toplevel.mkDecls.mapError toString let aiurSystem := Aiur.AiurSystem.build compiled.bytecode commitmentParameters friParameters return ⟨compiled, decls, aiurSystem, aiurSystem.circuitShapes⟩ diff --git a/Tests/Main.lean b/Tests/Main.lean index 59203355..32858d64 100644 --- a/Tests/Main.lean +++ b/Tests/Main.lean @@ -130,7 +130,17 @@ def primaryRunners : List (String × IO UInt32) := [ IO.println "aiur-prove" match AiurTestEnv.build (pure toplevel) with | .error e => IO.eprintln s!"Aiur setup failed: {e}"; return 1 - | .ok env => LSpec.lspecEachIO aiurTestCases fun tc => pure (env.runTestCase tc)), + | .ok env => do + let r1 ← LSpec.lspecEachIO aiurTestCases fun tc => pure (env.runTestCase tc) + -- The same toplevel with `testGroups` applied: the members share one + -- circuit, and the whole suite of grouped cases proves through it. + match AiurTestEnv.build (pure toplevel) testGroups with + | .error e => IO.eprintln s!"Aiur grouped setup failed: {e}"; return 1 + | .ok genv => do + let r2 ← LSpec.lspecEachIO groupedTestCases fun tc => pure (genv.runTestCase tc) + let r3 ← LSpec.lspecIO + (.ofList [("aiur-grouping", [groupingStructureChecks genv.compiled])]) [] + return if r1 == 0 && r2 == 0 && r3 == 0 then 0 else 1), ("aiur-hashes", do IO.println "aiur-hashes" let .ok blake3Env := AiurTestEnv.build (do @@ -172,7 +182,7 @@ def ignoredRunners (env : Lean.Environment) : List (String × IO UInt32) := [ -- committed kernel system). let kernelUnitTests := .exec `kernel_unit_tests let serdeTest ← serdeNatAddComm env - match AiurTestEnv.build IxVM.ixVM, AiurTestEnv.build IxVM.ixVMFull with + match AiurTestEnv.build IxVM.ixVM IxVM.coldGroups, AiurTestEnv.build IxVM.ixVMFull with | .error e, _ | _, .error e => IO.eprintln s!"IxVM env build failed: {e}"; return 1 | .ok v2Env, .ok v2FullEnv => diff --git a/Tests/MultiStark.lean b/Tests/MultiStark.lean index c46b9477..f59ded1e 100644 --- a/Tests/MultiStark.lean +++ b/Tests/MultiStark.lean @@ -177,7 +177,7 @@ def endToEndSuite : IO UInt32 := do let vTop ← match MultiStark.multiStark with | .error e => IO.eprintln s!"verifier toplevel merge failed: {e}"; return 1 | .ok t => pure t - let vCompiled ← match vTop.compile with + let vCompiled ← match vTop.compileWithGroups MultiStark.verifierColdGroups with | .error e => IO.eprintln s!"verifier compilation failed: {e}"; return 1 | .ok c => pure c let vIdx ← match vCompiled.getFuncIdx `verify_multi_stark_proof with diff --git a/crates/aiur/src/bytecode.rs b/crates/aiur/src/bytecode.rs index 57e426e8..3349906f 100644 --- a/crates/aiur/src/bytecode.rs +++ b/crates/aiur/src/bytecode.rs @@ -5,6 +5,23 @@ use super::G; pub struct Toplevel { pub functions: Vec, pub memory_sizes: Vec, + /// Circuit partition of the constrained functions, in first-occurrence + /// order. Computed by the Lean compiler (singletons by default; + /// `CompiledToplevel.groupFunctions` regroups); every constrained + /// function appears in exactly one circuit. + pub circuits: Vec, +} + +/// A circuit of the proving system, backing one or more functions. Ungrouped +/// functions get a singleton circuit; grouped functions share one circuit +/// whose branching selects the member function. +/// +/// `layout` is the merged layout: max `input_size`, sum of `selectors`, max +/// `auxiliaries` (which includes the single shared multiplicity column), max +/// `lookups` (slot 0 is the shared return lookup). +pub struct Circuit { + pub members: Vec, + pub layout: FunctionLayout, } pub struct Function { diff --git a/crates/aiur/src/constraints.rs b/crates/aiur/src/constraints.rs index 26488ce5..8691198b 100644 --- a/crates/aiur/src/constraints.rs +++ b/crates/aiur/src/constraints.rs @@ -6,7 +6,7 @@ use std::{array, ops::Range, sync::LazyLock}; use crate::{ FxIndexMap, G, - bytecode::{Block, Ctrl, Function, FunctionLayout, Op, Toplevel, ValIdx}, + bytecode::{Block, Ctrl, Op, Toplevel, ValIdx}, function_channel, gadgets::{ AiurGadget, @@ -65,11 +65,18 @@ pub struct Constraints { } struct ConstraintState { + /// Index of the circuit member currently being walked. function_index: G, - /// Exactly one selector: the function has a single leaf block (no - /// matches), so every lookup slot is written by exactly one branch. + /// Exactly one selector: the circuit backs a single function with a + /// single leaf block (no matches), so every lookup slot is written by + /// exactly one branch. branchless: bool, - layout: FunctionLayout, + /// Input size of the current member (inputs live in columns + /// `0..input_size` for every member; the circuit reserves the max). + input_size: usize, + /// Column of the current member's first selector: the circuit's input + /// block plus the selector counts of the members walked before it. + sel_base: usize, column: usize, lookup: usize, lookups: Vec>, @@ -88,7 +95,7 @@ struct SharedState { impl ConstraintState { fn selector_index(&self, sel: usize) -> usize { - sel + self.layout.input_size + sel + self.sel_base } /// Selector-gate a lookup argument. Lookup slots shared across branches @@ -131,34 +138,75 @@ impl ConstraintState { } impl Toplevel { + /// Build the constraints of one circuit. The circuit's members are walked + /// like branches of a single function: each walk restarts the auxiliary + /// column / lookup-slot counters (so members share those, like match arms + /// do), while selector columns are laid out consecutively per member. All + /// members fold their return message into the shared lookup slot 0, gated + /// by their own selectors and carrying their own function index, against + /// the single shared multiplicity column. pub fn build_constraints( &self, - function_index: usize, + circuit_index: usize, ) -> (Constraints, Vec>) { - let function = &self.functions[function_index]; + let circuit = &self.circuits[circuit_index]; + let layout = circuit.layout; let constraints = Constraints { zeros: vec![], - selectors: 0..0, - width: function.layout.width(), + selectors: layout.input_size..layout.input_size + layout.selectors, + width: layout.width(), }; let mut state = ConstraintState { - function_index: G::from_usize(function_index), - branchless: function.layout.selectors == 1, - layout: function.layout, + function_index: G::ZERO, + branchless: layout.selectors == 1, + input_size: 0, + sel_base: 0, column: 0, lookup: 0, map: vec![], - lookups: vec![empty_lookup(); function.layout.lookups], + lookups: vec![empty_lookup(); layout.lookups], constraints, yield_info: vec![], }; - function.build_constraints(&mut state); + // The shared multiplicity column: first auxiliary, right after the + // selectors. The return lookup occupies the first lookup slot. + let multiplicity = var(layout.input_size + layout.selectors); + state.lookups[0].multiplicity = -multiplicity; + let aux_start = layout.input_size + layout.selectors + 1; + let mut sel_base = layout.input_size; + let mut circuit_sel = Expr::from(G::ZERO); + for &member in &circuit.members { + let function = &self.functions[member]; + state.function_index = G::from_usize(member); + state.input_size = function.layout.input_size; + state.sel_base = sel_base; + state.column = aux_start; + state.lookup = 1; + state.map.clear(); + (0..function.layout.input_size).for_each(|i| state.map.push((var(i), 1))); + let body_sel = function.body.get_block_selector(&state); + circuit_sel = circuit_sel + body_sel.clone(); + function.body.collect_constraints(body_sel, &mut state); + debug_assert!(state.yield_info.is_empty()); + sel_base += function.layout.selectors; + } // The old `Air::eval` asserted each selector column boolean; the new // system compiles a constraint vector, so materialize those explicitly. for sel in state.constraints.selectors.clone() { let s = var(sel); state.constraints.zeros.push(s.clone() * (s - konst(G::ONE))); } + // Cross-member exclusivity: the circuit-level selector (the sum of the + // members' top-block selectors) must be boolean, so at most one member + // is active per row and the shared return lookup emits a single + // member's message. A singleton circuit already gets this from its top + // block's own boolean constraint. + if circuit.members.len() > 1 { + state + .constraints + .zeros + .push(circuit_sel.clone() * (Expr::from(G::ONE) - circuit_sel)); + } (state.constraints, state.lookups) } } @@ -167,26 +215,6 @@ fn empty_lookup() -> Lookup { Lookup { multiplicity: konst(G::ZERO), args: vec![] } } -impl Function { - fn build_constraints(&self, state: &mut ConstraintState) { - // the first columns are occupied by the input, which is also mapped - state.column += self.layout.input_size; - (0..self.layout.input_size).for_each(|i| state.map.push((var(i), 1))); - // then comes the selectors, which are not mapped - let init_sel = state.column; - let final_sel = state.column + self.layout.selectors; - state.constraints.selectors = init_sel..final_sel; - state.column = final_sel; - // the multiplicity occupies another column - let multiplicity = var(state.column); - state.column += 1; - // the return lookup occupies the first lookup slot - state.lookups[0].multiplicity = -multiplicity.clone(); - state.lookup += 1; - self.body.collect_constraints(self.body.get_block_selector(state), state); - } -} - impl Block { fn collect_constraints(&self, sel: Expr, state: &mut ConstraintState) { // Boolean constraint for this block's selector @@ -277,7 +305,7 @@ impl Ctrl { ]; // input args.extend( - (0..state.layout.input_size) + (0..state.input_size) .map(|arg| state.gate(&sel, state.map[arg].0.clone())), ); // output diff --git a/crates/aiur/src/synthesis.rs b/crates/aiur/src/synthesis.rs index 47d7ea37..6099213e 100644 --- a/crates/aiur/src/synthesis.rs +++ b/crates/aiur/src/synthesis.rs @@ -89,20 +89,17 @@ impl AiurSystem { }); }; - // Constrained functions (ascending index). - for i in 0..toplevel.functions.len() { - if !toplevel.functions[i].constrained { - continue; - } + // Function circuits, in partition order (singletons unless grouped). + for i in 0..toplevel.circuits.len() { let (constraints, lookups) = toplevel.build_constraints(i); - // A branchless function's lookup arguments are sent raw (degree 1; + // A branchless circuit's lookup arguments are sent raw (degree 1; // see `ConstraintState::gate`), so two lookups fit in one chained // accumulator step at degree 3 — within the degree the selector-gated - // constraints already pay for. Branching functions keep k = 1: their + // constraints already pay for. Branching circuits keep k = 1: their // superposed arguments are degree 2, and grouping would push the // logUp constraints past the quotient budget. let group_size = - if toplevel.functions[i].layout.selectors == 1 && lookups.len() >= 2 { + if toplevel.circuits[i].layout.selectors == 1 && lookups.len() >= 2 { 2 } else { 1 @@ -156,11 +153,8 @@ impl AiurSystem { /// order the circuits were chained in [`AiurSystem::build`], so index `i` /// of the returned `Vec` corresponds to `self.system.circuits[i]`. fn circuit_types(&self) -> Vec { - let functions = (0..self.toplevel.functions.len()).filter_map(|idx| { - self.toplevel.functions[idx] - .constrained - .then_some(CircuitType::Function { idx }) - }); + let functions = (0..self.toplevel.circuits.len()) + .map(|idx| CircuitType::Function { idx }); let memories = self .toplevel .memory_sizes @@ -384,6 +378,25 @@ mod tests { /// fresh auxiliary column pinned by `sel * (col - a*b)`. /// - `lookups = 1`: the function-provide (return) lookup in slot 0, which /// pulls the claim `[function_channel, fun_idx, a, b, a*b]`. + /// + /// Test-side singleton partition (production circuits come pre-built from + /// the Lean compiler). + fn with_singleton_circuits( + functions: Vec, + memory_sizes: Vec, + ) -> Toplevel { + let circuits = functions + .iter() + .enumerate() + .filter(|(_, f)| f.constrained) + .map(|(i, f)| crate::bytecode::Circuit { + members: vec![i], + layout: f.layout, + }) + .collect(); + Toplevel { functions, memory_sizes, circuits } + } + fn mul_toplevel() -> Toplevel { let body = Block { ops: vec![Op::Mul(0, 1)], ctrl: Ctrl::Return(0, vec![2]) }; @@ -398,7 +411,7 @@ mod tests { entry: true, constrained: true, }; - Toplevel { functions: vec![function], memory_sizes: vec![] } + with_singleton_circuits(vec![function], vec![]) } fn xor_splits_toplevel() -> Toplevel { @@ -417,7 +430,7 @@ mod tests { entry: true, constrained: true, }; - Toplevel { functions: vec![function], memory_sizes: vec![] } + with_singleton_circuits(vec![function], vec![]) } #[test] @@ -523,7 +536,7 @@ mod tests { constrained: true, }; - Toplevel { functions: vec![f, g], memory_sizes: vec![1] } + with_singleton_circuits(vec![f, g], vec![1]) } #[test] diff --git a/crates/aiur/src/trace.rs b/crates/aiur/src/trace.rs index e87dfe15..66e88cf9 100644 --- a/crates/aiur/src/trace.rs +++ b/crates/aiur/src/trace.rs @@ -12,7 +12,7 @@ use rayon::{ use crate::{ FxIndexMap, G, - bytecode::{Block, Ctrl, Function, Op, Toplevel}, + bytecode::{Block, Ctrl, Function, FunctionLayout, Op, Toplevel}, execute::{ IOBuffer, IOKeyInfo, QueryRecord, find_unconstrained_big_uint_div_mod, g_inverse_value, @@ -20,6 +20,7 @@ use crate::{ function_channel, gadgets::{bytes1::Bytes1, bytes2::Bytes2}, memory::Memory, + querymap::QueryRef, u8_add_channel, u8_and_channel, u8_bit_decomposition_channel, u8_less_than_channel, u8_mul_channel, u8_or_channel, u8_range_check_channel, u8_shift_left_channel, u8_shift_right_channel, u8_sub_channel, @@ -55,15 +56,23 @@ fn u32_sum(values: &[u64]) -> ([G; 4], G) { } impl<'a, 'b> ColumnMutSlice<'a, 'b> { + /// Slice a circuit row into the regions of one member function: the + /// member's inputs are a prefix of the circuit's input block, its + /// selectors a sub-range of the circuit's selector block at `sel_offset`, + /// and the auxiliary block is shared by all members. fn from_slice( function: &Function, + circuit_layout: &FunctionLayout, + sel_offset: usize, slice: &'a mut [G], lookups: &'a mut LookupRowMut<'b, G>, ) -> Self { - let (inputs, slice) = slice.split_at_mut(function.layout.input_size); - let (selectors, slice) = slice.split_at_mut(function.layout.selectors); - let (auxiliaries, slice) = slice.split_at_mut(function.layout.auxiliaries); - assert!(slice.is_empty()); + let (inputs, slice) = slice.split_at_mut(circuit_layout.input_size); + let (selectors, auxiliaries) = slice.split_at_mut(circuit_layout.selectors); + assert_eq!(auxiliaries.len(), circuit_layout.auxiliaries); + let inputs = &mut inputs[..function.layout.input_size]; + let selectors = + &mut selectors[sel_offset..sel_offset + function.layout.selectors]; Self { inputs, selectors, auxiliaries, lookups } } @@ -92,22 +101,49 @@ struct TraceContext<'a> { query_record: &'a QueryRecord, } +/// One row of a circuit trace: the member function it belongs to, the +/// member's selector offset within the circuit, its function index, and the +/// recorded query. +struct RowMeta<'a> { + function: &'a Function, + sel_offset: usize, + function_index: G, + inputs: &'a [G], + result: QueryRef<'a>, +} + impl Toplevel { pub fn witness_data( &self, - function_index: usize, + circuit_index: usize, query_record: &QueryRecord, io_buffer: &IOBuffer, slot_arg_widths: &[usize], ) -> (RowMajorMatrix, LookupValues) { - let func = &self.functions[function_index]; - let width = func.width(); - let unfiltered_queries = &query_record.function_queries[function_index]; - let queries = unfiltered_queries - .iter() - .filter(|(_, res)| !res.multiplicity.is_zero()) - .collect::>(); - let height_no_padding = queries.len(); + let circuit = &self.circuits[circuit_index]; + let layout = &circuit.layout; + let width = layout.width(); + // Concatenate the members' queried rows, in member order. + let mut rows_meta = Vec::new(); + let mut sel_offset = 0; + for &member in &circuit.members { + let function = &self.functions[member]; + let function_index = G::from_usize(member); + rows_meta.extend( + query_record.function_queries[member] + .iter() + .filter(|(_, res)| !res.multiplicity.is_zero()) + .map(|(inputs, result)| RowMeta { + function, + sel_offset, + function_index, + inputs, + result, + }), + ); + sel_offset += function.layout.selectors; + } + let height_no_padding = rows_meta.len(); // An unqueried circuit yields an EMPTY trace (not a padded height-1 one): // the prover deactivates it, so it is neither committed nor opened. let height = if height_no_padding == 0 { @@ -126,21 +162,27 @@ impl Toplevel { .zip(row_writers[..height_no_padding].par_iter_mut()) .enumerate() .for_each(|(i, (row, lookups))| { - let (inputs, result) = queries[i]; + let meta = &rows_meta[i]; let index = &mut ColumnIndex { auxiliary: 0, // we skip the first lookup, which is reserved for return lookup: 1, }; - let slice = &mut ColumnMutSlice::from_slice(func, row, lookups); + let slice = &mut ColumnMutSlice::from_slice( + meta.function, + layout, + meta.sel_offset, + row, + lookups, + ); let context = TraceContext { - function_index: G::from_usize(function_index), - inputs, - multiplicity: result.multiplicity, - output: result.output, + function_index: meta.function_index, + inputs: meta.inputs, + multiplicity: meta.result.multiplicity, + output: meta.result.output, query_record, }; - func.populate_row(index, slice, context, io_buffer); + meta.function.populate_row(index, slice, context, io_buffer); }); drop(row_writers); let trace = RowMajorMatrix::new(rows, width); diff --git a/crates/ffi/src/aiur/toplevel.rs b/crates/ffi/src/aiur/toplevel.rs index 4b88cdac..fbead94f 100644 --- a/crates/ffi/src/aiur/toplevel.rs +++ b/crates/ffi/src/aiur/toplevel.rs @@ -2,12 +2,15 @@ use multi_stark::p3_field::PrimeCharacteristicRing; use lean_ffi::object::{LeanBorrowed, LeanCtor, LeanRef}; +use crate::lean::LeanAiurCircuit; use crate::lean::LeanAiurFunction; use crate::lean::LeanAiurToplevel; use aiur::{ FxIndexMap, G, - bytecode::{Block, Ctrl, Function, FunctionLayout, Op, Toplevel, ValIdx}, + bytecode::{ + Block, Circuit, Ctrl, Function, FunctionLayout, Op, Toplevel, ValIdx, + }, }; use crate::aiur::{lean_unbox_g, lean_unbox_nat_as_usize}; @@ -294,14 +297,23 @@ fn decode_function(ctor: LeanCtor>) -> Function { Function { body, layout, entry, constrained } } +fn decode_circuit(ctor: LeanCtor>) -> Circuit { + let ctor = LeanAiurCircuit::from_ctor(ctor); + // Object field 0 is the circuit's display name (`String`), unused here. + let members = ctor.get_obj(1).as_array().map(|x| lean_unbox_nat_as_usize(&x)); + let layout = decode_function_layout(ctor.get_obj(2).as_ctor()); + Circuit { members, layout } +} + pub(crate) fn decode_toplevel( obj: &LeanAiurToplevel, ) -> Toplevel { let ctor = obj.as_ctor(); - let [functions_obj, memory_sizes_obj] = ctor.objs::<2>(); + let [functions_obj, memory_sizes_obj, circuits_obj] = ctor.objs::<3>(); let functions = functions_obj.as_array().map(|o| decode_function(o.as_ctor())); let memory_sizes = memory_sizes_obj.as_array().map(|x| lean_unbox_nat_as_usize(&x)); - Toplevel { functions, memory_sizes } + let circuits = circuits_obj.as_array().map(|o| decode_circuit(o.as_ctor())); + Toplevel { functions, memory_sizes, circuits } } diff --git a/crates/ffi/src/lean.rs b/crates/ffi/src/lean.rs index e9040c66..0cd5b9ed 100644 --- a/crates/ffi/src/lean.rs +++ b/crates/ffi/src/lean.rs @@ -270,8 +270,9 @@ lean_ffi::lean_inductive! { // --- Aiur types --- - LeanAiurToplevel [ { num_obj: 2 } ]; + LeanAiurToplevel [ { num_obj: 3 } ]; LeanAiurFunction [ { num_obj: 2, num_8: 2 } ]; + LeanAiurCircuit [ { num_obj: 3 } ]; // Aiur FFI result structures (`Ix/Aiur/Semantics/BytecodeFfi.lean`, // `Ix/Aiur/Protocol.lean`). `IOBuffer` hashmaps cross the boundary as