Add Columns and DerefColumns derive macros

Cargo.toml

@@ -3,6 +3,7 @@ members = ["mpt_trie",
     "proof_gen",
     "trace_decoder",
     "evm_arithmetization",
+    "derive",
     "zero_bin/leader",
     "zero_bin/worker",
     "zero_bin/common",
@@ -108,3 +109,9 @@ plonky2 = "0.2.2"
 plonky2_maybe_rayon = "0.2.0"
 plonky2_util = "0.2.0"
 starky = "0.4.0"
+
+# proc macro related dependencies
+syn = "2.0"
+proc-macro2 = "1.0"
+quote = "1.0"
+zk_evm_derive = { path = "derive" }

derive/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "zk_evm_derive"
+version = "0.1.0"
+publish = false
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[lib]
+proc-macro = true
+
+[dependencies]
+syn = { workspace = true }
+proc-macro2 = { workspace = true }
+quote = { workspace = true }

derive/src/common.rs

@@ -0,0 +1,41 @@
+use syn::punctuated::Punctuated;
+use syn::{token, Attribute, Meta};
+
+/// Prefixes an error message and generates a `syn::Error` from the message.
+macro_rules! span_err {
+    ($ast:expr, $msg:literal $(,)?) => {
+        ::syn::Error::new_spanned($ast, ::core::concat!("zk_evm_derive error: ", $msg))
+    };
+}
+pub(crate) use span_err;
+
+/// Checks the condition and returns early with the given error message if
+/// false.
+macro_rules! ensure {
+    ($cond:expr, $ast:expr, $msg:literal $(,)?) => {
+        if !$cond {
+            return Err($crate::common::span_err!($ast, $msg));
+        }
+    };
+}
+pub(crate) use ensure;
+
+/// Parses the `Meta` of a `repr` attribute and returns true if one of the
+/// elements is "C".
+fn is_meta_c(outer: &Meta) -> bool {
+    if let Meta::List(inner) = outer {
+        let parsed: Punctuated<Meta, token::Comma> = inner
+            .parse_args_with(Punctuated::parse_terminated)
+            .unwrap_or_default();
+        parsed.iter().any(|meta| meta.path().is_ident("C"))
+    } else {
+        false
+    }
+}
+
+/// Returns true if `#[repr(C)]` is contained in the attributes.
+pub(crate) fn is_repr_c<'a>(attrs: impl IntoIterator<Item = &'a Attribute>) -> bool {
+    attrs
+        .into_iter()
+        .any(|attr| attr.path().is_ident("repr") && is_meta_c(&attr.meta))
+}

derive/src/impls/columns.rs

@@ -0,0 +1,135 @@
+use quote::quote;
+use syn::{Data, DeriveInput, Result};
+
+use crate::common::{ensure, is_repr_c};
+
+/// Implements `Borrow`, `BorrowMut`, `From`, `Index`, `IndexMut`, and
+/// `Default`.
+pub(crate) fn try_derive(ast: DeriveInput) -> Result<proc_macro2::TokenStream> {
+    let is_struct = matches!(ast.data, Data::Struct(_));
+    ensure!(is_struct, &ast, "expected `struct`");
+
+    // Check that the struct is `#[repr(C)]`
+    let repr_c = is_repr_c(&ast.attrs);
+    ensure!(repr_c, &ast, "column struct must be `#[repr(C)]`");
+
+    // The name of the struct.
+    let name = &ast.ident;
+
+    // Safety: `u8` is guaranteed to have a `size_of` of 1.
+    // https://doc.rust-lang.org/reference/type-layout.html#primitive-data-layout
+    let num_columns = quote!(::core::mem::size_of::<#name<u8>>());
+
+    // Safety:
+    // A repr(C) struct generic over T has the same layout as an array [T; N] if:
+    // - Every field of the struct is either T or a type with the same alignment as
+    //   T and a size of `size_of::<T>() * M` where M <= N. i.e. every field is one
+    //   of T, [T; M], or a type with the same layout as [T; M].
+    // - The total number of elements of type T is N.
+    // https://doc.rust-lang.org/reference/type-layout.html#reprc-structs
+    // https://doc.rust-lang.org/reference/type-layout.html#array-layout
+    Ok(quote! {
+        impl<T> ::core::borrow::Borrow<#name<T>> for [T; #num_columns]
+        where
+            T: ::core::marker::Copy,
+        {
+            fn borrow(&self) -> &#name<T> {
+                unsafe { ::core::mem::transmute(self) }
+            }
+        }
+
+        impl<T> ::core::borrow::BorrowMut<#name<T>> for [T; #num_columns]
+        where
+            T: ::core::marker::Copy,
+        {
+            fn borrow_mut(&mut self) -> &mut #name<T> {
+                unsafe { ::core::mem::transmute(self) }
+            }
+        }
+
+        impl<T> ::core::borrow::Borrow<[T; #num_columns]> for #name<T>
+        where
+            T: ::core::marker::Copy,
+        {
+            fn borrow(&self) -> &[T; #num_columns] {
+                unsafe { ::core::mem::transmute(self) }
+            }
+        }
+
+        impl<T> ::core::borrow::BorrowMut<[T; #num_columns]> for #name<T>
+        where
+            T: ::core::marker::Copy,
+        {
+            fn borrow_mut(&mut self) -> &mut [T; #num_columns] {
+                unsafe { ::core::mem::transmute(self) }
+            }
+        }
+
+        impl<T> From<[T; #num_columns]> for #name<T>
+        where
+            T: ::core::marker::Copy,
+        {
+            fn from(value: [T; #num_columns]) -> Self {
+                debug_assert_eq!(
+                    ::core::mem::size_of::<#name<T>>(),
+                    ::core::mem::size_of::<[T; #num_columns]>()
+                );
+                // Need ManuallyDrop so that `value` is not dropped by this function.
+                let value = ::core::mem::ManuallyDrop::new(value);
+                // Copy the bit pattern. The original value is no longer safe to use.
+                unsafe { ::core::mem::transmute_copy(&value) }
+            }
+        }
+
+        impl<T> From<#name<T>> for [T; #num_columns]
+        where
+            T: ::core::marker::Copy,
+        {
+            fn from(value: #name<T>) -> Self {
+                debug_assert_eq!(
+                    ::core::mem::size_of::<#name<T>>(),
+                    ::core::mem::size_of::<[T; #num_columns]>()
+                );
+                // Need ManuallyDrop so that `value` is not dropped by this function.
+                let value = ::core::mem::ManuallyDrop::new(value);
+                // Copy the bit pattern. The original value is no longer safe to use.
+                unsafe { ::core::mem::transmute_copy(&value) }
+            }
+        }
+
+        impl<T, I> ::core::ops::Index<I> for #name<T>
+        where
+            T: ::core::marker::Copy,
+            [T]: ::core::ops::Index<I>,
+        {
+            type Output = <[T] as ::core::ops::Index<I>>::Output;
+
+            fn index(&self, index: I) -> &<Self as ::core::ops::Index<I>>::Output {
+                let arr = ::core::borrow::Borrow::<[T; #num_columns]>::borrow(self);
+                <[T] as ::core::ops::Index<I>>::index(arr, index)
+            }
+        }
+
+        impl<T, I> ::core::ops::IndexMut<I> for #name<T>
+        where
+            T: ::core::marker::Copy,
+            [T]: ::core::ops::IndexMut<I>,
+        {
+            fn index_mut(&mut self, index: I) -> &mut <Self as ::core::ops::Index<I>>::Output {
+                let arr = ::core::borrow::BorrowMut::<[T; #num_columns]>::borrow_mut(self);
+                <[T] as ::core::ops::IndexMut<I>>::index_mut(arr, index)
+            }
+        }
+
+        impl<T> ::core::default::Default for #name<T>
+        where
+            T: ::core::marker::Copy + ::core::default::Default,
+        {
+            fn default() -> Self {
+                ::core::convert::Into::<Self>::into(
+                    [<T as ::core::default::Default>::default(); #num_columns]
+                )
+            }
+        }
+    })
+}

derive/src/impls/deref_columns.rs

@@ -0,0 +1,45 @@
+use quote::quote;
+use syn::{Data, DeriveInput, Result};
+
+use crate::common::{ensure, is_repr_c};
+
+/// Implements `Deref` and `DerefMut`.
+pub(crate) fn try_derive(ast: DeriveInput) -> Result<proc_macro2::TokenStream> {
+    let is_struct = matches!(ast.data, Data::Struct(_));
+    ensure!(is_struct, &ast, "expected `struct`");
+
+    // Check that the struct is `#[repr(C)]`
+    let repr_c = is_repr_c(&ast.attrs);
+    ensure!(repr_c, &ast, "column struct must be `#[repr(C)]`");
+
+    // The name of the struct.
+    let name = &ast.ident;
+
+    // Safety: `u8` is guaranteed to have a `size_of` of 1.
+    // https://doc.rust-lang.org/reference/type-layout.html#primitive-data-layout
+    let num_columns = quote!(::core::mem::size_of::<#name<u8>>());
+
+    // Safety:
+    // A repr(C) struct generic over T has the same layout as an array [T; N] if:
+    // - Every field of the struct is either T or a type with the same alignment as
+    //   T and a size of `size_of::<T>() * M` where M <= N. i.e. every field is one
+    //   of T, [T; M], or a type with the same layout as [T; M].
+    // - The total number of elements of type T is N.
+    // https://doc.rust-lang.org/reference/type-layout.html#reprc-structs
+    // https://doc.rust-lang.org/reference/type-layout.html#array-layout
+    Ok(quote! {
+        impl<T: ::core::marker::Copy> ::core::ops::Deref for #name<T> {
+            type Target = [T; #num_columns];
+
+            fn deref(&self) -> &<Self as ::core::ops::Deref>::Target {
+                unsafe { ::core::mem::transmute(self) }
+            }
+        }
+
+        impl<T: ::core::marker::Copy> ::core::ops::DerefMut for #name<T> {
+            fn deref_mut(&mut self) -> &mut <Self as ::core::ops::Deref>::Target {
+                unsafe { ::core::mem::transmute(self) }
+            }
+        }
+    })
+}

derive/src/impls/mod.rs

@@ -0,0 +1,2 @@
+pub(crate) mod columns;
+pub(crate) mod deref_columns;

derive/src/lib.rs

@@ -0,0 +1,47 @@
+//! This library provides two convenient derive macros for interpreting arrays
+//! of field elements as structs representing an AIR.
+//!
+//! Deriving [`Columns`] on a struct `Struct<T>` implements the following
+//! conversion traits between `Struct<T>` and arrays `[T; N]` where `N` is the
+//! number of fields in the struct: [`Borrow`], [`BorrowMut`], and [`From`].
+//! Additionally, the traits [`Index`], [`IndexMut`], and [`Default`] are
+//! implemented for `Struct<T>`.
+//!
+//! Deriving [`DerefColumns`] for a struct generic over `T` implements [`Deref`]
+//! and [`DerefMut`] with target `[T; N]` where `N` is the number of fields in
+//! the struct.
+//!
+//! These implementations employ unsafe code and as such place a burden on the
+//! user to ensure their safe usage. Please see the respective macro
+//! implementations to understand the conditions that should be upheld by any
+//! struct deriving [`Columns`] or [`DerefColumns`]. In short, the struct must
+//! be `#[repr(C)]` and all fields must be one of `T`, `[T; M]`, or a type with
+//! the same layout as `[T; M]`.
+//!
+//! [`Borrow`]: ::core::borrow::Borrow
+//! [`BorrowMut`]: ::core::borrow::BorrowMut
+//! [`Index`]: ::core::ops::Index
+//! [`IndexMut`]: ::core::ops::IndexMut
+//! [`Deref`]: ::core::ops::Deref
+//! [`DerefMut`]: ::core::ops::DerefMut
+
+pub(crate) mod common;
+mod impls;
+
+use impls::{columns, deref_columns};
+
+#[proc_macro_derive(Columns)]
+pub fn derive_columns(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    let ast = syn::parse_macro_input!(input as syn::DeriveInput);
+    columns::try_derive(ast)
+        .unwrap_or_else(syn::Error::into_compile_error)
+        .into()
+}
+
+#[proc_macro_derive(DerefColumns)]
+pub fn derive_deref_columns(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    let ast = syn::parse_macro_input!(input as syn::DeriveInput);
+    deref_columns::try_derive(ast)
+        .unwrap_or_else(syn::Error::into_compile_error)
+        .into()
+}

evm_arithmetization/Cargo.toml

@@ -45,6 +45,7 @@ serde_json = { workspace = true }
 
 # Local dependencies
 mpt_trie = { workspace = true }
+zk_evm_derive = { workspace = true }
 
 [dev-dependencies]
 criterion = { workspace = true }