-- | Efficient, pure, standard 32-bit MT19937 Mersenne Twister implementation.

module MT19937.Pure
  ( MT19937(idx, mt), MT19937.Pure.init, init', extract, skip
  ) where

import MT19937.Internal ( temper, twist )

import Data.Word ( Word32 )
import Data.Bits

import Data.Vector.Unboxed qualified as VU
import Data.Vector.Unboxed.Mutable qualified as VUM

import Control.Monad ( replicateM_ )

-- | MT19937 state, holding the state vector and the current index.
data MT19937 = MT19937
  { MT19937 -> Int
idx :: Int
  , MT19937 -> Vector Word32
mt  :: VU.Vector Word32
  } deriving stock Int -> MT19937 -> ShowS
[MT19937] -> ShowS
MT19937 -> String
(Int -> MT19937 -> ShowS)
-> (MT19937 -> String) -> ([MT19937] -> ShowS) -> Show MT19937
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> MT19937 -> ShowS
showsPrec :: Int -> MT19937 -> ShowS
$cshow :: MT19937 -> String
show :: MT19937 -> String
$cshowList :: [MT19937] -> ShowS
showList :: [MT19937] -> ShowS
Show

-- | Initialize an MT19937 with the given seed.
init :: Word32 -> MT19937
init :: Word32 -> MT19937
init Word32
seed = MT19937 { idx :: Int
idx = Int
0, mt :: Vector Word32
mt = Vector Word32
initState }
  where
    initState :: Vector Word32
initState = (forall s. ST s (MVector s Word32)) -> Vector Word32
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Word32)) -> Vector Word32)
-> (forall s. ST s (MVector s Word32)) -> Vector Word32
forall a b. (a -> b) -> a -> b
$ do
        MVector s Word32
mt' <- Int -> ST s (MVector (PrimState (ST s)) Word32)
forall (m :: Type -> Type) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
624
        MVector (PrimState (ST s)) Word32 -> Int -> Word32 -> ST s ()
forall (m :: Type -> Type) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Word32
MVector (PrimState (ST s)) Word32
mt' Int
0 Word32
seed
        MVector (PrimState (ST s)) Word32 -> Word32 -> ST s ()
forall (m :: Type -> Type).
PrimMonad m =>
MVector (PrimState m) Word32 -> Word32 -> m ()
initHelper MVector s Word32
MVector (PrimState (ST s)) Word32
mt' Word32
1
        MVector (PrimState (ST s)) Word32 -> ST s ()
forall (m :: Type -> Type).
PrimMonad m =>
MVector (PrimState m) Word32 -> m ()
twist MVector s Word32
MVector (PrimState (ST s)) Word32
mt'
        MVector s Word32 -> ST s (MVector s Word32)
forall a. a -> ST s a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure MVector s Word32
mt'

-- | Initialize an MT19937 with the given seed, without pre-twisting.
init' :: Word32 -> MT19937
init' :: Word32 -> MT19937
init' Word32
seed = MT19937 { idx :: Int
idx = Int
624, mt :: Vector Word32
mt = Vector Word32
initState }
  where
    initState :: Vector Word32
initState = (forall s. ST s (MVector s Word32)) -> Vector Word32
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Word32)) -> Vector Word32)
-> (forall s. ST s (MVector s Word32)) -> Vector Word32
forall a b. (a -> b) -> a -> b
$ do
        MVector s Word32
mt' <- Int -> ST s (MVector (PrimState (ST s)) Word32)
forall (m :: Type -> Type) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
624
        MVector (PrimState (ST s)) Word32 -> Int -> Word32 -> ST s ()
forall (m :: Type -> Type) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Word32
MVector (PrimState (ST s)) Word32
mt' Int
0 Word32
seed
        MVector (PrimState (ST s)) Word32 -> Word32 -> ST s ()
forall (m :: Type -> Type).
PrimMonad m =>
MVector (PrimState m) Word32 -> Word32 -> m ()
initHelper MVector s Word32
MVector (PrimState (ST s)) Word32
mt' Word32
1
        MVector s Word32 -> ST s (MVector s Word32)
forall a. a -> ST s a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure MVector s Word32
mt'

-- separated out due to ST scoping issues idk
initHelper
    :: VUM.PrimMonad m => VU.MVector (VUM.PrimState m) Word32 -> Word32 -> m ()
initHelper :: forall (m :: Type -> Type).
PrimMonad m =>
MVector (PrimState m) Word32 -> Word32 -> m ()
initHelper MVector (PrimState m) Word32
mt' = \case
  Word32
624 -> () -> m ()
forall a. a -> m a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ()
  Word32
i   -> do
    Word32
prev <- MVector (PrimState m) Word32 -> Int -> m Word32
forall (m :: Type -> Type) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector (PrimState m) Word32
mt' (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word32
iWord32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
-Word32
1))
    let mti :: Word32
mti = Word32
f Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
* (Word32
prev Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
`xor` (Word32
prev Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`shiftR` (Int
wInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
2))) Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
i
    MVector (PrimState m) Word32 -> Int -> Word32 -> m ()
forall (m :: Type -> Type) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector (PrimState m) Word32
mt' (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
i) Word32
mti
    MVector (PrimState m) Word32 -> Word32 -> m ()
forall (m :: Type -> Type).
PrimMonad m =>
MVector (PrimState m) Word32 -> Word32 -> m ()
initHelper MVector (PrimState m) Word32
mt' (Word32
iWord32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+Word32
1)
  where
    f :: Word32
f = Word32
1812433253 -- 0x6c078965
    w :: Int
w = Int
32

-- | Extract the next random word and return the updated state.
extract :: MT19937 -> (Word32, MT19937)
extract :: MT19937 -> (Word32, MT19937)
extract (MT19937 Int
idx Vector Word32
mt) = do
    if Int
idx Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
624 then -- could do a >= check here if we're paranoid
        let mt' :: Vector Word32
mt' = (forall s. MVector s Word32 -> ST s ())
-> Vector Word32 -> Vector Word32
forall a.
Unbox a =>
(forall s. MVector s a -> ST s ()) -> Vector a -> Vector a
VU.modify MVector s Word32 -> ST s ()
MVector (PrimState (ST s)) Word32 -> ST s ()
forall s. MVector s Word32 -> ST s ()
forall (m :: Type -> Type).
PrimMonad m =>
MVector (PrimState m) Word32 -> m ()
twist Vector Word32
mt
            w :: Word32
w   = Word32 -> Word32
forall a. (Num a, Bits a) => a -> a
temper (Vector Word32
mt' Vector Word32 -> Int -> Word32
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
0)
        in  (Word32
w, Int -> Vector Word32 -> MT19937
MT19937 Int
1       Vector Word32
mt')
    else
        let w :: Word32
w   = Word32 -> Word32
forall a. (Num a, Bits a) => a -> a
temper (Vector Word32
mt Vector Word32 -> Int -> Word32
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
idx)
        in  (Word32
w, Int -> Vector Word32 -> MT19937
MT19937 (Int
idxInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Vector Word32
mt)

-- | Skip the given number of random words.
--
-- If the skips would result in multiple twists, we perform these in a single
-- pass (rather than copying the array every twist).
skip :: Int -> MT19937 -> MT19937
skip :: Int -> MT19937 -> MT19937
skip Int
n (MT19937 Int
idx Vector Word32
mt) =
    let (Int
twists, Int
idx') = (Int
idxInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
n) Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`quotRem` Int
624
    in  if   Int
twists Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
        then let mt' :: Vector Word32
mt' = (forall s. MVector s Word32 -> ST s ())
-> Vector Word32 -> Vector Word32
forall a.
Unbox a =>
(forall s. MVector s a -> ST s ()) -> Vector a -> Vector a
VU.modify (Int -> ST s () -> ST s ()
forall (m :: Type -> Type) a. Applicative m => Int -> m a -> m ()
replicateM_ Int
twists (ST s () -> ST s ())
-> (MVector s Word32 -> ST s ()) -> MVector s Word32 -> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MVector s Word32 -> ST s ()
MVector (PrimState (ST s)) Word32 -> ST s ()
forall (m :: Type -> Type).
PrimMonad m =>
MVector (PrimState m) Word32 -> m ()
twist) Vector Word32
mt
             in  Int -> Vector Word32 -> MT19937
MT19937 Int
idx' Vector Word32
mt'
        else Int -> Vector Word32 -> MT19937
MT19937 Int
idx' Vector Word32
mt