1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE CPP #-}
3
4module Data.HashTable.Internal.CheapPseudoRandomBitStream
5  ( BitStream
6  , newBitStream
7  , getNextBit
8  , getNBits
9  ) where
10
11import           Control.Applicative
12import           Control.Monad.ST
13import           Data.Bits                     ((.&.))
14import           Data.STRef
15import           Data.Vector.Unboxed           (Vector)
16import qualified Data.Vector.Unboxed           as V
17
18#if __GLASGOW_HASKELL__ >= 808
19import           Data.Word                     (Word32, Word64)
20#else
21import           Data.Word                     (Word, Word32, Word64)
22#endif
23
24import           Data.HashTable.Internal.Utils
25
26
27------------------------------------------------------------------------------
28-- Chosen by fair dice roll. Guaranteed random. More importantly, there are an
29-- equal number of 0 and 1 bits in both of these vectors.
30random32s :: Vector Word32
31random32s = V.fromList [ 0xe293c315
32                       , 0x82e2ff62
33                       , 0xcb1ef9ae
34                       , 0x78850172
35                       , 0x551ee1ce
36                       , 0x59d6bfd1
37                       , 0xb717ec44
38                       , 0xe7a3024e
39                       , 0x02bb8976
40                       , 0x87e2f94f
41                       , 0xfa156372
42                       , 0xe1325b17
43                       , 0xe005642a
44                       , 0xc8d02eb3
45                       , 0xe90c0a87
46                       , 0x4cb9e6e2
47                       ]
48
49
50------------------------------------------------------------------------------
51random64s :: Vector Word64
52random64s = V.fromList [ 0x62ef447e007e8732
53                       , 0x149d6acb499feef8
54                       , 0xca7725f9b404fbf8
55                       , 0x4b5dfad194e626a9
56                       , 0x6d76f2868359491b
57                       , 0x6b2284e3645dcc87
58                       , 0x5b89b485013eaa16
59                       , 0x6e2d4308250c435b
60                       , 0xc31e641a659e0013
61                       , 0xe237b85e9dc7276d
62                       , 0x0b3bb7fa40d94f3f
63                       , 0x4da446874d4ca023
64                       , 0x69240623fedbd26b
65                       , 0x76fb6810dcf894d3
66                       , 0xa0da4e0ce57c8ea7
67                       , 0xeb76b84453dc3873
68                       ]
69
70
71------------------------------------------------------------------------------
72numRandoms :: Int
73numRandoms = 16
74
75
76------------------------------------------------------------------------------
77randoms :: Vector Word
78randoms | wordSize == 32 = V.map fromIntegral random32s
79        | otherwise      = V.map fromIntegral random64s
80
81
82------------------------------------------------------------------------------
83data BitStream s = BitStream {
84      _curRandom :: !(STRef s Word)
85    , _bitsLeft  :: !(STRef s Int )
86    , _randomPos :: !(STRef s Int )
87    }
88
89
90------------------------------------------------------------------------------
91newBitStream :: ST s (BitStream s)
92newBitStream =
93    unwrapMonad $
94    BitStream <$> (WrapMonad $ newSTRef $ V.unsafeIndex randoms 0)
95              <*> (WrapMonad $ newSTRef wordSize)
96              <*> (WrapMonad $ newSTRef 1)
97
98
99------------------------------------------------------------------------------
100getNextBit :: BitStream s -> ST s Word
101getNextBit = getNBits 1
102
103
104------------------------------------------------------------------------------
105getNBits :: Int -> BitStream s -> ST s Word
106getNBits nbits (BitStream crRef blRef rpRef) = do
107    !bl <- readSTRef blRef
108    if bl < nbits
109      then newWord
110      else nextBits bl
111
112  where
113    newWord = do
114        !rp <- readSTRef rpRef
115        let r = V.unsafeIndex randoms rp
116        writeSTRef blRef $! wordSize - nbits
117        writeSTRef rpRef $! if rp == (numRandoms-1) then 0 else rp + 1
118        extractBits r
119
120    extractBits r = do
121        let !b = r .&. ((1 `shiftL` nbits) - 1)
122        writeSTRef crRef $! (r `shiftRL` nbits)
123        return b
124
125    nextBits bl = do
126        !r <- readSTRef crRef
127        writeSTRef blRef $! bl - nbits
128        extractBits r
129