Core.hs

{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE NoImplicitPrelude #-}

-- | Shared vocabulary of the streaming JSON decoder.
module HQ.JSON.Decoder.Core where

import Data.Char (digitToInt, isHexDigit)
import Data.Text qualified as T
import HQ.JSON.Decoder.Error (DecodeError)
import HQ.JSON.Event (JSONEvent (..))
import Relude hiding (Compose, id, many, some, state)
import Streaming (Of, Stream)

-- | Buffered string fragments (reversed for O(1) prepend).
newtype StringBuffer = StringBuffer [Text] deriving (Eq)

emptyStringBuffer :: StringBuffer
emptyStringBuffer = StringBuffer []

appendStringBuffer :: Text -> StringBuffer -> StringBuffer
appendStringBuffer t (StringBuffer ts)
  | T.null t = StringBuffer ts
  | otherwise = StringBuffer (t : ts)

appendCharStringBuffer :: Char -> StringBuffer -> StringBuffer
appendCharStringBuffer c = appendStringBuffer (T.singleton c)

finishStringBuffer :: StringBuffer -> Text
finishStringBuffer (StringBuffer ts) = T.concat (reverse ts)

-- | Nesting depth: the number of enclosing containers.
newtype NestDepth = NestDepth Int deriving (Eq, Ord)

initialDepth :: NestDepth
initialDepth = NestDepth 0

deeper :: NestDepth -> NestDepth
deeper (NestDepth n) = NestDepth (n + 1)

shallower :: NestDepth -> NestDepth
shallower (NestDepth n) = NestDepth (n - 1)

data DecoderState = DecoderState
  { decoderInput :: Text,
    decoderStack :: [DecodeContext],
    decoderNestDepth :: NestDepth,
    decoderPhase :: DecoderPhase
  }

data DecoderPhase
  = DecoderPhaseValue
  | DecoderPhaseObjectKey
  | DecoderPhaseObjectColon
  | DecoderPhaseObjectComma
  | DecoderPhaseArrayComma
  | DecoderPhaseString StringState
  | DecoderPhaseNumber NumberState
  | DecoderPhaseKeyword KeywordState
  | DecoderPhaseFinished
  deriving (Eq)

data DecodeContext = DecodeArray | DecodeObject

data NumberState = NumberState
  { numberNegative :: Bool,
    numberCoeff :: Integer,
    -- | Exponent digits accumulated so far.
    numberExp :: Integer,
    -- | Fraction digits folded into 'numberCoeff'.
    numberFracDigits :: Int,
    -- | '-' seen after 'e'/'E'.
    numberExpNeg :: Bool,
    numberPhase :: NumberPhase,
    -- Error payloads only: retained, never decoded, on the success path.
    numberFrags :: [Text]
  }
  deriving (Eq)

-- | Tracks how much of a JSON number has been consumed.
data NumberPhase
  = -- | Saw leading '-'
    NumberSign
  | -- | Saw '0' (must not be followed by more digits)
    NumberZero
  | -- | Saw [1-9] and possibly more digits
    NumberNonZero
  | -- | Saw '.' after integer part (need fractional digits)
    NumberDot
  | -- | Saw '.' + at least one digit
    NumberFraction
  | -- | Saw 'e'/'E' (waiting for optional sign or digit)
    NumberExpSign
  | -- | Saw 'e'/'E' then '+'/'-' (need digits)
    NumberExpAfterSign
  | -- | Saw exponent digits
    NumberExpDigit
  deriving (Eq)

data NumberStep
  = NumberStep NumberPhase
  | NumberEnd
  | NumberError

data StringTarget = StringKey | StringValue
  deriving (Eq)

data BufferedStringTarget = BufferedStringTarget StringTarget StringBuffer
  deriving (Eq)

data Unicode = Unicode {unicodeValue :: Int, unicodeDigits :: Int}
  deriving (Eq)

data StringState
  = InString BufferedStringTarget
  | AfterEscape BufferedStringTarget
  | InUnicodeEscape BufferedStringTarget Unicode
  | AfterHighSurrogate BufferedStringTarget Int
  | -- Mid low-surrogate digits; separate so split escapes resume in the low half.
    InLowSurrogateEscape BufferedStringTarget Int Unicode
  | -- Low escape's backslash at a chunk end; next chunk must start with @u@.
    AfterLowBackslash BufferedStringTarget Int
  deriving (Eq)

data KeywordState
  = KeywordNull Int
  | KeywordTrue Int
  | KeywordFalse Int
  deriving (Eq)

data DecoderResult
  = Emit JSONEvent DecoderState
  | NeedInput DecoderState
  | Done DecoderState

-- | Emit one event, continuing after the value (container ends pop
-- and mark the parent via 'finishValue'). Single policy for scalars
-- and container ends. Inlined into the call sites.
{-# INLINE emitEvent #-}
emitEvent :: JSONEvent -> Text -> DecoderState -> Either DecodeError DecoderResult
emitEvent event remaining decoder =
  let newDec = decoder {decoderInput = remaining}
   in Right $ Emit event $ finishValue newDec

finishValue :: DecoderState -> DecoderState
finishValue decoder = case decoderStack decoder of
  [] -> decoder {decoderPhase = DecoderPhaseFinished}
  DecodeArray : _ -> decoder {decoderPhase = DecoderPhaseArrayComma}
  DecodeObject : _ -> decoder {decoderPhase = DecoderPhaseObjectComma}

isWhitespace :: Char -> Bool
isWhitespace c = c `elem` [' ', '\t', '\n', '\r']

isJsonDelimiter :: Char -> Bool
isJsonDelimiter c = isWhitespace c || c `elem` [',', ']', '}']

isHighSurrogate :: Int -> Bool
isHighSurrogate x = x >= 0xD800 && x <= 0xDBFF

isLowSurrogate :: Int -> Bool
isLowSurrogate x = x >= 0xDC00 && x <= 0xDFFF

-- | Verbatim string bytes; shared so decoding, skipping and collection
-- split identically.
isStringChar :: Char -> Bool
isStringChar c = c /= '"' && c /= '\\' && ord c >= 0x20

-- | Simple (non-\u) escapes in the JSON string table.
isSimpleEscape :: Char -> Bool
isSimpleEscape c = c `elem` ['"', '\\', '/', 'b', 'f', 'n', 'r', 't']

-- | Decode a simple escape to its character. Caller must have checked
-- 'isSimpleEscape' (except 'u', handled separately).
decodeSimpleEscape :: Char -> Char
decodeSimpleEscape '"' = '"'
decodeSimpleEscape '\\' = '\\'
decodeSimpleEscape '/' = '/'
decodeSimpleEscape 'b' = '\b'
decodeSimpleEscape 'f' = '\f'
decodeSimpleEscape 'n' = '\n'
decodeSimpleEscape 'r' = '\r'
decodeSimpleEscape 't' = '\t'
decodeSimpleEscape c = c

-- | Accumulate up to @needed@ hex digits from the head of @input@.
-- Returns @(value', digits', rest)@ where @value'@ folds the new digits
-- into @value@ and @rest@ is the unconsumed remainder.
accumulateHex :: Int -> Int -> Text -> (Int, Int, Text)
accumulateHex value digits input =
  let needed = 4 - digits
      hex = T.take needed (T.takeWhile isHexDigit input)
      rest = T.drop (T.length hex) input
      value' = T.foldl' (\v c -> v * 16 + digitToInt c) value hex
      digits' = digits + T.length hex
   in (value', digits', rest)

data Next = EndOfInput | NextEvent NextEmitted

-- | One decoded event with the decoder and text positioned after it.
data NextEmitted = NextEmitted
  { nextEvent :: JSONEvent,
    nextDecoder :: DecoderState,
    nextText :: StreamIO Text ()
  }

-- | Stream over 'IO'. Failures abort via 'HQ.Early' exceptions.
type StreamIO s = Stream (Of s) IO