{-# LANGUAGE BangPatterns #-}

-- | Fast imperative CSV tokenizer over strict 'ByteString'.
--
-- Complementary to the combinator path in "Circuit.Parser.Csv": a single
-- 'unsafeIndex' pass emitting zero-copy slices. Commas are implied between
-- adjacent field tokens; 'CRowEnd' marks record boundaries. 'CQuoted'
-- payloads keep their doubled-quote escapes intact.
module Circuit.Parser.Csv.Lexer
  ( -- * Tokens
    CsvToken (..),

    -- * Running
    runCsvLexerBS,
  )
where

import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.ByteString.Unsafe qualified as BSU
import Data.Word (Word8)

-- $setup
-- >>> import Data.ByteString.Char8 qualified as C

-- | A flat CSV token. Payloads are zero-copy slices of the input.
data CsvToken
  = CField ByteString
  | CQuoted ByteString
  | CRowEnd
  deriving (CsvToken -> CsvToken -> Bool
(CsvToken -> CsvToken -> Bool)
-> (CsvToken -> CsvToken -> Bool) -> Eq CsvToken
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: CsvToken -> CsvToken -> Bool
== :: CsvToken -> CsvToken -> Bool
$c/= :: CsvToken -> CsvToken -> Bool
/= :: CsvToken -> CsvToken -> Bool
Eq, Int -> CsvToken -> ShowS
[CsvToken] -> ShowS
CsvToken -> [Char]
(Int -> CsvToken -> ShowS)
-> (CsvToken -> [Char]) -> ([CsvToken] -> ShowS) -> Show CsvToken
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> CsvToken -> ShowS
showsPrec :: Int -> CsvToken -> ShowS
$cshow :: CsvToken -> [Char]
show :: CsvToken -> [Char]
$cshowList :: [CsvToken] -> ShowS
showList :: [CsvToken] -> ShowS
Show)

-- | Tokenize a strict 'ByteString' as CSV.
--
-- On failure, returns the byte offset and a message.
--
-- >>> runCsvLexerBS (C.pack "a,b\r\n\"x\"\"y\",z")
-- Right [CField "a",CField "b",CRowEnd,CQuoted "x\"\"y",CField "z"]
--
-- >>> runCsvLexerBS (C.pack "a,\"b")
-- Left (4,"unterminated quoted field")
--
-- >>> runCsvLexerBS (C.pack "ab\"cd")
-- Left (2,"unexpected quote in plain field")
runCsvLexerBS :: ByteString -> Either (Int, String) [CsvToken]
runCsvLexerBS :: ByteString -> Either (Int, [Char]) [CsvToken]
runCsvLexerBS ByteString
bs = Int -> Either (Int, [Char]) [CsvToken]
fieldAt Int
0
  where
    !len :: Int
len = ByteString -> Int
BS.length ByteString
bs
    fieldAt :: Int -> Either (Int, [Char]) [CsvToken]
fieldAt !Int
i
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len = [CsvToken] -> Either (Int, [Char]) [CsvToken]
forall a b. b -> Either a b
Right []
      | ByteString -> Int -> Word8
BSU.unsafeIndex ByteString
bs Int
i Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
quote = Int -> Either (Int, [Char]) [CsvToken]
quotedAt Int
i
      | Bool
otherwise = Int -> Either (Int, [Char]) [CsvToken]
plainAt Int
i
    plainAt :: Int -> Either (Int, [Char]) [CsvToken]
plainAt !Int
i =
      let !j :: Int
j = Int -> Int
plainEnd Int
i
       in if Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len Bool -> Bool -> Bool
&& ByteString -> Int -> Word8
BSU.unsafeIndex ByteString
bs Int
j Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
quote
            then (Int, [Char]) -> Either (Int, [Char]) [CsvToken]
forall a b. a -> Either a b
Left (Int
j, [Char]
"unexpected quote in plain field")
            else (ByteString -> CsvToken
CField (Int -> Int -> ByteString
slice Int
i Int
j) CsvToken -> [CsvToken] -> [CsvToken]
forall a. a -> [a] -> [a]
:) ([CsvToken] -> [CsvToken])
-> Either (Int, [Char]) [CsvToken]
-> Either (Int, [Char]) [CsvToken]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Either (Int, [Char]) [CsvToken]
sep Int
j
    plainEnd :: Int -> Int
plainEnd !Int
i
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len = Int
i
      | Bool
otherwise = case ByteString -> Int -> Word8
BSU.unsafeIndex ByteString
bs Int
i of
          Word8
0x2C -> Int
i -- ,
          Word8
0x0D -> Int
i -- CR
          Word8
0x0A -> Int
i -- LF
          Word8
0x22 -> Int
i -- "
          Word8
_ -> Int -> Int
plainEnd (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    quotedAt :: Int -> Either (Int, [Char]) [CsvToken]
quotedAt !Int
i = case Int -> Maybe Int
quotedEnd (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) of
      Maybe Int
Nothing -> (Int, [Char]) -> Either (Int, [Char]) [CsvToken]
forall a b. a -> Either a b
Left (Int
len, [Char]
"unterminated quoted field")
      Just Int
j -> (ByteString -> CsvToken
CQuoted (Int -> Int -> ByteString
slice (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
j) CsvToken -> [CsvToken] -> [CsvToken]
forall a. a -> [a] -> [a]
:) ([CsvToken] -> [CsvToken])
-> Either (Int, [Char]) [CsvToken]
-> Either (Int, [Char]) [CsvToken]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Either (Int, [Char]) [CsvToken]
sep (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    -- index of the closing quote (doubled "" is an escape), or end
    quotedEnd :: Int -> Maybe Int
quotedEnd !Int
i
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len = Maybe Int
forall a. Maybe a
Nothing
      | Bool
otherwise = case ByteString -> Int -> Word8
BSU.unsafeIndex ByteString
bs Int
i of
          Word8
0x22 ->
            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 -> Maybe Int
quotedEnd (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)
              else Int -> Maybe Int
forall a. a -> Maybe a
Just Int
i
          Word8
_ -> Int -> Maybe Int
quotedEnd (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    -- after a field: comma, line ending, or end of input
    sep :: Int -> Either (Int, [Char]) [CsvToken]
sep !Int
i
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len = [CsvToken] -> Either (Int, [Char]) [CsvToken]
forall a b. b -> Either a b
Right []
      | Bool
otherwise = case ByteString -> Int -> Word8
BSU.unsafeIndex ByteString
bs Int
i of
          Word8
0x2C -> Int -> Either (Int, [Char]) [CsvToken]
fieldAt (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
          Word8
0x0D -> (CsvToken
CRowEnd CsvToken -> [CsvToken] -> [CsvToken]
forall a. a -> [a] -> [a]
:) ([CsvToken] -> [CsvToken])
-> Either (Int, [Char]) [CsvToken]
-> Either (Int, [Char]) [CsvToken]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Either (Int, [Char]) [CsvToken]
fieldAt (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
0x0A then Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2 else Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
          Word8
0x0A -> (CsvToken
CRowEnd CsvToken -> [CsvToken] -> [CsvToken]
forall a. a -> [a] -> [a]
:) ([CsvToken] -> [CsvToken])
-> Either (Int, [Char]) [CsvToken]
-> Either (Int, [Char]) [CsvToken]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Either (Int, [Char]) [CsvToken]
fieldAt (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
          Word8
w -> (Int, [Char]) -> Either (Int, [Char]) [CsvToken]
forall a b. a -> Either a b
Left (Int
i, [Char]
"unexpected byte " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Word8 -> [Char]
forall a. Show a => a -> [Char]
show Word8
w)
    slice :: Int -> Int -> ByteString
slice Int
s Int
e = Int -> ByteString -> ByteString
BSU.unsafeTake (Int
e Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
s) (Int -> ByteString -> ByteString
BSU.unsafeDrop Int
s ByteString
bs)

quote :: Word8
quote :: Word8
quote = Word8
0x22