{-# LANGUAGE BangPatterns #-}

{- | UTF-8 validation and @decodeUtf8Lenient@-parity slice decoding used by
'DataFrame.Internal.ColumnBuilder' to turn shared byte buffers into 'Text'.
-}
module DataFrame.Internal.Utf8 (
    isValidUtf8Slice,
    isUtf8Boundary,
    lenientDecodeSlice,
    sliceTextVector,
) where

import qualified Data.Text as T
import qualified Data.Text.Array as A
import qualified Data.Vector as VB
import qualified Data.Vector.Mutable as VBM
import qualified Data.Vector.Unboxed as VU

import Data.Text.Internal (Text (..))
import Data.Text.Internal.Encoding.Utf8 (
    DecoderResult (..),
    utf8DecodeContinue,
    utf8DecodeStart,
 )
import Data.Text.Internal.Validate (isValidUtf8ByteArray)
import Data.Word (Word8)

-- | Whether @len@ bytes starting at @off@ are well-formed UTF-8.
isValidUtf8Slice :: A.Array -> Int -> Int -> Bool
isValidUtf8Slice :: Array -> Int -> Int -> Bool
isValidUtf8Slice = Array -> Int -> Int -> Bool
isValidUtf8ByteArray
{-# INLINE isValidUtf8Slice #-}

{- | Whether a byte may start a code point (i.e. is not a continuation
byte). Field slices of a valid buffer are themselves valid iff every
field starts on a boundary.
-}
isUtf8Boundary :: Word8 -> Bool
isUtf8Boundary :: Word8 -> Bool
isUtf8Boundary Word8
w = Word8
w Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
< Word8
0x80 Bool -> Bool -> Bool
|| Word8
w Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
>= Word8
0xC0
{-# INLINE isUtf8Boundary #-}

{- | Decode a byte slice exactly like @decodeUtf8Lenient@: greedy decode at
each position; any byte that cannot begin a complete, valid sequence within
the slice becomes one U+FFFD and decoding resumes at the next byte.
-}
lenientDecodeSlice :: A.Array -> Int -> Int -> T.Text
lenientDecodeSlice :: Array -> Int -> Int -> Text
lenientDecodeSlice Array
arr Int
off Int
len = String -> Text
T.pack (Int -> String
go Int
off)
  where
    !end :: Int
end = Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len
    go :: Int -> String
go !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
end = []
        | Bool
otherwise = case Int -> Maybe (Char, Int)
tryDecode Int
i of
            Just (Char
c, Int
i') -> Char
c Char -> String -> String
forall a. a -> [a] -> [a]
: Int -> String
go Int
i'
            Maybe (Char, Int)
Nothing -> Char
'\xFFFD' Char -> String -> String
forall a. a -> [a] -> [a]
: Int -> String
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    tryDecode :: Int -> Maybe (Char, Int)
tryDecode !Int
i = DecoderResult -> Int -> Maybe (Char, Int)
loop (Word8 -> DecoderResult
utf8DecodeStart (Array -> Int -> Word8
A.unsafeIndex Array
arr Int
i)) (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
      where
        loop :: DecoderResult -> Int -> Maybe (Char, Int)
loop (Accept Char
c) !Int
j = (Char, Int) -> Maybe (Char, Int)
forall a. a -> Maybe a
Just (Char
c, Int
j)
        loop DecoderResult
Reject Int
_ = Maybe (Char, Int)
forall a. Maybe a
Nothing
        loop (Incomplete DecoderState
st CodePoint
cp) !Int
j
            | Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
end = Maybe (Char, Int)
forall a. Maybe a
Nothing
            | Bool
otherwise = DecoderResult -> Int -> Maybe (Char, Int)
loop (Word8 -> DecoderState -> CodePoint -> DecoderResult
utf8DecodeContinue (Array -> Int -> Word8
A.unsafeIndex Array
arr Int
j) DecoderState
st CodePoint
cp) (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

{- | Slice forced 'Text' values off a shared array; row @i@ spans bytes
@[offs!i, offs!(i+1))@. Fast path validates the whole span once when every field
starts on a code-point boundary; else per-field validation with lenient decode.
-}
sliceTextVector :: A.Array -> VU.Vector Int -> VB.Vector T.Text
sliceTextVector :: Array -> Vector Int -> Vector Text
sliceTextVector Array
arr Vector Int
offs = (forall s. ST s (MVector s Text)) -> Vector Text
forall a. (forall s. ST s (MVector s a)) -> Vector a
VB.create ((forall s. ST s (MVector s Text)) -> Vector Text)
-> (forall s. ST s (MVector s Text)) -> Vector Text
forall a b. (a -> b) -> a -> b
$ do
    MVector s Text
mv <- Int -> ST s (MVector (PrimState (ST s)) Text)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MVector (PrimState m) a)
VBM.unsafeNew Int
n
    let fill :: (Int -> Int -> Text) -> f ()
fill Int -> Int -> Text
dec = Int -> f ()
forall {f :: * -> *}. (PrimState f ~ s, PrimMonad f) => Int -> f ()
go Int
0
          where
            go :: Int -> f ()
go !Int
i
                | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = () -> f ()
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                | Bool
otherwise = do
                    let o :: Int
o = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
offs Int
i
                        !t :: Text
t = Int -> Int -> Text
dec Int
o (Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
offs (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
o)
                    MVector (PrimState f) Text -> Int -> Text -> f ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
VBM.unsafeWrite MVector s Text
MVector (PrimState f) Text
mv Int
i Text
t
                    Int -> f ()
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    if Bool
fast then (Int -> Int -> Text) -> ST s ()
forall {f :: * -> *}.
(PrimState f ~ s, PrimMonad f) =>
(Int -> Int -> Text) -> f ()
fill Int -> Int -> Text
mkSlice else (Int -> Int -> Text) -> ST s ()
forall {f :: * -> *}.
(PrimState f ~ s, PrimMonad f) =>
(Int -> Int -> Text) -> f ()
fill Int -> Int -> Text
decodeField
    MVector s Text -> ST s (MVector s Text)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s Text
mv
  where
    n :: Int
n = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
offs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
    base :: Int
base = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
offs Int
0
    used :: Int
used = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
offs Int
n
    boundariesOk :: Int -> Bool
boundariesOk !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = Bool
True
        | Bool
otherwise =
            let o :: Int
o = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
offs Int
i
             in (Int
o Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
used Bool -> Bool -> Bool
|| Word8 -> Bool
isUtf8Boundary (Array -> Int -> Word8
A.unsafeIndex Array
arr Int
o))
                    Bool -> Bool -> Bool
&& Int -> Bool
boundariesOk (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    fast :: Bool
fast = Array -> Int -> Int -> Bool
isValidUtf8Slice Array
arr Int
base (Int
used Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
base) Bool -> Bool -> Bool
&& Int -> Bool
boundariesOk Int
0
    mkSlice :: Int -> Int -> Text
mkSlice Int
o Int
l = if Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Text
T.empty else Array -> Int -> Int -> Text
Text Array
arr Int
o Int
l
    decodeField :: Int -> Int -> Text
decodeField Int
o Int
l
        | Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Text
T.empty
        | Array -> Int -> Int -> Bool
isValidUtf8Slice Array
arr Int
o Int
l = Array -> Int -> Int -> Text
Text Array
arr Int
o Int
l
        | Bool
otherwise = Array -> Int -> Int -> Text
lenientDecodeSlice Array
arr Int
o Int
l