use crate::{
arch::word::Word,
buffer::Buffer,
ibig::IBig,
ops::UnsignedAbs,
ubig::{Repr::*, UBig},
};
use rand::{
distributions::uniform::{SampleBorrow, SampleUniform, UniformSampler},
Rng,
};
impl SampleUniform for UBig {
type Sampler = UniformUBig;
}
impl SampleUniform for IBig {
type Sampler = UniformIBig;
}
impl UBig {
fn uniform<R>(range: &UBig, rng: &mut R) -> UBig
where
R: Rng + ?Sized,
{
debug_assert!(*range != UBig::from_word(0));
match range.repr() {
Small(word) => UBig::from_word(rng.gen_range(0..*word)),
Large(buffer) => UBig::uniform_large(buffer, rng),
}
}
fn uniform_large<R>(words: &[Word], rng: &mut R) -> UBig
where
R: Rng + ?Sized,
{
let mut buffer = Buffer::allocate(words.len());
buffer.push_zeros(words.len());
while !try_fill_uniform(words, rng, &mut buffer) {
}
buffer.into()
}
}
fn try_fill_uniform<R>(words: &[Word], rng: &mut R, result: &mut [Word]) -> bool
where
R: Rng + ?Sized,
{
let n = words.len();
debug_assert!(n > 0 && result.len() == n);
let mut i = n - 1;
result[i] = rng.gen_range(0..=words[i]);
while result[i] == words[i] {
if i == 0 {
return false;
}
i -= 1;
result[i] = rng.gen();
if result[i] > words[i] {
return false;
}
}
rng.fill(&mut result[..i]);
true
}
pub struct UniformUBig {
range: UBig,
offset: UBig,
}
impl UniformSampler for UniformUBig {
type X = UBig;
fn new<B1, B2>(low: B1, high: B2) -> UniformUBig
where
B1: SampleBorrow<UBig>,
B2: SampleBorrow<UBig>,
{
let range = high.borrow() - low.borrow();
if range == UBig::from_word(0) {
panic!("Empty range");
}
UniformUBig {
range,
offset: low.borrow().clone(),
}
}
fn new_inclusive<B1, B2>(low: B1, high: B2) -> UniformUBig
where
B1: SampleBorrow<UBig>,
B2: SampleBorrow<UBig>,
{
let range = high.borrow() - low.borrow() + UBig::from_word(1);
UniformUBig {
range,
offset: low.borrow().clone(),
}
}
fn sample<R>(&self, rng: &mut R) -> UBig
where
R: Rng + ?Sized,
{
UBig::uniform(&self.range, rng) + &self.offset
}
}
pub struct UniformIBig {
range: UBig,
offset: IBig,
}
impl UniformSampler for UniformIBig {
type X = IBig;
fn new<B1, B2>(low: B1, high: B2) -> UniformIBig
where
B1: SampleBorrow<IBig>,
B2: SampleBorrow<IBig>,
{
let range = high.borrow() - low.borrow();
if range <= IBig::from(0u8) {
panic!("Empty range");
}
UniformIBig {
range: range.unsigned_abs(),
offset: low.borrow().clone(),
}
}
fn new_inclusive<B1, B2>(low: B1, high: B2) -> UniformIBig
where
B1: SampleBorrow<IBig>,
B2: SampleBorrow<IBig>,
{
let range = high.borrow() - low.borrow() + IBig::from(1u8);
if range <= IBig::from(0u8) {
panic!("Empty range");
}
UniformIBig {
range: range.unsigned_abs(),
offset: low.borrow().clone(),
}
}
fn sample<R>(&self, rng: &mut R) -> IBig
where
R: Rng + ?Sized,
{
IBig::from(UBig::uniform(&self.range, rng)) + &self.offset
}
}