public
Created

Simple, pure MWC random number generator with 256 bits of state

  • Download Gist
RandomMWC8.hs
Haskell
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163
{-# LANGUAGE BangPatterns, MagicHash #-}
 
-- IMPORTANT: Assumes a 64 bit platform!
 
import GHC.Prim
import GHC.Word
import GHC.Types
import Data.Word
import Numeric
import Data.Bits
 
data PrngState = PrngState
{ prng_c :: Word#
, prng_x04 :: Word#
, prng_x15 :: Word#
, prng_x26 :: Word#
, prng_x37 :: Word#
}
 
instance Show PrngState where
show (PrngState c x04 x15 x26 x37) =
"PrngState " ++ show (W# c)
++ " " ++ showHex (W# x04) ""
++ " " ++ showHex (W# x15) ""
++ " " ++ showHex (W# x26) ""
++ " " ++ showHex (W# x37) ""
 
defaultPrng :: PrngState
defaultPrng = PrngState (int2Word# 362436#)
(int2Word# 0x54f7a8797042e8b3#)
(int2Word# 0x0474b18406f7f4c5#)
(int2Word# 0xb3f8f692789ea382#)
(int2Word# 0x4114ea356fb15ad8#)
 
-- TODO: Add version that takes a list of Ints or Integer
 
initPrng :: Int -> PrngState
initPrng (I# n) =
let !seed = int2Word# n in
case defaultPrng of
PrngState c x04 x15 x26 x37 ->
PrngState c (x04 `xor#` seed)
(x15 `xor#` seed)
(x26 `xor#` seed)
(x37 `xor#` seed)
 
uniformWord32 :: PrngState -> (Word32, PrngState)
uniformWord32 (PrngState c0 x04 x15 x26 x37) =
let !q = x04 `and#` int2Word# 0xffffffff#
!t = (int2Word# 716514398# `timesWord#` q) `plusWord#` c0
!c = uncheckedShiftRL# t 32#
!x = (t `and#` int2Word# 0xffffffff#) `plusWord#` c
-- TODO: Make the below branchless
!d = if x `ltWord#` c then int2Word# 1# else int2Word# 0#
!x' = x `plusWord#` d
!c' = c `plusWord#` d
!x'' = (int2Word# 0xfffffffe# `minusWord#` x') `and#` (int2Word# 0xffffffff#)
!x40 = (x04 `uncheckedShiftL#` 32#) `or#` x''
-- See "The Rotating Array Queue" below
!state' = PrngState c' x15 x26 x37 x40
in (W32# x'', state')
 
-- TODO: Splitting
-- TODO: Verify (how?)
 
randoms :: PrngState -> [Word32]
randoms st0 = go st0
where go st = case uniformWord32 st of
(x, st') -> x : go st'
 
{-
The Rotating Array Queue
------------------------
An MWC PRNG normally uses `2^r` seed words which are usually stored
in a mutable array. Since we want our PRNG to be pure this would
mean copying the array each time. Arrays have additional meta-data
so that's not ideal.
We can unpack the array into the state datatype. Since we need to
copy the state around anyway, we can also rotate it at the same
time. For example, let's say our state consists of six 32-bit
words:
> import Data.Bits
> import Data.Word
>
> data RotatingArray = A !Word32 !Word32 !Word32
> !Word32 !Word32 !Word32
>
> mutateAtCursor :: (Word32 -> Word32)
> -> RotatingArray -> (Word32, RotatingArray)
> mutateAtCursor f (A x0 x1 x2 x3 x4 x5) =
> let !x0' = f x0
> !a = A x1 x2 x3 x4 x5 x0'
> in (x0', a)
This works well enough and avoids the extra storage for the cursor
and array size meta data.
Now assume that we are on an architecture with 64 bit words, but
our state is logically still six 32 bit words. We could just
store each 32 bit word inside a 64 bit machine word, but this is
wasteful. It would be better to somehow pack these six 32 bit
words into three 64 bit words.
If we use this packed representation, then straightforward rotation
on these 64 bit words doesn't work anymore.
> data RotatingArray64 = A64 !Word64 !Word64 !Word64
 
> lo :: Word64 -> Word32
> lo w = fromIntegral w
 
> hi :: Word64 -> Word32
> hi w = fromIntegral (w `shiftR` 32)
 
> toLo :: Word32 -> Word64
> toLo w = fromIntegral w
 
> toHi :: Word32 -> Word64
> toHi w = fromIntegral w `shiftL` 32
 
> mutateAtCursorSlow :: (Word32 -> Word32)
> -> RotatingArray64 -> (Word32, RotatingArray64)
> mutateAtCursorSlow f (A64 x01 x23 x45) =
> let !x0 = lo x01
> !x0' = f x0
> !x12 = toLo (hi x01) .|. toHi (lo x23)
> !x34 = toLo (hi x23) .|. toHi (lo x45)
> !x50 = toLo (hi x45) .|. toHi x0'
> !a = A64 x12 x34 x50
> in (x0', a)
 
This works, but there's a fair amount of extra computation needed. We
can do better by changing the way we pack 32 bit words into 64 bit
words.
 
> mkRotatingArray64 :: Word32 -> Word32 -> Word32
> -> Word32 -> Word32 -> Word32
> -> RotatingArray64
> mkRotatingArray64 !x0 !x1 !x2 !x3 !x4 !x5 =
> let !x03 = toLo x0 .|. toHi x3
> !x14 = toLo x1 .|. toHi x4
> !x25 = toLo x2 .|. toHi x5
> in A64 x03 x14 x25
 
> mutateAtCursorFast :: (Word32 -> Word32)
> -> RotatingArray64 -> (Word32, RotatingArray64)
> mutateAtCursorFast f (A64 x03 x14 x25) =
> let !x0 = lo x03
> !x0' = f x0
> -- Swap hi and lo before pushing it at the end of the queue.
> !x30 = toLo (hi x03) .|. toHi x0'
> !a = A64 x14 x25 x30
> in (x0', a)
 
The cursor is always focused at the lower 32 bits of the first state
word. However, after we've computed the new version of that element
we put it back as the *upper* 32 bits of the last element. It will
thus come back as the focused elements after two rotations of the 64
bit words.
-}

Please sign in to comment on this gist.

Something went wrong with that request. Please try again.