add support for the RISC Zero SHA-256 accelerator

sha2/src/sha256.rs

@@ -17,6 +17,9 @@ cfg_if::cfg_if! {
         mod soft;
         mod aarch64;
         use aarch64::compress;
+    } else if #[cfg(all(target_os = "zkvm", target_arch = "riscv32"))] {
+        mod risc0;
+        use risc0::compress;
     } else {
         mod soft;
         use soft::compress;

sha2/src/sha256/risc0.rs

@@ -0,0 +1,61 @@
+extern crate alloc;
+
+const DIGEST_BYTES: usize = 32;
+const DIGEST_WORDS: usize = DIGEST_BYTES / 4;
+const BLOCK_BYTES: usize = DIGEST_BYTES * 2;
+const BLOCK_WORDS: usize = DIGEST_WORDS * 2;
+
+extern "C" {
+    fn sys_sha_buffer(
+        out_state: *mut [u32; DIGEST_WORDS],
+        in_state: *const [u32; DIGEST_WORDS],
+        buf: *const u8,
+        count: u32,
+    );
+}
+
+use alloc::vec::Vec;
+
+#[inline(always)]
+fn compress_words(state: &mut [u32; DIGEST_WORDS], blocks: &[[u32; BLOCK_WORDS]]) {
+    unsafe {
+        sys_sha_buffer(
+            state,
+            state,
+            blocks.as_ptr() as *const u8,
+            blocks.len() as u32,
+        );
+    }
+}
+
+// When the blocks are unaligned they must be copied in order to align them on a u32 word boundary
+// before they can be passed to sys_sha_buffer. This function does this, allocating a new Vec.
+fn read_unaligned_blocks(blocks: &[[u8; BLOCK_BYTES]]) -> Vec<[u32; BLOCK_WORDS]> {
+    blocks
+        .iter()
+        .map(|block| unsafe { (block.as_ptr() as *const [u32; BLOCK_WORDS]).read_unaligned() })
+        .collect::<Vec<_>>()
+}
+
+/// SHA-256 compress implementation which calls into the RISZ Zero SHA-256 accelerator circuit.
+/// Based on https://github.com/risc0/risc0/tree/main/risc0/zkvm/src/guest/sha.rs
+#[inline]
+pub fn compress(state: &mut [u32; DIGEST_WORDS], blocks: &[[u8; BLOCK_BYTES]]) {
+    // On little-endian architectures, flip from big-endian to little-endian.
+    // RISC Zero expects the state to always be big-endian, wheras this crate uses native-endian.
+    for word in state.iter_mut() {
+        *word = word.to_be();
+    }
+
+    // If aligned, reinterpret the u8 array blocks as u32 array blocks.
+    // If unaligned, the data needs to be copied.
+    match unsafe { blocks.align_to::<[u32; BLOCK_WORDS]>() } {
+        (&[], aligned_blocks, &[]) => compress_words(state, aligned_blocks),
+        _ => compress_words(state, &read_unaligned_blocks(&blocks)),
+    };
+
+    // On little-endian architectures, flip from big-endian to little-endian.
+    for word in state.iter_mut() {
+        *word = word.to_be();
+    }
+}