{-# LANGUAGE BangPatterns #-}

{- | RFC 4180 field scanner for the default CSV reader (Round-2 WS-D).
One pass over a strict 'BS.ByteString'; cassava-parity semantics pinned
by the @IO.CsvGolden@ suite: quotes open fields only at the first byte,
a double quote inside an unquoted field is an error, garbage after a
closing quote is an error, an unclosed quote consumes to EOF, and rows
end at @\\n@, @\\r\\n@ or a lone @\\r@.
-}
module DataFrame.IO.CSV.Internal.Scanner (
    Term,
    termSep,
    termEol,
    termEof,
    withField,
    withStripC8,
    withStripAscii,
    unescapeQuotes,
    isBlankRecordAt,
) where

import qualified Data.ByteString as BS
import qualified Data.ByteString.Internal as BSI
import qualified Data.ByteString.Unsafe as BSU

import Data.Word (Word8)
import Foreign.Ptr (Ptr)
import Foreign.Storable (pokeByteOff)

-- | What ended a field: another field follows, the row ended, or EOF.
type Term = Int

termSep, termEol, termEof :: Term
termSep :: Int
termSep = Int
0
termEol :: Int
termEol = Int
1
termEof :: Int
termEof = Int
2

quote, nl, cr :: Word8
quote :: Word8
quote = Word8
34
nl :: Word8
nl = Word8
10
cr :: Word8
cr = Word8
13

parseError :: Int -> String -> a
parseError :: forall a. Int -> String -> a
parseError Int
off String
what =
    String -> a
forall a. HasCallStack => String -> a
error (String
"CSV Parse Error: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
what String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" at byte " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
off)

{- | Scan one field starting at @pos@ and continue with
@k contentStart contentEnd needsUnescape term nextPos@. Quoted fields
yield the bytes between the outer quotes; @needsUnescape@ is set when
the content contains doubled quotes ('unescapeQuotes' collapses them).
-}
{-# INLINE withField #-}
withField ::
    BS.ByteString ->
    Int ->
    Word8 ->
    Int ->
    (Int -> Int -> Bool -> Term -> Int -> r) ->
    r
withField :: forall r.
ByteString
-> Int
-> Word8
-> Int
-> (Int -> Int -> Bool -> Int -> Int -> r)
-> r
withField ByteString
bs !Int
len !Word8
sep !Int
pos Int -> Int -> Bool -> Int -> Int -> r
k
    | Int
pos Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len Bool -> Bool -> Bool
&& ByteString -> Int -> Word8
BSU.unsafeIndex ByteString
bs Int
pos Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
quote = Int -> Bool -> r
quoted (Int
pos Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Bool
False
    | Bool
otherwise = Int -> r
unquoted Int
pos
  where
    -- Field terminator dispatch shared by both shapes.
    {-# INLINE finish #-}
    finish :: Int -> Int -> Bool -> Int -> r
finish !Int
cs !Int
ce !Bool
unesc !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len = Int -> Int -> Bool -> Int -> Int -> r
k Int
cs Int
ce Bool
unesc Int
termEof Int
i
        | Bool
otherwise = case ByteString -> Int -> Word8
BSU.unsafeIndex ByteString
bs Int
i of
            Word8
b
                | Word8
b Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
sep -> Int -> Int -> Bool -> Int -> Int -> r
k Int
cs Int
ce Bool
unesc Int
termSep (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                | Word8
b Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
nl -> Int -> Int -> Bool -> Int -> Int -> r
k Int
cs Int
ce Bool
unesc Int
termEol (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                | Word8
b Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
cr ->
                    if Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len Bool -> Bool -> Bool
&& ByteString -> Int -> Word8
BSU.unsafeIndex ByteString
bs (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
nl
                        then Int -> Int -> Bool -> Int -> Int -> r
k Int
cs Int
ce Bool
unesc Int
termEol (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)
                        else Int -> Int -> Bool -> Int -> Int -> r
k Int
cs Int
ce Bool
unesc Int
termEol (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                | Bool
otherwise ->
                    Int -> String -> r
forall a. Int -> String -> a
parseError Int
i String
"malformed quoted field (garbage after closing quote)"
    unquoted :: Int -> r
unquoted !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len = Int -> Int -> Bool -> Int -> Int -> r
k Int
pos Int
i Bool
False Int
termEof Int
i
        | Bool
otherwise = case ByteString -> Int -> Word8
BSU.unsafeIndex ByteString
bs Int
i of
            Word8
b
                | Word8
b Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
sep -> Int -> Int -> Bool -> Int -> Int -> r
k Int
pos Int
i Bool
False Int
termSep (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                | Word8
b Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
nl -> Int -> Int -> Bool -> Int -> Int -> r
k Int
pos Int
i Bool
False Int
termEol (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                | Word8
b Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
cr ->
                    if Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len Bool -> Bool -> Bool
&& ByteString -> Int -> Word8
BSU.unsafeIndex ByteString
bs (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
nl
                        then Int -> Int -> Bool -> Int -> Int -> r
k Int
pos Int
i Bool
False Int
termEol (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)
                        else Int -> Int -> Bool -> Int -> Int -> r
k Int
pos Int
i Bool
False Int
termEol (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                | Word8
b Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
quote -> Int -> String -> r
forall a. Int -> String -> a
parseError Int
i String
"stray double quote in unquoted field"
                | Bool
otherwise -> Int -> r
unquoted (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    quoted :: Int -> Bool -> r
quoted !Int
i !Bool
unesc
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len = Int -> Int -> Bool -> Int -> Int -> r
k (Int
pos Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
len Bool
unesc Int
termEof Int
len -- unclosed: rest of input
        | ByteString -> Int -> Word8
BSU.unsafeIndex ByteString
bs Int
i Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
quote =
            if Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len Bool -> Bool -> Bool
&& ByteString -> Int -> Word8
BSU.unsafeIndex ByteString
bs (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
quote
                then Int -> Bool -> r
quoted (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2) Bool
True
                else Int -> Int -> Bool -> Int -> r
finish (Int
pos Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
i Bool
unesc (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        | Bool
otherwise = Int -> Bool -> r
quoted (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Bool
unesc

{- | Whether the record starting at @pos@ is a blank record (a single
empty field, quoted or not, then end of row) — cassava drops these.
Assumes @pos < len@. Errors propagate from 'withField'.
-}
{-# INLINE isBlankRecordAt #-}
isBlankRecordAt :: BS.ByteString -> Int -> Word8 -> Int -> Bool
isBlankRecordAt :: ByteString -> Int -> Word8 -> Int -> Bool
isBlankRecordAt ByteString
bs !Int
len !Word8
sep !Int
pos =
    ByteString
-> Int
-> Word8
-> Int
-> (Int -> Int -> Bool -> Int -> Int -> Bool)
-> Bool
forall r.
ByteString
-> Int
-> Word8
-> Int
-> (Int -> Int -> Bool -> Int -> Int -> r)
-> r
withField ByteString
bs Int
len Word8
sep Int
pos (\Int
cs Int
ce Bool
_ Int
term Int
_ -> Int
cs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
ce Bool -> Bool -> Bool
&& Int
term Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
termSep)

{- | Collapse doubled quotes (@\"\"@ → @\"@) in @[s, e)@ into a fresh strict
'BS.ByteString'. Every quote in the range is the first of a pair.
-}
unescapeQuotes :: BS.ByteString -> Int -> Int -> BS.ByteString
unescapeQuotes :: ByteString -> Int -> Int -> ByteString
unescapeQuotes ByteString
bs !Int
s !Int
e = Int -> (Ptr Word8 -> IO Int) -> ByteString
BSI.unsafeCreateUptoN (Int
e Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
s) (Int -> Int -> Ptr Word8 -> IO Int
go Int
s Int
0)
  where
    go :: Int -> Int -> Ptr Word8 -> IO Int
go !Int
i !Int
j !Ptr Word8
dst
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
e = Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
j
        | Bool
otherwise = do
            let w :: Word8
w = ByteString -> Int -> Word8
BSU.unsafeIndex ByteString
bs Int
i
            Ptr Word8 -> Int -> Word8 -> IO ()
forall b. Ptr b -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff (Ptr Word8
dst :: Ptr Word8) Int
j Word8
w
            if Word8
w Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
quote then Int -> Int -> Ptr Word8 -> IO Int
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2) (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Ptr Word8
dst else Int -> Int -> Ptr Word8 -> IO Int
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Ptr Word8
dst

{- | Strip with @Data.ByteString.Char8.strip@ semantics (Latin-1
whitespace: HT..CR, space and NBSP 0xA0) and continue with the
stripped range.
-}
{-# INLINE withStripC8 #-}
withStripC8 :: BS.ByteString -> Int -> Int -> (Int -> Int -> r) -> r
withStripC8 :: forall r. ByteString -> Int -> Int -> (Int -> Int -> r) -> r
withStripC8 ByteString
bs !Int
s0 !Int
e0 Int -> Int -> r
k = Int -> Int -> r
k Int
s Int
e
  where
    isC8Space :: a -> Bool
isC8Space a
w = a
w a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
32 Bool -> Bool -> Bool
|| (a
w a -> a -> a
forall a. Num a => a -> a -> a
- a
9) a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
4 Bool -> Bool -> Bool
|| a
w a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
160
    s :: Int
s = Int -> Int
skipF Int
s0
    e :: Int
e = Int -> Int
skipB Int
e0
    skipF :: Int -> Int
skipF !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
e0 Bool -> Bool -> Bool
&& Word8 -> Bool
forall {a}. (Num a, Ord a) => a -> Bool
isC8Space (ByteString -> Int -> Word8
BSU.unsafeIndex ByteString
bs Int
i) = Int -> Int
skipF (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        | Bool
otherwise = Int
i
    skipB :: Int -> Int
skipB !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
s Bool -> Bool -> Bool
&& Word8 -> Bool
forall {a}. (Num a, Ord a) => a -> Bool
isC8Space (ByteString -> Int -> Word8
BSU.unsafeIndex ByteString
bs (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) = Int -> Int
skipB (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
        | Bool
otherwise = Int
i

{- | Strip ASCII whitespace (the ASCII subset of @Data.Text.strip@) and
continue with the stripped range. Callers must fall back to a full
'Data.Text.strip' when the stripped range still has non-ASCII edges.
-}
{-# INLINE withStripAscii #-}
withStripAscii :: BS.ByteString -> Int -> Int -> (Int -> Int -> r) -> r
withStripAscii :: forall r. ByteString -> Int -> Int -> (Int -> Int -> r) -> r
withStripAscii ByteString
bs !Int
s0 !Int
e0 Int -> Int -> r
k = Int -> Int -> r
k Int
s Int
e
  where
    isAsciiSpace :: a -> Bool
isAsciiSpace a
w = a
w a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
32 Bool -> Bool -> Bool
|| (a
w a -> a -> a
forall a. Num a => a -> a -> a
- a
9) a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
4
    s :: Int
s = Int -> Int
skipF Int
s0
    e :: Int
e = Int -> Int
skipB Int
e0
    skipF :: Int -> Int
skipF !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
e0 Bool -> Bool -> Bool
&& Word8 -> Bool
forall {a}. (Num a, Ord a) => a -> Bool
isAsciiSpace (ByteString -> Int -> Word8
BSU.unsafeIndex ByteString
bs Int
i) = Int -> Int
skipF (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        | Bool
otherwise = Int
i
    skipB :: Int -> Int
skipB !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
s Bool -> Bool -> Bool
&& Word8 -> Bool
forall {a}. (Num a, Ord a) => a -> Bool
isAsciiSpace (ByteString -> Int -> Word8
BSU.unsafeIndex ByteString
bs (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) = Int -> Int
skipB (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
        | Bool
otherwise = Int
i