chore: Use aux_cols_factory.generate_*_aux

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

@@ -41,8 +41,8 @@ use crate::{
         merkle::{MemoryMerkleBus, MemoryMerkleChip},
         offline::{MemoryRecord, OfflineMemory, INITIAL_TIMESTAMP},
         offline_checker::{
-            MemoryBridge, MemoryBus, MemoryReadAuxCols, MemoryReadOrImmediateAuxCols,
-            MemoryWriteAuxCols, AUX_LEN,
+            MemoryBaseAuxCols, MemoryBridge, MemoryBus, MemoryReadAuxCols,
+            MemoryReadOrImmediateAuxCols, MemoryWriteAuxCols, AUX_LEN,
         },
         online::{Memory, MemoryLogEntry},
         persistent::PersistentBoundaryChip,
@@ -719,12 +719,7 @@ impl<F: PrimeField32> MemoryAuxColsFactory<F> {
             !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,
-        );
+        self.generate_base_aux(read, &mut buffer.base);
     }
 
     pub fn generate_read_or_immediate_aux(
@@ -736,12 +731,7 @@ impl<F: PrimeField32> MemoryAuxColsFactory<F> {
             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,
-        );
+        self.generate_base_aux(read, &mut buffer.base);
     }
 
     pub fn generate_write_aux<const N: usize>(
@@ -751,11 +741,15 @@ impl<F: PrimeField32> MemoryAuxColsFactory<F> {
     ) {
         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_base_aux(write, &mut buffer.base);
+    }
+
+    pub fn generate_base_aux(&self, record: &MemoryRecord<F>, buffer: &mut MemoryBaseAuxCols<F>) {
+        buffer.prev_timestamp = F::from_canonical_u32(record.prev_timestamp);
         self.generate_timestamp_lt(
-            write.prev_timestamp,
-            write.timestamp,
-            &mut buffer.base.clk_lt_aux,
+            record.prev_timestamp,
+            record.timestamp,
+            &mut buffer.clk_lt_aux,
         );
     }
 

crates/vm/src/system/memory/offline_checker/columns.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(crate) prev_timestamp: T,
+    pub(in crate::system::memory) prev_timestamp: T,
     /// The auxiliary columns to perform the less than check.
-    pub(crate) clk_lt_aux: LessThanAuxCols<T, AUX_LEN>,
+    pub(in crate::system::memory) 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(crate) base: MemoryBaseAuxCols<T>,
+    pub(in crate::system::memory) base: MemoryBaseAuxCols<T>,
 }
 
 impl<F: PrimeField32> MemoryReadAuxCols<F> {

crates/vm/src/system/memory/tests.rs

@@ -168,27 +168,27 @@ fn generate_trace<F: PrimeField32>(
 
         match (record.data.len(), &record.prev_data) {
             (1, &None) => {
-                row.read_1_aux = aux_factory.make_read_aux_cols(&record);
+                aux_factory.generate_read_aux(&record, &mut row.read_1_aux);
                 row.data_1 = record.data.try_into().unwrap();
                 row.is_read_1 = F::ONE;
             }
             (1, &Some(_)) => {
-                row.write_1_aux = aux_factory.make_write_aux_cols(&record);
+                aux_factory.generate_write_aux(&record, &mut row.write_1_aux);
                 row.data_1 = record.data.try_into().unwrap();
                 row.is_write_1 = F::ONE;
             }
             (4, &None) => {
-                row.read_4_aux = aux_factory.make_read_aux_cols(&record);
+                aux_factory.generate_read_aux(&record, &mut row.read_4_aux);
                 row.data_4 = record.data.try_into().unwrap();
                 row.is_read_4 = F::ONE;
             }
             (4, &Some(_)) => {
-                row.write_4_aux = aux_factory.make_write_aux_cols(&record);
+                aux_factory.generate_write_aux(&record, &mut row.write_4_aux);
                 row.data_4 = record.data.try_into().unwrap();
                 row.is_write_4 = F::ONE;
             }
             (MAX, &None) => {
-                row.read_max_aux = aux_factory.make_read_aux_cols(&record);
+                aux_factory.generate_read_aux(&record, &mut row.read_max_aux);
                 row.data_max = record.data.try_into().unwrap();
                 row.is_read_max = F::ONE;
             }

crates/vm/src/system/native_adapter/mod.rs

@@ -279,22 +279,22 @@ impl<F: PrimeField32, const R: usize, const W: usize> VmAdapterChip<F>
 
         row_slice.from_state = write_record.from_state.map(F::from_canonical_u32);
 
-        row_slice.reads_aux = read_record.reads.map(|(id, _)| {
-            let record = memory.record_by_id(id);
-            let address = MemoryAddress::new(record.address_space, record.pointer);
-            NativeAdapterReadCols {
-                address,
-                read_aux: aux_cols_factory.make_read_or_immediate_aux_cols(record),
-            }
-        });
-        row_slice.writes_aux = write_record.writes.map(|(id, _)| {
-            let record = memory.record_by_id(id);
-            let address = MemoryAddress::new(record.address_space, record.pointer);
-            NativeAdapterWriteCols {
-                address,
-                write_aux: aux_cols_factory.make_write_aux_cols(record),
-            }
-        });
+        for (i, read) in read_record.reads.iter().enumerate() {
+            let (id, _) = read;
+            let record = memory.record_by_id(*id);
+            aux_cols_factory
+                .generate_read_or_immediate_aux(record, &mut row_slice.reads_aux[i].read_aux);
+            row_slice.reads_aux[i].address =
+                MemoryAddress::new(record.address_space, record.pointer);
+        }
+
+        for (i, write) in write_record.writes.iter().enumerate() {
+            let (id, _) = write;
+            let record = memory.record_by_id(*id);
+            aux_cols_factory.generate_write_aux(record, &mut row_slice.writes_aux[i].write_aux);
+            row_slice.writes_aux[i].address =
+                MemoryAddress::new(record.address_space, record.pointer);
+        }
     }
 
     fn air(&self) -> &Self::Air {

extensions/keccak256/circuit/src/trace.rs

@@ -219,15 +219,17 @@ where
                         );
                     }
                     for (i, id) in register_reads.into_iter().enumerate() {
-                        // TODO[jpw] make_read_aux_cols should directly write into slice
-                        first_row.mem_oc.register_aux[i] =
-                            aux_cols_factory.make_read_aux_cols(memory.record_by_id(id));
+                        aux_cols_factory.generate_read_aux(
+                            memory.record_by_id(id),
+                            &mut first_row.mem_oc.register_aux[i],
+                        );
                     }
                 }
                 for (i, id) in block.reads.into_iter().enumerate() {
-                    // TODO[jpw] make_read_aux_cols should directly write into slice
-                    first_row.mem_oc.absorb_reads[i] =
-                        aux_cols_factory.make_read_aux_cols(memory.record_by_id(id));
+                    aux_cols_factory.generate_read_aux(
+                        memory.record_by_id(id),
+                        &mut first_row.mem_oc.absorb_reads[i],
+                    );
                 }
 
                 let last_row: &mut KeccakVmCols<Val<SC>> =
@@ -239,8 +241,8 @@ where
                     for (i, record_id) in digest_writes.into_iter().enumerate() {
                         // TODO: these aux columns are only used for the last row - can we share them with aux reads in first row?
                         let record = memory.record_by_id(record_id);
-                        last_row.mem_oc.digest_writes[i] =
-                            aux_cols_factory.make_write_aux_cols(record);
+                        aux_cols_factory
+                            .generate_write_aux(record, &mut last_row.mem_oc.digest_writes[i]);
                     }
                 }
             });

extensions/native/circuit/src/adapters/branch_native_adapter.rs

@@ -187,14 +187,13 @@ impl<F: PrimeField32> VmAdapterChip<F> for BranchNativeAdapterChip<F> {
         let aux_cols_factory = memory.aux_cols_factory();
 
         row_slice.from_state = write_record.map(F::from_canonical_u32);
-        row_slice.reads_aux = read_record.reads.map(|x| {
+        for (i, x) in read_record.reads.iter().enumerate() {
             let read = memory.record_by_id(x.0);
-            let address = MemoryAddress::new(read.address_space, read.pointer);
-            BranchNativeAdapterReadCols {
-                address,
-                read_aux: aux_cols_factory.make_read_or_immediate_aux_cols(read),
-            }
-        });
+
+            row_slice.reads_aux[i].address = MemoryAddress::new(read.address_space, read.pointer);
+            aux_cols_factory
+                .generate_read_or_immediate_aux(read, &mut row_slice.reads_aux[i].read_aux);
+        }
     }
 
     fn air(&self) -> &Self::Air {

extensions/native/circuit/src/adapters/convert_adapter.rs

@@ -216,8 +216,8 @@ impl<F: PrimeField32, const READ_SIZE: usize, const WRITE_SIZE: usize> VmAdapter
         row_slice.b_pointer = read.pointer;
         row_slice.b_as = read.address_space;
 
-        row_slice.reads_aux = [aux_cols_factory.make_read_aux_cols(read)];
-        row_slice.writes_aux = [aux_cols_factory.make_write_aux_cols(write)];
+        aux_cols_factory.generate_read_aux(read, &mut row_slice.reads_aux[0]);
+        aux_cols_factory.generate_write_aux(write, &mut row_slice.writes_aux[0]);
     }
 
     fn air(&self) -> &Self::Air {

extensions/native/circuit/src/adapters/jal_native_adapter.rs

@@ -171,7 +171,7 @@ impl<F: PrimeField32> VmAdapterChip<F> for JalNativeAdapterChip<F> {
         row_slice.from_state = write_record.from_state.map(F::from_canonical_u32);
         row_slice.a_pointer = write.pointer;
         row_slice.a_as = write.address_space;
-        row_slice.writes_aux = aux_cols_factory.make_write_aux_cols(write);
+        aux_cols_factory.generate_write_aux(write, &mut row_slice.writes_aux);
     }
 
     fn air(&self) -> &Self::Air {

extensions/native/circuit/src/adapters/loadstore_native_adapter.rs

@@ -311,18 +311,18 @@ impl<F: PrimeField32, const NUM_CELLS: usize> VmAdapterChip<F>
 
         let data_read = read_record.data_read.map(|read| memory.record_by_id(read));
         if let Some(data_read) = data_read {
-            cols.data_read_aux_cols = aux_cols_factory.make_read_aux_cols(data_read);
-        } else {
-            cols.data_read_aux_cols = MemoryReadAuxCols::disabled();
+            aux_cols_factory.generate_read_aux(data_read, &mut cols.data_read_aux_cols);
         }
 
         let write = memory.record_by_id(write_record.write_id);
         cols.data_write_as = write.address_space;
         cols.data_write_pointer = write.pointer;
 
-        cols.pointer_read_aux_cols =
-            aux_cols_factory.make_read_aux_cols(memory.record_by_id(read_record.pointer_read));
-        cols.data_write_aux_cols = aux_cols_factory.make_write_aux_cols(write);
+        aux_cols_factory.generate_read_aux(
+            memory.record_by_id(read_record.pointer_read),
+            &mut cols.pointer_read_aux_cols,
+        );
+        aux_cols_factory.generate_write_aux(write, &mut cols.data_write_aux_cols);
     }
 
     fn air(&self) -> &Self::Air {

extensions/native/circuit/src/adapters/native_vectorized_adapter.rs

@@ -223,12 +223,9 @@ impl<F: PrimeField32, const N: usize> VmAdapterChip<F> for NativeVectorizedAdapt
         row_slice.b_pointer = b_record.pointer;
         row_slice.c_pointer = c_record.pointer;
         row_slice.c_as = c_record.address_space;
-
-        row_slice.reads_aux = [
-            aux_cols_factory.make_read_aux_cols(b_record),
-            aux_cols_factory.make_read_aux_cols(c_record),
-        ];
-        row_slice.writes_aux = [aux_cols_factory.make_write_aux_cols(a_record)];
+        aux_cols_factory.generate_read_aux(b_record, &mut row_slice.reads_aux[0]);
+        aux_cols_factory.generate_read_aux(c_record, &mut row_slice.reads_aux[1]);
+        aux_cols_factory.generate_write_aux(a_record, &mut row_slice.writes_aux[0]);
     }
 
     fn air(&self) -> &Self::Air {