    {-# LANGUAGE BangPatterns #-}
    {-# LANGUAGE CPP #-}

    module Data.HashTable.Internal.CheapPseudoRandomBitStream
      ( BitStream
      , newBitStream
      , getNextBit
      , getNBits
      ) where

    import           Control.Applicative
    import           Control.Monad.ST
    import           Data.Bits                     ((.&.))
    import           Data.STRef
    import           Data.Vector.Unboxed           (Vector)
    import qualified Data.Vector.Unboxed           as V

    import           Data.Word                     (Word32, Word64)
    #if !MIN_VERSION_base(4,8,0)
    import           Data.Word                     (Word)
    #endif

    import           Data.HashTable.Internal.Utils

    ------------------------------------------------------------------------------
    -- Chosen by fair dice roll. Guaranteed random. More importantly, there are an
    -- equal number of 0 and 1 bits in both of these vectors.
    random32s :: Vector Word32
    random32s :: Vector Word32
    random32s = [Word32] -> Vector Word32
    forall a. Unbox a => [a] -> Vector a
    V.fromList [ Word32
    0xe293c315
                           , Word32
    0x82e2ff62
                           , Word32
    0xcb1ef9ae
                           , Word32
    0x78850172
                           , Word32
    0x551ee1ce
                           , Word32
    0x59d6bfd1
                           , Word32
    0xb717ec44
                           , Word32
    0xe7a3024e
                           , Word32
    0x02bb8976
                           , Word32
    0x87e2f94f
                           , Word32
    0xfa156372
                           , Word32
    0xe1325b17
                           , Word32
    0xe005642a
                           , Word32
    0xc8d02eb3
                           , Word32
    0xe90c0a87
                           , Word32
    0x4cb9e6e2
                           ]

    ------------------------------------------------------------------------------
    random64s :: Vector Word64
    random64s :: Vector Word64
    random64s = [Word64] -> Vector Word64
    forall a. Unbox a => [a] -> Vector a
    V.fromList [ Word64
    0x62ef447e007e8732
                           , Word64
    0x149d6acb499feef8
                           , Word64
    0xca7725f9b404fbf8
                           , Word64
    0x4b5dfad194e626a9
                           , Word64
    0x6d76f2868359491b
                           , Word64
    0x6b2284e3645dcc87
                           , Word64
    0x5b89b485013eaa16
                           , Word64
    0x6e2d4308250c435b
                           , Word64
    0xc31e641a659e0013
                           , Word64
    0xe237b85e9dc7276d
                           , Word64
    0x0b3bb7fa40d94f3f
                           , Word64
    0x4da446874d4ca023
                           , Word64
    0x69240623fedbd26b
                           , Word64
    0x76fb6810dcf894d3
                           , Word64
    0xa0da4e0ce57c8ea7
                           , Word64
    0xeb76b84453dc3873
                           ]

    ------------------------------------------------------------------------------
    numRandoms :: Int
    numRandoms :: Int
    numRandoms = Int
    16

    ------------------------------------------------------------------------------
    randoms :: Vector Word
    randoms :: Vector Word
    randoms | Int
    wordSize Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    32 = (Word32 -> Word) -> Vector Word32 -> Vector Word
    forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
    V.map Word32 -> Word
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral Vector Word32
    random32s
            | Bool
    otherwise      = (Word64 -> Word) -> Vector Word64 -> Vector Word
    forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
    V.map Word64 -> Word
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral Vector Word64
    random64s

    ------------------------------------------------------------------------------
    data BitStream s = BitStream {
          forall s. BitStream s -> STRef s Word
    _curRandom :: !(STRef s Word)
        , forall s. BitStream s -> STRef s Int
    _bitsLeft  :: !(STRef s Int )
        , forall s. BitStream s -> STRef s Int
    _randomPos :: !(STRef s Int )
        }

    ------------------------------------------------------------------------------
    newBitStream :: ST s (BitStream s)
    newBitStream :: forall s. ST s (BitStream s)
    newBitStream =
        WrappedMonad (ST s) (BitStream s) -> ST s (BitStream s)
    forall (m :: * -> *) a. WrappedMonad m a -> m a
    unwrapMonad (WrappedMonad (ST s) (BitStream s) -> ST s (BitStream s))
    -> WrappedMonad (ST s) (BitStream s) -> ST s (BitStream s)
    forall a b. (a -> b) -> a -> b
    $
        STRef s Word -> STRef s Int -> STRef s Int -> BitStream s
    forall s. STRef s Word -> STRef s Int -> STRef s Int -> BitStream s
    BitStream (STRef s Word -> STRef s Int -> STRef s Int -> BitStream s)
    -> WrappedMonad (ST s) (STRef s Word)
    -> WrappedMonad (ST s) (STRef s Int -> STRef s Int -> BitStream s)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (ST s (STRef s Word) -> WrappedMonad (ST s) (STRef s Word)
    forall (m :: * -> *) a. m a -> WrappedMonad m a
    WrapMonad (ST s (STRef s Word) -> WrappedMonad (ST s) (STRef s Word))
    -> ST s (STRef s Word) -> WrappedMonad (ST s) (STRef s Word)
    forall a b. (a -> b) -> a -> b
    $ Word -> ST s (STRef s Word)
    forall a s. a -> ST s (STRef s a)
    newSTRef (Word -> ST s (STRef s Word)) -> Word -> ST s (STRef s Word)
    forall a b. (a -> b) -> a -> b
    $ Vector Word -> Int -> Word
    forall a. Unbox a => Vector a -> Int -> a
    V.unsafeIndex Vector Word
    randoms Int
    0)
                  WrappedMonad (ST s) (STRef s Int -> STRef s Int -> BitStream s)
    -> WrappedMonad (ST s) (STRef s Int)
    -> WrappedMonad (ST s) (STRef s Int -> BitStream s)
    forall a b.
    WrappedMonad (ST s) (a -> b)
    -> WrappedMonad (ST s) a -> WrappedMonad (ST s) b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (ST s (STRef s Int) -> WrappedMonad (ST s) (STRef s Int)
    forall (m :: * -> *) a. m a -> WrappedMonad m a
    WrapMonad (ST s (STRef s Int) -> WrappedMonad (ST s) (STRef s Int))
    -> ST s (STRef s Int) -> WrappedMonad (ST s) (STRef s Int)
    forall a b. (a -> b) -> a -> b
    $ Int -> ST s (STRef s Int)
    forall a s. a -> ST s (STRef s a)
    newSTRef Int
    wordSize)
                  WrappedMonad (ST s) (STRef s Int -> BitStream s)
    -> WrappedMonad (ST s) (STRef s Int)
    -> WrappedMonad (ST s) (BitStream s)
    forall a b.
    WrappedMonad (ST s) (a -> b)
    -> WrappedMonad (ST s) a -> WrappedMonad (ST s) b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (ST s (STRef s Int) -> WrappedMonad (ST s) (STRef s Int)
    forall (m :: * -> *) a. m a -> WrappedMonad m a
    WrapMonad (ST s (STRef s Int) -> WrappedMonad (ST s) (STRef s Int))
    -> ST s (STRef s Int) -> WrappedMonad (ST s) (STRef s Int)
    forall a b. (a -> b) -> a -> b
    $ Int -> ST s (STRef s Int)
    forall a s. a -> ST s (STRef s a)
    newSTRef Int
    1)

    ------------------------------------------------------------------------------
    getNextBit :: BitStream s -> ST s Word
    getNextBit :: forall s. BitStream s -> ST s Word
    getNextBit = Int -> BitStream s -> ST s Word
    forall s. Int -> BitStream s -> ST s Word
    getNBits Int
    1

    ------------------------------------------------------------------------------
    getNBits :: Int -> BitStream s -> ST s Word
    getNBits :: forall s. Int -> BitStream s -> ST s Word
    getNBits Int
    nbits (BitStream STRef s Word
    crRef STRef s Int
    blRef STRef s Int
    rpRef) = do
        !bl <- STRef s Int -> ST s Int
    forall s a. STRef s a -> ST s a
    readSTRef STRef s Int
    blRef
        if bl < nbits
          then newWord
          else nextBits bl

      where
        newWord :: ST s Word
    newWord = do
            !rp <- STRef s Int -> ST s Int
    forall s a. STRef s a -> ST s a
    readSTRef STRef s Int
    rpRef
            let r = Vector Word -> Int -> Word
    forall a. Unbox a => Vector a -> Int -> a
    V.unsafeIndex Vector Word
    randoms Int
    rp
            writeSTRef blRef $! wordSize - nbits
            writeSTRef rpRef $! if rp == (numRandoms-1) then 0 else rp + 1
            extractBits r

        extractBits :: Word -> ST s Word
    extractBits Word
    r = do
            let !b :: Word
    b = Word
    r Word -> Word -> Word
    forall a. Bits a => a -> a -> a
    .&. ((Word
    1 Word -> Int -> Word
    `shiftL` Int
    nbits) Word -> Word -> Word
    forall a. Num a => a -> a -> a
    - Word
    1)
            STRef s Word -> Word -> ST s ()
    forall s a. STRef s a -> a -> ST s ()
    writeSTRef STRef s Word
    crRef (Word -> ST s ()) -> Word -> ST s ()
    forall a b. (a -> b) -> a -> b
    $! (Word
    r Word -> Int -> Word
    `shiftRL` Int
    nbits)
            Word -> ST s Word
    forall a. a -> ST s a
    forall (m :: * -> *) a. Monad m => a -> m a
    return Word
    b

        nextBits :: Int -> ST s Word
    nextBits Int
    bl = do
            !r <- STRef s Word -> ST s Word
    forall s a. STRef s a -> ST s a
    readSTRef STRef s Word
    crRef
            writeSTRef blRef $! bl - nbits
            extractBits r