random-mhs-1.3.1: bench/Main.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Main (main) where

import Control.Monad
import Control.Monad.State.Strict
import Data.Int
import Data.List (sortOn)
import Data.Proxy
import Data.Typeable
import Data.Word
import Foreign.C.Types
import Numeric.Natural (Natural)
import System.Random.SplitMix as SM
import Test.Tasty.Bench
#if MIN_VERSION_primitive(0,7,1)
import Control.Monad.Primitive
import Data.Primitive.Types
import Data.Primitive.PrimArray
#endif

import System.Random.Stateful

seed :: Int
seed = 1337

main :: IO ()
main = do
  let !sz = 100000
      !sz100MiB = 100 * 1024 * 1024
      genLengths :: ([Int], StdGen)
      genLengths =
        -- create 5000 small lengths that are needed for ShortByteString generation
        runStateGen (mkStdGen 2020) $ \g -> replicateM 5000 (uniformRM (16 + 1, 16 + 7) g)
  setStdGen $ mkStdGen seed
  defaultMain
    [ bgroup
        "baseline"
        [ env (pure $ SM.mkSMGen $ fromIntegral seed) $ \smGen ->
            bench "nextWord32" $ whnf (genMany SM.nextWord32 smGen) sz
        , env (pure $ SM.mkSMGen $ fromIntegral seed) $ \smGen ->
            bench "nextWord64" $ whnf (genMany SM.nextWord64 smGen) sz
        , env (pure $ SM.mkSMGen $ fromIntegral seed) $ \smGen ->
            bench "nextInt" $ whnf (genMany SM.nextInt smGen) sz
        , env (pure $ SM.mkSMGen $ fromIntegral seed) $ \smGen ->
            bench "split" $ whnf (genMany SM.splitSMGen smGen) sz
        ]
    , bgroup
        "pure"
        [ bgroup
            "random"
            [ pureBench random sz (Proxy :: Proxy Word8)
            , pureBench random sz (Proxy :: Proxy Word16)
            , pureBench random sz (Proxy :: Proxy Word32)
            , pureBench random sz (Proxy :: Proxy Word64)
            , pureBench random sz (Proxy :: Proxy Int8)
            , pureBench random sz (Proxy :: Proxy Int16)
            , pureBench random sz (Proxy :: Proxy Int32)
            , pureBench random sz (Proxy :: Proxy Int64)
            , pureBench random sz (Proxy :: Proxy Bool)
            , pureBench random sz (Proxy :: Proxy Char)
            , pureBench random sz (Proxy :: Proxy Float)
            , pureBench random sz (Proxy :: Proxy Double)
            , pureBench random sz (Proxy :: Proxy Integer)
            ]
        , bgroup
            "uniform"
            [ pureBench uniform sz (Proxy :: Proxy Word8)
            , pureBench uniform sz (Proxy :: Proxy Word16)
            , pureBench uniform sz (Proxy :: Proxy Word32)
            , pureBench uniform sz (Proxy :: Proxy Word64)
            , pureBench uniform sz (Proxy :: Proxy Int8)
            , pureBench uniform sz (Proxy :: Proxy Int16)
            , pureBench uniform sz (Proxy :: Proxy Int32)
            , pureBench uniform sz (Proxy :: Proxy Int64)
            , pureBench uniform sz (Proxy :: Proxy Bool)
            , pureBench uniform sz (Proxy :: Proxy Char)
            , pureBench uniform sz (Proxy :: Proxy CChar)
            , pureBench uniform sz (Proxy :: Proxy CSChar)
            , pureBench uniform sz (Proxy :: Proxy CUChar)
            , pureBench uniform sz (Proxy :: Proxy CShort)
            , pureBench uniform sz (Proxy :: Proxy CUShort)
            , pureBench uniform sz (Proxy :: Proxy CInt)
            , pureBench uniform sz (Proxy :: Proxy CUInt)
            , pureBench uniform sz (Proxy :: Proxy CLong)
            , pureBench uniform sz (Proxy :: Proxy CULong)
            , pureBench uniform sz (Proxy :: Proxy CPtrdiff)
            , pureBench uniform sz (Proxy :: Proxy CSize)
            , pureBench uniform sz (Proxy :: Proxy CWchar)
            , pureBench uniform sz (Proxy :: Proxy CSigAtomic)
            , pureBench uniform sz (Proxy :: Proxy CLLong)
            , pureBench uniform sz (Proxy :: Proxy CULLong)
            , pureBench uniform sz (Proxy :: Proxy CIntPtr)
            , pureBench uniform sz (Proxy :: Proxy CUIntPtr)
            , pureBench uniform sz (Proxy :: Proxy CIntMax)
            , pureBench uniform sz (Proxy :: Proxy CUIntMax)
            ]
        , bgroup
            "uniformR"
            [ bgroup
                "full"
                [ pureUniformRFullBench (Proxy :: Proxy Word8) sz
                , pureUniformRFullBench (Proxy :: Proxy Word16) sz
                , pureUniformRFullBench (Proxy :: Proxy Word32) sz
                , pureUniformRFullBench (Proxy :: Proxy Word64) sz
                , pureUniformRFullBench (Proxy :: Proxy Word) sz
                , pureUniformRFullBench (Proxy :: Proxy Int8) sz
                , pureUniformRFullBench (Proxy :: Proxy Int16) sz
                , pureUniformRFullBench (Proxy :: Proxy Int32) sz
                , pureUniformRFullBench (Proxy :: Proxy Int64) sz
                , pureUniformRFullBench (Proxy :: Proxy Int) sz
                , pureUniformRFullBench (Proxy :: Proxy Char) sz
                , pureUniformRFullBench (Proxy :: Proxy Bool) sz
                , pureUniformRFullBench (Proxy :: Proxy CChar) sz
                , pureUniformRFullBench (Proxy :: Proxy CSChar) sz
                , pureUniformRFullBench (Proxy :: Proxy CUChar) sz
                , pureUniformRFullBench (Proxy :: Proxy CShort) sz
                , pureUniformRFullBench (Proxy :: Proxy CUShort) sz
                , pureUniformRFullBench (Proxy :: Proxy CInt) sz
                , pureUniformRFullBench (Proxy :: Proxy CUInt) sz
                , pureUniformRFullBench (Proxy :: Proxy CLong) sz
                , pureUniformRFullBench (Proxy :: Proxy CULong) sz
                , pureUniformRFullBench (Proxy :: Proxy CPtrdiff) sz
                , pureUniformRFullBench (Proxy :: Proxy CSize) sz
                , pureUniformRFullBench (Proxy :: Proxy CWchar) sz
                , pureUniformRFullBench (Proxy :: Proxy CSigAtomic) sz
                , pureUniformRFullBench (Proxy :: Proxy CLLong) sz
                , pureUniformRFullBench (Proxy :: Proxy CULLong) sz
                , pureUniformRFullBench (Proxy :: Proxy CIntPtr) sz
                , pureUniformRFullBench (Proxy :: Proxy CUIntPtr) sz
                , pureUniformRFullBench (Proxy :: Proxy CIntMax) sz
                , pureUniformRFullBench (Proxy :: Proxy CUIntMax) sz
                ]
            , bgroup
                "excludeMax"
                [ pureUniformRExcludeMaxBench (Proxy :: Proxy Word8) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy Word16) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy Word32) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy Word64) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy Word) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy Int8) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy Int16) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy Int32) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy Int64) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy Int) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy Char) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy Bool) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy CChar) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy CSChar) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy CUChar) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy CShort) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy CUShort) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy CInt) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy CUInt) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy CLong) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy CULong) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy CPtrdiff) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy CSize) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy CWchar) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy CSigAtomic) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy CLLong) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy CULLong) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy CIntPtr) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy CUIntPtr) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy CIntMax) sz
                , pureUniformRExcludeMaxBench (Proxy :: Proxy CUIntMax) sz
                ]
            , bgroup
                "includeHalf"
                [ pureUniformRIncludeHalfBench (Proxy :: Proxy Word8) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy Word16) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy Word32) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy Word64) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy Word) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy Int8) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy Int16) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy Int32) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy Int64) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy Int) sz
                , pureUniformRIncludeHalfEnumBench (Proxy :: Proxy Char) sz
                , pureUniformRIncludeHalfEnumBench (Proxy :: Proxy Bool) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy CChar) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy CSChar) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy CUChar) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy CShort) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy CUShort) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy CInt) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy CUInt) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy CLong) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy CULong) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy CPtrdiff) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy CSize) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy CWchar) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy CSigAtomic) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy CLLong) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy CULLong) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy CIntPtr) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy CUIntPtr) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy CIntMax) sz
                , pureUniformRIncludeHalfBench (Proxy :: Proxy CUIntMax) sz
                ]
            , bgroup
                "unbounded"
                [ pureUniformRBench (Proxy :: Proxy Float) (1.23e-4, 5.67e8) sz
                , pureUniformRBench (Proxy :: Proxy Double) (1.23e-4, 5.67e8) sz
                , let !i = (10 :: Integer) ^ (100 :: Integer)
                      !range = (-i - 1, i + 1)
                   in pureUniformRBench (Proxy :: Proxy Integer) range sz
                , let !n = (10 :: Natural) ^ (100 :: Natural)
                      !range = (1, n - 1)
                   in pureUniformRBench (Proxy :: Proxy Natural) range sz
                ]
            , bgroup "floating" $
                fillFloating sz
                  ++ [ bgroup
                         "State"
                         [ bgroup
                             "Float"
                             [ env getStdGen $
                                 bench "uniformRM" . nf (`runStateGen` (replicateM_ sz . uniformRM (0.1 :: Float, 1.1)))
                             , env getStdGen $
                                 bench "uniformFloat01M" . nf (`runStateGen` (replicateM_ sz . uniformFloat01M))
                             , env getStdGen $
                                 bench "uniformFloatPositive01M"
                                   . nf (`runStateGen` (replicateM_ sz . uniformFloatPositive01M))
                             ]
                         , bgroup
                             "Double"
                             [ env getStdGen $
                                 bench "uniformRM" . nf (`runStateGen` (replicateM_ sz . uniformRM (0.1 :: Double, 1.1)))
                             , env getStdGen $
                                 bench "uniformDouble01M" . nf (`runStateGen` (replicateM_ sz . uniformDouble01M))
                             , env getStdGen $
                                 bench "uniformDoublePositive01M"
                                   . nf (`runStateGen` (replicateM_ sz . uniformDoublePositive01M))
                             ]
                         ]
                     , bgroup
                         "pure"
                         [ bgroup
                             "Float"
                             [ env getStdGen $ \gen ->
                                 bench "uniformRM" $
                                   nf
                                     (genMany (runState $ uniformRM (0.1 :: Float, 1.1) (StateGenM :: StateGenM StdGen)) gen)
                                     sz
                             , env getStdGen $ \gen ->
                                 bench "uniformFloat01M" $
                                   nf
                                     (genMany (runState $ uniformFloat01M (StateGenM :: StateGenM StdGen)) gen)
                                     sz
                             , env getStdGen $ \gen ->
                                 bench "uniformFloatPositive01M" $
                                   nf
                                     (genMany (runState $ uniformFloatPositive01M (StateGenM :: StateGenM StdGen)) gen)
                                     sz
                             ]
                         , bgroup
                             "Double"
                             [ env getStdGen $ \gen ->
                                 bench "uniformRM" $
                                   nf
                                     (genMany (runState $ uniformRM (0.1 :: Double, 1.1) (StateGenM :: StateGenM StdGen)) gen)
                                     sz
                             , env getStdGen $ \gen ->
                                 bench "uniformDouble01M" $
                                   nf
                                     (genMany (runState $ uniformDouble01M (StateGenM :: StateGenM StdGen)) gen)
                                     sz
                             , env getStdGen $ \gen ->
                                 bench "uniformDoublePositive01M" $
                                   nf
                                     (genMany (runState $ uniformDoublePositive01M (StateGenM :: StateGenM StdGen)) gen)
                                     sz
                             ]
                         ]
                     ]
            ]
        , bgroup
            "Bytes"
            [ env (pure genLengths) $ \ ~(ns, gen) ->
                bench "uniformShortByteStringM" $
                  nfIO $
                    runStateGenT gen $
                      \g -> mapM (`uniformShortByteStringM` g) ns
            , env getStdGen $ \gen ->
                bench "uniformByteStringM 100MB" $
                  nf (runStateGen gen . uniformByteStringM) sz100MiB
            , env getStdGen $ \gen ->
                bench "uniformByteArray 100MB" $ nf (\n -> uniformByteArray False n gen) sz100MiB
            , env getStdGen $ \gen ->
                bench "uniformByteString 100MB" $ nf (`uniformByteString` gen) sz100MiB
            ]
        ]
    , env (pure [0 :: Integer .. 200000]) $ \xs ->
        bgroup
          "shuffle"
          [ env getStdGen $ bench "uniformShuffleList" . nf (uniformShuffleList xs)
          , env getStdGen $ bench "uniformShuffleListM" . nf (`runStateGen` uniformShuffleListM xs)
          , env getStdGen $ bench "naiveShuffleListM" . nf (`runStateGen` naiveShuffleListM xs)
          ]
    ]
  where
#if MIN_VERSION_primitive(0,7,1)
    fillFloating sz =
      [ bgroup "IO"
        [ bgroup "Float"
          [ env ((,) <$> getStdGen <*> newAlignedPinnedPrimArray sz) $ \ ~(gen, ma) ->
              bench "uniformRM" $
              nfIO (runStateGenT gen (fillMutablePrimArrayM (uniformRM (0 :: Float, 1.1)) ma))
          , env ((,) <$> getStdGen <*> newAlignedPinnedPrimArray sz) $ \ ~(gen, ma) ->
              bench "uniformFloat01M" $
              nfIO (runStateGenT gen (fillMutablePrimArrayM uniformFloat01M ma))
          , env ((,) <$> getStdGen <*> newAlignedPinnedPrimArray sz) $ \ ~(gen, ma) ->
              bench "uniformFloatPositive01M" $
              nfIO (runStateGenT gen (fillMutablePrimArrayM uniformFloatPositive01M ma))
          ]
        , bgroup "Double"
          [ env ((,) <$> getStdGen <*> newAlignedPinnedPrimArray sz) $ \ ~(gen, ma) ->
              bench "uniformRM" $
              nfIO (runStateGenT gen (fillMutablePrimArrayM (uniformRM (0 :: Double, 1.1)) ma))
          , env ((,) <$> getStdGen <*> newAlignedPinnedPrimArray sz) $ \ ~(gen, ma) ->
              bench "uniformDouble01M" $
              nfIO (runStateGenT gen (fillMutablePrimArrayM uniformDouble01M ma))
          , env ((,) <$> getStdGen <*> newAlignedPinnedPrimArray sz) $ \ ~(gen, ma) ->
              bench "uniformDoublePositive01M" $
              nfIO (runStateGenT gen (fillMutablePrimArrayM uniformDoublePositive01M ma))
          ]
        ]
      ]
#else
    fillFloating _ = []
#endif

pureUniformRFullBench ::
  forall a.
  (Typeable a, UniformRange a, Bounded a) =>
  Proxy a ->
  Int ->
  Benchmark
pureUniformRFullBench px =
  let range = (minBound :: a, maxBound :: a)
   in pureUniformRBench px range
{-# INLINE pureUniformRFullBench #-}

pureUniformRExcludeMaxBench ::
  forall a.
  (Typeable a, UniformRange a, Bounded a, Enum a) =>
  Proxy a ->
  Int ->
  Benchmark
pureUniformRExcludeMaxBench px =
  let range = (minBound :: a, pred (maxBound :: a))
   in pureUniformRBench px range
{-# INLINE pureUniformRExcludeMaxBench #-}

pureUniformRIncludeHalfBench ::
  forall a.
  (Typeable a, UniformRange a, Bounded a, Integral a) =>
  Proxy a ->
  Int ->
  Benchmark
pureUniformRIncludeHalfBench px =
  let range = ((minBound :: a) + 1, ((maxBound :: a) `div` 2) + 1)
   in pureUniformRBench px range
{-# INLINE pureUniformRIncludeHalfBench #-}

pureUniformRIncludeHalfEnumBench ::
  forall a.
  (Typeable a, UniformRange a, Bounded a, Enum a) =>
  Proxy a ->
  Int ->
  Benchmark
pureUniformRIncludeHalfEnumBench px =
  let range = (succ (minBound :: a), toEnum ((fromEnum (maxBound :: a) `div` 2) + 1))
   in pureUniformRBench px range
{-# INLINE pureUniformRIncludeHalfEnumBench #-}

pureUniformRBench ::
  forall a.
  (Typeable a, UniformRange a) =>
  Proxy a ->
  (a, a) ->
  Int ->
  Benchmark
pureUniformRBench px range@(!_, !_) sz = pureBench (uniformR range) sz px
{-# INLINE pureUniformRBench #-}

pureBench ::
  forall a.
  Typeable a =>
  (StdGen -> (a, StdGen)) ->
  Int ->
  Proxy a ->
  Benchmark
pureBench f sz px =
  env getStdGen $ \gen ->
    bench (showsTypeRep (typeRep px) "") $ whnf (genMany f gen) sz
{-# INLINE pureBench #-}

genMany :: (g -> (a, g)) -> g -> Int -> a
genMany f g0 n = go 0 $ f g0
  where
    go i (!y, !g)
      | i < n = go (i + 1) $ f g
      | otherwise = y

#if MIN_VERSION_primitive(0,7,1)
fillMutablePrimArrayM ::
     (Prim a, PrimMonad m)
  => (gen -> m a)
  -> MutablePrimArray (PrimState m) a
  -> gen
  -> m (PrimArray a)
fillMutablePrimArrayM f ma g = do
  n <- getSizeofMutablePrimArray ma
  let go i
        | i < n = f g >>= writePrimArray ma i >> go (i + 1)
        | otherwise = pure ()
  go 0
  unsafeFreezePrimArray ma
#endif

naiveShuffleListM :: StatefulGen g m => [a] -> g -> m [a]
naiveShuffleListM xs gen = do
  is <- uniformListM n gen
  pure $ map snd $ sortOn fst $ zip (is :: [Int]) xs
  where
    !n = length xs
{-# INLINE naiveShuffleListM #-}