Update docs related to SHA-2

Author: Avaneesh-axiom

docs/vocs/docs/pages/book/acceleration-using-extensions/sha-2.mdx

@@ -1,10 +1,11 @@
 # SHA-2
 
-The SHA-2 extension guest provides functions that are meant to be linked to other external libraries. The external libraries can use these functions as a hook for SHA-2 intrinsics. This is enabled only when the target is `zkvm`. We support the SHA-256, SHA-512, and SHA-384 hash functions.
+The SHA-2 extension guest provides functions that are meant to be linked to other external libraries. The external libraries can use these functions as a hook for SHA-2 intrinsics. This is enabled only when the target is `zkvm`. We support the SHA-256 and SHA-512 hash functions.
 
-- `zkvm_shaXXX_impl(input: *const u8, len: usize, output: *mut u8)` where XXX is one of `256`, `512`, or `384`. These functions have `C` ABI. They take in a pointer to the input, the length of the input, and a pointer to the output buffer.
+We provide the following functions to compute the SHA-2 compression function.
+- `zkvm_shaXXX_impl(state: *const u8, input: *const u8, output: *mut u8)` where `XXX` is `256` or `512`. These functions have `C` ABI. They take in a pointer to the current hasher state (`state`), a pointer to the next block of the message (`input`), and a pointer where the new hasher state will be written (`output`). `state` is expected to be a pointer to 8 little-endian words, even though its type is `*const u8`. For Sha256 that means it is a pointer to `[u32; 8]`, and for Sha512 it's `[u64; 8]`.
 
-In the external library, you can do the following:
+In the external library, you can do something like the following:
 
 ```rust
 extern "C" {
@@ -14,11 +15,20 @@ extern "C" {
 fn sha256(input: &[u8]) -> [u8; 32] {
     #[cfg(target_os = "zkvm")]
     {
-        let mut output = [0u8; 32];
-        unsafe {
-            zkvm_sha256_impl(input.as_ptr(), input.len(), output.as_mut_ptr() as *mut u8);
-        }
-        output
+        let mut state = [0u32; 8];
+        let padded_input = add_padding(input);
+        padded_input
+            .chunks_exact(32)
+            .for_each(|input_block| {
+                unsafe {
+                    zkvm_sha256_impl(state.as_ptr() as *const u8, input_block.as_ptr(), output.as_mut_ptr() as *mut u8);
+                }
+            })
+        state
+            .map(|word| word.to_be_bytes())
+            .collect::<Vec<_>>()
+            .try_into()
+            .unwrap()
     }
     #[cfg(not(target_os = "zkvm"))] {
         // Regular SHA-256 implementation

docs/vocs/docs/pages/book/getting-started/introduction.mdx

@@ -12,7 +12,7 @@ OpenVM is an open-source zero-knowledge virtual machine (zkVM) framework focused
 
   - RISC-V support via RV32IM
   - A native field arithmetic extension for proof recursion and aggregation
-  - The Keccak-256, SHA-256, SHA-512, and SHA-384 hash functions
+  - The Keccak-256, SHA-256 and SHA-512 hash functions
   - Int256 arithmetic
   - Modular arithmetic over arbitrary fields
   - Elliptic curve operations, including multi-scalar multiplication and ECDSA signature verification, including for the secp256k1 and secp256r1 curves

docs/vocs/docs/pages/book/guest-libraries/sha2.mdx

@@ -6,34 +6,30 @@ The OpenVM SHA-2 guest library provides access to a set of accelerated SHA-2 fam
 - SHA-512
 - SHA-384
 
-Refer [here](https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.180-4.pdf) for more details on the SHA-2 family of hash functions.
-
-For SHA-256, the SHA2 guest library provides two functions for use in your guest code:
-- `sha256(input: &[u8]) -> [u8; 32]`: Computes the SHA-256 hash of the input data and returns it as an array of 32 bytes.
-- `set_sha256(input: &[u8], output: &mut [u8; 32])`: Sets the output to the SHA-256 hash of the input data into the provided output buffer.
-
-For SHA-512, we provide:
-- `sha512(input: &[u8]) -> [u8; 46]`: Computes the SHA-512 hash of the input data and returns it as an array of 64 bytes.
-- `set_sha512(input: &[u8], output: &mut [u8; 64])`: Sets the output to the SHA-512 hash of the input data into the provided output buffer.
-
-For SHA-384, we provide:
-- `sha384(input: &[u8]) -> [u8; 48]`: Computes the SHA-384 hash of the input data and returns it as an array of 48 bytes.
-- `set_sha384(input: &[u8], output: &mut [u8; 48])`: Sets the output to the SHA-384 hash of the input data into the provided output buffer.
-
-See the full example [here](https://github.com/openvm-org/openvm/blob/feat/sha-512-new-execution/examples/sha2/src/main.rs).
-
+Refer to [the FIPS publication](https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.180-4.pdf) for more details on the SHA-2 family of hash functions.
+
+The SHA-2 guest library provides the `Sha256`, `Sha512`, and `Sha384` structs for use in your guest code.
+These structs mimic the identically-named structs found in the `sha2` crate, in that they implement the [`sha2::Digest` trait](https://docs.rs/sha2/latest/sha2/trait.Digest.html).
+Importantly, the `sha2::Digest` trait has the following methods, which can be used to incrementally hash a stream of bytes.
+```Rust
+fn update(&mut self, data: impl AsRef<[u8]>);
+fn finalize(self) -> GenericArray<u8, Self::OutputSize>;
+```
 ### Example
 
+The following example can be run as guest code or host code (i.e. run in the zkvm or natively).
+To run with guest code, use `cargo openvm run`, and for host code use `cargo run`.
+The implementations for `Sha256`, `Sha512`, and `Sha384` fall back to using `sha2` when running as host code.
 ```rust
-// [!include ~/snippets/examples/sha2/src/main.rs:imports]
-// [!include ~/snippets/examples/sha2/src/main.rs:main]
+[!include ~/snippets/examples/sha2/src/main.rs:imports]
+[!include ~/snippets/examples/sha2/src/main.rs:main]
 ```
 
 To be able to import the `shaXXX` functions and run the example, add the following to your `Cargo.toml` file:
 
 ```toml
-openvm-sha2 = { git = "https://github.com/openvm-org/openvm.git", tag = "v1.4.1" }
-hex = { version = "0.4.3" }
+openvm = { git = "https://github.com/openvm-org/openvm.git" }
+openvm-sha2 = { git = "https://github.com/openvm-org/openvm.git" }
 ```
 
 ### Config parameters

docs/vocs/docs/pages/specs/openvm/isa.mdx

@@ -6,7 +6,7 @@ This specification describes the overall architecture and default VM extensions
 - [RV32IM](#rv32im-extension): An extension supporting the 32-bit RISC-V ISA with multiplication.
 - [Native](#native-extension): An extension supporting native field arithmetic for proof recursion and aggregation.
 - [Keccak-256](#keccak-extension): An extension implementing the Keccak-256 hash function compatibly with RISC-V memory.
-- [SHA2](#sha-2-extension): An extension implementing the SHA-256, SHA-512, and SHA-384 hash functions compatibly with RISC-V memory.
+- [SHA2](#sha-2-extension): An extension implementing the SHA-256 and SHA-512 hash functions compatibly with RISC-V memory.
 - [BigInt](#bigint-extension): An extension supporting 256-bit signed and unsigned integer arithmetic, including
   multiplication. This extension respects the RISC-V memory format.
 - [Algebra](#algebra-extension): An extension supporting modular arithmetic over arbitrary fields and their complex
@@ -557,9 +557,8 @@ meaning all memory cells are constrained to be bytes.
 
 | Name        | Operands    | Description                                                                                                                                                              |
 | ----------- | ----------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
-| SHA256_RV32 | `a,b,c,1,2` | `[r32{0}(a):32]_2 = sha256([r32{0}(b)..r32{0}(b)+r32{0}(c)]_2)`. Does the necessary padding. Performs memory reads with block size `16` and writes with block size `32`. |
-| SHA512_RV32 | `a,b,c,1,2` | `[r32{0}(a):64]_2 = sha512([r32{0}(b)..r32{0}(b)+r32{0}(c)]_2)`. Does the necessary padding. Performs memory reads with block size `32` and writes with block size `32`. |
-| SHA384_RV32 | `a,b,c,1,2` | `[r32{0}(a):64]_2 = sha384([r32{0}(b)..r32{0}(b)+r32{0}(c)]_2)`. Does the necessary padding. Performs memory reads with block size `32` and writes with block size `32`. Writes 64 bytes to memory: the first 48 are the SHA-384 digest and the last 16 are zeros. |
+| SHA256_UPDATE_RV32 | `a,b,c,1,2` | `[r32{0}(a):32]_2 = compress256([r32{0}(b):32]_2, [r32{0}(c):64]_2)`. where `compress256(state, input)` is the SHA-256 block compression function which returns the updated state. This instruction performs memory reads and writes with block size `4`. |
+| SHA512_UPDATE_RV32 | `a,b,c,1,2` | `[r32{0}(a):64]_2 = compress512([r32{0}(b):64]_2, [r32{0}(c):128]_2)`. where `compress512(state, input)` is the SHA-512 block compression function which returns the updated state. This instruction performs memory reads and writes with block size `4`. |
 
 ### BigInt Extension
 

docs/vocs/docs/pages/specs/reference/instruction-reference.mdx

@@ -134,11 +134,10 @@ In the tables below, we provide the mapping between the `LocalOpcode` and `Phant
 
 #### Instructions
 
-| VM Extension | `LocalOpcode` | ISA Instruction |
-| ------------- | ---------- | ------------- |
-| SHA-2 | `Rv32Sha2Opcode::SHA256` | SHA256_RV32 |
-| SHA-2 | `Rv32Sha2Opcode::SHA512` | SHA512_RV32 |
-| SHA-2 | `Rv32Sha2Opcode::SHA384` | SHA384_RV32 |
+| VM Extension  | `LocalOpcode`            | ISA Instruction          |
+| ------------- | ------------------------ | ------------------------ |
+| SHA-2         | `Rv32Sha2Opcode::SHA256` | SHA256_UPDATE_RV32       |
+| SHA-2         | `Rv32Sha2Opcode::SHA512` | SHA512_UPDATE_RV32       |
 
 ## BigInt Extension
 

docs/vocs/docs/pages/specs/reference/riscv-custom-code.mdx

@@ -5,7 +5,7 @@ The default VM extensions that support transpilation are:
 
 - [RV32IM](#rv32im-extension): An extension supporting the 32-bit RISC-V ISA with multiplication.
 - [Keccak-256](#keccak-extension): An extension implementing the Keccak-256 hash function compatibly with RISC-V memory.
-- [SHA2](#sha-2-extension): An extension implementing the SHA-256, SHA-512, and SHA-384 hash functions compatibly with RISC-V memory.
+- [SHA2](#sha-2-extension): An extension implementing the SHA-256 and SHA-512 hash functions compatibly with RISC-V memory.
 - [BigInt](#bigint-extension): An extension supporting 256-bit signed and unsigned integer arithmetic, including multiplication. This extension respects the RISC-V memory format.
 - [Algebra](#algebra-extension): An extension supporting modular arithmetic over arbitrary fields and their complex field extensions. This extension respects the RISC-V memory format.
 - [Elliptic curve](#elliptic-curve-extension): An extension for elliptic curve operations over Weierstrass curves, including addition and doubling. This can be used to implement multi-scalar multiplication and ECDSA scalar multiplication. This extension respects the RISC-V memory format.
@@ -87,11 +87,10 @@ implementation is here. But we use `funct3 = 111` because the native extension h
 
 ## SHA-2 Extension
 
-| RISC-V Inst | FMT | opcode[6:0] | funct3 | funct7 | RISC-V description and notes             |
-| ----------- | --- | ----------- | ------ | ------ | ---------------------------------------- |
-| sha256      | R   | 0001011     | 100    | 0x1    | `[rd:32]_2 = sha256([rs1..rs1 + rs2]_2)` |
-| sha512      | R   | 0001011     | 100    | 0x2    | `[rd:64]_2 = sha512([rs1..rs1 + rs2]_2)` |
-| sha384      | R   | 0001011     | 100    | 0x3    | `[rd:64]_2 = sha384([rs1..rs1 + rs2]_2)`. Last 16 bytes will be set to zeros. |
+| RISC-V Inst   | FMT | opcode[6:0] | funct3 | funct7 | RISC-V description and notes                       |
+| ------------- | --- | ----------- | ------ | ------ | -------------------------------------------------- |
+| sha256_update | R   | 0001011     | 100    | 0x1    | `[rd:32]_2 = compress256([rs1:32]_2, [rs2:64]_2)`  |
+| sha512_update | R   | 0001011     | 100    | 0x2    | `[rd:64]_2 = compress512([rs1:64]_2, [rs2:128]_2)` |
 
 ## BigInt Extension
 

docs/vocs/docs/pages/specs/reference/transpiler.mdx

@@ -153,11 +153,10 @@ Each VM extension's behavior is specified below.
 
 ### SHA-2 Extension
 
-| RISC-V Inst | OpenVM Instruction                              |
-| ----------- | ----------------------------------------------- |
-| sha256      | SHA256_RV32 `ind(rd), ind(rs1), ind(rs2), 1, 2` |
-| sha512      | SHA512_RV32 `ind(rd), ind(rs1), ind(rs2), 1, 2` |
-| sha384      | SHA384_RV32 `ind(rd), ind(rs1), ind(rs2), 1, 2` |
+| RISC-V Inst   | OpenVM Instruction                                     |
+| ------------- | ------------------------------------------------------ |
+| sha256_update | SHA256_UPDATE_RV32 `ind(rd), ind(rs1), ind(rs2), 1, 2` |
+| sha512_update | SHA512_UPDATE_RV32 `ind(rd), ind(rs1), ind(rs2), 1, 2` |
 
 ### BigInt Extension
 

extensions/sha2/guest/src/lib.rs

@@ -21,7 +21,7 @@ pub enum Sha2BaseFunct7 {
 ///
 /// The VM accepts the previous hash state and the next block of input, and writes the
 /// new hash state.
-/// - `prev_state` must point to a buffer of at least 32 bytes, storing the previous hash state as 8
+/// - `state` must point to a buffer of at least 32 bytes, storing the previous hash state as 8
 ///   32-bit words in little-endian order
 /// - `input` must point to a buffer of at least 64 bytes
 /// - `output` must point to a buffer of at least 32 bytes. It will be filled with the new hash
@@ -31,21 +31,21 @@ pub enum Sha2BaseFunct7 {
 #[cfg(target_os = "zkvm")]
 #[inline(always)]
 #[no_mangle]
-pub extern "C" fn zkvm_sha256_impl(prev_state: *const u8, input: *const u8, output: *mut u8) {
+pub extern "C" fn zkvm_sha256_impl(state: *const u8, input: *const u8, output: *mut u8) {
     // SAFETY: we handle all cases where `prev_state`, `input`, or `output` are not aligned to 4
     // bytes.
 
     // The minimum alignment required for the buffers
     const MIN_ALIGN: usize = 4;
     unsafe {
-        let prev_state_is_aligned = prev_state as usize % MIN_ALIGN == 0;
+        let state_is_aligned = state as usize % MIN_ALIGN == 0;
         let input_is_aligned = input as usize % MIN_ALIGN == 0;
         let output_is_aligned = output as usize % MIN_ALIGN == 0;
 
-        let prev_state_ptr = if prev_state_is_aligned {
-            prev_state
+        let state_ptr = if state_is_aligned {
+            state
         } else {
-            AlignedBuf::new(prev_state, 32, MIN_ALIGN).ptr
+            AlignedBuf::new(state, 32, MIN_ALIGN).ptr
         };
 
         let input_ptr = if input_is_aligned {
@@ -60,7 +60,7 @@ pub extern "C" fn zkvm_sha256_impl(prev_state: *const u8, input: *const u8, outp
             AlignedBuf::uninit(32, MIN_ALIGN).ptr
         };
 
-        __native_sha256_compress(prev_state_ptr, input_ptr, output_ptr);
+        __native_sha256_compress(state_ptr, input_ptr, output_ptr);
 
         if !output_is_aligned {
             core::ptr::copy_nonoverlapping(output_ptr, output, 32);
@@ -73,7 +73,7 @@ pub extern "C" fn zkvm_sha256_impl(prev_state: *const u8, input: *const u8, outp
 ///
 /// The VM accepts the previous hash state and the next block of input, and writes the
 /// new hash state.
-/// - `prev_state` must point to a buffer of at least 64 bytes, storing the previous hash state as 8
+/// - `state` must point to a buffer of at least 64 bytes, storing the previous hash state as 8
 ///   64-bit words in little-endian order
 /// - `input` must point to a buffer of at least 128 bytes
 /// - `output` must point to a buffer of at least 64 bytes. It will be filled with the new hash
@@ -83,21 +83,21 @@ pub extern "C" fn zkvm_sha256_impl(prev_state: *const u8, input: *const u8, outp
 #[cfg(target_os = "zkvm")]
 #[inline(always)]
 #[no_mangle]
-pub extern "C" fn zkvm_sha512_impl(prev_state: *const u8, input: *const u8, output: *mut u8) {
+pub extern "C" fn zkvm_sha512_impl(state: *const u8, input: *const u8, output: *mut u8) {
     // SAFETY: we handle all cases where `prev_state`, `input`, or `output` are not aligned to 4
     // bytes.
 
     // The minimum alignment required for the buffers
     const MIN_ALIGN: usize = 4;
     unsafe {
-        let prev_state_is_aligned = prev_state as usize % MIN_ALIGN == 0;
+        let state_is_aligned = state as usize % MIN_ALIGN == 0;
         let input_is_aligned = input as usize % MIN_ALIGN == 0;
         let output_is_aligned = output as usize % MIN_ALIGN == 0;
 
-        let prev_state_ptr = if prev_state_is_aligned {
-            prev_state
+        let state_ptr = if state_is_aligned {
+            state
         } else {
-            AlignedBuf::new(prev_state, 64, MIN_ALIGN).ptr
+            AlignedBuf::new(state, 64, MIN_ALIGN).ptr
         };
 
         let input_ptr = if input_is_aligned {
@@ -112,7 +112,7 @@ pub extern "C" fn zkvm_sha512_impl(prev_state: *const u8, input: *const u8, outp
             AlignedBuf::uninit(64, MIN_ALIGN).ptr
         };
 
-        __native_sha512_compress(prev_state_ptr, input_ptr, output_ptr);
+        __native_sha512_compress(state_ptr, input_ptr, output_ptr);
 
         if !output_is_aligned {
             core::ptr::copy_nonoverlapping(output_ptr, output, 64);