Updated names

Author: Avaneesh-axiom

Cargo.lock

@@ -112,7 +112,7 @@ openvm-stark-sdk = { git = "https://github.com/openvm-org/stark-backend.git", ta
 openvm-sdk = { path = "crates/sdk", default-features = false }
 openvm-mod-circuit-builder = { path = "crates/circuits/mod-builder", default-features = false }
 openvm-poseidon2-air = { path = "crates/circuits/poseidon2-air", default-features = false }
-openvm-sha256-air = { path = "crates/circuits/sha256-air", default-features = false }
+openvm-sha-air = { path = "crates/circuits/sha-air", default-features = false }
 openvm-sha-macros = { path = "crates/circuits/sha-macros", default-features = false }
 openvm-circuit-primitives = { path = "crates/circuits/primitives", default-features = false }
 openvm-circuit-primitives-derive = { path = "crates/circuits/primitives/derive", default-features = false }

Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "openvm-sha256-air"
+name = "openvm-sha-air"
 version.workspace = true
 authors.workspace = true
 edition.workspace = true

crates/circuits/sha256-air/Cargo.toml renamed to crates/circuits/sha-air/Cargo.toml

@@ -6,23 +6,18 @@ use openvm_stark_backend::p3_field::FieldAlgebra;
 
 use crate::ShaConfig;
 
-/// 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, ColsRef)]
 pub struct ShaRoundCols<
@@ -35,7 +30,6 @@ pub struct ShaRoundCols<
     const ROW_VAR_CNT: usize,
 > {
     pub flags: ShaFlagsCols<T, ROW_VAR_CNT>,
-    /// Stores the current state of the working variables
     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>,
@@ -55,11 +49,7 @@ pub struct ShaDigestCols<
     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.
+    /// Will serve as previous hash values for the next block
     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>,
@@ -79,11 +69,10 @@ pub struct ShaMessageScheduleCols<
     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
+    /// The message schedule words as 32-bit intergers
     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 2 bit numbers represented using 2 cells as individual bits
+    /// Note: carries are represented as 2 bit numbers
     pub carry_or_buffer: [[T; WORD_U8S]; ROUNDS_PER_ROW],
 }
 
@@ -114,12 +103,10 @@ pub struct ShaMessageHelperCols<
     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
+    /// 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; WORD_U16S]; ROUNDS_PER_ROW],
     pub intermed_8: [[T; WORD_U16S]; ROUNDS_PER_ROW],
     pub intermed_12: [[T; WORD_U16S]; ROUNDS_PER_ROW],
@@ -128,38 +115,27 @@ pub struct ShaMessageHelperCols<
 #[repr(C)]
 #[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.
+    /// A flag that indicates if the current row is among the first 4 rows of a block
     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
     //#[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.
+    /// 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>, 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 {
         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,
@@ -168,14 +144,11 @@ impl<O, T: Copy + core::ops::Add<Output = O>, const ROW_VAR_CNT: usize>
     }
 }
 
-// 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,

crates/circuits/sha256-air/src/air.rs renamed to crates/circuits/sha-air/src/air.rs

@@ -81,6 +81,9 @@ pub trait ShaConfig: Send + Sync + Clone {
     };
 }
 
+/// We can notice that `carry_a`'s and `carry_e`'s are always the same on invalid rows
+/// To optimize the trace generation of invalid rows, we precompute those values.
+/// This trait also stores the constants K and H for the given SHA config.
 pub trait ShaPrecomputedValues<T> {
     // these should be appropirately sized for the config
     fn get_invalid_carry_a(round_num: usize) -> &'static [u32];
@@ -111,8 +114,6 @@ impl ShaConfig for Sha256Config {
     const ROW_VAR_CNT: usize = 5;
 }
 
-/// We can notice that `carry_a`'s and `carry_e`'s are always the same on invalid rows
-/// To optimize the trace generation of invalid rows, we have those values precomputed here
 pub const SHA256_INVALID_CARRY_A: [[u32; Sha256Config::WORD_U16S]; Sha256Config::ROUNDS_PER_ROW] = [
     [1230919683, 1162494304],
     [266373122, 1282901987],
@@ -179,8 +180,6 @@ impl ShaConfig for Sha512Config {
     const ROW_VAR_CNT: usize = 6;
 }
 
-/// We can notice that `carry_a`'s and `carry_e`'s are always the same on invalid rows
-/// To optimize the trace generation of invalid rows, we have those values precomputed here
 pub(crate) const SHA512_INVALID_CARRY_A: [[u32; Sha512Config::WORD_U16S];
     Sha512Config::ROUNDS_PER_ROW] = [
     [55971842, 827997017, 993005918, 512731953],
@@ -308,7 +307,7 @@ impl ShaPrecomputedValues<u64> for Sha512Config {
 }
 
 // Needed to avoid compile errors in utils.rs
-// don't ask me why this doesn't inf loop
+// not sure why this doesn't inf loop
 pub trait RotateRight {
     fn rotate_right(self, n: u32) -> Self;
 }

crates/circuits/sha256-air/src/columns.rs renamed to crates/circuits/sha-air/src/columns.rs

@@ -1,4 +1,4 @@
-//! Implementation of the SHA256 compression function without padding
+//! Implementation of the SHA256/SHA512 compression function without padding
 //! This this AIR doesn't constrain any of the message padding
 
 mod air;

crates/circuits/sha256-air/src/config.rs renamed to crates/circuits/sha-air/src/config.rs

@@ -19,7 +19,7 @@ use crate::{
     ShaRoundColsRef, WrappingAdd,
 };
 
-/// The trace generation of SHA256 should be done in two passes.
+/// The trace generation of SHA should be done in two passes.
 /// The first pass should do `get_block_trace` for every block and generate the invalid rows through `get_default_row`
 /// The second pass should go through all the blocks and call `generate_missing_cells`
 impl<C: ShaConfig + ShaPrecomputedValues<C::Word>> ShaAir<C> {
@@ -455,7 +455,7 @@ impl<C: ShaConfig + ShaPrecomputedValues<C::Word>> ShaAir<C> {
     /// This function should be called only after `generate_block_trace` was called for all blocks
     /// And [`Self::generate_default_row`] is called for all invalid rows
     /// Will populate the missing values of `trace`, where the width of the trace is `trace_width`
-    /// and the starting column for the `Sha256Air` is `trace_start_col`.
+    /// and the starting column for the `ShaAir` is `trace_start_col`.
     /// Note: `trace` needs to be the rows 1..C::ROWS_PER_BLOCK of a block and the first row of the next block
     pub fn generate_missing_cells<F: PrimeField32>(
         &self,
@@ -748,11 +748,8 @@ pub fn generate_trace<F: PrimeField32, C: ShaConfig + ShaPrecomputedValues<C::Wo
         debug_assert!(input.len() == C::BLOCK_U8S);
     }
 
-    println!("records.len(): {}", records.len());
     let non_padded_height = records.len() * C::ROWS_PER_BLOCK;
-    println!("non_padded_height: {}", non_padded_height);
     let height = next_power_of_two_or_zero(non_padded_height);
-    println!("height: {}", height);
     let width = <ShaAir<C> as BaseAir<F>>::width(sub_air);
     let mut values = F::zero_vec(height * width);
 

crates/circuits/sha256-air/src/lib.rs renamed to crates/circuits/sha-air/src/lib.rs

@@ -102,7 +102,7 @@ pub(crate) fn xor_bit<F: FieldAlgebra + Clone>(
         + (not::<F>(x) * not::<F>(y) * z)
 }
 
-/// Computes x ^ y ^ z, where x, y, z are [SHA256_WORD_BITS] bit numbers
+/// Computes x ^ y ^ z, where x, y, z are [C::WORD_BITS] bit numbers
 #[inline]
 pub(crate) fn xor<F: FieldAlgebra + Clone>(
     x: &[impl Into<F> + Clone],
@@ -114,13 +114,13 @@ pub(crate) fn xor<F: FieldAlgebra + Clone>(
         .collect()
 }
 
-/// Choose function from SHA256
+/// Choose function from the SHA spec
 #[inline]
 pub fn ch<C: ShaConfig>(x: C::Word, y: C::Word, z: C::Word) -> C::Word {
     (x & y) ^ ((!x) & z)
 }
 
-/// Computes Ch(x,y,z), where x, y, z are [SHA256_WORD_BITS] bit numbers
+/// Computes Ch(x,y,z), where x, y, z are [C::WORD_BITS] bit numbers
 #[inline]
 pub(crate) fn ch_field<F: FieldAlgebra>(
     x: &[impl Into<F> + Clone],
@@ -132,12 +132,12 @@ pub(crate) fn ch_field<F: FieldAlgebra>(
         .collect()
 }
 
-/// Majority function from SHA256
+/// Majority function from the SHA spec
 pub fn maj<C: ShaConfig>(x: C::Word, y: C::Word, z: C::Word) -> C::Word {
     (x & y) ^ (x & z) ^ (y & z)
 }
 
-/// Computes Maj(x,y,z), where x, y, z are [SHA256_WORD_BITS] bit numbers
+/// Computes Maj(x,y,z), where x, y, z are [C::WORD_BITS] bit numbers
 #[inline]
 pub(crate) fn maj_field<F: FieldAlgebra + Clone>(
     x: &[impl Into<F> + Clone],
@@ -157,7 +157,7 @@ pub(crate) fn maj_field<F: FieldAlgebra + Clone>(
         .collect()
 }
 
-/// Big sigma_0 function from SHA256
+/// Big sigma_0 function from the SHA spec
 pub fn big_sig0<C: ShaConfig>(x: C::Word) -> C::Word {
     if C::WORD_BITS == 32 {
         x.rotate_right(2) ^ x.rotate_right(13) ^ x.rotate_right(22)
@@ -178,7 +178,7 @@ pub(crate) fn big_sig0_field<F: FieldAlgebra + Clone, C: ShaConfig>(
     }
 }
 
-/// Big sigma_1 function from SHA256
+/// Big sigma_1 function from the SHA spec
 pub fn big_sig1<C: ShaConfig>(x: C::Word) -> C::Word {
     if C::WORD_BITS == 32 {
         x.rotate_right(6) ^ x.rotate_right(11) ^ x.rotate_right(25)
@@ -187,7 +187,7 @@ pub fn big_sig1<C: ShaConfig>(x: C::Word) -> C::Word {
     }
 }
 
-/// Computes BigSigma1(x), where x is a [SHA256_WORD_BITS] bit number in little-endian
+/// Computes BigSigma1(x), where x is a [C::WORD_BITS] bit number in little-endian
 #[inline]
 pub(crate) fn big_sig1_field<F: FieldAlgebra + Clone, C: ShaConfig>(
     x: &[impl Into<F> + Clone],
@@ -199,7 +199,7 @@ pub(crate) fn big_sig1_field<F: FieldAlgebra + Clone, C: ShaConfig>(
     }
 }
 
-/// Small sigma_0 function from SHA256
+/// Small sigma_0 function from the SHA spec
 pub fn small_sig0<C: ShaConfig>(x: C::Word) -> C::Word {
     if C::WORD_BITS == 32 {
         x.rotate_right(7) ^ x.rotate_right(18) ^ (x >> 3)
@@ -208,7 +208,7 @@ pub fn small_sig0<C: ShaConfig>(x: C::Word) -> C::Word {
     }
 }
 
-/// Computes SmallSigma0(x), where x is a [SHA256_WORD_BITS] bit number in little-endian
+/// Computes SmallSigma0(x), where x is a [C::WORD_BITS] bit number in little-endian
 #[inline]
 pub(crate) fn small_sig0_field<F: FieldAlgebra + Clone, C: ShaConfig>(
     x: &[impl Into<F> + Clone],
@@ -220,7 +220,7 @@ pub(crate) fn small_sig0_field<F: FieldAlgebra + Clone, C: ShaConfig>(
     }
 }
 
-/// Small sigma_1 function from SHA256
+/// Small sigma_1 function from the SHA spec
 pub fn small_sig1<C: ShaConfig>(x: C::Word) -> C::Word {
     if C::WORD_BITS == 32 {
         x.rotate_right(17) ^ x.rotate_right(19) ^ (x >> 10)
@@ -229,7 +229,7 @@ pub fn small_sig1<C: ShaConfig>(x: C::Word) -> C::Word {
     }
 }
 
-/// Computes SmallSigma1(x), where x is a [SHA256_WORD_BITS] bit number in little-endian
+/// Computes SmallSigma1(x), where x is a [C::WORD_BITS] bit number in little-endian
 #[inline]
 pub(crate) fn small_sig1_field<F: FieldAlgebra + Clone, C: ShaConfig>(
     x: &[impl Into<F> + Clone],
@@ -253,7 +253,7 @@ pub fn get_flag_pt_array(encoder: &Encoder, flag_idx: usize) -> Vec<u32> {
     encoder.get_flag_pt(flag_idx)
 }
 
-/// Constrain the addition of [SHA256_WORD_BITS] bit words in 16-bit limbs
+/// Constrain the addition of [C::WORD_BITS] bit words in 16-bit limbs
 /// It takes in the terms some in bits some in 16-bit limbs,
 /// the expected sum in bits and the carries
 pub fn constraint_word_addition<AB: AirBuilder, C: ShaConfig>(