Rng.hs

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

-- |
-- Module: Censor.Rng
-- Copyright: (c) 2026 Jared Tobin
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
-- A minimal splitmix64 PRNG for input samplers.
--
-- Samplers feeding 'Censor.runCT' want a pseudorandom source that is
-- deterministic under a fixed seed, cheap, and light on allocation,
-- so that sampling work does not itself perturb the measurement (see
-- the measurement-hygiene notes in "Censor"). This module provides
-- exactly that and nothing more. It is /not/ a cryptographic
-- generator; the sampled inputs only need to be independent of the
-- target's timing distribution, not unpredictable.

module Censor.Rng (
    -- * Stateful generator
    Rng
  , mkRng
  , nextWord
  , reseed

    -- * Byte fills
  , randomBytes
  , fillRandom

    -- * Pure step
  , splitMix
  ) where

import qualified Data.Bits as B
import qualified Data.ByteString as BS
import qualified Data.ByteString.Internal as BI
import Data.IORef
import Data.Word (Word8, Word64)
import Foreign.Ptr (Ptr)
import Foreign.Storable (pokeByteOff)

-- | splitmix64 state threaded through an 'IORef'.
newtype Rng = Rng (IORef Word64)

-- | Create a generator from a seed. Equal seeds yield equal streams.
mkRng :: Word64 -> IO Rng
mkRng s = Rng <$> newIORef s
{-# INLINE mkRng #-}

-- | Draw the next 'Word64' from the generator.
nextWord :: Rng -> IO Word64
nextWord (Rng ref) = atomicModifyIORef' ref splitMix
{-# INLINE nextWord #-}

-- | Reset a generator to a seed in place: after @reseed g s@ the
--   generator draws exactly the stream of a fresh @'mkRng' s@.
--
--   Lets one long-lived generator replay a pinned stream, or adopt
--   a fresh one, without allocating per use — so the two classes
--   of a fix-vs-random sampler can share a single generator object
--   and differ only in the seed value written into it.
reseed :: Rng -> Word64 -> IO ()
reseed (Rng ref) s = atomicWriteIORef ref s
{-# INLINE reseed #-}

-- | The pure splitmix64 step: from a state, produce the successor
--   state and an output word.
splitMix :: Word64 -> (Word64, Word64)
splitMix !s =
  let !s' = s + 0x9E3779B97F4A7C15
      !z0 = (s' `B.xor` (s' `B.shiftR` 30)) * 0xBF58476D1CE4E5B9
      !z1 = (z0 `B.xor` (z0 `B.shiftR` 27)) * 0x94D049BB133111EB
      !z2 =  z1 `B.xor` (z1 `B.shiftR` 31)
  in  (s', z2)
{-# INLINE splitMix #-}

-- | Draw @n@ pseudorandom bytes as a strict 'BS.ByteString'.
--
--   The result is a fresh, fully-materialised heap value -- no
--   thunks for the meter to pay inside the timed region -- so a
--   sampler can return it directly.
randomBytes :: Rng -> Int -> IO BS.ByteString
randomBytes !g !n
  | n <= 0    = pure BS.empty
  | otherwise = BI.create n $ \p -> fillRandom g p n

-- | Fill @n@ bytes at the pointer with pseudorandom bytes drawn
--   from the generator.
fillRandom :: Rng -> Ptr Word8 -> Int -> IO ()
fillRandom !g !dst !n = go 0
  where
    go !i
      | i >= n    = pure ()
      | otherwise = do
          !w <- nextWord g
          let !left  = n - i
              !chunk = if left > 8 then 8 else left
          writeWord dst i w chunk
          go (i + chunk)

-- write the low @lim@ bytes of a 'Word64', little-endian, at the
-- given offset.
writeWord :: Ptr Word8 -> Int -> Word64 -> Int -> IO ()
writeWord !dst !off !w !lim = go 0
  where
    go !j
      | j >= lim  = pure ()
      | otherwise = do
          let !byte = fromIntegral (w `B.shiftR` (8 * j)) :: Word8
          pokeByteOff dst (off + j) byte
          go (j + 1)