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