{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}

module Main (main) where

import Criterion.Main
import Data.Char
import qualified Data.Text as T
import qualified Data.Text.Array as TA
import qualified Data.Text.Internal as T
import qualified Data.Text.Lazy as TL (Text, fromStrict)
import Data.Bits
import Data.Word

isAsciiLetter :: Char -> Bool
isAsciiLetter c = isAsciiLower c || isAsciiUpper c

isAllowedStart :: Char -> Bool
isAllowedStart c = c == '_' || isAsciiLetter c

data MangledSize = MangledSize
    { _unmangledChars :: !Int
    , _mangledWord16s :: !Int
    }

ord' :: Char -> Word16
ord' = fromIntegral . ord

mangleIdentifier' :: T.Text -> Either String TL.Text
mangleIdentifier' txt = case T.foldl' f (MangledSize 0 0) txt of
    MangledSize 0 _ -> Left "Empty identifier"
    MangledSize chars word16s
      | chars == word16s -> Right $! TL.fromStrict txt
      | otherwise -> Right $! TL.fromStrict $
        let !arr = TA.run $ do
            a <- TA.new word16s
            let poke !j !minj !x
                  | j < minj = pure ()
                  | otherwise = do

                        let !(!x', !r) = quotRem x 16
                        -- let !x' = x `unsafeShiftR` 4
                        -- let !r = x .&. 0xF

                        TA.unsafeWrite a j (fromIntegral $ ord $ intToDigit r)
                        poke (j - 1) minj x'
                go !i t = case T.uncons t of
                    Nothing -> pure ()
                    Just (!c, !t')
                      | isAllowedStart c || i > 0 && isDigit c -> TA.unsafeWrite a i (fromIntegral $ ord c) >> go (i + 1) t'
                      | c == '$' -> do
                            TA.unsafeWrite a i (ord' '$')
                            TA.unsafeWrite a (i + 1) (ord' '$')
                            go (i + 2) t'
                      | ord c <= 0xFFFF -> do
                            TA.unsafeWrite a i (ord' '$')
                            TA.unsafeWrite a (i + 1) (ord' 'u')
                            poke (i + 5) (i + 2) (ord c)
                            go (i + 6) t'
                      | otherwise -> do
                            TA.unsafeWrite a i (ord' '$')
                            TA.unsafeWrite a (i + 1) (ord' 'U')
                            poke (i + 9) (i + 2) (ord c)
                            go (i + 10) t'
            go 0 txt
            pure a
        in T.text arr 0 word16s
    where f :: MangledSize -> Char -> MangledSize
          f (MangledSize chars word16s) c
            | isAllowedStart c || chars > 0 && isDigit c = MangledSize (chars + 1) (word16s + 1)
            | c == '$' = MangledSize 1 (word16s + 2)
            | ord c > 0xFFFF = MangledSize 1 (word16s + 10)
            | otherwise = MangledSize 1 (word16s + 6)

criterionMain :: IO ()
criterionMain = defaultMain [bench "mangling new" $ nf mangleIdentifier' "$ $ $ $ $ $ $ $ $ $ $ $ $ $ $ $ $ $ $ $ "]

main :: IO ()
main = criterionMain