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
@@ -26,77 +24,111 @@ use super::{
 ///
 /// Note that the `Sha256WorkVarsCols` field it is used for different purposes in the two structs.
 #[repr(C)]
-#[derive(Clone, Copy, Debug, AlignedBorrow)]
-pub struct Sha256RoundCols<T> {
-    pub flags: Sha256FlagsCols<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>,
     /// Stores the current state of the working variables
-    pub work_vars: Sha256WorkVarsCols<T>,
-    pub schedule_helper: Sha256MessageHelperCols<T>,
-    pub message_schedule: Sha256MessageScheduleCols<T>,
+    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.
     ///     - on non-last blocks, this is the final hash of the current block
     ///     - on last blocks, this is the initial state constants, SHA256_H.
     /// The work variables constraints are applied on all rows, so `carry_a` and `carry_e`
     /// must be filled in with dummy values to ensure these constraints hold.
-    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> {
-    /// The message schedule words as 32-bit integers
-    /// The first 16 words will be the message data
-    pub w: [[T; SHA256_WORD_BITS]; SHA256_ROUNDS_PER_ROW],
+#[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 C::WORD_BITS-bit integers
+    /// The first 16 rows will be the message data
+    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 2 bit numbers represented using 2 cells as individual bits
-    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` (message schedule word) from 3 rounds ago
     /// In general, `w_i` means `w` from `i` 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
     /// This is needed to constrain the message schedule, since we can only constrain on two rows at a time
-    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> {
-    /// A flag that indicates if the current row is among the first 16 rows of a block.
+#[derive(Clone, Copy, Debug, ColsRef)]
+pub struct ShaFlagsCols<T, const ROW_VAR_CNT: usize> {
+    /// A flag that indicates if the current row is among the first C::ROUND_ROWS rows of a block.
     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,
@@ -106,7 +138,8 @@ pub struct Sha256FlagsCols<T> {
     // This flag is only used in digest rows.
     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 index of the current block in the trace starting at 1.
     /// Set to 0 on padding rows.
     pub global_block_idx: T,
@@ -116,7 +149,9 @@ pub struct Sha256FlagsCols<T> {
     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>
+{
     // This refers to the padding rows that are added to the air to make the trace length a power of 2.
     // Not to be confused with the padding added to messages as part of the SHA hash function.
     pub fn is_not_padding_row(&self) -> O {
@@ -132,3 +167,19 @@ impl<O, T: Copy + core::ops::Add<Output = O>> Sha256FlagsCols<T> {
         not(self.is_not_padding_row())
     }
 }
+
+// We need to implement this for the ColsRef type as well
+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
+    }
+
+    // This refers to the padding rows that are added to the air to make the trace length a power of 2.
+    // Not to be confused with the padding added to messages as part of the SHA hash function.
+    pub fn is_padding_row(&self) -> O
+    where
+        O: FieldAlgebra,
+    {
+        not(self.is_not_padding_row())
+    }
+}

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

@@ -30,12 +30,14 @@ pub trait ShaConfig: Send + Sync + Clone {
     const BLOCK_BITS: usize = Self::BLOCK_WORDS * Self::WORD_BITS;
     /// Number of rows per block
     const ROWS_PER_BLOCK: usize;
-    /// Number of rounds per row
+    /// Number of rounds per row. Must divide Self::ROUNDS_PER_BLOCK
     const ROUNDS_PER_ROW: usize;
     /// Number of rounds per row minus one (needed for one of the column structs)
     const ROUNDS_PER_ROW_MINUS_ONE: usize = Self::ROUNDS_PER_ROW - 1;
-    /// Number of rounds per block
+    /// Number of rounds per block. Must be a multiple of Self::ROUNDS_PER_ROW
     const ROUNDS_PER_BLOCK: usize;
+    /// Number of rows used to constrain rounds
+    const ROUND_ROWS: usize = Self::ROUNDS_PER_BLOCK / Self::ROUNDS_PER_ROW;
     /// Number of words in a SHA hash
     const HASH_WORDS: usize;
     /// Number of vars needed to encode the row index with [Encoder]

crates/circuits/sha256-air/src/config.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::*;