-- |A few little functions I found myself writing inline over and over again.
module Data.Random.Internal.Words where

import Data.Bits
import Data.Word
import Foreign.Marshal  (allocaBytes)
import Foreign.Ptr      (castPtr)
import Foreign.Storable (peek, pokeByteOff)
import System.IO.Unsafe (unsafePerformIO)

-- TODO: add a build flag for endianness-invariance, or just find a way
-- to make sure these operations all do the right thing without costing 
-- anything extra at runtime

{-# INLINE buildWord16 #-}
-- |Build a word out of 2 bytes.  No promises are made regarding the order
-- in which the bytes are stuffed.  Note that this means that a 'RandomSource'
-- or 'MonadRandom' making use of the default definition of 'getRandomWord', etc.,
-- may return different random values on different platforms when started 
-- with the same seed, depending on the platform's endianness.
buildWord16 :: Word8 -> Word8 -> Word16
buildWord16 :: Word8 -> Word8 -> Word16
buildWord16 Word8
b0 Word8
b1
    = IO Word16 -> Word16
forall a. IO a -> a
unsafePerformIO (IO Word16 -> Word16)
-> ((Ptr Any -> IO Word16) -> IO Word16)
-> (Ptr Any -> IO Word16)
-> Word16
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> (Ptr Any -> IO Word16) -> IO Word16
forall a b. Int -> (Ptr a -> IO b) -> IO b
allocaBytes Int
2 ((Ptr Any -> IO Word16) -> Word16)
-> (Ptr Any -> IO Word16) -> Word16
forall a b. (a -> b) -> a -> b
$ \Ptr Any
p -> do
        Ptr Any -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Any
p Int
0 Word8
b0
        Ptr Any -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Any
p Int
1 Word8
b1
        Ptr Word16 -> IO Word16
forall a. Storable a => Ptr a -> IO a
peek (Ptr Any -> Ptr Word16
forall a b. Ptr a -> Ptr b
castPtr Ptr Any
p)

{-# INLINE buildWord32 #-}
-- |Build a word out of 4 bytes.  No promises are made regarding the order
-- in which the bytes are stuffed.  Note that this means that a 'RandomSource'
-- or 'MonadRandom' making use of the default definition of 'getRandomWord', etc.,
-- may return different random values on different platforms when started 
-- with the same seed, depending on the platform's endianness.
buildWord32 :: Word8 -> Word8 -> Word8 -> Word8 -> Word32
buildWord32 :: Word8 -> Word8 -> Word8 -> Word8 -> Word32
buildWord32 Word8
b0 Word8
b1 Word8
b2 Word8
b3
    = IO Word32 -> Word32
forall a. IO a -> a
unsafePerformIO (IO Word32 -> Word32)
-> ((Ptr Any -> IO Word32) -> IO Word32)
-> (Ptr Any -> IO Word32)
-> Word32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> (Ptr Any -> IO Word32) -> IO Word32
forall a b. Int -> (Ptr a -> IO b) -> IO b
allocaBytes Int
4 ((Ptr Any -> IO Word32) -> Word32)
-> (Ptr Any -> IO Word32) -> Word32
forall a b. (a -> b) -> a -> b
$ \Ptr Any
p -> do
        Ptr Any -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Any
p Int
0 Word8
b0
        Ptr Any -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Any
p Int
1 Word8
b1
        Ptr Any -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Any
p Int
2 Word8
b2
        Ptr Any -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Any
p Int
3 Word8
b3
        Ptr Word32 -> IO Word32
forall a. Storable a => Ptr a -> IO a
peek (Ptr Any -> Ptr Word32
forall a b. Ptr a -> Ptr b
castPtr Ptr Any
p)

{-# INLINE buildWord32' #-}
buildWord32' :: Word16 -> Word16 -> Word32
buildWord32' :: Word16 -> Word16 -> Word32
buildWord32' Word16
w0 Word16
w1
    = IO Word32 -> Word32
forall a. IO a -> a
unsafePerformIO (IO Word32 -> Word32)
-> ((Ptr Any -> IO Word32) -> IO Word32)
-> (Ptr Any -> IO Word32)
-> Word32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> (Ptr Any -> IO Word32) -> IO Word32
forall a b. Int -> (Ptr a -> IO b) -> IO b
allocaBytes Int
4 ((Ptr Any -> IO Word32) -> Word32)
-> (Ptr Any -> IO Word32) -> Word32
forall a b. (a -> b) -> a -> b
$ \Ptr Any
p -> do
        Ptr Any -> Int -> Word16 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Any
p Int
0 Word16
w0
        Ptr Any -> Int -> Word16 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Any
p Int
2 Word16
w1
        Ptr Word32 -> IO Word32
forall a. Storable a => Ptr a -> IO a
peek (Ptr Any -> Ptr Word32
forall a b. Ptr a -> Ptr b
castPtr Ptr Any
p)

{-# INLINE buildWord64 #-}
-- |Build a word out of 8 bytes.  No promises are made regarding the order
-- in which the bytes are stuffed.  Note that this means that a 'RandomSource'
-- or 'MonadRandom' making use of the default definition of 'getRandomWord', etc.,
-- may return different random values on different platforms when started 
-- with the same seed, depending on the platform's endianness.
buildWord64 :: Word8 -> Word8 -> Word8 -> Word8 -> Word8 -> Word8 -> Word8 -> Word8 -> Word64
buildWord64 :: Word8
-> Word8
-> Word8
-> Word8
-> Word8
-> Word8
-> Word8
-> Word8
-> Word64
buildWord64 Word8
b0 Word8
b1 Word8
b2 Word8
b3 Word8
b4 Word8
b5 Word8
b6 Word8
b7
    = IO Word64 -> Word64
forall a. IO a -> a
unsafePerformIO (IO Word64 -> Word64)
-> ((Ptr Any -> IO Word64) -> IO Word64)
-> (Ptr Any -> IO Word64)
-> Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> (Ptr Any -> IO Word64) -> IO Word64
forall a b. Int -> (Ptr a -> IO b) -> IO b
allocaBytes Int
8 ((Ptr Any -> IO Word64) -> Word64)
-> (Ptr Any -> IO Word64) -> Word64
forall a b. (a -> b) -> a -> b
$ \Ptr Any
p -> do
        Ptr Any -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Any
p Int
0 Word8
b0
        Ptr Any -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Any
p Int
1 Word8
b1
        Ptr Any -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Any
p Int
2 Word8
b2
        Ptr Any -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Any
p Int
3 Word8
b3
        Ptr Any -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Any
p Int
4 Word8
b4
        Ptr Any -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Any
p Int
5 Word8
b5
        Ptr Any -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Any
p Int
6 Word8
b6
        Ptr Any -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Any
p Int
7 Word8
b7
        Ptr Word64 -> IO Word64
forall a. Storable a => Ptr a -> IO a
peek (Ptr Any -> Ptr Word64
forall a b. Ptr a -> Ptr b
castPtr Ptr Any
p)

{-# INLINE buildWord64' #-}
buildWord64' :: Word16 -> Word16 -> Word16 -> Word16 -> Word64
buildWord64' :: Word16 -> Word16 -> Word16 -> Word16 -> Word64
buildWord64' Word16
w0 Word16
w1 Word16
w2 Word16
w3
    = IO Word64 -> Word64
forall a. IO a -> a
unsafePerformIO (IO Word64 -> Word64)
-> ((Ptr Any -> IO Word64) -> IO Word64)
-> (Ptr Any -> IO Word64)
-> Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> (Ptr Any -> IO Word64) -> IO Word64
forall a b. Int -> (Ptr a -> IO b) -> IO b
allocaBytes Int
8 ((Ptr Any -> IO Word64) -> Word64)
-> (Ptr Any -> IO Word64) -> Word64
forall a b. (a -> b) -> a -> b
$ \Ptr Any
p -> do
        Ptr Any -> Int -> Word16 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Any
p Int
0 Word16
w0
        Ptr Any -> Int -> Word16 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Any
p Int
2 Word16
w1
        Ptr Any -> Int -> Word16 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Any
p Int
4 Word16
w2
        Ptr Any -> Int -> Word16 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Any
p Int
6 Word16
w3
        Ptr Word64 -> IO Word64
forall a. Storable a => Ptr a -> IO a
peek (Ptr Any -> Ptr Word64
forall a b. Ptr a -> Ptr b
castPtr Ptr Any
p)

{-# INLINE buildWord64'' #-}
buildWord64'' :: Word32 -> Word32 -> Word64
buildWord64'' :: Word32 -> Word32 -> Word64
buildWord64'' Word32
w0 Word32
w1
    = IO Word64 -> Word64
forall a. IO a -> a
unsafePerformIO (IO Word64 -> Word64)
-> ((Ptr Any -> IO Word64) -> IO Word64)
-> (Ptr Any -> IO Word64)
-> Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> (Ptr Any -> IO Word64) -> IO Word64
forall a b. Int -> (Ptr a -> IO b) -> IO b
allocaBytes Int
8 ((Ptr Any -> IO Word64) -> Word64)
-> (Ptr Any -> IO Word64) -> Word64
forall a b. (a -> b) -> a -> b
$ \Ptr Any
p -> do
        Ptr Any -> Int -> Word32 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Any
p Int
0 Word32
w0
        Ptr Any -> Int -> Word32 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Any
p Int
4 Word32
w1
        Ptr Word64 -> IO Word64
forall a. Storable a => Ptr a -> IO a
peek (Ptr Any -> Ptr Word64
forall a b. Ptr a -> Ptr b
castPtr Ptr Any
p)

{-# INLINE word32ToFloat #-}
-- |Pack the low 23 bits from a 'Word32' into a 'Float' in the range [0,1).
-- Used to convert a 'stdUniform' 'Word32' to a 'stdUniform' 'Double'.
word32ToFloat :: Word32 -> Float
word32ToFloat :: Word32 -> Float
word32ToFloat Word32
x = (Integer -> Int -> Float
forall a. RealFloat a => Integer -> Int -> a
encodeFloat (Integer -> Int -> Float) -> Integer -> Int -> Float
forall a b. (a -> b) -> a -> b
$! Word32 -> Integer
forall a. Integral a => a -> Integer
toInteger (Word32
x Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.&. Word32
0x007fffff {- 2^23-1 -} )) (Int -> Float) -> Int -> Float
forall a b. (a -> b) -> a -> b
$ (-Int
23)

{-# INLINE word32ToFloatWithExcess #-}
-- |Same as word32ToFloat, but also return the unused bits (as the 9
-- least significant bits of a 'Word32')
word32ToFloatWithExcess :: Word32 -> (Float, Word32)
word32ToFloatWithExcess :: Word32 -> (Float, Word32)
word32ToFloatWithExcess Word32
x = (Word32 -> Float
word32ToFloat Word32
x, Word32
x Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`shiftR` Int
23)

{-# INLINE wordToFloat #-}
-- |Pack the low 23 bits from a 'Word64' into a 'Float' in the range [0,1).
-- Used to convert a 'stdUniform' 'Word64' to a 'stdUniform' 'Double'.
wordToFloat :: Word64 -> Float
wordToFloat :: Word64 -> Float
wordToFloat Word64
x = (Integer -> Int -> Float
forall a. RealFloat a => Integer -> Int -> a
encodeFloat (Integer -> Int -> Float) -> Integer -> Int -> Float
forall a b. (a -> b) -> a -> b
$! Word64 -> Integer
forall a. Integral a => a -> Integer
toInteger (Word64
x Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. Word64
0x007fffff {- 2^23-1 -} )) (Int -> Float) -> Int -> Float
forall a b. (a -> b) -> a -> b
$ (-Int
23)

{-# INLINE wordToFloatWithExcess #-}
-- |Same as wordToFloat, but also return the unused bits (as the 41
-- least significant bits of a 'Word64')
wordToFloatWithExcess :: Word64 -> (Float, Word64)
wordToFloatWithExcess :: Word64 -> (Float, Word64)
wordToFloatWithExcess Word64
x = (Word64 -> Float
wordToFloat Word64
x, Word64
x Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
23)

{-# INLINE wordToDouble #-}
-- |Pack the low 52 bits from a 'Word64' into a 'Double' in the range [0,1).
-- Used to convert a 'stdUniform' 'Word64' to a 'stdUniform' 'Double'.
wordToDouble :: Word64 -> Double
wordToDouble :: Word64 -> Double
wordToDouble Word64
x = (Integer -> Int -> Double
forall a. RealFloat a => Integer -> Int -> a
encodeFloat (Integer -> Int -> Double) -> Integer -> Int -> Double
forall a b. (a -> b) -> a -> b
$! Word64 -> Integer
forall a. Integral a => a -> Integer
toInteger (Word64
x Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. Word64
0x000fffffffffffff {- 2^52-1 -})) (Int -> Double) -> Int -> Double
forall a b. (a -> b) -> a -> b
$ (-Int
52)

{-# INLINE word32ToDouble #-}
-- |Pack a 'Word32' into a 'Double' in the range [0,1).  Note that a Double's 
-- mantissa is 52 bits, so this does not fill all of them.
word32ToDouble :: Word32 -> Double
word32ToDouble :: Word32 -> Double
word32ToDouble Word32
x = (Integer -> Int -> Double
forall a. RealFloat a => Integer -> Int -> a
encodeFloat (Integer -> Int -> Double) -> Integer -> Int -> Double
forall a b. (a -> b) -> a -> b
$! Word32 -> Integer
forall a. Integral a => a -> Integer
toInteger Word32
x) (Int -> Double) -> Int -> Double
forall a b. (a -> b) -> a -> b
$ (-Int
32)

{-# INLINE wordToDoubleWithExcess #-}
-- |Same as wordToDouble, but also return the unused bits (as the 12
-- least significant bits of a 'Word64')
wordToDoubleWithExcess :: Word64 -> (Double, Word64)
wordToDoubleWithExcess :: Word64 -> (Double, Word64)
wordToDoubleWithExcess Word64
x = (Word64 -> Double
wordToDouble Word64
x, Word64
x Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
52)