Support SHA-512

Author: Avaneesh-axiom

Cargo.lock

@@ -212,6 +212,7 @@ rrs-lib = "0.1.0"
 rand = { version = "0.8.5", default-features = false }
 hex = { version = "0.4.3", default-features = false }
 serde-big-array = "0.5.1"
+ndarray = "0.16.1"
 
 # default-features = false for no_std for use in guest programs
 itertools = { version = "0.14.0", default-features = false }

Cargo.toml

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

crates/circuits/sha-air/Cargo.toml

@@ -24,15 +24,15 @@ use crate::{
 
 /// Expects the message to be padded to a multiple of C::BLOCK_WORDS * C::WORD_BITS bits
 #[derive(Clone, Debug)]
-pub struct ShaAir<C: ShaConfig> {
+pub struct Sha2Air<C: ShaConfig> {
     pub bitwise_lookup_bus: BitwiseOperationLookupBus,
     pub row_idx_encoder: Encoder,
     /// Internal bus for self-interactions in this AIR.
     bus: PermutationCheckBus,
     _phantom: PhantomData<C>,
 }
 
-impl<C: ShaConfig> ShaAir<C> {
+impl<C: ShaConfig> Sha2Air<C> {
     pub fn new(bitwise_lookup_bus: BitwiseOperationLookupBus, self_bus_idx: BusIndex) -> Self {
         Self {
             bitwise_lookup_bus,
@@ -43,13 +43,13 @@ impl<C: ShaConfig> ShaAir<C> {
     }
 }
 
-impl<F, C: ShaConfig> BaseAir<F> for ShaAir<C> {
+impl<F, C: ShaConfig> BaseAir<F> for Sha2Air<C> {
     fn width(&self) -> usize {
         max(C::ROUND_WIDTH, C::DIGEST_WIDTH)
     }
 }
 
-impl<AB: InteractionBuilder, C: ShaConfig> SubAir<AB> for ShaAir<C> {
+impl<AB: InteractionBuilder, C: ShaConfig> SubAir<AB> for Sha2Air<C> {
     /// The start column for the sub-air to use
     type AirContext<'a>
         = usize
@@ -69,7 +69,7 @@ impl<AB: InteractionBuilder, C: ShaConfig> SubAir<AB> for ShaAir<C> {
     }
 }
 
-impl<C: ShaConfig> ShaAir<C> {
+impl<C: ShaConfig> Sha2Air<C> {
     /// Implements the single row constraints (i.e. imposes constraints only on local)
     /// Implements some sanity constraints on the row index, flags, and work variables
     fn eval_row<AB: InteractionBuilder>(&self, builder: &mut AB, start_col: usize) {
@@ -211,30 +211,35 @@ impl<C: ShaConfig> ShaAir<C> {
             for j in 0..C::WORD_U16S {
                 let work_var_limb = if i < C::ROUNDS_PER_ROW {
                     compose::<AB::Expr>(
-                        &local
+                        local
                             .work_vars
                             .a
                             .slice(s![C::ROUNDS_PER_ROW - 1 - i, j * 16..(j + 1) * 16])
-                            .as_slice().unwrap(),
+                            .as_slice()
+                            .unwrap(),
                         1,
                     )
                 } else {
                     compose::<AB::Expr>(
-                        &local
+                        local
                             .work_vars
                             .e
                             .slice(s![C::ROUNDS_PER_ROW + 3 - i, j * 16..(j + 1) * 16])
-                            .as_slice().unwrap(),
+                            .as_slice()
+                            .unwrap(),
                         1,
                     )
                 };
                 let final_hash_limb = compose::<AB::Expr>(
-                    &next.final_hash.slice(s![i, j * 2..(j + 1) * 2]).as_slice().unwrap(),
+                    next.final_hash
+                        .slice(s![i, j * 2..(j + 1) * 2])
+                        .as_slice()
+                        .unwrap(),
                     8,
                 );
 
                 carry = AB::Expr::from(AB::F::from_canonical_u32(1 << 16).inverse())
-                    * (next.prev_hash[[i,j]] + work_var_limb + carry - final_hash_limb);
+                    * (next.prev_hash[[i, j]] + work_var_limb + carry - final_hash_limb);
                 builder
                     .when(*next.flags.is_digest_row)
                     .assert_bool(carry.clone());

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

@@ -127,7 +127,6 @@ pub struct ShaFlagsCols<T, const ROW_VAR_CNT: usize> {
     pub is_digest_row: T,
     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 global index of the current block
     pub global_block_idx: T,
@@ -150,7 +149,7 @@ impl<O, T: Copy + core::ops::Add<Output = O>, const ROW_VAR_CNT: usize>
     }
 }
 
-impl<'a, O, T: Copy + core::ops::Add<Output = O>> ShaFlagsColsRef<'a, T> {
+impl<O, T: Copy + core::ops::Add<Output = O>> ShaFlagsColsRef<'_, T> {
     pub fn is_not_padding_row(&self) -> O {
         *self.is_round_row + *self.is_digest_row
     }

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

@@ -1,5 +1,7 @@
 use std::ops::{BitAnd, BitOr, BitXor, Not, Shl, Shr};
 
+use crate::{ShaDigestColsRef, ShaRoundColsRef};
+
 pub trait ShaConfig: Send + Sync + Clone {
     type Word: 'static
         + Shr<usize, Output = Self::Word>
@@ -44,48 +46,25 @@ pub trait ShaConfig: Send + Sync + Clone {
     const HASH_WORDS: usize;
     /// Number of vars needed to encode the row index with [Encoder]
     const ROW_VAR_CNT: usize;
-
     /// Width of the ShaRoundCols
-    const ROUND_WIDTH: usize = Self::FLAGS_WIDTH
-        + Self::WORK_VARS_WIDTH
-        + Self::MESSAGE_HELPER_WIDTH
-        + Self::MESSAGE_SCHEDULE_WIDTH;
+    const ROUND_WIDTH: usize = ShaRoundColsRef::<u8>::width::<Self>();
     /// Width of the ShaDigestCols
-    const DIGEST_WIDTH: usize = Self::FLAGS_WIDTH
-        + Self::WORK_VARS_WIDTH
-        + Self::MESSAGE_HELPER_WIDTH
-        + Self::MESSAGE_SCHEDULE_WIDTH
-        + Self::WORD_U8S * Self::HASH_WORDS
-        + Self::WORD_U16S * Self::HASH_WORDS;
-    /// Width of the ShaFlagsCols
-    const FLAGS_WIDTH: usize = Self::ROW_VAR_CNT + 6;
-    /// Width of the ShaWorkVarsCols
-    const WORK_VARS_WIDTH: usize =
-        2 * Self::WORD_BITS * Self::ROUNDS_PER_ROW + 2 * Self::WORD_U16S * Self::ROUNDS_PER_ROW;
-    // Width of the ShaMessageHelperCols
-    const MESSAGE_HELPER_WIDTH: usize =
-        Self::WORD_U8S * (Self::ROUNDS_PER_ROW - 1) + 3 * Self::WORD_U16S * Self::ROUNDS_PER_ROW;
-    /// Width of the ShaMessageScheduleCols
-    const MESSAGE_SCHEDULE_WIDTH: usize =
-        Self::WORD_BITS * Self::ROUNDS_PER_ROW + Self::WORD_U8S * Self::ROUNDS_PER_ROW;
-
-    /// Size of the buffer of the first 4 rows of a block (each row's size)
-    const BUFFER_SIZE: usize = Self::ROUNDS_PER_ROW * Self::WORD_U8S;
+    const DIGEST_WIDTH: usize = ShaDigestColsRef::<u8>::width::<Self>();
     /// Width of the ShaCols
     const WIDTH: usize = if Self::ROUND_WIDTH > Self::DIGEST_WIDTH {
         Self::ROUND_WIDTH
     } else {
         Self::DIGEST_WIDTH
     };
-}
+    /// Size of the buffer of the first 4 rows of a block (each row's size)
+    const BUFFER_SIZE: usize = Self::ROUNDS_PER_ROW * Self::WORD_U8S;
 
-/// 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> {
+    ///  To optimize the trace generation of invalid rows, we precompute those values.
     // these should be appropriately sized for the config
     fn get_invalid_carry_a(round_num: usize) -> &'static [u32];
     fn get_invalid_carry_e(round_num: usize) -> &'static [u32];
+
+    /// We also store the SHA constants K and H
     fn get_k() -> &'static [Self::Word];
     fn get_h() -> &'static [Self::Word];
 }

crates/circuits/sha-air/src/config.rs

@@ -1,5 +1,5 @@
+use crate::{ShaDigestColsRefMut, ShaRoundColsRef, ShaRoundColsRefMut};
 use std::{borrow::BorrowMut, cmp::max, sync::Arc};
-use crate::{ShaDigestColsRefMut, ShaRoundColsRefMut, ShaRoundColsRef};
 
 use openvm_circuit::arch::{
     instructions::riscv::RV32_CELL_BITS,
@@ -22,21 +22,19 @@ use openvm_stark_backend::{
 use openvm_stark_sdk::utils::create_seeded_rng;
 use rand::Rng;
 
-use crate::{
-    compose, small_sig0_field, Sha256Config, Sha512Config, ShaAir, ShaConfig,
-};
+use crate::{compose, small_sig0_field, Sha256Config, Sha512Config, Sha2Air, ShaConfig};
 
 // A wrapper AIR purely for testing purposes
 #[derive(Clone, Debug)]
 pub struct ShaTestAir<C: ShaConfig> {
-    pub sub_air: ShaAir<C>,
+    pub sub_air: Sha2Air<C>,
 }
 
 impl<F: Field, C: ShaConfig> BaseAirWithPublicValues<F> for ShaTestAir<C> {}
 impl<F: Field, C: ShaConfig> PartitionedBaseAir<F> for ShaTestAir<C> {}
 impl<F: Field, C: ShaConfig> BaseAir<F> for ShaTestAir<C> {
     fn width(&self) -> usize {
-        <ShaAir<C> as BaseAir<F>>::width(&self.sub_air)
+        <Sha2Air<C> as BaseAir<F>>::width(&self.sub_air)
     }
 }
 
@@ -103,7 +101,7 @@ fn rand_sha_test<C: ShaConfig + 'static>() {
         .collect();
     let chip = ShaTestChip {
         air: ShaTestAir {
-            sub_air: ShaAir::<C>::new(bitwise_bus, SELF_BUS_IDX),
+            sub_air: Sha2Air::<C>::new(bitwise_bus, SELF_BUS_IDX),
         },
         bitwise_lookup_chip: bitwise_chip.clone(),
         records: random_records,
@@ -125,14 +123,13 @@ fn rand_sha512_test() {
 
 // A wrapper Chip to test that the final_hash is properly constrained.
 // This chip implements a malicious trace gen that violates the final_hash constraints.
-pub struct ShaTestBadFinalHashChip<C: ShaConfig + ShaPrecomputedValues<C::Word>> {
+pub struct ShaTestBadFinalHashChip<C: ShaConfig> {
     pub air: ShaTestAir<C>,
     pub bitwise_lookup_chip: SharedBitwiseOperationLookupChip<8>,
     pub records: Vec<(Vec<u8>, bool)>, // length of inner vec should be C::BLOCK_U8S
 }
 
-impl<SC: StarkGenericConfig, C: ShaConfig + ShaPrecomputedValues<C::Word> + 'static> Chip<SC>
-    for ShaTestBadFinalHashChip<C>
+impl<SC: StarkGenericConfig, C: ShaConfig + 'static> Chip<SC> for ShaTestBadFinalHashChip<C>
 where
     Val<SC>: PrimeField32,
 {
@@ -170,12 +167,15 @@ where
                 let mut last_digest_row: crate::ShaRoundColsRefMut<Val<SC>> =
                     ShaRoundColsRefMut::from::<C>(last_digest_row.borrow_mut());
                 // fix the intermed_4 for the digest row
-                generate_intermed_4::<Val<SC>, C>(&ShaRoundColsRef::from_mut::<C>(&last_round_row), &mut last_digest_row);
+                generate_intermed_4::<Val<SC>, C>(
+                    &ShaRoundColsRef::from_mut::<C>(&last_round_row),
+                    &mut last_digest_row,
+                );
             }
         }
 
         let non_padded_height = self.records.len() * C::ROWS_PER_BLOCK;
-        let width = <ShaAir<C> as BaseAir<Val<SC>>>::width(&self.air.sub_air);
+        let width = <Sha2Air<C> as BaseAir<Val<SC>>>::width(&self.air.sub_air);
         // recalculate the missing cells (second pass of generate_trace)
         trace.values[width..]
             .par_chunks_mut(width * C::ROWS_PER_BLOCK)
@@ -190,7 +190,7 @@ where
 
 // Copy of private method in Sha256Air used for testing
 /// Puts the correct intermed_4 in the `next_row`
-fn generate_intermed_4<F: PrimeField32, C: ShaConfig + ShaPrecomputedValues<C::Word>>(
+fn generate_intermed_4<F: PrimeField32, C: ShaConfig>(
     local_cols: &ShaRoundColsRef<F>,
     next_cols: &mut ShaRoundColsRefMut<F>,
 ) {
@@ -209,10 +209,9 @@ fn generate_intermed_4<F: PrimeField32, C: ShaConfig + ShaPrecomputedValues<C::W
             .collect::<Vec<_>>(),
     ]
     .concat();
- 
-    
+
     // length of inner vec is C::WORD_U16S
-    let w_limbs: Vec<Vec<F>> = w 
+    let w_limbs: Vec<Vec<F>> = w
         .iter()
         .map(|x| {
             (0..C::WORD_U16S)
@@ -231,7 +230,7 @@ fn generate_intermed_4<F: PrimeField32, C: ShaConfig + ShaPrecomputedValues<C::W
     }
 }
 
-impl<C: ShaConfig + ShaPrecomputedValues<C::Word>> ChipUsageGetter for ShaTestBadFinalHashChip<C> {
+impl<C: ShaConfig> ChipUsageGetter for ShaTestBadFinalHashChip<C> {
     fn air_name(&self) -> String {
         get_air_name(&self.air)
     }
@@ -244,18 +243,25 @@ impl<C: ShaConfig + ShaPrecomputedValues<C::Word>> ChipUsageGetter for ShaTestBa
     }
 }
 
-fn test_sha_final_hash_constraints<C: ShaConfig + ShaPrecomputedValues<C::Word> + 'static>() {
+fn test_sha_final_hash_constraints<C: ShaConfig + '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(|_| ((0..C::BLOCK_U8S).map(|_| rng.gen::<u8>()).collect::<Vec<_>>(), true))
+        .map(|_| {
+            (
+                (0..C::BLOCK_U8S)
+                    .map(|_| rng.gen::<u8>())
+                    .collect::<Vec<_>>(),
+                true,
+            )
+        })
         .collect();
     let chip = ShaTestBadFinalHashChip {
         air: ShaTestAir {
-            sub_air: ShaAir::<C>::new(bitwise_bus, SELF_BUS_IDX),
+            sub_air: Sha2Air::<C>::new(bitwise_bus, SELF_BUS_IDX),
         },
         bitwise_lookup_chip: bitwise_chip.clone(),
         records: random_records,