<commit_before>{-# LANGUAGE BangPatterns, CPP, MagicHash, ScopedTypeVariables #-}
-- |
-- Module    : Statistics.RandomVariate
-- Copyright : (c) 2009 <NAME>
-- License   : BSD3
--
-- Maintainer  : <EMAIL>
-- Stability   : experimental
-- Portability : portable
--
-- Random variate generation.

module Statistics.RandomVariate
    (
     Gen
    , Variate(..)
    , create
    , initialize
    ) where

#if defined(__GLASGOW_HASKELL__) && !defined(__HADDOCK__)
#include "MachDeps.h"
#endif

import Control.Monad.ST -- (ST)
import Data.Array.Vector
import Data.Bits ((.&.), (.|.))
import Data.Int (Int8, Int16, Int32, Int64)
import Data.Word (Word, Word8, Word16, Word32, Word64)
import GHC.Base (Int(I#))
import GHC.Word (Word64(W64#), uncheckedShiftL64#, uncheckedShiftRL64#)

class Variate a where
    uniform :: Gen s -> ST s a

instance Variate Int where
#if WORD_SIZE_IN_BITS < 64
    uniform = uniform1 fromIntegral
#else
    uniform = uniform2 wordsTo64Bit
#endif
    {-# INLINE uniform #-}

instance Variate Int8 where
    uniform = uniform1 fromIntegral
    {-# INLINE uniform #-}

instance Variate Int16 where
    uniform = uniform1 fromIntegral
    {-# INLINE uniform #-}

instance Variate Int32 where
    uniform = uniform1 fromIntegral
    {-# INLINE uniform #-}

instance Variate Int64 where
    uniform = uniform2 wordsTo64Bit
    {-# INLINE uniform #-}

instance Variate Word where
#if WORD_SIZE_IN_BITS < 64
    uniform = uniform1 fromIntegral
#else
    uniform = uniform2 wordsTo64Bit
#endif
    {-# INLINE uniform #-}

instance Variate Word8 where
    uniform = uniform1 fromIntegral
    {-# INLINE uniform #-}

instance Variate Word16 where
    uniform = uniform1 fromIntegral
    {-# INLINE uniform #-}

instance Variate Word32 where
    uniform = uniformWord32
    {-# INLINE uniform #-}

instance Variate Word64 where
    uniform = uniform2 wordsTo64Bit
    {-# INLINE uniform #-}

wordsTo64Bit :: Integral a => Word32 -> Word32 -> a
wordsTo64Bit a b =
    fromIntegral ((fromIntegral a `shiftL` 32) .|. fromIntegral b)
{-# INLINE wordsTo64Bit #-}

instance Variate Bool where
    uniform = uniform1 f
        where f i = (i .&. 1) /= 0
    {-# INLINE uniform #-}

-- | 2**(-32)
m_inv_32 :: Fractional t => t
m_inv_32 = 2.3283064365386962890625e-10

wordToFloat :: Word32 -> Float
wordToFloat x = (fromIntegral i * m_inv_32) + 0.5 + m_inv_33
    where m_inv_33 = 1.16415321826934814453125e-10
          i = fromIntegral x :: Int32
{-# INLINE wordToFloat #-}

wordsToDouble :: Word32 -> Word32 -> Double
wordsToDouble x y = (fromIntegral a * m_inv_32 + (0.5 + m_inv_53) +
                    fromIntegral (b .&. 0xFFFFF) * m_inv_52) 
    where m_inv_52 = 2.220446049250313080847263336181640625e-16
          m_inv_53 = 1.1102230246251565404236316680908203125e-16
          a = fromIntegral x :: Int32
          b = fromIntegral y :: Int32
{-# INLINE wordsToDouble #-}

instance Variate Float where
    uniform = uniform1 wordToFloat
    {-# INLINE uniform #-}

instance Variate Double where
    uniform = uniform2 wordsToDouble
    {-# INLINE uniform #-}

-- | State of the pseudo-random number generator.
newtype Gen s = Gen (MUArr Word32 s)

ioff, coff :: Int
ioff = 256
coff = 257

create :: ST s (Gen s)
create = initialize defaultSeed
{-# INLINE create #-}

initialize :: UArr Word32 -> ST s (Gen s)
initialize seed = do
    q <- newMU 258
    fill q
    writeMU q ioff 255
    writeMU q coff 362436
    return (Gen q)
  where fill q = go 0 where
          go i | i == 256  = return ()
               | otherwise = writeMU q i (indexU s i) >> go (i+1)
            where s | i >= fini = defaultSeed
                    | otherwise = seed
        fini = lengthU seed
                               
-- | Unchecked 64-bit left shift.
shiftL :: Word64 -> Int -> Word64
shiftL (W64# x#) (I# i#) = W64# (x# `uncheckedShiftL64#` i#)

-- | Unchecked 64-bit right shift.
shiftR :: Word64 -> Int -> Word64
shiftR (W64# x#) (I# i#) = W64# (x# `uncheckedShiftRL64#` i#)

uniformWord32 :: Gen s -> ST s Word32
uniformWord32 (Gen q) = do
  let a = 809430660 :: Word64
  i <- (fromIntegral . (.&. 255) . succ) `fmap` readMU q ioff
  c <- fromIntegral `fmap` readMU q coff
  qi <- fromIntegral `fmap` readMU q i
  let t   = a * qi + c
      t32 = fromIntegral t
  writeMU q i t32
  writeMU q ioff (fromIntegral i)
  writeMU q coff (fromIntegral (t `shiftR` 32))
  return t32

uniform1 :: (Word32 -> a) -> Gen s -> ST s a
uniform1 f gen = do
  i <- uniformWord32 gen
  return $ f i
{-# INLINE uniform1 #-}

uniform2 :: (Word32 -> Word32 -> a) -> Gen s -> ST s a
uniform2 f (Gen q) = do
  let a = 809430660 :: Word64
  i <- (fromIntegral . (.&. 255) . succ) `fmap` readMU q ioff
  let j = (i + 1) .&. 255
  c <- fromIntegral `fmap` readMU q coff
  qi <- fromIntegral `fmap` readMU q i
  qj <- fromIntegral `fmap` readMU q j
  let t   = a * qi + c
      t32 = fromIntegral t
      c'  = t `shiftR` 32
      u   = a * qj + c'
      u32 = fromIntegral u
  writeMU q i t32
  writeMU q j u32
  writeMU q ioff (fromIntegral j)
  writeMU q coff (fromIntegral (u `shiftR` 32))
  return $! f t32 u32
{-# INLINE uniform2 #-}

uniformMU :: UA e => (Word32 -> e) -> Gen s -> MUArr e s -> ST s ()
uniformMU f (Gen q) mu = do
    i <- readMU q ioff
    c <- readMU q coff
    loop 0 i c
  where
    loop !k !j !c | k == n    = writeMU q ioff j >> writeMU q coff c
                  | otherwise = do
      let i = fromIntegral ((j+1) .&. 255)
      qi <- fromIntegral `fmap` readMU q i
      let t   = a * qi + fromIntegral c
          t32 = fromIntegral t
      writeMU q i t32
      writeMU mu k (f t32)<commit_msg>Small performance improvements.<commit_after>{-# LANGUAGE BangPatterns, CPP, MagicHash, ScopedTypeVariables #-}
-- |
-- Module    : Statistics.RandomVariate
-- Copyright : (c) 2009 <NAME>
-- License   : BSD3
--
-- Maintainer  : <EMAIL>
-- Stability   : experimental
-- Portability : portable
--
-- Random variate generation.

module Statistics.RandomVariate
    (
     Gen
    , Variate(..)
    , create
    , initialize
    ) where

#if defined(__GLASGOW_HASKELL__) && !defined(__HADDOCK__)
#include "MachDeps.h"
#endif

import Control.Monad.ST -- (ST)
import Data.Array.Vector
import Data.Bits ((.&.), (.|.))
import Data.Int (Int8, Int16, Int32, Int64)
import Data.Word (Word, Word8, Word16, Word32, Word64)
import GHC.Base (Int(I#))
import GHC.Word (Word64(W64#), uncheckedShiftL64#, uncheckedShiftRL64#)

class Variate a where
    uniform :: Gen s -> ST s a

instance Variate Int where
#if WORD_SIZE_IN_BITS < 64
    uniform = uniform1 fromIntegral
#else
    uniform = uniform2 wordsTo64Bit
#endif
    {-# INLINE uniform #-}

instance Variate Int8 where
    uniform = uniform1 fromIntegral
    {-# INLINE uniform #-}

instance Variate Int16 where
    uniform = uniform1 fromIntegral
    {-# INLINE uniform #-}

instance Variate Int32 where
    uniform = uniform1 fromIntegral
    {-# INLINE uniform #-}

instance Variate Int64 where
    uniform = uniform2 wordsTo64Bit
    {-# INLINE uniform #-}

instance Variate Word where
#if WORD_SIZE_IN_BITS < 64
    uniform = uniform1 fromIntegral
#else
    uniform = uniform2 wordsTo64Bit
#endif
    {-# INLINE uniform #-}

instance Variate Word8 where
    uniform = uniform1 fromIntegral
    {-# INLINE uniform #-}

instance Variate Word16 where
    uniform = uniform1 fromIntegral
    {-# INLINE uniform #-}

instance Variate Word32 where
    uniform = uniformWord32
    {-# INLINE uniform #-}

instance Variate Word64 where
    uniform = uniform2 wordsTo64Bit
    {-# INLINE uniform #-}

wordsTo64Bit :: Integral a => Word32 -> Word32 -> a
wordsTo64Bit a b =
    fromIntegral ((fromIntegral a `shiftL` 32) .|. fromIntegral b)
{-# INLINE wordsTo64Bit #-}

instance Variate Bool where
    uniform = uniform1 f
        where f i = (i .&. 1) /= 0
    {-# INLINE uniform #-}

wordToFloat :: Word32 -> Float
wordToFloat x = (fromIntegral i * m_inv_32) + 0.5 + m_inv_33
    where m_inv_33 = 1.16415321826934814453125e-10
          m_inv_32 = 2.3283064365386962890625e-10
          i = fromIntegral x :: Int32
{-# INLINE wordToFloat #-}

wordsToDouble :: Word32 -> Word32 -> Double
wordsToDouble x y = (fromIntegral a * m_inv_32 + (0.5 + m_inv_53) +
                    fromIntegral (b .&. 0xFFFFF) * m_inv_52) 
    where m_inv_52 = 2.220446049250313080847263336181640625e-16
          m_inv_53 = 1.1102230246251565404236316680908203125e-16
          m_inv_32 = 2.3283064365386962890625e-10
          a = fromIntegral x :: Int32
          b = fromIntegral y :: Int32
{-# INLINE wordsToDouble #-}

instance Variate Float where
    uniform = uniform1 wordToFloat
    {-# INLINE uniform #-}

instance Variate Double where
    uniform = uniform2 wordsToDouble
    {-# INLINE uniform #-}

-- | State of the pseudo-random number generator.
newtype Gen s = Gen (MUArr Word32 s)

ioff, coff :: Int
ioff = 256
coff = 257

create :: ST s (Gen s)
create = initialize defaultSeed
{-# INLINE create #-}

initialize :: UArr Word32 -> ST s (Gen s)
initialize seed = do
    q <- newMU 258
    fill q
    writeMU q ioff 255
    writeMU q coff 362436
    return (Gen q)
  where fill q = go 0 where
          go i | i == 256  = return ()
               | otherwise = writeMU q i (indexU s i) >> go (i+1)
            where s | i >= fini = defaultSeed
                    | otherwise = seed
        fini = lengthU seed
{-# INLINE initialize #-}
                               
-- | Unchecked 64-bit left shift.
shiftL :: Word64 -> Int -> Word64
shiftL (W64# x#) (I# i#) = W64# (x# `uncheckedShiftL64#` i#)

-- | Unchecked 64-bit right shift.
shiftR :: Word64 -> Int -> Word64
shiftR (W64# x#) (I# i#) = W64# (x# `uncheckedShiftRL64#` i#)

uniformWord32 :: Gen s -> ST s Word32
uniformWord32 (Gen q) = do
  let a = 809430660 :: Word64
  i <- (fromIntegral . (.&. 255) . (+1)) `fmap` readMU q ioff
  c <- fromIntegral `fmap` readMU q coff
  qi <- fromIntegral `fmap` readMU q i
  let t   = a * qi + c
      t32 = fromIntegral t
  writeMU q i t32
  writeMU q ioff (fromIntegral i)
  writeMU q coff (fromIntegral (t `shiftR` 32))
  return t32

uniform1 :: (Word32 -> a) -> Gen s -> ST s a
uniform1 f gen = do
  i <- uniformWord32 gen
  return $! f i
{-# INLINE uniform1 #-}

uniform2 :: (Word32 -> Word32 -> a) -> Gen s -> ST s a
uniform2 f (Gen q) = do
  let a = 809430660 :: Word64
  i <- (fromIntegral . (.&. 255) . (+1)) `fmap` readMU q ioff
  let j = (i + 1) .&. 255
  c <- fromIntegral `fmap` readMU q coff
  qi <- fromIntegral `fmap` readMU q i
  qj <- fromIntegral `fmap` readMU q j
  let t   = a * qi + c
      t32 = fromIntegral t
      c'  = t `shiftR` 32
      u   = a * qj + c'
      u32 = fromIntegral u
  writeMU q i t32
  writeMU q j u32
  writeMU q ioff (fromIntegral j)
  writeMU q coff (fromIntegral (u `shiftR` 32))
  return $! f t32 u32
{-# INLINE uniform2 #-}

uniformMU :: UA e => (Word32 -> e) -> Gen s -> MUArr e s -> ST s ()
uniformMU f (Gen q) mu = do
    i <- readMU q ioff
    c <- readMU q coff
    loop 0 i c
  where
    loop !k !j !c | k == n    = writeMU q ioff j >> writeMU q coff c
                  | otherwise = do
      let i = fromIntegral ((j+1) .&. 255)
      qi <- fromIntegral `fmap` readMU q i
      let t   = a * qi + fromIntegral c
          t32 = fromIntegral t
      writeMU q i t32
      writeMU mu k (f t32)
{'fraction_non_alphanumeric': 0.11476426799007444, 'fraction_numerical': 0.07178305565402339, 'mean_word_length': 2.515576323987539, 'pattern_counts': {'":': 0, '<': 35, '<?xml version=': 0, '>': 68, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1046353', 'n_tokens_mistral': 4611, 'n_tokens_neox': 4011, 'n_words': 1314}