{-# LANGUAGE BangPatterns, CPP #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

import Control.Applicative (many)
import Control.DeepSeq (NFData(rnf))
import Criterion.Main (bench, bgroup, defaultMain, nf)
import Data.Bits
import Data.Char (isAlpha)
import Data.Functor ((<$>))
import Data.Word (Word32)
import Data.Word (Word8)
import Numbers (numbers)
import Text.Parsec.Text ()
import Text.Parsec.Text.Lazy ()
import qualified AttoAeson
import qualified PicoAeson
import qualified Data.Attoparsec.ByteString as AB
import qualified Data.Attoparsec.ByteString.Char8 as AC
import qualified Data.Attoparsec.ByteString.Lazy as ABL
import qualified Data.Attoparsec.Text as AT
import qualified Data.Attoparsec.Text.Lazy as ATL
import qualified Data.Picoparsec as AM
import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 as BC
import qualified Data.ByteString.Lazy as BL
import qualified Data.ByteString.Lazy.Char8 as BLC
import qualified Data.Text as T
import qualified Data.Text.Lazy as TL
import qualified HeadersByteString
import qualified HeadersText
import qualified Links
import qualified Text.Parsec as P

import qualified Data.Monoid.Instances.ByteString.UTF8 as UTF8
import Data.Monoid.Instances.ByteString.Char8 ()

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

instance NFData P.ParseError where
    rnf = rnf . show

chunksOf :: Int -> [a] -> [[a]]
chunksOf k = go
  where go xs = case splitAt k xs of
                  ([],_)  -> []
                  (y, ys) -> y : go ys

main :: IO ()
main = do
  let s  = take 1024 . cycle $ ['a'..'z'] ++ ['A'..'Z']
      !b = BC.pack s
      !bl = BL.fromChunks . map BC.pack . chunksOf 4 $ s
      !t = T.pack s
      !tl = TL.fromChunks . map T.pack . chunksOf 4 $ s
      !utf8b = UTF8.ByteStringUTF8 b
  aesonA <- AttoAeson.aeson
  aesonP <- PicoAeson.aeson
  headersBS <- HeadersByteString.headers
  headersT <- HeadersText.headers
  defaultMain [
     bgroup "many" [
       bgroup "attoparsec" [
         bench "B" $ nf (AB.parse (many (AC.satisfy AC.isAlpha_ascii))) b
       , bench "BL" $ nf (ABL.parse (many (AC.satisfy AC.isAlpha_ascii))) bl
       , bench "T" $ nf (AT.parse (many (AT.satisfy AC.isAlpha_ascii))) t
       , bench "TL" $ nf (ATL.parse (many (AT.satisfy AC.isAlpha_ascii))) tl
       ]
     , bgroup "picoparsec" [
         bench "S" $ nf (AM.parse (many (AM.satisfyChar isAlpha))) s
       , bench "B" $ nf (AM.parse (many (AM.satisfy (isAlpha . BC.head)))) b
       , bench "BL" $ nf (AM.parse (many (AM.satisfy (isAlpha . BLC.head)))) bl
       , bench "T" $ nf (AM.parse (many (AM.satisfyChar AC.isAlpha_ascii))) t
       , bench "TL" $ nf (AM.parse (many (AM.satisfyChar AC.isAlpha_ascii))) tl
       ]
     , bgroup "parsec" [
         bench "S" $ nf (P.parse (many (P.satisfy AC.isAlpha_ascii)) "") s
       , bench "B" $ nf (P.parse (many (P.satisfy AC.isAlpha_ascii)) "") b
       , bench "BL" $ nf (P.parse (many (P.satisfy AC.isAlpha_ascii)) "") bl
       , bench "T" $ nf (P.parse (many (P.satisfy AC.isAlpha_ascii)) "") t
       , bench "TL" $ nf (P.parse (many (P.satisfy AC.isAlpha_ascii)) "") tl
       ]
     ]
   , bgroup "comparison" [
       bgroup "many-vs-takeWhile" [
         bgroup "attoparsec" [
           bench "many" $ nf (AB.parse (many (AC.satisfy AC.isAlpha_ascii))) b
         , bench "takeWhile" $ nf (AB.parse (AC.takeWhile AC.isAlpha_ascii)) b
         ]
       , bgroup "picoparsec" [
            bgroup "UTF8" [
               bench "many" $ nf (AM.parse (many (AM.satisfyChar AC.isAlpha_ascii))) utf8b
               , bench "takeWhile" $ nf (AM.parse (AM.takeCharsWhile AC.isAlpha_ascii)) utf8b
               ]
            , bgroup "Char8" [
               bench "many" $ nf (AM.parse (many (AM.satisfyChar AC.isAlpha_ascii))) b
               , bench "takeWhile" $ nf (AM.parse (AM.takeCharsWhile AC.isAlpha_ascii)) b
               ]
            ]
       ]
     , bgroup "letter-vs-isAlpha" [
         bgroup "attoparsec" [
           bench "letter" $ nf (AB.parse (many AC.letter_ascii)) b
         , bench "isAlpha" $ nf (AB.parse (many (AC.satisfy AC.isAlpha_ascii))) b
         ]
       , bgroup "picoparsec" [
           bench "letter" $ nf (AM.parse (AM.takeCharsWhile AC.isAlpha_ascii)) b
         , bench "isAlpha" $ nf (AM.parse (many (AM.satisfyChar isAlpha))) utf8b
         ]
       ]
     ]
   , bgroup "takeWhile" [
       bgroup "attoparsec" [
         bench "isAlpha" $ nf (ABL.parse (AC.takeWhile isAlpha)) bl
       , bench "isAlpha_ascii" $ nf (ABL.parse (AC.takeWhile AC.isAlpha_ascii)) bl
       , bench "isAlpha_iso8859_15" $ nf (ABL.parse (AC.takeWhile AC.isAlpha_iso8859_15)) bl
       ]
     , bgroup "picoparsec" [
          bgroup "UTF8" [
             bench "isAlpha" $ nf (AM.parse (AM.takeCharsWhile isAlpha)) utf8b
             , bench "isAlpha_ascii" $ nf (AM.parse (AM.takeCharsWhile AC.isAlpha_ascii)) utf8b
             , bench "isAlpha_iso8859_15" $ nf (AM.parse (AM.takeCharsWhile AC.isAlpha_iso8859_15)) utf8b
             ]
          , bgroup "Char8" [
             bench "isAlpha" $ nf (AM.parse (AM.takeCharsWhile isAlpha)) b
             , bench "isAlpha_ascii" $ nf (AM.parse (AM.takeCharsWhile AC.isAlpha_ascii)) b
             , bench "isAlpha_iso8859_15" $ nf (AM.parse (AM.takeCharsWhile AC.isAlpha_iso8859_15)) b
             ]
          ]
     ]
   , bgroup "word32LE" [
       bench "attoparsec" $ nf (AB.parse word32LE) b
     , bench "picoparsec" $ nf (AM.parse word32LE') b
     ]
   , bgroup "scan" [
       bench "short" $ nf (AB.parse quotedString) (BC.pack "abcdefghijk\"")
     , bench "long" $ nf (AB.parse quotedString) b
     ]
   , aesonA
   , aesonP
   , headersBS
   , headersT
   , Links.links
   , numbers
   ]

-- Benchmarks bind and (potential) bounds-check merging.
word32LE :: AB.Parser Word32
word32LE = do
    w1 <- AB.anyWord8
    w2 <- AB.anyWord8
    w3 <- AB.anyWord8
    w4 <- AB.anyWord8
    return $! (fromIntegral w1 :: Word32) +
        fromIntegral w2 `unsafeShiftL` 8 +
        fromIntegral w3 `unsafeShiftL` 16 +
        fromIntegral w4 `unsafeShiftL` 32

word32LE' :: AM.Parser B.ByteString Word32
word32LE' = do
    w1 <- B.head <$> AM.anyToken
    w2 <- B.head <$> AM.anyToken
    w3 <- B.head <$> AM.anyToken
    w4 <- B.head <$> AM.anyToken
    return $! (fromIntegral w1 :: Word32) +
        fromIntegral w2 `unsafeShiftL` 8 +
        fromIntegral w3 `unsafeShiftL` 16 +
        fromIntegral w4 `unsafeShiftL` 32

doubleQuote, backslash :: Word8
doubleQuote = 34
backslash = 92
{-# INLINE backslash #-}
{-# INLINE doubleQuote #-}

-- | Parse a string without a leading quote.
quotedString :: AB.Parser B.ByteString
quotedString = AB.scan False $ \s c -> if s then Just False
                                       else if c == doubleQuote
                                            then Nothing
                                            else Just (c == backslash)

#if !MIN_VERSION_base(4,5,0)
unsafeShiftL :: Bits a => a -> Int -> a
unsafeShiftL = shiftL
#endif