Fold.hs

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

module HQ.Runner.Fold where

import Data.Aeson (Value (..))
import Data.Aeson.Key qualified as Key
import Data.Aeson.KeyMap qualified as KeyMap
import Data.Vector qualified as Vector
import HQ.Early (Early, leave, orLeave)
import HQ.Error (HQError (..))
import HQ.JSON.Decoder (initialDecoder)
import HQ.JSON.Event
import HQ.Optic (Optic (..), PrismKind (..), prismPredicate)
import HQ.Runner.Cursor
import HQ.Runner.Error (RunnerError (..))
import HQ.Runner.Take (takeFirstValue, takeValue)
import HQ.Transformation (Transformation, runTransformation)
import HQ.Transformation.Error (TransformationError (..))
import Relude hiding (Compose, id, many, some, state)
import Streaming (Of (..))
import Streaming.Prelude qualified as S

-- | Materialize the value at the cursor into events and a 'Value',
-- returning the value, its events and the cursor after it.
-- Shared by folding and rewriting so @filter@ gates and @over@
-- replacements agree on error mapping.
data Materialized = Materialized
  { materializedValue :: Value,
    materializedEvents :: NonEmpty JSONEvent,
    materializedCursor :: Cursor
  }

materializeValue :: Early HQError -> Cursor -> IO Materialized
materializeValue early cur = do
  events :> afterValue <- S.toList (takeValue early cur)
  ne <- case nonEmpty events of
    Nothing -> leave early $ HQRunnerError EmptyValue
    Just ne -> pure ne
  v <- orLeave early $ first HQRunnerError $ eventsToValue (toList ne)
  pure (Materialized v ne afterValue)

-- | Test a @filter@ gate, mapping failures to 'HQError'.
-- Shared by 'runFold' and 'runRewrite'.
gateValue :: Optic -> Transformation -> Value -> Either HQError Bool
gateValue o t v = first HQTransformationError $ evalFilterGate o t v

-- | Test a @filter@ gate on a streamed value, returning whether it is
-- kept plus replay/after cursors. Shared by 'runFold' and 'runRewrite'
-- so both agree on gating and cursor rebuild.
data GateResult = GateResult
  { gateKeep :: Bool,
    gateFirstEvent :: JSONEvent,
    gateReplay :: Cursor,
    gateAfter :: Cursor
  }

gateTake :: Early HQError -> Optic -> Transformation -> Cursor -> IO GateResult
gateTake early o t cur = do
  Materialized v ne@(firstEv :| _) afterValue <- materializeValue early cur
  keep <- orLeave early (gateValue o t v)
  let Cursor _ dec txt = afterValue
  pure (GateResult keep firstEv (Cursor (toList ne) dec txt) afterValue)

-- | Apply a transformation, mapping failures to 'HQError'.
-- Shared by rewriting (@over@) replacements.
applyTransformation :: Transformation -> Value -> Either HQError Value
applyTransformation t v =
  first HQTransformationError $ runTransformation t v

onEventOrEnd ::
  Early HQError -> Cursor -> (Step -> EventStream Cursor) -> EventStream Cursor
onEventOrEnd early input f = do
  result <- lift (pullCursor early input)
  case result of
    Nothing -> pure input
    Just step -> f step

-- | Match the next event against a table of opens; anything else is
-- skipped as a whole value. Unifies the open-or-skip preamble in
-- Fold (@runField@/@runEach@/@runKeys@/@runValues@/@runIndex@).
-- Inlined so the table specializes per call site.
{-# INLINE onOpenOrSkip #-}
onOpenOrSkip ::
  Early HQError ->
  Cursor ->
  [(JSONEvent, Cursor -> EventStream Cursor)] ->
  EventStream Cursor
onOpenOrSkip early input arms =
  onEventOrEnd early input $ \(Step event rest) ->
    case [handle | (open, handle) <- arms, open == event] of
      (handle : _) -> handle rest
      [] -> skipValue early (pushCursor event rest)

-- | Loop pulling member steps until the container ends; @body@ handles
-- each step, returning the cursor to continue from. Shared by
-- 'traverseObject' and 'traverseArray'. Inlined so @next@/@body@
-- specialize per walk; the 'INLINE's on the traverse wrappers and
-- step pullers dissolve the whole walk into each caller, leaving a
-- concrete single-type loop.
{-# INLINE walkMembers #-}
walkMembers ::
  (Cursor -> IO (Either Cursor s)) ->
  (s -> EventStream Cursor) ->
  Cursor ->
  EventStream Cursor
walkMembers next body = go
  where
    go stream = do
      step <- lift (next stream)
      case step of
        Left rest -> pure rest
        Right pair -> body pair >>= go

-- | Walk an object body step by step; @body@ handles each member key,
-- returning the cursor to continue from. Emits nothing, returns the
-- cursor after @JSONEndObject@.
-- Inlined so the step puller dissolves into the caller's walk loop.
{-# INLINE traverseObject #-}
traverseObject ::
  Early HQError ->
  Cursor ->
  (ObjectStep -> EventStream Cursor) ->
  EventStream Cursor
traverseObject early input body = walkMembers (expectObjectStep early) body input

-- | Walk an array body step by step; @body@ handles each element event,
-- returning the cursor to continue from.
-- Inlined so the step puller dissolves into the caller's walk loop.
{-# INLINE traverseArray #-}
traverseArray ::
  Early HQError -> Cursor -> (Step -> EventStream Cursor) -> EventStream Cursor
traverseArray early input body = walkMembers (expectArrayStep early) body input

-- | Consume the rest of an object body (all remaining members).
skipRestOfObject :: Early HQError -> Cursor -> EventStream Cursor
skipRestOfObject early input = traverseObject early input
  $ \(ObjectStep _ rest) -> skipMemberValue early rest

-- | Consume the rest of an array body (all remaining elements).
skipRestOfArray :: Early HQError -> Cursor -> EventStream Cursor
skipRestOfArray early input = traverseArray early input
  $ \(Step event rest) -> skipValue early (pushCursor event rest)

-- | Interpret an optic against the JSON value at the cursor.
--
-- The continuation receives the cursor positioned immediately after
-- the value currently being interpreted.
runFold :: Early HQError -> Optic -> Continuation
runFold early optic = run optic (takeValue early)
  where
    run :: Optic -> Continuation -> Continuation
    run opticF k input = case opticF of
      Id -> k input
      Field name -> runField name k input
      Each -> runEach k input
      Keys -> runKeys k input
      Values -> runValues k input
      Filter o t -> runFilter o t k input
      Compose left right -> run left (run right k) input
      Prism kind -> runScalar (prismPredicate kind) k input
      PrismJust -> runJust k input
      Ix i -> runIndex i k input

    runField :: Text -> Continuation -> Continuation
    runField name k input =
      onOpenOrSkip early input [(JSONBeginObject, findField name k)]

    findField :: Text -> Continuation -> Continuation
    findField name k = go
      where
        go stream = do
          step <- lift (expectObjectStep early stream)
          case step of
            Left rest -> pure rest
            Right (ObjectStep key rest)
              | key == name -> k rest >>= skipRestOfObject early
              | otherwise -> skipMemberValue early rest >>= go

    ----------------------------------------------------------------
    -- Each
    ----------------------------------------------------------------

    runEach :: Continuation -> Continuation
    runEach k input = onOpenOrSkip early input eventsWithCont
      where
        eventsWithCont =
          [ ( JSONBeginArray,
              \rest -> traverseArray early rest
                $ \(Step event rest') -> k (pushCursor event rest')
            ),
            ( JSONBeginObject,
              \rest -> traverseObject early rest $ k . objectStepCursor
            )
          ]

    ----------------------------------------------------------------
    -- Keys: object keys as strings (objects only)
    ----------------------------------------------------------------

    runKeys :: Continuation -> Continuation
    runKeys k input = onOpenOrSkip early input eventsWithCont
      where
        eventsWithCont =
          [ ( JSONBeginObject,
              \rest -> traverseObject early rest $ \(ObjectStep key rest') -> do
                _ <- k $ Cursor [JSONString key] initialDecoder $ pure ()
                skipMemberValue early rest'
            )
          ]

    ----------------------------------------------------------------
    -- Values: object member values only (arrays focus on nothing)
    ----------------------------------------------------------------

    runValues :: Continuation -> Continuation
    runValues k input = onOpenOrSkip early input eventsWithCont
      where
        eventsWithCont =
          [ ( JSONBeginObject,
              \rest -> traverseObject early rest $ k . objectStepCursor
            )
          ]

    ----------------------------------------------------------------
    -- Filter: keep the value when the predicate holds of the
    -- sub-optic's focus (existential over the focused values)
    ----------------------------------------------------------------

    runFilter :: Optic -> Transformation -> Continuation -> Continuation
    runFilter o t k input = onEventOrEnd early input $ \(Step event rest0) -> do
      GateResult keep _ replay rest <-
        lift (gateTake early o t (pushCursor event rest0))
      if keep then k replay else pure rest

    ----------------------------------------------------------------
    -- Scalar prisms
    ----------------------------------------------------------------

    runScalar :: (JSONEvent -> Bool) -> Continuation -> Continuation
    runScalar predicate k input =
      onEventOrEnd early input $ \(Step event rest) ->
        if predicate event
          then k (pushCursor event rest)
          else skipValue early (pushCursor event rest) -- not a match; consume anyway

    ----------------------------------------------------------------
    -- PrismJust
    ----------------------------------------------------------------

    runJust :: Continuation -> Continuation
    runJust k input = onEventOrEnd early input $ \case
      Step JSONNull rest -> pure rest
      Step event rest -> k (pushCursor event rest)

    ----------------------------------------------------------------
    -- Array index
    ----------------------------------------------------------------

    runIndex :: Int -> Continuation -> Continuation
    runIndex index k input =
      onOpenOrSkip early input [(JSONBeginArray, findIndex index)]
      where
        findIndex n stream
          | n < 0 = skipRestOfArray early stream
          | otherwise =
              lift (expectArrayStep early stream) >>= \case
                Left rest -> pure rest
                Right (Step event rest) -> case n of
                  0 -> k (pushCursor event rest) >>= skipRestOfArray early
                  _ -> skipValue early (pushCursor event rest) >>= findIndex (n - 1)

--------------------------------------------------------------------------------
-- Pure navigation: the same optic semantics over in-memory values.
--
-- Navigation here is total: absent members, out-of-bounds indices and
-- prism mismatches focus on nothing (mirroring the streaming runners'
-- skip behavior). Only @filter@ predicates are partial, and their
-- failures propagate.
--------------------------------------------------------------------------------

-- | Focus an optic on an in-memory value, collecting every focused
-- value. Used to evaluate @filter@ gates without re-driving the
-- streaming decoder.
focusMany :: Optic -> Value -> Either TransformationError [Value]
focusMany o v = case o of
  Field name -> pure (lookupField name v)
  Each -> pure (eachValues v)
  Keys -> pure (keyValues v)
  Values -> pure (objectValues v)
  Id -> pure [v]
  Compose l r -> do
    ls <- focusMany l v
    mconcat <$> traverse (focusMany r) ls
  Prism kind -> pure $ matchPrism kind v
  PrismJust -> pure $ matchJust v
  Ix i -> pure $ matchIndex i v
  Filter o' t -> do
    keep <- evalFilterGate o' t v
    pure [v | keep]

-- | Evaluate a @filter@ gate on an in-memory value: true when the
-- transformation maps some focused sub-value to true.
evalFilterGate ::
  Optic -> Transformation -> Value -> Either TransformationError Bool
evalFilterGate o t v = focusMany o v >>= go
  where
    go :: [Value] -> Either TransformationError Bool
    go [] = pure False
    go (w : ws) = do
      b <- testValue w
      if b then pure True else go ws
    testValue :: Value -> Either TransformationError Bool
    testValue w = case runTransformation t w of
      Left err -> Left err
      Right (Bool b) -> pure b
      Right _ -> Left FilterNotBoolean

-- | Object member lookup: missing members and non-objects focus on
-- nothing.
lookupField :: Text -> Value -> [Value]
lookupField name v = case v of
  Object o -> maybeToList (KeyMap.lookup (Key.fromText name) o)
  _ -> []

-- | Object member values only; arrays and scalars focus on nothing.
-- Arrays and scalars focus on nothing; use 'Each' for arrays.
objectValues :: Value -> [Value]
objectValues v = case v of
  Object o -> KeyMap.elems o
  _ -> []

-- | Array elements plus object member values.
eachValues :: Value -> [Value]
eachValues v = case v of
  Array a -> Vector.toList a
  Object _ -> objectValues v
  _ -> []

-- | Object keys as strings; arrays focus on nothing.
keyValues :: Value -> [Value]
keyValues v = case v of
  Object o -> map (String . Key.toText) (KeyMap.keys o)
  _ -> []

-- | Match a type prism against an in-memory value directly.
matchPrism :: PrismKind -> Value -> [Value]
matchPrism kind v = case kind of
  PString -> case v of s@(String _) -> [s]; _ -> []
  PNumber -> case v of n@(Number _) -> [n]; _ -> []
  PBool -> case v of b@(Bool _) -> [b]; _ -> []
  PNull -> case v of Null -> [v]; _ -> []
  PArray -> case v of a@(Array _) -> [a]; _ -> []
  PObject -> case v of o@(Object _) -> [o]; _ -> []

matchJust :: Value -> [Value]
matchJust v = case v of
  Null -> []
  _ -> [v]

matchIndex :: Int -> Value -> [Value]
matchIndex i v = case v of
  Array a -> maybeToList (a Vector.!? i)
  _ -> []

-- | Execute an optic, emitting at most one value: the first one it
-- selects.
--
-- The fold is short-circuited: once the first value has been emitted,
-- no further input is read from the source. This is library-level
-- first-match semantics over any optic; the @preview@ CLI command
-- exposes it directly (any optic accepted, first focus wins).
runPreview :: Early HQError -> Optic -> Continuation
runPreview early optic input =
  takeFirstValue (runFold early optic input) $> input