{- cabal:
build-depends: base, random, text, tasty-bench
-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# OPTIONS_GHC -ddump-simpl -ddump-to-file #-}

import qualified Data.List as L
import qualified Data.Text as T
import System.Random (mkStdGen, uniforms)

import qualified Data.Text.Array as A
import Data.Text.Internal (Text(..))
import Data.Text.Internal.Encoding.Utf8 (utf8LengthByLeader, chr2, chr3, chr4)
import Data.Text.Internal.Unsafe.Char (unsafeChr8)
import Data.Text.Unsafe (iterArray, Iter(..))
import GHC.Word (Word8(..))
import GHC.Exts

import Test.Tasty.Bench

main :: IO ()
main = defaultMain
  [ env (pure (mk n)) $ \t ->
    bgroup name
      [ bgroup "sum"
        [ bench "current" $ whnf (mkSum utf8LengthByLeader) t
        , bcompare (currentName "sum") $
            bench "branching" $ whnf (mkSum utf8LengthByLeader_branching) t
        , bcompare (currentName "sum") $
            bench "simple" $ whnf (mkSum utf8LengthByLeader_simple) t
        , bcompare (currentName "sum") $
            bench "lookup" $ whnf (mkSum utf8LengthByLeader_lookup) t
        ]
      , bgroup "length"
        [ bench "current" $ whnf (mkLength utf8LengthByLeader) t
        , bcompare (currentName "length") $
            bench "branching" $ whnf (mkLength utf8LengthByLeader_branching) t
        , bcompare (currentName "length") $
            bench "simple" $ whnf (mkLength utf8LengthByLeader_simple) t
        , bcompare (currentName "length") $
            bench "lookup" $ whnf (mkLength utf8LengthByLeader_lookup) t
        ]
      ]
  | n <- [100, 100000]
  , (typ, mk) <- [("ascii", mkAscii), ("random", mkRandom)]
  , let name = show n ++ " " ++ typ
        currentName which = "All." ++ name ++ "." ++ which ++ ".current"
  ]

----------------------
-- Benchmarked funcs
----------------------

-- Branches on the char length
mkSum :: (Word8 -> Int) -> Text -> Int
mkSum f = \(Text arr off len) -> loop arr (off + len) off 0
  where
    loop !a !n !i !acc
      | i >= n = acc
      | otherwise = loop a n (i + d) (acc + fromEnum c)
      where
        m0 = A.unsafeIndex a i
        m1 = A.unsafeIndex a (i + 1)
        m2 = A.unsafeIndex a (i + 2)
        m3 = A.unsafeIndex a (i + 3)
        d = f m0
        c = case d of
          1 -> unsafeChr8 m0
          2 -> chr2 m0 m1
          3 -> chr3 m0 m1 m2
          _ -> chr4 m0 m1 m2 m3
{-# INLINE mkSum #-}

-- Does not branch on the char length
mkLength :: (Word8 -> Int) -> Text -> Int
mkLength f = \(Text arr off len) -> loop arr (off + len) off 0
  where
    loop !a !n !i !acc
      | i >= n = acc
      | otherwise = loop a n (i + f (A.unsafeIndex a i)) (acc + 1)
{-# INLINE mkLength #-}

---------
-- Data
---------

mkAscii :: Int -> Text
mkAscii n =
  T.pack $
  L.take n $
  L.cycle $
  "The quick brown fox jumps over the lazy dog\n"

-- Each Char's utf-8 byte length is random. This takes away any advantage
-- the branching impl might have due to branch prediction.
mkRandom :: Int -> Text
mkRandom n =
  T.pack $
  L.map mkC $
  L.take n $
  uniforms (mkStdGen 42)
  where
    mkC :: Int -> Char
    mkC i = toEnum $ case i `mod` 4 of
      0 -> 0x00007F
      1 -> 0x0007FF
      2 -> 0x00FFFF
      3 -> 0x10FFFF

-----------------------------
-- utf8LengthByLeader impls
-----------------------------

utf8LengthByLeader_branching :: Word8 -> Int
utf8LengthByLeader_branching w
  | w < 0x80  = 1
  | w < 0xE0  = 2
  | w < 0xF0  = 3
  | otherwise = 4
{-# INLINE utf8LengthByLeader_branching #-}

utf8LengthByLeader_simple :: Word8 -> Int
utf8LengthByLeader_simple (W8# w8#) =
  let w# = word8ToWord# w8#
  in   1
     + I# (w# `geWord#` 0x80##)
     + I# (w# `geWord#` 0xE0##)
     + I# (w# `geWord#` 0xF0##)
{-# INLINE utf8LengthByLeader_simple #-}

utf8LengthByLeader_lookup :: Word8 -> Int
utf8LengthByLeader_lookup (W8# w8#) =
  I# (int8ToInt# (indexInt8OffAddr# a# (word2Int# (uncheckedShiftRL# (word8ToWord# w8#) 4#))))
  where
    a# = "\1\1\1\1\1\1\1\1\1\1\1\1\2\2\3\4"#
    -- 0..7 -> 1
    -- 8..B -> invalid leading byte
    -- C..D -> 2
    -- E    -> 3
    -- F    -> 4
{-# INLINE utf8LengthByLeader_lookup #-}