1use rand::distr::uniform::{self, SampleBorrow, SampleUniform, UniformSampler};
4use rand::prelude::*;
5use rand::Rng;
6
7use crate::BigInt;
8use crate::BigUint;
9use crate::Sign::*;
10
11use crate::big_digit::BigDigit;
12use crate::bigint::{into_magnitude, magnitude};
13use crate::integer::Integer;
14#[cfg(feature = "prime")]
15use num_iter::range_step;
16use num_traits::Zero;
17#[cfg(feature = "prime")]
18use num_traits::{FromPrimitive, ToPrimitive};
19#[cfg(feature = "prime")]
20use once_cell::sync::Lazy;
21
22#[cfg(feature = "prime")]
23use crate::prime::probably_prime;
24
25pub trait RandBigInt {
26 fn gen_biguint(&mut self, bit_size: usize) -> BigUint;
28
29 fn gen_bigint(&mut self, bit_size: usize) -> BigInt;
31
32 fn gen_biguint_below(&mut self, bound: &BigUint) -> BigUint;
35
36 fn gen_biguint_range(&mut self, lbound: &BigUint, ubound: &BigUint) -> BigUint;
40
41 fn gen_bigint_range(&mut self, lbound: &BigInt, ubound: &BigInt) -> BigInt;
45}
46
47impl<R: Rng + ?Sized> RandBigInt for R {
48 fn gen_biguint(&mut self, bit_size: usize) -> BigUint {
49 use super::big_digit::BITS;
50 let (digits, rem) = bit_size.div_rem(&BITS);
51 let mut data = smallvec![BigDigit::default(); digits + (rem > 0) as usize];
52
53 self.fill(data.as_mut_slice());
56
57 if rem > 0 {
58 data[digits] >>= BITS - rem;
59 }
60 BigUint::new_native(data)
61 }
62
63 fn gen_bigint(&mut self, bit_size: usize) -> BigInt {
64 loop {
65 let biguint = self.gen_biguint(bit_size);
67 let sign = if biguint.is_zero() {
69 if self.random() {
74 continue;
75 } else {
76 NoSign
77 }
78 } else if self.random() {
79 Plus
80 } else {
81 Minus
82 };
83 return BigInt::from_biguint(sign, biguint);
84 }
85 }
86
87 fn gen_biguint_below(&mut self, bound: &BigUint) -> BigUint {
88 assert!(!bound.is_zero());
89 let bits = bound.bits();
90 loop {
91 let n = self.gen_biguint(bits);
92 if n < *bound {
93 return n;
94 }
95 }
96 }
97
98 fn gen_biguint_range(&mut self, lbound: &BigUint, ubound: &BigUint) -> BigUint {
99 assert!(*lbound < *ubound);
100 if lbound.is_zero() {
101 self.gen_biguint_below(ubound)
102 } else {
103 lbound + self.gen_biguint_below(&(ubound - lbound))
104 }
105 }
106
107 fn gen_bigint_range(&mut self, lbound: &BigInt, ubound: &BigInt) -> BigInt {
108 assert!(*lbound < *ubound);
109 if lbound.is_zero() {
110 BigInt::from(self.gen_biguint_below(magnitude(&ubound)))
111 } else if ubound.is_zero() {
112 lbound + BigInt::from(self.gen_biguint_below(magnitude(&lbound)))
113 } else {
114 let delta = ubound - lbound;
115 lbound + BigInt::from(self.gen_biguint_below(magnitude(&delta)))
116 }
117 }
118}
119
120#[derive(Clone, Debug)]
122pub struct UniformBigUint {
123 base: BigUint,
124 len: BigUint,
125}
126
127impl UniformSampler for UniformBigUint {
128 type X = BigUint;
129
130 #[inline]
131 fn new<B1, B2>(low_b: B1, high_b: B2) -> Result<Self, uniform::Error>
132 where
133 B1: SampleBorrow<Self::X> + Sized,
134 B2: SampleBorrow<Self::X> + Sized,
135 {
136 let low = low_b.borrow();
137 let high = high_b.borrow();
138
139 assert!(low < high);
140
141 Ok(UniformBigUint {
142 len: high - low,
143 base: low.clone(),
144 })
145 }
146
147 #[inline]
148 fn new_inclusive<B1, B2>(low_b: B1, high_b: B2) -> Result<Self, uniform::Error>
149 where
150 B1: SampleBorrow<Self::X> + Sized,
151 B2: SampleBorrow<Self::X> + Sized,
152 {
153 Self::new(low_b, high_b.borrow() + 1u32)
154 }
155
156 #[inline]
157 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Self::X {
158 &self.base + rng.gen_biguint_below(&self.len)
159 }
160
161 #[inline]
162 fn sample_single<R: Rng + ?Sized, B1, B2>(
163 low_b: B1,
164 high_b: B2,
165 rng: &mut R,
166 ) -> Result<Self::X, uniform::Error>
167 where
168 B1: SampleBorrow<Self::X> + Sized,
169 B2: SampleBorrow<Self::X> + Sized,
170 {
171 let low = low_b.borrow();
172 let high = high_b.borrow();
173
174 Ok(rng.gen_biguint_range(low, high))
175 }
176}
177
178impl SampleUniform for BigUint {
179 type Sampler = UniformBigUint;
180}
181
182#[derive(Clone, Debug)]
184pub struct UniformBigInt {
185 base: BigInt,
186 len: BigUint,
187}
188
189impl UniformSampler for UniformBigInt {
190 type X = BigInt;
191
192 #[inline]
193 fn new<B1, B2>(low_b: B1, high_b: B2) -> Result<Self, uniform::Error>
194 where
195 B1: SampleBorrow<Self::X> + Sized,
196 B2: SampleBorrow<Self::X> + Sized,
197 {
198 let low = low_b.borrow();
199 let high = high_b.borrow();
200
201 assert!(low < high);
202 Ok(UniformBigInt {
203 len: into_magnitude(high - low),
204 base: low.clone(),
205 })
206 }
207
208 #[inline]
209 fn new_inclusive<B1, B2>(low_b: B1, high_b: B2) -> Result<Self, uniform::Error>
210 where
211 B1: SampleBorrow<Self::X> + Sized,
212 B2: SampleBorrow<Self::X> + Sized,
213 {
214 let low = low_b.borrow();
215 let high = high_b.borrow();
216
217 assert!(low <= high);
218 Self::new(low, high + 1u32)
219 }
220
221 #[inline]
222 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Self::X {
223 &self.base + BigInt::from(rng.gen_biguint_below(&self.len))
224 }
225
226 #[inline]
227 fn sample_single<R: Rng + ?Sized, B1, B2>(
228 low_b: B1,
229 high_b: B2,
230 rng: &mut R,
231 ) -> Result<Self::X, uniform::Error>
232 where
233 B1: SampleBorrow<Self::X> + Sized,
234 B2: SampleBorrow<Self::X> + Sized,
235 {
236 let low = low_b.borrow();
237 let high = high_b.borrow();
238
239 Ok(rng.gen_bigint_range(low, high))
240 }
241}
242
243impl SampleUniform for BigInt {
244 type Sampler = UniformBigInt;
245}
246
247#[derive(Clone, Copy, Debug)]
249pub struct RandomBits {
250 bits: usize,
251}
252
253impl RandomBits {
254 #[inline]
255 pub fn new(bits: usize) -> RandomBits {
256 RandomBits { bits }
257 }
258}
259
260impl Distribution<BigUint> for RandomBits {
261 #[inline]
262 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> BigUint {
263 rng.gen_biguint(self.bits)
264 }
265}
266
267impl Distribution<BigInt> for RandomBits {
268 #[inline]
269 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> BigInt {
270 rng.gen_bigint(self.bits)
271 }
272}
273
274#[cfg_attr(feature = "std", doc = " ```")]
281#[cfg_attr(not(feature = "std"), doc = " ```ignore")]
282#[cfg(feature = "prime")]
294pub trait RandPrime {
295 fn gen_prime(&mut self, bits: usize) -> BigUint;
297}
298
299#[cfg(feature = "prime")]
305const SMALL_PRIMES: [u8; 15] = [3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53];
306
307#[cfg(feature = "prime")]
312static SMALL_PRIMES_PRODUCT: Lazy<BigUint> =
313 Lazy::new(|| BigUint::from_u64(16_294_579_238_595_022_365).unwrap());
314
315#[cfg(feature = "prime")]
316impl<R: Rng + ?Sized> RandPrime for R {
317 fn gen_prime(&mut self, bit_size: usize) -> BigUint {
318 if bit_size < 2 {
319 panic!("prime size must be at least 2-bit");
320 }
321
322 let mut b = bit_size % 8;
323 if b == 0 {
324 b = 8;
325 }
326
327 let bytes_len = (bit_size + 7) / 8;
328 let mut bytes = alloc::vec![0u8; bytes_len];
329
330 loop {
331 self.fill_bytes(&mut bytes);
332 bytes[0] &= ((1u32 << (b as u32)) - 1) as u8;
334
335 if b >= 2 {
340 bytes[0] |= 3u8.wrapping_shl(b as u32 - 2);
341 } else {
342 bytes[0] |= 1;
344 if bytes_len > 1 {
345 bytes[1] |= 0x80;
346 }
347 }
348
349 bytes[bytes_len - 1] |= 1u8;
351
352 let mut p = BigUint::from_bytes_be(&bytes);
353 let rem = (&p % &*SMALL_PRIMES_PRODUCT).to_u64().unwrap();
355
356 'next: for delta in range_step(0, 1 << 20, 2) {
357 let m = rem + delta;
358
359 for prime in &SMALL_PRIMES {
360 if m % u64::from(*prime) == 0 && (bit_size > 6 || m != u64::from(*prime)) {
361 continue 'next;
362 }
363 }
364
365 if delta > 0 {
366 p += BigUint::from_u64(delta).unwrap();
367 }
368
369 break;
370 }
371
372 if p.bits() == bit_size && probably_prime(&p, 20) {
375 return p;
376 }
377 }
378 }
379}