Emit.hs

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

-- | Byte-stream take for the rewrite side: traversal that transcribes
-- values straight to output bytes.
module HQ.Runner.Rewrite.Emit where

import Control.Exception (try)
import Data.Text qualified as T
import Data.Text.Encoding (encodeUtf8Builder)
import HQ.Early (Early, leave)
import HQ.JSON.Decoder
import HQ.JSON.Decoder.Number (isNumberStart)
import HQ.JSON.Decoder.Skip
import HQ.JSON.Encoder
import HQ.JSON.Event (JSONEvent (..), matchingClose)
import HQ.Runner.Cursor
import HQ.Runner.Error (RunnerError (..))
import HQ.Runner.Take (pullOne, takeContainerGeneric)
import Relude hiding (Compose, id, many, some, state)

-- | Rewrite output: byte stream returning advanced emit state and cursor.
type RewriteContinuation = Cursor -> EmitState -> BSStream IO RewriteResult

-- | The rewrite pipeline's accumulator: emit state threaded through
-- byte emission alongside the cursor.
data RewriteResult = RewriteResult
  { rewriteEncoder :: EmitState,
    rewriteCursor :: Cursor
  }

-- | Pull one event; on end return unchanged, otherwise run the handler.
-- Byte-stream analogue of 'HQ.Runner.Cursor.onEventOrEnd'.
onEventOrEndChunks ::
  Early HQError ->
  Cursor ->
  EmitState ->
  (Step -> BSStream IO RewriteResult) ->
  BSStream IO RewriteResult
onEventOrEndChunks early input st f = do
  result <- lift (pullCursor early input)
  case result of
    Nothing -> pure (RewriteResult st input)
    Just pulled -> f pulled

-- | Match the next event against a table of opens; anything else is
-- passed through as a whole value. Byte-stream analogue of
-- 'HQ.Runner.Cursor.onOpenOrSkip'. Inlined so the table specializes
-- per call site.
{-# INLINE onOpenOrSkipChunks #-}
onOpenOrSkipChunks ::
  Early HQError ->
  EncoderConfig ->
  Cursor ->
  EmitState ->
  [(JSONEvent, Cursor -> EmitState -> BSStream IO RewriteResult)] ->
  BSStream IO RewriteResult
onOpenOrSkipChunks early config input st arms =
  onEventOrEndChunks early input st $ \(Step event rest) ->
    case [handle | (open, handle) <- arms, open == event] of
      (handle : _) -> handle rest st
      [] -> takeValueChunks early config (pushCursor event rest) st

-- | Loop pulling member steps until the container ends, emitting
-- @close@; @body@ handles each step, threading the encoder state.
-- Shared by 'traverseObjectChunks' and 'traverseArrayChunks': the
-- close emission is the only difference. Inlined so @close@/@next@/
-- @body@ specialize per walk; the 'INLINE's on the traverse wrappers
-- and step pullers below dissolve the whole walk into each caller,
-- leaving a concrete single-type loop (no per-element dictionary or
-- higher-order call).
{-# INLINE walkMembersChunks #-}
walkMembersChunks ::
  JSONEvent ->
  EncoderConfig ->
  (Cursor -> IO (Either Cursor s)) ->
  (s -> EmitState -> BSStream IO RewriteResult) ->
  Cursor ->
  EmitState ->
  BSStream IO RewriteResult
walkMembersChunks close config next body = go
  where
    go stream s = do
      step' <- lift (next stream)
      case step' of
        Left rest -> do
          s' <- emitChunk config close s
          pure (RewriteResult s' rest)
        Right pair -> do
          RewriteResult s' after <- body pair s
          go after s'

-- | Walk an object body, emitting @JSONEndObject@ at the end.
-- Caller must have emitted @JSONBeginObject@ already.
-- Inlined so the close tag and step puller dissolve into the
-- caller's 'walkMembersChunks' loop.
{-# INLINE traverseObjectChunks #-}
traverseObjectChunks ::
  Early HQError ->
  EncoderConfig ->
  Cursor ->
  EmitState ->
  (ObjectStep -> EmitState -> BSStream IO RewriteResult) ->
  BSStream IO RewriteResult
traverseObjectChunks early config input st body =
  walkMembersChunks JSONEndObject config (expectObjectStep early) body input st

-- | Walk an array body, emitting @JSONEndArray@ at the end.
{-# INLINE traverseArrayChunks #-}
traverseArrayChunks ::
  Early HQError ->
  EncoderConfig ->
  Cursor ->
  EmitState ->
  (Step -> EmitState -> BSStream IO RewriteResult) ->
  BSStream IO RewriteResult
traverseArrayChunks early config input st body =
  walkMembersChunks JSONEndArray config (expectArrayStep early) body input st

-- | Emit a key then passthrough its value; the 4x pattern in
-- @rewriteMember@/pairs/allKeys@.
emitKeyAndTake ::
  Early HQError ->
  EncoderConfig ->
  Text ->
  Cursor ->
  EmitState ->
  BSStream IO RewriteResult
emitKeyAndTake early config key rest s = do
  s' <- emitChunk config (JSONObjectKey key) s
  takeValueChunks early config rest s'

-- | Take one complete value, transcribing it straight to output bytes.
-- Strings and numbers take a verbatim raw path ('peekRawScalar'),
-- and so do whole containers in compact mode ('takeRawContainerChunks'
-- below): chunk-side only. The event-side take ('takeValue') always
-- decodes, so both paths must agree on value boundaries; the raw
-- scanners ('skipStringCollect'/'skipNumberCollect'/
-- 'skipContainerCollect') mirror the decoder ('consumeString'/
-- 'stepNumber'/structural loop) exactly, and the split/adversarial
-- tests cover the agreement. Change one side only with a test on
-- the other.
takeValueChunks ::
  Early HQError ->
  EncoderConfig ->
  Cursor ->
  EmitState ->
  BSStream IO RewriteResult
takeValueChunks early config cursor st = case peekRawScalar cursor of
  Just (RawString afterQuote, dec, txt) ->
    takeRawStringChunks early config afterQuote dec txt st
  Just (RawNumber atNumber, dec, txt) ->
    takeRawNumberChunks early config atNumber dec txt st
  Nothing -> do
    (Step event cursor') <- lift (pullOne early cursor)
    let slowPath = do
          st' <- emitChunk config event st
          case matchingClose event of
            Just closing -> takeContainerChunks early config closing cursor' st'
            Nothing -> case event of
              JSONEndArray -> lift (leave early $ HQRunnerError UnexpectedEndOfArray)
              JSONEndObject -> lift (leave early $ HQRunnerError UnexpectedEndOfObject)
              JSONObjectKey _ -> lift (leave early $ HQRunnerError UnexpectedObjectKey)
              _ -> pure (RewriteResult st' cursor')
    -- Wholesale containers bypass event decoding only when nothing is
    -- buffered (buffered replay has no source text to collect) and the
    -- output is compact (verbatim bytes are already canonical there;
    -- pretty output must be generated).
    case event of
      JSONBeginArray -> case cursor' of
        Cursor [] dec txt
          | outputStyle config == Compact ->
              takeRawContainerChunks early config JSONEndArray JSONBeginArray dec txt st
        _ -> slowPath
      JSONBeginObject -> case cursor' of
        Cursor [] dec txt
          | outputStyle config == Compact ->
              takeRawContainerChunks early config JSONEndObject JSONBeginObject dec txt st
        _ -> slowPath
      _ -> slowPath

-- | A peeked string or number value: payload after the quote / at the
-- first digit, pending colon already skipped.
data RawScalar = RawString Text | RawNumber Text

-- | Peek a string/number value; decline anything else (including
-- buffered replay cursors) so existing paths report errors.
-- Inlined: runs once per value on the chunk take path, and dissolving
-- it exposes the cursor/decoder deconstruction to the caller.
{-# INLINE peekRawScalar #-}
peekRawScalar :: Cursor -> Maybe (RawScalar, DecoderState, StreamIO Text ())
peekRawScalar cursor
  | null buf,
    Just start <- peekValue phase input,
    Just (c, rest) <- T.uncons start = case c of
      '"' -> Just (RawString rest, decoder, txt)
      _ | isNumberStart c -> Just (RawNumber start, decoder, txt)
      _ -> Nothing
  | otherwise = Nothing
  where
    Cursor buf decoder txt = cursor
    DecoderState {decoderInput = input, decoderPhase = phase} = decoder
    -- Member values hide behind their colon; anything else must already
    -- be in value position.
    peekValue DecoderPhaseValue inp = Just (T.dropWhile isWhitespace inp)
    peekValue DecoderPhaseObjectColon inp =
      case T.uncons (T.dropWhile isWhitespace inp) of
        Just (':', rest) -> Just (T.dropWhile isWhitespace rest)
        _ -> Nothing
    peekValue _ _ = Nothing

-- | Build the output bytes from a raw body. A single right-nested
-- fold: 'toLazyByteString' then streams in constant space, where
-- left-nested @acc <> piece@ accumulation retained ~150x the
-- collected size (measured on a 1.3 MB untouched array).
-- Inlined: runs once per verbatim body, fusing the fragment fold
-- with the consuming buffer append.
{-# INLINE buildFrags #-}
buildFrags :: RawBody -> Builder
buildFrags = foldl' (\acc t -> encodeUtf8Builder t <> acc) mempty . rawFrags

-- | Transcribe a peeked string value verbatim (escapes preserved,
-- quotes included in the fragments).
takeRawStringChunks ::
  Early HQError ->
  EncoderConfig ->
  Text ->
  DecoderState ->
  StreamIO Text () ->
  EmitState ->
  BSStream IO RewriteResult
takeRawStringChunks early config afterQuote dec txt st = do
  (Collected (SkipRest remainder rest) body) <- lift (skipStringCollect early afterQuote txt)
  st' <- emitRawBytes config (buildFrags body) (rawSize body) st
  pure (RewriteResult st' (Cursor [] (finishValue dec {decoderInput = remainder}) rest))

-- | Transcribe a peeked number value verbatim (no 'Scientific' roundtrip).
takeRawNumberChunks ::
  Early HQError ->
  EncoderConfig ->
  Text ->
  DecoderState ->
  StreamIO Text () ->
  EmitState ->
  BSStream IO RewriteResult
takeRawNumberChunks early config atNumber dec txt st = do
  (Collected (SkipRest remainder rest) body) <- lift (skipNumberCollect early atNumber txt)
  st' <- emitRawBytes config (buildFrags body) (rawSize body) st
  pure (RewriteResult st' (Cursor [] (finishValue dec {decoderInput = remainder}) rest))

-- | Transcribe a whole container value verbatim in bounded segments.
-- The cursor must be unbuffered with the decoder just past the opening
-- bracket (same precondition as walking it with 'takeContainerChunks';
-- the open event itself was already pulled). Separators wrap the whole
-- body exactly once ('emitRawOpen'/'emitRawTail'), so segment emission
-- composes byte-identically to one uninterrupted collect while each
-- segment stays under 'collectBudget'.
takeRawContainerChunks ::
  Early HQError ->
  EncoderConfig ->
  JSONEvent ->
  JSONEvent ->
  DecoderState ->
  StreamIO Text () ->
  EmitState ->
  BSStream IO RewriteResult
takeRawContainerChunks early config closing open dec txt st = case closing of
  JSONEndArray -> segmented ExpectValue
  JSONEndObject -> segmented ExpectKey
  -- Unreachable: callers pass literal closings. Collect nothing.
  _ -> pure (RewriteResult st (Cursor [] dec txt))
  where
    segmented expect = do
      stBracket <- emitRawOpen config open st
      let depth = decoderNestDepth dec
          base = shallower depth
      (RewriteResult stBody cur) <-
        collectLoop stBracket expect base depth (decoderStack dec) (decoderInput dec) txt
      st' <- emitRawTail config stBody
      pure (RewriteResult st' cur)
    -- Collect one budget-bounded segment, emit it verbatim, resume
    -- until the container completes. The resume state is exact loop
    -- state, so pausing is invisible downstream.
    collectLoop st0 exp0 base0 dep0 stk0 inp0 stream0 = do
      mseg <- liftIO $ try (collectExpect early mempty base0 dep0 stk0 exp0 inp0 stream0)
      case mseg of
        Right (Collected (SkipRest rest stream') body) -> do
          st' <- emitBuilder (buildFrags body) (rawSize body) st0
          pure (RewriteResult st' (Cursor [] (finishContainer dec rest) stream'))
        Left (CollectPaused body exp1 base1 dep1 stk1 inp1 stream1) -> do
          st' <- emitBuilder (buildFrags body) (rawSize body) st0
          collectLoop st' exp1 base1 dep1 stk1 inp1 stream1

-- | Take a container body, fusing decode and format per event with no
-- intermediate event stream. Thin wrapper over 'takeContainerGeneric'
-- with 'emitChunk' as the output action.
takeContainerChunks ::
  Early HQError ->
  EncoderConfig ->
  JSONEvent ->
  Cursor ->
  EmitState ->
  BSStream IO RewriteResult
takeContainerChunks early config closing cursor st = do
  (cur, st') <- takeContainerGeneric early closing cursor st (emitChunk config)
  pure (RewriteResult st' cur)
{-# INLINE takeContainerChunks #-}