Updated sha256 air to use sha-macros (ColsRef structs)

Author: Avaneesh-axiom

crates/circuits/sha256-air/Cargo.toml

@@ -9,6 +9,8 @@ openvm-circuit-primitives = { workspace = true }
 openvm-stark-backend = { workspace = true }
 sha2 = { version = "0.10", features = ["compress"] }
 rand.workspace = true
+openvm-sha-macros = { workspace = true }
+ndarray = "0.16"
 
 [dev-dependencies]
 openvm-stark-sdk = { workspace = true }

crates/circuits/sha256-air/src/air.rs

@@ -1,12 +1,10 @@
 //! WARNING: the order of fields in the structs is important, do not change it
 
-use openvm_circuit_primitives::{utils::not, AlignedBorrow};
+use openvm_circuit_primitives::utils::not;
+use openvm_sha_macros::ColsRef;
 use openvm_stark_backend::p3_field::FieldAlgebra;
 
-use super::{
-    SHA256_HASH_WORDS, SHA256_ROUNDS_PER_ROW, SHA256_ROW_VAR_CNT, SHA256_WORD_BITS,
-    SHA256_WORD_U16S, SHA256_WORD_U8S,
-};
+use crate::ShaConfig;
 
 /// In each SHA256 block:
 /// - First 16 rows use Sha256RoundCols
@@ -21,82 +19,119 @@ use super::{
 /// 1. Common constraints to work on either struct type by accessing these shared fields
 /// 2. Specific constraints to use the appropriate struct, with flags helping to do conditional constraints
 #[repr(C)]
-#[derive(Clone, Copy, Debug, AlignedBorrow)]
-pub struct Sha256RoundCols<T> {
-    pub flags: Sha256FlagsCols<T>,
-    pub work_vars: Sha256WorkVarsCols<T>,
-    pub schedule_helper: Sha256MessageHelperCols<T>,
-    pub message_schedule: Sha256MessageScheduleCols<T>,
+#[derive(Clone, Copy, Debug, ColsRef)]
+pub struct ShaRoundCols<
+    T,
+    const WORD_BITS: usize,
+    const WORD_U8S: usize,
+    const WORD_U16S: usize,
+    const ROUNDS_PER_ROW: usize,
+    const ROUNDS_PER_ROW_MINUS_ONE: usize,
+    const ROW_VAR_CNT: usize,
+> {
+    pub flags: ShaFlagsCols<T, ROW_VAR_CNT>,
+    pub work_vars: ShaWorkVarsCols<T, WORD_BITS, ROUNDS_PER_ROW, WORD_U16S>,
+    pub schedule_helper:
+        ShaMessageHelperCols<T, WORD_U16S, ROUNDS_PER_ROW, ROUNDS_PER_ROW_MINUS_ONE>,
+    pub message_schedule: ShaMessageScheduleCols<T, WORD_BITS, ROUNDS_PER_ROW, WORD_U8S>,
 }
 
 #[repr(C)]
-#[derive(Clone, Copy, Debug, AlignedBorrow)]
-pub struct Sha256DigestCols<T> {
-    pub flags: Sha256FlagsCols<T>,
+#[derive(Clone, Copy, Debug, ColsRef)]
+pub struct ShaDigestCols<
+    T,
+    const WORD_BITS: usize,
+    const WORD_U8S: usize,
+    const WORD_U16S: usize,
+    const HASH_WORDS: usize,
+    const ROUNDS_PER_ROW: usize,
+    const ROUNDS_PER_ROW_MINUS_ONE: usize,
+    const ROW_VAR_CNT: usize,
+> {
+    pub flags: ShaFlagsCols<T, ROW_VAR_CNT>,
     /// Will serve as previous hash values for the next block
-    pub hash: Sha256WorkVarsCols<T>,
-    pub schedule_helper: Sha256MessageHelperCols<T>,
+    pub hash: ShaWorkVarsCols<T, WORD_BITS, ROUNDS_PER_ROW, WORD_U16S>,
+    pub schedule_helper:
+        ShaMessageHelperCols<T, WORD_U16S, ROUNDS_PER_ROW, ROUNDS_PER_ROW_MINUS_ONE>,
     /// The actual final hash values of the given block
     /// Note: the above `hash` will be equal to `final_hash` unless we are on the last block
-    pub final_hash: [[T; SHA256_WORD_U8S]; SHA256_HASH_WORDS],
+    pub final_hash: [[T; WORD_U8S]; HASH_WORDS],
     /// The final hash of the previous block
     /// Note: will be constrained using interactions with the chip itself
-    pub prev_hash: [[T; SHA256_WORD_U16S]; SHA256_HASH_WORDS],
+    pub prev_hash: [[T; WORD_U16S]; HASH_WORDS],
 }
 
 #[repr(C)]
-#[derive(Clone, Copy, Debug, AlignedBorrow)]
-pub struct Sha256MessageScheduleCols<T> {
+#[derive(Clone, Copy, Debug, ColsRef)]
+pub struct ShaMessageScheduleCols<
+    T,
+    const WORD_BITS: usize,
+    const ROUNDS_PER_ROW: usize,
+    const WORD_U8S: usize,
+> {
     /// The message schedule words as 32-bit intergers
-    pub w: [[T; SHA256_WORD_BITS]; SHA256_ROUNDS_PER_ROW],
+    pub w: [[T; WORD_BITS]; ROUNDS_PER_ROW],
     /// Will be message schedule carries for rows 4..16 and a buffer for rows 0..4 to be used freely by wrapper chips
     /// Note: carries are represented as 2 bit numbers
-    pub carry_or_buffer: [[T; SHA256_WORD_U8S]; SHA256_ROUNDS_PER_ROW],
+    pub carry_or_buffer: [[T; WORD_U8S]; ROUNDS_PER_ROW],
 }
 
 #[repr(C)]
-#[derive(Clone, Copy, Debug, AlignedBorrow)]
-pub struct Sha256WorkVarsCols<T> {
+#[derive(Clone, Copy, Debug, ColsRef)]
+pub struct ShaWorkVarsCols<
+    T,
+    const WORD_BITS: usize,
+    const ROUNDS_PER_ROW: usize,
+    const WORD_U16S: usize,
+> {
     /// `a` and `e` after each iteration as 32-bits
-    pub a: [[T; SHA256_WORD_BITS]; SHA256_ROUNDS_PER_ROW],
-    pub e: [[T; SHA256_WORD_BITS]; SHA256_ROUNDS_PER_ROW],
+    pub a: [[T; WORD_BITS]; ROUNDS_PER_ROW],
+    pub e: [[T; WORD_BITS]; ROUNDS_PER_ROW],
     /// The carry's used for addition during each iteration when computing `a` and `e`
-    pub carry_a: [[T; SHA256_WORD_U16S]; SHA256_ROUNDS_PER_ROW],
-    pub carry_e: [[T; SHA256_WORD_U16S]; SHA256_ROUNDS_PER_ROW],
+    pub carry_a: [[T; WORD_U16S]; ROUNDS_PER_ROW],
+    pub carry_e: [[T; WORD_U16S]; ROUNDS_PER_ROW],
 }
 
 /// These are the columns that are used to help with the message schedule additions
 /// Note: these need to be correctly assigned for every row even on padding rows
 #[repr(C)]
-#[derive(Clone, Copy, Debug, AlignedBorrow)]
-pub struct Sha256MessageHelperCols<T> {
+#[derive(Clone, Copy, Debug, ColsRef)]
+pub struct ShaMessageHelperCols<
+    T,
+    const WORD_U16S: usize,
+    const ROUNDS_PER_ROW: usize,
+    const ROUNDS_PER_ROW_MINUS_ONE: usize,
+> {
     /// The following are used to move data forward to constrain the message schedule additions
     /// The value of `w` from 3 rounds ago
-    pub w_3: [[T; SHA256_WORD_U16S]; SHA256_ROUNDS_PER_ROW - 1],
+    pub w_3: [[T; WORD_U16S]; ROUNDS_PER_ROW_MINUS_ONE],
     /// Here intermediate(i) =  w_i + sig_0(w_{i+1})
     /// Intermed_t represents the intermediate t rounds ago
-    pub intermed_4: [[T; SHA256_WORD_U16S]; SHA256_ROUNDS_PER_ROW],
-    pub intermed_8: [[T; SHA256_WORD_U16S]; SHA256_ROUNDS_PER_ROW],
-    pub intermed_12: [[T; SHA256_WORD_U16S]; SHA256_ROUNDS_PER_ROW],
+    pub intermed_4: [[T; WORD_U16S]; ROUNDS_PER_ROW],
+    pub intermed_8: [[T; WORD_U16S]; ROUNDS_PER_ROW],
+    pub intermed_12: [[T; WORD_U16S]; ROUNDS_PER_ROW],
 }
 
 #[repr(C)]
-#[derive(Clone, Copy, Debug, AlignedBorrow)]
-pub struct Sha256FlagsCols<T> {
+#[derive(Clone, Copy, Debug, ColsRef)]
+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
     pub is_first_4_rows: T,
     pub is_digest_row: T,
     pub is_last_block: T,
     /// We will encode the row index [0..17) using 5 cells
-    pub row_idx: [T; SHA256_ROW_VAR_CNT],
+    //#[length(ROW_VAR_CNT)]
+    pub row_idx: [T; ROW_VAR_CNT],
     /// The global index of the current block
     pub global_block_idx: T,
     /// Will store the index of the current block in the current message starting from 0
     pub local_block_idx: T,
 }
 
-impl<O, T: Copy + core::ops::Add<Output = O>> Sha256FlagsCols<T> {
+impl<O, T: Copy + core::ops::Add<Output = O>, const ROW_VAR_CNT: usize>
+    ShaFlagsCols<T, ROW_VAR_CNT>
+{
     pub fn is_not_padding_row(&self) -> O {
         self.is_round_row + self.is_digest_row
     }
@@ -108,3 +143,16 @@ impl<O, T: Copy + core::ops::Add<Output = O>> Sha256FlagsCols<T> {
         not(self.is_not_padding_row())
     }
 }
+
+impl<'a, O, T: Copy + core::ops::Add<Output = O>> ShaFlagsColsRef<'a, T> {
+    pub fn is_not_padding_row(&self) -> O {
+        *self.is_round_row + *self.is_digest_row
+    }
+
+    pub fn is_padding_row(&self) -> O
+    where
+        O: FieldAlgebra,
+    {
+        not(self.is_not_padding_row())
+    }
+}

crates/circuits/sha256-air/src/columns.rs

@@ -3,11 +3,13 @@
 
 mod air;
 mod columns;
+mod config;
 mod trace;
 mod utils;
 
 pub use air::*;
 pub use columns::*;
+pub use config::*;
 pub use trace::*;
 pub use utils::*;
 

crates/circuits/sha256-air/src/lib.rs

@@ -12,46 +12,57 @@ use openvm_stark_backend::{
     config::{StarkGenericConfig, Val},
     interaction::InteractionBuilder,
     p3_air::{Air, BaseAir},
-    p3_field::{Field, PrimeField32},
+    p3_field::{Field, FieldAlgebra, PrimeField32},
     prover::types::AirProofInput,
     rap::{get_air_name, BaseAirWithPublicValues, PartitionedBaseAir},
     AirRef, Chip, ChipUsageGetter,
 };
-use openvm_stark_sdk::utils::create_seeded_rng;
+use openvm_stark_sdk::{p3_baby_bear::BabyBear, utils::create_seeded_rng};
 use rand::Rng;
+use sha2::Sha256;
 
 use crate::{
-    Sha256Air, SHA256_BLOCK_U8S, SHA256_DIGEST_WIDTH, SHA256_ROUND_WIDTH, SHA256_ROWS_PER_BLOCK,
+    Sha256Config, Sha512Config, ShaAir, ShaConfig, ShaFlagsColsRef, ShaFlagsColsRefMut,
+    ShaPrecomputedValues,
 };
 
 // A wrapper AIR purely for testing purposes
 #[derive(Clone, Debug)]
-pub struct Sha256TestAir {
-    pub sub_air: Sha256Air,
+pub struct ShaTestAir<C: ShaConfig + ShaPrecomputedValues<C::Word>> {
+    pub sub_air: ShaAir<C>,
 }
 
-impl<F: Field> BaseAirWithPublicValues<F> for Sha256TestAir {}
-impl<F: Field> PartitionedBaseAir<F> for Sha256TestAir {}
-impl<F: Field> BaseAir<F> for Sha256TestAir {
+impl<F: Field, C: ShaConfig + ShaPrecomputedValues<C::Word>> BaseAirWithPublicValues<F>
+    for ShaTestAir<C>
+{
+}
+impl<F: Field, C: ShaConfig + ShaPrecomputedValues<C::Word>> PartitionedBaseAir<F>
+    for ShaTestAir<C>
+{
+}
+impl<F: Field, C: ShaConfig + ShaPrecomputedValues<C::Word>> BaseAir<F> for ShaTestAir<C> {
     fn width(&self) -> usize {
-        <Sha256Air as BaseAir<F>>::width(&self.sub_air)
+        <ShaAir<C> as BaseAir<F>>::width(&self.sub_air)
     }
 }
 
-impl<AB: InteractionBuilder> Air<AB> for Sha256TestAir {
+impl<AB: InteractionBuilder, C: ShaConfig + ShaPrecomputedValues<C::Word>> Air<AB>
+    for ShaTestAir<C>
+{
     fn eval(&self, builder: &mut AB) {
         self.sub_air.eval(builder, 0);
     }
 }
 
 // A wrapper Chip purely for testing purposes
-pub struct Sha256TestChip {
-    pub air: Sha256TestAir,
+pub struct ShaTestChip<C: ShaConfig + ShaPrecomputedValues<C::Word>> {
+    pub air: ShaTestAir<C>,
     pub bitwise_lookup_chip: SharedBitwiseOperationLookupChip<8>,
-    pub records: Vec<([u8; SHA256_BLOCK_U8S], bool)>,
+    pub records: Vec<(Vec<u8>, bool)>, // length of inner vec is BLOCK_U8S
 }
 
-impl<SC: StarkGenericConfig> Chip<SC> for Sha256TestChip
+impl<SC: StarkGenericConfig, C: ShaConfig + ShaPrecomputedValues<C::Word> + 'static> Chip<SC>
+    for ShaTestChip<C>
 where
     Val<SC>: PrimeField32,
 {
@@ -60,7 +71,7 @@ where
     }
 
     fn generate_air_proof_input(self) -> AirProofInput<SC> {
-        let trace = crate::generate_trace::<Val<SC>>(
+        let trace = crate::generate_trace::<Val<SC>, C>(
             &self.air.sub_air,
             self.bitwise_lookup_chip.clone(),
             self.records,
@@ -69,33 +80,39 @@ where
     }
 }
 
-impl ChipUsageGetter for Sha256TestChip {
+impl<C: ShaConfig + ShaPrecomputedValues<C::Word>> ChipUsageGetter for ShaTestChip<C> {
     fn air_name(&self) -> String {
         get_air_name(&self.air)
     }
     fn current_trace_height(&self) -> usize {
-        self.records.len() * SHA256_ROWS_PER_BLOCK
+        self.records.len() * C::ROWS_PER_BLOCK
     }
 
     fn trace_width(&self) -> usize {
-        max(SHA256_ROUND_WIDTH, SHA256_DIGEST_WIDTH)
+        max(C::ROUND_WIDTH, C::DIGEST_WIDTH)
     }
 }
 
 const SELF_BUS_IDX: usize = 28;
-#[test]
-fn rand_sha256_test() {
+fn rand_sha_test<C: ShaConfig + ShaPrecomputedValues<C::Word> + 'static>() {
     let mut rng = create_seeded_rng();
     let tester = VmChipTestBuilder::default();
     let bitwise_bus = BitwiseOperationLookupBus::new(BITWISE_OP_LOOKUP_BUS);
     let bitwise_chip = SharedBitwiseOperationLookupChip::<RV32_CELL_BITS>::new(bitwise_bus);
     let len = rng.gen_range(1..100);
     let random_records: Vec<_> = (0..len)
-        .map(|_| (array::from_fn(|_| rng.gen::<u8>()), true))
+        .map(|_| {
+            (
+                (0..C::BLOCK_U8S)
+                    .map(|_| rng.gen::<u8>())
+                    .collect::<Vec<_>>(),
+                true,
+            )
+        })
         .collect();
-    let chip = Sha256TestChip {
-        air: Sha256TestAir {
-            sub_air: Sha256Air::new(bitwise_bus, SELF_BUS_IDX),
+    let chip = ShaTestChip {
+        air: ShaTestAir {
+            sub_air: ShaAir::<C>::new(bitwise_bus, SELF_BUS_IDX),
         },
         bitwise_lookup_chip: bitwise_chip.clone(),
         records: random_records,
@@ -104,3 +121,8 @@ fn rand_sha256_test() {
     let tester = tester.build().load(chip).load(bitwise_chip).finalize();
     tester.simple_test().expect("Verification failed");
 }
+
+#[test]
+fn rand_sha256_test() {
+    rand_sha_test::<Sha256Config>();
+}