6846e1d177feb8c4d5dfbf8167da3dba00174007
[folly.git] / folly / Random.h
1 /*
2  * Copyright 2014 Facebook, Inc.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *   http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16
17 #ifndef FOLLY_RANDOM_H_
18 #define FOLLY_RANDOM_H_
19
20 #include <type_traits>
21 #include <random>
22 #include <stdint.h>
23 #include <folly/ThreadLocal.h>
24
25 #if __GNUC_PREREQ(4, 8) && !defined(ANDROID)
26 #include <ext/random>
27 #define FOLLY_USE_SIMD_PRNG 1
28 #endif
29
30 namespace folly {
31
32 /**
33  * A PRNG with one instance per thread. This PRNG uses a mersenne twister random
34  * number generator and is seeded from /dev/urandom. It should not be used for
35  * anything which requires security, only for statistical randomness.
36  *
37  * An instance of this class represents the current threads PRNG. This means
38  * copying an instance of this class across threads will result in corruption
39  *
40  * Most users will use the Random class which implicitly creates this class.
41  * However, if you are worried about performance, you can memoize the TLS
42  * lookups that get the per thread state by manually using this class:
43  *
44  * ThreadLocalPRNG rng = Random::threadLocalPRNG()
45  * for (...) {
46  *   Random::rand32(rng);
47  * }
48  */
49 class ThreadLocalPRNG {
50  public:
51   typedef uint32_t result_type;
52
53   uint32_t operator()() {
54     // Using a static method allows the compiler to avoid allocating stack space
55     // for this class.
56     return getImpl(local_);
57   }
58
59   static constexpr result_type min() {
60     return std::numeric_limits<result_type>::min();
61   }
62   static constexpr result_type max() {
63     return std::numeric_limits<result_type>::max();
64   }
65   friend class Random;
66
67   ThreadLocalPRNG();
68
69  private:
70   class LocalInstancePRNG;
71
72   static result_type getImpl(LocalInstancePRNG* local);
73   LocalInstancePRNG* local_;
74 };
75
76
77 class Random {
78
79  private:
80   template<class RNG>
81   using ValidRNG = typename std::enable_if<
82    std::is_unsigned<typename std::result_of<RNG&()>::type>::value,
83    RNG>::type;
84
85  public:
86   // Default generator type.
87 #if FOLLY_USE_SIMD_PRNG
88   typedef __gnu_cxx::sfmt19937 DefaultGenerator;
89 #else
90   typedef std::mt19937 DefaultGenerator;
91 #endif
92
93   /**
94    * Get secure random bytes. (On Linux and OSX, this means /dev/urandom).
95    */
96   static void secureRandom(void* data, size_t len);
97
98   /**
99    * Shortcut to get a secure random value of integral type.
100    */
101   template <class T>
102   static typename std::enable_if<
103     std::is_integral<T>::value && !std::is_same<T,bool>::value,
104     T>::type
105   secureRandom() {
106     T val;
107     secureRandom(&val, sizeof(val));
108     return val;
109   }
110
111   /**
112    * (Re-)Seed an existing RNG with a good seed.
113    *
114    * Note that you should usually use ThreadLocalPRNG unless you need
115    * reproducibility (such as during a test), in which case you'd want
116    * to create a RNG with a good seed in production, and seed it yourself
117    * in test.
118    */
119   template <class RNG = DefaultGenerator>
120   static void seed(ValidRNG<RNG>& rng);
121
122   /**
123    * Create a new RNG, seeded with a good seed.
124    *
125    * Note that you should usually use ThreadLocalPRNG unless you need
126    * reproducibility (such as during a test), in which case you'd want
127    * to create a RNG with a good seed in production, and seed it yourself
128    * in test.
129    */
130   template <class RNG = DefaultGenerator>
131   static ValidRNG<RNG> create();
132
133   /**
134    * Returns a random uint32_t
135    */
136   template<class RNG = ThreadLocalPRNG>
137   static uint32_t rand32(ValidRNG<RNG> rng = RNG()) {
138     uint32_t r = rng.operator()();
139     return r;
140   }
141
142   /**
143    * Returns a random uint32_t in [0, max). If max == 0, returns 0.
144    */
145   template<class RNG = ThreadLocalPRNG>
146   static uint32_t rand32(uint32_t max, ValidRNG<RNG> rng = RNG()) {
147     if (max == 0) {
148       return 0;
149     }
150
151     return std::uniform_int_distribution<uint32_t>(0, max - 1)(rng);
152   }
153
154   /**
155    * Returns a random uint32_t in [min, max). If min == max, returns 0.
156    */
157   template<class RNG = ThreadLocalPRNG>
158   static uint32_t rand32(uint32_t min,
159                          uint32_t max,
160                          ValidRNG<RNG> rng = RNG()) {
161     if (min == max) {
162       return 0;
163     }
164
165     return std::uniform_int_distribution<uint32_t>(min, max - 1)(rng);
166   }
167
168   /**
169    * Returns a random uint64_t
170    */
171   template<class RNG = ThreadLocalPRNG>
172   static uint64_t rand64(ValidRNG<RNG> rng = RNG()) {
173     return ((uint64_t) rng() << 32) | rng();
174   }
175
176   /**
177    * Returns a random uint64_t in [0, max). If max == 0, returns 0.
178    */
179   template<class RNG = ThreadLocalPRNG>
180   static uint64_t rand64(uint64_t max, ValidRNG<RNG> rng = RNG()) {
181     if (max == 0) {
182       return 0;
183     }
184
185     return std::uniform_int_distribution<uint64_t>(0, max - 1)(rng);
186   }
187
188   /**
189    * Returns a random uint64_t in [min, max). If min == max, returns 0.
190    */
191   template<class RNG = ThreadLocalPRNG>
192   static uint64_t rand64(uint64_t min,
193                          uint64_t max,
194                          ValidRNG<RNG> rng = RNG()) {
195     if (min == max) {
196       return 0;
197     }
198
199     return std::uniform_int_distribution<uint64_t>(min, max - 1)(rng);
200   }
201
202   /**
203    * Returns true 1/n of the time. If n == 0, always returns false
204    */
205   template<class RNG = ThreadLocalPRNG>
206   static bool oneIn(uint32_t n, ValidRNG<RNG> rng = RNG()) {
207     if (n == 0) {
208       return false;
209     }
210
211     return rand32(n, rng) == 0;
212   }
213
214   /**
215    * Returns a double in [0, 1)
216    */
217   template<class RNG = ThreadLocalPRNG>
218   static double randDouble01(ValidRNG<RNG> rng = RNG()) {
219     return std::generate_canonical<double, std::numeric_limits<double>::digits>
220       (rng);
221   }
222
223   /**
224     * Returns a double in [min, max), if min == max, returns 0.
225     */
226   template<class RNG = ThreadLocalPRNG>
227   static double randDouble(double min, double max, ValidRNG<RNG> rng = RNG()) {
228     if (std::fabs(max - min) < std::numeric_limits<double>::epsilon()) {
229       return 0;
230     }
231     return std::uniform_real_distribution<double>(min, max)(rng);
232   }
233
234 };
235
236 /*
237  * Return a good seed for a random number generator.
238  * Note that this is a legacy function, as it returns a 32-bit value, which
239  * is too small to be useful as a "real" RNG seed. Use the functions in class
240  * Random instead.
241  */
242 inline uint32_t randomNumberSeed() {
243   return Random::rand32();
244 }
245
246 }
247
248 #include <folly/Random-inl.h>
249
250 #endif