{-# 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 {
      BitStream s -> STRef s Word
_curRandom :: !(STRef s Word)
    , BitStream s -> STRef s Int
_bitsLeft  :: !(STRef s Int )
    , BitStream s -> STRef s Int
_randomPos :: !(STRef s Int )
    }


------------------------------------------------------------------------------
newBitStream :: ST s (BitStream s)
newBitStream :: 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 (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 (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 :: 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 :: Int -> BitStream s -> ST s Word
getNBits Int
nbits (BitStream STRef s Word
crRef STRef s Int
blRef STRef s Int
rpRef) = do
    !Int
bl <- STRef s Int -> ST s Int
forall s a. STRef s a -> ST s a
readSTRef STRef s Int
blRef
    if Int
bl Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
nbits
      then ST s Word
newWord
      else Int -> ST s Word
nextBits Int
bl

  where
    newWord :: ST s Word
newWord = do
        !Int
rp <- STRef s Int -> ST s Int
forall s a. STRef s a -> ST s a
readSTRef STRef s Int
rpRef
        let r :: Word
r = Vector Word -> Int -> Word
forall a. Unbox a => Vector a -> Int -> a
V.unsafeIndex Vector Word
randoms Int
rp
        STRef s Int -> Int -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s Int
blRef (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$! Int
wordSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
nbits
        STRef s Int -> Int -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s Int
rpRef (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$! if Int
rp Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== (Int
numRandomsInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) then Int
0 else Int
rp Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
        Word -> ST s Word
extractBits Word
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 (m :: * -> *) a. Monad m => a -> m a
return Word
b

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