random.h (2802B)
1 // Copyright (c) the JPEG XL Project Authors. All rights reserved. 2 // 3 // Use of this source code is governed by a BSD-style 4 // license that can be found in the LICENSE file. 5 6 #ifndef LIB_JXL_BASE_RANDOM_ 7 #define LIB_JXL_BASE_RANDOM_ 8 9 // Random number generator + distributions. 10 // We don't use <random> because the implementation (and thus results) differs 11 // between libstdc++ and libc++. 12 13 #include <stdint.h> 14 #include <string.h> 15 16 #include <algorithm> 17 #include <cmath> 18 19 #include "lib/jxl/base/status.h" 20 21 namespace jxl { 22 struct Rng { 23 explicit Rng(uint64_t seed) 24 : s{static_cast<uint64_t>(0x94D049BB133111EBull), 25 static_cast<uint64_t>(0xBF58476D1CE4E5B9ull) + seed} {} 26 27 // Xorshift128+ adapted from xorshift128+-inl.h 28 uint64_t operator()() { 29 uint64_t s1 = s[0]; 30 const uint64_t s0 = s[1]; 31 const uint64_t bits = s1 + s0; // b, c 32 s[0] = s0; 33 s1 ^= s1 << 23; 34 s1 ^= s0 ^ (s1 >> 18) ^ (s0 >> 5); 35 s[1] = s1; 36 return bits; 37 } 38 39 // Uniformly distributed int64_t in [begin, end), under the assumption that 40 // `end-begin` is significantly smaller than 1<<64, otherwise there is some 41 // bias. 42 int64_t UniformI(int64_t begin, int64_t end) { 43 JXL_DASSERT(end > begin); 44 return static_cast<int64_t>((*this)() % 45 static_cast<uint64_t>(end - begin)) + 46 begin; 47 } 48 49 // Same as UniformI, but for uint64_t. 50 uint64_t UniformU(uint64_t begin, uint64_t end) { 51 JXL_DASSERT(end > begin); 52 return (*this)() % (end - begin) + begin; 53 } 54 55 // Uniformly distributed float in [begin, end) range. Note: only 23 bits of 56 // randomness. 57 float UniformF(float begin, float end) { 58 float f; 59 // Bits of a random [1, 2) float. 60 uint32_t u = ((*this)() >> (64 - 23)) | 0x3F800000; 61 static_assert(sizeof(f) == sizeof(u), 62 "Float and U32 must have the same size"); 63 memcpy(&f, &u, sizeof(f)); 64 // Note: (end-begin) * f + (2*begin-end) may fail to return a number >= 65 // begin. 66 return (end - begin) * (f - 1.0f) + begin; 67 } 68 69 // Bernoulli trial 70 bool Bernoulli(float p) { return UniformF(0, 1) < p; } 71 72 // State for geometric distributions. 73 // The stored value is inv_log_1mp 74 using GeometricDistribution = float; 75 static GeometricDistribution MakeGeometric(float p) { 76 return 1.0 / std::log(1 - p); 77 } 78 79 uint32_t Geometric(const GeometricDistribution& dist) { 80 float f = UniformF(0, 1); 81 float inv_log_1mp = dist; 82 float log = std::log(1 - f) * inv_log_1mp; 83 return static_cast<uint32_t>(log); 84 } 85 86 template <typename T> 87 void Shuffle(T* t, size_t n) { 88 for (size_t i = 0; i + 1 < n; i++) { 89 size_t a = UniformU(i, n); 90 std::swap(t[a], t[i]); 91 } 92 } 93 94 private: 95 uint64_t s[2]; 96 }; 97 98 } // namespace jxl 99 #endif // LIB_JXL_BASE_RANDOM_