chore: make Archon witness population more flexible (#127)

Accept pushes of arbitrary numbers of uints to entries of witness
modules instead of always requiring 128 bits of data. The data is
accumulated on buffers and gets concretized into underliers in
chunks of 128 bits (the size of an `OptimalUnderlier`).

This API is propagated to Lean via unsafe bindings.

Author: arthurpaulino

Ix/Archon/Witness.lean

@@ -33,8 +33,24 @@ opaque addEntryWithCapacity : WitnessModule → (logBits : UInt8) → EntryId ×
 opaque bindOracleTo : WitnessModule → OracleIdx → EntryId → TowerField → WitnessModule
 
 /-- **Invalidates** the input `WitnessModule` -/
-@[never_extract, extern "c_rs_witness_module_push_u128_to"]
-opaque pushUInt128To : WitnessModule → @& UInt128 → EntryId → WitnessModule
+@[never_extract, extern "c_rs_witness_module_push_u8s_to"]
+opaque pushUInt8sTo : WitnessModule → @& Array UInt8 → EntryId → WitnessModule
+
+/-- **Invalidates** the input `WitnessModule` -/
+@[never_extract, extern "c_rs_witness_module_push_u16s_to"]
+opaque pushUInt16sTo : WitnessModule → @& Array UInt16 → EntryId → WitnessModule
+
+/-- **Invalidates** the input `WitnessModule` -/
+@[never_extract, extern "c_rs_witness_module_push_u32s_to"]
+opaque pushUInt32sTo : WitnessModule → @& Array UInt32 → EntryId → WitnessModule
+
+/-- **Invalidates** the input `WitnessModule` -/
+@[never_extract, extern "c_rs_witness_module_push_u64s_to"]
+opaque pushUInt64sTo : WitnessModule → @& Array UInt64 → EntryId → WitnessModule
+
+/-- **Invalidates** the input `WitnessModule` -/
+@[never_extract, extern "c_rs_witness_module_push_u128s_to"]
+opaque pushUInt128sTo : WitnessModule → @& Array UInt128 → EntryId → WitnessModule
 
 /-- **Invalidates** the input `WitnessModule` -/
 @[never_extract, extern "c_rs_witness_module_populate"]

Tests/Archon.lean

@@ -45,7 +45,7 @@ def linearCombination : TestSeq := Id.run do
     witnessModule := newWitnessModule
     for j in [0 : 1 <<< i ] do
       let u128 := UInt128.ofNatWrap (j * j + 17)
-      witnessModule := witnessModule.pushUInt128To u128 entryId
+      witnessModule := witnessModule.pushUInt128sTo #[u128] entryId
     match i with
     | 0 => witnessModule := witnessModule.bindOracleTo oracleId entryId .b1
     | 1 => witnessModule := witnessModule.bindOracleTo oracleId entryId .b2
@@ -66,7 +66,7 @@ def packed : TestSeq :=
   let circuitModule := circuitModule.freezeOracles
   let witnessModule := circuitModule.initWitnessModule
   let (entryId, witnessModule) := witnessModule.addEntryWithCapacity 7
-  let witnessModule := witnessModule.pushUInt128To 2347928368726 entryId
+  let witnessModule := witnessModule.pushUInt128sTo #[2347928368726] entryId
   let witnessModule := witnessModule.bindOracleTo x entryId .b1
   populateAndValidate circuitModule witnessModule 128 "Archon packed works"
 
@@ -77,10 +77,40 @@ def shifted : TestSeq :=
   let circuitModule := circuitModule.freezeOracles
   let witnessModule := circuitModule.initWitnessModule
   let (entryId, witnessModule) := witnessModule.addEntryWithCapacity 7
-  let witnessModule := witnessModule.pushUInt128To 2347928368726 entryId
+  let witnessModule := witnessModule.pushUInt128sTo #[2347928368726] entryId
   let witnessModule := witnessModule.bindOracleTo x entryId .b1
   populateAndValidate circuitModule witnessModule 128 "Archon shifted works"
 
+def pushData : TestSeq :=
+  let circuitModule := CircuitModule.new 0
+  let (u8s, circuitModule) := circuitModule.addCommitted "u8s" .b1
+  let (u16s, circuitModule) := circuitModule.addCommitted "u16s" .b1
+  let (u32s, circuitModule) := circuitModule.addCommitted "u32s" .b1
+  let (u64s, circuitModule) := circuitModule.addCommitted "u64s" .b1
+  let (u128s, circuitModule) := circuitModule.addCommitted "u128s" .b1
+  let circuitModule := circuitModule.assertZero "u8s xor u16s" #[] $ .add (.oracle u8s) (.oracle u16s)
+  let circuitModule := circuitModule.assertZero "u16s xor u32s" #[] $ .add (.oracle u16s) (.oracle u32s)
+  let circuitModule := circuitModule.assertZero "u32s xor u64s" #[] $ .add (.oracle u32s) (.oracle u64s)
+  let circuitModule := circuitModule.assertZero "u64s xor u128s" #[] $ .add (.oracle u64s) (.oracle u128s)
+  let circuitModule := circuitModule.freezeOracles
+  let witnessModule := circuitModule.initWitnessModule
+  let (u8sEntry, witnessModule) := witnessModule.addEntryWithCapacity 7
+  let (u16sEntry, witnessModule) := witnessModule.addEntryWithCapacity 7
+  let (u32sEntry, witnessModule) := witnessModule.addEntryWithCapacity 7
+  let (u64sEntry, witnessModule) := witnessModule.addEntryWithCapacity 7
+  let (u128sEntry, witnessModule) := witnessModule.addEntryWithCapacity 7
+  let witnessModule := witnessModule.pushUInt8sTo ⟨(List.replicate 16 0)⟩ u8sEntry
+  let witnessModule := witnessModule.pushUInt16sTo ⟨(List.replicate 8 0)⟩ u16sEntry
+  let witnessModule := witnessModule.pushUInt32sTo #[0, 0, 0, 0] u32sEntry
+  let witnessModule := witnessModule.pushUInt64sTo #[0, 0] u64sEntry
+  let witnessModule := witnessModule.pushUInt128sTo #[0] u128sEntry
+  let witnessModule := witnessModule.bindOracleTo u8s u8sEntry .b1
+  let witnessModule := witnessModule.bindOracleTo u16s u16sEntry .b1
+  let witnessModule := witnessModule.bindOracleTo u32s u32sEntry .b1
+  let witnessModule := witnessModule.bindOracleTo u64s u64sEntry .b1
+  let witnessModule := witnessModule.bindOracleTo u128s u128sEntry .b1
+  populateAndValidate circuitModule witnessModule 128 "Archon witness data pushes work"
+
 def proveAndVerify : TestSeq :=
   let circuitModule := CircuitModule.new 0
   let (x, circuitModule) := circuitModule.addCommitted "x" .b1
@@ -89,7 +119,7 @@ def proveAndVerify : TestSeq :=
   let circuitModule := circuitModule.freezeOracles
   let witnessModule := circuitModule.initWitnessModule
   let (entryId, witnessModule) := witnessModule.addEntryWithCapacity 7
-  let witnessModule := witnessModule.pushUInt128To (.ofNatCore $ UInt128.size - 1) entryId
+  let witnessModule := witnessModule.pushUInt128sTo #[(.ofNatCore $ UInt128.size - 1)] entryId
   let witnessModule := witnessModule.bindOracleTo x entryId .b1
   let witnessModule := witnessModule.bindOracleTo y entryId .b1
   let height := 128
@@ -115,6 +145,7 @@ def Tests.Archon.suite := [
     linearCombination,
     packed,
     shifted,
+    pushData,
     proveAndVerify,
     versionCircuitModules,
   ]

c/archon.c

@@ -47,21 +47,83 @@ extern lean_obj_res c_rs_witness_module_bind_oracle_to(
     return alloc_lean_linear_object(new_linear);
 }
 
-extern lean_obj_res c_rs_witness_module_push_u128_to(
+static inline lean_obj_res c_rs_witness_module_push_data_to(
     lean_obj_arg l_witness,
-    b_lean_obj_arg u128,
-    size_t entry_id
+    b_lean_obj_arg data,
+    size_t entry_id,
+    void (*rs_fn)(void*, b_lean_obj_arg, size_t)
 ) {
     linear_object *linear = validated_linear(l_witness);
-    rs_witness_module_push_u128_to(
-        get_object_ref(linear),
-        lean_get_external_data(u128),
-        entry_id
-    );
+    rs_fn(get_object_ref(linear), data, entry_id);
     linear_object *new_linear = linear_bump(linear);
     return alloc_lean_linear_object(new_linear);
 }
 
+extern lean_obj_res c_rs_witness_module_push_u8s_to(
+    lean_obj_arg l_witness,
+    b_lean_obj_arg u8s,
+    size_t entry_id
+) {
+    return c_rs_witness_module_push_data_to(
+        l_witness,
+        u8s,
+        entry_id,
+        rs_witness_module_push_u8s_to
+    );
+}
+
+extern lean_obj_res c_rs_witness_module_push_u16s_to(
+    lean_obj_arg l_witness,
+    b_lean_obj_arg u16s,
+    size_t entry_id
+) {
+    return c_rs_witness_module_push_data_to(
+        l_witness,
+        u16s,
+        entry_id,
+        rs_witness_module_push_u16s_to
+    );
+}
+
+extern lean_obj_res c_rs_witness_module_push_u32s_to(
+    lean_obj_arg l_witness,
+    b_lean_obj_arg u32s,
+    size_t entry_id
+) {
+    return c_rs_witness_module_push_data_to(
+        l_witness,
+        u32s,
+        entry_id,
+        rs_witness_module_push_u32s_to
+    );
+}
+
+extern lean_obj_res c_rs_witness_module_push_u64s_to(
+    lean_obj_arg l_witness,
+    b_lean_obj_arg u64s,
+    size_t entry_id
+) {
+    return c_rs_witness_module_push_data_to(
+        l_witness,
+        u64s,
+        entry_id,
+        rs_witness_module_push_u64s_to
+    );
+}
+
+extern lean_obj_res c_rs_witness_module_push_u128s_to(
+    lean_obj_arg l_witness,
+    b_lean_obj_arg u128s,
+    size_t entry_id
+) {
+    return c_rs_witness_module_push_data_to(
+        l_witness,
+        u128s,
+        entry_id,
+        rs_witness_module_push_u128s_to
+    );
+}
+
 extern lean_obj_res c_rs_witness_module_populate(lean_obj_arg l_witness, uint64_t height) {
     linear_object *linear = validated_linear(l_witness);
     rs_witness_module_populate(get_object_ref(linear), height);

c/rust.h

@@ -13,7 +13,11 @@ void rs_witness_module_free(void*);
 size_t rs_witness_module_add_entry(void*);
 size_t rs_witness_module_add_entry_with_capacity(void*, uint8_t);
 void rs_witness_module_bind_oracle_to(void*, size_t, size_t, uint8_t);
-void rs_witness_module_push_u128_to(void*, void*, size_t);
+void rs_witness_module_push_u8s_to(void*, b_lean_obj_arg, size_t);
+void rs_witness_module_push_u16s_to(void*, b_lean_obj_arg, size_t);
+void rs_witness_module_push_u32s_to(void*, b_lean_obj_arg, size_t);
+void rs_witness_module_push_u64s_to(void*, b_lean_obj_arg, size_t);
+void rs_witness_module_push_u128s_to(void*, b_lean_obj_arg, size_t);
 void rs_witness_module_populate(void*, uint64_t);
 void *rs_compile_witness_modules(void**, b_lean_obj_arg);
 

src/archon/precompiles/blake3.rs

@@ -137,10 +137,7 @@ fn state_transition_module(
 
     for (xy, trace) in traces.clone().state_trace.into_iter().enumerate() {
         let entry_id_xy = witness_module.new_entry();
-
-        for chunk in trace.chunks(4) {
-            witness_module.push_u32s_to(chunk.try_into()?, entry_id_xy);
-        }
+        witness_module.push_u32s_to(trace, entry_id_xy);
         witness_module.bind_oracle_to::<B32>(state_transitions[xy], entry_id_xy);
     }
 
@@ -299,18 +296,10 @@ fn additions_xor_rotates_module(
         let yin_entry = witness_module.new_entry();
         let zout_entry = witness_module.new_entry();
 
-        for chunk in xin_traces[xy].chunks(4) {
-            witness_module.push_u32s_to(chunk.try_into()?, xin_entry)
-        }
-        for chunk in yin_traces[xy].chunks(4) {
-            witness_module.push_u32s_to(chunk.try_into()?, yin_entry)
-        }
-        for chunk in traces.cout_trace[xy].chunks(4) {
-            witness_module.push_u32s_to(chunk.try_into()?, cout_entry)
-        }
-        for chunk in zout_traces[xy].chunks(4) {
-            witness_module.push_u32s_to(chunk.try_into()?, zout_entry)
-        }
+        witness_module.push_u32s_to(xin_traces[xy].clone(), xin_entry);
+        witness_module.push_u32s_to(yin_traces[xy].clone(), yin_entry);
+        witness_module.push_u32s_to(traces.cout_trace[xy].clone(), cout_entry);
+        witness_module.push_u32s_to(zout_traces[xy].clone(), zout_entry);
 
         witness_module.bind_oracle_to::<B1>(xins[xy], xin_entry);
         witness_module.bind_oracle_to::<B1>(yins[xy], yin_entry);
@@ -320,9 +309,7 @@ fn additions_xor_rotates_module(
 
     // not to forget about d_in
     let d_in_entry = witness_module.new_entry();
-    for chunk in traces.d_in_trace.chunks(4) {
-        witness_module.push_u32s_to(chunk.try_into()?, d_in_entry)
-    }
+    witness_module.push_u32s_to(traces.d_in_trace.clone(), d_in_entry);
     witness_module.bind_oracle_to::<B1>(d_in, d_in_entry);
 
     witness_module.populate(height)?;
@@ -364,17 +351,9 @@ fn cv_output_module(
     let state_i_entry = witness_module.new_entry();
     let state_i_8_entry = witness_module.new_entry();
 
-    for cv_vals in traces.cv_trace.chunks(4) {
-        witness_module.push_u32s_to(cv_vals.try_into().unwrap(), cv_entry);
-    }
-
-    for state_i_vals in traces.state_i_trace.chunks(4) {
-        witness_module.push_u32s_to(state_i_vals.try_into().unwrap(), state_i_entry);
-    }
-
-    for state_i_8_vals in traces.state_i_8_trace.chunks(4) {
-        witness_module.push_u32s_to(state_i_8_vals.try_into().unwrap(), state_i_8_entry);
-    }
+    witness_module.push_u32s_to(traces.cv_trace.clone(), cv_entry);
+    witness_module.push_u32s_to(traces.state_i_trace.clone(), state_i_entry);
+    witness_module.push_u32s_to(traces.state_i_8_trace.clone(), state_i_8_entry);
 
     witness_module.bind_oracle_to::<B32>(cv, cv_entry);
     witness_module.bind_oracle_to::<B32>(state_i, state_i_entry);
@@ -951,13 +930,13 @@ pub mod tests {
         let a_0_entry = witness_modules[witness_module_id].new_entry();
         let d_in_entry = witness_modules[witness_module_id].new_entry();
 
-        for chunk in traces.a_0_trace.chunks(4) {
-            witness_modules[witness_module_id].push_u32s_to(chunk.try_into().unwrap(), a_0_entry)
-        }
-
-        for chunk in traces.d_in_trace.chunks(4) {
-            witness_modules[witness_module_id].push_u32s_to(chunk.try_into().unwrap(), d_in_entry)
-        }
+        let Trace {
+            a_0_trace,
+            d_in_trace,
+            ..
+        } = traces;
+        witness_modules[witness_module_id].push_u32s_to(a_0_trace, a_0_entry);
+        witness_modules[witness_module_id].push_u32s_to(d_in_trace, d_in_entry);
 
         witness_modules[witness_module_id].bind_oracle_to::<B1>(a_0, a_0_entry);
         witness_modules[witness_module_id].bind_oracle_to::<B1>(d_in, d_in_entry);

src/archon/protocol.rs

@@ -192,7 +192,7 @@ mod tests {
 
     fn populate_a_xor_b_witness_with_ones(witness_module: &mut WitnessModule, oracles: &Oracles) {
         let ones = witness_module.new_entry_with_capacity(7);
-        witness_module.push_u128_to(u128::MAX, ones);
+        witness_module.push_u128s_to([u128::MAX], ones);
         witness_module.bind_oracle_to::<B1>(oracles.s, ones);
         witness_module.bind_oracle_to::<B1>(oracles.a, ones);
         witness_module.bind_oracle_to::<B1>(oracles.b, ones);
@@ -204,8 +204,8 @@ mod tests {
         let mut witness_module = circuit_module.init_witness_module().unwrap();
         let zeros = witness_module.new_entry_with_capacity(7);
         let ones = witness_module.new_entry_with_capacity(7);
-        witness_module.push_u128_to(0, zeros);
-        witness_module.push_u128_to(u128::MAX, ones);
+        witness_module.push_u128s_to([0], zeros);
+        witness_module.push_u128s_to([u128::MAX], ones);
         witness_module.bind_oracle_to::<B1>(oracles.s, ones);
         witness_module.bind_oracle_to::<B1>(oracles.a, ones);
         witness_module.bind_oracle_to::<B1>(oracles.b, zeros);