From 14e7db606e26103f9cf1f90d9c2facf76293d4f4 Mon Sep 17 00:00:00 2001
From: Brendan Hickey <brendan@bhickey.net>
Date: Wed, 24 Aug 2016 18:35:26 -0400
Subject: [PATCH] Speed up range sampling.

This change removes division from the rejection sampler in random_range().
Replace divisions in range sampling with bitshifting. Instead of finding a
well fitting range, we generate a number [0,2^n) and reject out of range
values.

In synthetic benchmarks, this approximately doubles the throughput.
---
 benches/distributions/mod.rs   |  1 +
 benches/distributions/range.rs | 18 ++++++++++++++++++
 src/distributions/range.rs     | 27 +++++++++------------------
 3 files changed, 28 insertions(+), 18 deletions(-)
 create mode 100644 benches/distributions/range.rs

diff --git a/benches/distributions/mod.rs b/benches/distributions/mod.rs
index 49f6bd9c06c..e88d51c3901 100644
--- a/benches/distributions/mod.rs
+++ b/benches/distributions/mod.rs
@@ -1,3 +1,4 @@
 mod exponential;
 mod normal;
 mod gamma;
+mod range;
diff --git a/benches/distributions/range.rs b/benches/distributions/range.rs
new file mode 100644
index 00000000000..e03c7f4acf8
--- /dev/null
+++ b/benches/distributions/range.rs
@@ -0,0 +1,18 @@
+use std::mem::size_of;
+use std::cmp;
+use test::Bencher;
+use rand;
+use rand::distributions::Sample;
+use rand::distributions::Range;
+
+#[bench]
+fn rand_range(b: &mut Bencher) {
+    let mut rng = rand::weak_rng();
+    let mut range = Range::new(10, 10000);
+    b.iter(|| {
+        for _ in 0..::RAND_BENCH_N {
+            range.sample(&mut rng);
+        }
+    });
+    b.bytes = size_of::<u32>() as u64 * ::RAND_BENCH_N;
+}
diff --git a/src/distributions/range.rs b/src/distributions/range.rs
index 3cf47be0207..d4dc7812a02 100644
--- a/src/distributions/range.rs
+++ b/src/distributions/range.rs
@@ -50,7 +50,7 @@ use distributions::{Sample, IndependentSample};
 pub struct Range<X> {
     low: X,
     range: X,
-    accept_zone: X
+    lz: X,
 }
 
 impl<X: SampleRange + PartialOrd> Range<X> {
@@ -98,31 +98,22 @@ macro_rules! integer_impl {
 
             fn construct_range(low: $ty, high: $ty) -> Range<$ty> {
                 let range = (w(high as $unsigned) - w(low as $unsigned)).0;
-                let unsigned_max: $unsigned = ::std::$unsigned::MAX;
-
-                // this is the largest number that fits into $unsigned
-                // that `range` divides evenly, so, if we've sampled
-                // `n` uniformly from this region, then `n % range` is
-                // uniform in [0, range)
-                let zone = unsigned_max - unsigned_max % range;
 
                 Range {
                     low: low,
                     range: range as $ty,
-                    accept_zone: zone as $ty
+                    lz: range.leading_zeros() as $ty,
                 }
             }
             #[inline]
             fn sample_range<R: Rng>(r: &Range<$ty>, rng: &mut R) -> $ty {
                 loop {
-                    // rejection sample
-                    let v = rng.gen::<$unsigned>();
-                    // until we find something that fits into the
-                    // region which r.range evenly divides (this will
-                    // be uniformly distributed)
-                    if v < r.accept_zone as $unsigned {
-                        // and return it, with some adjustments
-                        return (w(r.low) + w((v % r.range as $unsigned) as $ty)).0;
+                    // rejection sample. Slough off the low bits so the maximum
+                    // generated value is the next largest power of 2.
+                    let v = rng.gen::<$unsigned>() >> r.lz;
+                    if v < r.range as $unsigned {
+                        // offset the value onto our range.
+                        return (w(r.low) + w((v as $unsigned) as $ty)).0;
                     }
                 }
             }
@@ -148,7 +139,7 @@ macro_rules! float_impl {
                 Range {
                     low: low,
                     range: high - low,
-                    accept_zone: 0.0 // unused
+                    lz: 0.0,
                 }
             }
             fn sample_range<R: Rng>(r: &Range<$ty>, rng: &mut R) -> $ty {