content: {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE MagicHash #-}
  {-# LANGUAGE UnboxedTuples #-}
  {-# OPTIONS_GHC -fno-warn-name-shadowing #-}
  
  module Data.Utf8
    ( Utf8 (..),
      isEmpty,
      empty,
      size,
      contains,
      startsWith,
      startsWithChar,
      endsWithWord8,
      split,
      join,
      --
      getUnder256,
      putUnder256,
      --
      getVeryLong,
      putVeryLong,
      --
      toChars,
      toByteString,
      toBuilder,
      toEscapedBuilder,
      --
      fromPtr,
      fromSnippet,
      fromChars,
      --
      MBA,
      newByteArray,
      copyFromPtr,
      writeWord8,
      freeze,
    )
  where
  
  import Data.Binary (Get, Put, get, getWord8, put, putWord8)
  import Data.Binary.Get.Internal (readN)
  import Data.Binary.Put (putBuilder)
  import Data.Bits (shiftR, (.&.))
  import Data.ByteString.Builder qualified as Builder
  import Data.ByteString.Builder.Internal qualified as B
  import Data.ByteString.Internal qualified as B
  import Data.ByteString.Lazy qualified as LazyByteString
  import Data.Char qualified as Char
  import Data.List qualified as List
  import Foreign.ForeignPtr (touchForeignPtr)
  import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
  import Foreign.Ptr (minusPtr, plusPtr)
  import GHC.Exts
    ( Char (C#),
      Int (I#),
      Ptr (Ptr),
      isTrue#,
    )
  import GHC.IO
  import GHC.Prim
  import GHC.ST (ST (ST), runST)
  import GHC.Word (Word8 (W8#))
  import Parse.Primitives qualified as P
  import Prelude hiding (String, all, any, concat)
  
  -- UTF-8
  
  data Utf8 tipe
    = Utf8 ByteArray#
  
  instance Show (Utf8 tipe) where
    show str = '"' : toChars str ++ "\""
  
  -- EMPTY
  
  empty :: Utf8 t
  empty =
    runST (freeze =<< newByteArray 0)
  
  isEmpty :: Utf8 t -> Bool
  isEmpty (Utf8 ba#) =
    isTrue# (sizeofByteArray# ba# ==# 0#)
  
  -- SIZE
  
  size :: Utf8 t -> Int
  size (Utf8 ba#) =
    I# (sizeofByteArray# ba#)
  
  -- CONTAINS
  
  contains :: Word8 -> Utf8 t -> Bool
  contains (W8# word#) (Utf8 ba#) =
    containsHelp word# ba# 0# (sizeofByteArray# ba#)
  
  containsHelp :: Word8# -> ByteArray# -> Int# -> Int# -> Bool
  containsHelp word# ba# !offset# len# =
    if isTrue# (offset# <# len#)
      then
        if isTrue# (eqWord8# word# (indexWord8Array# ba# offset#))
          then True
          else containsHelp word# ba# (offset# +# 1#) len#
      else False
  
  -- STARTS WITH
  
  startsWith :: Utf8 t -> Utf8 t -> Bool
  startsWith (Utf8 ba1#) (Utf8 ba2#) =
    let !len1# = sizeofByteArray# ba1#
        !len2# = sizeofByteArray# ba2#
     in isTrue# (len1# <=# len2#)
          && isTrue# (0# ==# compareByteArrays# ba1# 0# ba2# 0# len1#)
  
  -- STARTS WITH CHAR
  
  startsWithChar :: (Char -> Bool) -> Utf8 t -> Bool
  startsWithChar isGood bytes@(Utf8 ba#) =
    if isEmpty bytes
      then False
      else
        let !w# = word8ToWord# (indexWord8Array# ba# 0#)
            !char
              | isTrue# (ltWord# w# 0xC0##) = C# (chr# (word2Int# w#))
              | isTrue# (ltWord# w# 0xE0##) = chr2 ba# 0# w#
              | isTrue# (ltWord# w# 0xF0##) = chr3 ba# 0# w#
              | True = chr4 ba# 0# w#
         in isGood char
  
  -- ENDS WITH WORD
  
  endsWithWord8 :: Word8 -> Utf8 t -> Bool
  endsWithWord8 (W8# w#) (Utf8 ba#) =
    let len# = sizeofByteArray# ba#
     in isTrue# (len# ># 0#)
          && isTrue# (eqWord8# w# (indexWord8Array# ba# (len# -# 1#)))
  
  -- SPLIT
  
  split :: Word8 -> Utf8 t -> [Utf8 t]
  split (W8# divider#) str@(Utf8 ba#) =
    splitHelp str 0 (findDividers divider# ba# 0# (sizeofByteArray# ba#) [])
  
  splitHelp :: Utf8 t -> Int -> [Int] -> [Utf8 t]
  splitHelp str start offsets =
    case offsets of
      [] ->
        [unsafeSlice str start (size str)]
      offset : offsets ->
        unsafeSlice str start offset : splitHelp str (offset + 1) offsets
  
  findDividers :: Word8# -> ByteArray# -> Int# -> Int# -> [Int] -> [Int]
  findDividers divider# ba# !offset# len# revOffsets =
    if isTrue# (offset# <# len#)
      then
        findDividers divider# ba# (offset# +# 1#) len# $
          if isTrue# (eqWord8# divider# (indexWord8Array# ba# offset#))
            then I# offset# : revOffsets
            else revOffsets
      else reverse revOffsets
  
  unsafeSlice :: Utf8 t -> Int -> Int -> Utf8 t
  unsafeSlice str start end =
    let !len = end - start
     in if len == 0
          then empty
          else runST $
            do
              mba <- newByteArray len
              copy str start mba 0 len
              freeze mba
  
  -- JOIN
  
  join :: Word8 -> [Utf8 t] -> Utf8 t
  join sep strings =
    case strings of
      [] ->
        empty
      str : strs ->
        runST $
          do
            let !len = List.foldl' (\w s -> w + 1 + size s) (size str) strs
            mba <- newByteArray len
            joinHelp sep mba 0 str strs
            freeze mba
  
  joinHelp :: Word8 -> MBA s -> Int -> Utf8 t -> [Utf8 t] -> ST s ()
  joinHelp sep mba offset str strings =
    let !len = size str
     in case strings of
          [] ->
            copy str 0 mba offset len
          s : ss ->
            do
              copy str 0 mba offset len
              let !dotOffset = offset + len
              writeWord8 mba dotOffset sep
              let !newOffset = dotOffset + 1
              joinHelp sep mba newOffset s ss
  
  -- EQUAL
  
  instance Eq (Utf8 t) where
    (==) (Utf8 ba1#) (Utf8 ba2#) =
      let !len1# = sizeofByteArray# ba1#
          !len2# = sizeofByteArray# ba2#
       in isTrue# (len1# ==# len2#)
            && isTrue# (0# ==# compareByteArrays# ba1# 0# ba2# 0# len1#)
  
  -- COMPARE
  --
  -- TODO: is it fine to sort by length and only compare bytes on length ties?
  --
  
  instance Ord (Utf8 t) where
    compare (Utf8 ba1#) (Utf8 ba2#) =
      let !len1# = sizeofByteArray# ba1#
          !len2# = sizeofByteArray# ba2#
          !len# = if isTrue# (len1# <# len2#) then len1# else len2#
          !cmp# = compareByteArrays# ba1# 0# ba2# 0# len#
       in case () of
            _
              | isTrue# (cmp# <# 0#) -> LT
              | isTrue# (cmp# ># 0#) -> GT
              | isTrue# (len1# <# len2#) -> LT
              | isTrue# (len1# ># len2#) -> GT
              | True -> EQ
  
  -- FROM STRING
  
  fromChars :: [Char] -> Utf8 t
  fromChars chars =
    runST
      ( do
          mba <- newByteArray (sum (map getWidth chars))
          writeChars mba 0 chars
      )
  
  writeChars :: MBA s -> Int -> [Char] -> ST s (Utf8 t)
  writeChars !mba !offset chars =
    case chars of
      [] ->
        freeze mba
      char : chars
        | n < 0x80 ->
            do
              writeWord8 mba offset (fromIntegral n)
              writeChars mba (offset + 1) chars
        | n < 0x800 ->
            do
              writeWord8 mba offset (fromIntegral ((shiftR n 6) + 0xC0))
              writeWord8 mba (offset + 1) (fromIntegral ((n .&. 0x3F) + 0x80))
              writeChars mba (offset + 2) chars
        | n < 0x10000 ->
            do
              writeWord8 mba offset (fromIntegral ((shiftR n 12) + 0xE0))
              writeWord8 mba (offset + 1) (fromIntegral ((shiftR n 6 .&. 0x3F) + 0x80))
              writeWord8 mba (offset + 2) (fromIntegral ((n .&. 0x3F) + 0x80))
              writeChars mba (offset + 3) chars
        | otherwise ->
            do
              writeWord8 mba offset (fromIntegral ((shiftR n 18) + 0xF0))
              writeWord8 mba (offset + 1) (fromIntegral ((shiftR n 12 .&. 0x3F) + 0x80))
              writeWord8 mba (offset + 2) (fromIntegral ((shiftR n 6 .&. 0x3F) + 0x80))
              writeWord8 mba (offset + 3) (fromIntegral ((n .&. 0x3F) + 0x80))
              writeChars mba (offset + 4) chars
        where
          n = Char.ord char
  
  getWidth :: Char -> Int
  getWidth char
    | code < 0x80 = 1
    | code < 0x800 = 2
    | code < 0x10000 = 3
    | otherwise = 4
    where
      code = Char.ord char
  
  -- TO CHARS
  
  toChars :: Utf8 t -> [Char]
  toChars (Utf8 ba#) =
    toCharsHelp ba# 0# (sizeofByteArray# ba#)
  
  toCharsHelp :: ByteArray# -> Int# -> Int# -> [Char]
  toCharsHelp ba# offset# len# =
    if isTrue# (offset# >=# len#)
      then []
      else
        let !w# = word8ToWord# (indexWord8Array# ba# offset#)
            !(# char, width# #)
              | isTrue# (ltWord# w# 0xC0##) = (# C# (chr# (word2Int# w#)), 1# #)
              | isTrue# (ltWord# w# 0xE0##) = (# chr2 ba# offset# w#, 2# #)
              | isTrue# (ltWord# w# 0xF0##) = (# chr3 ba# offset# w#, 3# #)
              | True = (# chr4 ba# offset# w#, 4# #)
  
            !newOffset# = offset# +# width#
         in char : toCharsHelp ba# newOffset# len#
  
  chr2 :: ByteArray# -> Int# -> Word# -> Char
  chr2 ba# offset# firstWord# =
    let !i1# = word2Int# firstWord#
        !i2# = word8ToInt# (indexWord8Array# ba# (offset# +# 1#))
        !c1# = uncheckedIShiftL# (i1# -# 0xC0#) 6#
        !c2# = i2# -# 0x80#
     in C# (chr# (c1# +# c2#))
  
  chr3 :: ByteArray# -> Int# -> Word# -> Char
  chr3 ba# offset# firstWord# =
    let !i1# = word2Int# firstWord#
        !i2# = word8ToInt# (indexWord8Array# ba# (offset# +# 1#))
        !i3# = word8ToInt# (indexWord8Array# ba# (offset# +# 2#))
        !c1# = uncheckedIShiftL# (i1# -# 0xE0#) 12#
        !c2# = uncheckedIShiftL# (i2# -# 0x80#) 6#
        !c3# = i3# -# 0x80#
     in C# (chr# (c1# +# c2# +# c3#))
  
  chr4 :: ByteArray# -> Int# -> Word# -> Char
  chr4 ba# offset# firstWord# =
    let !i1# = word2Int# firstWord#
        !i2# = word8ToInt# (indexWord8Array# ba# (offset# +# 1#))
        !i3# = word8ToInt# (indexWord8Array# ba# (offset# +# 2#))
        !i4# = word8ToInt# (indexWord8Array# ba# (offset# +# 3#))
        !c1# = uncheckedIShiftL# (i1# -# 0xF0#) 18#
        !c2# = uncheckedIShiftL# (i2# -# 0x80#) 12#
        !c3# = uncheckedIShiftL# (i3# -# 0x80#) 6#
        !c4# = i4# -# 0x80#
     in C# (chr# (c1# +# c2# +# c3# +# c4#))
  
  word8ToInt# :: Word8# -> Int#
  word8ToInt# word8 =
    word2Int# (word8ToWord# word8)
  
  -- TO BYTESTRING
  
  toByteString :: Utf8 t -> B.ByteString
  toByteString bytes =
    LazyByteString.toStrict $ Builder.toLazyByteString $ toBuilder bytes
  
  -- TO BUILDER
  
  toBuilder :: Utf8 t -> B.Builder
  toBuilder bytes =
    B.builder (toBuilderHelp bytes)
  
  toBuilderHelp :: Utf8 t -> B.BuildStep a -> B.BuildStep a
  toBuilderHelp !bytes@(Utf8 ba#) k =
    go 0 (I# (sizeofByteArray# ba#))
    where
      go !offset !end !(B.BufferRange bOffset bEnd) =
        let !bLen = minusPtr bEnd bOffset
            !len = end - offset
         in if len <= bLen
              then do
                copyToPtr bytes offset bOffset len
                let !br' = B.BufferRange (plusPtr bOffset len) bEnd
                k br'
              else do
                copyToPtr bytes offset bOffset bLen
                let !offset' = offset + bLen
                return $ B.bufferFull 1 bEnd (go offset' end)
  
  -- TO ESCAPED BUILDER
  
  toEscapedBuilder :: Word8 -> Word8 -> Utf8 t -> B.Builder
  toEscapedBuilder before after =
    \name -> B.builder (toEscapedBuilderHelp before after name)
  
  toEscapedBuilderHelp :: Word8 -> Word8 -> Utf8 t -> B.BuildStep a -> B.BuildStep a
  toEscapedBuilderHelp before after !name@(Utf8 ba#) k =
    go 0 (I# (sizeofByteArray# ba#))
    where
      go !offset !len !(B.BufferRange bOffset bEnd) =
        let !bLen = minusPtr bEnd bOffset
         in if len <= bLen
              then do
                -- PERF test if writing word-by-word is faster
                copyToPtr name offset bOffset len
                escape before after bOffset name offset len 0
                let !newBufferRange = B.BufferRange (plusPtr bOffset len) bEnd
                k newBufferRange
              else do
                copyToPtr name offset bOffset bLen
                escape before after bOffset name offset bLen 0
                let !newOffset = offset + bLen
                let !newLength = len - bLen
                return $ B.bufferFull 1 bEnd (go newOffset newLength)
  
  escape :: Word8 -> Word8 -> Ptr a -> Utf8 t -> Int -> Int -> Int -> IO ()
  escape before@(W8# before#) after ptr name@(Utf8 ba#) offset@(I# offset#) len@(I# len#) i@(I# i#) =
    if isTrue# (i# <# len#)
      then
        if isTrue# (eqWord8# before# (indexWord8Array# ba# (offset# +# i#)))
          then do
            writeWordToPtr ptr i after
            escape before after ptr name offset len (i + 1)
          else do escape before after ptr name offset len (i + 1)
      else return ()
  
  -- FROM PTR
  
  fromPtr :: Ptr Word8 -> Ptr Word8 -> Utf8 t
  fromPtr pos end =
    unsafeDupablePerformIO
      ( stToIO
          ( do
              let !len = minusPtr end pos
              mba <- newByteArray len
              copyFromPtr pos mba 0 len
              freeze mba
          )
      )
  
  -- FROM SNIPPET
  
  fromSnippet :: P.Snippet -> Utf8 t
  fromSnippet (P.Snippet fptr off len _ _) =
    unsafeDupablePerformIO
      ( stToIO
          ( do
              mba <- newByteArray len
              let !pos = plusPtr (unsafeForeignPtrToPtr fptr) off
              copyFromPtr pos mba 0 len
              freeze mba
          )
      )
  
  -- BINARY
  
  putUnder256 :: Utf8 t -> Put
  putUnder256 bytes =
    do
      putWord8 (fromIntegral (size bytes))
      putBuilder (toBuilder bytes)
  
  getUnder256 :: Get (Utf8 t)
  getUnder256 =
    do
      word <- getWord8
      let !n = fromIntegral word
      readN n (copyFromByteString n)
  
  putVeryLong :: Utf8 t -> Put
  putVeryLong bytes =
    do
      put (size bytes)
      putBuilder (toBuilder bytes)
  
  getVeryLong :: Get (Utf8 t)
  getVeryLong =
    do
      n <- get
      if n > 0
        then readN n (copyFromByteString n)
        else return empty
  
  -- COPY FROM BYTESTRING
  
  copyFromByteString :: Int -> B.ByteString -> Utf8 t
  copyFromByteString len (B.PS fptr offset _) =
    unsafeDupablePerformIO
      ( do
          mba <- stToIO (newByteArray len)
          stToIO (copyFromPtr (unsafeForeignPtrToPtr fptr `plusPtr` offset) mba 0 len)
          touchForeignPtr fptr
          stToIO (freeze mba)
      )
  
  -- PRIMITIVES
  
  data MBA s
    = MBA# (MutableByteArray# s)
  
  newByteArray :: Int -> ST s (MBA s) -- PERF see if newPinnedByteArray for len > 256 is positive
  newByteArray (I# len#) =
    ST $ \s ->
      case newByteArray# len# s of
        (# s, mba# #) -> (# s, MBA# mba# #)
  
  freeze :: MBA s -> ST s (Utf8 t)
  freeze (MBA# mba#) =
    ST $ \s ->
      case unsafeFreezeByteArray# mba# s of
        (# s, ba# #) -> (# s, Utf8 ba# #)
  
  copy :: Utf8 t -> Int -> MBA s -> Int -> Int -> ST s ()
  copy (Utf8 ba#) (I# offset#) (MBA# mba#) (I# i#) (I# len#) =
    ST $ \s ->
      case copyByteArray# ba# offset# mba# i# len# s of
        s -> (# s, () #)
  
  copyFromPtr :: Ptr a -> MBA RealWorld -> Int -> Int -> ST RealWorld ()
  copyFromPtr (Ptr src#) (MBA# mba#) (I# offset#) (I# len#) =
    ST $ \s ->
      case copyAddrToByteArray# src# mba# offset# len# s of
        s -> (# s, () #)
  
  copyToPtr :: Utf8 t -> Int -> Ptr a -> Int -> IO ()
  copyToPtr (Utf8 ba#) (I# offset#) (Ptr mba#) (I# len#) =
    IO $ \s ->
      case copyByteArrayToAddr# ba# offset# mba# len# s of
        s -> (# s, () #)
  
  writeWord8 :: MBA s -> Int -> Word8 -> ST s ()
  writeWord8 (MBA# mba#) (I# offset#) (W8# w#) =
    ST $ \s ->
      case writeWord8Array# mba# offset# w# s of
        s -> (# s, () #)
  
  writeWordToPtr :: Ptr a -> Int -> Word8 -> IO ()
  writeWordToPtr (Ptr addr#) (I# offset#) (W8# word#) =
    IO $ \s ->
      case writeWord8OffAddr# addr# offset# word# s of
        s -> (# s, () #)
  