From bedaffe436e5c17ae9f659bd342ce2b882e426ef Mon Sep 17 00:00:00 2001
From: Jonathan Wang <31040440+jonathanpwang@users.noreply.github.com>
Date: Tue, 31 Dec 2024 03:22:43 -0500
Subject: [PATCH] chore: update book

---
 book/src/custom-extensions/keccak.md | 10 ++++++----
 book/src/custom-extensions/sha256.md | 14 ++++++++------
 docs/specs/RISCV.md                  | 13 +++++++------
 extensions/sha256/guest/src/lib.rs   |  9 ++++-----
 4 files changed, 25 insertions(+), 21 deletions(-)

diff --git a/book/src/custom-extensions/keccak.md b/book/src/custom-extensions/keccak.md
index 40c76b9b6f..5d3b733eb7 100644
--- a/book/src/custom-extensions/keccak.md
+++ b/book/src/custom-extensions/keccak.md
@@ -1,7 +1,7 @@
 # OpenVM Keccak256
 
-The OpenVm Keccak256 extension provides tools for using the Keccak-256 hash function. 
-The functional part is provided by the `openvm-keccak-guest` crate, which is a guest library that can be used in any OpenVM program. 
+The OpenVm Keccak256 extension provides tools for using the Keccak-256 hash function.
+The functional part is provided by the `openvm-keccak-guest` crate, which is a guest library that can be used in any OpenVM program.
 
 ## Functions for guest code
 
@@ -13,6 +13,7 @@ The OpenVM Keccak256 Guest extension provides two functions for using in your gu
 See the full example [here](https://github.com/openvm-org/openvm/blob/main/crates/toolchain/tests/programs/examples/keccak.rs).
 
 ### Example:
+
 ```rust
 use hex::FromHex;
 use openvm_keccak256_guest::keccak256;
@@ -54,8 +55,9 @@ extern "C" {
 }
 
 fn keccak256(input: &[u8]) -> [u8; 32] {
-    #[cfg(target_os = "zkvm")] {
-    let mut output = [0u8; 32];
+    #[cfg(target_os = "zkvm")]
+    {
+        let mut output = [0u8; 32];
         unsafe {
             native_keccak256(input.as_ptr(), input.len(), output.as_mut_ptr() as *mut u8);
         }
diff --git a/book/src/custom-extensions/sha256.md b/book/src/custom-extensions/sha256.md
index aa6cf3859d..d5478f5fb8 100644
--- a/book/src/custom-extensions/sha256.md
+++ b/book/src/custom-extensions/sha256.md
@@ -1,7 +1,7 @@
 # OpenVM Sha256
 
 The OpenVm Sha256 extension provides tools for using the Sha256 hash function. Refer [here][https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.180-4.pdf] for more details on the Sha256.
-The functional part is provided by the `openvm-sha256-guest` crate, which is a guest library that can be used in any OpenVM program. 
+The functional part is provided by the `openvm-sha256-guest` crate, which is a guest library that can be used in any OpenVM program.
 
 ## Functions for guest code
 
@@ -13,6 +13,7 @@ The OpenVM Sha256 Guest extension provides two functions for using in your guest
 See the full example [here](https://github.com/openvm-org/openvm/blob/main/examples/sha256).
 
 ### Example:
+
 ```rust
 extern crate alloc;
 
@@ -51,20 +52,21 @@ hex = { version = "0.4.3", default-features = false, features = ["alloc"] }
 
 Keccak guest extension also provides another way to use the native Sha256 implementation. It provides a function that is meant to be linked to other external libraries. The external libraries can use this function as a hook for the Sha256 native implementation. Enabled only when the target is `zkvm`.
 
-- `native_sha256(input: *const u8, len: usize, output: *mut u8)`: This function has `C` ABI. It takes in a pointer to the input, the length of the input, and a pointer to the output buffer.
+- `zkvm_sha256_impl(input: *const u8, len: usize, output: *mut u8)`: This function has `C` ABI. It takes in a pointer to the input, the length of the input, and a pointer to the output buffer.
 
 In the external library, you can do the following:
 
 ```rust
 extern "C" {
-    fn native_sha256(input: *const u8, len: usize, output: *mut u8);
+    fn zkvm_sha256_impl(input: *const u8, len: usize, output: *mut u8);
 }
 
 fn sha256(input: &[u8]) -> [u8; 32] {
-    #[cfg(target_os = "zkvm")] {
-    let mut output = [0u8; 32];
+    #[cfg(target_os = "zkvm")]
+    {
+        let mut output = [0u8; 32];
         unsafe {
-            native_sha256(input.as_ptr(), input.len(), output.as_mut_ptr() as *mut u8);
+            zkvm_sha256_impl(input.as_ptr(), input.len(), output.as_mut_ptr() as *mut u8);
         }
         output
     }
diff --git a/docs/specs/RISCV.md b/docs/specs/RISCV.md
index 78dc894178..d0273436d2 100644
--- a/docs/specs/RISCV.md
+++ b/docs/specs/RISCV.md
@@ -28,7 +28,8 @@ We start with the instructions using _custom-0_ opcode[6:0] prefix **0001011**..
 | RISC-V Inst | FMT | opcode[6:0] | funct3 | funct7 | RISC-V description and notes                |
 | ----------- | --- | ----------- | ------ | ------ | ------------------------------------------- |
 | keccak256   | R   | 0001011     | 100    | 0x0    | `[rd:32]_2 = keccak256([rs1..rs1 + rs2]_2)` |
-| sha256      | R   | 0001011     | 111    | 0x0    | `[rd:32]_2 = sha256([rs1..rs1 + rs2]_2)` |
+| sha256      | R   | 0001011     | 100    | 0x1    | `[rd:32]_2 = sha256([rs1..rs1 + rs2]_2)`    |
+
 ## 256-bit Integers
 
 | RISC-V Inst | FMT | opcode[6:0] | funct3 | funct7 | RISC-V description and notes                              |
@@ -87,12 +88,12 @@ Since `funct7` is 7-bits, up to 16 moduli can be supported simultaneously. We us
 
 Short Weierstrass elliptic curve arithmetic depends on elliptic curve `C`. The instruction set and VM can be simultaneously configured _ahead of time_ to support a fixed ordered list of supported curves. We use `config.curve_idx(C)` to denote the index of `C` in this list. In the list below, `idx` denotes `config.curve_idx(C)`.
 
-| RISC-V Inst     | FMT | opcode[6:0] | funct3 | funct7    | RISC-V description and notes                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                            |
-| --------------- | --- | ----------- | ------ | --------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| sw_add_ne\<C\>  | R   | 0101011     | 001    | `idx*8`   | `EcPoint([rd:2*C::COORD_SIZE]_2) = EcPoint([rs1:2*C::COORD_SIZE]_2) + EcPoint([rs2:2*C::COORD_SIZE]_2)`. Assumes that input affine points are not identity and do not have same x-coordinate.                                                                                                                                                                                                                                                                                                                                                                                           |
-| sw_double\<C\>  | R   | 0101011     | 001    | `idx*8+1` | `EcPoint([rd:2*C::COORD_SIZE]_2) = 2 * EcPoint([rs1:2*C::COORD_SIZE]_2)`. Assumes that input affine point is not identity. `rs2` is unused and must be set to `x0`.                                                                                                                                                                                                                                                                                                                                                                                                                     |
+| RISC-V Inst     | FMT | opcode[6:0] | funct3 | funct7    | RISC-V description and notes                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                        |
+| --------------- | --- | ----------- | ------ | --------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| sw_add_ne\<C\>  | R   | 0101011     | 001    | `idx*8`   | `EcPoint([rd:2*C::COORD_SIZE]_2) = EcPoint([rs1:2*C::COORD_SIZE]_2) + EcPoint([rs2:2*C::COORD_SIZE]_2)`. Assumes that input affine points are not identity and do not have same x-coordinate.                                                                                                                                                                                                                                                                                                                                                                       |
+| sw_double\<C\>  | R   | 0101011     | 001    | `idx*8+1` | `EcPoint([rd:2*C::COORD_SIZE]_2) = 2 * EcPoint([rs1:2*C::COORD_SIZE]_2)`. Assumes that input affine point is not identity. `rs2` is unused and must be set to `x0`.                                                                                                                                                                                                                                                                                                                                                                                                 |
 | setup\<C\>      | R   | 0101011     | 001    | `idx*8+2` | `assert([rs1: C::COORD_SIZE]_2 == C::MODULUS)` in the chip defined by the register index of `rs2`. For the sake of implementation convenience it also writes something (can be anything) into `[rd: 2*C::COORD_SIZE]_2`. If `ind(rs2) != 0`, then this instruction is setup for `sw_add_ne`. Otherwise it is setup for `sw_double`. When `ind(rs2) != 0` (add_ne), it is required for proper functionality that `[rs2: C::COORD_SIZE]_2 != [rs1: C::COORD_SIZE]_2`; otherwise (double), it is required that `[rs1 + C::COORD_SIZE: C::COORD_SIZE]_2 != C::Fp::ZERO` |
-| hint_decompress | R   | 0101011     | 001    | `idx*8+3` | Read `x: C::Fp` from `[rs1: C::COORD_SIZE]_2` and `rec_id: u8` from `[rs2]_2`. Reset the hint stream to equal the unique `y: C::Fp` such that `(x, y)` is a point on `C` and `y` has the same parity as `rec_id`, if it exists. Otherwise reset hint stream to arbitrary `C::Fp`. `rd` should be `x0`.                                                                                                                                                                                                                                                                                  |
+| hint_decompress | R   | 0101011     | 001    | `idx*8+3` | Read `x: C::Fp` from `[rs1: C::COORD_SIZE]_2` and `rec_id: u8` from `[rs2]_2`. Reset the hint stream to equal the unique `y: C::Fp` such that `(x, y)` is a point on `C` and `y` has the same parity as `rec_id`, if it exists. Otherwise reset hint stream to arbitrary `C::Fp`. `rd` should be `x0`.                                                                                                                                                                                                                                                              |
 
 Since `funct7` is 7-bits, up to 16 curves can be supported simultaneously. We use `idx*8` to leave some room for future expansion.
 
diff --git a/extensions/sha256/guest/src/lib.rs b/extensions/sha256/guest/src/lib.rs
index 3034a8ff0c..a43ee3dccc 100644
--- a/extensions/sha256/guest/src/lib.rs
+++ b/extensions/sha256/guest/src/lib.rs
@@ -1,13 +1,12 @@
 #![cfg_attr(not(feature = "std"), no_std)]
 
-// TODO[arayi]: Revisit this
+#[cfg(target_os = "zkvm")]
+use core::mem::MaybeUninit;
+
 /// This is custom-0 defined in RISC-V spec document
 pub const OPCODE: u8 = 0x0b;
 pub const SHA256_FUNCT3: u8 = 0b100;
-pub const SHA256_FUNCT7: u8 = 0b1;
-
-#[cfg(target_os = "zkvm")]
-use core::mem::MaybeUninit;
+pub const SHA256_FUNCT7: u8 = 0x1;
 
 /// The sha256 cryptographic hash function.
 #[inline(always)]