chore: MemoryAuxColsFactory functions to mutate trace rows in place (#1258)

* perf: MemoryAuxColsFactory mutate rows in place

* clean

Author: jonathanpwang

Cargo.lock

@@ -94,7 +94,7 @@ pub trait VmExtension<F: PrimeField32> {
 }
 
 /// SystemPort combines system resources needed by most extensions
-#[derive(Clone)]
+#[derive(Clone, Copy)]
 pub struct SystemPort {
     pub execution_bus: ExecutionBus,
     pub program_bus: ProgramBus,

crates/vm/src/arch/extensions.rs

@@ -714,6 +714,65 @@ pub struct MemoryAuxColsFactory<T> {
 
 // NOTE[jpw]: The `make_*_aux_cols` functions should be thread-safe so they can be used in parallelized trace generation.
 impl<F: PrimeField32> MemoryAuxColsFactory<F> {
+    pub fn generate_read_aux(&self, read: &MemoryRecord<F>, buffer: &mut MemoryReadAuxCols<F>) {
+        assert!(
+            !read.address_space.is_zero(),
+            "cannot make `MemoryReadAuxCols` for address space 0"
+        );
+        buffer.base.prev_timestamp = F::from_canonical_u32(read.prev_timestamp);
+        self.generate_timestamp_lt(
+            read.prev_timestamp,
+            read.timestamp,
+            &mut buffer.base.clk_lt_aux,
+        );
+    }
+
+    pub fn generate_read_or_immediate_aux(
+        &self,
+        read: &MemoryRecord<F>,
+        buffer: &mut MemoryReadOrImmediateAuxCols<F>,
+    ) {
+        IsZeroSubAir.generate_subrow(
+            read.address_space,
+            (&mut buffer.is_zero_aux, &mut buffer.is_immediate),
+        );
+        buffer.base.prev_timestamp = F::from_canonical_u32(read.prev_timestamp);
+        self.generate_timestamp_lt(
+            read.prev_timestamp,
+            read.timestamp,
+            &mut buffer.base.clk_lt_aux,
+        );
+    }
+
+    pub fn generate_write_aux<const N: usize>(
+        &self,
+        write: &MemoryRecord<F>,
+        buffer: &mut MemoryWriteAuxCols<F, N>,
+    ) {
+        let prev_data = write.prev_data.clone().unwrap();
+        buffer.prev_data = prev_data.try_into().unwrap();
+        buffer.base.prev_timestamp = F::from_canonical_u32(write.prev_timestamp);
+        self.generate_timestamp_lt(
+            write.prev_timestamp,
+            write.timestamp,
+            &mut buffer.base.clk_lt_aux,
+        );
+    }
+
+    fn generate_timestamp_lt(
+        &self,
+        prev_timestamp: u32,
+        timestamp: u32,
+        buffer: &mut LessThanAuxCols<F, AUX_LEN>,
+    ) {
+        debug_assert!(prev_timestamp < timestamp);
+        self.timestamp_lt_air.generate_subrow(
+            (self.range_checker.as_ref(), prev_timestamp, timestamp),
+            &mut buffer.lower_decomp,
+        );
+    }
+
+    // TODO[jpw]: delete these functions and use the ones that fill buffers above.
     pub fn make_read_aux_cols(&self, read: &MemoryRecord<F>) -> MemoryReadAuxCols<F> {
         assert!(
             !read.address_space.is_zero(),

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

@@ -14,9 +14,9 @@ use crate::system::memory::offline_checker::bridge::AUX_LEN;
 #[derive(Clone, Copy, Debug, AlignedBorrow)]
 pub struct MemoryBaseAuxCols<T> {
     /// The previous timestamps in which the cells were accessed.
-    pub(super) prev_timestamp: T,
+    pub(crate) prev_timestamp: T,
     /// The auxiliary columns to perform the less than check.
-    pub(super) clk_lt_aux: LessThanAuxCols<T, AUX_LEN>,
+    pub(crate) clk_lt_aux: LessThanAuxCols<T, AUX_LEN>,
 }
 
 #[repr(C)]
@@ -69,7 +69,7 @@ impl<const N: usize, F: FieldAlgebra> MemoryWriteAuxCols<F, N> {
 #[repr(C)]
 #[derive(Clone, Copy, Debug, AlignedBorrow)]
 pub struct MemoryReadAuxCols<T> {
-    pub(super) base: MemoryBaseAuxCols<T>,
+    pub(crate) base: MemoryBaseAuxCols<T>,
 }
 
 impl<F: PrimeField32> MemoryReadAuxCols<F> {
@@ -99,9 +99,9 @@ impl<F: FieldAlgebra + Copy> MemoryReadAuxCols<F> {
 #[repr(C)]
 #[derive(Clone, Debug, AlignedBorrow)]
 pub struct MemoryReadOrImmediateAuxCols<T> {
-    pub(super) base: MemoryBaseAuxCols<T>,
-    pub(super) is_immediate: T,
-    pub(super) is_zero_aux: T,
+    pub(crate) base: MemoryBaseAuxCols<T>,
+    pub(crate) is_immediate: T,
+    pub(crate) is_zero_aux: T,
 }
 
 impl<T> MemoryReadOrImmediateAuxCols<T> {

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

@@ -29,6 +29,7 @@ num-bigint.workspace = true
 num-integer.workspace = true
 serde = { workspace = true, features = ["derive", "std"] }
 serde-big-array.workspace = true
+once_cell.workspace = true
 
 [dev-dependencies]
 openvm-stark-sdk = { workspace = true }