From 23f0a5139bb91099b0011b0cd1e1972696e959cf Mon Sep 17 00:00:00 2001
From: Diggory Hardy <git@dhardy.name>
Date: Fri, 15 Feb 2019 09:24:22 +0000
Subject: [PATCH] =?UTF-8?q?Implement=20Standard=20support=20for=20NonZero*?=
 =?UTF-8?q?=20types=20(Rustc=20=E2=89=A5=201.28)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 benches/distributions.rs     | 27 +++++++++++++++++++++++++++
 build.rs                     |  1 +
 src/distributions/integer.rs | 23 +++++++++++++++++++++++
 3 files changed, 51 insertions(+)

diff --git a/benches/distributions.rs b/benches/distributions.rs
index 069a82856a5..56a8d43b791 100644
--- a/benches/distributions.rs
+++ b/benches/distributions.rs
@@ -14,6 +14,8 @@ extern crate rand;
 const RAND_BENCH_N: u64 = 1000;
 
 use std::mem::size_of;
+#[cfg(rustc_1_28)]
+use std::num::{NonZeroU8, NonZeroU16, NonZeroU32, NonZeroU64, NonZeroU128};
 use test::Bencher;
 use std::time::Duration;
 
@@ -41,6 +43,26 @@ macro_rules! distr_int {
     }
 }
 
+macro_rules! distr_nz_int {
+    ($fnn:ident, $tynz:ty, $ty:ty, $distr:expr) => {
+        #[bench]
+        fn $fnn(b: &mut Bencher) {
+            let mut rng = SmallRng::from_entropy();
+            let distr = $distr;
+
+            b.iter(|| {
+                let mut accum = 0 as $ty;
+                for _ in 0..::RAND_BENCH_N {
+                    let x: $tynz = distr.sample(&mut rng);
+                    accum = accum.wrapping_add(x.get());
+                }
+                accum
+            });
+            b.bytes = size_of::<$ty>() as u64 * ::RAND_BENCH_N;
+        }
+    }
+}
+
 macro_rules! distr_float {
     ($fnn:ident, $ty:ty, $distr:expr) => {
         #[bench]
@@ -156,6 +178,11 @@ distr_int!(distr_standard_i16, i16, Standard);
 distr_int!(distr_standard_i32, i32, Standard);
 distr_int!(distr_standard_i64, i64, Standard);
 distr_int!(distr_standard_i128, i128, Standard);
+#[cfg(rustc_1_28)] distr_nz_int!(distr_standard_nz8, NonZeroU8, u8, Standard);
+#[cfg(rustc_1_28)] distr_nz_int!(distr_standard_nz16, NonZeroU16, u16, Standard);
+#[cfg(rustc_1_28)] distr_nz_int!(distr_standard_nz32, NonZeroU32, u32, Standard);
+#[cfg(rustc_1_28)] distr_nz_int!(distr_standard_nz64, NonZeroU64, u64, Standard);
+#[cfg(rustc_1_28)] distr_nz_int!(distr_standard_nz128, NonZeroU128, u128, Standard);
 
 distr!(distr_standard_bool, bool, Standard);
 distr!(distr_standard_alphanumeric, char, Alphanumeric);
diff --git a/build.rs b/build.rs
index a554ad9875e..c5085c98041 100644
--- a/build.rs
+++ b/build.rs
@@ -7,4 +7,5 @@ fn main() {
     ac.emit_rustc_version(1, 25);
     ac.emit_rustc_version(1, 26);
     ac.emit_rustc_version(1, 27);
+    ac.emit_rustc_version(1, 28);
 }
diff --git a/src/distributions/integer.rs b/src/distributions/integer.rs
index 7e408dbf1b2..973f8940cac 100644
--- a/src/distributions/integer.rs
+++ b/src/distributions/integer.rs
@@ -10,6 +10,8 @@
 
 use {Rng};
 use distributions::{Distribution, Standard};
+#[cfg(rustc_1_28)]
+use core::num::{NonZeroU8, NonZeroU16, NonZeroU32, NonZeroU64, NonZeroU128, NonZeroUsize};
 #[cfg(feature="simd_support")]
 use packed_simd::*;
 #[cfg(all(target_arch = "x86", feature="nightly"))]
@@ -88,6 +90,27 @@ impl_int_from_uint! { i64, u64 }
 #[cfg(all(rustc_1_26, not(target_os = "emscripten")))] impl_int_from_uint! { i128, u128 }
 impl_int_from_uint! { isize, usize }
 
+macro_rules! impl_nzint {
+    ($ty:ty, $new:path) => {
+        impl Distribution<$ty> for Standard {
+            fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> $ty {
+                loop {
+                    if let Some(nz) = $new(rng.gen()) {
+                        break nz;
+                    }
+                }
+            }
+        }
+    }
+}
+
+#[cfg(rustc_1_28)] impl_nzint!(NonZeroU8, NonZeroU8::new);
+#[cfg(rustc_1_28)] impl_nzint!(NonZeroU16, NonZeroU16::new);
+#[cfg(rustc_1_28)] impl_nzint!(NonZeroU32, NonZeroU32::new);
+#[cfg(rustc_1_28)] impl_nzint!(NonZeroU64, NonZeroU64::new);
+#[cfg(all(rustc_1_28, not(target_os = "emscripten")))] impl_nzint!(NonZeroU128, NonZeroU128::new);
+#[cfg(rustc_1_28)] impl_nzint!(NonZeroUsize, NonZeroUsize::new);
+
 #[cfg(feature="simd_support")]
 macro_rules! simd_impl {
     ($(($intrinsic:ident, $vec:ty),)+) => {$(