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