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

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

@@ -1,7 +1,7 @@
 use alloc::vec::Vec;
 
 use rand_utils::rand_value;
-use vm_core::{Felt, FieldElement};
+use vm_core::{Felt, FieldElement, WORD_SIZE};
 
 use super::{
     EvaluationFrame, MEMORY_BATCH_COL_IDX, MEMORY_CLK_COL_IDX, MEMORY_CTX_COL_IDX,
@@ -11,7 +11,7 @@ use crate::{
     chiplets::memory,
     trace::{
         chiplets::{
-            memory::{MEMORY_ACCESS_WORD, MEMORY_READ, MEMORY_WRITE, NUM_ELEMENTS_IN_BATCH},
+            memory::{MEMORY_ACCESS_WORD, MEMORY_READ, MEMORY_WRITE},
             MEMORY_ELEMENT_OR_WORD_COL_IDX, MEMORY_FLAG_SAME_BATCH_AND_CONTEXT,
             MEMORY_IDX0_COL_IDX, MEMORY_IDX1_COL_IDX, MEMORY_READ_WRITE_COL_IDX,
         },
@@ -150,7 +150,7 @@ fn get_test_frame(
     next[MEMORY_CLK_COL_IDX] = Felt::new(delta_row[2]);
 
     // Set the old and new values.
-    for idx in 0..NUM_ELEMENTS_IN_BATCH {
+    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;

air/src/trace/chiplets/memory.rs

@@ -1,3 +1,5 @@
+use vm_core::WORD_SIZE;
+
 use super::{create_range, Felt, Range, ONE, ZERO};
 
 // CONSTANTS
@@ -6,25 +8,6 @@ use super::{create_range, Felt, Range, ONE, ZERO};
 /// Number of columns needed to record an execution trace of the memory chiplet.
 pub const TRACE_WIDTH: usize = 15;
 
-// TODO(plafer): get rid of all "selector" constants
-/// Number of selector columns in the trace.
-pub const NUM_SELECTORS: usize = 2;
-
-/// Type for Memory trace selectors.
-///
-/// These selectors are used to define which operation and memory state update (init & read / copy &
-/// read / write) is to be applied at a specific row of the memory execution trace.
-pub type Selectors = [Felt; NUM_SELECTORS];
-
-/// Specifies an operation that initializes new memory and then reads it.
-pub const MEMORY_INIT_READ: Selectors = [ONE, ZERO];
-
-/// Specifies an operation that copies existing memory and then reads it.
-pub const MEMORY_COPY_READ: Selectors = [ONE, ONE];
-
-/// Specifies a memory write operation.
-pub const MEMORY_WRITE_SELECTOR: Selectors = [ZERO, ZERO];
-
 // --- OPERATION SELECTORS ------------------------------------------------------------------------
 
 /// Specifies the value of the `READ_WRITE` column when the operation is a write.
@@ -36,20 +19,26 @@ pub const MEMORY_ACCESS_ELEMENT: Felt = ZERO;
 /// Specifies the value of the `ELEMENT_OR_WORD` column when the operation is over a word.
 pub const MEMORY_ACCESS_WORD: Felt = ONE;
 
-// TODO(plafer): figure out the new labels
+// --- BUS LABELS ------------------------------------------------------------------------
 
-/// Unique label computed as 1 plus the full chiplet selector with the bits reversed.
-/// mem_read selector=[1, 1, 0, 1], rev(selector)=[1, 0, 1, 1], +1=[1, 1, 0, 0]
-pub const MEMORY_READ_LABEL: u8 = 0b1100;
+// All bus labels encode the chiplet selector (1, 1, 0), as well as the read/write and element/word
+// columns. The purpose of the label is to force the chiplet to assign the correct values to the
+// read/write and element/word columns. We also include the chiplet selector as a "namespace" for
+// memory chiplet labels (to really ensure they don't collide with labels from other chiplets).
 
-/// Unique label computed as 1 plus the full chiplet selector with the bits reversed.
-/// mem_write selector=[1, 1, 0, 0] rev(selector)=[0, 0, 1, 1] +1=[0, 1, 0, 0]
-pub const MEMORY_WRITE_LABEL: u8 = 0b0100;
+/// Unique label when r/w=0 and e/w=0.
+pub const MEMORY_WRITE_ELEMENT_LABEL: u8 = 0b11000;
 
-// --- COLUMN ACCESSOR INDICES WITHIN THE CHIPLET -------------------------------------------------
+/// Unique label when r/w=0 and e/w=1.
+pub const MEMORY_WRITE_WORD_LABEL: u8 = 0b11001;
+
+/// Unique label when r/w=1 and e/w=0.
+pub const MEMORY_READ_ELEMENT_LABEL: u8 = 0b11010;
 
-/// The number of elements accessible in one read or write memory access.
-pub const NUM_ELEMENTS_IN_BATCH: usize = 4;
+/// Unique label when r/w=1 and e/w=1.
+pub const MEMORY_READ_WORD_LABEL: u8 = 0b11011;
+
+// --- COLUMN ACCESSOR INDICES WITHIN THE CHIPLET -------------------------------------------------
 
 /// Column to hold the whether the operation is a read or write.
 pub const READ_WRITE_COL_IDX: usize = 0;
@@ -67,7 +56,7 @@ pub const IDX1_COL_IDX: usize = IDX0_COL_IDX + 1;
 pub const CLK_COL_IDX: usize = IDX1_COL_IDX + 1;
 /// Columns to hold the values stored at a given memory context, address, and clock cycle after
 /// the memory operation. When reading from a new address, these are initialized to zero.
-pub const V_COL_RANGE: Range<usize> = create_range(CLK_COL_IDX + 1, NUM_ELEMENTS_IN_BATCH);
+pub const V_COL_RANGE: Range<usize> = create_range(CLK_COL_IDX + 1, WORD_SIZE);
 /// Column for the lower 16-bits of the delta between two consecutive context IDs, addresses, or
 /// clock cycles.
 pub const D0_COL_IDX: usize = V_COL_RANGE.end;

air/src/trace/main_trace.rs

@@ -10,7 +10,7 @@ use super::{
         hasher::{DIGEST_LEN, HASH_CYCLE_LEN, STATE_WIDTH},
         BITWISE_A_COL_IDX, BITWISE_B_COL_IDX, BITWISE_OUTPUT_COL_IDX, HASHER_NODE_INDEX_COL_IDX,
         HASHER_STATE_COL_RANGE, MEMORY_BATCH_COL_IDX, MEMORY_CLK_COL_IDX, MEMORY_CTX_COL_IDX,
-        MEMORY_V_COL_RANGE,
+        MEMORY_IDX0_COL_IDX, MEMORY_IDX1_COL_IDX, MEMORY_V_COL_RANGE,
     },
     decoder::{
         GROUP_COUNT_COL_IDX, HASHER_STATE_OFFSET, IN_SPAN_COL_IDX, IS_CALL_FLAG_COL_IDX,
@@ -371,10 +371,20 @@ impl MainTrace {
     }
 
     /// Returns the i-th row of the chiplet column containing memory address.
-    pub fn chiplet_memory_addr(&self, i: RowIndex) -> Felt {
+    pub fn chiplet_memory_batch(&self, i: RowIndex) -> Felt {
         self.columns.get_column(MEMORY_BATCH_COL_IDX)[i]
     }
 
+    /// Returns the i-th row of the chiplet column containing 0th bit of the batch index.
+    pub fn chiplet_memory_idx0(&self, i: RowIndex) -> Felt {
+        self.columns.get_column(MEMORY_IDX0_COL_IDX)[i]
+    }
+
+    /// Returns the i-th row of the chiplet column containing 1st bit of the batch index.
+    pub fn chiplet_memory_idx1(&self, i: RowIndex) -> Felt {
+        self.columns.get_column(MEMORY_IDX1_COL_IDX)[i]
+    }
+
     /// Returns the i-th row of the chiplet column containing clock cycle.
     pub fn chiplet_memory_clk(&self, i: RowIndex) -> Felt {
         self.columns.get_column(MEMORY_CLK_COL_IDX)[i]