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

RandomMWC8.hs

{-# 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.
-}