1use super::{Digest, ElementHasher, Felt, FieldElement, Hasher, StarkField};
  2use crate::utils::{
  3    bytes_to_hex_string, hex_to_bytes, string::String, ByteReader, ByteWriter, Deserializable,
  4    DeserializationError, HexParseError, Serializable,
  5};
  6use core::{
  7    mem::{size_of, transmute, transmute_copy},
  8    ops::Deref,
  9    slice::from_raw_parts,
 10};
 11
 12#[cfg(test)]
 13mod tests;
 14
 15// CONSTANTS
 16// ================================================================================================
 17
 18const DIGEST32_BYTES: usize = 32;
 19const DIGEST24_BYTES: usize = 24;
 20const DIGEST20_BYTES: usize = 20;
 21
 22// BLAKE3 N-BIT OUTPUT
 23// ================================================================================================
 24
 25/// N-bytes output of a blake3 function.
 26///
 27/// Note: `N` can't be greater than `32` because [`Digest::as_bytes`] currently supports only 32
 28/// bytes.
 29#[derive(Debug, Copy, Clone, Eq, PartialEq)]
 30#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
 31#[cfg_attr(feature = "serde", serde(into = "String", try_from = "&str"))]
 32pub struct Blake3Digest<const N: usize>([u8; N]);
 33
 34impl<const N: usize> Default for Blake3Digest<N> {
 35    fn default() -> Self {
 36        Self([0; N])
 37    }
 38}
 39
 40impl<const N: usize> Deref for Blake3Digest<N> {
 41    type Target = [u8];
 42
 43    fn deref(&self) -> &Self::Target {
 44        &self.0
 45    }
 46}
 47
 48impl<const N: usize> From<Blake3Digest<N>> for [u8; N] {
 49    fn from(value: Blake3Digest<N>) -> Self {
 50        value.0
 51    }
 52}
 53
 54impl<const N: usize> From<[u8; N]> for Blake3Digest<N> {
 55    fn from(value: [u8; N]) -> Self {
 56        Self(value)
 57    }
 58}
 59
 60impl<const N: usize> From<Blake3Digest<N>> for String {
 61    fn from(value: Blake3Digest<N>) -> Self {
 62        bytes_to_hex_string(value.as_bytes())
 63    }
 64}
 65
 66impl<const N: usize> TryFrom<&str> for Blake3Digest<N> {
 67    type Error = HexParseError;
 68
 69    fn try_from(value: &str) -> Result<Self, Self::Error> {
 70        hex_to_bytes(value).map(|v| v.into())
 71    }
 72}
 73
 74impl<const N: usize> Serializable for Blake3Digest<N> {
 75    fn write_into<W: ByteWriter>(&self, target: &mut W) {
 76        target.write_bytes(&self.0);
 77    }
 78}
 79
 80impl<const N: usize> Deserializable for Blake3Digest<N> {
 81    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
 82        source.read_array().map(Self)
 83    }
 84}
 85
 86impl<const N: usize> Digest for Blake3Digest<N> {
 87    fn as_bytes(&self) -> [u8; 32] {
 88        // compile-time assertion
 89        assert!(N <= 32, "digest currently supports only 32 bytes!");
 90        expand_bytes(&self.0)
 91    }
 92}
 93
 94// BLAKE3 256-BIT OUTPUT
 95// ================================================================================================
 96
 97/// 256-bit output blake3 hasher.
 98#[derive(Debug, Copy, Clone, Eq, PartialEq)]
 99pub struct Blake3_256;
100
101impl Hasher for Blake3_256 {
102    /// Blake3 collision resistance is 128-bits for 32-bytes output.
103    const COLLISION_RESISTANCE: u32 = 128;
104
105    type Digest = Blake3Digest<32>;
106
107    fn hash(bytes: &[u8]) -> Self::Digest {
108        Blake3Digest(blake3::hash(bytes).into())
109    }
110
111    fn merge(values: &[Self::Digest; 2]) -> Self::Digest {
112        Self::hash(prepare_merge(values))
113    }
114
115    fn merge_with_int(seed: Self::Digest, value: u64) -> Self::Digest {
116        let mut hasher = blake3::Hasher::new();
117        hasher.update(&seed.0);
118        hasher.update(&value.to_le_bytes());
119        Blake3Digest(hasher.finalize().into())
120    }
121}
122
123impl ElementHasher for Blake3_256 {
124    type BaseField = Felt;
125
126    fn hash_elements<E>(elements: &[E]) -> Self::Digest
127    where
128        E: FieldElement<BaseField = Self::BaseField>,
129    {
130        Blake3Digest(hash_elements(elements))
131    }
132}
133
134impl Blake3_256 {
135    /// Returns a hash of the provided sequence of bytes.
136    #[inline(always)]
137    pub fn hash(bytes: &[u8]) -> Blake3Digest<DIGEST32_BYTES> {
138        <Self as Hasher>::hash(bytes)
139    }
140
141    /// Returns a hash of two digests. This method is intended for use in construction of
142    /// Merkle trees and verification of Merkle paths.
143    #[inline(always)]
144    pub fn merge(values: &[Blake3Digest<DIGEST32_BYTES>; 2]) -> Blake3Digest<DIGEST32_BYTES> {
145        <Self as Hasher>::merge(values)
146    }
147
148    /// Returns a hash of the provided field elements.
149    #[inline(always)]
150    pub fn hash_elements<E>(elements: &[E]) -> Blake3Digest<DIGEST32_BYTES>
151    where
152        E: FieldElement<BaseField = Felt>,
153    {
154        <Self as ElementHasher>::hash_elements(elements)
155    }
156}
157
158// BLAKE3 192-BIT OUTPUT
159// ================================================================================================
160
161/// 192-bit output blake3 hasher.
162#[derive(Debug, Copy, Clone, Eq, PartialEq)]
163pub struct Blake3_192;
164
165impl Hasher for Blake3_192 {
166    /// Blake3 collision resistance is 96-bits for 24-bytes output.
167    const COLLISION_RESISTANCE: u32 = 96;
168
169    type Digest = Blake3Digest<24>;
170
171    fn hash(bytes: &[u8]) -> Self::Digest {
172        Blake3Digest(*shrink_bytes(&blake3::hash(bytes).into()))
173    }
174
175    fn merge(values: &[Self::Digest; 2]) -> Self::Digest {
176        Self::hash(prepare_merge(values))
177    }
178
179    fn merge_with_int(seed: Self::Digest, value: u64) -> Self::Digest {
180        let mut hasher = blake3::Hasher::new();
181        hasher.update(&seed.0);
182        hasher.update(&value.to_le_bytes());
183        Blake3Digest(*shrink_bytes(&hasher.finalize().into()))
184    }
185}
186
187impl ElementHasher for Blake3_192 {
188    type BaseField = Felt;
189
190    fn hash_elements<E>(elements: &[E]) -> Self::Digest
191    where
192        E: FieldElement<BaseField = Self::BaseField>,
193    {
194        Blake3Digest(hash_elements(elements))
195    }
196}
197
198impl Blake3_192 {
199    /// Returns a hash of the provided sequence of bytes.
200    #[inline(always)]
201    pub fn hash(bytes: &[u8]) -> Blake3Digest<DIGEST24_BYTES> {
202        <Self as Hasher>::hash(bytes)
203    }
204
205    /// Returns a hash of two digests. This method is intended for use in construction of
206    /// Merkle trees and verification of Merkle paths.
207    #[inline(always)]
208    pub fn merge(values: &[Blake3Digest<DIGEST24_BYTES>; 2]) -> Blake3Digest<DIGEST24_BYTES> {
209        <Self as Hasher>::merge(values)
210    }
211
212    /// Returns a hash of the provided field elements.
213    #[inline(always)]
214    pub fn hash_elements<E>(elements: &[E]) -> Blake3Digest<DIGEST24_BYTES>
215    where
216        E: FieldElement<BaseField = Felt>,
217    {
218        <Self as ElementHasher>::hash_elements(elements)
219    }
220}
221
222// BLAKE3 160-BIT OUTPUT
223// ================================================================================================
224
225/// 160-bit output blake3 hasher.
226#[derive(Debug, Copy, Clone, Eq, PartialEq)]
227pub struct Blake3_160;
228
229impl Hasher for Blake3_160 {
230    /// Blake3 collision resistance is 80-bits for 20-bytes output.
231    const COLLISION_RESISTANCE: u32 = 80;
232
233    type Digest = Blake3Digest<20>;
234
235    fn hash(bytes: &[u8]) -> Self::Digest {
236        Blake3Digest(*shrink_bytes(&blake3::hash(bytes).into()))
237    }
238
239    fn merge(values: &[Self::Digest; 2]) -> Self::Digest {
240        Self::hash(prepare_merge(values))
241    }
242
243    fn merge_with_int(seed: Self::Digest, value: u64) -> Self::Digest {
244        let mut hasher = blake3::Hasher::new();
245        hasher.update(&seed.0);
246        hasher.update(&value.to_le_bytes());
247        Blake3Digest(*shrink_bytes(&hasher.finalize().into()))
248    }
249}
250
251impl ElementHasher for Blake3_160 {
252    type BaseField = Felt;
253
254    fn hash_elements<E>(elements: &[E]) -> Self::Digest
255    where
256        E: FieldElement<BaseField = Self::BaseField>,
257    {
258        Blake3Digest(hash_elements(elements))
259    }
260}
261
262impl Blake3_160 {
263    /// Returns a hash of the provided sequence of bytes.
264    #[inline(always)]
265    pub fn hash(bytes: &[u8]) -> Blake3Digest<DIGEST20_BYTES> {
266        <Self as Hasher>::hash(bytes)
267    }
268
269    /// Returns a hash of two digests. This method is intended for use in construction of
270    /// Merkle trees and verification of Merkle paths.
271    #[inline(always)]
272    pub fn merge(values: &[Blake3Digest<DIGEST20_BYTES>; 2]) -> Blake3Digest<DIGEST20_BYTES> {
273        <Self as Hasher>::merge(values)
274    }
275
276    /// Returns a hash of the provided field elements.
277    #[inline(always)]
278    pub fn hash_elements<E>(elements: &[E]) -> Blake3Digest<DIGEST20_BYTES>
279    where
280        E: FieldElement<BaseField = Felt>,
281    {
282        <Self as ElementHasher>::hash_elements(elements)
283    }
284}
285
286// HELPER FUNCTIONS
287// ================================================================================================
288
289/// Zero-copy ref shrink to array.
290fn shrink_bytes<const M: usize, const N: usize>(bytes: &[u8; M]) -> &[u8; N] {
291    // compile-time assertion
292    assert!(M >= N, "N should fit in M so it can be safely transmuted into a smaller slice!");
293    // safety: bytes len is asserted
294    unsafe { transmute(bytes) }
295}
296
297/// Hash the elements into bytes and shrink the output.
298fn hash_elements<const N: usize, E>(elements: &[E]) -> [u8; N]
299where
300    E: FieldElement<BaseField = Felt>,
301{
302    // don't leak assumptions from felt and check its actual implementation.
303    // this is a compile-time branch so it is for free
304    let digest = if Felt::IS_CANONICAL {
305        blake3::hash(E::elements_as_bytes(elements))
306    } else {
307        let mut hasher = blake3::Hasher::new();
308
309        // BLAKE3 state is 64 bytes - so, we can absorb 64 bytes into the state in a single
310        // permutation. we move the elements into the hasher via the buffer to give the CPU
311        // a chance to process multiple element-to-byte conversions in parallel
312        let mut buf = [0_u8; 64];
313        let mut chunk_iter = E::slice_as_base_elements(elements).chunks_exact(8);
314        for chunk in chunk_iter.by_ref() {
315            for i in 0..8 {
316                buf[i * 8..(i + 1) * 8].copy_from_slice(&chunk[i].as_int().to_le_bytes());
317            }
318            hasher.update(&buf);
319        }
320
321        for element in chunk_iter.remainder() {
322            hasher.update(&element.as_int().to_le_bytes());
323        }
324
325        hasher.finalize()
326    };
327    *shrink_bytes(&digest.into())
328}
329
330/// Owned bytes expansion.
331fn expand_bytes<const M: usize, const N: usize>(bytes: &[u8; M]) -> [u8; N] {
332    // compile-time assertion
333    assert!(M <= N, "M should fit in N so M can be expanded!");
334    // this branch is constant so it will be optimized to be either one of the variants in release
335    // mode
336    if M == N {
337        // safety: the sizes are checked to be the same
338        unsafe { transmute_copy(bytes) }
339    } else {
340        let mut expanded = [0u8; N];
341        expanded[..M].copy_from_slice(bytes);
342        expanded
343    }
344}
345
346// Cast the slice into contiguous bytes.
347fn prepare_merge<const N: usize, D>(args: &[D; N]) -> &[u8]
348where
349    D: Deref<Target = [u8]>,
350{
351    // compile-time assertion
352    assert!(N > 0, "N shouldn't represent an empty slice!");
353    let values = args.as_ptr() as *const u8;
354    let len = size_of::<D>() * N;
355    // safety: the values are tested to be contiguous
356    let bytes = unsafe { from_raw_parts(values, len) };
357    debug_assert_eq!(args[0].deref(), &bytes[..len / N]);
358    bytes
359}