Function groups

Benchmarks/AggregatePolicy.lean

@@ -83,11 +83,12 @@ def parseShardIds (value : String) : Except String (Array Nat) := do
   pure ids
 
 def compileToplevel (label : String)
-    (source : Except Aiur.Global Aiur.Source.Toplevel) :
+    (source : Except Aiur.Global Aiur.Source.Toplevel)
+    (groups : Array (String × Array String)) :
     IO (Except String Aiur.CompiledToplevel) := do
   match source with
   | .error error => pure (.error s!"{label} toplevel merge failed: {error}")
-  | .ok top => match top.compile with
+  | .ok top => match top.compileWithGroups groups with
     | .error error => pure (.error s!"{label} compilation failed: {error}")
     | .ok compiled => pure (.ok compiled)
 
@@ -604,13 +605,13 @@ def main (args : List String) : IO UInt32 := do
   writeReport jsonPath? metadata0 rows #[] "preparing"
   TracingTexray.startSampler 25
   IO.println "[aggregate-policy] compiling IxVM and ixAggr systems"
-  let ixvmCompiled ← match ← compileToplevel "IxVM" IxVM.ixVM with
+  let ixvmCompiled ← match ← compileToplevel "IxVM" IxVM.ixVM IxVM.functionGroups with
     | .error error =>
       writeReport jsonPath? metadata0 rows #[] "error" (error? := some error)
       IO.eprintln error
       return 1
     | .ok compiled => pure compiled
-  let aggrCompiled ← match ← compileToplevel "ixAggr recursion" Aggr.ixAggr with
+  let aggrCompiled ← match ← compileToplevel "ixAggr recursion" Aggr.ixAggr Aggr.functionGroups with
     | .error error =>
       writeReport jsonPath? metadata0 rows #[] "error" (error? := some error)
       IO.eprintln error

Benchmarks/RecursionDebug.lean

@@ -74,7 +74,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.functionGroups
     | IO.eprintln "IxVM compile failed"; return none
   let entrypoint := if skipDeps then `verify_const else `verify_claim
   let some funIdx := compiled.getFuncIdx entrypoint
@@ -165,7 +165,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.verifierFunctionGroups
       | IO.eprintln "multi-stark verifier compile failed"; return 1
     let entries := vCompiled.nameMap.toArray.qsort (·.2 < ·.2)
     for (g, i) in entries do
@@ -199,7 +199,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.verifierFunctionGroups
       | 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

Benchmarks/Typecheck.lean

@@ -450,7 +450,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.functionGroups
     | throw (IO.userError "Compilation of IxVM kernel failed")
   let entrypoint := if skipDeps then `verify_const else `verify_claim
   let some funIdx := compiled.getFuncIdx entrypoint
@@ -479,7 +479,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.verifierFunctionGroups
         | 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")
@@ -493,7 +493,7 @@ def runTypecheckCmd (p : Cli.Parsed) : IO UInt32 := do
     if !join then pure none else do
       let .ok aggrTop := Aggr.ixAggr
         | throw (IO.userError "Merging ix_aggr failed")
-      let .ok aggrCompiled := aggrTop.compile
+      let .ok aggrCompiled := aggrTop.compileWithGroups Aggr.functionGroups
         | throw (IO.userError "Compilation of ix_aggr failed")
       let some aggrIdx := aggrCompiled.getFuncIdx `ix_aggr
         | throw (IO.userError "ix_aggr entrypoint missing")

Ix/Aggr.lean

@@ -1,6 +1,7 @@
 module
 public import Blake3.Rust
 public import Ix.Aggr.Circuit
+public import Ix.Aggr.FunctionGroups
 public import Ix.Aggr.Host
 public import Ix.AssumptionTree
 public import Ix.MultiStark

Ix/Aggr/FunctionGroups.lean

@@ -0,0 +1,20 @@
+module
+
+/-!
+Function-grouping data for the production aggregation toplevel
+(`Aggr.ixAggr`), 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 Aggr
+
+def functionGroups : Array (String × Array String) := #[]
+
+end Aggr
+
+end

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,46 @@ 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!"<fn {i}>" }
   pure (CompiledToplevel.mk t bytecode nameMap)
 
+/-- Name of the environment variable that switches function grouping OFF
+process-wide: when it is set to anything but `0` or the empty string,
+`compileWithGroups` ignores its grouping and compiles the singleton
+partition. For testing and measurement only - the prover and the verifier
+must see the same setting, since grouping changes the verifying key. -/
+def noFunctionGroupsEnvVar : String := "IX_NO_FUNCTION_GROUPS"
+
+unsafe def functionGroupsDisabledUnsafe (_ : Unit) : Bool :=
+  match unsafeBaseIO (IO.getEnv noFunctionGroupsEnvVar) with
+  | some v => v != "0" && v != ""
+  | none => false
+
+@[implemented_by functionGroupsDisabledUnsafe]
+opaque functionGroupsDisabledImpl (_ : Unit) : Bool := false
+
+/-- Whether `IX_NO_FUNCTION_GROUPS` disables function grouping in this
+process (see `noFunctionGroupsEnvVar`). -/
+def functionGroupsDisabled : Bool := functionGroupsDisabledImpl ()
+
+/-- `compile`, then apply a function grouping (see
+`CompiledToplevel.groupFunctions`) - unless `IX_NO_FUNCTION_GROUPS` is set
+(see `noFunctionGroupsEnvVar`), in which case the grouping is ignored and
+every constrained function keeps its singleton circuit. -/
+def Source.Toplevel.compileWithGroups (t : Source.Toplevel)
+    (groups : Array (String × Array String)) : Except String CompiledToplevel := do
+  let compiled ← t.compile
+  if functionGroupsDisabled then pure compiled else compiled.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

Ix/Aiur/Compiler/Lower.lean

@@ -633,7 +633,7 @@ def Concrete.Decls.toBytecode (decls : Concrete.Decls) :
         let memSizes := layoutMState.memSizes.foldl (·.insert ·) memSizes
         pure (functions.push function, memSizes, nameMap)
       | _ => pure acc
-  pure (⟨functions, memSizes.toArray⟩, nameMap)
+  pure (⟨functions, memSizes.toArray, #[]⟩, nameMap)
 
 end Aiur
 

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

Ix/Aiur/Statistics.lean

@@ -81,44 +81,37 @@ 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!"<fn {i}>"
-        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
   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)