{-# 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 :: Word64 -> IO Rng
mkRng Word64
s = IORef Word64 -> Rng
Rng (IORef Word64 -> Rng) -> IO (IORef Word64) -> IO Rng
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Word64 -> IO (IORef Word64)
forall a. a -> IO (IORef a)
newIORef Word64
s
{-# INLINE mkRng #-}

-- | Draw the next 'Word64' from the generator.
nextWord :: Rng -> IO Word64
nextWord :: Rng -> IO Word64
nextWord (Rng IORef Word64
ref) = IORef Word64 -> (Word64 -> (Word64, Word64)) -> IO Word64
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Word64
ref Word64 -> (Word64, Word64)
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 -> Word64 -> IO ()
reseed (Rng IORef Word64
ref) Word64
s = IORef Word64 -> Word64 -> IO ()
forall a. IORef a -> a -> IO ()
atomicWriteIORef IORef Word64
ref Word64
s
{-# INLINE reseed #-}

-- | The pure splitmix64 step: from a state, produce the successor
--   state and an output word.
splitMix :: Word64 -> (Word64, Word64)
splitMix :: Word64 -> (Word64, Word64)
splitMix !Word64
s =
  let !s' :: Word64
s' = Word64
s Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
0x9E3779B97F4A7C15
      !z0 :: Word64
z0 = (Word64
s' Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`B.xor` (Word64
s' Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`B.shiftR` Int
30)) Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
0xBF58476D1CE4E5B9
      !z1 :: Word64
z1 = (Word64
z0 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`B.xor` (Word64
z0 Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`B.shiftR` Int
27)) Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
0x94D049BB133111EB
      !z2 :: Word64
z2 =  Word64
z1 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`B.xor` (Word64
z1 Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`B.shiftR` Int
31)
  in  (Word64
s', Word64
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 :: Rng -> Int -> IO ByteString
randomBytes !Rng
g !Int
n
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0    = ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
BS.empty
  | Bool
otherwise = Int -> (Ptr Word8 -> IO ()) -> IO ByteString
BI.create Int
n ((Ptr Word8 -> IO ()) -> IO ByteString)
-> (Ptr Word8 -> IO ()) -> IO ByteString
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
p -> Rng -> Ptr Word8 -> Int -> IO ()
fillRandom Rng
g Ptr Word8
p Int
n

-- | Fill @n@ bytes at the pointer with pseudorandom bytes drawn
--   from the generator.
fillRandom :: Rng -> Ptr Word8 -> Int -> IO ()
fillRandom :: Rng -> Ptr Word8 -> Int -> IO ()
fillRandom !Rng
g !Ptr Word8
dst !Int
n = Int -> IO ()
go Int
0
  where
    go :: Int -> IO ()
go !Int
i
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n    = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      | Bool
otherwise = do
          !w <- Rng -> IO Word64
nextWord Rng
g
          let !left  = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i
              !chunk = if Int
left Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
8 then Int
8 else Int
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 :: Ptr Word8 -> Int -> Word64 -> Int -> IO ()
writeWord !Ptr Word8
dst !Int
off !Word64
w !Int
lim = Int -> IO ()
go Int
0
  where
    go :: Int -> IO ()
go !Int
j
      | Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
lim  = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      | Bool
otherwise = do
          let !byte :: Word8
byte = Word64 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64
w Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`B.shiftR` (Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
j)) :: Word8
          Ptr Word8 -> Int -> Word8 -> IO ()
forall b. Ptr b -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Word8
dst (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j) Word8
byte
          Int -> IO ()
go (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)