rand/distr/
integer.rs

1// Copyright 2018 Developers of the Rand project.
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8
9//! The implementations of the `StandardUniform` distribution for integer types.
10
11use crate::distr::{Distribution, StandardUniform};
12use crate::Rng;
13#[cfg(all(target_arch = "x86", feature = "simd_support"))]
14use core::arch::x86::__m512i;
15#[cfg(target_arch = "x86")]
16use core::arch::x86::{__m128i, __m256i};
17#[cfg(all(target_arch = "x86_64", feature = "simd_support"))]
18use core::arch::x86_64::__m512i;
19#[cfg(target_arch = "x86_64")]
20use core::arch::x86_64::{__m128i, __m256i};
21use core::num::{
22    NonZeroI128, NonZeroI16, NonZeroI32, NonZeroI64, NonZeroI8, NonZeroU128, NonZeroU16,
23    NonZeroU32, NonZeroU64, NonZeroU8,
24};
25#[cfg(feature = "simd_support")]
26use core::simd::*;
27
28impl Distribution<u8> for StandardUniform {
29    #[inline]
30    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u8 {
31        rng.next_u32() as u8
32    }
33}
34
35impl Distribution<u16> for StandardUniform {
36    #[inline]
37    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u16 {
38        rng.next_u32() as u16
39    }
40}
41
42impl Distribution<u32> for StandardUniform {
43    #[inline]
44    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u32 {
45        rng.next_u32()
46    }
47}
48
49impl Distribution<u64> for StandardUniform {
50    #[inline]
51    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 {
52        rng.next_u64()
53    }
54}
55
56impl Distribution<u128> for StandardUniform {
57    #[inline]
58    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u128 {
59        // Use LE; we explicitly generate one value before the next.
60        let x = u128::from(rng.next_u64());
61        let y = u128::from(rng.next_u64());
62        (y << 64) | x
63    }
64}
65
66macro_rules! impl_int_from_uint {
67    ($ty:ty, $uty:ty) => {
68        impl Distribution<$ty> for StandardUniform {
69            #[inline]
70            fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> $ty {
71                rng.random::<$uty>() as $ty
72            }
73        }
74    };
75}
76
77impl_int_from_uint! { i8, u8 }
78impl_int_from_uint! { i16, u16 }
79impl_int_from_uint! { i32, u32 }
80impl_int_from_uint! { i64, u64 }
81impl_int_from_uint! { i128, u128 }
82
83macro_rules! impl_nzint {
84    ($ty:ty, $new:path) => {
85        impl Distribution<$ty> for StandardUniform {
86            fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> $ty {
87                loop {
88                    if let Some(nz) = $new(rng.random()) {
89                        break nz;
90                    }
91                }
92            }
93        }
94    };
95}
96
97impl_nzint!(NonZeroU8, NonZeroU8::new);
98impl_nzint!(NonZeroU16, NonZeroU16::new);
99impl_nzint!(NonZeroU32, NonZeroU32::new);
100impl_nzint!(NonZeroU64, NonZeroU64::new);
101impl_nzint!(NonZeroU128, NonZeroU128::new);
102
103impl_nzint!(NonZeroI8, NonZeroI8::new);
104impl_nzint!(NonZeroI16, NonZeroI16::new);
105impl_nzint!(NonZeroI32, NonZeroI32::new);
106impl_nzint!(NonZeroI64, NonZeroI64::new);
107impl_nzint!(NonZeroI128, NonZeroI128::new);
108
109#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
110impl Distribution<__m128i> for StandardUniform {
111    #[inline]
112    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> __m128i {
113        // NOTE: It's tempting to use the u128 impl here, but confusingly this
114        // results in different code (return via rdx, r10 instead of rax, rdx
115        // with u128 impl) and is much slower (+130 time). This version calls
116        // impls::fill_bytes_via_next but performs well.
117
118        let mut buf = [0_u8; core::mem::size_of::<__m128i>()];
119        rng.fill_bytes(&mut buf);
120        // x86 is little endian so no need for conversion
121
122        // SAFETY: All byte sequences of `buf` represent values of the output type.
123        unsafe { core::mem::transmute(buf) }
124    }
125}
126
127#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
128impl Distribution<__m256i> for StandardUniform {
129    #[inline]
130    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> __m256i {
131        let mut buf = [0_u8; core::mem::size_of::<__m256i>()];
132        rng.fill_bytes(&mut buf);
133        // x86 is little endian so no need for conversion
134
135        // SAFETY: All byte sequences of `buf` represent values of the output type.
136        unsafe { core::mem::transmute(buf) }
137    }
138}
139
140#[cfg(all(
141    any(target_arch = "x86", target_arch = "x86_64"),
142    feature = "simd_support"
143))]
144impl Distribution<__m512i> for StandardUniform {
145    #[inline]
146    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> __m512i {
147        let mut buf = [0_u8; core::mem::size_of::<__m512i>()];
148        rng.fill_bytes(&mut buf);
149        // x86 is little endian so no need for conversion
150
151        // SAFETY: All byte sequences of `buf` represent values of the output type.
152        unsafe { core::mem::transmute(buf) }
153    }
154}
155
156#[cfg(feature = "simd_support")]
157macro_rules! simd_impl {
158    ($($ty:ty),+) => {$(
159        /// Requires nightly Rust and the [`simd_support`] feature
160        ///
161        /// [`simd_support`]: https://github.com/rust-random/rand#crate-features
162        #[cfg(feature = "simd_support")]
163        impl<const LANES: usize> Distribution<Simd<$ty, LANES>> for StandardUniform
164        where
165            LaneCount<LANES>: SupportedLaneCount,
166        {
167            #[inline]
168            fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Simd<$ty, LANES> {
169                let mut vec = Simd::default();
170                rng.fill(vec.as_mut_array().as_mut_slice());
171                vec
172            }
173        }
174    )+};
175}
176
177#[cfg(feature = "simd_support")]
178simd_impl!(u8, i8, u16, i16, u32, i32, u64, i64);
179
180#[cfg(test)]
181mod tests {
182    use super::*;
183
184    #[test]
185    fn test_integers() {
186        let mut rng = crate::test::rng(806);
187
188        rng.sample::<i8, _>(StandardUniform);
189        rng.sample::<i16, _>(StandardUniform);
190        rng.sample::<i32, _>(StandardUniform);
191        rng.sample::<i64, _>(StandardUniform);
192        rng.sample::<i128, _>(StandardUniform);
193
194        rng.sample::<u8, _>(StandardUniform);
195        rng.sample::<u16, _>(StandardUniform);
196        rng.sample::<u32, _>(StandardUniform);
197        rng.sample::<u64, _>(StandardUniform);
198        rng.sample::<u128, _>(StandardUniform);
199    }
200
201    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
202    #[test]
203    fn x86_integers() {
204        let mut rng = crate::test::rng(807);
205
206        rng.sample::<__m128i, _>(StandardUniform);
207        rng.sample::<__m256i, _>(StandardUniform);
208        #[cfg(feature = "simd_support")]
209        rng.sample::<__m512i, _>(StandardUniform);
210    }
211
212    #[test]
213    fn value_stability() {
214        fn test_samples<T: Copy + core::fmt::Debug + PartialEq>(zero: T, expected: &[T])
215        where
216            StandardUniform: Distribution<T>,
217        {
218            let mut rng = crate::test::rng(807);
219            let mut buf = [zero; 3];
220            for x in &mut buf {
221                *x = rng.sample(StandardUniform);
222            }
223            assert_eq!(&buf, expected);
224        }
225
226        test_samples(0u8, &[9, 247, 111]);
227        test_samples(0u16, &[32265, 42999, 38255]);
228        test_samples(0u32, &[2220326409, 2575017975, 2018088303]);
229        test_samples(
230            0u64,
231            &[
232                11059617991457472009,
233                16096616328739788143,
234                1487364411147516184,
235            ],
236        );
237        test_samples(
238            0u128,
239            &[
240                296930161868957086625409848350820761097,
241                145644820879247630242265036535529306392,
242                111087889832015897993126088499035356354,
243            ],
244        );
245
246        test_samples(0i8, &[9, -9, 111]);
247        // Skip further i* types: they are simple reinterpretation of u* samples
248
249        #[cfg(feature = "simd_support")]
250        {
251            // We only test a sub-set of types here and make assumptions about the rest.
252
253            test_samples(
254                u8x4::default(),
255                &[
256                    u8x4::from([9, 126, 87, 132]),
257                    u8x4::from([247, 167, 123, 153]),
258                    u8x4::from([111, 149, 73, 120]),
259                ],
260            );
261            test_samples(
262                u8x8::default(),
263                &[
264                    u8x8::from([9, 126, 87, 132, 247, 167, 123, 153]),
265                    u8x8::from([111, 149, 73, 120, 68, 171, 98, 223]),
266                    u8x8::from([24, 121, 1, 50, 13, 46, 164, 20]),
267                ],
268            );
269
270            test_samples(
271                i64x8::default(),
272                &[
273                    i64x8::from([
274                        -7387126082252079607,
275                        -2350127744969763473,
276                        1487364411147516184,
277                        7895421560427121838,
278                        602190064936008898,
279                        6022086574635100741,
280                        -5080089175222015595,
281                        -4066367846667249123,
282                    ]),
283                    i64x8::from([
284                        9180885022207963908,
285                        3095981199532211089,
286                        6586075293021332726,
287                        419343203796414657,
288                        3186951873057035255,
289                        5287129228749947252,
290                        444726432079249540,
291                        -1587028029513790706,
292                    ]),
293                    i64x8::from([
294                        6075236523189346388,
295                        1351763722368165432,
296                        -6192309979959753740,
297                        -7697775502176768592,
298                        -4482022114172078123,
299                        7522501477800909500,
300                        -1837258847956201231,
301                        -586926753024886735,
302                    ]),
303                ],
304            );
305        }
306    }
307}