criterion-1.3.0.0: examples/Maps.hs

-- Benchmark the cost of creating various types of map.

{-# LANGUAGE CPP, FlexibleContexts, ScopedTypeVariables #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

import Criterion.Main
import Data.ByteString (ByteString, pack)
import Data.Hashable (Hashable)
import System.Random.MWC
import qualified Data.CritBit.Map.Lazy as C
import qualified Data.HashMap.Lazy as H
import qualified Data.IntMap as I
import qualified Data.Map as M
import qualified Data.Vector as V
import qualified Data.Vector.Algorithms.Intro as I
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U

#if !MIN_VERSION_bytestring(0,10,0)
import Control.DeepSeq (NFData (..))
#endif

type V = U.Vector Int
type B = V.Vector ByteString

numbers :: IO (V, V, V)
numbers = do
  random <- withSystemRandom . asGenIO $ \gen -> uniformVector gen 40000
  let sorted    = G.modify I.sort random
      revsorted = G.reverse sorted
  return (random, sorted, revsorted)

strings :: IO (B, B, B)
strings = do
  random <- withSystemRandom . asGenIO $ \gen ->
    V.replicateM 10000 $
      (pack . U.toList) `fmap` (uniformVector gen =<< uniformR (1,16) gen)
  let sorted    = G.modify I.sort random
      revsorted = G.reverse sorted
  return (random, sorted, revsorted)

main :: IO ()
main = defaultMain [
         env numbers $ \ ~(random,sorted,revsorted) ->
         bgroup "Int" [
           bgroup "IntMap" [
             bench "sorted"    $ whnf intmap sorted
           , bench "random"    $ whnf intmap random
           , bench "revsorted" $ whnf intmap revsorted
           ]
         , bgroup "Map" [
             bench "sorted"    $ whnf mmap sorted
           , bench "random"    $ whnf mmap random
           , bench "revsorted" $ whnf mmap revsorted
           ]
         , bgroup "HashMap" [
             bench "sorted"    $ whnf hashmap sorted
           , bench "random"    $ whnf hashmap random
           , bench "revsorted" $ whnf hashmap revsorted
           ]
         ]
       , env strings $ \ ~(random,sorted,revsorted) ->
         bgroup "ByteString" [
           bgroup "Map" [
             bench "sorted"    $ whnf mmap sorted
           , bench "random"    $ whnf mmap random
           , bench "revsorted" $ whnf mmap revsorted
           ]
         , bgroup "HashMap" [
             bench "sorted"    $ whnf hashmap sorted
           , bench "random"    $ whnf hashmap random
           , bench "revsorted" $ whnf hashmap revsorted
           ]
         , bgroup "CritBit" [
             bench "sorted"    $ whnf critbit sorted
           , bench "random"    $ whnf critbit random
           , bench "revsorted" $ whnf critbit revsorted
           ]
         ]
       ]

critbit :: (G.Vector v k, C.CritBitKey k) => v k -> C.CritBit k Int
critbit xs = G.foldl' (\m k -> C.insert k value m) C.empty xs

hashmap :: (G.Vector v k, Hashable k, Eq k) => v k -> H.HashMap k Int
hashmap xs = G.foldl' (\m k -> H.insert k value m) H.empty xs

intmap :: G.Vector v Int => v Int -> I.IntMap Int
intmap xs = G.foldl' (\m k -> I.insert k value m) I.empty xs

mmap :: (G.Vector v k, Ord k) => v k -> M.Map k Int
mmap xs = G.foldl' (\m k -> M.insert k value m) M.empty xs

value :: Int
value = 31337

#if !MIN_VERSION_bytestring(0,10,0)
instance NFData ByteString where
    rnf bs = bs `seq` ()
#endif