feat: access adapters (#1614)

closes INT-3839

---------

Co-authored-by: Ayush Shukla <[email protected]>

Author: jonathanpwang

crates/vm/src/arch/execution.rs

@@ -6,7 +6,7 @@ use openvm_instructions::{
 };
 use openvm_stark_backend::{
     interaction::{BusIndex, InteractionBuilder, PermutationCheckBus},
-    p3_field::{Field, FieldAlgebra, PrimeField32},
+    p3_field::{FieldAlgebra, PrimeField32},
 };
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
@@ -86,10 +86,10 @@ pub struct VmStateMut<'a, MEM, CTX> {
     pub ctx: &'a mut CTX,
 }
 
-impl<CTX> VmStateMut<'_, TracingMemory, CTX> {
+impl<F: PrimeField32, CTX> VmStateMut<'_, TracingMemory<F>, CTX> {
     // TODO: store as u32 directly
     #[inline(always)]
-    pub fn ins_start<F: Field>(&self, from_state: &mut ExecutionState<F>) {
+    pub fn ins_start(&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);
     }

crates/vm/src/arch/integration_api.rs

@@ -224,7 +224,7 @@ pub struct AdapterAirContext<T, I: VmAdapterInterface<T>> {
 pub trait SingleTraceStep<F, CTX> {
     fn execute(
         &mut self,
-        state: VmStateMut<TracingMemory, CTX>,
+        state: VmStateMut<TracingMemory<F>, CTX>,
         instruction: &Instruction<F>,
         row_slice: &mut [F],
     ) -> Result<()>;
@@ -391,18 +391,18 @@ pub trait AdapterTraceStep<F, CTX> {
     where
         Self: 'a;
 
-    fn start(pc: u32, memory: &TracingMemory, adapter_row: &mut [F]);
+    fn start(pc: u32, memory: &TracingMemory<F>, adapter_row: &mut [F]);
 
     fn read(
         &self,
-        memory: &mut TracingMemory,
+        memory: &mut TracingMemory<F>,
         instruction: &Instruction<F>,
         adapter_row: &mut [F],
     ) -> Self::ReadData;
 
     fn write(
         &self,
-        memory: &mut TracingMemory,
+        memory: &mut TracingMemory<F>,
         instruction: &Instruction<F>,
         adapter_row: &mut [F],
         data: &Self::WriteData,

crates/vm/src/system/memory/adapter/mod.rs

@@ -1,4 +1,4 @@
-use std::{borrow::BorrowMut, cmp::max, sync::Arc};
+use std::{borrow::BorrowMut, cmp::max, io::Cursor, sync::Arc};
 
 pub use air::*;
 pub use columns::*;
@@ -145,6 +145,39 @@ impl<F> AccessAdapterInventory<F> {
             None
         }
     }
+
+    pub(crate) fn execute_split(
+        &mut self,
+        address: MemoryAddress<u32, u32>,
+        values: &[F],
+        timestamp: u32,
+        row_slice: &mut [F],
+    ) where
+        F: PrimeField32,
+    {
+        let index = get_chip_index(values.len());
+        self.chips[index].execute_split(address, values, timestamp, row_slice);
+    }
+
+    pub(crate) fn execute_merge(
+        &mut self,
+        address: MemoryAddress<u32, u32>,
+        values: &[F],
+        left_timestamp: u32,
+        right_timestamp: u32,
+        row_slice: &mut [F],
+    ) where
+        F: PrimeField32,
+    {
+        let index = get_chip_index(values.len());
+        self.chips[index].execute_merge(
+            address,
+            values,
+            left_timestamp,
+            right_timestamp,
+            row_slice,
+        );
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -173,6 +206,25 @@ pub trait GenericAccessAdapterChipTrait<F> {
     fn generate_trace(self) -> RowMajorMatrix<F>
     where
         F: PrimeField32;
+
+    fn execute_split(
+        &mut self,
+        address: MemoryAddress<u32, u32>,
+        values: &[F],
+        timestamp: u32,
+        row_slice: &mut [F],
+    ) where
+        F: PrimeField32;
+
+    fn execute_merge(
+        &mut self,
+        address: MemoryAddress<u32, u32>,
+        values: &[F],
+        left_timestamp: u32,
+        right_timestamp: u32,
+        row_slice: &mut [F],
+    ) where
+        F: PrimeField32;
 }
 
 #[derive(Chip, ChipUsageGetter)]
@@ -216,6 +268,7 @@ impl<F> GenericAccessAdapterChip<F> {
         }
     }
 }
+
 pub struct AccessAdapterChip<F, const N: usize> {
     air: AccessAdapterAir<N>,
     range_checker: SharedVariableRangeCheckerChip,
@@ -294,6 +347,72 @@ impl<F, const N: usize> GenericAccessAdapterChipTrait<F> for AccessAdapterChip<F
             });
         RowMajorMatrix::new(values, width)
     }
+
+    fn execute_split(
+        &mut self,
+        address: MemoryAddress<u32, u32>,
+        values: &[F],
+        timestamp: u32,
+        row_slice: &mut [F],
+    ) where
+        F: PrimeField32,
+    {
+        let row: &mut AccessAdapterCols<F, N> = row_slice.borrow_mut();
+        row.is_valid = F::ONE;
+        row.is_split = F::ONE;
+        row.address = MemoryAddress::new(
+            F::from_canonical_u32(address.address_space),
+            F::from_canonical_u32(address.pointer),
+        );
+        let timestamp = F::from_canonical_u32(timestamp);
+        row.left_timestamp = timestamp;
+        row.right_timestamp = timestamp;
+        row.is_right_larger = F::ZERO;
+        debug_assert_eq!(
+            values.len(),
+            N,
+            "Input values slice length must match the access adapter type"
+        );
+
+        // SAFETY: `values` slice is asserted to have length N. `row.values` is an array of length
+        // N. Pointers are valid and regions do not overlap because exactly one of them is a
+        // part of the trace.
+        unsafe {
+            std::ptr::copy_nonoverlapping(values.as_ptr(), row.values.as_mut_ptr(), N);
+        }
+    }
+
+    fn execute_merge(
+        &mut self,
+        address: MemoryAddress<u32, u32>,
+        values: &[F],
+        left_timestamp: u32,
+        right_timestamp: u32,
+        row_slice: &mut [F],
+    ) where
+        F: PrimeField32,
+    {
+        let row: &mut AccessAdapterCols<F, N> = row_slice.borrow_mut();
+        row.is_valid = F::ONE;
+        row.is_split = F::ZERO;
+        row.address = MemoryAddress::new(
+            F::from_canonical_u32(address.address_space),
+            F::from_canonical_u32(address.pointer),
+        );
+        row.left_timestamp = F::from_canonical_u32(left_timestamp);
+        row.right_timestamp = F::from_canonical_u32(right_timestamp);
+        debug_assert_eq!(
+            values.len(),
+            N,
+            "Input values slice length must match the access adapter type"
+        );
+        // SAFETY: `values` slice is asserted to have length N. `row.values` is an array of length
+        // N. Pointers are valid and regions do not overlap because exactly one of them is a
+        // part of the trace.
+        unsafe {
+            std::ptr::copy_nonoverlapping(values.as_ptr(), row.values.as_mut_ptr(), N);
+        }
+    }
 }
 
 impl<SC: StarkGenericConfig, const N: usize> Chip<SC> for AccessAdapterChip<Val<SC>, N>
@@ -328,3 +447,42 @@ impl<F, const N: usize> ChipUsageGetter for AccessAdapterChip<F, N> {
 fn air_name(n: usize) -> String {
     format!("AccessAdapter<{}>", n)
 }
+
+#[inline(always)]
+pub fn get_chip_index(block_size: usize) -> usize {
+    assert!(
+        block_size.is_power_of_two() && block_size >= 2,
+        "Invalid block size {} for split operation",
+        block_size
+    );
+    let index = block_size.trailing_zeros() - 1;
+    index as usize
+}
+
+pub struct AdapterInventoryTraceCursor<F> {
+    // [AG] TODO: replace with a pre-allocated space
+    cursors: Vec<Cursor<Vec<F>>>,
+    widths: Vec<usize>,
+}
+
+impl<F: PrimeField32> AdapterInventoryTraceCursor<F> {
+    pub fn new(as_cnt: usize) -> Self {
+        let cursors = vec![Cursor::new(Vec::new()); as_cnt];
+        let widths = vec![
+            size_of::<AccessAdapterCols<u8, 2>>(),
+            size_of::<AccessAdapterCols<u8, 4>>(),
+            size_of::<AccessAdapterCols<u8, 8>>(),
+            size_of::<AccessAdapterCols<u8, 16>>(),
+            size_of::<AccessAdapterCols<u8, 32>>(),
+        ];
+        Self { cursors, widths }
+    }
+
+    pub fn get_row_slice(&mut self, block_size: usize) -> &mut [F] {
+        let index = get_chip_index(block_size);
+        let begin = self.cursors[index].position() as usize;
+        let end = begin + self.widths[index];
+        self.cursors[index].set_position(end as u64);
+        &mut self.cursors[index].get_mut()[begin..end]
+    }
+}

crates/vm/src/system/memory/controller/mod.rs

@@ -101,7 +101,7 @@ pub struct MemoryController<F> {
     // Store separately to avoid smart pointer reference each time
     range_checker_bus: VariableRangeCheckerBus,
     // addr_space -> Memory data structure
-    pub(crate) memory: TracingMemory,
+    pub(crate) memory: TracingMemory<F>,
     /// A reference to the `OfflineMemory`. Will be populated after `finalize()`.
     pub offline_memory: Arc<Mutex<OfflineMemory<F>>>,
     pub access_adapters: AccessAdapterInventory<F>,
@@ -246,7 +246,7 @@ impl<F: PrimeField32> MemoryController<F> {
                     range_checker.clone(),
                 ),
             },
-            memory: TracingMemory::new(&mem_config),
+            memory: TracingMemory::new(&mem_config, range_checker.clone(), memory_bus),
             offline_memory: Arc::new(Mutex::new(OfflineMemory::new(
                 initial_memory,
                 1,
@@ -297,8 +297,8 @@ impl<F: PrimeField32> MemoryController<F> {
             memory_bus,
             mem_config,
             interface_chip,
-            memory: TracingMemory::new(&mem_config), /* it is expected that the memory will be
-                                                      * set later */
+            memory: TracingMemory::new(&mem_config, range_checker.clone(), memory_bus), /* it is expected that the memory will be
+                                                                                         * set later */
             offline_memory: Arc::new(Mutex::new(OfflineMemory::new(
                 AddressMap::from_mem_config(&mem_config),
                 CHUNK,
@@ -355,7 +355,12 @@ impl<F: PrimeField32> MemoryController<F> {
         let mut offline_memory = self.offline_memory.lock().unwrap();
         offline_memory.set_initial_memory(memory.clone(), self.mem_config);
 
-        self.memory = TracingMemory::from_image(memory.clone(), self.mem_config.access_capacity);
+        self.memory = TracingMemory::new(
+            &self.mem_config,
+            self.range_checker.clone(),
+            self.memory_bus,
+        )
+        .with_image(memory.clone(), self.mem_config.access_capacity);
 
         match &mut self.interface_chip {
             MemoryInterface::Volatile { .. } => {