refactor: use stwo_air_utils for par air generation

- follows current design in stwo-cairo for this opcode.
- inline write_trace_row
- add LookupData that wraps InteractionClaimGenerator fields, which
  also holds n_calls.
- fix edge cases if input is smaller than LOG_N_LANES

Author: Gilad Chase

Cargo.lock

@@ -22,13 +22,18 @@ num-traits = "0.2.17"
 paste = "1.0"
 pretty_assertions = "1.4.1"
 rand = "0.8.5"
+rayon = { version = "1.10.0" }
 ruint = "1.12.3"
 serde = "1.0.207"
 serde_json = "1.0.1"
 sonic-rs = "0.3.10"
 starknet-ff = "0.3.7"
+stwo-air-utils = { git = "https://github.com/starkware-libs/stwo", rev = "bede666c" }
+stwo-air-utils-derive = { git = "https://github.com/starkware-libs/stwo", rev = "bede666c" }
 # TODO(ShaharS): take stwo version from the source repository.
-stwo-prover = { git = "https://github.com/starkware-libs/stwo", rev = "bede666c" }
+stwo-prover = { git = "https://github.com/starkware-libs/stwo", rev = "bede666c", features = [
+    "parallel",
+] }
 thiserror = "1.0.63"
 tracing = "0.1.40"
 tracing-subscriber = "0.3.18"

Cargo.toml

@@ -10,8 +10,11 @@ cairo-vm.workspace = true
 hex.workspace = true
 itertools.workspace = true
 num-traits.workspace = true
+rayon.workspace = true
 serde.workspace = true
 sonic-rs.workspace = true
+stwo-air-utils.workspace = true
+stwo-air-utils-derive.workspace = true
 stwo-prover.workspace = true
 thiserror.workspace = true
 tracing.workspace = true

crates/prover/Cargo.toml

@@ -39,7 +39,7 @@ impl Eval {
 
 impl FrameworkEval for Eval {
     fn log_size(&self) -> u32 {
-        std::cmp::max(self.claim.n_calls.next_power_of_two().ilog2(), LOG_N_LANES)
+        std::cmp::max(self.claim.n_rows.next_power_of_two().ilog2(), LOG_N_LANES)
     }
 
     fn max_constraint_log_degree_bound(&self) -> u32 {
@@ -153,15 +153,15 @@ impl FrameworkEval for Eval {
 
 #[derive(Clone, Serialize, Deserialize)]
 pub struct Claim {
-    pub n_calls: usize,
+    pub n_rows: usize,
 }
 impl Claim {
     pub fn mix_into(&self, channel: &mut impl Channel) {
-        channel.mix_u64(self.n_calls as u64);
+        channel.mix_u64(self.n_rows as u64);
     }
 
     pub fn log_sizes(&self) -> TreeVec<Vec<u32>> {
-        let log_size = self.n_calls.next_power_of_two().ilog2();
+        let log_size = self.n_rows.next_power_of_two().ilog2();
         let preprocessed_log_sizes = vec![log_size];
         let interaction_1_log_sizes = vec![log_size; N_TRACE_COLUMNS];
         let interaction_2_log_sizes = vec![log_size; SECURE_EXTENSION_DEGREE * 3];

crates/prover/src/components/add_mul_opcode/component.rs

@@ -71,21 +71,21 @@ impl ClaimGenerator {
         });
         tree_builder.extend_evals(trace);
         let claim = Claim {
-            n_calls: self.inputs.len() * N_LANES,
+            n_rows: self.inputs.len() * N_LANES,
         };
         (claim, interaction_claim_generator)
     }
 }
 
 pub struct InteractionClaimGenerator {
-    pub n_calls: usize,
+    pub n_rows: usize,
     pub memory: [Vec<[PackedM31; N_MEMORY_ELEMS]>; N_MEMORY_LOOKUPS],
     pub state: [Vec<[PackedM31; STATE_SIZE]>; N_STATE_LOOKUPS],
 }
 impl InteractionClaimGenerator {
     pub fn with_capacity(capacity: usize) -> Self {
         Self {
-            n_calls: capacity,
+            n_rows: capacity,
             memory: [
                 Vec::with_capacity(capacity),
                 Vec::with_capacity(capacity),
@@ -138,13 +138,13 @@ impl InteractionClaimGenerator {
         }
         col2.finalize_col();
 
-        let (trace, total_sum, claimed_sum) = if self.n_calls == 1 << log_size {
+        let (trace, total_sum, claimed_sum) = if self.n_rows == 1 << log_size {
             let (trace, claimed_sum) = logup_gen.finalize_last();
             (trace, claimed_sum, None)
         } else {
             let (trace, [total_sum, claimed_sum]) =
-                logup_gen.finalize_at([(1 << log_size) - 1, self.n_calls - 1]);
-            (trace, total_sum, Some((claimed_sum, self.n_calls - 1)))
+                logup_gen.finalize_at([(1 << log_size) - 1, self.n_rows - 1]);
+            (trace, total_sum, Some((claimed_sum, self.n_rows - 1)))
         };
 
         tree_builder.extend_evals(trace);

crates/prover/src/components/add_mul_opcode/prover.rs

@@ -7,6 +7,7 @@ pub const LOOKUP_INTERACTION_PHASE: usize = 1;
 
 #[cfg(test)]
 mod tests {
+    use itertools::Itertools;
     use pretty_assertions::assert_eq;
     use stwo_prover::core::backend::simd::SimdBackend;
     use stwo_prover::core::fields::m31::BaseField;
@@ -62,8 +63,11 @@ mod tests {
             ret_claim_generator.write_trace(&mut tree_builder, &mut memory_claim_generator);
 
         let output = ret_interaction_prover
+            .lookup_data
             .memory_outputs
-            .map(|output| output[0].to_array()[0].to_m31_array()[0]);
+            .into_iter()
+            .map(|output| output.map(|x| x.to_array()[0].to_m31_array()[0]))
+            .collect_vec()[0];
         assert_eq!(output, *expected_output);
     }
 }