{-# LANGUAGE Rank2Types, DataKinds #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  System.Random.Dice.Internal
-- Copyright   :  Peter Robinson 2014
-- License     :  LGPL
--
-- Maintainer  :  Peter Robinson <peter.robinson@monoid.at>
-- Stability   :  experimental
-- Portability :  portable
--
-----------------------------------------------------------------------------
module System.Random.Dice.Internal
where
import System.Entropy
import Control.Monad.IO.Class
import Control.Monad
import Control.Exception
import qualified Data.ByteString as B
import Data.Word
import Data.Conduit
import qualified Data.Conduit.List as CL

-- | Converts a number to its base-2 representation (as a list of bits)
-- and prepends zeros to ensure the minimal size.
integralToBits :: (Integral n,Integral m)
               => Int  -- ^ minimal number of bits @b@
               -> n    -- ^ the number @n@
               -> [m]  -- ^ bit representation of @n@, length >= @b@
integralToBits :: forall n m. (Integral n, Integral m) => Int -> n -> [m]
integralToBits Int
b n
x = [m] -> [m]
forall a. [a] -> [a]
reverse ([m] -> [m]) -> [m] -> [m]
forall a b. (a -> b) -> a -> b
$ Int -> n -> [m]
forall {t} {a}. (Num a, Integral t) => Int -> t -> [a]
integralToBits' Int
0 n
x
  where
  integralToBits' :: Int -> t -> [a]
integralToBits' Int
ns t
0 = Int -> a -> [a]
forall a. Int -> a -> [a]
replicate (Int
bInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
ns) a
0
  integralToBits' Int
ns t
y =
    let (t
a,t
res) = t -> t -> (t, t)
forall a. Integral a => a -> a -> (a, a)
quotRem t
y t
2 in
    t -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral t
res a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Int -> t -> [a]
integralToBits' (Int
nsInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) t
a


-- | Convert a list of bits to an integral
bitsToIntegral :: (Integral n) =>[n] -> n
bitsToIntegral :: forall n. Integral n => [n] -> n
bitsToIntegral = n -> [n] -> n
forall n. Integral n => n -> [n] -> n
extendIntegralWithBits n
0


extendIntegralWithBits :: (Integral n) => n -> [n] -> n
extendIntegralWithBits :: forall n. Integral n => n -> [n] -> n
extendIntegralWithBits n
n = (n -> n -> n) -> n -> [n] -> n
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\n
c n
r -> n
2n -> n -> n
forall a. Num a => a -> a -> a
*n
r n -> n -> n
forall a. Num a => a -> a -> a
+ n
c) n
n ([n] -> n) -> ([n] -> [n]) -> [n] -> n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [n] -> [n]
forall a. [a] -> [a]
reverse


-- | Upper bound on the number of sides that a random dice can have.
upperBound :: Word64
upperBound :: Word64
upperBound = Word64
2Word64 -> Int -> Word64
forall a b. (Num a, Integral b) => a -> b -> a
^(Int
55 :: Int)


-- | Generates @k@ rolls of an @n@ sided dice.
getDiceRolls :: Int  -- ^ @n:@ number of sides
             -> Int  -- ^ @k:@ number of rolls
             -> IO [Int]
getDiceRolls :: Int -> Int -> IO [Int]
getDiceRolls Int
n Int
len =
  ConduitT () Word8 IO ()
Producer IO Word8
systemEntropy ConduitT () Word8 IO () -> Sink Word8 IO [Int] -> IO [Int]
forall (m :: * -> *) a b.
Monad m =>
Source m a -> Sink a m b -> m b
$$ Int -> Conduit Word8 IO Int
diceRolls Int
n Conduit Word8 IO Int
-> ConduitT Int Void IO [Int] -> Sink Word8 IO [Int]
forall (m :: * -> *) a b c r.
Monad m =>
Conduit a m b -> ConduitT b c m r -> ConduitT a c m r
=$= Int -> ConduitT Int Void IO [Int]
forall (m :: * -> *) a o. Monad m => Int -> ConduitT a o m [a]
CL.take Int
len


-- | Generates a list of random integer values in the specified range.
getRandomRs :: (Int,Int) -- ^ (inclusive) range
         -> Int          -- ^ number of samples
         -> IO [Int]
getRandomRs :: (Int, Int) -> Int -> IO [Int]
getRandomRs (Int, Int)
range Int
len =
  ConduitT () Word8 IO ()
Producer IO Word8
systemEntropy ConduitT () Word8 IO () -> Sink Word8 IO [Int] -> IO [Int]
forall (m :: * -> *) a b.
Monad m =>
Source m a -> Sink a m b -> m b
$$ (Int, Int) -> Conduit Word8 IO Int
randomRs (Int, Int)
range Conduit Word8 IO Int
-> ConduitT Int Void IO [Int] -> Sink Word8 IO [Int]
forall (m :: * -> *) a b c r.
Monad m =>
Conduit a m b -> ConduitT b c m r -> ConduitT a c m r
=$= Int -> ConduitT Int Void IO [Int]
forall (m :: * -> *) a o. Monad m => Int -> ConduitT a o m [a]
CL.take Int
len


-- | Produces a stream of random integer values in the range @[0,n-1]@, for a
-- given @n <= 2^55@.
-- This conduit needs to be attached to an entropy source such as
-- 'systemEntropy'.
diceRolls :: Int -> Conduit Word8 IO Int
diceRolls :: Int -> Conduit Word8 IO Int
diceRolls Int
n
  | Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
> Word64
upperBound Bool -> Bool -> Bool
|| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
    = AssertionFailed -> Conduit Word8 IO Int
forall a e. Exception e => e -> a
throw (AssertionFailed -> Conduit Word8 IO Int)
-> AssertionFailed -> Conduit Word8 IO Int
forall a b. (a -> b) -> a -> b
$ String -> AssertionFailed
AssertionFailed String
"diceRolls: n-sided dice are supported, for 1 <= n < 2^55."
  | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1
    = [Int] -> Conduit Word8 IO Int
forall (m :: * -> *) a i. Monad m => [a] -> ConduitT i a m ()
CL.sourceList [Int
0,Int
0..]
  | Bool
otherwise
    = Word64 -> Word64 -> Word64 -> Int -> Conduit Word8 IO (Int, Int)
dRoll (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) Word64
1 Word64
0 Int
0 Conduit Word8 IO (Int, Int)
-> ConduitT (Int, Int) Int IO () -> Conduit Word8 IO Int
forall (m :: * -> *) a b c r.
Monad m =>
Conduit a m b -> ConduitT b c m r -> ConduitT a c m r
=$= ((Int, Int) -> Int) -> ConduitT (Int, Int) Int IO ()
forall (m :: * -> *) a b. Monad m => (a -> b) -> ConduitT a b m ()
CL.map (Int, Int) -> Int
forall a b. (a, b) -> a
fst


-- | Produces a stream of random integer values within a range.
-- This conduit needs to be attached to an entropy source such as
-- 'systemEntropy'.
randomRs :: (Int,Int)             -- ^ range (inclusive)
         -> Conduit Word8 IO Int
randomRs :: (Int, Int) -> Conduit Word8 IO Int
randomRs (Int
low,Int
up) = Int -> Conduit Word8 IO Int
diceRolls (Int
upInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
lowInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Conduit Word8 IO Int
-> ConduitT Int Int IO () -> Conduit Word8 IO Int
forall (m :: * -> *) a b c r.
Monad m =>
Conduit a m b -> ConduitT b c m r -> ConduitT a c m r
=$= (Int -> Int) -> ConduitT Int Int IO ()
forall (m :: * -> *) a b. Monad m => (a -> b) -> ConduitT a b m ()
CL.map (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
low)


-- | A source of entropy. By default, we use the 'getEntropy' function from
-- the entropy package, see 'systemEntropy'.
--
-- /Warning:/ When combining a source of entropy with other conduits, it is
-- important that there is no \"backflow\" due to leftover values that
-- are being returned to the
-- source from the conduit. This can be done by fusing the conduit with the
-- identity map, e.g: @myEntropySrc $$ Data.Conduit.List.map id =$= myConduit@
--
systemEntropy :: Producer IO Word8
systemEntropy :: Producer IO Word8
systemEntropy = do
  [Word8]
bytes <- ByteString -> [Word8]
B.unpack (ByteString -> [Word8])
-> ConduitT i Word8 IO ByteString -> ConduitT i Word8 IO [Word8]
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` IO ByteString -> ConduitT i Word8 IO ByteString
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO  (Int -> IO ByteString
getEntropy Int
8)
  [Word8]
-> (Word8 -> ConduitT i Word8 IO ()) -> ConduitT i Word8 IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Word8]
bytes Word8 -> ConduitT i Word8 IO ()
forall (m :: * -> *) o i. Monad m => o -> ConduitT i o m ()
yield
  ConduitT i Word8 IO ()
Producer IO Word8
systemEntropy



-- | Internal function. Should not be invoked directly.
dRoll :: Word64 -> Word64 -> Word64 -> Int -> Conduit Word8 IO (Int,Int)
dRoll :: Word64 -> Word64 -> Word64 -> Int -> Conduit Word8 IO (Int, Int)
dRoll Word64
n Word64
m Word64
r Int
cnt = do
  let k :: Int
k = Double -> Int
forall a b. (RealFrac a, Integral b) => a -> b
ceiling (Double -> Int) -> Double -> Int
forall a b. (a -> b) -> a -> b
$ (Double -> Double -> Double
forall a. Floating a => a -> a -> a
logBase Double
2 (Word64 -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
upperBound) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double -> Double
forall a. Floating a => a -> a -> a
logBase Double
2 (Word64 -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
m :: Double)) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
8
  let m' :: Word64
m' = Word64
2Word64 -> Int -> Word64
forall a b. (Num a, Integral b) => a -> b -> a
^(Int
8Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
k) Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
m
  [Word64]
bits <- ((Word8 -> [Word64]) -> [Word8] -> [Word64]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Int -> Word8 -> [Word64]
forall n m. (Integral n, Integral m) => Int -> n -> [m]
integralToBits Int
8) ([Word8] -> [Word64])
-> (ByteString -> [Word8]) -> ByteString -> [Word64]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> [Word8]
B.unpack)
          (ByteString -> [Word64])
-> ConduitT Word8 (Int, Int) IO ByteString
-> ConduitT Word8 (Int, Int) IO [Word64]
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` (if Int
kInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>Int
0 then IO ByteString -> ConduitT Word8 (Int, Int) IO ByteString
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ByteString -> ConduitT Word8 (Int, Int) IO ByteString)
-> IO ByteString -> ConduitT Word8 (Int, Int) IO ByteString
forall a b. (a -> b) -> a -> b
$ Int -> IO ByteString
getEntropy Int
k else ByteString -> ConduitT Word8 (Int, Int) IO ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> ConduitT Word8 (Int, Int) IO ByteString)
-> ByteString -> ConduitT Word8 (Int, Int) IO ByteString
forall a b. (a -> b) -> a -> b
$ [Word8] -> ByteString
B.pack [])
  let w64 :: Word64
w64 = Word64 -> [Word64] -> Word64
forall n. Integral n => n -> [n] -> n
extendIntegralWithBits Word64
r [Word64]
bits
  let q :: Word64
q = Word64
m' Word64 -> Word64 -> Word64
forall a. Integral a => a -> a -> a
`div` Word64
n
  if Word64
w64 Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
n Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
q
    then do
      (Int, Int) -> Conduit Word8 IO (Int, Int)
forall (m :: * -> *) o i. Monad m => o -> ConduitT i o m ()
yield (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64 -> Int) -> Word64 -> Int
forall a b. (a -> b) -> a -> b
$ Word64
w64 Word64 -> Word64 -> Word64
forall a. Integral a => a -> a -> a
`mod` Word64
n,Int
k)
      Word64 -> Word64 -> Word64 -> Int -> Conduit Word8 IO (Int, Int)
dRoll Word64
n Word64
q (Word64
w64 Word64 -> Word64 -> Word64
forall a. Integral a => a -> a -> a
`div` Word64
n) (Int
cnt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k)
    else  Word64 -> Word64 -> Word64 -> Int -> Conduit Word8 IO (Int, Int)
dRoll Word64
n (Word64
m' Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
nWord64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
*Word64
q) (Word64
w64 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
nWord64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
*Word64
q) (Int
cnt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k)

{-
-- | Compute the performance of the algorithm in terms of used random bits
-- versus produced random values.
testPerformance :: Int   -- ^ number of sides of dice
                -> Int   -- ^ number of samples used for computing average.
                -> IO ()
testPerformance n len
  | fromIntegral n > upperBound
    = throw $ AssertionFailed "dice: range must be within Word64 bounds."
  | otherwise = do
    nbits <- systemEntropy $= dRoll (fromIntegral n) 1 0 0
                           $$ CL.take len
                           >>= return . sum . map snd
    putStrLn $ "Generated " ++ show len
            ++ " random samples in range [0," ++ show (n-1) ++ "]"
    putStrLn $ "Average number of bits used: "
            ++ show (8*fromIntegral nbits/ fromIntegral len :: Double)
    let lbound = logBase 2 (fromIntegral n) :: Double
    putStrLn $ "Entropy lower bound on the number of required bits: "
            ++ show lbound
    putStrLn $ "Performance ratio: " ++ show (((8*fromIntegral nbits
                                    / fromIntegral len) ::Double) / lbound)
-}