Merge branch 'main' of github.com:argumentcomputer/ix into jcb/ix-tests

Author: johnchandlerburnham

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);