risc0_zkp/core/hash/sha/
rng.rs1use alloc::boxed::Box;
18
19use rand_core::{impls, RngCore};
20use risc0_core::field::{Elem, Field};
21
22use super::{Digest, Impl, Sha256, DIGEST_WORDS};
23use crate::core::hash::Rng;
24
25#[derive(Clone, Debug)]
27pub struct ShaRng {
28 pool0: Box<Digest>,
30 pool1: Box<Digest>,
32 pool_used: usize,
33}
34
35impl Default for ShaRng {
36 fn default() -> Self {
37 Self::new()
38 }
39}
40
41impl ShaRng {
42 pub fn new() -> Self {
44 Self {
45 pool0: (*Impl::hash_bytes(b"Hello")).into(),
46 pool1: (*Impl::hash_bytes(b"World")).into(),
47 pool_used: 0,
48 }
49 }
50
51 pub fn inner_mix(&mut self, val: &Digest) {
53 for i in 0..DIGEST_WORDS {
54 self.pool0.as_mut_words()[i] = self.pool0.as_words()[i] ^ val.as_words()[i];
55 }
56 self.step();
57 }
58
59 fn step(&mut self) {
60 self.pool0 = (*Impl::hash_pair(&self.pool0, &self.pool1)).into();
61 self.pool1 = (*Impl::hash_pair(&self.pool0, &self.pool1)).into();
62 self.pool_used = 0;
63 }
64}
65
66impl RngCore for ShaRng {
67 fn next_u32(&mut self) -> u32 {
68 if self.pool_used == DIGEST_WORDS {
69 self.step();
70 }
71 let out = self.pool0.as_words()[self.pool_used];
72 self.pool_used += 1;
74 out
75 }
76
77 fn next_u64(&mut self) -> u64 {
78 ((self.next_u32() as u64) << 32) | (self.next_u32() as u64)
79 }
80
81 fn fill_bytes(&mut self, dest: &mut [u8]) {
82 impls::fill_bytes_via_next(self, dest);
83 }
84}
85
86impl<F: Field> Rng<F> for ShaRng {
87 fn mix(&mut self, val: &Digest) {
88 self.inner_mix(val);
89 }
90
91 fn random_bits(&mut self, bits: usize) -> u32 {
92 ((1 << bits) - 1) & self.next_u32()
93 }
94
95 fn random_elem(&mut self) -> F::Elem {
96 F::Elem::random(self)
97 }
98
99 fn random_ext_elem(&mut self) -> F::ExtElem {
100 F::ExtElem::random(self)
101 }
102}
103
104#[cfg(test)]
105mod tests {
106 use rand_core::RngCore;
107
108 use super::ShaRng;
109 use crate::core::hash::sha::{cpu::Impl, Sha256};
110
111 #[test]
114 fn test_sha_rng_impl() {
115 let mut x = ShaRng::new();
116 for _ in 0..10 {
117 x.next_u32();
118 }
119 assert_eq!(x.next_u32(), 785921476);
120 x.inner_mix(&Impl::hash_bytes(b"foo"));
121 assert_eq!(x.next_u32(), 4167871101);
122 }
123}