chore: fix lack of Context

src/lib.rs

@@ -4,31 +4,35 @@
 //! the beacon chain. The hash function changed during the specification process, so defining it
 //! once in this crate made it easy to replace.
 //!
-//! Now this crate serves primarily as a wrapper over the `sha2` crate.
+//! Now this crate serves primarily as a wrapper over the SHA256 crate `sha2`.
 
-use sha2::{Digest, Sha256};
+mod sha2_impl;
+
+pub use self::DynamicContext as Context;
+
+use sha2_impl::Sha2CrateImpl;
 
 #[cfg(feature = "zero_hash_cache")]
 use std::sync::LazyLock;
 
 /// Length of a SHA256 hash in bytes.
 pub const HASH_LEN: usize = 32;
 
-/// Returns the digest of `input` using the `sha2` implementation.
+/// Returns the digest of `input` using the best available implementation.
 pub fn hash(input: &[u8]) -> Vec<u8> {
-    Sha2Impl.hash(input)
+    DynamicImpl::best().hash(input)
 }
 
 /// Hash function returning a fixed-size array (to save on allocations).
 ///
-/// Uses the `sha2` implementation.
+/// Uses the best available implementation based on CPU features.
 pub fn hash_fixed(input: &[u8]) -> [u8; HASH_LEN] {
-    Sha2Impl.hash_fixed(input)
+    DynamicImpl::best().hash_fixed(input)
 }
 
 /// Compute the hash of two slices concatenated.
-pub fn hash32_concat(h1: &[u8], h2: &[u8]) -> [u8; HASH_LEN] {
-    let mut ctxt = Sha2Context::new();
+pub fn hash32_concat(h1: &[u8], h2: &[u8]) -> [u8; 32] {
+    let mut ctxt = DynamicContext::new();
     ctxt.update(h1);
     ctxt.update(h2);
     ctxt.finalize()
@@ -43,58 +47,93 @@ pub trait Sha256Context {
     fn finalize(self) -> [u8; HASH_LEN];
 }
 
-/// Implementation of SHA256 using the `sha2` crate.
-pub struct Sha2Context {
-    hasher: Sha256,
-}
+/// Top-level trait implemented by the `sha2` implementation.
+pub trait Sha256 {
+    type Context: Sha256Context;
 
-impl Sha256Context for Sha2Context {
-    fn new() -> Self {
-        Self {
-            hasher: Sha256::new(),
-        }
-    }
+    fn hash(&self, input: &[u8]) -> Vec<u8>;
 
-    fn update(&mut self, bytes: &[u8]) {
-        self.hasher.update(bytes);
-    }
+    fn hash_fixed(&self, input: &[u8]) -> [u8; HASH_LEN];
+}
 
-    fn finalize(self) -> [u8; HASH_LEN] {
-        let result = self.hasher.finalize();
-        let mut output = [0u8; HASH_LEN];
-        output.copy_from_slice(&result);
-        output
-    }
+/// Default dynamic implementation that switches between available implementations.
+pub enum DynamicImpl {
+    Sha2,
 }
 
-/// Top-level trait implemented by the `sha2` implementation.
-pub trait Sha256Trait {
-    type Context: Sha256Context;
+// Runtime latch for detecting the availability of SHA extensions on x86_64.
+//
+// Inspired by the runtime switch within the `sha2` crate itself.
+#[cfg(target_arch = "x86_64")]
+cpufeatures::new!(x86_sha_extensions, "sha", "sse2", "ssse3", "sse4.1");
 
-    fn hash(&self, input: &[u8]) -> Vec<u8>;
+#[inline(always)]
+pub fn have_sha_extensions() -> bool {
+    #[cfg(target_arch = "x86_64")]
+    return x86_sha_extensions::get();
 
-    fn hash_fixed(&self, input: &[u8]) -> [u8; HASH_LEN];
+    #[cfg(not(target_arch = "x86_64"))]
+    return false;
 }
 
-/// Implementation of the `Sha256Trait` using the `sha2` crate.
-pub struct Sha2Impl;
+impl DynamicImpl {
+    /// Choose the best available implementation based on the currently executing CPU.
+    #[inline(always)]
+    pub fn best() -> Self {
+        #[cfg(target_arch = "x86_64")]
+        if have_sha_extensions() {
+            Self::Sha2
+        } else {
+            Self::Sha2
+        }
+
+        #[cfg(not(target_arch = "x86_64"))]
+        Self::Sha2
+    }
+}
 
-impl Sha256Trait for Sha2Impl {
-    type Context = Sha2Context;
+impl Sha256 for DynamicImpl {
+    type Context = DynamicContext;
 
+    #[inline(always)]
     fn hash(&self, input: &[u8]) -> Vec<u8> {
-        let mut hasher = Sha256::new();
-        hasher.update(input);
-        hasher.finalize().to_vec()
+        match self {
+            Self::Sha2 => Sha2CrateImpl.hash(input),
+        }
     }
 
+    #[inline(always)]
     fn hash_fixed(&self, input: &[u8]) -> [u8; HASH_LEN] {
-        let mut hasher = Sha256::new();
-        hasher.update(input);
-        let result = hasher.finalize();
-        let mut output = [0u8; HASH_LEN];
-        output.copy_from_slice(&result);
-        output
+        match self {
+            Self::Sha2 => Sha2CrateImpl.hash_fixed(input),
+        }
+    }
+}
+
+/// Context encapsulating all implementation contexts.
+///
+/// This enum ends up being 8 bytes larger than the largest inner context.
+pub enum DynamicContext {
+    Sha2(sha2::Sha256),
+}
+
+impl Sha256Context for DynamicContext {
+    fn new() -> Self {
+        match DynamicImpl::best() {
+            DynamicImpl::Sha2 => Self::Sha2(Sha256Context::new()),
+        }
+    }
+
+    fn update(&mut self, bytes: &[u8]) {
+        match self {
+            Self::Sha2(ctxt) => Sha256Context::update(ctxt, bytes),
+        }
+    }
+
+    fn finalize(self) -> [u8; HASH_LEN] {
+        match self {
+            Self::Sha2(ctxt) => Sha256Context::finalize(ctxt),
+        }
     }
 }
 
@@ -119,7 +158,11 @@ mod tests {
     use super::*;
     use rustc_hex::FromHex;
 
-    #[test]
+    #[cfg(target_arch = "wasm32")]
+    use wasm_bindgen_test::*;
+
+    #[cfg_attr(not(target_arch = "wasm32"), test)]
+    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
     fn test_hashing() {
         let input: Vec<u8> = b"hello world".as_ref().into();
 

src/sha2_impl.rs

@@ -1,6 +1,5 @@
 // This implementation should only be compiled on x86_64 due to its dependency on the `sha2` and
 // `cpufeatures` crates which do not compile on some architectures like RISC-V.
-#![cfg(target_arch = "x86_64")]
 
 use crate::{Sha256, Sha256Context, HASH_LEN};
 use sha2::Digest;