From 1affc47613a64ec15d0bacd756d8c64cbc219b95 Mon Sep 17 00:00:00 2001 From: Zach Langley Date: Tue, 7 Jan 2025 13:56:47 -0500 Subject: [PATCH 1/2] feat: Remove legacy flatten/from_iter APIs --- .../system/memory/offline_checker/columns.rs | 247 ++---------------- 1 file changed, 27 insertions(+), 220 deletions(-) diff --git a/crates/vm/src/system/memory/offline_checker/columns.rs b/crates/vm/src/system/memory/offline_checker/columns.rs index c6e09f236f..2b850dc756 100644 --- a/crates/vm/src/system/memory/offline_checker/columns.rs +++ b/crates/vm/src/system/memory/offline_checker/columns.rs @@ -1,8 +1,6 @@ //! Defines auxiliary columns for memory operations: `MemoryReadAuxCols`, //! `MemoryReadWithImmediateAuxCols`, and `MemoryWriteAuxCols`. -use std::{array, borrow::Borrow, iter}; - use openvm_circuit_primitives::is_less_than::LessThanAuxCols; use openvm_circuit_primitives_derive::AlignedBorrow; use openvm_stark_backend::p3_field::{FieldAlgebra, PrimeField32}; @@ -21,29 +19,6 @@ pub struct MemoryBaseAuxCols { pub(super) clk_lt_aux: LessThanAuxCols, } -impl MemoryBaseAuxCols { - /// TODO[arayi]: Since we have AlignedBorrow, should remove all from_slice, from_iterator, and flatten in a future PR. - pub fn from_slice(slc: &[T]) -> Self { - let base_aux_cols: &MemoryBaseAuxCols = slc.borrow(); - base_aux_cols.clone() - } - - pub fn from_iterator>(iter: &mut I) -> Self { - let sm = iter.take(Self::width()).collect::>(); - let base_aux_cols: &MemoryBaseAuxCols = sm[..].borrow(); - base_aux_cols.clone() - } -} - -impl MemoryBaseAuxCols { - pub fn flatten(self) -> Vec { - iter::empty() - .chain(iter::once(self.prev_timestamp)) - .chain(self.clk_lt_aux.lower_decomp) - .collect() - } -} - #[repr(C)] #[derive(Clone, Copy, Debug, AlignedBorrow)] pub struct MemoryWriteAuxCols { @@ -63,22 +38,7 @@ impl MemoryWriteAuxCols { } } -impl MemoryWriteAuxCols { - pub fn from_slice(slc: &[T]) -> Self { - let width = MemoryBaseAuxCols::::width(); - Self { - base: MemoryBaseAuxCols::from_slice(&slc[..width]), - prev_data: array::from_fn(|i| slc[width + i].clone()), - } - } - - pub fn from_iterator>(iter: &mut I) -> Self { - Self { - base: MemoryBaseAuxCols::from_iterator(iter), - prev_data: array::from_fn(|_| iter.next().unwrap()), - } - } - +impl MemoryWriteAuxCols { pub fn from_base(base: MemoryBaseAuxCols, prev_data: [T; N]) -> Self { Self { base, prev_data } } @@ -88,19 +48,17 @@ impl MemoryWriteAuxCols { } } -impl MemoryWriteAuxCols { - pub fn flatten(self) -> Vec { - iter::empty() - .chain(self.base.flatten()) - .chain(self.prev_data) - .collect() - } -} - -impl MemoryWriteAuxCols { +impl MemoryWriteAuxCols { pub fn disabled() -> Self { - let width = MemoryWriteAuxCols::::width(); - MemoryWriteAuxCols::from_slice(&F::zero_vec(width)) + Self { + base: MemoryBaseAuxCols { + prev_timestamp: F::ZERO, + clk_lt_aux: LessThanAuxCols { + lower_decomp: [F::ZERO; AUX_LEN], + }, + }, + prev_data: [F::ZERO; N], + } } } @@ -125,101 +83,17 @@ impl MemoryReadAuxCols { } } -impl MemoryReadAuxCols { - pub fn from_slice(slc: &[T]) -> Self { - Self { - base: MemoryBaseAuxCols::from_slice(slc), - } - } - - pub fn from_iterator>(iter: &mut I) -> Self { - Self { - base: MemoryBaseAuxCols::from_iterator(iter), - } - } -} - -impl MemoryReadAuxCols { - pub fn flatten(self) -> Vec { - self.base.flatten() - } -} - impl MemoryReadAuxCols { pub fn disabled() -> Self { - let width = MemoryReadAuxCols::::width(); - MemoryReadAuxCols::from_slice(&F::zero_vec(width)) - } -} - -#[repr(C)] -#[derive(Clone, Debug, AlignedBorrow)] -pub struct MemoryHeapReadAuxCols { - pub address: MemoryReadAuxCols, - pub data: MemoryReadAuxCols, -} - -impl MemoryHeapReadAuxCols { - pub fn from_iterator>(iter: &mut I) -> Self { Self { - address: MemoryReadAuxCols::from_iterator(iter), - data: MemoryReadAuxCols::from_iterator(iter), - } - } - - pub fn flatten(self) -> Vec { - iter::empty() - .chain(self.address.flatten()) - .chain(self.data.flatten()) - .collect() - } -} - -impl MemoryHeapReadAuxCols { - pub fn disabled() -> Self { - let width = MemoryReadAuxCols::::width(); - let address = MemoryReadAuxCols::from_slice(&F::zero_vec(width)); - let width = MemoryReadAuxCols::::width(); - let data = MemoryReadAuxCols::from_slice(&F::zero_vec(width)); - MemoryHeapReadAuxCols { address, data } - } -} - -#[repr(C)] -#[derive(Clone, Debug)] -pub struct MemoryHeapWriteAuxCols { - pub address: MemoryReadAuxCols, - pub data: MemoryWriteAuxCols, -} - -impl MemoryHeapWriteAuxCols { - pub fn from_iterator>(iter: &mut I) -> Self { - Self { - address: MemoryReadAuxCols::from_iterator(iter), - data: MemoryWriteAuxCols::from_iterator(iter), + base: MemoryBaseAuxCols { + prev_timestamp: F::ZERO, + clk_lt_aux: LessThanAuxCols { + lower_decomp: [F::ZERO; AUX_LEN], + }, + }, } } - - pub fn flatten(self) -> Vec { - iter::empty() - .chain(self.address.flatten()) - .chain(self.data.flatten()) - .collect() - } - - pub const fn width() -> usize { - MemoryReadAuxCols::::width() + MemoryWriteAuxCols::::width() - } -} - -impl MemoryHeapWriteAuxCols { - pub fn disabled() -> Self { - let width = MemoryReadAuxCols::::width(); - let address = MemoryReadAuxCols::from_slice(&F::zero_vec(width)); - let width = MemoryWriteAuxCols::::width(); - let data = MemoryWriteAuxCols::from_slice(&F::zero_vec(width)); - MemoryHeapWriteAuxCols { address, data } - } } #[repr(C)] @@ -248,84 +122,17 @@ impl MemoryReadOrImmediateAuxCols { } } -impl MemoryReadOrImmediateAuxCols { - pub fn from_slice(slc: &[T]) -> Self { - let width = MemoryBaseAuxCols::::width(); - Self { - base: MemoryBaseAuxCols::from_slice(&slc[..width]), - is_immediate: slc[width].clone(), - is_zero_aux: slc[width + 1].clone(), - } - } - - pub fn from_iterator>(iter: &mut I) -> Self { - Self { - base: MemoryBaseAuxCols::from_iterator(iter), - is_immediate: iter.next().unwrap(), - is_zero_aux: iter.next().unwrap(), - } - } -} - -impl MemoryReadOrImmediateAuxCols { - pub fn flatten(self) -> Vec { - iter::empty() - .chain(self.base.flatten()) - .chain(iter::once(self.is_immediate)) - .chain(iter::once(self.is_zero_aux)) - .collect() - } -} - impl MemoryReadOrImmediateAuxCols { pub fn disabled() -> Self { - let width = MemoryReadOrImmediateAuxCols::::width(); - MemoryReadOrImmediateAuxCols::from_slice(&F::zero_vec(width)) - } -} - -#[cfg(test)] -mod tests { - use openvm_stark_sdk::p3_baby_bear::BabyBear; - - use super::*; - - #[test] - fn test_write_aux_cols_width() { - type F = BabyBear; - - let disabled = MemoryWriteAuxCols::::disabled(); - assert_eq!( - disabled.flatten().len(), - MemoryWriteAuxCols::::width() - ); - - let disabled = MemoryWriteAuxCols::::disabled(); - assert_eq!( - disabled.flatten().len(), - MemoryWriteAuxCols::::width() - ); - } - - #[test] - fn test_read_aux_cols_width() { - type F = BabyBear; - - let disabled = MemoryReadAuxCols::::disabled(); - assert_eq!(disabled.flatten().len(), MemoryReadAuxCols::::width()); - - let disabled = MemoryReadAuxCols::::disabled(); - assert_eq!(disabled.flatten().len(), MemoryReadAuxCols::::width()); - } - - #[test] - fn test_read_or_immediate_aux_cols_width() { - type F = BabyBear; - - let disabled = MemoryReadOrImmediateAuxCols::::disabled(); - assert_eq!( - disabled.flatten().len(), - MemoryReadOrImmediateAuxCols::::width() - ); + MemoryReadOrImmediateAuxCols { + base: MemoryBaseAuxCols { + prev_timestamp: F::ZERO, + clk_lt_aux: LessThanAuxCols { + lower_decomp: [F::ZERO; AUX_LEN], + }, + }, + is_immediate: F::ZERO, + is_zero_aux: F::ZERO, + } } } From 4a9d88682e229144ef576a0681f6954ede546bf4 Mon Sep 17 00:00:00 2001 From: Jonathan Wang <31040440+jonathanpwang@users.noreply.github.com> Date: Tue, 7 Jan 2025 11:13:17 -0800 Subject: [PATCH 2/2] Apply suggestions from code review --- crates/vm/src/system/memory/offline_checker/columns.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/crates/vm/src/system/memory/offline_checker/columns.rs b/crates/vm/src/system/memory/offline_checker/columns.rs index 2b850dc756..e9036017f3 100644 --- a/crates/vm/src/system/memory/offline_checker/columns.rs +++ b/crates/vm/src/system/memory/offline_checker/columns.rs @@ -49,7 +49,7 @@ impl MemoryWriteAuxCols { } impl MemoryWriteAuxCols { - pub fn disabled() -> Self { + pub const fn disabled() -> Self { Self { base: MemoryBaseAuxCols { prev_timestamp: F::ZERO, @@ -84,7 +84,7 @@ impl MemoryReadAuxCols { } impl MemoryReadAuxCols { - pub fn disabled() -> Self { + pub const fn disabled() -> Self { Self { base: MemoryBaseAuxCols { prev_timestamp: F::ZERO, @@ -123,7 +123,7 @@ impl MemoryReadOrImmediateAuxCols { } impl MemoryReadOrImmediateAuxCols { - pub fn disabled() -> Self { + pub const fn disabled() -> Self { MemoryReadOrImmediateAuxCols { base: MemoryBaseAuxCols { prev_timestamp: F::ZERO,