WIP: update vm chip to support sha512

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"
 
 # default-features = false for no_std for use in guest programs
 itertools = { version = "0.14.0", default-features = false }

Cargo.toml

@@ -77,15 +77,14 @@ pub trait ShaConfig: Send + Sync + Clone {
     } else {
         Self::DIGEST_WIDTH
     };
-}
 
-/// 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> {
+    /// 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.
     // 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,9 +22,7 @@ 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, ShaAir, ShaConfig};
 
 // A wrapper AIR purely for testing purposes
 #[derive(Clone, Debug)]
@@ -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,7 +167,10 @@ 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,
+                );
             }
         }
 
@@ -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,14 +243,21 @@ 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 {

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

@@ -292,7 +292,7 @@ fn get_const_cols_ref_fields(
     let slice_var = format_ident!("{}_slice", f.ident.clone().unwrap());
 
     if f.attrs.iter().any(|attr| attr.path().is_ident("plain")) {
-        // treat the field as a struct that derives AlignedBorrow
+        // treat the field as a struct that derives AlignedBorrow (and doesn't depend on the config)
         let f_ty = &f.ty;
         return Ok(FieldInfo {
             ty: parse_quote! {

crates/circuits/sha-macros/src/lib.rs

@@ -16,13 +16,15 @@ openvm-instructions = { workspace = true }
 openvm-sha256-transpiler = { workspace = true }
 openvm-rv32im-circuit = { workspace = true }
 openvm-sha-air = { workspace = true }
+openvm-sha-macros = { workspace = true }
 
 derive-new.workspace = true
 derive_more = { workspace = true, features = ["from"] }
 rand.workspace = true
 serde.workspace = true
 sha2 = { version = "0.10", default-features = false }
 strum = { workspace = true }
+ndarray.workspace = true
 
 [dev-dependencies]
 openvm-stark-sdk = { workspace = true }

extensions/sha256/circuit/Cargo.toml

@@ -17,14 +17,15 @@ use openvm_rv32im_circuit::{
     Rv32MExecutor, Rv32MPeriphery,
 };
 use openvm_sha256_transpiler::Rv32Sha256Opcode;
+use openvm_sha_air::Sha256Config;
 use openvm_stark_backend::p3_field::PrimeField32;
 use serde::{Deserialize, Serialize};
 use strum::IntoEnumIterator;
 
 use crate::*;
 
 #[derive(Clone, Debug, VmConfig, derive_new::new, Serialize, Deserialize)]
-pub struct Sha256Rv32Config {
+pub struct Sha2Rv32Config {
     #[system]
     pub system: SystemConfig,
     #[extension]
@@ -34,38 +35,39 @@ pub struct Sha256Rv32Config {
     #[extension]
     pub io: Rv32Io,
     #[extension]
-    pub sha256: Sha256,
+    pub sha2: Sha2,
 }
 
-impl Default for Sha256Rv32Config {
+impl Default for Sha2Rv32Config {
     fn default() -> Self {
         Self {
             system: SystemConfig::default().with_continuations(),
             rv32i: Rv32I,
             rv32m: Rv32M::default(),
             io: Rv32Io,
-            sha256: Sha256,
+            sha2: Sha2,
         }
     }
 }
 
 #[derive(Clone, Copy, Debug, Default, Serialize, Deserialize)]
-pub struct Sha256;
+pub struct Sha2;
 
 #[derive(ChipUsageGetter, Chip, InstructionExecutor, From, AnyEnum)]
-pub enum Sha256Executor<F: PrimeField32> {
-    Sha256(Sha256VmChip<F>),
+pub enum Sha2Executor<F: PrimeField32> {
+    Sha256(ShaVmChip<F, Sha256Config>),
+    Sha512(ShaVmChip<F, Sha512Config>),
 }
 
 #[derive(From, ChipUsageGetter, Chip, AnyEnum)]
-pub enum Sha256Periphery<F: PrimeField32> {
+pub enum ShaPeriphery<F: PrimeField32> {
     BitwiseOperationLookup(SharedBitwiseOperationLookupChip<8>),
     Phantom(PhantomChip<F>),
 }
 
-impl<F: PrimeField32> VmExtension<F> for Sha256 {
-    type Executor = Sha256Executor<F>;
-    type Periphery = Sha256Periphery<F>;
+impl<F: PrimeField32> VmExtension<F> for Sha2 {
+    type Executor = ShaExecutor<F>;
+    type Periphery = ShaPeriphery<F>;
 
     fn build(
         &self,
@@ -84,7 +86,7 @@ impl<F: PrimeField32> VmExtension<F> for Sha256 {
             chip
         };
 
-        let sha256_chip = Sha256VmChip::new(
+        let sha256_chip = ShaVmChip::new(
             builder.system_port(),
             builder.system_config().memory_config.pointer_max_bits,
             bitwise_lu_chip,