threefry.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: BSD-2-Clause
 
// References:
//  Authors:    John K. Salmon, Mark A. Moraes, Ron O. Dror, David E. Shaw
//  Title:      Parallel Random Numbers: As Easy as 1, 2, 3
//  Year:       2011
//  URL:        https://doi.org/10.1145/2063384.2063405
 
#ifndef pRC_CORE_RANDOM_THREEFRY_H
#define pRC_CORE_RANDOM_THREEFRY_H
 
#include <prc/core/basic/functions/bit_rotate_left.hpp>
#include <prc/core/basic/limits.hpp>
#include <prc/core/basic/range.hpp>
#include <prc/core/basic/type_traits.hpp>
#include <prc/core/container/array_stack.hpp>
#include <prc/core/random/seq.hpp>
#include <prc/core/random/type_traits.hpp>
 
namespace pRC
{
    template<Size Rounds, class U>
    class Threefry
    {
    public:
        using result_type = U;
 
    public:
        static constexpr auto min()
        {
            return NumericLimits<U>::min();
        }
        static constexpr auto max()
        {
            return NumericLimits<U>::max();
        }
 
        static constexpr auto samplesPerBlock()
        {
            return 32 / sizeof(U);
        }
 
    public:
        ~Threefry() = default;
        constexpr Threefry(Threefry &&) = default;
        constexpr Threefry &operator=(Threefry const &) & = default;
        constexpr Threefry &operator=(Threefry &&) & = default;
        constexpr Threefry(Threefry const &) = default;
 
        Threefry()
        {
            seed();
        }
 
        template<Size S>
        Threefry(SeedSequence<S> const &seq)
        {
            seed(seq);
        }
 
        constexpr void seed()
        {
            seed(SeedSequence<0>().generate<8>());
        }
 
        template<Size S>
        constexpr void seed(SeedSequence<S> const &seq)
        {
            seed(seq.template generate<8>());
        }
 
        constexpr U operator()()
        {
            ++mChunkCounter;
            mChunkCounter %= samplesPerBlock();
 
            if(mChunkCounter == 0)
            {
                incrementCounter();
                encryptCounter();
            }
 
            return reinterpret_cast<U const *>(mCipher.data())[mChunkCounter];
        }
 
        constexpr void discard(std::uint64_t z)
        {
            if(z < samplesPerBlock() - mChunkCounter)
            {
                mChunkCounter += z;
                return;
            }
 
            z -= samplesPerBlock() - mChunkCounter;
            mChunkCounter = z % samplesPerBlock();
 
            z -= mChunkCounter;
            z /= samplesPerBlock();
 
            incrementCounter(++z);
            encryptCounter();
        }
 
        constexpr void advance(std::uint64_t const z)
        {
            if(z > NumericLimits<std::uint64_t>::max() - mCounter[1])
            {
                ++mCounter[2];
                if(mCounter[2] == 0)
                {
                    ++mCounter[3];
                }
            }
            mCounter[1] += z;
 
            encryptCounter();
        }
 
        constexpr auto operator==(const Threefry &rhs)
        {
            if(mChunkCounter != rhs.mChunkCounter)
            {
                return false;
            }
 
            for(Index i = 0; i < 4U; ++i)
            {
                if(mCounter[i] != rhs.mCounter[i])
                {
                    return false;
                }
                if(mKey[i] != rhs.mKey[i])
                {
                    return false;
                }
                if(mCipher[i] != rhs.mCipher[i])
                {
                    return false;
                }
            }
            return true;
        }
 
        constexpr auto operator!=(const Threefry &rhs)
        {
            return !(*this == rhs);
        }
 
    private:
        constexpr void seed(StackArray<Seed, 8> const &seq)
        {
            for(Index i = 0; i < 4U; ++i)
            {
                mKey[i] = (static_cast<std::uint64_t>(seq[2 * i]) << 32) |
                    seq[2 * i + 1];
            }
            mKey[4] =
                0x1BD11BDAA9FC1A22 ^ mKey[0] ^ mKey[1] ^ mKey[2] ^ mKey[3];
 
            resetCounter();
            encryptCounter();
        }
 
        constexpr void resetCounter()
        {
            for(Index i = 0; i < 4U; ++i)
            {
                mCounter[i] = 0;
            }
            mChunkCounter = 0;
        }
 
        constexpr void incrementCounter()
        {
            ++mCounter[0];
            if(mCounter[0] != 0)
            {
                return;
            }
 
            ++mCounter[1];
            if(mCounter[1] != 0)
            {
                return;
            }
 
            ++mCounter[2];
            if(mCounter[2] != 0)
            {
                return;
            }
 
            ++mCounter[3];
        }
 
        constexpr void incrementCounter(std::uint64_t const z)
        {
            if(z > NumericLimits<std::uint64_t>::max() - mCounter[0])
            {
                ++mCounter[1];
                if(mCounter[1] == 0)
                {
                    ++mCounter[2];
                    if(mCounter[2] == 0)
                    {
                        ++mCounter[3];
                    }
                }
            }
            mCounter[0] += z;
        }
 
        constexpr void encryptCounter()
        {
            auto const mix = [](auto &x, auto &y, std::uint8_t const bits)
            {
                x += y;
                y = bitRotateLeft(y, bits);
                y ^= x;
            };
 
            auto const add = [](auto &cipher, auto const &key,
                                 auto const offset, auto const cycle)
            {
                range<4>(
                    [&cipher, &key, &offset](auto const i)
                    {
                        cipher[i] += key[(offset + i) % 5];
                    });
 
                cipher[3] += cycle;
            };
 
            range<4>(
                [this](auto const i)
                {
                    mCipher[i] = mCounter[i];
                });
 
            range<4>(
                [this](auto const i)
                {
                    mCipher[i] += mKey[i];
                });
 
            Size cycle = 0;
 
            range<Context::CompileTime, Rounds>(
                [this, &mix, &add, &cycle](auto const i)
                {
                    if(i % 8 == 0)
                    {
                        mix(mCipher[0], mCipher[1], 14U);
                        mix(mCipher[2], mCipher[3], 16U);
                    }
 
                    if(i % 8 == 1)
                    {
                        mix(mCipher[0], mCipher[3], 52U);
                        mix(mCipher[2], mCipher[1], 57U);
                    }
 
                    if(i % 8 == 2)
                    {
                        mix(mCipher[0], mCipher[1], 23U);
                        mix(mCipher[2], mCipher[3], 40U);
                    }
 
                    if(i % 8 == 3)
                    {
                        mix(mCipher[0], mCipher[3], 5U);
                        mix(mCipher[2], mCipher[1], 37U);
                        add(mCipher, mKey, (i + 1) / 4, ++cycle);
                    }
 
                    if(i % 8 == 4)
                    {
                        mix(mCipher[0], mCipher[1], 25U);
                        mix(mCipher[2], mCipher[3], 33U);
                    }
 
                    if(i % 8 == 5)
                    {
                        mix(mCipher[0], mCipher[3], 46U);
                        mix(mCipher[2], mCipher[1], 12U);
                    }
 
                    if(i % 8 == 6)
                    {
                        mix(mCipher[0], mCipher[1], 58U);
                        mix(mCipher[2], mCipher[3], 22U);
                    }
 
                    if(i % 8 == 7)
                    {
                        mix(mCipher[0], mCipher[3], 32U);
                        mix(mCipher[2], mCipher[1], 32U);
                        add(mCipher, mKey, (i + 1) / 4, ++cycle);
                    }
                });
        }
 
    private:
        StackArray<std::uint64_t, 5> mKey;
        StackArray<std::uint64_t, 4> mCounter;
        StackArray<std::uint64_t, 4> mCipher;
        std::uint8_t mChunkCounter;
    };
}
#endif // pRC_CORE_RANDOM_THREEFRY_H