Add types module, move PackedCasmState there

Author: Gilad Chase

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

@@ -18,20 +18,14 @@ use crate::components::memory;
 use crate::input::instructions::VmState;
 use crate::relations::{MemoryRelation, StateRelation, N_MEMORY_ELEMS, STATE_SIZE};
 use crate::utils::prover::decode_opcode;
+use crate::utils::types::PackedCasmState;
 use crate::utils::{Selector, SelectorTrait};
 
 const N_MEMORY_LOOKUPS: usize = 4;
 const N_STATE_LOOKUPS: usize = 2;
 
-// TODO(Ohad): take from prover_types and remove.
-pub struct PackedVmState {
-    pub pc: PackedM31,
-    pub ap: PackedM31,
-    pub fp: PackedM31,
-}
-
 pub struct ClaimGenerator {
-    pub inputs: Vec<PackedVmState>,
+    pub inputs: Vec<PackedCasmState>,
 }
 impl ClaimGenerator {
     pub fn new(mut inputs: Vec<VmState>) -> Self {
@@ -44,7 +38,7 @@ impl ClaimGenerator {
         let inputs = inputs
             .into_iter()
             .array_chunks::<N_LANES>()
-            .map(|chunk| PackedVmState {
+            .map(|chunk| PackedCasmState {
                 pc: PackedM31::from_array(std::array::from_fn(|i| {
                     M31::from_u32_unchecked(chunk[i].pc)
                 })),
@@ -157,7 +151,7 @@ impl InteractionClaimGenerator {
 }
 
 fn write_trace_simd(
-    inputs: &[PackedVmState],
+    inputs: &[PackedCasmState],
     memory_trace_generator: &memory::ClaimGenerator,
 ) -> (
     Vec<CircleEvaluation<SimdBackend, M31, BitReversedOrder>>,
@@ -199,7 +193,7 @@ fn write_trace_simd(
 // | State (3) | flags (5) | offsets (3) | addrs (3) | values (3 * 4) |
 fn write_trace_row(
     trace: &mut [Col<SimdBackend, M31>],
-    input: &PackedVmState,
+    input: &PackedCasmState,
     row_index: usize,
     interaction_claim_generator: &mut InteractionClaimGenerator,
     memory_trace_generator: &memory::ClaimGenerator,

crates/prover/src/components/ret_opcode/prover.rs

@@ -18,17 +18,10 @@ use crate::components::memory;
 use crate::components::ret_opcode::component::RET_N_TRACE_CELLS;
 use crate::input::instructions::VmState;
 use crate::relations::MemoryRelation;
+use crate::utils::types::PackedCasmState;
 
 const N_MEMORY_CALLS: usize = 3;
 
-// TODO(Ohad): take from prover_types and remove.
-#[derive(Debug)]
-pub struct PackedCasmState {
-    pub pc: PackedM31,
-    pub ap: PackedM31,
-    pub fp: PackedM31,
-}
-
 #[derive(Debug)]
 pub struct ClaimGenerator {
     pub inputs: Vec<PackedCasmState>,

crates/prover/src/utils/mod.rs

@@ -1,5 +1,6 @@
 pub mod component;
 pub mod prover;
+pub mod types;
 
 use std::ops::{Add, Mul, Sub};
 

crates/prover/src/utils/types.rs

@@ -0,0 +1,8 @@
+use stwo_prover::core::backend::simd::m31::PackedM31;
+
+#[derive(Debug, Clone)]
+pub struct PackedCasmState {
+    pub pc: PackedM31,
+    pub ap: PackedM31,
+    pub fp: PackedM31,
+}