PRNG.hs

{-# OPTIONS_HADDOCK prune #-}
{-# LANGUAGE BangPatterns #-}

-- |
-- Module: Crypto.PRNG
-- Copyright: (c) 2026 Jared Tobin
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
-- A fast, non-cryptographic pseudorandom number generator:
-- [xoshiro256++](https://prng.di.unimi.it/), by David Blackman and
-- Sebastiano Vigna.
--
-- The generator is mutable, living in any 'PrimMonad' (e.g. 'IO' or
-- 'ST s'), and is not thread-safe. Use one generator per thread,
-- deriving non-overlapping streams via 'copy' and 'jump' or
-- 'long_jump'.
--
-- Its output is predictable from a small sample; use "Crypto.CSPRNG"
-- for anything that must be unpredictable. "Crypto.PRNG.Pure" offers
-- the same generator, with the same output, as an immutable value.

module Crypto.PRNG (
    -- * Generator
    Gen
  , new
  , new_system
  , from_state

    -- * Output
  , word64
  , word64_range
  , int_range
  , double

    -- * Streams
  , copy
  , jump
  , long_jump
  ) where

import Control.Monad.Primitive (PrimMonad, PrimState)
import Control.Monad.ST (ST)
import qualified Crypto.Range as Range
import Crypto.Xoshiro (S(..))
import qualified Crypto.Xoshiro as X
import Data.Bits ((.>>.))
import qualified Data.Primitive.PrimArray as PA
import Data.Word (Word64)
import GHC.Exts (RealWorld)

-- utils ----------------------------------------------------------------------

fi :: (Integral a, Num b) => a -> b
fi = fromIntegral
{-# INLINE fi #-}

-- generator ------------------------------------------------------------------

-- | A xoshiro256++ generator.
newtype Gen s = Gen (PA.MutablePrimArray s Word64)

read_s :: PrimMonad m => Gen (PrimState m) -> m S
read_s (Gen g) = do
  !s0 <- PA.readPrimArray g 0
  !s1 <- PA.readPrimArray g 1
  !s2 <- PA.readPrimArray g 2
  !s3 <- PA.readPrimArray g 3
  pure $! S s0 s1 s2 s3
{-# INLINE read_s #-}

write_s :: PrimMonad m => Gen (PrimState m) -> S -> m ()
write_s (Gen g) (S s0 s1 s2 s3) = do
  PA.writePrimArray g 0 s0
  PA.writePrimArray g 1 s1
  PA.writePrimArray g 2 s2
  PA.writePrimArray g 3 s3
{-# INLINE write_s #-}

alloc :: PrimMonad m => S -> m (Gen (PrimState m))
alloc s = do
  g <- fmap Gen (PA.newPrimArray 4)
  write_s g s
  pure g
{-# INLINE alloc #-}

-- | Create a generator from a 64-bit seed, expanded to the full
--   256-bit state via splitmix64, per the xoshiro authors'
--   recommendation.
--
--   >>> gen <- new 42
--   >>> word64 gen
--   15021278609987233951
new :: PrimMonad m => Word64 -> m (Gen (PrimState m))
new = alloc . X.seed
{-# INLINABLE new #-}
{-# SPECIALIZE new :: Word64 -> IO (Gen RealWorld) #-}
{-# SPECIALIZE new :: Word64 -> ST s (Gen s) #-}

-- | Create a generator from the given 256-bit state, returning
--   'Nothing' if it is all zero.
--
--   >>> Just gen <- from_state 1 2 3 4
--   >>> word64 gen
--   41943041
from_state
  :: PrimMonad m
  => Word64
  -> Word64
  -> Word64
  -> Word64
  -> m (Maybe (Gen (PrimState m)))
from_state s0 s1 s2 s3 = traverse alloc (X.from_words s0 s1 s2 s3)
{-# INLINABLE from_state #-}

-- | Create a generator with a 256-bit state drawn from the operating
--   system's entropy source, returning 'Nothing' if that fails.
--
--   >>> Just gen <- new_system
new_system :: IO (Maybe (Gen RealWorld))
new_system = X.system >>= traverse alloc

-- output ---------------------------------------------------------------------

-- | Generate a uniformly-distributed 'Word64'.
--
--   >>> Just gen <- from_state 1 2 3 4
--   >>> word64 gen
--   41943041
--   >>> word64 gen
--   58720359
word64 :: PrimMonad m => Gen (PrimState m) -> m Word64
word64 gen = do
  !s <- read_s gen
  write_s gen (X.next s)
  pure $! X.output s
{-# INLINE word64 #-}

-- | Generate a uniformly-distributed 'Word64' in the given inclusive
--   range, whose bounds may be given in either order.
--
--   The result is unbiased; see 'int_range'.
--
--   >>> gen <- new 42
--   >>> word64_range 1 6 gen
--   5
word64_range
  :: PrimMonad m => Word64 -> Word64 -> Gen (PrimState m) -> m Word64
word64_range a b gen = do
  let !lo = min a b
      !hi = max a b
  !r <- Range.upto (word64 gen) (hi - lo)
  pure $! lo + r
{-# INLINE word64_range #-}

-- | Generate a uniformly-distributed 'Int' in the given inclusive
--   range, whose bounds may be given in either order.
--
--   The result is unbiased, via Lemire's method: an occasional draw is
--   rejected and replaced, so the number of words consumed varies.
--
--   >>> gen <- new 42
--   >>> int_range (-10) 10 gen
--   7
int_range :: PrimMonad m => Int -> Int -> Gen (PrimState m) -> m Int
int_range a b gen = do
  let !lo = min a b
      !hi = max a b
  !r <- Range.upto (word64 gen) (fi hi - fi lo)
  pure $! lo + fi r
{-# INLINE int_range #-}

-- | Generate a uniformly-distributed 'Double' in [0, 1), from the
--   upper 53 bits of a 'Word64'.
--
--   >>> Just gen <- from_state 1 2 3 4
--   >>> double gen
--   2.2737367544323206e-12
double :: PrimMonad m => Gen (PrimState m) -> m Double
double gen = do
  !w <- word64 gen
  -- 2 ^ -53
  pure $! fi (w .>>. 11) * 1.1102230246251565e-16
{-# INLINE double #-}

-- streams --------------------------------------------------------------------

-- | Copy a generator. The copy produces the same output as the
--   original, independently of it.
--
--   >>> gen <- new 42
--   >>> cop <- copy gen
--   >>> (==) <$> word64 gen <*> word64 cop
--   True
copy :: PrimMonad m => Gen (PrimState m) -> m (Gen (PrimState m))
copy gen = read_s gen >>= alloc
{-# INLINE copy #-}

-- | Advance a generator by 2 ^ 128 outputs.
--
--   Copying and then jumping repeatedly yields up to 2 ^ 128
--   non-overlapping streams of 2 ^ 128 outputs each, e.g. one per
--   thread.
--
--   >>> gen <- new 42
--   >>> one <- copy gen
--   >>> jump gen
--   >>> two <- copy gen
--   >>> jump gen
jump :: PrimMonad m => Gen (PrimState m) -> m ()
jump gen = do
  !s <- read_s gen
  write_s gen (X.jump s)
{-# INLINABLE jump #-}

-- | Advance a generator by 2 ^ 192 outputs.
--
--   Copying and long-jumping repeatedly yields up to 2 ^ 64 starting
--   points, from each of which 'jump' yields 2 ^ 64 non-overlapping
--   streams.
--
--   >>> gen <- new 42
--   >>> long_jump gen
long_jump :: PrimMonad m => Gen (PrimState m) -> m ()
long_jump gen = do
  !s <- read_s gen
  write_s gen (X.long_jump s)
{-# INLINABLE long_jump #-}