Add Columns and DerefColumns derive macros

[gio256]

Adds two derive macros, Columns and DerefColumns. The expectation is that column view structs will derive one or both of these in order to cut down on the boilerplate for transmuting to/from arrays.

• Deriving Columns implements Borrow, BorrowMut, and From in both directions as well as Index, IndexMut and Default.

• Deriving DerefColumns implements Deref and DerefMut. This was made a separate macro as an extra opt-in step to avoid any confusion resulting from deref coercions.

Because this hides unsafe code in a proc macro, the primary goal was auditability. But, we could also trade a bit more complexity in the macros for more flexibility in their usage. This might be useful, for example, if any column structs were expected to use associated constants or const generics in the future as in SP1.

This PR relates to discussion in #312 around implementing column view structs for starks that currently use const indices into their columns. See also #311 for some discussion of whether adding a separate sub-crate for proc macros is actually desirable.

[0xaatif]

The derive code is fine, but I'm worried this is the wrong direction - is there any reason we can't use zerocopy for most of this functionality?

[0xaatif]

nit: I'd rather this was called zk-evm-proc-macro

[gio256]

Happy to call this anything you want, but would underscores be more consistent with the rest of the package names in the workspace?

Also, would you suggest renaming the directory from derive/ to proc_macro/ in this case?

[atanmarko]

@gio256 Naming it zk_evm_proc_macro to be consistent and renaming derive to proc_macro folder would be great I think. If we would need another project wide proc macro we would then add it here.

[atanmarko]

Maybe we should put derive macro in some subfolder, e.g. common or macros. I feel we will overcrowd the root folder.

[atanmarko]

Would be nice to add some integration test (in the derive/tests directory) with some dummy structs to regression test basic usecases and demonstrate what macros could do. This is generic macro, it may be used elsewhere.

[0xaatif]

If you do, use trybuild

[atanmarko]

The derive code is fine, but I'm worried this is the wrong direction - is there any reason we can't use zerocopy for most of this functionality?

@0xaatif Zerocopy only works with byte array, right? Here you index and manipulate array of generic type (could be complex, just needs to be Copy) (e.g. OpsColumnsView<Option<u32>>)

[gio256]

is there any reason we can't use zerocopy for most of this functionality?

@0xaatif I'm not sure if zerocopy would be able to help here or not. I suspect we might run into the same issue as you would if you tried to impl From<[T; N_COLS]> for Struct<T> using the safer transmute rather than transmute_copy:

Any struct with generic fields is going to be dependently-sized, and you can't safely transmute between dependently-sized types. The reason we can do this is that we know (though the macro does not enforce) that every field is effectively T, and therefore the sizes of Struct<T> and [T; N_COLS] depend on T in the exact same way.

If zerocopy isn't able to help here, another option would be to see if we can get more compile-time guarantees out of the macros here using some of the techniques from zerocopy.

[0xaatif]

@gio256, @atanmarko, I could have been more explicit in my suggestion :)

Here's a proof-of-concept of what I mean:

use std::mem;

use zerocopy::{AsBytes, FromBytes, FromZeroes};

#[derive(FromBytes, FromZeroes, PartialEq, Debug)]
#[repr(C)]
pub struct View<T> {
    x: T,
    y: T,
    z: [T; 5],
}

const WIDTH: usize = 7;

pub fn view<T: AsBytes + FromBytes>(arr: [T; WIDTH]) -> View<T> {
    const { assert!(mem::size_of::<View<T>>() == mem::size_of::<[T; WIDTH]>()) };
    match View::read_from(arr.as_bytes()) {
        Some(it) => it,
        None => unreachable!("we statically assert that the width works"),
    }
}

// pub fn view_ref
// pub fn view_mut

#[test]
fn test() {
    assert_eq!(
        view([0, 1, 2, 3, 4, 5, 6]),
        View {
            x: 0,
            y: 1,
            z: [2, 3, 4, 5, 6]
        }
    );
}

• we don't have to write/maintain our own derives
• we don't have to write any unsafe
• we get all the functionality of this PR

AIUI, the Field impls look like this:

pub struct Secp256K1Scalar(pub [u64; 4]);

Which should be fine to #[derive(zerocopy::...)]

You'd have write view and WIDTH for each struct because const generics aren't where they need to be yet, but I think that's fine - we can write a (decl) macro for that if we really want.

Are there any reasons why this won't work?

[atanmarko]

If zerocopy isn't able to help here, another option would be to see if we can get more compile-time guarantees out of the macros here using some of the techniques from zerocopy.

One check of memory byte size equivalence between #name<T> and [T; #num_columns] should be the most simple guarantee? Looking at the struct tokens to check if the the type of all struct members is T should also not be so hard?

[atanmarko]

Here's a proof-of-concept of what I mean

@0xaatif Looks nice. Issue could be that you can not do compile time checks what you can do in the macro, like if structure is with C byte alignment, or if all struct members are of T type.

EDIT: Another question, how would you implement with zerocopy without copying bytes if you want to index 3rd member of the struct and mutate it in place?

EDIT2: OK, I see as_bytes_mut() then cast it to array of T elements then index it

[0xaatif]

I'm satisfied that this isn't the correct approach given our discussion

[gio256]

@0xaatif that looks like a reasonable option to me. It might be worth benchmarking, but from a brief reading of the zerocopy code we would end up with some size/alignment checks followed by a ptr::read similar to using core::mem::transmute.

Some points that are probably obvious or insignificant:

• It would require adding zerobin as a dependency to plonky2_field.
• We would still need transmute for the const make_col_map functions (not a part of the macros currently). This is arguably the least safe of all the uses because it takes place in a const context.
• I would argue that a derive macro would be even better suited if we used zerobin, since we would get all the convenience without the safety concerns.

[gio256]

@0xaatif Reading through the Zerocopy docs (and playing around with a toy example on the side), it looks like we might run into some challenges going in the opposite direction of your example (ColumnsView<T> -> [T; N_COLS]).

I assume that to do this we would need ColumnsView<T> to derive zerocopy::AsBytes.

The analysis done by Zerocopy for deriving AsBytes states that a struct must have no padding (I read this as no inter-field padding). This means Zerocopy only allows a generic struct to derive AsBytes if it is repr(transparent) or repr(packed).

Obviously column view structs have more than one non-zero-sized field and therefore can't be repr(transparent), so everything would have to be repr(packed). Making all column view structs repr(packed) seems like a non-starter, as passing references to fields of a repr(packed) struct is unsafe, and with good reaason.

Because of the way padding is determined for structs, it seems to me that this comes back to the problem that there is no way for Zerocopy to know that all fields of our struct are effectively T.

We could implement AsBytes ourselves, but that seems more unsafe that using transmute directly, because it also requires us to understand the assumptions made by the Zerocopy crate.

[0xaatif]

@0xaatif Reading through the Zerocopy docs (and playing around with a toy example on the side), it looks like we might run into some challenges going in the opposite direction of your example (ColumnsView<T> -> [T; N_COLS]).

Makes sense! Looks like those call sites can't be refactored - how about using zerocopy for the obvious direction, and for the opposite, how do you feel about the following options:

• a custom AsBytes derive, maybe leaning on https://docs.rs/static_assertions/latest/static_assertions/macro.assert_type_eq_all.html ?
• a different trait

[gio256]

@0xaatif Can you explain what you mean by "a different trait"?

I think using Zerocopy in one direction sounds like a good option. My only concern would be that it might be a tough sell to add a plonky2 dependency so that we can cast an array of field elements as bytes.

In the other direction, we could certainly hack something together using Zerocopy and a custom AsBytes. But, it seems like it would be less auditable than the current version and at least as unsafe. It would require understanding the guarantees that Zerocopy makes and guaranteeing that the deriving struct is valid as &[u8] in addition to &[T; NUM_COLS]. (I confess there is a bit of circular logic here, as the hard part in both cases is guaranteeing that the struct contains no padding bytes, which may be uninitialized and therefore UB if read as u*).

Independently of what we decide on using Zerocopy, I think we have the same problem Zerocopy has -- it's hard to make guarantees about the layout of a generic struct. No matter how we go from ColumnsView<T> to [T; NUM_COLS], it will be based on an assumption about the deriving struct that we rely on the user to uphold. In this sense I think there are two safety levers we can pull here -- the implementation (Zerocopy vs. transmute), and the compile-time restrictions on the struct itself.

[0xaatif]

notblocking: is there a reason we test crates individually, rather than just cargo test --workspace?

[atanmarko]

This is technical debt I believe, we should unify them all.

[atanmarko]

New issue #342

[0xaatif]

@0xaatif Can you explain what you mean by "a different trait"?

Defining and implementing a different trait to zerocopy::AsBytes, which does what we need it to with different invariants.

I see where we're stuck r.e uninitialized padding bytes, as I recall zerocopy's authors were waiting from some guarantees from the lang team before making progress on this.

I still feel warm about using zerocopy for the forward direction, but not going backwards with the same approach complicates things...

How big a refactor do you thing changing the callsites will be?
i.e from

let foo: ColumnStruct;
let bar: [FieldTy; WIDTH] = foo.into();

to this

let bar: [FieldTy; WIDTH] = Default::default();
let foo: &mut ColumnStruct = as_mut(&mut bar);

We'd have to be less rusty and pass the &mut in lots of places instead of move-ing, but we would effectively sidestep the issue, with a readability tradeoff.

[gio256]

Refactoring to only convert in one direction seems like it could get messy, but maybe that is the "honest" route here.

Alternatively, the following is a little hacky but could give us some nice guarantees and allow leveraging zerocopy (still just a sketch).

#[repr(C)]
#[derive(Columns)]
pub struct ColumnsView<T> {
    a: T,
    b: [T; 2],
    c: T,
}

Turns into this:

impl<T> ColumnsView<T> {
    // This is just the struct layout algorithm with the loop unrolled.
    // https://doc.rust-lang.org/reference/type-layout.html#reprc-structs
    const fn __has_padding() -> bool {

        const fn __needs_padding<T>(ost: usize) -> bool {
            ost % ::core::mem::align_of::<T>() > 0
        }

        let mut offset = 0;
        if __needs_padding::<T>(offset) {
            return true;
        }
        offset += ::core::mem::size_of::<T>();
        if __needs_padding::<[T; 2]>(offset) {
            return true;
        }
        offset += ::core::mem::size_of::<[T; 2]>();
        if __needs_padding::<T>(offset) {
            return true;
        }
        offset += ::core::mem::size_of::<T>();
        __needs_padding::<Self>(offset)
    }

    // Assert that there is no inter-field padding in the struct layout.
    const fn __assert_valid() {
        assert!(!Self::__has_padding());
    }
}

// This is where this becomes pretty hacky. With `feature[(generic_const_exprs)]`
// this assert could be part of the trait bounds and prevent all misuse. Without
// that, the best I can think of is to hardcode the field and assert that any
// struct deriving `Columns` is safe to use with Goldilocks.
const _: () = ColumnsView::<
    ::plonky2::field::goldilocks_field::GoldilocksField,
>::__assert_valid();

// This should be safe now (I think).
unsafe impl<T> ::zerocopy::AsBytes for ColumnsView<T>
where
    T: ::zerocopy::AsBytes,
    [T; 2]: ::zerocopy::AsBytes,
    T: ::zerocopy::AsBytes,
{
    fn only_derive_is_allowed_to_implement_this_trait() {}
}

// ... other impls

[atanmarko]

Refactoring to only convert in one direction seems like it could get messy, but maybe that is the "honest" route here.

Alternatively, the following is a little hacky but could give us some nice guarantees and allow leveraging zerocopy (still just a sketch).

Plain and simple, we should not take the zerocopy direction. There is really no reason to force it at any cost.

[0xaatif]

let bar: [FieldTy; WIDTH] = Default::default();
let foo: &mut ColumnStruct = as_mut(&mut bar);

We can't write as_mut from here #315 (comment) without AsBytes, so that's ruled out. If we wanted build on zerocopy, we'd need to implement AsBytes as you've suggested here: #315 (comment)

As for which direction you go, I trust your judgement @gio256

[gio256]

Here's my best attempt at a summary of our discussion so far:

Currently, all conversions to/from column views are implemented manually using transmute or transmute_copy. This is only safe because all of the column view structs meet some very specific criteria. The only thing ensuring this criteria is met is manual review (and the fact that violating these criteria would likely create a nonsensical AIR).

This PR proposes moving these manual impls into a derive macro. It does very little to ensure that these impls are safe for the structs that derive them. The focus is on keeping the macro simple and auditable, but column view structs must be defined and reviewed with the same care as before.

It makes sense to also try to make these conversions safer, and there are two different levers we can pull on to do that:

• Prevent any structs in evm_arithmetization from implementing these conversions unless they meet the criteria required to make the conversions safe.
• Let someone else write the unsafe code for us (i.e. use zerocopy). Unfortunately, we currently have at least two cases that zerocopy is unable to handle -- generic structs and generic unions.

My guess is zerocopy won't be able to support either of these cases until generic_const_exprs is stabilized. In the meantime, we could hack together our own compile-time analysis with zerocopy traits and get something that's maybe 90% there for our usecase. But, this would require a significant effort and the end result would be very similar to "just don't make bad column structs". Hopefully we can come up with a middle-ground that has the guardrails we need without overcomplicating things.

My thoughts on the way forward would be:

• Consider this PR approximately as-is (or close it and wait for a safer alternative).
• Open an issue to improve the compile-time analysis and implementation safety of these conversions.
• Carefully review any future column view impls (see StarkColumnsView<T> structs for other STARKs #312).

[0xaatif]

Reflecting on the AsBytes blocker, I think @atanmarko is right that leaning on zerocopy might not be worth it.

I think I'd propose a refactor with a different contract, and keep all the unsafe code out of macros, and away from repeat implementations:

use core::{
    mem::{self, size_of, MaybeUninit},
    ptr,
};
use std::mem::align_of;

/// # Safety
/// - Implementors must be interpetable as an array of `WIDTH` [`Atom`](ArrayLayout::Atom)s.
/// - Implementors must NOT override any of the default items.
///
/// ## Notes on layout
/// - [Size](size_of) is always `>=` [alignment](align_of).
/// - Arrays are laid out with no padding between elements[^1].
/// - `#[repr(C)] struct`s are laid out with each field [aligned](align_of)[^2].
///
/// [^1]: https://github.com/rust-lang/reference/blob/6a0c84af7126dea04c418480216d752385f76493/src/type-layout.md?plain=1#L78
/// [^2]: https://github.com/rust-lang/reference/blob/6a0c84af7126dea04c418480216d752385f76493/src/type-layout.md?plain=1#L205
pub unsafe trait ArrayLayout<const WIDTH: usize>: Sized {
    type Atom;
    // make life easier for callers.
    const WIDTH: usize = WIDTH;
    fn as_array_ref(&self) -> &[Self::Atom; WIDTH] {
        const { assert::<WIDTH, Self::Atom, Self>() }
        // Safety: implementor upholds contract
        unsafe { mem::transmute(self) }
    }
    fn as_array_mut(&mut self) -> &mut [Self::Atom; WIDTH] {
        const { assert::<WIDTH, Self::Atom, Self>() }
        // Safety: implementor upholds contract
        unsafe { mem::transmute(self) }
    }
    fn into_array(self) -> [Self::Atom; WIDTH] {
        const { assert::<WIDTH, Self::Atom, Self>() }
        let mut array = MaybeUninit::<[Self::Atom; WIDTH]>::uninit();
        // Safety: implementor upholds contract
        unsafe { ptr::write(array.as_mut_ptr().cast::<Self>(), self) };
        unsafe { array.assume_init() }
    }
    fn from_array(array: [Self::Atom; WIDTH]) -> Self {
        const { assert::<WIDTH, Self::Atom, Self>() }
        let mut this = MaybeUninit::<Self>::uninit();
        // Safety: implementor upholds contract
        unsafe { ptr::write(this.as_mut_ptr().cast::<[Self::Atom; WIDTH]>(), array) };
        unsafe { this.assume_init() }
    }
    fn from_array_ref(array: &[Self::Atom; WIDTH]) -> &Self {
        const { assert::<WIDTH, Self::Atom, Self>() }
        // Safety: implementor upholds contract
        unsafe { mem::transmute(array) }
    }
    fn from_array_mut(array: &mut [Self::Atom; WIDTH]) -> &mut Self {
        const { assert::<WIDTH, Self::Atom, Self>() }
        // Safety: implementor upholds contract
        unsafe { mem::transmute(array) }
    }
}
unsafe impl<T, const N: usize> ArrayLayout<N> for [T; N] {
    type Atom = T;
}

#[track_caller]
const fn assert<const WIDTH: usize, AtomT, ArrayLayoutT>() {
    if size_of::<[AtomT; WIDTH]>() != size_of::<ArrayLayoutT>() {
        panic!("an `impl ArrayLayout` must have the same size as its underlying array")
    }
    if align_of::<[AtomT; WIDTH]>() != align_of::<ArrayLayoutT>() {
        panic!("an `impl ArrayLayout` must have the same alignment as its underlying array")
    }
}

If we still want a macro, we can have a simple one which does this:

const MORE: usize = 3;
#[repr(C)]
struct View<T> {
    one: T,
    two: [T; 2],
    more: [T; MORE]
}
unsafe impl<T> ArrayLayout<{1 + 2 + MORE}> for View<T> {
    type Atom = T;
}

Which:

• counts the fields/arrays, which we know to be correct.
• checks for #[repr(C)]

The const { assert(...) } do the required checks from there.
The macro might even be overkill at that point.
We can also drop the T: Copy bounds everywhere, I don't think they're right, and I don't like bounds on structs in any case.

[gio256]

I can see the attraction of using a trait given that you can enforce not overriding methods within the crate. Just a few notes:

• I didn't realize the #[track_caller] attribute works outside the std lib!
• Because of zero-sized types, the size of a type is not necessarily >= to its alignment.
• The const { assert::<..>() } assertions look like they rely on post-monomorphization errors, meaning they will panic on cargo build but not on cargo check (which is still better than runtime in my opinion). If you're willing to rely on post-monomorphization errors, you could also do something similar in the AsBytes impl in the example above.
• I'm not sure the macro as you've specified it is all that simple. syn::Type is a bit of a labyrinth, and you also have to handle structs, unions, and [MemoryChannelView<T>; N]. Additionally, I don't think there is a foolproof syntactic way to determine the size of a type.
• I believe the T: Copy bounds have proliferated because union fields must be Copy (or wrapped in ManuallyDrop<T>).

[0xaatif]

you also have to handle structs, unions, and [MemoryChannelView; N]. Additionally, I don't think there is a foolproof syntactic way to determine the size of a type.

Argh! I was using the keccak as my reference, but now that I see CpuColumnsView, I think you're right - a derive macro wouldn't really help (unless you gave it #[hint]s...)

[gio256]

@0xaatif I hadn't thought about requiring annotations on the struct fields. Perhaps there is something useful there.

I will leave this PR open for now, but please feel free to close and move discussion elsewhere if you feel confident that another route is better here.

[0xaatif]

I'm happy to merge

[atanmarko]

I think that this PR is clear and good improvement of the existing code, without introduction of the breaking changes. We should discuss further alternatives in the other issue tickets.

Just one thing @gio256 - could you rename derive folder and package as discussed above? I think proc_macro folder and zk_evm_proc_macro package would be great.

P.S. @muursh or @Nashtare need to approve this PR to be merged.

[Nashtare]

Trusting your judgement on this!

[atanmarko]

@gio256 Thank you for all your effort!

[gio256]

Thanks for the careful review and brainstorming @atanmarko @0xaatif!