Extra.hs

-- |
-- Module      : System.Random.Extra
-- Copyright   : (c) 2026 Institute for Digital Autonomy
-- License     : EUPL-1.2
-- Maintainer  : IDA
--
-- Thread-safe, low-contention generation of random values.
module System.Random.Extra where

import Control.Concurrent (myThreadId, threadCapability)
import Data.IORef (IORef, atomicModifyIORef', newIORef)
import Data.Tuple (swap)
import Data.Vector (Vector)
import Data.Vector qualified as Vector
import GHC.Conc (getNumCapabilities)
import System.IO.Unsafe (unsafePerformIO)
import System.Random (StdGen, newStdGen)
import System.Random.Stateful (Uniform, runStateGen, uniformM)
import Prelude

-- | One 'StdGen' per capability, each seeded independently from the global one.
generators :: Vector (IORef StdGen)
generators = unsafePerformIO do
    capabilities <- getNumCapabilities
    Vector.replicateM capabilities $ newIORef =<< newStdGen
{-# NOINLINE generators #-}

-- | Draw a uniformly random value from the current capability's generator.
uniformIO :: (Uniform a) => IO a
uniformIO = do
    (capability, _pinned) <- threadCapability =<< myThreadId
    let generator = generators Vector.! (capability `mod` Vector.length generators)
    atomicModifyIORef' generator \g -> swap $ runStateGen g uniformM
{-# INLINE uniformIO #-}