Skip to content

Commit 00869d7

Browse files
dhardydhardy
committed
Impl ONeill, Canon, Canon-Lemire and Bitmask methods for integer types
This is based on @TheIronBorn's work (#1154, #1172), with some changes.
1 parent 255ff71 commit 00869d7

File tree

1 file changed

+182
-16
lines changed

1 file changed

+182
-16
lines changed

‎src/distributions/uniform.rs

Lines changed: 182 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -433,7 +433,7 @@ pub struct UniformInt<X> {
433433
}
434434

435435
macro_rules! uniform_int_impl {
436-
($ty:ty, $unsigned:ident, $u_large:ident) => {
436+
($ty:ty, $unsigned:ident, $u_large:ident, $u_extra_large:ident) => {
437437
impl SampleUniform for $ty {
438438
type Sampler = UniformInt<$ty>;
439439
}
@@ -536,9 +536,8 @@ macro_rules! uniform_int_impl {
536536
"UniformSampler::sample_single_inclusive: low > high"
537537
);
538538
let range = high.wrapping_sub(low).wrapping_add(1) as $unsigned as $u_large;
539-
// If the above resulted in wrap-around to 0, the range is $ty::MIN..=$ty::MAX,
540-
// and any integer will do.
541539
if range == 0 {
540+
// Range is MAX+1 (unrepresentable), so we need a special case
542541
return rng.gen();
543542
}
544543

@@ -564,21 +563,188 @@ macro_rules! uniform_int_impl {
564563
}
565564
}
566565
}
566+
567+
impl UniformInt<$ty> {
568+
/// Sample single inclusive, using ONeill's method
569+
#[inline]
570+
pub fn sample_single_inclusive_oneill<R: Rng + ?Sized, B1, B2>(
571+
low_b: B1, high_b: B2, rng: &mut R,
572+
) -> $ty
573+
where
574+
B1: SampleBorrow<$ty> + Sized,
575+
B2: SampleBorrow<$ty> + Sized,
576+
{
577+
let low = *low_b.borrow();
578+
let high = *high_b.borrow();
579+
assert!(
580+
low <= high,
581+
"UniformSampler::sample_single_inclusive: low > high"
582+
);
583+
let range = high.wrapping_sub(low).wrapping_add(1) as $unsigned as $u_large;
584+
if range == 0 {
585+
// Range is MAX+1 (unrepresentable), so we need a special case
586+
return rng.gen();
587+
}
588+
589+
// we use the "Debiased Int Mult (t-opt, m-opt)" rejection sampling method
590+
// described here https://www.pcg-random.org/posts/bounded-rands.html
591+
// and here https://github.com/imneme/bounded-rands
592+
593+
let (mut hi, mut lo) = rng.gen::<$u_large>().wmul(range);
594+
if lo < range {
595+
let mut threshold = range.wrapping_neg();
596+
// this shortcut works best with large ranges
597+
if threshold >= range {
598+
threshold -= range;
599+
if threshold >= range {
600+
threshold %= range;
601+
}
602+
}
603+
while lo < threshold {
604+
let (new_hi, new_lo) = rng.gen::<$u_large>().wmul(range);
605+
hi = new_hi;
606+
lo = new_lo;
607+
}
608+
}
609+
low.wrapping_add(hi as $ty)
610+
}
611+
612+
/// Sample single inclusive, using Canon's method
613+
#[inline]
614+
pub fn sample_single_inclusive_canon<R: Rng + ?Sized, B1, B2>(
615+
low_b: B1, high_b: B2, rng: &mut R,
616+
) -> $ty
617+
where
618+
B1: SampleBorrow<$ty> + Sized,
619+
B2: SampleBorrow<$ty> + Sized,
620+
{
621+
let low = *low_b.borrow();
622+
let high = *high_b.borrow();
623+
assert!(
624+
low <= high,
625+
"UniformSampler::sample_single_inclusive: low > high"
626+
);
627+
let range = high.wrapping_sub(low).wrapping_add(1) as $unsigned as $u_extra_large;
628+
if range == 0 {
629+
// Range is MAX+1 (unrepresentable), so we need a special case
630+
return rng.gen();
631+
}
632+
633+
// generate a sample using a sensible integer type
634+
let (mut result, lo_order) = rng.gen::<$u_extra_large>().wmul(range);
635+
636+
// if the sample is biased...
637+
if lo_order > range.wrapping_neg() {
638+
// ...generate a new sample with 64 more bits, enough that bias is undetectable
639+
let (new_hi_order, _) =
640+
(rng.gen::<$u_extra_large>()).wmul(range as $u_extra_large);
641+
// and adjust if needed
642+
result += lo_order
643+
.checked_add(new_hi_order as $u_extra_large)
644+
.is_none() as $u_extra_large;
645+
}
646+
647+
low.wrapping_add(result as $ty)
648+
}
649+
650+
/// Sample single inclusive, using Canon's method with Lemire's early-out
651+
#[inline]
652+
pub fn sample_inclusive_canon_lemire<R: Rng + ?Sized, B1, B2>(
653+
low_b: B1, high_b: B2, rng: &mut R,
654+
) -> $ty
655+
where
656+
B1: SampleBorrow<$ty> + Sized,
657+
B2: SampleBorrow<$ty> + Sized,
658+
{
659+
let low = *low_b.borrow();
660+
let high = *high_b.borrow();
661+
assert!(
662+
low <= high,
663+
"UniformSampler::sample_single_inclusive: low > high"
664+
);
665+
let range = high.wrapping_sub(low).wrapping_add(1) as $unsigned as $u_extra_large;
666+
if range == 0 {
667+
// Range is MAX+1 (unrepresentable), so we need a special case
668+
return rng.gen();
669+
}
670+
671+
// generate a sample using a sensible integer type
672+
let (mut result, lo_order) = rng.gen::<$u_extra_large>().wmul(range);
673+
674+
// if the sample is biased... (since range won't be changing we can further
675+
// improve this check with a modulo)
676+
if lo_order < range.wrapping_neg() % range {
677+
// ...generate a new sample with 64 more bits, enough that bias is undetectable
678+
let (new_hi_order, _) =
679+
(rng.gen::<$u_extra_large>()).wmul(range as $u_extra_large);
680+
// and adjust if needed
681+
result += lo_order
682+
.checked_add(new_hi_order as $u_extra_large)
683+
.is_none() as $u_extra_large;
684+
}
685+
686+
low.wrapping_add(result as $ty)
687+
}
688+
689+
/// Sample single inclusive, using the Bitmask method
690+
#[inline]
691+
pub fn sample_single_inclusive_bitmask<R: Rng + ?Sized, B1, B2>(
692+
low_b: B1, high_b: B2, rng: &mut R,
693+
) -> $ty
694+
where
695+
B1: SampleBorrow<$ty> + Sized,
696+
B2: SampleBorrow<$ty> + Sized,
697+
{
698+
let low = *low_b.borrow();
699+
let high = *high_b.borrow();
700+
assert!(
701+
low <= high,
702+
"UniformSampler::sample_single_inclusive: low > high"
703+
);
704+
let mut range = high.wrapping_sub(low).wrapping_add(1) as $unsigned as $u_large;
705+
if range == 0 {
706+
// Range is MAX+1 (unrepresentable), so we need a special case
707+
return rng.gen();
708+
}
709+
710+
// the old impl use a mix of methods for different integer sizes, we only use
711+
// the lz method here for a better comparison.
712+
713+
let mut mask = $u_large::max_value();
714+
range -= 1;
715+
mask >>= (range | 1).leading_zeros();
716+
loop {
717+
let x = rng.gen::<$u_large>() & mask;
718+
if x <= range {
719+
return low.wrapping_add(x as $ty);
720+
}
721+
}
722+
}
723+
}
567724
};
568725
}
569-
570-
uniform_int_impl! { i8, u8, u32 }
571-
uniform_int_impl! { i16, u16, u32 }
572-
uniform_int_impl! { i32, u32, u32 }
573-
uniform_int_impl! { i64, u64, u64 }
574-
uniform_int_impl! { i128, u128, u128 }
575-
uniform_int_impl! { isize, usize, usize }
576-
uniform_int_impl! { u8, u8, u32 }
577-
uniform_int_impl! { u16, u16, u32 }
578-
uniform_int_impl! { u32, u32, u32 }
579-
uniform_int_impl! { u64, u64, u64 }
580-
uniform_int_impl! { usize, usize, usize }
581-
uniform_int_impl! { u128, u128, u128 }
726+
uniform_int_impl! { i8, u8, u32, u64 }
727+
uniform_int_impl! { i16, u16, u32, u64 }
728+
uniform_int_impl! { i32, u32, u32, u64 }
729+
uniform_int_impl! { i64, u64, u64, u64 }
730+
uniform_int_impl! { i128, u128, u128, u128 }
731+
uniform_int_impl! { u8, u8, u32, u64 }
732+
uniform_int_impl! { u16, u16, u32, u64 }
733+
uniform_int_impl! { u32, u32, u32, u64 }
734+
uniform_int_impl! { u64, u64, u64, u64 }
735+
uniform_int_impl! { u128, u128, u128, u128 }
736+
#[cfg(any(target_pointer_width = "16", target_pointer_width = "32",))]
737+
mod isize_int_impls {
738+
use super::*;
739+
uniform_int_impl! { isize, usize, usize, u64 }
740+
uniform_int_impl! { usize, usize, usize, u64 }
741+
}
742+
#[cfg(not(any(target_pointer_width = "16", target_pointer_width = "32",)))]
743+
mod isize_int_impls {
744+
use super::*;
745+
uniform_int_impl! { isize, usize, usize, usize }
746+
uniform_int_impl! { usize, usize, usize, usize }
747+
}
582748

583749
#[cfg(feature = "simd_support")]
584750
macro_rules! uniform_simd_int_impl {

0 commit comments

Comments
 (0)