Skip to content

Commit a87f4e9

Browse files
dhardydhardy
committed
UniformInt(single): add ONeill, Canon, Canon-Lemire and Bitmask methods
This is based on @TheIronBorn's work (#1154, #1172), with some changes.
1 parent f080e47 commit a87f4e9

File tree

1 file changed

+182
-16
lines changed

1 file changed

+182
-16
lines changed

‎src/distributions/uniform/uniform_int.rs

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

5757
macro_rules! uniform_int_impl {
58-
($ty:ty, $unsigned:ident, $u_large:ident) => {
58+
($ty:ty, $unsigned:ident, $u_large:ident, $u_extra_large:ident) => {
5959
impl SampleUniform for $ty {
6060
type Sampler = UniformInt<$ty>;
6161
}
@@ -158,9 +158,8 @@ macro_rules! uniform_int_impl {
158158
"UniformSampler::sample_single_inclusive: low > high"
159159
);
160160
let range = high.wrapping_sub(low).wrapping_add(1) as $unsigned as $u_large;
161-
// If the above resulted in wrap-around to 0, the range is $ty::MIN..=$ty::MAX,
162-
// and any integer will do.
163161
if range == 0 {
162+
// Range is MAX+1 (unrepresentable), so we need a special case
164163
return rng.gen();
165164
}
166165

@@ -186,21 +185,188 @@ macro_rules! uniform_int_impl {
186185
}
187186
}
188187
}
188+
189+
impl UniformInt<$ty> {
190+
/// Sample single inclusive, using ONeill's method
191+
#[inline]
192+
pub fn sample_single_inclusive_oneill<R: Rng + ?Sized, B1, B2>(
193+
low_b: B1, high_b: B2, rng: &mut R,
194+
) -> $ty
195+
where
196+
B1: SampleBorrow<$ty> + Sized,
197+
B2: SampleBorrow<$ty> + Sized,
198+
{
199+
let low = *low_b.borrow();
200+
let high = *high_b.borrow();
201+
assert!(
202+
low <= high,
203+
"UniformSampler::sample_single_inclusive: low > high"
204+
);
205+
let range = high.wrapping_sub(low).wrapping_add(1) as $unsigned as $u_large;
206+
if range == 0 {
207+
// Range is MAX+1 (unrepresentable), so we need a special case
208+
return rng.gen();
209+
}
210+
211+
// we use the "Debiased Int Mult (t-opt, m-opt)" rejection sampling method
212+
// described here https://www.pcg-random.org/posts/bounded-rands.html
213+
// and here https://github.com/imneme/bounded-rands
214+
215+
let (mut hi, mut lo) = rng.gen::<$u_large>().wmul(range);
216+
if lo < range {
217+
let mut threshold = range.wrapping_neg();
218+
// this shortcut works best with large ranges
219+
if threshold >= range {
220+
threshold -= range;
221+
if threshold >= range {
222+
threshold %= range;
223+
}
224+
}
225+
while lo < threshold {
226+
let (new_hi, new_lo) = rng.gen::<$u_large>().wmul(range);
227+
hi = new_hi;
228+
lo = new_lo;
229+
}
230+
}
231+
low.wrapping_add(hi as $ty)
232+
}
233+
234+
/// Sample single inclusive, using Canon's method
235+
#[inline]
236+
pub fn sample_single_inclusive_canon<R: Rng + ?Sized, B1, B2>(
237+
low_b: B1, high_b: B2, rng: &mut R,
238+
) -> $ty
239+
where
240+
B1: SampleBorrow<$ty> + Sized,
241+
B2: SampleBorrow<$ty> + Sized,
242+
{
243+
let low = *low_b.borrow();
244+
let high = *high_b.borrow();
245+
assert!(
246+
low <= high,
247+
"UniformSampler::sample_single_inclusive: low > high"
248+
);
249+
let range = high.wrapping_sub(low).wrapping_add(1) as $unsigned as $u_extra_large;
250+
if range == 0 {
251+
// Range is MAX+1 (unrepresentable), so we need a special case
252+
return rng.gen();
253+
}
254+
255+
// generate a sample using a sensible integer type
256+
let (mut result, lo_order) = rng.gen::<$u_extra_large>().wmul(range);
257+
258+
// if the sample is biased...
259+
if lo_order > range.wrapping_neg() {
260+
// ...generate a new sample with 64 more bits, enough that bias is undetectable
261+
let (new_hi_order, _) =
262+
(rng.gen::<$u_extra_large>()).wmul(range as $u_extra_large);
263+
// and adjust if needed
264+
result += lo_order
265+
.checked_add(new_hi_order as $u_extra_large)
266+
.is_none() as $u_extra_large;
267+
}
268+
269+
low.wrapping_add(result as $ty)
270+
}
271+
272+
/// Sample single inclusive, using Canon's method with Lemire's early-out
273+
#[inline]
274+
pub fn sample_inclusive_canon_lemire<R: Rng + ?Sized, B1, B2>(
275+
low_b: B1, high_b: B2, rng: &mut R,
276+
) -> $ty
277+
where
278+
B1: SampleBorrow<$ty> + Sized,
279+
B2: SampleBorrow<$ty> + Sized,
280+
{
281+
let low = *low_b.borrow();
282+
let high = *high_b.borrow();
283+
assert!(
284+
low <= high,
285+
"UniformSampler::sample_single_inclusive: low > high"
286+
);
287+
let range = high.wrapping_sub(low).wrapping_add(1) as $unsigned as $u_extra_large;
288+
if range == 0 {
289+
// Range is MAX+1 (unrepresentable), so we need a special case
290+
return rng.gen();
291+
}
292+
293+
// generate a sample using a sensible integer type
294+
let (mut result, lo_order) = rng.gen::<$u_extra_large>().wmul(range);
295+
296+
// if the sample is biased... (since range won't be changing we can further
297+
// improve this check with a modulo)
298+
if lo_order < range.wrapping_neg() % range {
299+
// ...generate a new sample with 64 more bits, enough that bias is undetectable
300+
let (new_hi_order, _) =
301+
(rng.gen::<$u_extra_large>()).wmul(range as $u_extra_large);
302+
// and adjust if needed
303+
result += lo_order
304+
.checked_add(new_hi_order as $u_extra_large)
305+
.is_none() as $u_extra_large;
306+
}
307+
308+
low.wrapping_add(result as $ty)
309+
}
310+
311+
/// Sample single inclusive, using the Bitmask method
312+
#[inline]
313+
pub fn sample_single_inclusive_bitmask<R: Rng + ?Sized, B1, B2>(
314+
low_b: B1, high_b: B2, rng: &mut R,
315+
) -> $ty
316+
where
317+
B1: SampleBorrow<$ty> + Sized,
318+
B2: SampleBorrow<$ty> + Sized,
319+
{
320+
let low = *low_b.borrow();
321+
let high = *high_b.borrow();
322+
assert!(
323+
low <= high,
324+
"UniformSampler::sample_single_inclusive: low > high"
325+
);
326+
let mut range = high.wrapping_sub(low).wrapping_add(1) as $unsigned as $u_large;
327+
if range == 0 {
328+
// Range is MAX+1 (unrepresentable), so we need a special case
329+
return rng.gen();
330+
}
331+
332+
// the old impl use a mix of methods for different integer sizes, we only use
333+
// the lz method here for a better comparison.
334+
335+
let mut mask = $u_large::max_value();
336+
range -= 1;
337+
mask >>= (range | 1).leading_zeros();
338+
loop {
339+
let x = rng.gen::<$u_large>() & mask;
340+
if x <= range {
341+
return low.wrapping_add(x as $ty);
342+
}
343+
}
344+
}
345+
}
189346
};
190347
}
191-
192-
uniform_int_impl! { i8, u8, u32 }
193-
uniform_int_impl! { i16, u16, u32 }
194-
uniform_int_impl! { i32, u32, u32 }
195-
uniform_int_impl! { i64, u64, u64 }
196-
uniform_int_impl! { i128, u128, u128 }
197-
uniform_int_impl! { isize, usize, usize }
198-
uniform_int_impl! { u8, u8, u32 }
199-
uniform_int_impl! { u16, u16, u32 }
200-
uniform_int_impl! { u32, u32, u32 }
201-
uniform_int_impl! { u64, u64, u64 }
202-
uniform_int_impl! { usize, usize, usize }
203-
uniform_int_impl! { u128, u128, u128 }
348+
uniform_int_impl! { i8, u8, u32, u64 }
349+
uniform_int_impl! { i16, u16, u32, u64 }
350+
uniform_int_impl! { i32, u32, u32, u64 }
351+
uniform_int_impl! { i64, u64, u64, u64 }
352+
uniform_int_impl! { i128, u128, u128, u128 }
353+
uniform_int_impl! { u8, u8, u32, u64 }
354+
uniform_int_impl! { u16, u16, u32, u64 }
355+
uniform_int_impl! { u32, u32, u32, u64 }
356+
uniform_int_impl! { u64, u64, u64, u64 }
357+
uniform_int_impl! { u128, u128, u128, u128 }
358+
#[cfg(any(target_pointer_width = "16", target_pointer_width = "32",))]
359+
mod isize_int_impls {
360+
use super::*;
361+
uniform_int_impl! { isize, usize, usize, u64 }
362+
uniform_int_impl! { usize, usize, usize, u64 }
363+
}
364+
#[cfg(not(any(target_pointer_width = "16", target_pointer_width = "32",)))]
365+
mod isize_int_impls {
366+
use super::*;
367+
uniform_int_impl! { isize, usize, usize, usize }
368+
uniform_int_impl! { usize, usize, usize, usize }
369+
}
204370

205371
#[cfg(test)]
206372
mod tests {

0 commit comments

Comments
 (0)