{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}

{- | Shared byte-level helpers for the fast slice parsers. All operate on
@(buf, start, end)@ slices; the caller guarantees
@0 <= start <= end <= length buf@ so everything below uses 'unsafeIndex'.
-}
module DataFrame.Internal.Parsing.Fast.Common (
    isStripByte,
    isDigitByte,
    skipStrip,
    skipStripEnd,
    skipZeroes,
    takeDigits64,
) where

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

import Data.Word (Word64, Word8)
import GHC.Exts (
    Int (..),
    Int#,
    RuntimeRep,
    TYPE,
    Word64#,
    isTrue#,
    plusWord64#,
    timesWord64#,
    wordToWord64#,
    (+#),
    (>=#),
 )
import GHC.Word (Word64 (..))

{- | Bytes removed by 'Data.ByteString.Char8.strip' (Latin-1 'isSpace'):
HT LF VT FF CR SP and NBSP 0xA0. Probed over all 256 bytes; parity
with the strip-based reference parsers depends on this exact set.
-}
isStripByte :: Word8 -> Bool
isStripByte :: Word8 -> Bool
isStripByte Word8
w = Word8
w Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x20 Bool -> Bool -> Bool
|| (Word8
w Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
- Word8
0x09) Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word8
4 Bool -> Bool -> Bool
|| Word8
w Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0xA0
{-# INLINE isStripByte #-}

isDigitByte :: Word8 -> Bool
isDigitByte :: Word8 -> Bool
isDigitByte Word8
w = (Word8
w Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
- Word8
0x30) Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word8
9
{-# INLINE isDigitByte #-}

-- | Index of the first non-strip byte in @[i, end)@.
skipStrip :: BS.ByteString -> Int -> Int -> Int
skipStrip :: ByteString -> Int -> Int -> Int
skipStrip ByteString
bs = Int -> Int -> Int
go
  where
    go :: Int -> Int -> Int
go !Int
i !Int
end
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
end Bool -> Bool -> Bool
&& Word8 -> Bool
isStripByte (ByteString -> Int -> Word8
BSU.unsafeIndex ByteString
bs Int
i) = Int -> Int -> Int
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
end
        | Bool
otherwise = Int
i
{-# INLINE skipStrip #-}

-- | New exclusive end after dropping trailing strip bytes in @[i, end)@.
skipStripEnd :: BS.ByteString -> Int -> Int -> Int
skipStripEnd :: ByteString -> Int -> Int -> Int
skipStripEnd ByteString
bs = Int -> Int -> Int
go
  where
    go :: Int -> Int -> Int
go !Int
i !Int
end
        | Int
end Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
i Bool -> Bool -> Bool
&& Word8 -> Bool
isStripByte (ByteString -> Int -> Word8
BSU.unsafeIndex ByteString
bs (Int
end Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) = Int -> Int -> Int
go Int
i (Int
end Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
        | Bool
otherwise = Int
end
{-# INLINE skipStripEnd #-}

-- | Index of the first non-@\'0\'@ byte in @[i, end)@.
skipZeroes :: BS.ByteString -> Int -> Int -> Int
skipZeroes :: ByteString -> Int -> Int -> Int
skipZeroes ByteString
bs = Int -> Int -> Int
go
  where
    go :: Int -> Int -> Int
go !Int
i !Int
end
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
end Bool -> Bool -> Bool
&& ByteString -> Int -> Word8
BSU.unsafeIndex ByteString
bs Int
i Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x30 = Int -> Int -> Int
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
end
        | Bool
otherwise = Int
i
{-# INLINE skipZeroes #-}

{- | Consume ASCII digits from @i@, passing the stop index and the wrapping
'Word64' accumulation to the continuation. Callers must bound the
significant digit count before trusting the value.
-}
takeDigits64 ::
    forall (rep :: RuntimeRep) (r :: TYPE rep).
    BS.ByteString ->
    Int ->
    Int ->
    (Int -> Word64 -> r) ->
    r
takeDigits64 :: forall r. ByteString -> Int -> Int -> (Int -> Word64 -> r) -> r
takeDigits64 ByteString
bs (I# Int#
i0) (I# Int#
end) Int -> Word64 -> r
k = Int# -> Word64# -> r
go Int#
i0 (Word# -> Word64#
wordToWord64# Word#
0##)
  where
    go :: Int# -> Word64# -> r
    go :: Int# -> Word64# -> r
go Int#
i Word64#
acc
        | Int# -> Bool
isTrue# (Int#
i Int# -> Int# -> Int#
>=# Int#
end) = Int -> Word64 -> r
k (Int# -> Int
I# Int#
i) (Word64# -> Word64
W64# Word64#
acc)
        | Word8 -> Bool
isDigitByte Word8
w =
            case Word8 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8
w Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
- Word8
0x30) :: Word64 of
                W64# Word64#
d ->
                    Int# -> Word64# -> r
go
                        (Int#
i Int# -> Int# -> Int#
+# Int#
1#)
                        ((Word64#
acc Word64# -> Word64# -> Word64#
`timesWord64#` Word# -> Word64#
wordToWord64# Word#
10##) Word64# -> Word64# -> Word64#
`plusWord64#` Word64#
d)
        | Bool
otherwise = Int -> Word64 -> r
k (Int# -> Int
I# Int#
i) (Word64# -> Word64
W64# Word64#
acc)
      where
        w :: Word8
w = ByteString -> Int -> Word8
BSU.unsafeIndex ByteString
bs (Int# -> Int
I# Int#
i)
{-# INLINE takeDigits64 #-}