risc0_zkp/core/hash/poseidon2/
rng.rs

1// Copyright 2024 RISC Zero, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! A Poseidon2 based CRNG used in Fiat-Shamir.
16
17use risc0_core::field::{
18    baby_bear::{BabyBear, BabyBearExtElem, Elem},
19    ExtElem,
20};
21
22use super::{consts::CELLS, poseidon2_mix, CELLS_OUT, CELLS_RATE};
23use crate::core::{digest::Digest, hash::Rng};
24
25/// A random number generator driven by Poseidon2
26#[derive(Clone, Debug)]
27pub struct Poseidon2Rng {
28    // The cells of the sponge
29    cells: [Elem; CELLS],
30    // How many cells have used so far
31    pool_used: usize,
32}
33
34impl Default for Poseidon2Rng {
35    fn default() -> Self {
36        Self::new()
37    }
38}
39
40impl Poseidon2Rng {
41    /// Construct a new Poseidon2Rng
42    pub fn new() -> Self {
43        Self {
44            cells: [Elem::new(0); CELLS],
45            pool_used: 0,
46        }
47    }
48}
49
50impl Rng<BabyBear> for Poseidon2Rng {
51    fn mix(&mut self, val: &Digest) {
52        // if switching from squeezing, do a mix
53        if self.pool_used != 0 {
54            poseidon2_mix(&mut self.cells);
55            self.pool_used = 0;
56        }
57        // Add in CELLS_OUT elements (also # of digest words)
58        for i in 0..CELLS_OUT {
59            self.cells[i] += Elem::new_raw(val.as_words()[i]);
60        }
61        // Mix
62        poseidon2_mix(&mut self.cells);
63    }
64
65    fn random_bits(&mut self, bits: usize) -> u32 {
66        let mut val = self.random_elem().as_u32();
67        for _ in 0..3 {
68            let new_val = self.random_elem().as_u32();
69            if val == 0 {
70                val = new_val;
71            }
72        }
73        ((1 << bits) - 1) & val
74    }
75
76    fn random_elem(&mut self) -> Elem {
77        if self.pool_used == CELLS_RATE {
78            poseidon2_mix(&mut self.cells);
79            self.pool_used = 0;
80        }
81        let out = self.cells[self.pool_used];
82        self.pool_used += 1;
83        out
84    }
85
86    fn random_ext_elem(&mut self) -> BabyBearExtElem {
87        ExtElem::from_subelems((0..4).map(|_| self.random_elem()))
88    }
89}
risc0_zkp/core/hash/poseidon2/rng.rs

risc0_zkp/core/hash/poseidon2/
rng.rs
