refactor: introduce a union type for opcode_specific_columns

.vscode/settings.json

@@ -55,6 +55,8 @@
         "tests/sha-compress/Cargo.toml",
         "tests/sha-extend/Cargo.toml",
         "tests/sha2/Cargo.toml",
+        // Eval.
+        "eval/Cargo.toml"
     ],
     "rust-analyzer.showUnlinkedFileNotification": false
 }

cli/src/build.rs

@@ -10,7 +10,8 @@ use std::{
 
 fn get_docker_image() -> String {
     // Get the docker image name from the environment variable
-    std::env::var("SP1_DOCKER_IMAGE").unwrap_or_else(|_| "ghcr.io/succinctlabs/sp1:latest".to_string())
+    std::env::var("SP1_DOCKER_IMAGE")
+        .unwrap_or_else(|_| "ghcr.io/succinctlabs/sp1:latest".to_string())
 }
 
 #[derive(Parser)]

core/src/air/word.rs

@@ -2,7 +2,6 @@ use std::array::IntoIter;
 use std::mem::size_of;
 use std::ops::{Index, IndexMut};
 
-use core::borrow::{Borrow, BorrowMut};
 use p3_air::AirBuilder;
 use p3_field::AbstractField;
 use p3_field::Field;

core/src/bytes/columns.rs

@@ -1,5 +1,3 @@
-use core::borrow::Borrow;
-use core::borrow::BorrowMut;
 use sp1_derive::AlignedBorrow;
 use std::mem::size_of;
 

core/src/cpu/air/branch.rs

@@ -1,10 +1,8 @@
-use core::borrow::Borrow;
-
 use p3_air::AirBuilder;
 use p3_field::AbstractField;
 
 use crate::air::{BaseAirBuilder, SP1AirBuilder, Word, WordAirBuilder};
-use crate::cpu::columns::{BranchCols, CpuCols, OpcodeSelectorCols, NUM_BRANCH_COLS};
+use crate::cpu::columns::{CpuCols, OpcodeSelectorCols};
 use crate::{cpu::CpuChip, runtime::Opcode};
 
 impl CpuChip {
@@ -37,8 +35,7 @@ impl CpuChip {
         next: &CpuCols<AB::Var>,
     ) {
         // Get the branch specific columns.
-        let branch_cols: BranchCols<AB::Var> =
-            *local.opcode_specific_columns[..NUM_BRANCH_COLS].borrow();
+        let branch_cols = local.opcode_specific_columns.branch();
 
         // Evaluate program counter constraints.
         {

core/src/cpu/air/memory.rs

@@ -1,10 +1,8 @@
-use core::borrow::Borrow;
-
 use p3_air::AirBuilder;
 use p3_field::AbstractField;
 
 use crate::air::{BaseAirBuilder, SP1AirBuilder, Word, WordAirBuilder};
-use crate::cpu::columns::{CpuCols, MemoryColumns, OpcodeSelectorCols, NUM_MEMORY_COLUMNS};
+use crate::cpu::columns::{CpuCols, MemoryColumns, OpcodeSelectorCols};
 use crate::cpu::CpuChip;
 use crate::memory::MemoryCols;
 use crate::runtime::Opcode;
@@ -52,23 +50,17 @@ impl CpuChip {
         local: &CpuCols<AB::Var>,
     ) {
         // Get the memory specific columns.
-        let memory_columns: MemoryColumns<AB::Var> =
-            *local.opcode_specific_columns[..NUM_MEMORY_COLUMNS].borrow();
+        let memory_columns = local.opcode_specific_columns.memory();
 
         // Compute whether this is a load instruction.
         let is_load = self.is_load_instruction::<AB>(&local.selectors);
 
         // Get the unsigned memory value.
-        self.eval_unsigned_mem_value(builder, &memory_columns, local);
+        self.eval_unsigned_mem_value(builder, memory_columns, local);
 
         // If it's a signed operation (such as LB or LH), then we need verify the bit decomposition
         // of the most significant byte to get it's sign.
-        self.eval_most_sig_byte_bit_decomp(
-            builder,
-            &memory_columns,
-            local,
-            &local.unsigned_mem_val,
-        );
+        self.eval_most_sig_byte_bit_decomp(builder, memory_columns, local, &local.unsigned_mem_val);
 
         // Assert that if `is_lb` and `is_lh` are both true, then the most significant byte is
         // matches the value of `local.mem_value_is_neg`.
@@ -107,11 +99,10 @@ impl CpuChip {
         builder: &mut AB,
         local: &CpuCols<AB::Var>,
     ) {
-        let memory_columns: MemoryColumns<AB::Var> =
-            *local.opcode_specific_columns[..NUM_MEMORY_COLUMNS].borrow();
+        let memory_columns = local.opcode_specific_columns.memory();
 
         // Get the memory offset flags.
-        self.eval_offset_value_flags(builder, &memory_columns, local);
+        self.eval_offset_value_flags(builder, memory_columns, local);
         let offset_is_zero = AB::Expr::one()
             - memory_columns.offset_is_one
             - memory_columns.offset_is_two

core/src/cpu/air/mod.rs

@@ -8,10 +8,9 @@ use p3_air::BaseAir;
 use p3_field::AbstractField;
 use p3_matrix::MatrixRowSlices;
 
-use super::columns::{NUM_AUIPC_COLS, NUM_JUMP_COLS, NUM_MEMORY_COLUMNS};
 use crate::air::{SP1AirBuilder, WordAirBuilder};
 use crate::cpu::columns::OpcodeSelectorCols;
-use crate::cpu::columns::{AuipcCols, CpuCols, JumpCols, MemoryColumns, NUM_CPU_COLS};
+use crate::cpu::columns::{CpuCols, NUM_CPU_COLS};
 use crate::cpu::CpuChip;
 use crate::memory::MemoryCols;
 use crate::runtime::{AccessPosition, Opcode};
@@ -88,8 +87,7 @@ where
 
         // For operations that require reading from memory (not registers), we need to read the
         // value into the memory columns.
-        let memory_columns: MemoryColumns<AB::Var> =
-            *local.opcode_specific_columns[..NUM_MEMORY_COLUMNS].borrow();
+        let memory_columns = local.opcode_specific_columns.memory();
         builder.constraint_memory_access(
             local.shard,
             local.clk + AB::F::from_canonical_u32(AccessPosition::Memory as u32),
@@ -175,8 +173,7 @@ impl CpuChip {
         next: &CpuCols<AB::Var>,
     ) {
         // Get the jump specific columns
-        let jump_columns: JumpCols<AB::Var> =
-            *local.opcode_specific_columns[..NUM_JUMP_COLS].borrow();
+        let jump_columns = local.opcode_specific_columns.jump();
 
         // Verify that the local.pc + 4 is saved in op_a for both jump instructions.
         builder
@@ -219,8 +216,7 @@ impl CpuChip {
     /// Constraints related to the AUIPC opcode.
     pub(crate) fn auipc_eval<AB: SP1AirBuilder>(&self, builder: &mut AB, local: &CpuCols<AB::Var>) {
         // Get the auipc specific columns.
-        let auipc_columns: AuipcCols<AB::Var> =
-            *local.opcode_specific_columns[..NUM_AUIPC_COLS].borrow();
+        let auipc_columns = local.opcode_specific_columns.auipc();
 
         // Verify that the word form of local.pc is correct.
         builder

core/src/cpu/columns/auipc.rs

@@ -1,4 +1,3 @@
-use core::borrow::{Borrow, BorrowMut};
 use sp1_derive::AlignedBorrow;
 use std::mem::size_of;
 

core/src/cpu/columns/branch.rs

@@ -1,4 +1,3 @@
-use core::borrow::{Borrow, BorrowMut};
 use sp1_derive::AlignedBorrow;
 use std::mem::size_of;
 

core/src/cpu/columns/instruction.rs

@@ -1,4 +1,3 @@
-use core::borrow::{Borrow, BorrowMut};
 use p3_field::PrimeField;
 use sp1_derive::AlignedBorrow;
 use std::mem::size_of;

core/src/cpu/columns/jump.rs

@@ -1,4 +1,3 @@
-use core::borrow::{Borrow, BorrowMut};
 use sp1_derive::AlignedBorrow;
 use std::mem::size_of;
 

core/src/cpu/columns/memory.rs

@@ -1,4 +1,3 @@
-use core::borrow::{Borrow, BorrowMut};
 use sp1_derive::AlignedBorrow;
 use std::mem::size_of;
 

core/src/cpu/columns/mod.rs

@@ -4,15 +4,16 @@ mod instruction;
 mod jump;
 mod memory;
 mod opcode;
+mod opcode_specific;
 
 pub use auipc::*;
 pub use branch::*;
 pub use instruction::*;
 pub use jump::*;
 pub use memory::*;
 pub use opcode::*;
+pub use opcode_specific::*;
 
-use core::borrow::{Borrow, BorrowMut};
 use p3_util::indices_arr;
 use sp1_derive::AlignedBorrow;
 use std::mem::{size_of, transmute};
@@ -26,12 +27,10 @@ pub const NUM_CPU_COLS: usize = size_of::<CpuCols<u8>>();
 
 pub const CPU_COL_MAP: CpuCols<usize> = make_col_map();
 
-pub const OPCODE_SPECIFIC_COLUMNS_SIZE: usize = size_of_opcode_specific_columns();
-
 /// The column layout for the CPU.
 #[derive(AlignedBorrow, Default, Debug, Clone, Copy)]
 #[repr(C)]
-pub struct CpuCols<T> {
+pub struct CpuCols<T: Copy> {
     /// The current shard.
     pub shard: T,
 
@@ -52,8 +51,7 @@ pub struct CpuCols<T> {
     pub op_b_access: MemoryReadCols<T>,
     pub op_c_access: MemoryReadCols<T>,
 
-    /// This is transmuted to MemoryColumns, BranchColumns, JumpColumns, or AUIPCColumns
-    pub opcode_specific_columns: [T; OPCODE_SPECIFIC_COLUMNS_SIZE],
+    pub opcode_specific_columns: OpcodeSpecificCols<T>,
 
     /// Selector to label whether this row is a non padded row.
     pub is_real: T,
@@ -84,20 +82,20 @@ pub struct CpuCols<T> {
     pub unsigned_mem_val: Word<T>,
 }
 
-impl<T: Clone> CpuCols<T> {
+impl<T: Copy> CpuCols<T> {
     /// Gets the value of the first operand.
     pub fn op_a_val(&self) -> Word<T> {
-        self.op_a_access.value().clone()
+        *self.op_a_access.value()
     }
 
     /// Gets the value of the second operand.
     pub fn op_b_val(&self) -> Word<T> {
-        self.op_b_access.value().clone()
+        *self.op_b_access.value()
     }
 
     /// Gets the value of the third operand.
     pub fn op_c_val(&self) -> Word<T> {
-        self.op_c_access.value().clone()
+        *self.op_c_access.value()
     }
 }
 
@@ -106,21 +104,3 @@ const fn make_col_map() -> CpuCols<usize> {
     let indices_arr = indices_arr::<NUM_CPU_COLS>();
     unsafe { transmute::<[usize; NUM_CPU_COLS], CpuCols<usize>>(indices_arr) }
 }
-
-/// Returns the size of the opcode specific columns and makes sure each opcode specific column fits.
-const fn size_of_opcode_specific_columns() -> usize {
-    let memory_columns_size = NUM_MEMORY_COLUMNS;
-    let branch_columns_size = NUM_BRANCH_COLS;
-    let jump_columns_size = NUM_JUMP_COLS;
-    let aui_pc_columns_size = NUM_AUIPC_COLS;
-
-    if branch_columns_size > memory_columns_size {
-        panic!("BranchColumns is too large to fit in the opcode_specific_columns array.");
-    } else if jump_columns_size > memory_columns_size {
-        panic!("JumpColumns is too large to fit in the opcode_specific_columns array.");
-    } else if aui_pc_columns_size > memory_columns_size {
-        panic!("AUIPCColumns is too large to fit in the opcode_specific_columns array.");
-    }
-
-    memory_columns_size
-}

core/src/cpu/columns/opcode.rs

@@ -1,4 +1,3 @@
-use core::borrow::{Borrow, BorrowMut};
 use p3_field::PrimeField;
 use sp1_derive::AlignedBorrow;
 use std::mem::size_of;

core/src/cpu/columns/opcode_specific.rs

@@ -0,0 +1,59 @@
+use crate::cpu::columns::{AuipcCols, BranchCols, JumpCols, MemoryColumns};
+use std::fmt::{Debug, Formatter};
+use std::mem::{size_of, transmute};
+
+pub const NUM_OPCODE_SPECIFIC_COLS: usize = size_of::<OpcodeSpecificCols<u8>>();
+
+/// Shared columns whose interpretation depends on the instruction being executed.
+#[derive(Clone, Copy)]
+#[repr(C)]
+pub union OpcodeSpecificCols<T: Copy> {
+    memory: MemoryColumns<T>,
+    branch: BranchCols<T>,
+    jump: JumpCols<T>,
+    auipc: AuipcCols<T>,
+}
+
+impl<T: Copy + Default> Default for OpcodeSpecificCols<T> {
+    fn default() -> Self {
+        OpcodeSpecificCols {
+            memory: MemoryColumns::default(),
+        }
+    }
+}
+
+impl<T: Copy + Debug> Debug for OpcodeSpecificCols<T> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
+        // SAFETY: repr(C) ensures uniform fields are in declaration order with no padding.
+        let self_arr: &[T; NUM_OPCODE_SPECIFIC_COLS] = unsafe { transmute(self) };
+        Debug::fmt(self_arr, f)
+    }
+}
+
+// SAFETY: Each view is a valid interpretation of the underlying array.
+impl<T: Copy> OpcodeSpecificCols<T> {
+    pub fn memory(&self) -> &MemoryColumns<T> {
+        unsafe { &self.memory }
+    }
+    pub fn memory_mut(&mut self) -> &mut MemoryColumns<T> {
+        unsafe { &mut self.memory }
+    }
+    pub fn branch(&self) -> &BranchCols<T> {
+        unsafe { &self.branch }
+    }
+    pub fn branch_mut(&mut self) -> &mut BranchCols<T> {
+        unsafe { &mut self.branch }
+    }
+    pub fn jump(&self) -> &JumpCols<T> {
+        unsafe { &self.jump }
+    }
+    pub fn jump_mut(&mut self) -> &mut JumpCols<T> {
+        unsafe { &mut self.jump }
+    }
+    pub fn auipc(&self) -> &AuipcCols<T> {
+        unsafe { &self.auipc }
+    }
+    pub fn auipc_mut(&mut self) -> &mut AuipcCols<T> {
+        unsafe { &mut self.auipc }
+    }
+}

core/src/cpu/trace.rs

@@ -1,12 +1,9 @@
-use super::columns::{
-    AuipcCols, BranchCols, JumpCols, CPU_COL_MAP, NUM_AUIPC_COLS, NUM_BRANCH_COLS, NUM_CPU_COLS,
-    NUM_JUMP_COLS, NUM_MEMORY_COLUMNS,
-};
+use super::columns::{CPU_COL_MAP, NUM_CPU_COLS};
 use super::{CpuChip, CpuEvent};
 use crate::air::MachineAir;
 use crate::alu::{self, AluEvent};
 use crate::bytes::{ByteLookupEvent, ByteOpcode};
-use crate::cpu::columns::{CpuCols, MemoryColumns};
+use crate::cpu::columns::CpuCols;
 use crate::cpu::memory::MemoryRecordEnum;
 use crate::disassembler::WORD_SIZE;
 use crate::field::event::FieldEvent;
@@ -157,8 +154,7 @@ impl CpuChip {
 
         // Populate memory accesses for reading from memory.
         assert_eq!(event.memory_record.is_some(), event.memory.is_some());
-        let memory_columns: &mut MemoryColumns<F> =
-            cols.opcode_specific_columns[..NUM_MEMORY_COLUMNS].borrow_mut();
+        let memory_columns = cols.opcode_specific_columns.memory_mut();
         if let Some(record) = event.memory_record {
             memory_columns
                 .memory_access
@@ -200,8 +196,7 @@ impl CpuChip {
         }
 
         // Populate addr_word and addr_aligned columns.
-        let memory_columns: &mut MemoryColumns<F> =
-            cols.opcode_specific_columns[0..NUM_MEMORY_COLUMNS].borrow_mut();
+        let memory_columns = cols.opcode_specific_columns.memory_mut();
         let memory_addr = event.b.wrapping_add(event.c);
         memory_columns.addr_word = memory_addr.into();
         memory_columns.addr_aligned =
@@ -308,8 +303,7 @@ impl CpuChip {
         alu_events: &mut HashMap<Opcode, Vec<alu::AluEvent>>,
     ) {
         if event.instruction.is_branch_instruction() {
-            let branch_columns: &mut BranchCols<F> =
-                cols.opcode_specific_columns[..NUM_BRANCH_COLS].borrow_mut();
+            let branch_columns = cols.opcode_specific_columns.branch_mut();
 
             let a_eq_b = event.a == event.b;
 
@@ -404,8 +398,7 @@ impl CpuChip {
         alu_events: &mut HashMap<Opcode, Vec<alu::AluEvent>>,
     ) {
         if event.instruction.is_jump_instruction() {
-            let jump_columns: &mut JumpCols<F> =
-                cols.opcode_specific_columns[..NUM_JUMP_COLS].borrow_mut();
+            let jump_columns = cols.opcode_specific_columns.jump_mut();
 
             match event.instruction.opcode {
                 Opcode::JAL => {
@@ -456,8 +449,7 @@ impl CpuChip {
         alu_events: &mut HashMap<Opcode, Vec<alu::AluEvent>>,
     ) {
         if matches!(event.instruction.opcode, Opcode::AUIPC) {
-            let auipc_columns: &mut AuipcCols<F> =
-                cols.opcode_specific_columns[..NUM_AUIPC_COLS].borrow_mut();
+            let auipc_columns = cols.opcode_specific_columns.auipc_mut();
 
             auipc_columns.pc = event.pc.into();
 

core/src/memory/columns.rs

@@ -1,4 +1,3 @@
-use core::borrow::{Borrow, BorrowMut};
 use sp1_derive::AlignedBorrow;
 use std::mem::size_of;