Update SHA-256 subair to support SHA-512 and SHA_384

Author: Avaneesh-axiom

crates/circuits/sha2-air/Cargo.toml

@@ -1,14 +1,16 @@
 [package]
-name = "openvm-sha256-air"
+name = "openvm-sha2-air"
 version.workspace = true
 authors.workspace = true
 edition.workspace = true
 
 [dependencies]
 openvm-circuit-primitives = { workspace = true }
 openvm-stark-backend = { workspace = true }
+openvm-circuit-primitives-derive = { workspace = true }
 sha2 = { version = "0.10", features = ["compress"] }
 rand.workspace = true
+ndarray.workspace = true
 
 [dev-dependencies]
 openvm-stark-sdk = { workspace = true }

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

@@ -1,130 +1,159 @@
 //! 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_circuit_primitives_derive::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::Sha2Config;
 
-/// In each SHA256 block:
-/// - First 16 rows use Sha256RoundCols
-/// - Final row uses Sha256DigestCols
+/// In each SHA block:
+/// - First C::ROUND_ROWS rows use ShaRoundCols
+/// - Final row uses ShaDigestCols
 ///
-/// Note that for soundness, we require that there is always a padding row after the last digest row in the trace.
-/// Right now, this is true because the unpadded height is a multiple of 17, and thus not a power of 2.
-///
-/// Sha256RoundCols and Sha256DigestCols share the same first 3 fields:
+/// ShaRoundCols and ShaDigestCols share the same first 3 fields:
 /// - flags
 /// - work_vars/hash (same type, different name)
 /// - schedule_helper
 ///
 /// This design allows for:
 /// 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
-///
-/// 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>,
-    /// Stores the current state of the working variables
-    pub work_vars: Sha256WorkVarsCols<T>,
-    pub schedule_helper: Sha256MessageHelperCols<T>,
-    pub message_schedule: Sha256MessageScheduleCols<T>,
+#[derive(Clone, Copy, Debug, ColsRef)]
+#[config(Sha2Config)]
+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: Sha2FlagsCols<T, ROW_VAR_CNT>,
+    pub work_vars: ShaWorkVarsCols<T, WORD_BITS, ROUNDS_PER_ROW, WORD_U16S>,
+    pub schedule_helper:
+        Sha2MessageHelperCols<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>,
-    /// 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>,
+#[derive(Clone, Copy, Debug, ColsRef)]
+#[config(Sha2Config)]
+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: Sha2FlagsCols<T, ROW_VAR_CNT>,
+    /// Will serve as previous hash values for the next block
+    pub hash: ShaWorkVarsCols<T, WORD_BITS, ROUNDS_PER_ROW, WORD_U16S>,
+    pub schedule_helper:
+        Sha2MessageHelperCols<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)]
+#[config(Sha2Config)]
+pub struct ShaMessageScheduleCols<
+    T,
+    const WORD_BITS: usize,
+    const ROUNDS_PER_ROW: usize,
+    const WORD_U8S: usize,
+> {
     /// 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],
-    /// 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 w: [[T; WORD_BITS]; ROUNDS_PER_ROW],
+    /// Will be message schedule carries for rows 4..C::ROUND_ROWS 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; WORD_U8S]; ROUNDS_PER_ROW],
 }
 
 #[repr(C)]
-#[derive(Clone, Copy, Debug, AlignedBorrow)]
-pub struct Sha256WorkVarsCols<T> {
+#[derive(Clone, Copy, Debug, ColsRef)]
+#[config(Sha2Config)]
+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)]
+#[config(Sha2Config)]
+pub struct Sha2MessageHelperCols<
+    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],
+    /// The value of `w` from 3 rounds ago
+    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)]
+#[config(Sha2Config)]
+pub struct Sha2FlagsCols<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.
+    /// A flag that indicates if the current row is among the first 4 rows of a block (the message rows)
     pub is_first_4_rows: T,
-    /// A flag that indicates if the current row is the last (17th) row of a block.
     pub is_digest_row: T,
-    // A flag that indicates if the current row is the last block of the message.
-    // 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],
-    /// The index of the current block in the trace starting at 1.
-    /// Set to 0 on padding rows.
+    pub row_idx: [T; ROW_VAR_CNT],
+    /// The global index of the current block
     pub global_block_idx: T,
-    /// The index of the current block in the current message starting at 0.
-    /// Resets after every message.
-    /// Set to 0 on padding rows.
+    /// 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> {
-    // 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.
+impl<O, T: Copy + core::ops::Add<Output = O>, const ROW_VAR_CNT: usize>
+    Sha2FlagsCols<T, ROW_VAR_CNT>
+{
     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())
+    }
+}
+
+impl<O, T: Copy + core::ops::Add<Output = O>> Sha2FlagsColsRef<'_, 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,