arkhe_rand/
range.rs

1//! Lemire `nearlydivisionless` unbiased range sampling.
2//!
3//! [`gen_range`] / [`gen_range_inclusive`] produce uniform integers
4//! from a `Range` / `RangeInclusive` without modulo bias. Each call
5//! consumes 4 or 8 bytes per accepted sample (rejection rate is
6//! bounded by `2 * range_size / 2^32` for u32 and `2 * range_size /
7//! 2^64` for u64); seed bytes are sourced via
8//! [`RngSource::fill_bytes`].
9//!
10//! # Type support
11//!
12//! The sealed [`RandInt`] trait is implemented for `u8` / `u16` /
13//! `u32` / `u64` / `usize`. Signed and `u128` are excluded —
14//! Lemire's widening multiplication is only defined for u32 → u64
15//! and u64 → u128.
16
17use core::ops::{Range, RangeInclusive};
18
19use crate::RngSource;
20
21mod sealed {
22    pub trait Sealed {}
23}
24
25/// Integer types accepted by [`gen_range`] / [`gen_range_inclusive`].
26///
27/// Sealed: only `u8` / `u16` / `u32` / `u64` / `usize` implement this
28/// trait.
29pub trait RandInt: sealed::Sealed + Copy + PartialOrd {
30    #[doc(hidden)]
31    fn lemire_bounded(rng: &mut RngSource, range_size: Self) -> Self;
32
33    #[doc(hidden)]
34    fn add(self, n: Self) -> Self;
35
36    #[doc(hidden)]
37    fn sub(self, n: Self) -> Self;
38
39    #[doc(hidden)]
40    fn one() -> Self;
41}
42
43macro_rules! impl_rand_int_via_u32 {
44    ($($t:ty),* $(,)?) => {
45        $(
46            impl sealed::Sealed for $t {}
47            impl RandInt for $t {
48                #[inline]
49                fn lemire_bounded(rng: &mut RngSource, range_size: Self) -> Self {
50                    lemire_u32(rng, range_size as u32) as Self
51                }
52                #[inline]
53                fn add(self, n: Self) -> Self { self + n }
54                #[inline]
55                fn sub(self, n: Self) -> Self { self - n }
56                #[inline]
57                fn one() -> Self { 1 }
58            }
59        )*
60    };
61}
62
63impl_rand_int_via_u32!(u8, u16, u32);
64
65impl sealed::Sealed for u64 {}
66impl RandInt for u64 {
67    #[inline]
68    fn lemire_bounded(rng: &mut RngSource, range_size: Self) -> Self {
69        lemire_u64(rng, range_size)
70    }
71    #[inline]
72    fn add(self, n: Self) -> Self {
73        self + n
74    }
75    #[inline]
76    fn sub(self, n: Self) -> Self {
77        self - n
78    }
79    #[inline]
80    fn one() -> Self {
81        1
82    }
83}
84
85impl sealed::Sealed for usize {}
86impl RandInt for usize {
87    #[inline]
88    fn lemire_bounded(rng: &mut RngSource, range_size: Self) -> Self {
89        #[cfg(target_pointer_width = "32")]
90        {
91            lemire_u32(rng, range_size as u32) as usize
92        }
93        #[cfg(target_pointer_width = "64")]
94        {
95            lemire_u64(rng, range_size as u64) as usize
96        }
97    }
98    #[inline]
99    fn add(self, n: Self) -> Self {
100        self + n
101    }
102    #[inline]
103    fn sub(self, n: Self) -> Self {
104        self - n
105    }
106    #[inline]
107    fn one() -> Self {
108        1
109    }
110}
111
112/// Sample a uniformly-random value from `range` ∈ `[start, end)`.
113///
114/// # Panics
115///
116/// Debug builds panic on empty range (`start >= end`). Release builds
117/// elide the assert (std `Range` API symmetry).
118pub fn gen_range<T: RandInt>(rng: &mut RngSource, range: Range<T>) -> T {
119    debug_assert!(range.start < range.end, "gen_range: empty range");
120    let size = range.end.sub(range.start);
121    range.start.add(T::lemire_bounded(rng, size))
122}
123
124/// Sample a uniformly-random value from `range` ∈ `[start, end]`.
125///
126/// # Panics
127///
128/// Debug builds panic on inverted range (`start > end`) or when the
129/// inclusive size overflows `T` (`start == T::MIN && end == T::MAX`
130/// edge case for tightly-packed integer types). Use [`gen_range`]
131/// with manual high-bound handling if you need full type range.
132pub fn gen_range_inclusive<T: RandInt>(rng: &mut RngSource, range: RangeInclusive<T>) -> T {
133    let (start, end) = range.into_inner();
134    debug_assert!(start <= end, "gen_range_inclusive: inverted range");
135    let size = end.sub(start).add(T::one());
136    start.add(T::lemire_bounded(rng, size))
137}
138
139/// Lemire's `nearlydivisionless` for u32. Returns `r ∈ [0, n)` uniformly.
140#[inline]
141fn lemire_u32(rng: &mut RngSource, n: u32) -> u32 {
142    debug_assert!(n > 0, "lemire_u32: n == 0");
143    let mut buf = [0u8; 4];
144    loop {
145        rng.fill_bytes(&mut buf);
146        let x = u32::from_le_bytes(buf);
147        let m = u64::from(x) * u64::from(n);
148        let l = m as u32;
149        if l < n {
150            let t = n.wrapping_neg() % n;
151            if l < t {
152                continue;
153            }
154        }
155        return (m >> 32) as u32;
156    }
157}
158
159/// Lemire's `nearlydivisionless` for u64. Returns `r ∈ [0, n)` uniformly.
160#[inline]
161fn lemire_u64(rng: &mut RngSource, n: u64) -> u64 {
162    debug_assert!(n > 0, "lemire_u64: n == 0");
163    let mut buf = [0u8; 8];
164    loop {
165        rng.fill_bytes(&mut buf);
166        let x = u64::from_le_bytes(buf);
167        let m = u128::from(x) * u128::from(n);
168        let l = m as u64;
169        if l < n {
170            let t = n.wrapping_neg() % n;
171            if l < t {
172                continue;
173            }
174        }
175        return (m >> 64) as u64;
176    }
177}
arkhe_rand/range.rs

arkhe_rand/
range.rs
