{-# LANGUAGE CPP #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE BangPatterns #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

import qualified Data.ByteString as BS
import qualified Data.HashMap.Strict as HM
import qualified Data.Map as M
import qualified Data.Sequence as S
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as U

import Data.Vector.Generic.Lens
import Data.ByteString.Lens

import Control.Lens
import Criterion.Main
import Criterion.Types

#if !(MIN_VERSION_bytestring(0,10,0))
import Control.DeepSeq (NFData(..))
import qualified Data.ByteString.Internal as BS
#endif

#if !(MIN_VERSION_containers(0,5,0))
import qualified Data.Foldable as F
#endif

#if !(MIN_VERSION_bytestring(0,10,0))
instance NFData BS.ByteString where
    rnf (BS.PS _ _ _) = ()
#endif

#if !(MIN_VERSION_containers(0,5,0))
-- Sadly, containers doesn't export the constructor for Seq on older versions,
-- so we'll have to settle for this inefficient implementation of rnf.
instance NFData a => NFData (S.Seq a) where
    rnf = rnf . F.toList
#endif

main :: IO ()
main = defaultMainWith config
  [
    bgroup "vector"
    [ bgroup "map"
      [ bench "native"     $ nf (V.map (+100)) v
      , bench "itraversed" $ nf (over itraversed (+100)) v
      ]
    , bgroup "imap"
      [ bench "native"     $ nf (V.imap           (\i x -> x + i +100)) v
      , bench "imap"       $ nf (imap             (\i x -> x + i +100)) v
      , bench "itraversed" $ nf (iover itraversed (\i x -> x + i +100)) v
      ]
    ]
  , bgroup "unboxed-vector"
    [ bgroup "map"
      [ bench "native"     $ nf (U.map (+100)) u
      , bench "itraversed" $ nf (over each (+100)) u
      ]
    , bgroup "imap"
      [ bench "native"     $ nf (U.imap (\i x -> x + i +100)) u
      , bench "itraversed" $ nf (iover vectorTraverse (\i x -> x + i) :: U.Vector Int -> U.Vector Int) u
      ]
    ]
  , bgroup "sequence"
    [ bgroup "map"
      [ bench "native" $ nf (fmap            (+100)) s
      , bench "each"   $ nf (over each       (+100)) s
      ]
    , bgroup "imap"
      [ bench "native" $ nf (S.mapWithIndex    (\i x -> x + i +100)) s
      , bench "imap"   $ nf (imap              (\i x -> x + i +100)) s
      ]
    ]
  , bgroup "bytestring"
    [ bgroup "map"
      [ bench "native" $ nf (BS.map     (+100)) b
      , bench "each"   $ nf (over each  (+100)) b
      ]
    , bgroup "imap"
      [
        bench "bytes" $ nf (iover bytes (\i x -> x + fromIntegral i +100)) b
      ]
    ]
  , bgroup "list"
    [ bgroup "map"
      [ bench "native" $ nf (map       (+100)) l
      , bench "each"   $ nf (over each (+100)) l
      ]
    , bgroup "imap"
      [ bench "imap" $ nf (imap (\i x -> x + i +100)) l
      ]
    ]
  , bgroup "map"
    [ bgroup "map"
      [ bench "native"     $ nf (fmap            (+100)) m
      , bench "each"       $ nf (over each       (+100)) m
      , bench "itraversed" $ nf (over itraversed (+100)) m
      ]
    , bgroup "imap"
      [ bench "native" $ nf (M.mapWithKey (\i x -> x + i +100)) m
      , bench "each"   $ nf (imap         (\i x -> x + i +100)) m
      ]
    ]
  , bgroup "hash map"
    [ bgroup "map"
      [ bench "native" $ nf (HM.map    (+100)) h
      , bench "each"   $ nf (over each (+100)) h
      ]
    , bgroup "imap"
      [ bench "native" $ nf (HM.mapWithKey (\i x -> x + i +100)) h
      , bench "imap"   $ nf (imap          (\i x -> x + i +100)) h
      ]
    ]
  ]
  where
    config = defaultConfig { timeLimit = 1 }
    l  = [0..10000] :: [Int]
    xl = [0..100000] :: [Int]
    b  = BS.pack $ map fromIntegral xl
    h  = HM.fromList $ zip l l
    m  = M.fromList $ zip l l
    s  = S.fromList l
    u  = U.fromList xl
    v  = V.fromList l