Xoshiro.hs

{-# OPTIONS_HADDOCK hide #-}
{-# LANGUAGE BangPatterns #-}

-- |
-- Module: Crypto.Xoshiro
-- Copyright: (c) 2026 Jared Tobin
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
-- The xoshiro256++ state, step function, scrambler, jumps, and
-- seeding, shared by the mutable and pure generators.

module Crypto.Xoshiro (
    S(..)
  , output
  , next
  , jump
  , long_jump
  , seed
  , from_words
  , system
  ) where

import qualified Crypto.Entropy as Entropy
import Data.Bits ((.|.), (.^.), (.<<.), (.>>.))
import qualified Data.Bits as B
import Data.Word (Word64)
import Foreign.Marshal.Alloc (allocaBytes)
import Foreign.Ptr (Ptr, castPtr)
import Foreign.Storable (peekElemOff)

rotl :: Word64 -> Int -> Word64
rotl x k = (x .<<. k) .|. (x .>>. (64 - k))
{-# INLINE rotl #-}

-- | A xoshiro256++ state, which must not be all zero.
data S = S
  {-# UNPACK #-} !Word64
  {-# UNPACK #-} !Word64
  {-# UNPACK #-} !Word64
  {-# UNPACK #-} !Word64

-- | The ++ scrambler: the output for a state.
output :: S -> Word64
output (S s0 _ _ s3) = rotl (s0 + s3) 23 + s0
{-# INLINE output #-}

-- | The linear engine: the next state.
next :: S -> S
next (S s0 s1 s2 s3) =
  let !t   = s1 .<<. 17
      !a2  = s2 .^. s0
      !a3  = s3 .^. s1
      !b1  = s1 .^. a2
      !b0  = s0 .^. a3
      !b2  = a2 .^. t
      !b3  = rotl a3 45
  in  S b0 b1 b2 b3
{-# INLINE next #-}

xor_s :: S -> S -> S
xor_s (S a0 a1 a2 a3) (S b0 b1 b2 b3) =
  S (a0 .^. b0) (a1 .^. b1) (a2 .^. b2) (a3 .^. b3)
{-# INLINE xor_s #-}

data P = P !S !S

-- advance the state by the jump polynomial with the given coefficients
leap :: Word64 -> Word64 -> Word64 -> Word64 -> S -> S
leap j0 j1 j2 j3 s0 =
  let !(P acc _) =
        coef j3 (coef j2 (coef j1 (coef j0 (P (S 0 0 0 0) s0))))
  in  acc
  where
    coef :: Word64 -> P -> P
    coef j = loop 0 where
      loop :: Int -> P -> P
      loop !b p@(P acc s)
        | b == 64   = p
        | otherwise =
            let !nacc | B.testBit j b = xor_s acc s
                      | otherwise     = acc
            in  loop (b + 1) (P nacc (next s))

-- | Advance a state by 2 ^ 128 steps.
jump :: S -> S
jump = leap 0x180ec6d33cfd0aba 0xd5a61266f0c9392c
            0xa9582618e03fc9aa 0x39abdc4529b1661c

-- | Advance a state by 2 ^ 192 steps.
long_jump :: S -> S
long_jump = leap 0x76e15d3efefdcbbf 0xc5004e441c522fb3
                 0x77710069854ee241 0x39109bb02acbe635

data SM = SM {-# UNPACK #-} !Word64 {-# UNPACK #-} !Word64

splitmix64 :: Word64 -> SM
splitmix64 x =
  let !nx = x + 0x9e3779b97f4a7c15
      !z0 = (nx .^. (nx .>>. 30)) * 0xbf58476d1ce4e5b9
      !z1 = (z0 .^. (z0 .>>. 27)) * 0x94d049bb133111eb
  in  SM nx (z1 .^. (z1 .>>. 31))
{-# INLINE splitmix64 #-}

-- | A state expanded from a 64-bit seed via splitmix64, per the
--   xoshiro authors' recommendation. It is never all zero.
seed :: Word64 -> S
seed x =
  let !(SM x0 s0) = splitmix64 x
      !(SM x1 s1) = splitmix64 x0
      !(SM x2 s2) = splitmix64 x1
      !(SM _  s3) = splitmix64 x2
  in  S s0 s1 s2 s3

-- | A state from the given words, unless they are all zero.
from_words :: Word64 -> Word64 -> Word64 -> Word64 -> Maybe S
from_words s0 s1 s2 s3
  | s0 .|. s1 .|. s2 .|. s3 == 0 = Nothing
  | otherwise = Just (S s0 s1 s2 s3)
{-# INLINE from_words #-}

-- | A state drawn from the operating system's entropy source, unless
--   that fails.
system :: IO (Maybe S)
system = allocaBytes 32 $ \buf -> do
  ok <- Entropy.fill buf 32
  if   not ok
  then pure Nothing
  else do
    let p = castPtr buf :: Ptr Word64
    s0 <- peekElemOff p 0
    s1 <- peekElemOff p 1
    s2 <- peekElemOff p 2
    s3 <- peekElemOff p 3
    pure $! from_words s0 s1 s2 s3