module Main where

import Criterion.Main (defaultMain)
import Data.ByteString (ByteString)
import Hedgehog.Corpus (muppets, simpsons, vegetables, viruses)
import Criterion
import qualified Data.Map.Strict as M
import qualified Data.ByteString.Char8 as B
import qualified Data.Trie as Trie

setupEnv :: IO ([(ByteString, ())], [(ByteString, ())])
setupEnv = pure (small, big)
  where
    small =
      [ (firstname <> "." <> surname <> "@" <> domain <> "." <> tld, ())
      | firstname <- simpsons
      , surname <- muppets
      , domain <- vegetables
      , tld <- ["com", "fi", "co.uk", "info", "org"]
      ]
    big =
      [ (firstname <> "." <> surname <> "@" <> domain1 <> "." <> domain2 <> "." <> tld, ())
      | firstname <- simpsons
      , surname <- muppets
      , domain1 <- vegetables
      , domain2 <- viruses
      , tld <- ["com", "fi", "co.uk", "info", "org"]
      ]

main :: IO ()
main = defaultMain
  [ env setupEnv $ \ ~(small,big) -> bgroup "main"
    [ bgroup "small"
      [ bgroup "Data.Map"
        [ bench "length . toList . fromList" $ whnf (length . M.toList . M.fromList) small
        , bench "length . filter (\"homer\" `isPrefixOf`) . keys . fromList" $ whnf (length . M.keys . M.filterWithKey (\k _ -> "homer" `B.isPrefixOf` k) . M.fromList) small
        ]
      , bgroup "Data.Trie"
        [ bench "length . toList . fromList" $ whnf (length . Trie.toList . Trie.fromList) small
        , bench "length . keys . submap \"homer\" . fromList" $ whnf (length . Trie.keys . Trie.submap "homer" .  Trie.fromList) small
        ]
      ]
    , bgroup "big"
      [ bgroup "Data.Map"
        [ bench "length . toList . fromList" $ whnf (length . M.toList . M.fromList) big
        , bench "length . filter (\"homer\" `isPrefixOf`) . keys . fromList" $ whnf (length . M.keys . M.filterWithKey (\k _ -> "homer" `B.isPrefixOf` k) . M.fromList) big
        ]
      , bgroup "Data.Trie"
        [ bench "length . toList . fromList" $ whnf (length . Trie.toList . Trie.fromList) big
        , bench "length . keys . submap \"homer\" . fromList" $ whnf (length . Trie.keys . Trie.submap "homer" .  Trie.fromList) big
        ]
      ]
    ]
  ]