use crate::{rbig::RBig, repr::Repr, Relaxed};
use dashu_base::Gcd;
use dashu_int::{
rand::{UniformBelow, UniformIBig},
DoubleWord, IBig, UBig,
};
use rand_v08::{
distributions::{
uniform::{SampleBorrow, SampleUniform, UniformSampler},
Open01, OpenClosed01, Standard,
},
prelude::Distribution,
Rng,
};
pub struct Uniform01<'a> {
limit: &'a UBig,
include_zero: bool, include_one: bool, }
impl<'a> Uniform01<'a> {
#[inline]
pub fn new(limit: &'a UBig) -> Self {
Self {
limit,
include_zero: true,
include_one: false,
}
}
#[inline]
pub fn new_closed(limit: &'a UBig) -> Self {
Self {
limit,
include_zero: true,
include_one: true,
}
}
#[inline]
pub fn new_open(limit: &'a UBig) -> Self {
Self {
limit,
include_zero: false,
include_one: false,
}
}
#[inline]
pub fn new_open_closed(limit: &'a UBig) -> Self {
Self {
limit,
include_zero: false,
include_one: true,
}
}
}
impl<'a> Distribution<Repr> for Uniform01<'a> {
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Repr {
match (self.include_zero, self.include_one) {
(true, false) => {
let num: UBig = UniformBelow::new(self.limit).sample(rng);
Repr {
numerator: num.into(),
denominator: self.limit.clone(),
}
}
(true, true) => {
let num: UBig = UniformBelow::new(self.limit).sample(rng);
Repr {
numerator: num.into(),
denominator: self.limit.clone() - UBig::ONE,
}
}
(false, false) => {
let num = loop {
let n: UBig = UniformBelow::new(self.limit).sample(rng);
if !n.is_zero() {
break n;
}
};
Repr {
numerator: num.into(),
denominator: self.limit.clone(),
}
}
(false, true) => {
let num: UBig = UniformBelow::new(self.limit).sample(rng);
Repr {
numerator: (num + UBig::ONE).into(),
denominator: self.limit.clone(),
}
}
}
}
}
impl<'a> Distribution<RBig> for Uniform01<'a> {
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> RBig {
RBig(Distribution::<Repr>::sample(self, rng).reduce())
}
}
impl<'a> Distribution<Relaxed> for Uniform01<'a> {
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Relaxed {
Relaxed(Distribution::<Repr>::sample(self, rng).reduce2())
}
}
pub struct UniformRBig {
num_sampler: UniformIBig,
den: UBig,
}
impl UniformRBig {
fn parse_bounds<B1, B2>(low: B1, high: B2) -> (IBig, IBig, UBig)
where
B1: SampleBorrow<RBig> + Sized,
B2: SampleBorrow<RBig> + Sized,
{
let (low_d, high_d) = (low.borrow().denominator(), high.borrow().denominator());
let g = low_d.gcd(high_d);
let low_n = high_d / &g * low.borrow().numerator();
let high_n = low_d / &g * high.borrow().numerator();
let den = low_d / g * high_d;
(low_n, high_n, den)
}
}
impl UniformSampler for UniformRBig {
type X = RBig;
#[inline]
fn new<B1, B2>(low: B1, high: B2) -> Self
where
B1: SampleBorrow<Self::X> + Sized,
B2: SampleBorrow<Self::X> + Sized,
{
let (low_n, high_n, den) = Self::parse_bounds(low, high);
UniformRBig {
num_sampler: UniformIBig::new(low_n, high_n),
den,
}
}
#[inline]
fn new_inclusive<B1, B2>(low: B1, high: B2) -> Self
where
B1: SampleBorrow<Self::X> + Sized,
B2: SampleBorrow<Self::X> + Sized,
{
let (low_n, high_n, den) = Self::parse_bounds(low, high);
UniformRBig {
num_sampler: UniformIBig::new_inclusive(low_n, high_n),
den,
}
}
#[inline]
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> RBig {
RBig::from_parts(self.num_sampler.sample(rng), self.den.clone())
}
}
impl SampleUniform for RBig {
type Sampler = UniformRBig;
}
macro_rules! impl_builtin_distr {
(impl $trait:ident for $t:ty, $ctor:ident) => {
impl Distribution<$t> for $trait {
#[inline]
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> $t {
let limit = UBig::from_dword(DoubleWord::MAX);
Uniform01::$ctor(&limit).sample(rng)
}
}
};
}
impl_builtin_distr!(impl Standard for RBig, new);
impl_builtin_distr!(impl Standard for Relaxed, new);
impl_builtin_distr!(impl Open01 for RBig, new_open);
impl_builtin_distr!(impl Open01 for Relaxed, new_open);
impl_builtin_distr!(impl OpenClosed01 for RBig, new_open_closed);
impl_builtin_distr!(impl OpenClosed01 for Relaxed, new_open_closed);