Move multiplicy column code to utils

stwo_cairo_prover/crates/prover/src/components/add_ap_opcode_is_imm_f_op_1_base_fp_f/prover.rs

@@ -27,7 +27,8 @@ use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
-use crate::components::{memory_address_to_id, memory_id_to_big, pack_values, verify_instruction};
+use crate::components::utils::pack_values;
+use crate::components::{memory_address_to_id, memory_id_to_big, verify_instruction};
 use crate::relations;
 
 pub type InputType = CasmState;

stwo_cairo_prover/crates/prover/src/components/add_ap_opcode_is_imm_f_op_1_base_fp_t/prover.rs

@@ -27,7 +27,8 @@ use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
-use crate::components::{memory_address_to_id, memory_id_to_big, pack_values, verify_instruction};
+use crate::components::utils::pack_values;
+use crate::components::{memory_address_to_id, memory_id_to_big, verify_instruction};
 use crate::relations;
 
 pub type InputType = CasmState;

stwo_cairo_prover/crates/prover/src/components/add_ap_opcode_is_imm_t_op_1_base_fp_f/prover.rs

@@ -27,7 +27,8 @@ use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
-use crate::components::{memory_address_to_id, memory_id_to_big, pack_values, verify_instruction};
+use crate::components::utils::pack_values;
+use crate::components::{memory_address_to_id, memory_id_to_big, verify_instruction};
 use crate::relations;
 
 pub type InputType = CasmState;

stwo_cairo_prover/crates/prover/src/components/add_opcode_is_small_f_is_imm_f/prover.rs

@@ -27,7 +27,8 @@ use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
-use crate::components::{memory_address_to_id, memory_id_to_big, pack_values, verify_instruction};
+use crate::components::utils::pack_values;
+use crate::components::{memory_address_to_id, memory_id_to_big, verify_instruction};
 use crate::relations;
 
 pub type InputType = CasmState;

stwo_cairo_prover/crates/prover/src/components/add_opcode_is_small_f_is_imm_t/prover.rs

@@ -27,7 +27,8 @@ use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
-use crate::components::{memory_address_to_id, memory_id_to_big, pack_values, verify_instruction};
+use crate::components::utils::pack_values;
+use crate::components::{memory_address_to_id, memory_id_to_big, verify_instruction};
 use crate::relations;
 
 pub type InputType = CasmState;

stwo_cairo_prover/crates/prover/src/components/add_opcode_is_small_t_is_imm_f/prover.rs

@@ -27,7 +27,8 @@ use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
-use crate::components::{memory_address_to_id, memory_id_to_big, pack_values, verify_instruction};
+use crate::components::utils::pack_values;
+use crate::components::{memory_address_to_id, memory_id_to_big, verify_instruction};
 use crate::relations;
 
 pub type InputType = CasmState;

stwo_cairo_prover/crates/prover/src/components/add_opcode_is_small_t_is_imm_t/prover.rs

@@ -27,7 +27,8 @@ use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
-use crate::components::{memory_address_to_id, memory_id_to_big, pack_values, verify_instruction};
+use crate::components::utils::pack_values;
+use crate::components::{memory_address_to_id, memory_id_to_big, verify_instruction};
 use crate::relations;
 
 pub type InputType = CasmState;

stwo_cairo_prover/crates/prover/src/components/assert_eq_opcode_is_double_deref_f_is_imm_f/prover.rs

@@ -27,7 +27,8 @@ use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
-use crate::components::{memory_address_to_id, memory_id_to_big, pack_values, verify_instruction};
+use crate::components::utils::pack_values;
+use crate::components::{memory_address_to_id, memory_id_to_big, verify_instruction};
 use crate::relations;
 
 pub type InputType = CasmState;

stwo_cairo_prover/crates/prover/src/components/assert_eq_opcode_is_double_deref_f_is_imm_t/prover.rs

@@ -27,7 +27,8 @@ use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
-use crate::components::{memory_address_to_id, memory_id_to_big, pack_values, verify_instruction};
+use crate::components::utils::pack_values;
+use crate::components::{memory_address_to_id, memory_id_to_big, verify_instruction};
 use crate::relations;
 
 pub type InputType = CasmState;

stwo_cairo_prover/crates/prover/src/components/assert_eq_opcode_is_double_deref_t_is_imm_f/prover.rs

@@ -27,7 +27,8 @@ use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
-use crate::components::{memory_address_to_id, memory_id_to_big, pack_values, verify_instruction};
+use crate::components::utils::pack_values;
+use crate::components::{memory_address_to_id, memory_id_to_big, verify_instruction};
 use crate::relations;
 
 pub type InputType = CasmState;

stwo_cairo_prover/crates/prover/src/components/call_opcode_is_rel_f_op_1_base_fp_f/prover.rs

@@ -27,7 +27,8 @@ use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
-use crate::components::{memory_address_to_id, memory_id_to_big, pack_values, verify_instruction};
+use crate::components::utils::pack_values;
+use crate::components::{memory_address_to_id, memory_id_to_big, verify_instruction};
 use crate::relations;
 
 pub type InputType = CasmState;

stwo_cairo_prover/crates/prover/src/components/call_opcode_is_rel_f_op_1_base_fp_t/prover.rs

@@ -27,7 +27,8 @@ use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
-use crate::components::{memory_address_to_id, memory_id_to_big, pack_values, verify_instruction};
+use crate::components::utils::pack_values;
+use crate::components::{memory_address_to_id, memory_id_to_big, verify_instruction};
 use crate::relations;
 
 pub type InputType = CasmState;

stwo_cairo_prover/crates/prover/src/components/call_opcode_is_rel_t_op_1_base_fp_f/prover.rs

@@ -27,7 +27,8 @@ use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
-use crate::components::{memory_address_to_id, memory_id_to_big, pack_values, verify_instruction};
+use crate::components::utils::pack_values;
+use crate::components::{memory_address_to_id, memory_id_to_big, verify_instruction};
 use crate::relations;
 
 pub type InputType = CasmState;

stwo_cairo_prover/crates/prover/src/components/generic_opcode/prover.rs

@@ -29,7 +29,7 @@ use stwo_prover::core::poly::circle::{CanonicCoset, CircleEvaluation};
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
-use crate::components::pack_values;
+use crate::components::utils::pack_values;
 use crate::relations;
 use crate::components::memory_address_to_id;use crate::components::memory_id_to_big;use crate::components::verify_instruction;use crate::components::range_check_9_9;use crate::components::range_check_19;
 

stwo_cairo_prover/crates/prover/src/components/jnz_opcode_is_taken_f_dst_base_fp_f/prover.rs

@@ -27,7 +27,8 @@ use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
-use crate::components::{memory_address_to_id, memory_id_to_big, pack_values, verify_instruction};
+use crate::components::utils::pack_values;
+use crate::components::{memory_address_to_id, memory_id_to_big, verify_instruction};
 use crate::relations;
 
 pub type InputType = CasmState;

stwo_cairo_prover/crates/prover/src/components/jnz_opcode_is_taken_f_dst_base_fp_t/prover.rs

@@ -27,7 +27,8 @@ use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
-use crate::components::{memory_address_to_id, memory_id_to_big, pack_values, verify_instruction};
+use crate::components::utils::pack_values;
+use crate::components::{memory_address_to_id, memory_id_to_big, verify_instruction};
 use crate::relations;
 
 pub type InputType = CasmState;

stwo_cairo_prover/crates/prover/src/components/jnz_opcode_is_taken_t_dst_base_fp_f/prover.rs

@@ -27,7 +27,8 @@ use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
-use crate::components::{memory_address_to_id, memory_id_to_big, pack_values, verify_instruction};
+use crate::components::utils::pack_values;
+use crate::components::{memory_address_to_id, memory_id_to_big, verify_instruction};
 use crate::relations;
 
 pub type InputType = CasmState;

stwo_cairo_prover/crates/prover/src/components/jnz_opcode_is_taken_t_dst_base_fp_t/prover.rs

@@ -27,7 +27,8 @@ use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
-use crate::components::{memory_address_to_id, memory_id_to_big, pack_values, verify_instruction};
+use crate::components::utils::pack_values;
+use crate::components::{memory_address_to_id, memory_id_to_big, verify_instruction};
 use crate::relations;
 
 pub type InputType = CasmState;

stwo_cairo_prover/crates/prover/src/components/jump_opcode_is_rel_f_is_imm_f_is_double_deref_f/prover.rs

@@ -27,7 +27,8 @@ use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
-use crate::components::{memory_address_to_id, memory_id_to_big, pack_values, verify_instruction};
+use crate::components::utils::pack_values;
+use crate::components::{memory_address_to_id, memory_id_to_big, verify_instruction};
 use crate::relations;
 
 pub type InputType = CasmState;

stwo_cairo_prover/crates/prover/src/components/jump_opcode_is_rel_f_is_imm_f_is_double_deref_t/prover.rs

@@ -27,7 +27,8 @@ use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
-use crate::components::{memory_address_to_id, memory_id_to_big, pack_values, verify_instruction};
+use crate::components::utils::pack_values;
+use crate::components::{memory_address_to_id, memory_id_to_big, verify_instruction};
 use crate::relations;
 
 pub type InputType = CasmState;

stwo_cairo_prover/crates/prover/src/components/jump_opcode_is_rel_t_is_imm_f_is_double_deref_f/prover.rs

@@ -27,7 +27,8 @@ use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
-use crate::components::{memory_address_to_id, memory_id_to_big, pack_values, verify_instruction};
+use crate::components::utils::pack_values;
+use crate::components::{memory_address_to_id, memory_id_to_big, verify_instruction};
 use crate::relations;
 
 pub type InputType = CasmState;

stwo_cairo_prover/crates/prover/src/components/jump_opcode_is_rel_t_is_imm_t_is_double_deref_f/prover.rs

@@ -27,7 +27,8 @@ use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
-use crate::components::{memory_address_to_id, memory_id_to_big, pack_values, verify_instruction};
+use crate::components::utils::pack_values;
+use crate::components::{memory_address_to_id, memory_id_to_big, verify_instruction};
 use crate::relations;
 
 pub type InputType = CasmState;

stwo_cairo_prover/crates/prover/src/components/memory/memory_address_to_id/prover.rs

@@ -19,7 +19,7 @@ use super::component::{Claim, InteractionClaim, N_SPLIT_CHUNKS};
 use crate::components::memory_address_to_id::component::{
     N_ID_AND_MULT_COLUMNS_PER_CHUNK, N_TRACE_COLUMNS,
 };
-use crate::components::AtomicMultiplicityColumn;
+use crate::components::utils::AtomicMultiplicityColumn;
 use crate::input::memory::Memory;
 use crate::relations;
 

stwo_cairo_prover/crates/prover/src/components/memory/memory_id_to_big/prover.rs

@@ -24,7 +24,7 @@ use super::component::{
 };
 use crate::components::memory::MEMORY_ADDRESS_BOUND;
 use crate::components::range_check_vector::range_check_9_9;
-use crate::components::AtomicMultiplicityColumn;
+use crate::components::utils::AtomicMultiplicityColumn;
 use crate::felt::split_f252_simd;
 use crate::input::memory::{
     u128_to_4_limbs, EncodedMemoryValueId, Memory, MemoryValueId, LARGE_MEMORY_VALUE_ID_TAG,

stwo_cairo_prover/crates/prover/src/components/mod.rs

@@ -1,11 +1,3 @@
-use std::mem::transmute;
-use std::simd::Simd;
-use std::sync::atomic::{AtomicU32, Ordering};
-
-use stwo_prover::core::backend::simd::column::BaseColumn;
-use stwo_prover::core::backend::simd::conversion::Pack;
-use stwo_prover::core::backend::simd::m31::{PackedM31, N_LANES};
-
 pub mod add_ap_opcode_is_imm_f_op_1_base_fp_f;
 pub mod add_ap_opcode_is_imm_f_op_1_base_fp_t;
 pub mod add_ap_opcode_is_imm_t_op_1_base_fp_f;
@@ -31,6 +23,7 @@ pub mod jump_opcode_is_rel_t_is_imm_t_is_double_deref_f;
 pub mod memory;
 pub mod range_check_vector;
 pub mod ret_opcode;
+pub mod utils;
 pub mod verify_instruction;
 
 // TODO(Ohad): mul small.
@@ -39,106 +32,3 @@ pub mod mul_opcode_is_small_f_is_imm_t;
 
 pub use memory::{memory_address_to_id, memory_id_to_big};
 pub use range_check_vector::{range_check_19, range_check_4_3, range_check_7_2_5, range_check_9_9};
-use stwo_prover::core::fields::m31::M31;
-
-// When padding is needed, the inputs must be arranged in the order defined by the neighbor
-// function. This order allows using the partial sum mechanism to sum only the first n_call inputs.
-// After getting the `SubComponentInputs` we apply the permutation again to ignore padded values at
-// the tail of the vector.
-// TODO(Ohad): generalize the padding logic, and move above doc to the relevant function.
-
-pub fn pack_values<T: Pack>(values: &[T]) -> Vec<T::SimdType> {
-    values
-        .array_chunks::<N_LANES>()
-        .map(|c| T::pack(*c))
-        .collect()
-}
-
-// TODO(Gali): Move to stwo-air-utils.
-/// A column of multiplicities for lookup arguments. Allows increasing the multiplicity at a given
-/// index.
-pub struct MultiplicityColumn {
-    data: Vec<Simd<u32, N_LANES>>,
-}
-impl MultiplicityColumn {
-    /// Creates a new `MultiplicityColumn` with the given size. The elements are initialized to 0.
-    pub fn new(size: usize) -> Self {
-        let vec_size = size.div_ceil(N_LANES);
-        Self {
-            data: vec![unsafe { std::mem::zeroed() }; vec_size],
-        }
-    }
-
-    pub fn increase_at(&mut self, address: u32) {
-        self.data[address as usize / N_LANES][address as usize % N_LANES] += 1;
-    }
-
-    /// Returns the internal data as a Vec<PackedM31>. The last element of the vector is padded with
-    /// zeros if needed.
-    pub fn into_simd_vec(self) -> Vec<PackedM31> {
-        // Safe because the data is aligned to the size of PackedM31 and the size of the data is a
-        // multiple of N_LANES.
-        unsafe { transmute(self.data) }
-    }
-}
-
-/// A column of multiplicities for lookup arguments. Allows increasing the multiplicity at a given
-/// index. This version uses atomic operations to increase the multiplicity, and is `Send`.
-pub struct AtomicMultiplicityColumn {
-    data: Vec<AtomicU32>,
-}
-impl AtomicMultiplicityColumn {
-    /// Creates a new `AtomicMultiplicityColumn` with the given size. The elements are initialized
-    /// to 0.
-    pub fn new(size: usize) -> Self {
-        Self {
-            data: (0..size as u32).map(|_| AtomicU32::new(0)).collect(),
-        }
-    }
-
-    pub fn increase_at(&self, address: u32) {
-        self.data[address as usize].fetch_add(1, Ordering::Relaxed);
-    }
-
-    /// Returns the internal data as a Vec<PackedM31>. The last element of the vector is padded with
-    /// zeros if needed. This function performs a copy on the inner data, If atomics are not
-    /// necessary, use [`MultiplicityColumn`] instead.
-    pub fn into_simd_vec(self) -> Vec<PackedM31> {
-        // Safe because the data is aligned to the size of PackedM31 and the size of the data is a
-        // multiple of N_LANES.
-        BaseColumn::from_iter(
-            self.data
-                .into_iter()
-                .map(|a| M31(a.load(Ordering::Relaxed))),
-        )
-        .data
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use num_traits::{One, Zero};
-    use rand::rngs::SmallRng;
-    use rand::{Rng, SeedableRng};
-    use stwo_prover::core::backend::simd::m31::N_LANES;
-    use stwo_prover::core::fields::m31::M31;
-
-    #[test]
-    fn test_multiplicities_column() {
-        let mut rng = SmallRng::seed_from_u64(0u64);
-        let mut multiplicity_column = super::MultiplicityColumn::new(6 * N_LANES - 2);
-        let mut expected = vec![M31::zero(); 6 * N_LANES];
-
-        (0..10 * N_LANES).for_each(|_| {
-            let addr = rng.gen_range(0..N_LANES * 6);
-            multiplicity_column.increase_at(addr as u32);
-            expected[addr] += M31::one();
-        });
-        let res = multiplicity_column.into_simd_vec();
-
-        assert!(res.len() == 6);
-        for (res_chunk, expected_chunk) in res.iter().zip(expected.chunks(N_LANES)) {
-            assert!(res_chunk.to_array() == expected_chunk);
-        }
-    }
-}

stwo_cairo_prover/crates/prover/src/components/mul_opcode_is_small_f_is_imm_f/prover.rs

@@ -27,8 +27,9 @@ use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
+use crate::components::utils::pack_values;
 use crate::components::{
-    memory_address_to_id, memory_id_to_big, pack_values, range_check_19, verify_instruction,
+    memory_address_to_id, memory_id_to_big, range_check_19, verify_instruction,
 };
 use crate::relations;
 

stwo_cairo_prover/crates/prover/src/components/mul_opcode_is_small_f_is_imm_t/prover.rs

@@ -27,8 +27,9 @@ use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::bit_reverse_coset_to_circle_domain_order;
 
 use super::component::{Claim, InteractionClaim};
+use crate::components::utils::pack_values;
 use crate::components::{
-    memory_address_to_id, memory_id_to_big, pack_values, range_check_19, verify_instruction,
+    memory_address_to_id, memory_id_to_big, range_check_19, verify_instruction,
 };
 use crate::relations;
 

stwo_cairo_prover/crates/prover/src/components/range_check_vector/component_prover.rs

@@ -19,7 +19,7 @@ macro_rules! range_check_prover {
             use stwo_prover::core::poly::BitReversedOrder;
             use stwo_prover::core::poly::circle::{CanonicCoset, CircleEvaluation};
 
-            use $crate::components::AtomicMultiplicityColumn;
+            use $crate::components::utils::AtomicMultiplicityColumn;
             use $crate::components::range_check_vector::{generate_partitioned_enumeration,
                                             partition_into_bit_segments, SIMD_ENUMERATION_0};