Generalize ColsRef macro

- allow using traits other than ShaConfig for const parameters
- support fields types that derive AlignedBorrow (not recursed into)
- support nested-array-type fields with literal array lengths
- support nested-array-type fields with constant array lengths

Author: Avaneesh-axiom

crates/circuits/sha-macros/Cargo.toml

@@ -14,7 +14,8 @@ proc-macro2 = "1.0"
 
 [dev-dependencies]
 openvm-sha-air = { workspace = true }
-ndarray = "0.16"
+openvm-circuit-primitives-derive = { workspace = true }
+ndarray.workspace = true
 
 [lib]
 proc-macro = true

crates/circuits/sha-macros/src/lib.rs

@@ -0,0 +1,19 @@
+use openvm_sha_air::{Sha256Config, ShaConfig};
+use openvm_sha_macros::ColsRef;
+
+#[derive(ColsRef)]
+#[config(ShaConfig)]
+struct ArrayTest<T> {
+    a: T,
+    b: [T; 4],
+    c: [[T; 4]; 4],
+}
+
+#[test]
+fn arrays() {
+    let input = [1; 1 + 4 + 4 * 4];
+    let test: ArrayTestRef<u32> = ArrayTestRef::from::<Sha256Config>(&input);
+    println!("{}", test.a);
+    println!("{}", test.b);
+    println!("{}", test.c);
+}

crates/circuits/sha-macros/tests/arrays.rs

@@ -0,0 +1,27 @@
+use openvm_sha_macros::ColsRef;
+
+const ONE: usize = 1;
+const TWO: usize = 2;
+const THREE: usize = 3;
+
+mod test_config;
+use test_config::{TestConfig, TestConfigImpl};
+
+#[derive(ColsRef)]
+#[config(TestConfig)]
+struct ConstLenArrayTest<T, const N: usize> {
+    a: T,
+    b: [T; N],
+    c: [[T; ONE]; TWO],
+    d: [[[T; ONE]; TWO]; THREE],
+}
+
+#[test]
+fn const_len_arrays() {
+    let input = [1; 1 + TestConfigImpl::N * 2 + 1 * 2 * 3];
+    let test: ConstLenArrayTestRef<u32> = ConstLenArrayTestRef::from::<TestConfigImpl>(&input);
+    println!("{}", test.a);
+    println!("{}", test.b);
+    println!("{}", test.c);
+    println!("{}", test.d);
+}

crates/circuits/sha-macros/tests/const-len-arrays.rs

@@ -3,6 +3,7 @@ use openvm_sha_macros::ColsRef;
 
 #[repr(C)]
 #[derive(Clone, Copy, Debug, ColsRef)]
+#[config(ShaConfig)]
 pub struct ShaFlagsCols<T, const ROW_VAR_CNT: usize> {
     pub is_round_row: T,
     /// A flag that indicates if the current row is among the first 4 rows of a block

crates/circuits/sha-macros/tests/flags.rs

@@ -2,12 +2,14 @@ use openvm_sha_air::{Sha256Config, ShaConfig};
 use openvm_sha_macros::ColsRef;
 
 #[derive(ColsRef)]
+#[config(ShaConfig)]
 struct Test1Cols<T, const WORD_BITS: usize> {
     pub a: T,
     pub nested: Test2Cols<T, WORD_BITS>,
 }
 
 #[derive(ColsRef)]
+#[config(ShaConfig)]
 struct Test2Cols<T, const WORD_BITS: usize> {
     pub b: T,
     pub c: [T; WORD_BITS],

crates/circuits/sha-macros/tests/nested.rs

@@ -0,0 +1,53 @@
+use openvm_circuit_primitives_derive::AlignedBorrow;
+use openvm_sha_macros::ColsRef;
+
+mod test_config;
+use test_config::{TestConfig, TestConfigImpl};
+
+#[derive(ColsRef)]
+#[config(TestConfig)]
+struct TestCols<T, const N: usize> {
+    a: [T; N],
+    // Forces the field to be treated as a struct that derives AlignedBorrow.
+    // In particular, ignores the fact that it ends with `Cols` and doesn't
+    // expect a `PlainTestColsRef` type.
+    #[plain]
+    b: PlainCols<T>,
+}
+
+#[derive(Clone, Copy, Debug, AlignedBorrow)]
+struct PlainCols<T> {
+    a: T,
+    b: [T; 4],
+}
+
+#[test]
+fn plain() {
+    let input = [1; TestConfigImpl::N + 1 + 4];
+    let test: TestColsRef<u32> = TestColsRef::from::<TestConfigImpl>(&input);
+    println!("{}", test.a);
+    println!("{:?}", test.b);
+}
+
+#[test]
+fn plain_mut() {
+    let mut input = [1; TestConfigImpl::N + 1 + 4];
+    let mut test: TestColsRefMut<u32> = TestColsRefMut::from::<TestConfigImpl>(&mut input);
+    test.a[0] = 1;
+    test.b.a = 1;
+    test.b.b[0] = 1;
+    println!("{}", test.a);
+    println!("{:?}", test.b);
+}
+
+#[test]
+fn plain_from_mut() {
+    let mut input = [1; TestConfigImpl::N + 1 + 4];
+    let mut test: TestColsRefMut<u32> = TestColsRefMut::from::<TestConfigImpl>(&mut input);
+    test.a[0] = 1;
+    test.b.a = 1;
+    test.b.b[0] = 1;
+    let test2: TestColsRef<u32> = TestColsRef::from_mut::<TestConfigImpl>(&mut test);
+    println!("{}", test2.a);
+    println!("{:?}", test2.b);
+}

crates/circuits/sha-macros/tests/plain.rs

@@ -2,6 +2,7 @@ use openvm_sha_air::{Sha256Config, ShaConfig};
 use openvm_sha_macros::ColsRef;
 
 #[derive(ColsRef)]
+#[config(ShaConfig)]
 struct Test<T, const WORD_BITS: usize, const ROUNDS_PER_ROW: usize, const WORD_U16S: usize> {
     a: T,
     b: [T; WORD_BITS],

crates/circuits/sha-macros/tests/simple.rs

@@ -0,0 +1,7 @@
+pub trait TestConfig {
+    const N: usize;
+}
+pub struct TestConfigImpl;
+impl TestConfig for TestConfigImpl {
+    const N: usize = 4;
+}

crates/circuits/sha-macros/tests/test_config.rs

@@ -3,6 +3,7 @@ use openvm_sha_macros::ColsRef;
 
 #[repr(C)]
 #[derive(Clone, Copy, Debug, ColsRef)]
+#[config(ShaConfig)]
 pub struct ShaWorkVarsCols<
     T,
     const WORD_BITS: usize,