feat: e1,e3 implementation for sha2

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

@@ -15,11 +15,11 @@ use openvm_stark_backend::{
 
 use super::{
     big_sig0_field, big_sig1_field, ch_field, compose, maj_field, small_sig0_field,
-    small_sig1_field, u32_into_limbs, Sha256DigestCols, Sha256RoundCols, SHA256_DIGEST_WIDTH,
-    SHA256_H, SHA256_HASH_WORDS, SHA256_K, SHA256_ROUNDS_PER_ROW, SHA256_ROUND_WIDTH,
-    SHA256_WORD_BITS, SHA256_WORD_U16S, SHA256_WORD_U8S,
+    small_sig1_field, Sha256DigestCols, Sha256RoundCols, SHA256_DIGEST_WIDTH, SHA256_H,
+    SHA256_HASH_WORDS, SHA256_K, SHA256_ROUNDS_PER_ROW, SHA256_ROUND_WIDTH, SHA256_WORD_BITS,
+    SHA256_WORD_U16S, SHA256_WORD_U8S,
 };
-use crate::constraint_word_addition;
+use crate::{constraint_word_addition, u32_into_u16s};
 
 /// Expects the message to be padded to a multiple of 512 bits
 #[derive(Clone, Debug)]
@@ -154,7 +154,7 @@ impl Sha256Air {
                     .assert_eq(
                         a_limb,
                         AB::Expr::from_canonical_u32(
-                            u32_into_limbs::<2>(SHA256_H[SHA256_ROUNDS_PER_ROW - i - 1])[j],
+                            u32_into_u16s(SHA256_H[SHA256_ROUNDS_PER_ROW - i - 1])[j],
                         ),
                     );
 
@@ -166,7 +166,7 @@ impl Sha256Air {
                     .assert_eq(
                         e_limb,
                         AB::Expr::from_canonical_u32(
-                            u32_into_limbs::<2>(SHA256_H[SHA256_ROUNDS_PER_ROW - i + 3])[j],
+                            u32_into_u16s(SHA256_H[SHA256_ROUNDS_PER_ROW - i + 3])[j],
                         ),
                     );
             }
@@ -561,9 +561,8 @@ impl Sha256Air {
                         .map(|rw_idx| {
                             (
                                 rw_idx,
-                                u32_into_limbs::<SHA256_WORD_U16S>(
-                                    SHA256_K[rw_idx * SHA256_ROUNDS_PER_ROW + i],
-                                )[j] as usize,
+                                u32_into_u16s(SHA256_K[rw_idx * SHA256_ROUNDS_PER_ROW + i])[j]
+                                    as usize,
                             )
                         })
                         .collect::<Vec<_>>(),

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

@@ -13,18 +13,19 @@ use openvm_stark_backend::{
     interaction::{BusIndex, InteractionBuilder},
     p3_air::{Air, BaseAir},
     p3_field::{Field, FieldAlgebra, PrimeField32},
-    p3_maybe_rayon::prelude::{IndexedParallelIterator, ParallelIterator, ParallelSliceMut},
+    p3_matrix::{dense::RowMajorMatrix, Matrix},
     prover::types::AirProofInput,
     rap::{get_air_name, BaseAirWithPublicValues, PartitionedBaseAir},
+    utils::disable_debug_builder,
+    verifier::VerificationError,
     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 crate::{
-    compose, small_sig0_field, Sha256Air, Sha256RoundCols, SHA256_BLOCK_U8S, SHA256_DIGEST_WIDTH,
-    SHA256_HASH_WORDS, SHA256_ROUNDS_PER_ROW, SHA256_ROUND_WIDTH, SHA256_ROWS_PER_BLOCK,
-    SHA256_WORD_U16S, SHA256_WORD_U8S,
+    Sha256Air, Sha256DigestCols, Sha256StepHelper, SHA256_BLOCK_U8S, SHA256_DIGEST_WIDTH,
+    SHA256_HASH_WORDS, SHA256_ROUND_WIDTH, SHA256_ROWS_PER_BLOCK, SHA256_WORD_U8S,
 };
 
 // A wrapper AIR purely for testing purposes
@@ -50,6 +51,7 @@ impl<AB: InteractionBuilder> Air<AB> for Sha256TestAir {
 // A wrapper Chip purely for testing purposes
 pub struct Sha256TestChip {
     pub air: Sha256TestAir,
+    pub step: Sha256StepHelper,
     pub bitwise_lookup_chip: SharedBitwiseOperationLookupChip<8>,
     pub records: Vec<([u8; SHA256_BLOCK_U8S], bool)>,
 }
@@ -64,8 +66,9 @@ where
 
     fn generate_air_proof_input(self) -> AirProofInput<SC> {
         let trace = crate::generate_trace::<Val<SC>>(
-            &self.air.sub_air,
-            self.bitwise_lookup_chip.clone(),
+            &self.step,
+            self.bitwise_lookup_chip.as_ref(),
+            <Sha256Air as BaseAir<Val<SC>>>::width(&self.air.sub_air),
             self.records,
         );
         AirProofInput::simple_no_pis(trace)
@@ -86,10 +89,10 @@ impl ChipUsageGetter for Sha256TestChip {
 }
 
 const SELF_BUS_IDX: BusIndex = 28;
-#[test]
-fn rand_sha256_test() {
+type F = BabyBear;
+
+fn create_chip_with_rand_records() -> (Sha256TestChip, SharedBitwiseOperationLookupChip<8>) {
     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);
@@ -105,129 +108,47 @@ fn rand_sha256_test() {
         air: Sha256TestAir {
             sub_air: Sha256Air::new(bitwise_bus, SELF_BUS_IDX),
         },
+        step: Sha256StepHelper::new(),
         bitwise_lookup_chip: bitwise_chip.clone(),
         records: random_records,
     };
+    (chip, bitwise_chip)
+}
 
+#[test]
+fn rand_sha256_test() {
+    let tester = VmChipTestBuilder::default();
+    let (chip, bitwise_chip) = create_chip_with_rand_records();
     let tester = tester.build().load(chip).load(bitwise_chip).finalize();
     tester.simple_test().expect("Verification failed");
 }
 
-// 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 Sha256TestBadFinalHashChip {
-    pub air: Sha256TestAir,
-    pub bitwise_lookup_chip: SharedBitwiseOperationLookupChip<8>,
-    pub records: Vec<([u8; SHA256_BLOCK_U8S], bool)>,
-}
-
-impl<SC: StarkGenericConfig> Chip<SC> for Sha256TestBadFinalHashChip
-where
-    Val<SC>: PrimeField32,
-{
-    fn air(&self) -> AirRef<SC> {
-        Arc::new(self.air.clone())
-    }
-
-    fn generate_air_proof_input(self) -> AirProofInput<SC> {
-        let mut trace = crate::generate_trace::<Val<SC>>(
-            &self.air.sub_air,
-            self.bitwise_lookup_chip.clone(),
-            self.records.clone(),
-        );
-
-        // Set the final_hash in the digest row of the last block of each hash to zero.
-        // That is, every hash that this chip does will result in a final_hash of zero.
-        for (i, row) in self.records.iter().enumerate() {
-            if row.1 {
-                let last_digest_row_idx = (i + 1) * SHA256_ROWS_PER_BLOCK - 1;
-                let last_digest_row: &mut crate::Sha256DigestCols<Val<SC>> =
-                    trace.row_mut(last_digest_row_idx)[..SHA256_DIGEST_WIDTH].borrow_mut();
-                // Set the final_hash to all zeros
+#[test]
+fn negative_sha256_test_bad_final_hash() {
+    let tester = VmChipTestBuilder::default();
+    let (chip, bitwise_chip) = create_chip_with_rand_records();
+
+    // Set the final_hash to all zeros
+    let modify_trace = |trace: &mut RowMajorMatrix<F>| {
+        trace.row_chunks_exact_mut(1).for_each(|row| {
+            let mut row_slice = row.row_slice(0).to_vec();
+            let cols: &mut Sha256DigestCols<F> = row_slice[..SHA256_DIGEST_WIDTH].borrow_mut();
+            if cols.flags.is_last_block.is_one() && cols.flags.is_digest_row.is_one() {
                 for i in 0..SHA256_HASH_WORDS {
                     for j in 0..SHA256_WORD_U8S {
-                        last_digest_row.final_hash[i][j] = Val::<SC>::ZERO;
+                        cols.final_hash[i][j] = F::ZERO;
                     }
                 }
-
-                let (last_round_row, last_digest_row) =
-                    trace.row_pair_mut(last_digest_row_idx - 1, last_digest_row_idx);
-                let last_round_row: &mut crate::Sha256RoundCols<Val<SC>> =
-                    last_round_row.borrow_mut();
-                let last_digest_row: &mut crate::Sha256RoundCols<Val<SC>> =
-                    last_digest_row.borrow_mut();
-                // fix the intermed_4 for the digest row
-                generate_intermed_4(last_round_row, last_digest_row);
+                row.values.copy_from_slice(&row_slice);
             }
-        }
-
-        let non_padded_height = self.records.len() * SHA256_ROWS_PER_BLOCK;
-        let width = <Sha256Air 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 * SHA256_ROWS_PER_BLOCK)
-            .take(non_padded_height / SHA256_ROWS_PER_BLOCK)
-            .for_each(|chunk| {
-                self.air.sub_air.generate_missing_cells(chunk, width, 0);
-            });
-
-        AirProofInput::simple_no_pis(trace)
-    }
-}
-
-// Copy of private method in Sha256Air used for testing
-/// Puts the correct intermed_4 in the `next_row`
-fn generate_intermed_4<F: PrimeField32>(
-    local_cols: &Sha256RoundCols<F>,
-    next_cols: &mut Sha256RoundCols<F>,
-) {
-    let w = [local_cols.message_schedule.w, next_cols.message_schedule.w].concat();
-    let w_limbs: Vec<[F; SHA256_WORD_U16S]> = w
-        .iter()
-        .map(|x| array::from_fn(|i| compose::<F>(&x[i * 16..(i + 1) * 16], 1)))
-        .collect();
-    for i in 0..SHA256_ROUNDS_PER_ROW {
-        let sig_w = small_sig0_field::<F>(&w[i + 1]);
-        let sig_w_limbs: [F; SHA256_WORD_U16S] =
-            array::from_fn(|j| compose::<F>(&sig_w[j * 16..(j + 1) * 16], 1));
-        for (j, sig_w_limb) in sig_w_limbs.iter().enumerate() {
-            next_cols.schedule_helper.intermed_4[i][j] = w_limbs[i][j] + *sig_w_limb;
-        }
-    }
-}
-
-impl ChipUsageGetter for Sha256TestBadFinalHashChip {
-    fn air_name(&self) -> String {
-        get_air_name(&self.air)
-    }
-    fn current_trace_height(&self) -> usize {
-        self.records.len() * SHA256_ROWS_PER_BLOCK
-    }
-
-    fn trace_width(&self) -> usize {
-        max(SHA256_ROUND_WIDTH, SHA256_DIGEST_WIDTH)
-    }
-}
-
-#[test]
-#[should_panic]
-fn test_sha256_final_hash_constraints() {
-    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))
-        .collect();
-    let chip = Sha256TestBadFinalHashChip {
-        air: Sha256TestAir {
-            sub_air: Sha256Air::new(bitwise_bus, SELF_BUS_IDX),
-        },
-        bitwise_lookup_chip: bitwise_chip.clone(),
-        records: random_records,
+        });
     };
 
-    let tester = tester.build().load(chip).load(bitwise_chip).finalize();
-    tester.simple_test().expect("Verification failed");
+    disable_debug_builder();
+    let tester = tester
+        .build()
+        .load_and_prank_trace(chip, modify_trace)
+        .load(bitwise_chip)
+        .finalize();
+    tester.simple_test_with_expected_error(VerificationError::OodEvaluationMismatch);
 }