module PureSAT.LCG (
    LCG,
    newLCG,
    nextLCG,
) where

import Control.Monad.ST       (ST)
import Data.Primitive.PrimVar
import Data.Word              (Word64)

-- $setup
-- >>> import Control.Monad.ST (runST)

-- | Park-Miller LCG
--
-- >>> runST $ do { lcg <- newLCG 42; traverse (\_ -> nextLCG lcg) [1..10] }
-- [2027382,1226992407,551494037,961371815,1404753842,2076553157,1350734175,1538354858,90320905,488601845]
--
newtype LCG s = LCG (PrimVar s Word64)

newLCG :: Word64 -> ST s (LCG s)
newLCG :: forall s. Word64 -> ST s (LCG s)
newLCG Word64
seed = PrimVar s Word64 -> LCG s
forall s. PrimVar s Word64 -> LCG s
LCG (PrimVar s Word64 -> LCG s)
-> ST s (PrimVar s Word64) -> ST s (LCG s)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Word64 -> ST s (PrimVar (PrimState (ST s)) Word64)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
a -> m (PrimVar (PrimState m) a)
newPrimVar (Word64 -> Word64 -> Word64
forall a. Ord a => a -> a -> a
max Word64
1 (Word64 -> Word64 -> Word64
forall a. Ord a => a -> a -> a
min Word64
0x7fffffe Word64
seed))

nextLCG :: LCG s -> ST s Word64
nextLCG :: forall s. LCG s -> ST s Word64
nextLCG (LCG PrimVar s Word64
state) = do
    Word64
s <- PrimVar (PrimState (ST s)) Word64 -> ST s Word64
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> m a
readPrimVar PrimVar s Word64
PrimVar (PrimState (ST s)) Word64
state
    let !s' :: Word64
s' = Word64 -> Word64 -> Word64
forall a. Integral a => a -> a -> a
mod (Word64
s Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
48271) Word64
0x7fffffff
    PrimVar (PrimState (ST s)) Word64 -> Word64 -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> a -> m ()
writePrimVar PrimVar s Word64
PrimVar (PrimState (ST s)) Word64
state Word64
s'
    Word64 -> ST s Word64
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return Word64
s'