{-# 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 = 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 = 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 forall a. Eq a => a -> a -> Bool
== Int
32 = forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
V.map forall a b. (Integral a, Num b) => a -> b
fromIntegral Vector Word32
random32s
        | Bool
otherwise      = forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
V.map 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 =
    forall (m :: * -> *) a. WrappedMonad m a -> m a
unwrapMonad forall a b. (a -> b) -> a -> b
$
    forall s. STRef s Word -> STRef s Int -> STRef s Int -> BitStream s
BitStream forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall (m :: * -> *) a. m a -> WrappedMonad m a
WrapMonad forall a b. (a -> b) -> a -> b
$ forall a s. a -> ST s (STRef s a)
newSTRef forall a b. (a -> b) -> a -> b
$ forall a. Unbox a => Vector a -> Int -> a
V.unsafeIndex Vector Word
randoms Int
0)
              forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (forall (m :: * -> *) a. m a -> WrappedMonad m a
WrapMonad forall a b. (a -> b) -> a -> b
$ forall a s. a -> ST s (STRef s a)
newSTRef Int
wordSize)
              forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (forall (m :: * -> *) a. m a -> WrappedMonad m a
WrapMonad forall a b. (a -> b) -> a -> b
$ 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 = 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
    !Int
bl <- forall s a. STRef s a -> ST s a
readSTRef STRef s Int
blRef
    if Int
bl 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 <- forall s a. STRef s a -> ST s a
readSTRef STRef s Int
rpRef
        let r :: Word
r = forall a. Unbox a => Vector a -> Int -> a
V.unsafeIndex Vector Word
randoms Int
rp
        forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s Int
blRef forall a b. (a -> b) -> a -> b
$! Int
wordSize forall a. Num a => a -> a -> a
- Int
nbits
        forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s Int
rpRef forall a b. (a -> b) -> a -> b
$! if Int
rp forall a. Eq a => a -> a -> Bool
== (Int
numRandomsforall a. Num a => a -> a -> a
-Int
1) then Int
0 else Int
rp 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 forall a. Bits a => a -> a -> a
.&. ((Word
1 Word -> Int -> Word
`shiftL` Int
nbits) forall a. Num a => a -> a -> a
- Word
1)
        forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s Word
crRef forall a b. (a -> b) -> a -> b
$! (Word
r Word -> Int -> Word
`shiftRL` Int
nbits)
        forall (m :: * -> *) a. Monad m => a -> m a
return Word
b

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