1use crate::{
  2    ByteReader, ByteWriter, Deserializable, DeserializationError, Digest, ElementHasher, Felt,
  3    FieldElement, HashFn, Serializable, StarkField,
  4};
  5use core::{
  6    mem::{size_of, transmute, transmute_copy},
  7    ops::Deref,
  8    slice::from_raw_parts,
  9};
 10
 11#[cfg(test)]
 12mod tests;
 13
 14// BLAKE3 N-BIT OUTPUT
 15// ================================================================================================
 16
 17/// N-bytes output of a blake3 function.
 18///
 19/// Note: `N` can't be greater than `32` because [`Digest::as_bytes`] currently supports only 32
 20/// bytes.
 21#[derive(Debug, Copy, Clone, Eq, PartialEq)]
 22pub struct Blake3Digest<const N: usize>([u8; N]);
 23
 24impl<const N: usize> Default for Blake3Digest<N> {
 25    fn default() -> Self {
 26        Self([0; N])
 27    }
 28}
 29
 30impl<const N: usize> Deref for Blake3Digest<N> {
 31    type Target = [u8];
 32
 33    fn deref(&self) -> &Self::Target {
 34        &self.0
 35    }
 36}
 37
 38impl<const N: usize> Serializable for Blake3Digest<N> {
 39    fn write_into<W: ByteWriter>(&self, target: &mut W) {
 40        target.write_u8_slice(&self.0);
 41    }
 42}
 43
 44impl<const N: usize> Deserializable for Blake3Digest<N> {
 45    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
 46        source.read_u8_array().map(Self)
 47    }
 48}
 49
 50impl<const N: usize> Digest for Blake3Digest<N> {
 51    fn as_bytes(&self) -> [u8; 32] {
 52        // compile-time assertion
 53        assert!(N <= 32, "digest currently supports only 32 bytes!");
 54        expand_bytes(&self.0)
 55    }
 56}
 57
 58// BLAKE3 256-BIT OUTPUT
 59// ================================================================================================
 60
 61/// 256-bit output blake3 hasher.
 62pub struct Blake3_256;
 63
 64impl HashFn for Blake3_256 {
 65    type Digest = Blake3Digest<32>;
 66
 67    fn hash(bytes: &[u8]) -> Self::Digest {
 68        Blake3Digest(blake3::hash(bytes).into())
 69    }
 70
 71    fn merge(values: &[Self::Digest; 2]) -> Self::Digest {
 72        Self::hash(prepare_merge(values))
 73    }
 74
 75    fn merge_with_int(seed: Self::Digest, value: u64) -> Self::Digest {
 76        let mut hasher = blake3::Hasher::new();
 77        hasher.update(&seed.0);
 78        hasher.update(&value.to_le_bytes());
 79        Blake3Digest(hasher.finalize().into())
 80    }
 81}
 82
 83impl ElementHasher for Blake3_256 {
 84    type BaseField = Felt;
 85
 86    fn hash_elements<E>(elements: &[E]) -> Self::Digest
 87    where
 88        E: FieldElement<BaseField = Self::BaseField>,
 89    {
 90        Blake3Digest(hash_elements(elements))
 91    }
 92}
 93
 94// BLAKE3 192-BIT OUTPUT
 95// ================================================================================================
 96
 97/// 192-bit output blake3 hasher.
 98pub struct Blake3_192;
 99
100impl HashFn for Blake3_192 {
101    type Digest = Blake3Digest<24>;
102
103    fn hash(bytes: &[u8]) -> Self::Digest {
104        Blake3Digest(*shrink_bytes(&blake3::hash(bytes).into()))
105    }
106
107    fn merge(values: &[Self::Digest; 2]) -> Self::Digest {
108        Self::hash(prepare_merge(values))
109    }
110
111    fn merge_with_int(seed: Self::Digest, value: u64) -> Self::Digest {
112        let mut hasher = blake3::Hasher::new();
113        hasher.update(&seed.0);
114        hasher.update(&value.to_le_bytes());
115        Blake3Digest(*shrink_bytes(&hasher.finalize().into()))
116    }
117}
118
119impl ElementHasher for Blake3_192 {
120    type BaseField = Felt;
121
122    fn hash_elements<E>(elements: &[E]) -> Self::Digest
123    where
124        E: FieldElement<BaseField = Self::BaseField>,
125    {
126        Blake3Digest(hash_elements(elements))
127    }
128}
129
130// BLAKE3 160-BIT OUTPUT
131// ================================================================================================
132
133/// 160-bit output blake3 hasher.
134pub struct Blake3_160;
135
136impl HashFn for Blake3_160 {
137    type Digest = Blake3Digest<20>;
138
139    fn hash(bytes: &[u8]) -> Self::Digest {
140        Blake3Digest(*shrink_bytes(&blake3::hash(bytes).into()))
141    }
142
143    fn merge(values: &[Self::Digest; 2]) -> Self::Digest {
144        Self::hash(prepare_merge(values))
145    }
146
147    fn merge_with_int(seed: Self::Digest, value: u64) -> Self::Digest {
148        let mut hasher = blake3::Hasher::new();
149        hasher.update(&seed.0);
150        hasher.update(&value.to_le_bytes());
151        Blake3Digest(*shrink_bytes(&hasher.finalize().into()))
152    }
153}
154
155impl ElementHasher for Blake3_160 {
156    type BaseField = Felt;
157
158    fn hash_elements<E>(elements: &[E]) -> Self::Digest
159    where
160        E: FieldElement<BaseField = Self::BaseField>,
161    {
162        Blake3Digest(hash_elements(elements))
163    }
164}
165
166// HELPER FUNCTIONS
167// ================================================================================================
168
169/// Zero-copy ref shrink to array.
170fn shrink_bytes<const M: usize, const N: usize>(bytes: &[u8; M]) -> &[u8; N] {
171    // compile-time assertion
172    assert!(
173        M >= N,
174        "N should fit in M so it can be safely transmuted into a smaller slice!"
175    );
176    // safety: bytes len is asserted
177    unsafe { transmute(bytes) }
178}
179
180/// Hash the elements into bytes and shrink the output.
181fn hash_elements<const N: usize, E>(elements: &[E]) -> [u8; N]
182where
183    E: FieldElement<BaseField = Felt>,
184{
185    // don't leak assumptions from felt and check its actual implementation.
186    // this is a compile-time branch so it is for free
187    let digest = if Felt::IS_CANONICAL {
188        blake3::hash(E::elements_as_bytes(elements))
189    } else {
190        E::as_base_elements(elements)
191            .iter()
192            .fold(blake3::Hasher::new(), |mut hasher, felt| {
193                hasher.update(&felt.as_int().to_le_bytes());
194                hasher
195            })
196            .finalize()
197    };
198    *shrink_bytes(&digest.into())
199}
200
201/// Owned bytes expansion.
202fn expand_bytes<const M: usize, const N: usize>(bytes: &[u8; M]) -> [u8; N] {
203    // compile-time assertion
204    assert!(M <= N, "M should fit in N so M can be expanded!");
205    // this branch is constant so it will be optimized to be either one of the variants in release
206    // mode
207    if M == N {
208        // safety: the sizes are checked to be the same
209        unsafe { transmute_copy(bytes) }
210    } else {
211        let mut expanded = [0u8; N];
212        expanded[..M].copy_from_slice(bytes);
213        expanded
214    }
215}
216
217// Cast the slice into contiguous bytes.
218fn prepare_merge<const N: usize, D>(args: &[D; N]) -> &[u8]
219where
220    D: Deref<Target = [u8]>,
221{
222    // compile-time assertion
223    assert!(N > 0, "N shouldn't represent an empty slice!");
224    let values = args.as_ptr() as *const u8;
225    let len = size_of::<D>() * N;
226    // safety: the values are tested to be contiguous
227    let bytes = unsafe { from_raw_parts(values, len) };
228    debug_assert_eq!(args[0].deref(), &bytes[..len / N]);
229    bytes
230}