Implement element-addressable memory (#1598)

* feat: update memory chiplet to be element-addressable

* feat(air): fix air constraints for new memory chiplet

* feat(bus): fix bus for element-addressable memory

* fix: fix stdlib after element-addressable memory

* changelog

* more test fixes

* fix: recursive verifier

* fix sha256

* PR fixes

* nomenclature: change `batch` -> `word`

* more PR fixes

* clippy

* cleanup

* feat: make the number of locals element-addressable

* fix: fix stdlib after locals change

* more PR fixes

* more PR fixes

* fix rounding up of locals

* PR fixes

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Changelog
 
+#### Highlights
+- [BREAKING] Memory is now element-addressable (#1598)
+
 #### Changes
 - [BREAKING] `Process` no longer takes ownership of the `Host` (#1571).
 - [BREAKING] `ProcessState` was converted from a trait to a struct (#1571).

air/src/constraints/chiplets/memory/tests.rs

@@ -1,59 +1,49 @@
 use alloc::vec::Vec;
 
 use rand_utils::rand_value;
+use vm_core::{Felt, FieldElement, WORD_SIZE};
 
 use super::{
-    EvaluationFrame, MEMORY_ADDR_COL_IDX, MEMORY_CLK_COL_IDX, MEMORY_CTX_COL_IDX,
-    MEMORY_D0_COL_IDX, MEMORY_D1_COL_IDX, MEMORY_D_INV_COL_IDX, MEMORY_V_COL_RANGE, NUM_ELEMENTS,
+    EvaluationFrame, MEMORY_CLK_COL_IDX, MEMORY_CTX_COL_IDX, MEMORY_D0_COL_IDX, MEMORY_D1_COL_IDX,
+    MEMORY_D_INV_COL_IDX, MEMORY_V_COL_RANGE, MEMORY_WORD_COL_IDX,
 };
 use crate::{
     chiplets::memory,
     trace::{
         chiplets::{
-            memory::{Selectors, MEMORY_COPY_READ, MEMORY_INIT_READ, MEMORY_WRITE},
-            MEMORY_TRACE_OFFSET,
+            memory::{MEMORY_ACCESS_WORD, MEMORY_READ, MEMORY_WRITE},
+            MEMORY_FLAG_SAME_CONTEXT_AND_WORD, MEMORY_IDX0_COL_IDX, MEMORY_IDX1_COL_IDX,
+            MEMORY_IS_READ_COL_IDX, MEMORY_IS_WORD_ACCESS_COL_IDX,
         },
         TRACE_WIDTH,
     },
-    Felt, FieldElement, ONE, ZERO,
+    ONE, ZERO,
 };
 
 // UNIT TESTS
 // ================================================================================================
 
 #[test]
 fn test_memory_write() {
-    let expected = [ZERO; memory::NUM_CONSTRAINTS];
+    let expected_constraint_evals = [ZERO; memory::NUM_CONSTRAINTS];
 
-    let old_values = vec![0, 0, 0, 0];
-    let new_values = vec![1, 0, 0, 0];
+    let old_word = vec![0, 0, 0, 0];
+    let new_word = vec![1, 0, 0, 0];
 
     // Write to a new context.
-    let result = get_constraint_evaluation(
-        MEMORY_WRITE,
-        MemoryTestDeltaType::Context,
-        &old_values,
-        &new_values,
-    );
-    assert_eq!(expected, result);
+    let result =
+        get_constraint_evaluation(MEMORY_WRITE, MemoryTestDeltaType::Context, &old_word, &new_word);
+    assert_eq!(expected_constraint_evals, result);
 
     // Write to a new address in the same context.
-    let result = get_constraint_evaluation(
-        MEMORY_WRITE,
-        MemoryTestDeltaType::Address,
-        &old_values,
-        &new_values,
-    );
-    assert_eq!(expected, result);
+    let result =
+        get_constraint_evaluation(MEMORY_WRITE, MemoryTestDeltaType::Word, &old_word, &new_word);
+    assert_eq!(expected_constraint_evals, result);
 
     // Write to the same context and address at a new clock cycle.
-    let result = get_constraint_evaluation(
-        MEMORY_WRITE,
-        MemoryTestDeltaType::Clock,
-        &old_values,
-        &new_values,
-    );
-    assert_eq!(expected, result);
+    let result =
+        get_constraint_evaluation(MEMORY_WRITE, MemoryTestDeltaType::Clock, &old_word, &new_word);
+    assert_eq!(expected_constraint_evals, result);
 }
 
 #[test]
@@ -65,7 +55,7 @@ fn test_memory_read() {
 
     // Read from a new context.
     let result = get_constraint_evaluation(
-        MEMORY_INIT_READ,
+        MEMORY_READ,
         MemoryTestDeltaType::Context,
         &old_values,
         &init_values,
@@ -74,16 +64,16 @@ fn test_memory_read() {
 
     // Read from a new address in the same context.
     let result = get_constraint_evaluation(
-        MEMORY_INIT_READ,
-        MemoryTestDeltaType::Address,
+        MEMORY_READ,
+        MemoryTestDeltaType::Word,
         &old_values,
         &init_values,
     );
     assert_eq!(expected, result);
 
     // Read from the same context and address at a new clock cycle.
     let result = get_constraint_evaluation(
-        MEMORY_COPY_READ,
+        MEMORY_READ,
         MemoryTestDeltaType::Clock,
         &old_values,
         &old_values,
@@ -100,7 +90,7 @@ fn test_memory_read() {
 /// - Clock: when the delta occurs in the clock column, context and address must stay fixed.
 enum MemoryTestDeltaType {
     Context,
-    Address,
+    Word,
     Clock,
 }
 
@@ -113,17 +103,17 @@ enum MemoryTestDeltaType {
 /// - To test a valid read, the `delta_type` must be Clock and the `old_values` and `new_values`
 ///   must be equal.
 fn get_constraint_evaluation(
-    selectors: Selectors,
+    read_write: Felt,
     delta_type: MemoryTestDeltaType,
     old_values: &[u32],
     new_values: &[u32],
 ) -> [Felt; memory::NUM_CONSTRAINTS] {
     let delta_row = get_test_delta_row(&delta_type);
-    let frame = get_test_frame(selectors, &delta_type, &delta_row, old_values, new_values);
+    let frame = get_test_frame(read_write, &delta_type, &delta_row, old_values, new_values);
 
     let mut result = [ZERO; memory::NUM_CONSTRAINTS];
 
-    memory::enforce_constraints(&frame, &mut result, ONE);
+    memory::enforce_constraints(&frame, &mut result, ONE, ONE, ZERO);
 
     result
 }
@@ -141,7 +131,7 @@ fn get_constraint_evaluation(
 ///   row.
 /// - `new_values`: specifies the new values, which are placed in the value columns of the next row.
 fn get_test_frame(
-    selectors: Selectors,
+    read_write: Felt,
     delta_type: &MemoryTestDeltaType,
     delta_row: &[u64],
     old_values: &[u32],
@@ -151,16 +141,16 @@ fn get_test_frame(
     let mut next = vec![ZERO; TRACE_WIDTH];
 
     // Set the operation in the next row.
-    next[MEMORY_TRACE_OFFSET] = selectors[0];
-    next[MEMORY_TRACE_OFFSET + 1] = selectors[1];
+    next[MEMORY_IS_READ_COL_IDX] = read_write;
+    next[MEMORY_IS_WORD_ACCESS_COL_IDX] = MEMORY_ACCESS_WORD;
 
     // Set the context, addr, and clock columns in the next row to the values in the delta row.
     next[MEMORY_CTX_COL_IDX] = Felt::new(delta_row[0]);
-    next[MEMORY_ADDR_COL_IDX] = Felt::new(delta_row[1]);
+    next[MEMORY_WORD_COL_IDX] = Felt::new(delta_row[1]);
     next[MEMORY_CLK_COL_IDX] = Felt::new(delta_row[2]);
 
     // Set the old and new values.
-    for idx in 0..NUM_ELEMENTS {
+    for idx in 0..WORD_SIZE {
         let old_value = Felt::new(old_values[idx] as u64);
         // Add a write for the old values to the current row.
         current[MEMORY_V_COL_RANGE.start + idx] = old_value;
@@ -177,27 +167,41 @@ fn get_test_frame(
     let delta: u64 = match delta_type {
         MemoryTestDeltaType::Clock => delta_row[MemoryTestDeltaType::Clock as usize] - 1,
         MemoryTestDeltaType::Context => delta_row[MemoryTestDeltaType::Context as usize],
-        MemoryTestDeltaType::Address => delta_row[MemoryTestDeltaType::Address as usize],
+        MemoryTestDeltaType::Word => delta_row[MemoryTestDeltaType::Word as usize],
     };
     next[MEMORY_D0_COL_IDX] = Felt::new(delta as u16 as u64);
     next[MEMORY_D1_COL_IDX] = Felt::new(delta >> 16);
     next[MEMORY_D_INV_COL_IDX] = (Felt::new(delta)).inv();
 
+    // since we're always writing a word, the idx0 and idx1 columns should be zero
+    next[MEMORY_IDX0_COL_IDX] = ZERO;
+    next[MEMORY_IDX1_COL_IDX] = ZERO;
+
+    // If the context or word columns are changed, the "same context and word" flag should be
+    // zero.
+    if delta_row[MemoryTestDeltaType::Word as usize] > 0
+        || delta_row[MemoryTestDeltaType::Context as usize] > 0
+    {
+        next[MEMORY_FLAG_SAME_CONTEXT_AND_WORD] = ZERO;
+    } else {
+        next[MEMORY_FLAG_SAME_CONTEXT_AND_WORD] = ONE;
+    }
+
     EvaluationFrame::<Felt>::from_rows(current, next)
 }
 
-/// Generates a row of valid test values for the context, address, and clock columns according to
+/// Generates a row of valid test values for the context, word, and clock columns according to
 /// the specified delta type, which determines the column over which the delta and delta inverse
 /// values of the trace would be calculated.
 ///
-/// - When the delta type is Context, the address and clock columns can be anything.
-/// - When the delta type is Address, the context must remain unchanged but the clock can change.
-/// - When the delta type is Clock, both the context and address columns must remain unchanged.
+/// - When the delta type is Context, the word and clock columns can be anything.
+/// - When the delta type is Word, the context must remain unchanged but the clock can change.
+/// - When the delta type is Clock, both the context and word columns must remain unchanged.
 fn get_test_delta_row(delta_type: &MemoryTestDeltaType) -> Vec<u64> {
-    let delta_value = rand_value::<u32>() as u64;
+    let delta_value = word_aligned_rand_value() as u64;
     let mut row = vec![0; 3];
     let ctx_idx = MemoryTestDeltaType::Context as usize;
-    let addr_idx = MemoryTestDeltaType::Address as usize;
+    let word_addr_idx = MemoryTestDeltaType::Word as usize;
     let clk_idx = MemoryTestDeltaType::Clock as usize;
 
     // Set the context, addr, and clock columns according to the specified delta type.
@@ -207,13 +211,13 @@ fn get_test_delta_row(delta_type: &MemoryTestDeltaType) -> Vec<u64> {
             row[ctx_idx] = delta_value;
 
             // Set addr and clock in the row column to random values.
-            row[addr_idx] = rand_value::<u32>() as u64;
+            row[word_addr_idx] = word_aligned_rand_value() as u64;
             row[clk_idx] = rand_value::<u32>() as u64;
         },
-        MemoryTestDeltaType::Address => {
+        MemoryTestDeltaType::Word => {
             // Keep the context value the same in current and row rows (leave it as ZERO).
             // Set the row value for the address.
-            row[addr_idx] = delta_value;
+            row[word_addr_idx] = delta_value;
 
             // Set clock in the row column to a random value.
             row[clk_idx] = rand_value::<u32>() as u64;
@@ -227,3 +231,9 @@ fn get_test_delta_row(delta_type: &MemoryTestDeltaType) -> Vec<u64> {
 
     row
 }
+
+/// Returns a random value that is divisible by 4 (i.e. "word aligned" when treated as an address).
+fn word_aligned_rand_value() -> u32 {
+    let value = rand_value::<u32>();
+    value - (value % 4)
+}

air/src/constraints/chiplets/mod.rs

@@ -87,7 +87,13 @@ pub fn enforce_constraints<E: FieldElement<BaseField = Felt>>(
     constraint_offset += bitwise::get_transition_constraint_count();
 
     // memory transition constraints
-    memory::enforce_constraints(frame, &mut result[constraint_offset..], frame.memory_flag(false));
+    memory::enforce_constraints(
+        frame,
+        &mut result[constraint_offset..],
+        frame.memory_flag(),
+        frame.memory_flag_not_last_row(),
+        frame.memory_flag_first_row(),
+    );
 }
 
 // TRANSITION CONSTRAINT HELPERS
@@ -143,12 +149,21 @@ trait EvaluationFrameExt<E: FieldElement> {
     /// Flag to indicate whether the frame is in the bitwise portion of the Chiplets trace.
     fn bitwise_flag(&self) -> E;
 
-    /// Flag to indicate whether the frame is in the memory portion of the Chiplets trace.
-    /// When `include_last_row` is true, the memory flag is true for every row where the memory
-    /// selectors are set. When false, the last row is excluded. When this flag is used for
-    /// transition constraints with `include_last_row = false`, they will not be applied to the
-    /// final row of the memory trace.
-    fn memory_flag(&self, include_last_row: bool) -> E;
+    /// Flag to indicate whether the current row of the frame is in the memory portion of the
+    /// Chiplets trace.
+    fn memory_flag(&self) -> E;
+
+    /// Flag to indicate whether the current row of the frame is in the memory portion of the
+    /// Chiplets trace, except for the last memory chiplet row.
+    fn memory_flag_not_last_row(&self) -> E;
+
+    /// Flag to indicate whether the next row of the frame is in the memory portion of the Chiplets
+    /// trace.
+    fn memory_flag_next(&self) -> E;
+
+    /// Flag to indicate whether the next row of the frame is the first row of the memory portion of
+    /// the Chiplets trace.
+    fn memory_flag_first_row(&self) -> E;
 }
 
 impl<E: FieldElement> EvaluationFrameExt<E> for &EvaluationFrame<E> {
@@ -175,12 +190,23 @@ impl<E: FieldElement> EvaluationFrameExt<E> for &EvaluationFrame<E> {
     }
 
     #[inline(always)]
-    fn memory_flag(&self, include_last_row: bool) -> E {
-        if include_last_row {
-            self.s(0) * self.s(1) * binary_not(self.s(2))
-        } else {
-            self.s(0) * self.s(1) * binary_not(self.s_next(2))
-        }
+    fn memory_flag(&self) -> E {
+        self.s(0) * self.s(1) * binary_not(self.s(2))
+    }
+
+    #[inline(always)]
+    fn memory_flag_not_last_row(&self) -> E {
+        self.s(0) * self.s(1) * binary_not(self.s_next(2))
+    }
+
+    #[inline(always)]
+    fn memory_flag_next(&self) -> E {
+        self.s_next(0) * self.s_next(1) * binary_not(self.s_next(2))
+    }
+
+    #[inline(always)]
+    fn memory_flag_first_row(&self) -> E {
+        self.hasher_flag() * self.memory_flag_next()
     }
 }
 
@@ -196,6 +222,6 @@ pub trait ChipletsFrameExt<E: FieldElement> {
 impl<E: FieldElement> ChipletsFrameExt<E> for &EvaluationFrame<E> {
     #[inline(always)]
     fn chiplets_memory_flag(&self) -> E {
-        self.memory_flag(true)
+        self.memory_flag()
     }
 }