From d8fbb77c606c021cd56302c51ee897c0a6c25ee4 Mon Sep 17 00:00:00 2001
From: Gio <102917377+gio256@users.noreply.github.com>
Date: Tue, 9 Jul 2024 14:17:30 -0600
Subject: [PATCH] fix(evm_arithmetization): Adjust layout of
 `CpuGeneralColumnsView` (#355)

* Adjust layout of CpuGeneralColumnsView

* Rename misleading UNION_FIELD_SIZE

* Clarify unused columns in CpuStackView

* Remove NUM_UNION_COLUMNS constant
---
 .../src/cpu/columns/general.rs                | 40 +++++++++++++++++--
 1 file changed, 36 insertions(+), 4 deletions(-)

diff --git a/evm_arithmetization/src/cpu/columns/general.rs b/evm_arithmetization/src/cpu/columns/general.rs
index 8b2c2bc41..a3aeed13f 100644
--- a/evm_arithmetization/src/cpu/columns/general.rs
+++ b/evm_arithmetization/src/cpu/columns/general.rs
@@ -2,8 +2,11 @@ use core::borrow::{Borrow, BorrowMut};
 use core::fmt::{Debug, Formatter};
 use core::mem::{size_of, transmute};
 
+use static_assertions::const_assert;
+
 /// General purpose columns, which can have different meanings depending on what
 /// CTL or other operation is occurring at this row.
+#[repr(C)]
 #[derive(Clone, Copy)]
 pub(crate) union CpuGeneralColumnsView<T: Copy> {
     exception: CpuExceptionView<T>,
@@ -108,14 +111,21 @@ impl<T: Copy> BorrowMut<[T; NUM_SHARED_COLUMNS]> for CpuGeneralColumnsView<T> {
 }
 
 /// View of the first three `CpuGeneralColumns` containing exception code bits.
+#[repr(C)]
 #[derive(Copy, Clone)]
 pub(crate) struct CpuExceptionView<T: Copy> {
     /// Exception code as little-endian bits.
     pub(crate) exc_code_bits: [T; 3],
+    /// Reserve the unused columns.
+    _padding_columns: [T; NUM_SHARED_COLUMNS - 3],
 }
 
 /// View of the `CpuGeneralColumns` storing pseudo-inverses used to prove logic
 /// operations.
+///
+/// Because this is the largest field of the [`CpuGeneralColumnsView`] union,
+/// we don't add any padding columns.
+#[repr(C)]
 #[derive(Copy, Clone)]
 pub(crate) struct CpuLogicView<T: Copy> {
     /// Pseudoinverse of `(input0 - input1)`. Used prove that they are unequal.
@@ -125,30 +135,40 @@ pub(crate) struct CpuLogicView<T: Copy> {
 
 /// View of the first two `CpuGeneralColumns` storing a flag and a pseudoinverse
 /// used to prove jumps.
+#[repr(C)]
 #[derive(Copy, Clone)]
 pub(crate) struct CpuJumpsView<T: Copy> {
     /// A flag indicating whether a jump should occur.
     pub(crate) should_jump: T,
     /// Pseudoinverse of `cond.iter().sum()`. Used to check `should_jump`.
     pub(crate) cond_sum_pinv: T,
+    /// Reserve the unused columns.
+    _padding_columns: [T; NUM_SHARED_COLUMNS - 2],
 }
 
 /// View of the first `CpuGeneralColumns` storing a pseudoinverse used to prove
 /// shift operations.
+#[repr(C)]
 #[derive(Copy, Clone)]
 pub(crate) struct CpuShiftView<T: Copy> {
     /// For a shift amount of displacement: [T], this is the inverse of
     /// sum(displacement[1..]) or zero if the sum is zero.
     pub(crate) high_limb_sum_inv: T,
+    /// Reserve the unused columns.
+    _padding_columns: [T; NUM_SHARED_COLUMNS - 1],
 }
 
 /// View of the last four `CpuGeneralColumns` storing stack-related variables.
 /// The first three are used for conditionally enabling and disabling channels
 /// when reading the next `stack_top`, and the fourth one is used to check for
 /// stack overflow.
+#[repr(C)]
 #[derive(Copy, Clone)]
 pub(crate) struct CpuStackView<T: Copy> {
-    _unused: [T; 4],
+    /// Reserve the unused columns at the beginning. This allows `Self` to
+    /// coexist with any view that uses only the first four columns (i.e. all
+    /// except `CpuLogicView`).
+    _unused: [T; NUM_SHARED_COLUMNS - 4],
     /// Pseudoinverse of `stack_len - num_pops`.
     pub(crate) stack_inv: T,
     /// stack_inv * stack_len.
@@ -160,6 +180,18 @@ pub(crate) struct CpuStackView<T: Copy> {
     pub(crate) stack_len_bounds_aux: T,
 }
 
-/// Number of columns shared by all the views of `CpuGeneralColumnsView`.
-/// `u8` is guaranteed to have a `size_of` of 1.
-pub(crate) const NUM_SHARED_COLUMNS: usize = size_of::<CpuGeneralColumnsView<u8>>();
+/// The number of columns shared by all views of [`CpuGeneralColumnsView`].
+/// This is defined in terms of the largest view in order to determine the
+/// number of padding columns to add to each field without creating a cycle
+/// for rustc.
+/// NB: `u8` is guaranteed to have a `size_of` of 1.
+pub(crate) const NUM_SHARED_COLUMNS: usize = size_of::<CpuLogicView<u8>>();
+const_assert!(NUM_SHARED_COLUMNS == size_of::<CpuGeneralColumnsView<u8>>());
+
+/// Assert that each field of the [`CpuGeneralColumnsView`] union contains the
+/// correct number of columns.
+const_assert!(size_of::<CpuExceptionView<u8>>() == NUM_SHARED_COLUMNS);
+const_assert!(size_of::<CpuLogicView<u8>>() == NUM_SHARED_COLUMNS);
+const_assert!(size_of::<CpuJumpsView<u8>>() == NUM_SHARED_COLUMNS);
+const_assert!(size_of::<CpuShiftView<u8>>() == NUM_SHARED_COLUMNS);
+const_assert!(size_of::<CpuStackView<u8>>() == NUM_SHARED_COLUMNS);