Skip to main content

stwo_gpu/core/vcs/
blake2_hash.rs

1use core::fmt;
2
3use blake2::{Blake2s256, Digest};
4use bytemuck::{Pod, Zeroable};
5use serde::{Deserialize, Serialize};
6use std_shims::Vec;
7
8use crate::core::fields::m31::M31;
9
10// Wrapper for the blake2s hash type.
11#[repr(C, align(32))]
12#[derive(Clone, Copy, PartialEq, Default, Eq, Pod, Zeroable, Deserialize, Serialize)]
13pub struct Blake2sHash(pub [u8; 32]);
14
15impl From<Blake2sHash> for Vec<u8> {
16    fn from(value: Blake2sHash) -> Self {
17        Vec::from(value.0)
18    }
19}
20
21impl From<Vec<u8>> for Blake2sHash {
22    fn from(value: Vec<u8>) -> Self {
23        Self(
24            value
25                .try_into()
26                .expect("Failed converting Vec<u8> to Blake2Hash type"),
27        )
28    }
29}
30
31impl From<&[u8]> for Blake2sHash {
32    fn from(value: &[u8]) -> Self {
33        Self(
34            value
35                .try_into()
36                .expect("Failed converting &[u8] to Blake2sHash Type!"),
37        )
38    }
39}
40
41impl AsRef<[u8]> for Blake2sHash {
42    fn as_ref(&self) -> &[u8] {
43        &self.0
44    }
45}
46
47impl From<Blake2sHash> for [u8; 32] {
48    fn from(val: Blake2sHash) -> Self {
49        val.0
50    }
51}
52
53impl fmt::Display for Blake2sHash {
54    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
55        f.write_str(&hex::encode(self.0))
56    }
57}
58
59impl fmt::Debug for Blake2sHash {
60    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
61        <Blake2sHash as fmt::Display>::fmt(self, f)
62    }
63}
64
65impl super::hash::Hash for Blake2sHash {}
66
67pub type Blake2sHasher = Blake2sHasherGeneric<false>;
68/// Same as [Blake2sHasher], expect that the hash output is taken modulo M31::P.
69pub type Blake2sM31Hasher = Blake2sHasherGeneric<true>;
70
71// Wrapper for the blake2s Hashing functionalities.
72#[derive(Clone, Debug, Default)]
73pub struct Blake2sHasherGeneric<const IS_M31_OUTPUT: bool> {
74    state: Blake2s256,
75}
76
77impl<const IS_M31_OUTPUT: bool> Blake2sHasherGeneric<IS_M31_OUTPUT> {
78    pub fn new() -> Self {
79        Self {
80            state: Blake2s256::new(),
81        }
82    }
83
84    pub fn update(&mut self, data: &[u8]) {
85        blake2::Digest::update(&mut self.state, data);
86    }
87
88    pub fn finalize(self) -> Blake2sHash {
89        let mut r: [u8; 32] = self.state.finalize().into();
90        if IS_M31_OUTPUT {
91            r = reduce_to_m31(r);
92        }
93        Blake2sHash(r)
94    }
95
96    pub fn concat_and_hash(v1: &Blake2sHash, v2: &Blake2sHash) -> Blake2sHash {
97        let mut hasher = Self::new();
98        hasher.update(v1.as_ref());
99        hasher.update(v2.as_ref());
100        hasher.finalize()
101    }
102
103    pub fn hash(data: &[u8]) -> Blake2sHash {
104        let mut hasher = Self::new();
105        hasher.update(data);
106        hasher.finalize()
107    }
108}
109
110/// Reduces each u32 in the input (interpreted as little-endian) modulo M31::P.
111pub fn reduce_to_m31(value: [u8; 32]) -> [u8; 32] {
112    let mut res = [0u8; 32];
113    for (i, c) in value.chunks(4).enumerate() {
114        let val = M31::reduce(u32::from_le_bytes(c.try_into().unwrap()).into());
115        res[i * 4..(i + 1) * 4].copy_from_slice(&val.0.to_le_bytes());
116    }
117    res
118}
119
120#[cfg(test)]
121mod tests {
122    use blake2::Digest;
123    use std_shims::ToString;
124
125    use super::{Blake2sHash, Blake2sHasher};
126
127    impl Blake2sHasher {
128        fn finalize_reset(&mut self) -> Blake2sHash {
129            Blake2sHash(self.state.finalize_reset().into())
130        }
131    }
132
133    #[test]
134    fn single_hash_test() {
135        let hash_a = Blake2sHasher::hash(b"a");
136        assert_eq!(
137            hash_a.to_string(),
138            "4a0d129873403037c2cd9b9048203687f6233fb6738956e0349bd4320fec3e90"
139        );
140    }
141
142    #[test]
143    fn hash_state_test() {
144        let mut state = Blake2sHasher::new();
145        state.update(b"a");
146        state.update(b"b");
147        let hash = state.finalize_reset();
148        let hash_empty = state.finalize();
149
150        assert_eq!(hash.to_string(), Blake2sHasher::hash(b"ab").to_string());
151        assert_eq!(hash_empty.to_string(), Blake2sHasher::hash(b"").to_string());
152    }
153}