1pub use winter_crypto::{RandomCoin, RandomCoinError};
  2use winter_math::{FieldElement, StarkField};
  3
  4use super::{Felt, FeltRng, Word, ZERO};
  5use crate::hash::rpo::{Rpo256, RpoDigest};
  6use crate::utils::collections::Vec;
  7use crate::utils::vec;
  8
  9// CONSTANTS
 10// ================================================================================================
 11
 12const STATE_WIDTH: usize = Rpo256::STATE_WIDTH;
 13const RATE_START: usize = Rpo256::RATE_RANGE.start;
 14const RATE_END: usize = Rpo256::RATE_RANGE.end;
 15const HALF_RATE_WIDTH: usize = (Rpo256::RATE_RANGE.end - Rpo256::RATE_RANGE.start) / 2;
 16
 17// RPO RANDOM COIN
 18// ================================================================================================
 19/// A simplified version of the `SPONGE_PRG` reseedable pseudo-random number generator algorithm
 20/// described in https://eprint.iacr.org/2011/499.pdf. The simplification is related to
 21/// the following facts:
 22/// 1. A call to the reseed method implies one and only one call to the permutation function.
 23///  This is possible because in our case we never reseed with more than 4 field elements.
 24/// 2. As a result of the previous point, we dont make use of an input buffer to accumulate seed
 25///  material.
 26pub struct RpoRandomCoin {
 27    state: [Felt; STATE_WIDTH],
 28    current: usize,
 29}
 30
 31impl RpoRandomCoin {
 32    fn draw_basefield(&mut self) -> Felt {
 33        if self.current == RATE_END {
 34            Rpo256::apply_permutation(&mut self.state);
 35            self.current = RATE_START;
 36        }
 37
 38        self.current += 1;
 39        self.state[self.current - 1]
 40    }
 41}
 42
 43impl RandomCoin for RpoRandomCoin {
 44    type BaseField = Felt;
 45    type Hasher = Rpo256;
 46
 47    fn new(seed: &[Self::BaseField]) -> Self {
 48        let mut state = [ZERO; STATE_WIDTH];
 49        let digest: Word = Rpo256::hash_elements(seed).into();
 50
 51        for i in 0..HALF_RATE_WIDTH {
 52            state[RATE_START + i] += digest[i];
 53        }
 54
 55        // Absorb
 56        Rpo256::apply_permutation(&mut state);
 57
 58        RpoRandomCoin { state, current: RATE_START }
 59    }
 60
 61    fn reseed(&mut self, data: RpoDigest) {
 62        // Reset buffer
 63        self.current = RATE_START;
 64
 65        // Add the new seed material to the first half of the rate portion of the RPO state
 66        let data: Word = data.into();
 67
 68        self.state[RATE_START] += data[0];
 69        self.state[RATE_START + 1] += data[1];
 70        self.state[RATE_START + 2] += data[2];
 71        self.state[RATE_START + 3] += data[3];
 72
 73        // Absorb
 74        Rpo256::apply_permutation(&mut self.state);
 75    }
 76
 77    fn check_leading_zeros(&self, value: u64) -> u32 {
 78        let value = Felt::new(value);
 79        let mut state_tmp = self.state;
 80
 81        state_tmp[RATE_START] += value;
 82
 83        Rpo256::apply_permutation(&mut state_tmp);
 84
 85        let first_rate_element = state_tmp[RATE_START].as_int();
 86        first_rate_element.trailing_zeros()
 87    }
 88
 89    fn draw<E: FieldElement<BaseField = Felt>>(&mut self) -> Result<E, RandomCoinError> {
 90        let ext_degree = E::EXTENSION_DEGREE;
 91        let mut result = vec![ZERO; ext_degree];
 92        for r in result.iter_mut().take(ext_degree) {
 93            *r = self.draw_basefield();
 94        }
 95
 96        let result = E::slice_from_base_elements(&result);
 97        Ok(result[0])
 98    }
 99
100    fn draw_integers(
101        &mut self,
102        num_values: usize,
103        domain_size: usize,
104        nonce: u64,
105    ) -> Result<Vec<usize>, RandomCoinError> {
106        assert!(domain_size.is_power_of_two(), "domain size must be a power of two");
107        assert!(num_values < domain_size, "number of values must be smaller than domain size");
108
109        // absorb the nonce
110        let nonce = Felt::new(nonce);
111        self.state[RATE_START] += nonce;
112        Rpo256::apply_permutation(&mut self.state);
113
114        // reset the buffer
115        self.current = RATE_START;
116
117        // determine how many bits are needed to represent valid values in the domain
118        let v_mask = (domain_size - 1) as u64;
119
120        // draw values from PRNG until we get as many unique values as specified by num_queries
121        let mut values = Vec::new();
122        for _ in 0..1000 {
123            // get the next pseudo-random field element
124            let value = self.draw_basefield().as_int();
125
126            // use the mask to get a value within the range
127            let value = (value & v_mask) as usize;
128
129            values.push(value);
130            if values.len() == num_values {
131                break;
132            }
133        }
134
135        if values.len() < num_values {
136            return Err(RandomCoinError::FailedToDrawIntegers(num_values, values.len(), 1000));
137        }
138
139        Ok(values)
140    }
141}
142
143impl FeltRng for RpoRandomCoin {
144    fn draw_element(&mut self) -> Felt {
145        self.draw_basefield()
146    }
147
148    fn draw_word(&mut self) -> Word {
149        let mut output = [ZERO; 4];
150        for o in output.iter_mut() {
151            *o = self.draw_basefield();
152        }
153        output
154    }
155}
156
157// TESTS
158// ================================================================================================
159
160#[cfg(all(test, feature = "std"))]
161mod tests {
162    use super::{FeltRng, RandomCoin, RpoRandomCoin, ZERO};
163
164    #[test]
165    fn test_randfeltsgen_felt() {
166        let mut rpocoin = RpoRandomCoin::new(&[ZERO; 4]);
167        let output = rpocoin.draw_element();
168
169        let mut rpocoin = RpoRandomCoin::new(&[ZERO; 4]);
170        let expected = rpocoin.draw_basefield();
171
172        assert_eq!(output, expected);
173    }
174
175    #[test]
176    fn test_randfeltsgen_word() {
177        let mut rpocoin = RpoRandomCoin::new(&[ZERO; 4]);
178        let output = rpocoin.draw_word();
179
180        let mut rpocoin = RpoRandomCoin::new(&[ZERO; 4]);
181        let mut expected = [ZERO; 4];
182        for o in expected.iter_mut() {
183            *o = rpocoin.draw_basefield();
184        }
185
186        assert_eq!(output, expected);
187    }
188}