libcxx/include/__random/uniform_int_distribution.h

   1 //===----------------------------------------------------------------------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8
   9 #ifndef _LIBCPP___RANDOM_UNIFORM_INT_DISTRIBUTION_H
  10 #define _LIBCPP___RANDOM_UNIFORM_INT_DISTRIBUTION_H
  11
  12 #include <__bits>
  13 #include <__config>
  14 #include <__random/log2.h>
  15 #include <bit>
  16 #include <cstddef>
  17 #include <cstdint>
  18 #include <iosfwd>
  19 #include <limits>
  20 #include <type_traits>
  21
  22 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
  23 #pragma GCC system_header
  24 #endif
  25
  26 _LIBCPP_PUSH_MACROS
  27 #include <__undef_macros>
  28
  29 _LIBCPP_BEGIN_NAMESPACE_STD
  30
  31 template<class _Engine, class _UIntType>
  32 class __independent_bits_engine
  33 {
  34 public:
  35     // types
  36     typedef _UIntType result_type;
  37
  38 private:
  39     typedef typename _Engine::result_type _Engine_result_type;
  40     typedef typename conditional
  41         <
  42             sizeof(_Engine_result_type) <= sizeof(result_type),
  43                 result_type,
  44                 _Engine_result_type
  45         >::type _Working_result_type;
  46
  47     _Engine& __e_;
  48     size_t __w_;
  49     size_t __w0_;
  50     size_t __n_;
  51     size_t __n0_;
  52     _Working_result_type __y0_;
  53     _Working_result_type __y1_;
  54     _Engine_result_type __mask0_;
  55     _Engine_result_type __mask1_;
  56
  57 #ifdef _LIBCPP_CXX03_LANG
  58     static const _Working_result_type _Rp = _Engine::_Max - _Engine::_Min
  59                                           + _Working_result_type(1);
  60 #else
  61     static _LIBCPP_CONSTEXPR const _Working_result_type _Rp = _Engine::max() - _Engine::min()
  62                                                       + _Working_result_type(1);
  63 #endif
  64     static _LIBCPP_CONSTEXPR const size_t __m = __log2<_Working_result_type, _Rp>::value;
  65     static _LIBCPP_CONSTEXPR const size_t _WDt = numeric_limits<_Working_result_type>::digits;
  66     static _LIBCPP_CONSTEXPR const size_t _EDt = numeric_limits<_Engine_result_type>::digits;
  67
  68 public:
  69     // constructors and seeding functions
  70     __independent_bits_engine(_Engine& __e, size_t __w);
  71
  72     // generating functions
  73     result_type operator()() {return __eval(integral_constant<bool, _Rp != 0>());}
  74
  75 private:
  76     result_type __eval(false_type);
  77     result_type __eval(true_type);
  78 };
  79
  80 template<class _Engine, class _UIntType>
  81 __independent_bits_engine<_Engine, _UIntType>
  82     ::__independent_bits_engine(_Engine& __e, size_t __w)
  83         : __e_(__e),
  84           __w_(__w)
  85 {
  86     __n_ = __w_ / __m + (__w_ % __m != 0);
  87     __w0_ = __w_ / __n_;
  88     if (_Rp == 0)
  89         __y0_ = _Rp;
  90     else if (__w0_ < _WDt)
  91         __y0_ = (_Rp >> __w0_) << __w0_;
  92     else
  93         __y0_ = 0;
  94     if (_Rp - __y0_ > __y0_ / __n_)
  95     {
  96         ++__n_;
  97         __w0_ = __w_ / __n_;
  98         if (__w0_ < _WDt)
  99             __y0_ = (_Rp >> __w0_) << __w0_;
 100         else
 101             __y0_ = 0;
 102     }
 103     __n0_ = __n_ - __w_ % __n_;
 104     if (__w0_ < _WDt - 1)
 105         __y1_ = (_Rp >> (__w0_ + 1)) << (__w0_ + 1);
 106     else
 107         __y1_ = 0;
 108     __mask0_ = __w0_ > 0 ? _Engine_result_type(~0) >> (_EDt - __w0_) :
 109                           _Engine_result_type(0);
 110     __mask1_ = __w0_ < _EDt - 1 ?
 111                                _Engine_result_type(~0) >> (_EDt - (__w0_ + 1)) :
 112                                _Engine_result_type(~0);
 113 }
 114
 115 template<class _Engine, class _UIntType>
 116 inline
 117 _UIntType
 118 __independent_bits_engine<_Engine, _UIntType>::__eval(false_type)
 119 {
 120     return static_cast<result_type>(__e_() & __mask0_);
 121 }
 122
 123 template<class _Engine, class _UIntType>
 124 _UIntType
 125 __independent_bits_engine<_Engine, _UIntType>::__eval(true_type)
 126 {
 127     const size_t _WRt = numeric_limits<result_type>::digits;
 128     result_type _Sp = 0;
 129     for (size_t __k = 0; __k < __n0_; ++__k)
 130     {
 131         _Engine_result_type __u;
 132         do
 133         {
 134             __u = __e_() - _Engine::min();
 135         } while (__u >= __y0_);
 136         if (__w0_ < _WRt)
 137             _Sp <<= __w0_;
 138         else
 139             _Sp = 0;
 140         _Sp += __u & __mask0_;
 141     }
 142     for (size_t __k = __n0_; __k < __n_; ++__k)
 143     {
 144         _Engine_result_type __u;
 145         do
 146         {
 147             __u = __e_() - _Engine::min();
 148         } while (__u >= __y1_);
 149         if (__w0_ < _WRt - 1)
 150             _Sp <<= __w0_ + 1;
 151         else
 152             _Sp = 0;
 153         _Sp += __u & __mask1_;
 154     }
 155     return _Sp;
 156 }
 157
 158 template<class _IntType = int> // __int128_t is also supported as an extension here
 159 class uniform_int_distribution
 160 {
 161 public:
 162     // types
 163     typedef _IntType result_type;
 164
 165     class param_type
 166     {
 167         result_type __a_;
 168         result_type __b_;
 169     public:
 170         typedef uniform_int_distribution distribution_type;
 171
 172         explicit param_type(result_type __a = 0,
 173                             result_type __b = numeric_limits<result_type>::max())
 174             : __a_(__a), __b_(__b) {}
 175
 176         result_type a() const {return __a_;}
 177         result_type b() const {return __b_;}
 178
 179         friend bool operator==(const param_type& __x, const param_type& __y)
 180             {return __x.__a_ == __y.__a_ && __x.__b_ == __y.__b_;}
 181         friend bool operator!=(const param_type& __x, const param_type& __y)
 182             {return !(__x == __y);}
 183     };
 184
 185 private:
 186     param_type __p_;
 187
 188 public:
 189     // constructors and reset functions
 190 #ifndef _LIBCPP_CXX03_LANG
 191     uniform_int_distribution() : uniform_int_distribution(0) {}
 192     explicit uniform_int_distribution(
 193         result_type __a, result_type __b = numeric_limits<result_type>::max())
 194         : __p_(param_type(__a, __b)) {}
 195 #else
 196     explicit uniform_int_distribution(
 197         result_type __a = 0,
 198         result_type __b = numeric_limits<result_type>::max())
 199         : __p_(param_type(__a, __b)) {}
 200 #endif
 201     explicit uniform_int_distribution(const param_type& __p) : __p_(__p) {}
 202     void reset() {}
 203
 204     // generating functions
 205     template<class _URNG> result_type operator()(_URNG& __g)
 206         {return (*this)(__g, __p_);}
 207     template<class _URNG> result_type operator()(_URNG& __g, const param_type& __p);
 208
 209     // property functions
 210     result_type a() const {return __p_.a();}
 211     result_type b() const {return __p_.b();}
 212
 213     param_type param() const {return __p_;}
 214     void param(const param_type& __p) {__p_ = __p;}
 215
 216     result_type min() const {return a();}
 217     result_type max() const {return b();}
 218
 219     friend bool operator==(const uniform_int_distribution& __x,
 220                            const uniform_int_distribution& __y)
 221         {return __x.__p_ == __y.__p_;}
 222     friend bool operator!=(const uniform_int_distribution& __x,
 223                            const uniform_int_distribution& __y)
 224             {return !(__x == __y);}
 225 };
 226
 227 template<class _IntType>
 228 template<class _URNG>
 229 typename uniform_int_distribution<_IntType>::result_type
 230 uniform_int_distribution<_IntType>::operator()(_URNG& __g, const param_type& __p)
 231 _LIBCPP_DISABLE_UBSAN_UNSIGNED_INTEGER_CHECK
 232 {
 233     typedef typename conditional<sizeof(result_type) <= sizeof(uint32_t), uint32_t,
 234                                  typename make_unsigned<result_type>::type>::type _UIntType;
 235     const _UIntType _Rp = _UIntType(__p.b()) - _UIntType(__p.a()) + _UIntType(1);
 236     if (_Rp == 1)
 237         return __p.a();
 238     const size_t _Dt = numeric_limits<_UIntType>::digits;
 239     typedef __independent_bits_engine<_URNG, _UIntType> _Eng;
 240     if (_Rp == 0)
 241         return static_cast<result_type>(_Eng(__g, _Dt)());
 242     size_t __w = _Dt - __countl_zero(_Rp) - 1;
 243     if ((_Rp & (numeric_limits<_UIntType>::max() >> (_Dt - __w))) != 0)
 244         ++__w;
 245     _Eng __e(__g, __w);
 246     _UIntType __u;
 247     do
 248     {
 249         __u = __e();
 250     } while (__u >= _Rp);
 251     return static_cast<result_type>(__u + __p.a());
 252 }
 253
 254 template <class _CharT, class _Traits, class _IT>
 255 basic_ostream<_CharT, _Traits>&
 256 operator<<(basic_ostream<_CharT, _Traits>& __os,
 257            const uniform_int_distribution<_IT>& __x)
 258 {
 259     __save_flags<_CharT, _Traits> __lx(__os);
 260     typedef basic_ostream<_CharT, _Traits> _Ostream;
 261     __os.flags(_Ostream::dec | _Ostream::left);
 262     _CharT __sp = __os.widen(' ');
 263     __os.fill(__sp);
 264     return __os << __x.a() << __sp << __x.b();
 265 }
 266
 267 template <class _CharT, class _Traits, class _IT>
 268 basic_istream<_CharT, _Traits>&
 269 operator>>(basic_istream<_CharT, _Traits>& __is,
 270            uniform_int_distribution<_IT>& __x)
 271 {
 272     typedef uniform_int_distribution<_IT> _Eng;
 273     typedef typename _Eng::result_type result_type;
 274     typedef typename _Eng::param_type param_type;
 275     __save_flags<_CharT, _Traits> __lx(__is);
 276     typedef basic_istream<_CharT, _Traits> _Istream;
 277     __is.flags(_Istream::dec | _Istream::skipws);
 278     result_type __a;
 279     result_type __b;
 280     __is >> __a >> __b;
 281     if (!__is.fail())
 282         __x.param(param_type(__a, __b));
 283     return __is;
 284 }
 285
 286 _LIBCPP_END_NAMESPACE_STD
 287
 288 _LIBCPP_POP_MACROS
 289
 290 #endif // _LIBCPP___RANDOM_UNIFORM_INT_DISTRIBUTION_H