chore: simplify archon (#175)

*Archon flush values are now of `OracleOrConst`, following the Binius
design more closely.
* Unify Archon exponentials by using `OracleOrConst` as the type for
the base.

Author: arthurpaulino

Ix/Aiur/Constraints.lean

@@ -46,7 +46,7 @@ structure Constraints where
   recvs : Array (Channel × ArithExpr × Array ArithExpr)
   requires : Array (Channel × ArithExpr × OracleIdx × Array ArithExpr)
   topmostSelector : ArithExpr
-  io : Array OracleIdx
+  io : Array OracleOrConst
   multiplicity : OracleIdx
 
 def blockSelector (block : Bytecode.Block) (columns : Columns) : ArithExpr :=
@@ -63,7 +63,7 @@ def new (function : Bytecode.Function) (layout : Layout) (columns : Columns) : C
     recvs := #[]
     requires := #[]
     topmostSelector := blockSelector function.body columns
-    io := columns.inputs ++ columns.outputs
+    io := columns.inputs ++ columns.outputs |>.map .oracle
     multiplicity := columns.multiplicity }
 
 @[inline] def pushUnique (constraints : Constraints) (expr : ArithExpr) : Constraints :=

Ix/Aiur/Synthesis.lean

@@ -35,31 +35,26 @@ abbrev SynthM := StateM SynthState
   (stt.circuitModule.freezeOracles, a)
 
 @[inline] def flush (direction : FlushDirection) (channelId : ChannelId)
-    (selector : OracleIdx) (args : @& Array OracleIdx) (multiplicity : UInt64) : SynthM Unit :=
+    (selector : OracleIdx) (args : @& Array OracleOrConst) (multiplicity : UInt64) : SynthM Unit :=
   modify fun stt =>
     let circuitModule := stt.circuitModule.flush direction channelId selector
       args multiplicity
     { stt with circuitModule }
 
-@[inline] def send (channelId : ChannelId) (args : @& Array OracleIdx) : SynthM Unit :=
+@[inline] def send (channelId : ChannelId) (args : @& Array OracleOrConst) : SynthM Unit :=
   flush .push channelId CircuitModule.selector args 1
 
-@[inline] def recv (channelId : ChannelId) (args : @& Array OracleIdx) : SynthM Unit :=
+@[inline] def recv (channelId : ChannelId) (args : @& Array OracleOrConst) : SynthM Unit :=
   flush .pull channelId CircuitModule.selector args 1
 
 @[inline] def assertZero (name : @& String) (expr : @& ArithExpr) : SynthM Unit :=
   modify fun stt =>
     { stt with circuitModule := stt.circuitModule.assertZero name #[] expr }
 
-@[inline] def assertDynamicExp (expBits : @& Array OracleIdx) (result base : OracleIdx) :
+@[inline] def assertExp (expBits : @& Array OracleIdx) (result : OracleIdx) (base : @& OracleOrConst) :
     SynthM Unit :=
   modify fun stt =>
-    { stt with circuitModule := stt.circuitModule.assertDynamicExp expBits result base }
-
-@[inline] def assertStaticExp (expBits : @& Array OracleIdx) (result : OracleIdx)
-    (base : @& UInt128) (baseTF : TowerField): SynthM Unit :=
-  modify fun stt =>
-    { stt with circuitModule := stt.circuitModule.assertStaticExp expBits result base baseTF }
+    { stt with circuitModule := stt.circuitModule.assertExp expBits result base }
 
 def addCommitted (name : @& String) (tf : TowerField) (relativeHeight : RelativeHeight) :
     SynthM OracleIdx :=
@@ -97,15 +92,6 @@ def addProjected (name : @& String) (inner : OracleIdx) (selection : UInt64)
     let (o, circuitModule) := stt.circuitModule.addProjected name inner selection chunkSize
     (o, { stt with circuitModule })
 
-def cacheConst (value : UInt128) : SynthM OracleIdx :=
-  modifyGet fun stt => match stt.cachedOracles.find? (.const value) with
-  | some o => (o, stt)
-  | none =>
-    let (o, circuitModule) := stt.circuitModule.addTransparent
-      s!"cached-const-{value}" (.const value) .base
-    let cachedOracles := stt.cachedOracles.insert (.const value) o
-    (o, ⟨circuitModule, cachedOracles⟩)
-
 def cacheLc (expr : ArithExpr) : SynthM OracleIdx :=
   let key := .lc expr
   modifyGet fun stt => match stt.cachedOracles.find? key with
@@ -127,16 +113,15 @@ where
     | _ => unreachable!
 
 def provide (channelId : ChannelId) (multiplicity : OracleIdx)
-    (args : Array OracleIdx) : SynthM Unit := do
-  let ones ← cacheConst 1
-  send channelId (args.push ones)
-  recv channelId (args.push multiplicity)
+    (args : Array OracleOrConst) : SynthM Unit := do
+  send channelId (args.push (.const 1 .b1))
+  recv channelId (args.push (.oracle multiplicity))
 
-def require (channelId : ChannelId) (prevIdx : OracleIdx) (args : Array OracleIdx)
+def require (channelId : ChannelId) (prevIdx : OracleIdx) (args : Array OracleOrConst)
     (sel : OracleIdx) : SynthM Unit := do
   let idx ← addLinearCombination s!"index-{channelId.toUSize}" 0 #[(prevIdx, B64_MULT_GEN)] .base
-  flush .pull channelId sel (args.push prevIdx) 1
-  flush .push channelId sel (args.push idx) 1
+  flush .pull channelId sel (args.push (.oracle prevIdx)) 1
+  flush .push channelId sel (args.push (.oracle idx)) 1
 
 def synthesizeFunction (funcIdx : FuncIdx) (function : Bytecode.Function)
     (layout : Layout) (aiurChannels : AiurChannels) : SynthM Columns := do
@@ -154,13 +139,15 @@ def synthesizeFunction (funcIdx : FuncIdx) (function : Bytecode.Function)
   constraints.recvs.forM fun (channel, sel, args) => do
     let sel ← cacheLc sel
     let args ← args.mapM cacheLc
+    let args := args.map .oracle
     match channel with
     | .add => flush .pull aiurChannels.add sel args 1
     | .mul => flush .pull aiurChannels.mul sel args 1
     | _ => unreachable!
   constraints.requires.forM fun (channel, sel, prevIdx, values) => do
     let sel ← cacheLc sel
     let values ← values.mapM cacheLc
+    let values := values.map .oracle
     match channel with
     | .func funcIdx =>
       let funcChannel := aiurChannels.func[funcIdx]!
@@ -192,7 +179,7 @@ def synthesizeAdd (channelId : ChannelId) : SynthM AddColumns := do
   let coutProjected ← addProjected "add-cout-projected" cout 63 64
   assertZero "add-sum" $ xin + yin + cin - zout
   assertZero "add-carry" $ (xin + cin) * (yin + cin) + cin - cout
-  send channelId #[xinPacked, yinPacked, zoutPacked, coutProjected]
+  send channelId $ #[xinPacked, yinPacked, zoutPacked, coutProjected].map .oracle
   pure { xin, yin, zout, cout }
 
 structure MulColumns where
@@ -224,7 +211,7 @@ def synthesizeMul (channelId : ChannelId) : SynthM MulColumns := do
     ← mul xinBits yinBits
   let zoutLow := zoutBits.extract (stop := 64)
   bitDecomposition "mul-bit-decomposition-zout" zoutLow zout
-  send channelId #[xin, yin, zout]
+  send channelId $ #[xin, yin, zout].map .oracle
   pure {
     xin,
     yin,
@@ -262,10 +249,10 @@ where
     let zoutLow := zoutBits.extract (stop := outSize)
     let zoutHigh := zoutBits.extract (start := outSize)
 
-    assertStaticExp xinBits xinExpResult B128_MULT_GEN .b128
-    assertDynamicExp yinBits yinExpResult xinExpResult
-    assertStaticExp zoutLow zoutLowExpResult B128_MULT_GEN .b128
-    assertStaticExp zoutHigh zoutHighExpResult B128GenPow2To64 .b128
+    assertExp xinBits xinExpResult (.const B128_MULT_GEN .b128)
+    assertExp yinBits yinExpResult (.oracle xinExpResult)
+    assertExp zoutLow zoutLowExpResult (.const B128_MULT_GEN .b128)
+    assertExp zoutHigh zoutHighExpResult (.const B128GenPow2To64 .b128)
 
     pure (xinExpResult, yinExpResult, zoutLowExpResult, zoutHighExpResult, zoutBits)
 
@@ -277,7 +264,7 @@ def synthesizeMemory (channelId : ChannelId) (width : Nat) : SynthM MemoryColumn
   let address ← addTransparent s!"mem-{width}-address" .incremental .base
   let values ← Array.range width |>.mapM (addCommitted s!"mem-{width}-value-{·}" .b64 .base)
   let multiplicity ← addCommitted s!"mem-{width}-multiplicity" .b64 .base
-  let args := #[address] ++ values
+  let args := #[address] ++ values |>.map .oracle
   provide channelId multiplicity args
   pure { values, multiplicity }
 

Ix/Archon/Circuit.lean

@@ -2,6 +2,7 @@ import Blake3
 import Ix.Archon.ArithExpr
 import Ix.Archon.OracleIdx
 import Ix.Archon.RelativeHeight
+import Ix.Archon.OracleOrConst
 import Ix.Archon.Transparent
 import Ix.Archon.Witness
 import Ix.Binius.Common
@@ -32,7 +33,7 @@ opaque initWitnessModule : @& CircuitModule → WitnessModule
 /-- **Invalidates** the input `CircuitModule` -/
 @[never_extract, extern "c_rs_circuit_module_flush"]
 opaque flush : CircuitModule → Binius.FlushDirection → Binius.ChannelId →
-  (selector : OracleIdx) → @& Array OracleIdx → (multiplicity : UInt64) → CircuitModule
+  (selector : OracleIdx) → @& Array OracleOrConst → (multiplicity : UInt64) → CircuitModule
 
 /-- **Invalidates** the input `CircuitModule` -/
 @[never_extract, extern "c_rs_circuit_module_assert_zero"]
@@ -44,14 +45,9 @@ opaque assertZero : CircuitModule → @& String → @& Array OracleIdx →
 opaque assertNotZero : CircuitModule → OracleIdx → CircuitModule
 
 /-- **Invalidates** the input `CircuitModule` -/
-@[never_extract, extern "c_rs_circuit_module_assert_dynamic_exp"]
-opaque assertDynamicExp : CircuitModule → (expBits : @& Array OracleIdx) →
-  (result : OracleIdx) → (base : OracleIdx) → CircuitModule
-
-/-- **Invalidates** the input `CircuitModule` -/
-@[never_extract, extern "c_rs_circuit_module_assert_static_exp"]
-opaque assertStaticExp : CircuitModule → (expBits : @& Array OracleIdx) →
-  (result : OracleIdx) → (base : @& UInt128) → (baseTF : TowerField) → CircuitModule
+@[never_extract, extern "c_rs_circuit_module_assert_exp"]
+opaque assertExp : CircuitModule → (expBits : @& Array OracleIdx) →
+  (result : OracleIdx) → (base : @& OracleOrConst) → CircuitModule
 
 /-- **Invalidates** the input `CircuitModule` -/
 @[never_extract, extern "c_rs_circuit_module_add_committed"]

Ix/Archon/OracleOrConst.lean

@@ -0,0 +1,33 @@
+import Ix.Archon.OracleIdx
+import Ix.Archon.TowerField
+import Ix.Unsigned
+
+namespace Archon
+
+inductive OracleOrConst
+  | oracle : OracleIdx → OracleOrConst
+  | const : UInt128 → TowerField → OracleOrConst
+  deriving Inhabited
+
+namespace OracleOrConst
+
+def toString : OracleOrConst → String
+  | oracle o => s!"Oracle({o.toUSize})"
+  | const base tf => s!"Const({base}, {tf})"
+
+instance : ToString OracleOrConst := ⟨toString⟩
+
+def toBytes : @& OracleOrConst → ByteArray
+  | oracle o => ⟨#[0]⟩ ++ o.toUSize.toLEBytes
+  | const base tf => ⟨#[1]⟩ ++ base.toLEBytes |>.push tf.logDegree.toUInt8
+
+/--
+Function meant for testing that tells whether the Lean→Rust mapping of OracleOrConst
+results on the same expression as deserializing the provided bytes.
+-/
+@[extern "rs_oracle_or_const_is_equivalent_to_bytes"]
+opaque isEquivalentToBytes : @& OracleOrConst → @& ByteArray → Bool
+
+end OracleOrConst
+
+end Archon

Ix/Archon/TowerField.lean

@@ -4,6 +4,7 @@ namespace Archon
 
 inductive TowerField
   | b1 | b2 | b4 | b8 | b16 | b32 | b64 | b128
+  deriving Inhabited
 
 def TowerField.logDegree : TowerField → USize
   | .b1 => 0

Ix/Binius/Common.lean

@@ -6,7 +6,6 @@ structure ChannelId where
   toUSize : USize
   deriving Inhabited
 
--- We can delete this later if we don't need it
 structure ChannelAllocator where
   nextId : USize
 

Tests/Aiur.lean

@@ -145,12 +145,11 @@ where
     let (x, circuitModule) := circuitModule.addPacked "x-bits-packed" xBits 6
     let (y, circuitModule) := circuitModule.addPacked "y-bits-packed" yBits 6
     let (z, circuitModule) := circuitModule.addPacked "z-bits-packed" zBits 6
-    let args := #[x, y, z]
-    let (ones, circuitModule) := circuitModule.addTransparent "ones" (.const 1) .base
+    let args := #[x, y, z].map .oracle
     let circuitModule := circuitModule.flush .push channelId CircuitModule.selector
-      (args.push ones) 1
+      (args.push (.const 1 .b1)) 1
     let circuitModule := circuitModule.flush .pull channelId CircuitModule.selector
-      (args.push multiplicity) 1
+      (args.push (.oracle multiplicity)) 1
     (circuitModule.popNamespace, #[xBits, yBits, multiplicity])
   populate entries oracles witnessModule :=
     if entries.isEmpty then (witnessModule, .inactive) else

Tests/FFIConsistency.lean

@@ -2,6 +2,7 @@ import LSpec
 import Tests.Common
 import Ix.Binius.Boundary
 import Ix.Archon.ArithExpr
+import Ix.Archon.OracleOrConst
 import Ix.Archon.ModuleMode
 import Ix.Archon.RelativeHeight
 import Ix.Archon.Transparent
@@ -88,6 +89,28 @@ instance : Repr ModuleMode where
 
 instance : SampleableExt ModuleMode := SampleableExt.mkSelfContained genModuleMode
 
+/- OracleOrConst -/
+
+def genOracleIdx : Gen OracleIdx :=
+  OracleIdx.mk <$> genUSize
+
+def genTowerField : Gen TowerField :=
+  elements #[.b1, .b2, .b4, .b8, .b16, .b32, .b64, .b128]
+
+def genOracleOrConst : Gen OracleOrConst :=
+  frequency [
+      (5, .oracle <$> genOracleIdx),
+      (10, .const <$> genUInt128 <*> genTowerField),
+    ]
+
+instance : Shrinkable OracleOrConst where
+  shrink _ := []
+
+instance : Repr OracleOrConst where
+  reprPrec oc _ := oc.toString
+
+instance : SampleableExt OracleOrConst := SampleableExt.mkSelfContained genOracleOrConst
+
 /- RelativeHeight -/
 
 def genRelativeHeight : Gen RelativeHeight :=
@@ -132,6 +155,8 @@ def Tests.FFIConsistency.suite := [
       (∀ boundary : Boundary, boundary.isEquivalentToBytes boundary.toBytes),
     check "ModuleMode Lean->Rust mapping matches the deserialized bytes"
       (∀ moduleMode : ModuleMode, moduleMode.isEquivalentToBytes moduleMode.toBytes),
+    check "OracleOrConst Lean->Rust mapping matches the deserialized bytes"
+      (∀ oc : OracleOrConst, oc.isEquivalentToBytes oc.toBytes),
     check "RelativeHeight Lean->Rust mapping matches the deserialized bytes"
       (∀ relativeHeight : RelativeHeight, relativeHeight.isEquivalentToBytes relativeHeight.toBytes),
     check "Transparent Lean->Rust mapping matches the deserialized bytes"

c/archon.c

@@ -216,7 +216,7 @@ extern lean_obj_res c_rs_circuit_module_flush(
     bool direction_pull,
     size_t channel_id,
     size_t selector,
-    b_lean_obj_arg oracle_idxs,
+    b_lean_obj_arg values,
     uint64_t multiplicity
 ) {
     linear_object *linear = validated_linear(l_circuit);
@@ -225,7 +225,7 @@ extern lean_obj_res c_rs_circuit_module_flush(
         direction_pull,
         channel_id,
         selector,
-        oracle_idxs,
+        values,
         multiplicity
     );
     linear_object *new_linear = linear_bump(linear);
@@ -260,14 +260,14 @@ extern lean_obj_res c_rs_circuit_module_assert_not_zero(
     return alloc_lean_linear_object(new_linear);
 }
 
-extern lean_obj_res c_rs_circuit_module_assert_dynamic_exp(
+extern lean_obj_res c_rs_circuit_module_assert_exp(
     lean_obj_arg l_circuit,
     b_lean_obj_arg exp_bits,
     size_t result,
-    size_t base
+    b_lean_obj_arg base
 ) {
     linear_object *linear = validated_linear(l_circuit);
-    rs_circuit_module_assert_dynamic_exp(
+    rs_circuit_module_assert_exp(
         get_object_ref(linear),
         exp_bits,
         result,
@@ -277,25 +277,6 @@ extern lean_obj_res c_rs_circuit_module_assert_dynamic_exp(
     return alloc_lean_linear_object(new_linear);
 }
 
-extern lean_obj_res c_rs_circuit_module_assert_static_exp(
-    lean_obj_arg l_circuit,
-    b_lean_obj_arg exp_bits,
-    size_t result,
-    b_lean_obj_arg base,
-    uint8_t base_tower_level
-) {
-    linear_object *linear = validated_linear(l_circuit);
-    rs_circuit_module_assert_static_exp(
-        get_object_ref(linear),
-        exp_bits,
-        result,
-        lean_get_external_data(base),
-        base_tower_level
-    );
-    linear_object *new_linear = linear_bump(linear);
-    return alloc_lean_linear_object(new_linear);
-}
-
 extern lean_obj_res c_rs_circuit_module_add_committed(
     lean_obj_arg l_circuit,
     b_lean_obj_arg name,

c/rust.h

@@ -37,8 +37,7 @@ void rs_circuit_module_assert_zero(
     void*, char const*, b_lean_obj_arg, b_lean_obj_arg
 );
 void rs_circuit_module_assert_not_zero(void*, size_t);
-void rs_circuit_module_assert_dynamic_exp(void*, b_lean_obj_arg, size_t, size_t);
-void rs_circuit_module_assert_static_exp(void*, b_lean_obj_arg, size_t, uint8_t*, uint8_t);
+void rs_circuit_module_assert_exp(void*, b_lean_obj_arg, size_t, b_lean_obj_arg);
 size_t rs_circuit_module_add_committed(void*, char const *, uint8_t, b_lean_obj_arg);
 size_t rs_circuit_module_add_transparent(void*, char const *, b_lean_obj_arg, b_lean_obj_arg);
 size_t rs_circuit_module_add_linear_combination(