zkp_u256/
rand.rs

1use crate::{Binary, U256};
2use rand::{
3    distributions::{
4        uniform::{SampleBorrow, SampleUniform, UniformSampler},
5        Distribution, Standard,
6    },
7    Rng,
8};
9
10/// Draw from a uniform distribution over all values.
11impl Distribution<U256> for Standard {
12    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> U256 {
13        U256::from_limbs([rng.gen(), rng.gen(), rng.gen(), rng.gen()])
14    }
15}
16
17/// Helper struct for uniform sampling using `rand`
18#[derive(Clone, Debug)]
19pub enum UniformU256 {
20    Full,
21    Ranged { low: U256, range: U256 },
22}
23
24impl SampleUniform for U256 {
25    type Sampler = UniformU256;
26}
27
28impl UniformSampler for UniformU256 {
29    type X = U256;
30
31    fn new<B1, B2>(low: B1, high: B2) -> Self
32    where
33        B1: SampleBorrow<U256> + Sized,
34        B2: SampleBorrow<U256> + Sized,
35    {
36        let low = low.borrow().clone();
37        let range = high.borrow() - &low;
38        Self::Ranged { low, range }
39    }
40
41    fn new_inclusive<B1, B2>(low: B1, high: B2) -> Self
42    where
43        B1: SampleBorrow<U256> + Sized,
44        B2: SampleBorrow<U256> + Sized,
45    {
46        if low.borrow() == &U256::ZERO && high.borrow() == &U256::MAX {
47            Self::Full
48        } else {
49            let low = low.borrow().clone();
50            let range = high.borrow() - &low + U256::ONE;
51            Self::Ranged { low, range }
52        }
53    }
54
55    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> U256 {
56        match self {
57            Self::Full => rng.gen(),
58            Self::Ranged { low, range } => {
59                // Strategy: bitshift to within 2x then rejection sample.
60                let shift = range.leading_zeros();
61                let mut result = U256::MAX;
62                while result >= *range {
63                    result = rng.gen();
64                    result >>= shift;
65                }
66                result + low
67            }
68        }
69    }
70}
71
72// TODO: Tests
zkp_u256/rand.rs

zkp_u256/
rand.rs
