stwo_gpu/core/channel/
blake2s.rs

1use core::{array, iter};
2
3use itertools::Itertools;
4use std_shims::Vec;
5
6use super::Channel;
7use crate::core::fields::m31::{BaseField, P};
8use crate::core::fields::qm31::{SecureField, SECURE_EXTENSION_DEGREE};
9use crate::core::vcs::blake2_hash::{Blake2sHash, Blake2sHasherGeneric};
10
11pub const BLAKE_BYTES_PER_HASH: usize = 32;
12pub const FELTS_PER_HASH: usize = 8;
13
14pub type Blake2sChannel = Blake2sChannelGeneric<false>;
15/// Same as [Blake2sChannel], expect that the hash output is taken modulo M31::P.
16pub type Blake2sM31Channel = Blake2sChannelGeneric<true>;
17
18/// A channel that can be used to draw random elements from a [Blake2sHash] digest.
19#[derive(Default, Clone, Debug)]
20pub struct Blake2sChannelGeneric<const IS_M31_OUTPUT: bool> {
21    digest: Blake2sHash,
22    n_draws: u32,
23}
24
25impl<const IS_M31_OUTPUT: bool> Blake2sChannelGeneric<IS_M31_OUTPUT> {
26    pub const POW_PREFIX: u32 = 0x12345678;
27
28    pub const fn digest(&self) -> Blake2sHash {
29        self.digest
30    }
31    pub const fn update_digest(&mut self, new_digest: Blake2sHash) {
32        self.digest = new_digest;
33        self.n_draws = 0;
34    }
35    /// Generates a uniform random vector of BaseField elements.
36    fn draw_base_felts(&mut self) -> [BaseField; FELTS_PER_HASH] {
37        // Repeats hashing with an increasing counter until getting a good result.
38        // Retry probability for each round is ~ 2^(-28).
39        loop {
40            let u32s: [u32; FELTS_PER_HASH] = self.draw_u32s().try_into().unwrap();
41
42            // Retry if not all the u32 are in the range [0, 2P).
43            if u32s.iter().all(|x| *x < 2 * P) {
44                return u32s
45                    .into_iter()
46                    .map(|x| BaseField::reduce(x as u64))
47                    .collect::<Vec<_>>()
48                    .try_into()
49                    .unwrap();
50            }
51        }
52    }
53}
54
55impl<const IS_M31_OUTPUT: bool> Channel for Blake2sChannelGeneric<IS_M31_OUTPUT> {
56    const BYTES_PER_HASH: usize = BLAKE_BYTES_PER_HASH;
57
58    fn mix_felts(&mut self, felts: &[SecureField]) {
59        let felts_bytes = felts
60            .iter()
61            .flat_map(|qm31| qm31.to_m31_array())
62            .flat_map(|m31| m31.0.to_le_bytes())
63            .collect_vec();
64        let mut hasher = Blake2sHasherGeneric::<IS_M31_OUTPUT>::new();
65        hasher.update(self.digest.as_ref());
66        hasher.update(&felts_bytes);
67
68        self.update_digest(hasher.finalize());
69    }
70
71    fn mix_u32s(&mut self, data: &[u32]) {
72        let mut hasher = Blake2sHasherGeneric::<IS_M31_OUTPUT>::new();
73        hasher.update(self.digest.as_ref());
74        for word in data {
75            hasher.update(&word.to_le_bytes());
76        }
77
78        self.update_digest(hasher.finalize());
79    }
80
81    fn mix_u64(&mut self, value: u64) {
82        self.mix_u32s(&[value as u32, (value >> 32) as u32])
83    }
84
85    fn draw_secure_felt(&mut self) -> SecureField {
86        let felts: [BaseField; FELTS_PER_HASH] = self.draw_base_felts();
87        SecureField::from_m31_array(felts[..SECURE_EXTENSION_DEGREE].try_into().unwrap())
88    }
89
90    fn draw_secure_felts(&mut self, n_felts: usize) -> Vec<SecureField> {
91        let mut felts = iter::from_fn(|| Some(self.draw_base_felts())).flatten();
92        let secure_felts = iter::from_fn(|| {
93            Some(SecureField::from_m31_array([
94                felts.next()?,
95                felts.next()?,
96                felts.next()?,
97                felts.next()?,
98            ]))
99        });
100        secure_felts.take(n_felts).collect()
101    }
102
103    fn draw_u32s(&mut self) -> Vec<u32> {
104        let mut hash_input = self.digest.as_ref().to_vec();
105
106        // Append counter bytes directly (4 bytes for u32).
107        let counter_bytes = self.n_draws.to_le_bytes();
108        hash_input.extend_from_slice(&counter_bytes);
109
110        // Append a zero byte for domain separation between generating randomness and mixing a
111        // single u32.
112        hash_input.push(0_u8);
113
114        self.n_draws += 1;
115        Blake2sHasherGeneric::<IS_M31_OUTPUT>::hash(&hash_input)
116            .0
117            .chunks_exact(4)
118            .map(|chunk| u32::from_le_bytes(chunk.try_into().unwrap()))
119            .collect()
120    }
121
122    /// Verifies that `H(H(POW_PREFIX, [0_u8; 12], digest, n_bits), nonce)` has at least `n_bits`
123    /// many leading zeros.
124    fn verify_pow_nonce(&self, n_bits: u32, nonce: u64) -> bool {
125        let digest = self.digest();
126        // Compute H(POW_PREFIX, [0_u8; 12], digest, n_bits).
127        let mut hasher = Blake2sHasherGeneric::<IS_M31_OUTPUT>::default();
128        hasher.update(&Self::POW_PREFIX.to_le_bytes());
129        hasher.update(&[0_u8; 12]);
130        hasher.update(&digest.0[..]);
131        hasher.update(&n_bits.to_le_bytes());
132        let prefixed_digest = hasher.finalize();
133        // Compute `H(prefixed_digest, nonce)`.
134        let mut hasher = Blake2sHasherGeneric::<IS_M31_OUTPUT>::default();
135        hasher.update(prefixed_digest.as_ref());
136        hasher.update(&nonce.to_le_bytes());
137        let res = hasher.finalize();
138        let n_zeros = u128::from_le_bytes(array::from_fn(|i| res.0[i])).trailing_zeros();
139        n_zeros >= n_bits
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use itertools::Itertools;
146    use std_shims::BTreeSet;
147
148    use crate::core::channel::blake2s::Blake2sChannel;
149    use crate::core::channel::Channel;
150    use crate::core::fields::qm31::SecureField;
151    use crate::m31;
152
153    #[test]
154    fn test_channel_draws() {
155        let mut channel = Blake2sChannel::default();
156
157        assert_eq!(channel.n_draws, 0);
158
159        channel.draw_u32s();
160        assert_eq!(channel.n_draws, 1);
161
162        channel.draw_secure_felts(9);
163        assert_eq!(channel.n_draws, 6);
164    }
165
166    #[test]
167    fn test_draw_u32s() {
168        let mut channel = Blake2sChannel::default();
169
170        let first_random_words = channel.draw_u32s();
171
172        // Assert that next random words are different.
173        assert_ne!(first_random_words, channel.draw_u32s());
174    }
175
176    #[test]
177    pub fn test_draw_secure_felt() {
178        let mut channel = Blake2sChannel::default();
179
180        let first_random_felt = channel.draw_secure_felt();
181
182        // Assert that next random felt is different.
183        assert_ne!(first_random_felt, channel.draw_secure_felt());
184    }
185
186    #[test]
187    pub fn test_draw_secure_felts() {
188        let mut channel = Blake2sChannel::default();
189
190        let mut random_felts = channel.draw_secure_felts(5);
191        random_felts.extend(channel.draw_secure_felts(4));
192
193        // Assert that all the random felts are unique.
194        assert_eq!(
195            random_felts.len(),
196            random_felts.iter().collect::<BTreeSet<_>>().len()
197        );
198    }
199
200    #[test]
201    pub fn test_mix_felts() {
202        let mut channel = Blake2sChannel::default();
203        let initial_digest = channel.digest;
204        let felts = (0..2)
205            .map(|i| SecureField::from(m31!(i + 1923782)))
206            .collect_vec();
207
208        channel.mix_felts(felts.as_slice());
209
210        assert_ne!(initial_digest, channel.digest);
211    }
212
213    #[test]
214    pub fn test_mix_u64() {
215        let mut channel = Blake2sChannel::default();
216        channel.mix_u64(0x1111222233334444);
217        let digest_64 = channel.digest;
218
219        let mut channel = Blake2sChannel::default();
220        channel.mix_u32s(&[0x33334444, 0x11112222]);
221
222        assert_eq!(digest_64, channel.digest);
223        let digest_bytes: [u8; 32] = digest_64.into();
224        assert_eq!(
225            digest_bytes,
226            [
227                0xbc, 0x9e, 0x3f, 0xc1, 0xd2, 0x4e, 0x88, 0x97, 0x95, 0x6d, 0x33, 0x59, 0x32, 0x73,
228                0x97, 0x24, 0x9d, 0x6b, 0xca, 0xcd, 0x22, 0x4d, 0x92, 0x74, 0x4, 0xe7, 0xba, 0x4a,
229                0x77, 0xdc, 0x6e, 0xce
230            ]
231        )
232    }
233
234    #[test]
235    pub fn test_mix_u32s() {
236        let mut channel = Blake2sChannel::default();
237        channel.mix_u32s(&[1, 2, 3, 4, 5, 6, 7, 8, 9]);
238        let digest: [u8; 32] = channel.digest.into();
239        assert_eq!(
240            digest,
241            [
242                0x70, 0x91, 0x76, 0x83, 0x57, 0xbb, 0x1b, 0xb3, 0x34, 0x6f, 0xda, 0xb6, 0xb3, 0x57,
243                0xd7, 0xfa, 0x46, 0xb8, 0xfb, 0xe3, 0x2c, 0x2e, 0x43, 0x24, 0xa0, 0xff, 0xc2, 0x94,
244                0xcb, 0xf9, 0xa1, 0xc7
245            ]
246        );
247    }
248}
stwo_gpu/core/channel/blake2s.rs

stwo_gpu/core/channel/
blake2s.rs
