fix rebase issues

Avaneesh-axiom · Avaneesh-axiom · commit eee18cd7059c · 2025-03-28T12:35:30.000-04:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/circuits/sha-air/src/air.rs b/crates/circuits/sha-air/src/air.rs
@@ -209,34 +209,44 @@ impl<C: ShaConfig + ShaPrecomputedValues<C::Word>> ShaAir<C> {
         // Assert that the previous hash + work vars == final hash.
         // That is, `next.prev_hash[i] + local.work_vars[i] == next.final_hash[i]`
         // where addition is done modulo 2^32
-        for i in 0..SHA256_HASH_WORDS {
+        for i in 0..C::HASH_WORDS {
             let mut carry = AB::Expr::ZERO;
-            for j in 0..SHA256_WORD_U16S {
-                let work_var_limb = if i < SHA256_ROUNDS_PER_ROW {
+            for j in 0..C::WORD_U16S {
+                let work_var_limb = if i < C::ROUNDS_PER_ROW {
                     compose::<AB::Expr>(
-                        &local.work_vars.a[SHA256_ROUNDS_PER_ROW - 1 - i][j * 16..(j + 1) * 16],
+                        &local
+                            .work_vars
+                            .a
+                            .slice(s![C::ROUNDS_PER_ROW - 1 - i, j * 16..(j + 1) * 16])
+                            .as_slice().unwrap(),
                         1,
                     )
                 } else {
                     compose::<AB::Expr>(
-                        &local.work_vars.e[SHA256_ROUNDS_PER_ROW + 3 - i][j * 16..(j + 1) * 16],
+                        &local
+                            .work_vars
+                            .e
+                            .slice(s![C::ROUNDS_PER_ROW + 3 - i, j * 16..(j + 1) * 16])
+                            .as_slice().unwrap(),
                         1,
                     )
                 };
-                let final_hash_limb =
-                    compose::<AB::Expr>(&next.final_hash[i][j * 2..(j + 1) * 2], 8);
+                let final_hash_limb = compose::<AB::Expr>(
+                    &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)
+                    .when(*next.flags.is_digest_row)
                     .assert_bool(carry.clone());
             }
             // constrain the final hash limbs two at a time since we can do two checks per interaction
-            for chunk in next.final_hash[i].chunks(2) {
+            for chunk in next.final_hash.row(i).as_slice().unwrap().chunks(2) {
                 self.bitwise_lookup_bus
                     .send_range(chunk[0], chunk[1])
-                    .eval(builder, next.flags.is_digest_row);
+                    .eval(builder, *next.flags.is_digest_row);
             }
         }
     }
@@ -364,7 +374,7 @@ impl<C: ShaConfig + ShaPrecomputedValues<C::Word>> ShaAir<C> {
             .assert_zero(*next_cols.flags.local_block_idx);
 
         self.eval_message_schedule(builder, local_cols.clone(), next_cols.clone());
-        self.eval_work_vars(builder, local_cols, next_cols);
+        self.eval_work_vars(builder, local_cols.clone(), next_cols);
         let next_cols: ShaDigestColsRef<AB::Var> =
             ShaDigestColsRef::from::<C>(&next[start_col..start_col + C::DIGEST_WIDTH]);
         self.eval_digest_row(builder, local_cols, next_cols);
diff --git a/crates/circuits/sha-air/src/columns.rs b/crates/circuits/sha-air/src/columns.rs
@@ -69,7 +69,7 @@ pub struct ShaMessageScheduleCols<
     const ROUNDS_PER_ROW: usize,
     const WORD_U8S: usize,
 > {
-    /// The message schedule words as 32-bit intergers
+    /// The message schedule words as 32-bit integers
     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
diff --git a/crates/circuits/sha-air/src/config.rs b/crates/circuits/sha-air/src/config.rs
@@ -40,8 +40,6 @@ pub trait ShaConfig: Send + Sync + Clone {
     const ROUNDS_PER_ROW_MINUS_ONE: usize = Self::ROUNDS_PER_ROW - 1;
     /// 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]
@@ -85,7 +83,7 @@ pub trait ShaConfig: Send + Sync + Clone {
 /// 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
+    // 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];
     fn get_k() -> &'static [T];
diff --git a/crates/circuits/sha-air/src/tests.rs b/crates/circuits/sha-air/src/tests.rs
@@ -1,4 +1,5 @@
-use std::{array, borrow::BorrowMut, cmp::max, sync::Arc};
+use std::{borrow::BorrowMut, cmp::max, sync::Arc};
+use crate::{ShaDigestColsRefMut, ShaRoundColsRefMut, ShaRoundColsRef};
 
 use openvm_circuit::arch::{
     instructions::riscv::RV32_CELL_BITS,
@@ -22,8 +23,7 @@ use openvm_stark_sdk::utils::create_seeded_rng;
 use rand::Rng;
 
 use crate::{
-    compose, small_sig0_field, Sha256Config, Sha512Config, ShaAir, ShaConfig, ShaFlagsColsRef,
-    ShaFlagsColsRefMut, ShaPrecomputedValues,
+    compose, small_sig0_field, Sha256Config, Sha512Config, ShaAir, ShaConfig, ShaPrecomputedValues,
 };
 
 // A wrapper AIR purely for testing purposes
@@ -101,7 +101,7 @@ fn rand_sha_test<C: ShaConfig + ShaPrecomputedValues<C::Word> + 'static>() {
     let bitwise_chip = SharedBitwiseOperationLookupChip::<RV32_CELL_BITS>::new(bitwise_bus);
     let len = rng.gen_range(1..100);
     let random_records: Vec<_> = (0..len)
-        .map(|_| {
+        .map(|i| {
             (
                 (0..C::BLOCK_U8S)
                     .map(|_| rng.gen::<u8>())
@@ -137,7 +137,7 @@ fn rand_sha512_test() {
 pub struct ShaTestBadFinalHashChip<C: ShaConfig + ShaPrecomputedValues<C::Word>> {
     pub air: ShaTestAir<C>,
     pub bitwise_lookup_chip: SharedBitwiseOperationLookupChip<8>,
-    pub records: Vec<(Vec<u8>, bool)>, // length of inner vec is C::BLOCK_U8S
+    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>
@@ -150,7 +150,7 @@ where
     }
 
     fn generate_air_proof_input(self) -> AirProofInput<SC> {
-        let mut trace = crate::generate_trace::<Val<SC>>(
+        let mut trace = crate::generate_trace::<Val<SC>, C>(
             &self.air.sub_air,
             self.bitwise_lookup_chip.clone(),
             self.records.clone(),
@@ -161,7 +161,7 @@ where
         for (i, row) in self.records.iter().enumerate() {
             if row.1 {
                 let last_digest_row_idx = (i + 1) * C::ROWS_PER_BLOCK - 1;
-                let last_digest_row: crate::ShaDigestColsRefMut<Val<SC>> =
+                let mut last_digest_row: crate::ShaDigestColsRefMut<Val<SC>> =
                     ShaDigestColsRefMut::from::<C>(
                         trace.row_mut(last_digest_row_idx)[..C::DIGEST_WIDTH].borrow_mut(),
                     );
@@ -176,10 +176,10 @@ where
                     trace.row_pair_mut(last_digest_row_idx - 1, last_digest_row_idx);
                 let last_round_row: crate::ShaRoundColsRefMut<Val<SC>> =
                     ShaRoundColsRefMut::from::<C>(last_round_row.borrow_mut());
-                let last_digest_row: crate::ShaRoundColsRefMut<Val<SC>> =
+                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(last_round_row, last_digest_row);
+                generate_intermed_4::<Val<SC>, C>(&ShaRoundColsRef::from_mut::<C>(&last_round_row), &mut last_digest_row);
             }
         }
 
@@ -199,52 +199,72 @@ 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>(
-    local_cols: &Sha256RoundCols<F>,
-    next_cols: &mut Sha256RoundCols<F>,
+fn generate_intermed_4<F: PrimeField32, C: ShaConfig + ShaPrecomputedValues<C::Word>>(
+    local_cols: &ShaRoundColsRef<F>,
+    next_cols: &mut ShaRoundColsRefMut<F>,
 ) {
-    let w = [local_cols.message_schedule.w, next_cols.message_schedule.w].concat();
-    let w_limbs: Vec<[F; SHA256_WORD_U16S]> = w
+    let w = [
+        local_cols
+            .message_schedule
+            .w
+            .rows()
+            .into_iter()
+            .collect::<Vec<_>>(),
+        next_cols
+            .message_schedule
+            .w
+            .rows()
+            .into_iter()
+            .collect::<Vec<_>>(),
+    ]
+    .concat();
+ 
+    
+    // length of inner vec is C::WORD_U16S
+    let w_limbs: Vec<Vec<F>> = w 
         .iter()
-        .map(|x| array::from_fn(|i| compose::<F>(&x[i * 16..(i + 1) * 16], 1)))
+        .map(|x| {
+            (0..C::WORD_U16S)
+                .map(|i| compose::<F>(&x.as_slice().unwrap()[i * 16..(i + 1) * 16], 1))
+                .collect::<Vec<F>>()
+        })
         .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 i in 0..C::ROUNDS_PER_ROW {
+        let sig_w = small_sig0_field::<F, C>(w[i + 1].as_slice().unwrap());
+        let sig_w_limbs: Vec<F> = (0..C::WORD_U16S)
+            .map(|j| compose::<F>(&sig_w[j * 16..(j + 1) * 16], 1))
+            .collect();
         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;
+            next_cols.schedule_helper.intermed_4[[i, j]] = w_limbs[i][j] + *sig_w_limb;
         }
     }
 }
 
-impl ChipUsageGetter for Sha256TestBadFinalHashChip {
+impl<C: ShaConfig + ShaPrecomputedValues<C::Word>> ChipUsageGetter for ShaTestBadFinalHashChip<C> {
     fn air_name(&self) -> String {
         get_air_name(&self.air)
     }
     fn current_trace_height(&self) -> usize {
-        self.records.len() * SHA256_ROWS_PER_BLOCK
+        self.records.len() * C::ROWS_PER_BLOCK
     }
 
     fn trace_width(&self) -> usize {
-        max(SHA256_ROUND_WIDTH, SHA256_DIGEST_WIDTH)
+        max(C::ROUND_WIDTH, C::DIGEST_WIDTH)
     }
 }
 
-#[test]
-#[should_panic]
-fn test_sha256_final_hash_constraints() {
+fn test_sha_final_hash_constraints<C: ShaConfig + ShaPrecomputedValues<C::Word> + '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(|_| (array::from_fn(|_| rng.gen::<u8>()), true))
+        .map(|_| ((0..C::BLOCK_U8S).map(|_| rng.gen::<u8>()).collect::<Vec<_>>(), true))
         .collect();
-    let chip = Sha256TestBadFinalHashChip {
-        air: Sha256TestAir {
-            sub_air: Sha256Air::new(bitwise_bus, SELF_BUS_IDX),
+    let chip = ShaTestBadFinalHashChip {
+        air: ShaTestAir {
+            sub_air: ShaAir::<C>::new(bitwise_bus, SELF_BUS_IDX),
         },
         bitwise_lookup_chip: bitwise_chip.clone(),
         records: random_records,
@@ -253,3 +273,15 @@ fn test_sha256_final_hash_constraints() {
     let tester = tester.build().load(chip).load(bitwise_chip).finalize();
     tester.simple_test().expect("Verification failed");
 }
+
+#[test]
+#[should_panic]
+fn test_sha256_final_hash_constraints() {
+    test_sha_final_hash_constraints::<Sha256Config>();
+}
+
+#[test]
+#[should_panic]
+fn test_sha512_final_hash_constraints() {
+    test_sha_final_hash_constraints::<Sha512Config>();
+}
diff --git a/crates/circuits/sha-air/src/trace.rs b/crates/circuits/sha-air/src/trace.rs
@@ -339,9 +339,10 @@ impl<C: ShaConfig + ShaPrecomputedValues<C::Word>> ShaAir<C> {
                 let final_hash: Vec<C::Word> = (0..C::HASH_WORDS)
                     .map(|i| work_vars[i].wrapping_add(prev_hash[i]))
                     .collect();
-                let final_hash_limbs: Vec<Vec<C::Word>> = final_hash.iter().map(|i| 
-                    word_into_u8_limbs::<C>(final_hash[i])
-                ) 
+                let final_hash_limbs: Vec<Vec<u32>> = final_hash
+                    .iter()
+                    .map(|word| word_into_u8_limbs::<C>(*word))
+                    .collect();
                 // need to ensure final hash limbs are bytes, in order for
                 //   prev_hash[i] + work_vars[i] == final_hash[i]
                 // to be constrained correctly
@@ -368,7 +369,7 @@ impl<C: ShaConfig + ShaPrecomputedValues<C::Word>> ShaAir<C> {
                             .collect::<Vec<_>>()
                     }))
                     .for_each(|(x, y)| *x = y);
- 
+
                 let hash = if is_last_block {
                     C::get_h()
                         .iter()
@@ -476,11 +477,20 @@ impl<C: ShaConfig + ShaPrecomputedValues<C::Word>> ShaAir<C> {
             &mut next_block_first_row[trace_start_col..trace_start_col + C::ROUND_WIDTH],
         );
         // Fill in the last round row's `intermed_12` with dummy values so the message schedule constraints holds on row 16
-        Self::generate_intermed_12(&mut cols_last_round_row, ShaRoundColsRef::from_mut::<C>(&cols_digest_row));
+        Self::generate_intermed_12(
+            &mut cols_last_round_row,
+            ShaRoundColsRef::from_mut::<C>(&cols_digest_row),
+        );
         // Fill in the digest row's `intermed_12` with dummy values so the message schedule constraints holds on the next block's row 0
-        Self::generate_intermed_12(&mut cols_digest_row, ShaRoundColsRef::from_mut::<C>(&cols_next_block_first_row));
+        Self::generate_intermed_12(
+            &mut cols_digest_row,
+            ShaRoundColsRef::from_mut::<C>(&cols_next_block_first_row),
+        );
         // Fill in the next block's first row's `intermed_4` with dummy values so the message schedule constraints holds on that row
-        Self::generate_intermed_4(ShaRoundColsRef::from_mut::<C>(&cols_digest_row), &mut cols_next_block_first_row);
+        Self::generate_intermed_4(
+            ShaRoundColsRef::from_mut::<C>(&cols_digest_row),
+            &mut cols_next_block_first_row,
+        );
     }
 
     /// Fills the `cols` as a padding row
@@ -796,9 +806,8 @@ pub fn generate_trace<F: PrimeField32, C: ShaConfig + ShaPrecomputedValues<C::Wo
             let input_words = (0..C::BLOCK_WORDS)
                 .map(|i| {
                     limbs_into_word::<C>(
-                       &(0..C::WORD_U8S).map(|j|
-                            input[((i + 1) * C::WORD_U8S - j - 1]
-                        )
+                        &(0..C::WORD_U8S)
+                            .map(|j| input[(i + 1) * C::WORD_U8S - j - 1] as u32)
                             .collect::<Vec<_>>(),
                     )
                 })
diff --git a/crates/circuits/sha-air/src/utils.rs b/crates/circuits/sha-air/src/utils.rs
@@ -1,5 +1,3 @@
-use std::array;
-
 pub use openvm_circuit_primitives::utils::compose;
 use openvm_circuit_primitives::{
     encoder::Encoder,
diff --git a/crates/circuits/sha-macros/src/lib.rs b/crates/circuits/sha-macros/src/lib.rs
@@ -1,9 +1,6 @@
-#![feature(proc_macro_diagnostic)]
-
 extern crate proc_macro;
 
 use itertools::Itertools;
-//use openvm_macros_common::MacroArgs;
 use proc_macro::TokenStream;
 use quote::{format_ident, quote};
 use syn::{parse_macro_input, parse_quote, DeriveInput, Expr};

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,3 @@`
`1`		`-use std::array;`
`2`		`-`
`3`	`1`	`pub use openvm_circuit_primitives::utils::compose;`
`4`	`2`	`use openvm_circuit_primitives::{`
`5`	`3`	`encoder::Encoder,`