feat(new-execution): AdapterTraceStep trait and rv32 ALU refactor (#1589)

Co-authored-by: Alexander Golovanov <Golovanov12345@gmail.com>

Author: jonathanpwang

crates/vm/src/arch/execution.rs

@@ -6,14 +6,17 @@ use openvm_instructions::{
 };
 use openvm_stark_backend::{
     interaction::{BusIndex, InteractionBuilder, PermutationCheckBus},
-    p3_field::{FieldAlgebra, PrimeField32},
+    p3_field::{Field, FieldAlgebra, PrimeField32},
 };
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
 use super::{Streams, VmExecutionState};
 use crate::system::{
-    memory::{online::GuestMemory, MemoryController},
+    memory::{
+        online::{GuestMemory, TracingMemory},
+        MemoryController,
+    },
     program::ProgramBus,
 };
 
@@ -82,6 +85,15 @@ pub struct VmStateMut<'a, MEM, CTX> {
     pub ctx: &'a mut CTX,
 }
 
+impl<CTX> VmStateMut<'_, TracingMemory, CTX> {
+    // TODO: store as u32 directly
+    #[inline(always)]
+    pub fn ins_start<F: Field>(&self, from_state: &mut ExecutionState<F>) {
+        from_state.pc = F::from_canonical_u32(*self.pc);
+        from_state.timestamp = F::from_canonical_u32(self.memory.timestamp);
+    }
+}
+
 // TODO: old
 pub trait InstructionExecutor<F> {
     /// Runtime execution of the instruction, if the instruction is owned by the

crates/vm/src/arch/integration_api.rs

@@ -246,27 +246,27 @@ pub trait SingleTraceStep<F, CTX> {
     fn get_opcode_name(&self, opcode: usize) -> String;
 }
 
-pub struct NewVmChipWrapper<F, AIR, C> {
+pub struct NewVmChipWrapper<F, AIR, STEP> {
     pub air: AIR,
-    pub inner: C,
+    pub step: STEP,
     pub trace_buffer: Vec<F>,
     width: usize,
     buffer_idx: usize,
     mem_helper: SharedMemoryHelper<F>,
 }
 
-impl<F, AIR, C> NewVmChipWrapper<F, AIR, C>
+impl<F, AIR, STEP> NewVmChipWrapper<F, AIR, STEP>
 where
     F: Field,
     AIR: BaseAir<F>,
 {
-    pub fn new(air: AIR, inner: C, height: usize, mem_helper: SharedMemoryHelper<F>) -> Self {
+    pub fn new(air: AIR, step: STEP, height: usize, mem_helper: SharedMemoryHelper<F>) -> Self {
         assert!(height == 0 || height.is_power_of_two());
         let width = air.width();
         let trace_buffer = F::zero_vec(height * width);
         Self {
             air,
-            inner,
+            step,
             trace_buffer,
             width,
             buffer_idx: 0,
@@ -275,10 +275,10 @@ where
     }
 }
 
-impl<F, Air, C> InstructionExecutor<F> for NewVmChipWrapper<F, Air, C>
+impl<F, AIR, STEP> InstructionExecutor<F> for NewVmChipWrapper<F, AIR, STEP>
 where
     F: PrimeField32,
-    C: SingleTraceStep<F, ()>, // TODO: CTX?
+    STEP: SingleTraceStep<F, ()>, // TODO: CTX?
 {
     fn execute(
         &mut self,
@@ -305,27 +305,27 @@ where
             self.trace_buffer
                 .get_unchecked_mut(start_idx..self.buffer_idx)
         };
-        self.inner.execute(state, instruction, row_slice)?;
+        self.step.execute(state, instruction, row_slice)?;
         Ok(ExecutionState {
             pc,
             timestamp: memory.memory.timestamp,
         })
     }
 
     fn get_opcode_name(&self, opcode: usize) -> String {
-        self.inner.get_opcode_name(opcode)
+        self.step.get_opcode_name(opcode)
     }
 }
 
 // Note[jpw]: the statement we want is:
 // - `Air` is an `Air<AB>` for all `AB: AirBuilder`s needed by stark-backend
 // which is equivalent to saying it implements AirRef<SC>
 // The where clauses to achieve this statement is unfortunately really verbose.
-impl<SC, AIR, C> Chip<SC> for NewVmChipWrapper<Val<SC>, AIR, C>
+impl<SC, AIR, STEP> Chip<SC> for NewVmChipWrapper<Val<SC>, AIR, STEP>
 where
     SC: StarkGenericConfig,
     Val<SC>: PrimeField32,
-    C: SingleTraceStep<Val<SC>, ()> + Send + Sync,
+    STEP: SingleTraceStep<Val<SC>, ()> + Send + Sync,
     AIR: Clone + AnyRap<SC> + 'static,
 {
     fn air(&self) -> AirRef<SC> {
@@ -346,13 +346,13 @@ where
         self.trace_buffer[..rows_used * self.width]
             .par_chunks_exact_mut(self.width)
             .for_each(|row_slice| {
-                self.inner.fill_trace_row(&mem_helper, row_slice);
+                self.step.fill_trace_row(&mem_helper, row_slice);
             });
         drop(self.mem_helper);
         let trace = RowMajorMatrix::new(self.trace_buffer, self.width);
         // self.inner.finalize(&mut trace, num_records);
 
-        AirProofInput::simple(trace, self.inner.generate_public_values())
+        AirProofInput::simple(trace, self.step.generate_public_values())
     }
 }
 
@@ -371,6 +371,47 @@ where
     }
 }
 
+// TODO[jpw]: switch read,write to store into abstract buffer, then fill_trace_row using buffer
+/// A helper trait for expressing generic state accesses within the implementation of
+/// [SingleTraceStep]. Note that this is only a helper trait when the same interface of state access
+/// is reused or shared by multiple implementations. It is not required to implement this trait if
+/// it is easier to implement the [SingleTraceStep] trait directly without this trait.
+pub trait AdapterTraceStep<F, CTX> {
+    /// Adapter row width
+    const WIDTH: usize;
+    type ReadData;
+    type WriteData;
+    /// The minimal amount of information needed to generate the sub-row of the trace matrix.
+    /// This type has a lifetime so other context, such as references to other chips, can be
+    /// provided.
+    type TraceContext<'a>
+    where
+        Self: 'a;
+
+    fn start(pc: u32, memory: &TracingMemory, adapter_row: &mut [F]);
+
+    fn read(
+        memory: &mut TracingMemory,
+        instruction: &Instruction<F>,
+        adapter_row: &mut [F],
+    ) -> Self::ReadData;
+
+    fn write(
+        memory: &mut TracingMemory,
+        instruction: &Instruction<F>,
+        adapter_row: &mut [F],
+        data: &Self::WriteData,
+    );
+
+    // Note[jpw]: should we reuse TraceSubRowGenerator trait instead?
+    /// Post-execution filling of rest of adapter row.
+    fn fill_trace_row(
+        mem_helper: &MemoryAuxColsFactory<F>,
+        ctx: Self::TraceContext<'_>,
+        adapter_row: &mut [F],
+    );
+}
+
 pub struct VmChipWrapper<F, A: VmAdapterChip<F>, C: VmCoreChip<F, A::Interface>> {
     pub adapter: A,
     pub core: C,
@@ -440,18 +481,18 @@ where
     }
 }
 
-impl<Mem, Ctx, F, A, C> InsExecutorE1<Mem, Ctx, F> for NewVmChipWrapper<F, A, C>
+impl<Mem, Ctx, F, A, S> InsExecutorE1<Mem, Ctx, F> for NewVmChipWrapper<F, A, S>
 where
     Mem: GuestMemory,
     F: PrimeField32,
-    C: InsExecutorE1<Mem, Ctx, F>,
+    S: InsExecutorE1<Mem, Ctx, F>,
 {
     fn execute_e1(
         &mut self,
         state: &mut VmExecutionState<Mem, Ctx>,
         instruction: &Instruction<F>,
     ) -> Result<()> {
-        self.inner.execute_e1(state, instruction)
+        self.step.execute_e1(state, instruction)
     }
 }
 

crates/vm/src/system/memory/offline_checker/columns.rs

@@ -88,6 +88,11 @@ impl<F: PrimeField32> MemoryReadAuxCols<F> {
     pub fn get_base(self) -> MemoryBaseAuxCols<F> {
         self.base
     }
+
+    /// Sets the previous timestamp **without** updating the less than auxiliary columns.
+    pub fn set_prev(&mut self, timestamp: F) {
+        self.base.prev_timestamp = timestamp;
+    }
 }
 
 #[repr(C)]
@@ -114,3 +119,9 @@ impl<T, const N: usize> AsMut<MemoryBaseAuxCols<T>> for MemoryWriteAuxCols<T, N>
         &mut self.base
     }
 }
+
+impl<T> AsMut<MemoryBaseAuxCols<T>> for MemoryReadAuxCols<T> {
+    fn as_mut(&mut self) -> &mut MemoryBaseAuxCols<T> {
+        &mut self.base
+    }
+}