SplitMix.hs

-- | An implementation of SplitMix: <https://doi.org/10.1145/2660193.2660195>
module Mini.Random.SplitMix (
  -- * Type
  SplitMix,

  -- * Construction
  seed,

  -- * Operations
  nextWord32,
  nextWord64,
  split,
) where

import Data.Bits (
  popCount,
  shiftR,
  xor,
  (.|.),
 )
import Data.Word (
  Word32,
  Word64,
 )
import Mini.Data.Recursion (
  bool,
 )
import Prelude (
  Eq,
  Ord,
  Show,
  fromIntegral,
  ($),
  (*),
  (+),
  (<),
 )

-- Type

-- | Abstract representation of a SplitMix generator
data SplitMix = SplitMix Word64 Word64
  deriving (Eq, Ord, Show)

-- Construction

-- | Make a generator from a seed
seed :: Word64 -> SplitMix
seed s = SplitMix (mix64 s) (mixGamma $ s + 0x9e3779b97f4a7c15)

-- Operations

-- | Generate the next 32-bit word
nextWord32 :: SplitMix -> (Word32, SplitMix)
nextWord32 (SplitMix s g) =
  let s' = s + g
   in (mix32 s', SplitMix s' g)

-- | Generate the next 64-bit word
nextWord64 :: SplitMix -> (Word64, SplitMix)
nextWord64 (SplitMix s g) =
  let s' = s + g
   in (mix64 s', SplitMix s' g)

-- | Split a generator into two (seemingly) independent generators
split :: SplitMix -> (SplitMix, SplitMix)
split (SplitMix s g) =
  let s' = s + g
      s'' = s' + g
      sm = SplitMix s'' g
      sm' = SplitMix (mix64 s') (mixGamma s'')
   in (sm, sm')

-- Helpers

mix64 :: Word64 -> Word64
mix64 z0 =
  let z1 = z0 `xor` (z0 `shiftR` 33)
      z2 = z1 * 0xff51afd7ed558ccd
      z3 = z2 `xor` (z2 `shiftR` 33)
      z4 = z3 * 0xc4ceb9fe1a85ec53
      z5 = z4 `xor` (z4 `shiftR` 33)
   in z5

mix32 :: Word64 -> Word32
mix32 z0 =
  let z1 = z0 `xor` (z0 `shiftR` 33)
      z2 = z1 * 0xff51afd7ed558ccd
      z3 = z2 `xor` (z2 `shiftR` 33)
      z4 = z3 * 0xc4ceb9fe1a85ec53
      z5 = z4 `shiftR` 32
   in fromIntegral z5

mix64variant13 :: Word64 -> Word64
mix64variant13 z0 =
  let z1 = z0 `xor` (z0 `shiftR` 30)
      z2 = z1 * 0xbf58476d1ce4e5b9
      z3 = z2 `xor` (z2 `shiftR` 27)
      z4 = z3 * 0x94d049bb133111eb
      z5 = z4 `xor` (z4 `shiftR` 31)
   in z5

mixGamma :: Word64 -> Word64
mixGamma z0 =
  let z1 = mix64variant13 z0 .|. 1
      n = popCount $ z1 `xor` (z1 `shiftR` 1)
   in -- based on the text on p. 466, not the code with inverted logic on p. 465
      bool z1 (z1 `xor` 0xaaaaaaaaaaaaaaaa) $ n < 24