use crate::{
arch::word::{DoubleWord, Word},
buffer::Buffer,
ibig::IBig,
math::ceil_div,
ops::UnsignedAbs,
primitive::{DWORD_BITS_USIZE, WORD_BITS_USIZE},
repr::{Repr, TypedReprRef::*},
ubig::UBig,
};
use dashu_base::Sign;
use rand_v08::{
distributions::uniform::{SampleBorrow, SampleUniform, UniformSampler},
prelude::Distribution,
Rng,
};
pub struct UniformBits {
bits: usize,
}
impl UniformBits {
#[inline]
pub const fn new(bits: usize) -> Self {
UniformBits { bits }
}
}
impl Distribution<UBig> for UniformBits {
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> UBig {
if self.bits == 0 {
UBig::ZERO
} else if self.bits <= DWORD_BITS_USIZE {
let dword: DoubleWord = rng.gen();
UBig::from_dword(dword >> (DWORD_BITS_USIZE - self.bits))
} else {
let num_words = ceil_div(self.bits, WORD_BITS_USIZE);
let mut buffer = Buffer::allocate(num_words);
buffer.push_zeros(num_words);
rng.fill(buffer.as_mut());
let rem = self.bits % WORD_BITS_USIZE;
if rem != 0 {
*buffer.last_mut().unwrap() >>= WORD_BITS_USIZE - rem;
}
UBig(Repr::from_buffer(buffer))
}
}
}
impl Distribution<IBig> for UniformBits {
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> IBig {
loop {
let mag: UBig = self.sample(rng);
let neg = rng.gen();
if mag.is_zero() && neg {
continue;
}
break IBig::from_parts(Sign::from(neg), mag);
}
}
}
pub struct UniformBelow<'a> {
limit: &'a UBig,
}
impl<'a> UniformBelow<'a> {
#[inline]
pub const fn new(limit: &'a UBig) -> Self {
Self { limit }
}
}
impl<'a> Distribution<UBig> for UniformBelow<'a> {
#[inline]
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> UBig {
UBig::uniform(self.limit, rng)
}
}
impl<'a> Distribution<IBig> for UniformBelow<'a> {
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> IBig {
loop {
let mag: UBig = UBig::uniform(self.limit, rng);
let neg = rng.gen();
if mag.is_zero() && neg {
continue;
}
break IBig::from_parts(Sign::from(neg), mag);
}
}
}
impl UBig {
#[inline]
fn uniform<R>(range: &UBig, rng: &mut R) -> UBig
where
R: Rng + ?Sized,
{
debug_assert!(!range.is_zero());
match range.repr() {
RefSmall(dword) => UBig::from(rng.gen_range(0..dword)),
RefLarge(words) => UBig::uniform_large(words, 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) {
}
UBig(Repr::from_buffer(buffer))
}
}
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,
}
const fn panic_empty_range() -> ! {
panic!("empty range for random generation")
}
impl UniformSampler for UniformUBig {
type X = UBig;
#[inline]
fn new<B1, B2>(low: B1, high: B2) -> UniformUBig
where
B1: SampleBorrow<UBig>,
B2: SampleBorrow<UBig>,
{
if high.borrow() <= low.borrow() {
panic_empty_range()
}
let range = high.borrow() - low.borrow();
UniformUBig {
range,
offset: low.borrow().clone(),
}
}
#[inline]
fn new_inclusive<B1, B2>(low: B1, high: B2) -> UniformUBig
where
B1: SampleBorrow<UBig>,
B2: SampleBorrow<UBig>,
{
if high.borrow() < low.borrow() {
panic_empty_range()
}
let range = high.borrow() - low.borrow() + UBig::ONE;
UniformUBig {
range,
offset: low.borrow().clone(),
}
}
#[inline]
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;
#[inline]
fn new<B1, B2>(low: B1, high: B2) -> UniformIBig
where
B1: SampleBorrow<IBig>,
B2: SampleBorrow<IBig>,
{
if high.borrow() <= low.borrow() {
panic_empty_range()
}
let range = high.borrow() - low.borrow();
UniformIBig {
range: range.unsigned_abs(),
offset: low.borrow().clone(),
}
}
#[inline]
fn new_inclusive<B1, B2>(low: B1, high: B2) -> UniformIBig
where
B1: SampleBorrow<IBig>,
B2: SampleBorrow<IBig>,
{
if high.borrow() < low.borrow() {
panic_empty_range()
}
let range = high.borrow() - low.borrow() + IBig::ONE;
UniformIBig {
range: range.unsigned_abs(),
offset: low.borrow().clone(),
}
}
#[inline]
fn sample<R>(&self, rng: &mut R) -> IBig
where
R: Rng + ?Sized,
{
IBig::from(UBig::uniform(&self.range, rng)) + &self.offset
}
}
impl SampleUniform for UBig {
type Sampler = UniformUBig;
}
impl SampleUniform for IBig {
type Sampler = UniformIBig;
}