[feat] sha2 new design

Cargo.toml

@@ -51,9 +51,9 @@ members = [
     "extensions/keccak256/circuit",
     "extensions/keccak256/transpiler",
     "extensions/keccak256/guest",
-    "extensions/sha256/circuit",
-    "extensions/sha256/transpiler",
-    "extensions/sha256/guest",
+    "extensions/sha2/circuit",
+    "extensions/sha2/transpiler",
+    "extensions/sha2/guest",
     "extensions/ecc/circuit",
     "extensions/ecc/transpiler",
     "extensions/ecc/guest",
@@ -123,7 +123,7 @@ openvm-cuda-common = { git = "https://github.com/openvm-org/stark-backend.git",
 openvm-sdk = { path = "crates/sdk", default-features = false }
 openvm-mod-circuit-builder = { path = "crates/circuits/mod-builder", default-features = false }
 openvm-poseidon2-air = { path = "crates/circuits/poseidon2-air", default-features = false }
-openvm-sha256-air = { path = "crates/circuits/sha256-air", default-features = false }
+openvm-sha2-air = { path = "crates/circuits/sha2-air", default-features = false }
 openvm-circuit-primitives = { path = "crates/circuits/primitives", default-features = false }
 openvm-circuit-primitives-derive = { path = "crates/circuits/primitives/derive", default-features = false }
 openvm = { path = "crates/toolchain/openvm", default-features = false }
@@ -153,9 +153,9 @@ openvm-native-transpiler = { path = "extensions/native/transpiler", default-feat
 openvm-keccak256-circuit = { path = "extensions/keccak256/circuit", default-features = false }
 openvm-keccak256-transpiler = { path = "extensions/keccak256/transpiler", default-features = false }
 openvm-keccak256-guest = { path = "extensions/keccak256/guest", default-features = false }
-openvm-sha256-circuit = { path = "extensions/sha256/circuit", default-features = false }
-openvm-sha256-transpiler = { path = "extensions/sha256/transpiler", default-features = false }
-openvm-sha256-guest = { path = "extensions/sha256/guest", default-features = false }
+openvm-sha2-circuit = { path = "extensions/sha2/circuit", default-features = false }
+openvm-sha2-transpiler = { path = "extensions/sha2/transpiler", default-features = false }
+openvm-sha2-guest = { path = "extensions/sha2/guest", default-features = false }
 openvm-bigint-circuit = { path = "extensions/bigint/circuit", default-features = false }
 openvm-bigint-transpiler = { path = "extensions/bigint/transpiler", default-features = false }
 openvm-bigint-guest = { path = "extensions/bigint/guest", default-features = false }
@@ -234,6 +234,8 @@ memmap2 = "0.9.5"
 libc = "0.2.175"
 tracing-subscriber = { version = "0.3.20", features = ["std", "env-filter"] }
 tokio = "1" # >=1.0.0 to allow downstream flexibility
+ndarray = { version = "0.16.1", default-features = false }
+num_enum = { version = "0.7.4", default-features = false }
 
 # default-features = false for no_std for use in guest programs
 itertools = { version = "0.14.0", default-features = false }

benchmarks/execute/Cargo.toml

@@ -26,8 +26,8 @@ openvm-keccak256-circuit.workspace = true
 openvm-keccak256-transpiler.workspace = true
 openvm-rv32im-circuit.workspace = true
 openvm-rv32im-transpiler.workspace = true
-openvm-sha256-circuit.workspace = true
-openvm-sha256-transpiler.workspace = true
+openvm-sha2-circuit.workspace = true
+openvm-sha2-transpiler.workspace = true
 openvm-continuations = { workspace = true }
 openvm-native-recursion = { workspace = true }
 openvm-sdk = { workspace = true }

benchmarks/execute/benches/execute.rs

@@ -47,8 +47,8 @@ use openvm_sdk::{
     commit::VmCommittedExe,
     config::{AggregationConfig, DEFAULT_NUM_CHILDREN_INTERNAL, DEFAULT_NUM_CHILDREN_LEAF},
 };
-use openvm_sha256_circuit::{Sha256, Sha256Executor, Sha2CpuProverExt};
-use openvm_sha256_transpiler::Sha256TranspilerExtension;
+use openvm_sha2_circuit::{Sha2, Sha2CpuProverExt, Sha2Executor};
+use openvm_sha2_transpiler::Sha2TranspilerExtension;
 use openvm_stark_sdk::{
     config::{baby_bear_poseidon2::BabyBearPoseidon2Engine, FriParameters},
     engine::{StarkEngine, StarkFriEngine},
@@ -106,7 +106,7 @@ pub struct ExecuteConfig {
     #[extension]
     pub keccak: Keccak256,
     #[extension]
-    pub sha256: Sha256,
+    pub sha2: Sha2,
     #[extension]
     pub modular: ModularExtension,
     #[extension]
@@ -127,7 +127,7 @@ impl Default for ExecuteConfig {
             io: Rv32Io,
             bigint: Int256::default(),
             keccak: Keccak256,
-            sha256: Sha256,
+            sha2: Sha2,
             modular: ModularExtension::new(vec![
                 bn_config.modulus.clone(),
                 bn_config.scalar.clone(),
@@ -215,7 +215,7 @@ fn create_default_transpiler() -> Transpiler<BabyBear> {
         .with_extension(Rv32MTranspilerExtension)
         .with_extension(Int256TranspilerExtension)
         .with_extension(Keccak256TranspilerExtension)
-        .with_extension(Sha256TranspilerExtension)
+        .with_extension(Sha2TranspilerExtension)
         .with_extension(ModularTranspilerExtension)
         .with_extension(Fp2TranspilerExtension)
         .with_extension(EccTranspilerExtension)

benchmarks/guest/kitchen-sink/openvm.toml

@@ -2,7 +2,7 @@
 [app_vm_config.rv32m]
 [app_vm_config.io]
 [app_vm_config.keccak]
-[app_vm_config.sha256]
+[app_vm_config.sha2]
 [app_vm_config.bigint]
 
 [app_vm_config.modular]

benchmarks/guest/sha256/openvm.toml

@@ -1,4 +1,4 @@
 [app_vm_config.rv32i]
 [app_vm_config.rv32m]
 [app_vm_config.io]
-[app_vm_config.sha256]
+[app_vm_config.sha2]

benchmarks/guest/sha256_iter/openvm.toml

@@ -1,4 +1,4 @@
 [app_vm_config.rv32i]
 [app_vm_config.rv32m]
 [app_vm_config.io]
-[app_vm_config.sha256]
+[app_vm_config.sha2]

benchmarks/prove/src/bin/kitchen_sink.rs

@@ -67,6 +67,63 @@ fn verify_native_max_trace_heights(
     Ok(())
 }
 
+fn verify_native_max_trace_heights(
+    sdk: &Sdk,
+    app_pk: Arc<AppProvingKey<SdkVmConfig>>,
+    app_committed_exe: Arc<NonRootCommittedExe>,
+    leaf_vm_pk: Arc<VmProvingKey<SC, NativeConfig>>,
+    num_children_leaf: usize,
+) -> Result<()> {
+    let app_proof =
+        sdk.generate_app_proof(app_pk.clone(), app_committed_exe.clone(), StdIn::default())?;
+    let leaf_inputs =
+        LeafVmVerifierInput::chunk_continuation_vm_proof(&app_proof, num_children_leaf);
+    let vm_vk = leaf_vm_pk.vm_pk.get_vk();
+
+    leaf_inputs.iter().for_each(|leaf_input| {
+        let executor = {
+            let mut executor = SingleSegmentVmExecutor::new(leaf_vm_pk.vm_config.clone());
+            executor
+                .set_trace_height_constraints(leaf_vm_pk.vm_pk.trace_height_constraints.clone());
+            executor
+        };
+        let max_trace_heights = executor
+            .execute_metered(
+                app_pk.leaf_committed_exe.exe.clone(),
+                leaf_input.write_to_stream(),
+                &vm_vk.num_interactions(),
+            )
+            .expect("execute_metered failed");
+        println!("max_trace_heights: {:?}", max_trace_heights);
+
+        let actual_trace_heights = executor
+            .execute_and_generate(
+                app_pk.leaf_committed_exe.clone(),
+                leaf_input.write_to_stream(),
+                &max_trace_heights,
+            )
+            .expect("execute_and_generate failed")
+            .per_air
+            .iter()
+            .map(|(_, air)| air.raw.height())
+            .collect::<Vec<usize>>();
+        println!("actual_trace_heights: {:?}", actual_trace_heights);
+
+        actual_trace_heights
+            .iter()
+            .zip(NATIVE_MAX_TRACE_HEIGHTS)
+            .for_each(|(&actual, &expected)| {
+                assert!(
+                    actual <= (expected as usize),
+                    "Actual trace height {} exceeds expected height {}",
+                    actual,
+                    expected
+                );
+            });
+    });
+    Ok(())
+}
+
 fn main() -> Result<()> {
     let args = BenchmarkCli::parse();
 

crates/circuits/primitives/derive/Cargo.toml

@@ -12,6 +12,13 @@ license.workspace = true
 proc-macro = true
 
 [dependencies]
-syn = { version = "2.0", features = ["parsing"] }
+syn = { version = "2.0", features = ["full", "parsing", "extra-traits"] }
 quote = "1.0"
-itertools = { workspace = true }
+itertools = { workspace = true, default-features = true }
+proc-macro2 = "1.0"
+
+[dev-dependencies]
+ndarray.workspace = true
+
+[package.metadata.cargo-shear]
+ignored = ["ndarray"]

crates/circuits/primitives/derive/src/cols_ref/README.md

@@ -0,0 +1,113 @@
+# ColsRef macro
+
+The `ColsRef` procedural macro is used in constraint generation to create column structs that have dynamic sizes.
+
+Note: this macro was originally created for use in the SHA-2 VM extension, where we reuse the same constraint generation code for three different circuits (SHA-256, SHA-512, and SHA-384).
+See the [SHA-2 VM extension](../../../../../../extensions/sha2/circuit/src/sha2_chip/air.rs) for an example of how to use the `ColsRef` macro to reuse constraint generation code over multiple circuits.
+
+## Overview
+
+As an illustrative example, consider the following columns struct:
+```rust
+struct ExampleCols<T, const N: usize> {
+    arr: [T; N],
+    sum: T,
+}
+```
+Let's say we want to constrain `sum` to be the sum of the elements of `arr`, and `N` can be either 5 or 10.
+We can define a trait that stores the config parameters.
+```rust
+pub trait ExampleConfig {
+    const N: usize;
+}
+```
+and then implement it for the two different configs.
+```rust
+pub struct ExampleConfigImplA;
+impl ExampleConfig for ExampleConfigImplA {
+    const N: usize = 5;
+}
+pub struct ExampleConfigImplB;
+impl ExampleConfig for ExampleConfigImplB {
+    const N: usize = 10;
+}
+```
+Then we can use the `ColsRef` macro like this
+```rust
+#[derive(ColsRef)]
+#[config(ExampleConfig)]
+struct ExampleCols<T, const N: usize> {
+    arr: [T; N],
+    sum: T,
+}
+```
+which will generate a columns struct that uses references to the fields.
+```rust
+struct ExampleColsRef<'a, T, const N: usize> {
+    arr: ndarray::ArrayView1<'a, T>, // an n-dimensional view into the input slice (ArrayView2 for 2D arrays, etc.)
+    sum: &'a T,
+}
+```
+The `ColsRef` macro will also generate a `from` method that takes a slice of the correct length and returns an instance of the columns struct.
+The `from` method is parameterized by a struct that implements the `ExampleConfig` trait, and it uses the associated constants to determine how to split the input slice into the fields of the columns struct.
+
+So, the constraint generation code can be written as
+```rust
+impl<AB: InteractionBuilder, C: ExampleConfig> Air<AB> for ExampleAir<C> {
+    fn eval(&self, builder: &mut AB) {
+        let main = builder.main();
+        let (local, _) = (main.row_slice(0), main.row_slice(1));
+        let local_cols = ExampleColsRef::<AB::Var>::from::<C>(&local[..C::N + 1]);
+        let sum = local_cols.arr.iter().sum();
+        builder.assert_eq(local_cols.sum, sum);
+    }
+}
+```
+Notes:
+- the `arr` and `sum` fields of `ExampleColsRef` are references to the elements of the `local` slice.
+- the name, `N`, of the const generic parameter must match the name of the associated constant `N` in the `ExampleConfig` trait.
+
+The `ColsRef` macro also generates a `ExampleColsRefMut` struct that stores mutable references to the fields, for use in trace generation.
+
+The `ColsRef` macro supports more than just variable-length array fields.
+The field types can also be:
+- any type that derives `AlignedBorrow` via `#[derive(AlignedBorrow)]`
+- any type that derives `ColsRef` via `#[derive(ColsRef)]`
+- (possibly nested) arrays of `T` or (possibly nested) arrays of a type that derives `AlignedBorrow`
+
+Note that we currently do not support arrays of types that derive `ColsRef`.
+
+## Specification
+
+Annotating a struct named `ExampleCols` with `#[derive(ColsRef)]` and `#[config(ExampleConfig)]` produces two structs, `ExampleColsRef` and `ExampleColsRefMut`.
+- we assume `ExampleCols` has exactly one generic type parameter, typically named `T`, and any number of const generic parameters. Each const generic parameter must have a name that matches an associated constant in the `ExampleConfig` trait
+
+The fields of `ExampleColsRef` have the same names as the fields of `ExampleCols`, but their types are transformed as follows:
+- type `T` becomes `&T`
+- type `[T; LEN]` becomes `&ArrayView1<T>` (see [ndarray](https://docs.rs/ndarray/latest/ndarray/index.html)) where `LEN` is an associated constant in `ExampleConfig`
+    - the `ExampleColsRef::from` method will correctly infer the length of the array from the config
+- fields with names that end in `Cols` are assumed to be a columns struct that derives `ColsRef` and are transformed into the appropriate `ColsRef` type recursively
+    - one restriction is that any nested `ColsRef` type must have the same config as the outer `ColsRef` type
+- fields that are annotated with `#[aligned_borrow]` are assumed to derive `AlignedBorrow` and are borrowed from the input slice. The new type is a reference to the `AlignedBorrow` type
+    - if a field whose name ends in `Cols` is annotated with `#[aligned_borrow]`, then the aligned borrow takes precedence, and the field is not transformed into an `ArrayView`
+- nested arrays of `U` become `&ArrayViewX<U>` where `X` is the number of dimensions in the nested array type
+    - `U` can be either the generic type `T` or a type that derives `AlignedBorrow`. In the latter case, the field must be annotated with `#[aligned_borrow]`
+    - the `ArrayViewX` type provides a `X`-dimensional view into the row slice 
+
+The fields of `ExampleColsRefMut` are almost the same as the fields of `ExampleColsRef`, but they are mutable references.
+- the `ArrayViewMutX` type is used instead of `ArrayViewX` for the array fields.
+- fields that derive `ColsRef` are transformed into the appropriate `ColsRefMut` type recursively.
+
+Each of the `ExampleColsRef` and `ExampleColsRefMut` types has the following methods implemented:
+```rust
+// Takes a slice of the correct length and returns an instance of the columns struct.
+pub const fn from<C: ExampleConfig>(slice: &[T]) -> Self;
+// Returns the number of cells in the struct
+pub const fn width<C: ExampleConfig>() -> usize;
+```
+Note that the `width` method on both structs returns the same value.
+
+Additionally, the `ExampleColsRef` struct has a `from_mut` method that takes a `ExampleColsRefMut` and returns a `ExampleColsRef`.
+This may be useful in trace generation to pass a `ExampleColsRefMut` to a function that expects a `ExampleColsRef`.
+
+See the [tests](../../tests/test_cols_ref.rs) for concrete examples of how the `ColsRef` macro handles each of the supported field types.