Run.hs

module Engine.Run
  ( runStack
  , run
  , step
  ) where

import RIO

import GHC.Clock (getMonotonicTimeNSec)
import Graphics.UI.GLFW qualified as GLFW
import RIO.App (App(..), appEnv)
import RIO.State (get, put)
import UnliftIO.Resource (ReleaseKey, release)
import Vulkan.Core10 qualified as Vk

import Engine.DataRecycler (DataRecycler(..))
import Engine.DataRecycler qualified as DataRecycler
import Engine.Frame qualified as Frame
import Engine.Render (renderFrame)
import Engine.Setup.Window qualified as Window
import Engine.StageSwitch (getNextStage)
import Engine.Types (Frame, NextStage(..), RecycledResources, StageStack, StackStage(..), Stage(..), StageRIO)
import Engine.Types qualified as Engine
import Engine.Types.Options (optionsRecyclerWait, optionsMaxFPS)
import Engine.Types.RefCounted (releaseRefCounted)
import Engine.Vulkan.Swapchain (SwapchainResources(srRelease), threwSwapchainError)
import Engine.Vulkan.Types (RenderPass, getDevice)
import Engine.Worker qualified as Worker

runStack :: StageStack -> StageRIO (Maybe SwapchainResources) ()
runStack = \case
  [] -> do
    logInfo "Stage stack finished"
    get >>= \case
      Nothing ->
        pure ()
      Just oldSR ->
        releaseRefCounted (srRelease oldSR)
    asks getDevice >>= Vk.deviceWaitIdle
    threadDelay 0.5e6 -- XXX: this is arbitrary and may not always satisfy the validator

  StackStage stage : rest -> do
    logInfo $ "Setting up stage " <> display (sTitle stage)
    (stageRelease, runner) <- prepareStage (sInitialRS stage)
    oldSR <- get
    runner (run oldSR stageRelease stage) >>= proceed stageRelease rest

  StackStageContinue stKey state stage : rest -> do
    logInfo $ "Resuming stage " <> display (sTitle stage)
    (stageRelease, runner) <- prepareStage (pure (stKey, state))
    oldSR <- get
    runner (run oldSR stageRelease stage) >>= proceed stageRelease rest

type StageResult = (SwapchainResources, StageAction)

proceed :: ReleaseKey -> StageStack -> StageResult -> StageRIO (Maybe SwapchainResources) ()
proceed stageRelease rest (oldSR, stageAction) = do
  put (Just oldSR)
  case stageAction of
    StageDone -> do
      releaseStage "Stage done"
      runStack rest
    StageReplace nextStage -> do
      releaseStage "Stage replaced"
      runStack (nextStage : rest)
    StagePush frozenStage nextStage -> do
      case frozenStage of
        StackStage{} ->
          logDebug "Restarting stage pushed"
        StackStageContinue{} ->
          logDebug "Frozen stage pushed"
      runStack (nextStage : frozenStage : rest)
  where
    releaseStage label = do
      logDebug label
      void $! async do
        {-
          XXX: wait for the current render job to finish.

          Since the new frames wouldn't be submitted until the control flow
          gets back to the render loop, this should be enough of a signal that
          the stage resources aren't used anymore.

          The call is costly, but its effects aren't seen.
        -}
        asks getDevice >>= Vk.deviceWaitIdle
        release stageRelease
        logDebug $ label <> ", released"

prepareStage
  :: StageRIO env (ReleaseKey, st)
  -> StageRIO env (ReleaseKey, StageRIO st a -> StageRIO env a)
prepareStage initialRS = do
  (key, rs) <- initialRS
  freshStateVar <- newSomeRef rs
  App{..} <- ask
  let
    stageApp = App
      { appState = freshStateVar
      , ..
      }
  pure (key, runRIO stageApp)

run
  :: RenderPass rp
  => Maybe SwapchainResources
  -> ReleaseKey
  -> Stage rp p rr st
  -> StageRIO st StageResult
run oldSR stKey stage@Stage{..} = do
  logInfo $ "Starting stage: " <> display sTitle

  recycler <- DataRecycler.new

  startFrame <- Frame.initial oldSR (drDump recycler) stage

  Engine.GlobalHandles{ghWindow, ghOptions} <- asks appEnv
  quit <- liftIO $ GLFW.windowShouldClose ghWindow
  if quit then do
    logDebug $ "Forcing stage unwind for " <> display sTitle
    pure
      ( Frame.fSwapchainResources startFrame
      , StageDone
      )
  else do
    logInfo $ "Entering stage loop: " <> display sTitle
    startTime <- liftIO getMonotonicTimeNSec
    (finalFrame, stageAction) <- bracket sBeforeLoop sAfterLoop \_stagePrivates ->
      stageLoop (step stKey stage recycler (optionsMaxFPS ghOptions)) startFrame
    endTime <- liftIO getMonotonicTimeNSec
    let
      frames  = Frame.fIndex finalFrame
      seconds = fromIntegral (endTime - startTime) / 1e9 :: Double

    logInfo $ "Stage finished: " <> display sTitle
    logInfo $ "Running time: " <> display seconds
    logInfo $ "Average FPS: " <> display (fromIntegral frames / seconds)
    releaseRefCounted $ fst (Frame.fStageResources finalFrame)
    pure
      ( Frame.fSwapchainResources finalFrame
      , stageAction
      )

step
  :: RenderPass rp
  => ReleaseKey
  -> Stage rp p rr st
  -> DataRecycler (RecycledResources rr)
  -> Maybe Int
  -> Frame rp p rr
  -> StageRIO st (LoopAction (Frame rp p rr))
step stKey stage@Stage{..} DataRecycler{..} maxFPS_ frame = do
  startTime <- liftIO getMonotonicTimeNSec
  liftIO GLFW.pollEvents -- XXX: runs callback processing on a frame budget

  -- XXX: hard unwind all the stages, rendering nothing
  quit <- liftIO $ GLFW.windowShouldClose (Frame.fWindow frame)
  if quit then
    pure LoopQuit
  else
    getNextStage >>= \case
      Nothing -> do
        Engine.GlobalHandles{ghOptions, ghScreenVar, ghWindow, ghWindowSize} <- asks appEnv
        needsNewSwapchain <- threwSwapchainError do
          let recyclerWait = optionsRecyclerWait ghOptions
          rs <- get
          Frame.run drDump recyclerWait (renderFrame (sUpdateBuffers rs) sRecordCommands) frame
        nextFrame <- Frame.advance drWait frame needsNewSwapchain
        when needsNewSwapchain do
          newSize <- liftIO (Window.getSize ghWindow)
          Worker.pushInput ghScreenVar $ const newSize.framebufferExtent
          Worker.pushInput ghWindowSize $ const newSize
        for_ maxFPS_ \maxFPS -> do
          -- Wait until the end of allocated time.
          endTime <- liftIO getMonotonicTimeNSec
          let
            elapsedUS = fromIntegral (endTime - startTime) `div` 1e3
            fpsCapUS = 1e6 `div` maxFPS
          threadDelay $ fpsCapUS - min fpsCapUS elapsedUS
        pure $ LoopNextFrame nextFrame
      Just Finish ->
        -- XXX: finish the stage and proceed with the remaining stack
        pure LoopQuit
      Just (Replace nextStage) ->
        pure $ LoopReplaceStage frame nextStage
      Just (PushRestart nextStage) ->
        pure $ LoopPushStage frame (StackStage stage) nextStage
      Just (PushFreeze nextStage) -> do
        frozen <- freeze stKey stage
        pure $ LoopPushStage frame frozen nextStage

freeze :: RenderPass rp => ReleaseKey -> Stage rp p rr st -> StageRIO st StackStage
freeze stKey stage = do
  st <- get
  pure $ StackStageContinue stKey st stage

data StageAction
  = StageReplace StackStage
  | StagePush StackStage StackStage
  | StageDone

data LoopAction f
  = LoopNextFrame f
  | LoopReplaceStage f StackStage
  | LoopPushStage f StackStage StackStage
  | LoopQuit

{-# INLINE stageLoop #-}
stageLoop :: (f -> StageRIO st (LoopAction f)) -> f -> StageRIO st (f, StageAction)
stageLoop action current =
  action current >>= \case
    LoopNextFrame nextFrame ->
      stageLoop action nextFrame
    LoopQuit ->
      pure (current, StageDone)
    LoopReplaceStage lastFrame nextStage ->
      pure (lastFrame, StageReplace nextStage)
    LoopPushStage lastFrame frozenStage nextStage ->
      pure (lastFrame, StagePush frozenStage nextStage)