stwo_gpu/core/vcs/
blake2_hash.rs1use 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#[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>;
68pub type Blake2sM31Hasher = Blake2sHasherGeneric<true>;
70
71#[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
110pub 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}