Skip to content

Instantly share code, notes, and snippets.

@nominolo
Created October 19, 2012 16:09
Show Gist options
  • Save nominolo/3919067 to your computer and use it in GitHub Desktop.
Save nominolo/3919067 to your computer and use it in GitHub Desktop.
Simple, pure MWC random number generator with 256 bits of state
{-# 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.
-}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment