Merge ShaPrecomputedValues trait into ShaConfig trait for simplicity

Author: Avaneesh-axiom

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

@@ -19,21 +19,20 @@ use super::{
     small_sig1_field,
 };
 use crate::{
-    constraint_word_addition, word_into_u16_limbs, ShaConfig, ShaDigestColsRef,
-    ShaPrecomputedValues, ShaRoundColsRef,
+    constraint_word_addition, word_into_u16_limbs, ShaConfig, ShaDigestColsRef, ShaRoundColsRef,
 };
 
 /// 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 + ShaPrecomputedValues<C::Word>> {
+pub struct ShaAir<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 + ShaPrecomputedValues<C::Word>> ShaAir<C> {
+impl<C: ShaConfig> ShaAir<C> {
     pub fn new(bitwise_lookup_bus: BitwiseOperationLookupBus, self_bus_idx: BusIndex) -> Self {
         Self {
             bitwise_lookup_bus,
@@ -44,15 +43,13 @@ impl<C: ShaConfig + ShaPrecomputedValues<C::Word>> ShaAir<C> {
     }
 }
 
-impl<F, C: ShaConfig + ShaPrecomputedValues<C::Word>> BaseAir<F> for ShaAir<C> {
+impl<F, C: ShaConfig> BaseAir<F> for ShaAir<C> {
     fn width(&self) -> usize {
         max(C::ROUND_WIDTH, C::DIGEST_WIDTH)
     }
 }
 
-impl<AB: InteractionBuilder, C: ShaConfig + ShaPrecomputedValues<C::Word>> SubAir<AB>
-    for ShaAir<C>
-{
+impl<AB: InteractionBuilder, C: ShaConfig> SubAir<AB> for ShaAir<C> {
     /// The start column for the sub-air to use
     type AirContext<'a>
         = usize
@@ -72,7 +69,7 @@ impl<AB: InteractionBuilder, C: ShaConfig + ShaPrecomputedValues<C::Word>> SubAi
     }
 }
 
-impl<C: ShaConfig + ShaPrecomputedValues<C::Word>> ShaAir<C> {
+impl<C: ShaConfig> ShaAir<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) {

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

@@ -20,6 +20,7 @@ use crate::ShaConfig;
 /// 2. Specific constraints to use the appropriate struct, with flags helping to do conditional constraints
 #[repr(C)]
 #[derive(Clone, Copy, Debug, ColsRef)]
+#[config(ShaConfig)]
 pub struct ShaRoundCols<
     T,
     const WORD_BITS: usize,
@@ -38,6 +39,7 @@ pub struct ShaRoundCols<
 
 #[repr(C)]
 #[derive(Clone, Copy, Debug, ColsRef)]
+#[config(ShaConfig)]
 pub struct ShaDigestCols<
     T,
     const WORD_BITS: usize,
@@ -63,6 +65,7 @@ pub struct ShaDigestCols<
 
 #[repr(C)]
 #[derive(Clone, Copy, Debug, ColsRef)]
+#[config(ShaConfig)]
 pub struct ShaMessageScheduleCols<
     T,
     const WORD_BITS: usize,
@@ -78,6 +81,7 @@ pub struct ShaMessageScheduleCols<
 
 #[repr(C)]
 #[derive(Clone, Copy, Debug, ColsRef)]
+#[config(ShaConfig)]
 pub struct ShaWorkVarsCols<
     T,
     const WORD_BITS: usize,
@@ -96,6 +100,7 @@ pub struct ShaWorkVarsCols<
 /// Note: these need to be correctly assigned for every row even on padding rows
 #[repr(C)]
 #[derive(Clone, Copy, Debug, ColsRef)]
+#[config(ShaConfig)]
 pub struct ShaMessageHelperCols<
     T,
     const WORD_U16S: usize,
@@ -114,6 +119,7 @@ pub struct ShaMessageHelperCols<
 
 #[repr(C)]
 #[derive(Clone, Copy, Debug, ColsRef)]
+#[config(ShaConfig)]
 pub struct ShaFlagsCols<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

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

@@ -86,8 +86,8 @@ pub trait ShaPrecomputedValues<T> {
     // 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];
-    fn get_h() -> &'static [T];
+    fn get_k() -> &'static [Self::Word];
+    fn get_h() -> &'static [Self::Word];
 }
 
 #[derive(Clone)]
@@ -110,6 +110,19 @@ impl ShaConfig for Sha256Config {
     const HASH_WORDS: usize = 8;
     /// Number of vars needed to encode the row index with [Encoder]
     const ROW_VAR_CNT: usize = 5;
+
+    fn get_invalid_carry_a(round_num: usize) -> &'static [u32] {
+        &SHA256_INVALID_CARRY_A[round_num]
+    }
+    fn get_invalid_carry_e(round_num: usize) -> &'static [u32] {
+        &SHA256_INVALID_CARRY_E[round_num]
+    }
+    fn get_k() -> &'static [u32] {
+        &SHA256_K
+    }
+    fn get_h() -> &'static [u32] {
+        &SHA256_H
+    }
 }
 
 pub const SHA256_INVALID_CARRY_A: [[u32; Sha256Config::WORD_U16S]; Sha256Config::ROUNDS_PER_ROW] = [
@@ -141,21 +154,6 @@ pub const SHA256_H: [u32; 8] = [
     0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19,
 ];
 
-impl ShaPrecomputedValues<u32> for Sha256Config {
-    fn get_invalid_carry_a(round_num: usize) -> &'static [u32] {
-        &SHA256_INVALID_CARRY_A[round_num]
-    }
-    fn get_invalid_carry_e(round_num: usize) -> &'static [u32] {
-        &SHA256_INVALID_CARRY_E[round_num]
-    }
-    fn get_k() -> &'static [u32] {
-        &SHA256_K
-    }
-    fn get_h() -> &'static [u32] {
-        &SHA256_H
-    }
-}
-
 #[derive(Clone)]
 pub struct Sha512Config;
 
@@ -176,6 +174,19 @@ impl ShaConfig for Sha512Config {
     const HASH_WORDS: usize = 8;
     /// Number of vars needed to encode the row index with [Encoder]
     const ROW_VAR_CNT: usize = 6;
+
+    fn get_invalid_carry_a(round_num: usize) -> &'static [u32] {
+        &SHA512_INVALID_CARRY_A[round_num]
+    }
+    fn get_invalid_carry_e(round_num: usize) -> &'static [u32] {
+        &SHA512_INVALID_CARRY_E[round_num]
+    }
+    fn get_k() -> &'static [u64] {
+        &SHA512_K
+    }
+    fn get_h() -> &'static [u64] {
+        &SHA512_H
+    }
 }
 
 pub(crate) const SHA512_INVALID_CARRY_A: [[u32; Sha512Config::WORD_U16S];
@@ -289,21 +300,6 @@ pub const SHA512_H: [u64; 8] = [
     0x5be0cd19137e2179,
 ];
 
-impl ShaPrecomputedValues<u64> for Sha512Config {
-    fn get_invalid_carry_a(round_num: usize) -> &'static [u32] {
-        &SHA512_INVALID_CARRY_A[round_num]
-    }
-    fn get_invalid_carry_e(round_num: usize) -> &'static [u32] {
-        &SHA512_INVALID_CARRY_E[round_num]
-    }
-    fn get_k() -> &'static [u64] {
-        &SHA512_K
-    }
-    fn get_h() -> &'static [u64] {
-        &SHA512_H
-    }
-}
-
 // Needed to avoid compile errors in utils.rs
 // not sure why this doesn't inf loop
 pub trait RotateRight {

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

@@ -23,46 +23,37 @@ use openvm_stark_sdk::utils::create_seeded_rng;
 use rand::Rng;
 
 use crate::{
-    compose, small_sig0_field, Sha256Config, Sha512Config, ShaAir, ShaConfig, ShaPrecomputedValues,
+    compose, small_sig0_field, Sha256Config, Sha512Config, ShaAir, ShaConfig,
 };
 
 // A wrapper AIR purely for testing purposes
 #[derive(Clone, Debug)]
-pub struct ShaTestAir<C: ShaConfig + ShaPrecomputedValues<C::Word>> {
+pub struct ShaTestAir<C: ShaConfig> {
     pub sub_air: ShaAir<C>,
 }
 
-impl<F: Field, C: ShaConfig + ShaPrecomputedValues<C::Word>> BaseAirWithPublicValues<F>
-    for ShaTestAir<C>
-{
-}
-impl<F: Field, C: ShaConfig + ShaPrecomputedValues<C::Word>> PartitionedBaseAir<F>
-    for ShaTestAir<C>
-{
-}
-impl<F: Field, C: ShaConfig + ShaPrecomputedValues<C::Word>> BaseAir<F> for ShaTestAir<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)
     }
 }
 
-impl<AB: InteractionBuilder, C: ShaConfig + ShaPrecomputedValues<C::Word>> Air<AB>
-    for ShaTestAir<C>
-{
+impl<AB: InteractionBuilder, C: ShaConfig> Air<AB> for ShaTestAir<C> {
     fn eval(&self, builder: &mut AB) {
         self.sub_air.eval(builder, 0);
     }
 }
 
 // A wrapper Chip purely for testing purposes
-pub struct ShaTestChip<C: ShaConfig + ShaPrecomputedValues<C::Word>> {
+pub struct ShaTestChip<C: ShaConfig> {
     pub air: ShaTestAir<C>,
     pub bitwise_lookup_chip: SharedBitwiseOperationLookupChip<8>,
     pub records: Vec<(Vec<u8>, bool)>, // length of inner vec is BLOCK_U8S
 }
 
-impl<SC: StarkGenericConfig, C: ShaConfig + ShaPrecomputedValues<C::Word> + 'static> Chip<SC>
-    for ShaTestChip<C>
+impl<SC: StarkGenericConfig, C: ShaConfig + 'static> Chip<SC> for ShaTestChip<C>
 where
     Val<SC>: PrimeField32,
 {
@@ -80,7 +71,7 @@ where
     }
 }
 
-impl<C: ShaConfig + ShaPrecomputedValues<C::Word>> ChipUsageGetter for ShaTestChip<C> {
+impl<C: ShaConfig> ChipUsageGetter for ShaTestChip<C> {
     fn air_name(&self) -> String {
         get_air_name(&self.air)
     }
@@ -94,7 +85,7 @@ impl<C: ShaConfig + ShaPrecomputedValues<C::Word>> ChipUsageGetter for ShaTestCh
 }
 
 const SELF_BUS_IDX: BusIndex = 28;
-fn rand_sha_test<C: ShaConfig + ShaPrecomputedValues<C::Word> + 'static>() {
+fn rand_sha_test<C: ShaConfig + 'static>() {
     let mut rng = create_seeded_rng();
     let tester = VmChipTestBuilder::default();
     let bitwise_bus = BitwiseOperationLookupBus::new(BITWISE_OP_LOOKUP_BUS);

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

@@ -15,14 +15,14 @@ use super::{
 };
 use crate::{
     big_sig0, big_sig1, ch, limbs_into_word, maj, small_sig0, small_sig1, word_into_bits,
-    word_into_u16_limbs, word_into_u8_limbs, ShaConfig, ShaDigestColsRefMut, ShaPrecomputedValues,
-    ShaRoundColsRef, WrappingAdd,
+    word_into_u16_limbs, word_into_u8_limbs, ShaConfig, ShaDigestColsRefMut, ShaRoundColsRef,
+    WrappingAdd,
 };
 
 /// 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> {
+impl<C: ShaConfig> ShaAir<C> {
     /// This function takes the input_message (padding not handled), the previous hash,
     /// and returns the new hash after processing the block input
     pub fn get_block_hash(prev_hash: &[C::Word], input: Vec<u8>) -> Vec<C::Word> {
@@ -749,7 +749,7 @@ impl<C: ShaConfig + ShaPrecomputedValues<C::Word>> ShaAir<C> {
 }
 
 /// `records` consists of pairs of `(input_block, is_last_block)`.
-pub fn generate_trace<F: PrimeField32, C: ShaConfig + ShaPrecomputedValues<C::Word>>(
+pub fn generate_trace<F: PrimeField32, C: ShaConfig>(
     sub_air: &ShaAir<C>,
     bitwise_lookup_chip: SharedBitwiseOperationLookupChip<8>,
     records: Vec<(Vec<u8>, bool)>,