fragr-0.1.0.0: test/Utils.hs

{-# LANGUAGE TypeFamilies #-}

-- | Test resource types and the mock backends the suites drive them through.
module Utils
  ( -- * Event-log resources
    Env (..)
  , newEnv
  , Event (..)
  , push
  , getEvents
  , ran
  , ranHere
  , Tex (..)
  , TexDesc (..)
  , tex
  , Buf (..)

    -- * Simulated device
  , Device (..)
  , newDevice
  , logQ
  , logHere
  , queueLog
  , Image (..)
  , ImgDesc (..)
  , Layout (..)
  , layoutName
  , img
  , runDevice
  , runHere

    -- * Driving a frame
  , mkBackend
  , runQueued
  , assertHandoffsDrained

    -- * Assertions
  , assertFatal
  , has
  , syncOf
  , tshow
  ) where

import Fragr

import Control.Exception (try)
import Control.Monad (unless, when)
import Control.Monad.IO.Class (liftIO)
import Data.Foldable (find, for_)
import Data.IORef (IORef, atomicModifyIORef', modifyIORef', newIORef, readIORef, writeIORef)
import Data.IntMap.Strict (IntMap)
import Data.IntMap.Strict qualified as IntMap
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Maybe (fromJust)
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Word (Word64)
import Test.Tasty.HUnit (assertBool, assertFailure)

import Fragr qualified as FG

-- * Test resource types

{- | Everything observable funnels into one event log; the allocator also
carries a creation counter (checklist: "observable via a creation
counter in the test resource type").
-}
data Env = Env
  { counter :: IORef Int
  , events :: IORef [Event]
  }

newEnv :: IO Env
newEnv = Env <$> newIORef 0 <*> newIORef []

data Event
  = ECreate Text Int
  | EDestroy Text
  | EPreRead Text Word64
  | EPreWrite Text Word64
  | EFlush Text
  | EPostFlush Text
  | ERun Text
  deriving stock (Eq, Show)

push :: Env -> Event -> IO ()
push env e = modifyIORef' env.events (e :)

getEvents :: Env -> IO [Event]
getEvents env = reverse <$> readIORef env.events

ran :: Env -> Text -> IO ()
ran env passName = push env (ERun passName)

-- | 'ran' from inside an execution callback.
ranHere :: Text -> FG.Exec Env Env ()
ranHere passName = do
  env <- FG.askCtx
  liftIO (ran env passName)

newtype Tex = Tex Int
  deriving stock (Eq, Show)

data TexDesc = TexDesc
  { tag :: Text
  , size :: Int
  }
  deriving stock (Eq, Show)

instance FG.Resource Tex where
  type Desc Tex = TexDesc
  type Alloc Tex = Env
  type Ctx Tex = Env
  type Flags Tex = Word64

  createResource desc env = do
    n <- atomicModifyIORef' env.counter \n -> (n + 1, n)
    push env (ECreate desc.tag n)
    pure (Tex n)

  destroyResource desc env _ = push env (EDestroy desc.tag)

  preRead _ desc w env _ = push env (EPreRead desc.tag w)

  preWrite _ desc w env _ = push env (EPreWrite desc.tag w)

  describeDesc desc = desc.tag

{- | A second resource type, for type-mismatch tests. No hooks: relies on
the default no-op 'preRead' / 'preWrite', empty 'describeDesc' and unit
'Flags'.
-}
data Buf = Buf
  deriving stock (Eq, Show)

instance FG.Resource Buf where
  type Desc Buf = Int
  type Alloc Buf = Env
  type Ctx Buf = Env
  createResource _ _ = pure Buf
  destroyResource _ _ _ = pure ()

tex :: Text -> TexDesc
tex t = TexDesc{tag = t, size = 4}

assertFatal :: IO a -> IO ()
assertFatal act =
  try @FragrError (act >> pure ()) >>= \case
    Left _ -> pure ()
    Right () -> assertFailure "expected a FragrError"

-- | Assert the rendered output contains the needle.
has :: Text -> Text -> IO ()
has out needle = assertBool (Text.unpack needle <> " in output") (needle `Text.isInfixOf` out)

-- * Simulated Vulkan-like backend (test-only, no GPU vocabulary in src/)

{- | A mock multi-queue device: per-queue command logs, per-queue timeline
counters, a set of signaled events, and the queue currently recording (so
the layout-tracking image hooks know where to log their barriers).
-}
data Device = Device
  { devLog :: IORef (IntMap [Text])
  -- ^ queue -> reversed command log
  , devTimeline :: IORef (IntMap Word64)
  -- ^ queue -> highest signaled (== reached, since we run synchronously)
  , devEvents :: IORef (Set Int)
  , devQueue :: IORef QueueId
  , devHandoffs :: IORef (Map (Text, Int) (Int, Word64))
  -- ^ (handle, dst queue) -> (src queue, its signal): released, not yet acquired
  , devWaited :: IORef (Map (Int, Int) Word64)
  -- ^ (queue, foreign queue) -> the highest value it has waited for so far
  }

newDevice :: IO Device
newDevice =
  Device
    <$> newIORef IntMap.empty
    <*> newIORef IntMap.empty
    <*> newIORef Set.empty
    <*> newIORef (QueueId 0)
    <*> newIORef Map.empty
    <*> newIORef Map.empty

logQ :: Device -> QueueId -> Text -> IO ()
logQ dev (QueueId q) msg = modifyIORef' dev.devLog (IntMap.insertWith (++) q [msg])

logHere :: Device -> Text -> IO ()
logHere dev msg = do
  q <- readIORef dev.devQueue
  logQ dev q msg

queueLog :: Device -> QueueId -> IO [Text]
queueLog dev (QueueId q) = reverse . IntMap.findWithDefault [] q <$> readIORef dev.devLog

{- | A simulated image whose layout state lives in an 'IORef'; the
'FG.preRead' / 'FG.preWrite' hooks diff requested-vs-current and record a
barrier command, exactly like a real backend would.
-}
data Image = Image (IORef Text) Bool

newtype ImgDesc = ImgDesc Text

-- | Image layout states: a real ADT reaching the hooks, no bit packing.
data Layout = ShaderRead | ColorAttachment | General
  deriving stock (Eq, Ord, Show)

layoutName :: Layout -> Text
layoutName = \case
  ShaderRead -> "shader-read"
  ColorAttachment -> "color-attachment"
  General -> "general"

instance FG.Resource Image where
  type Desc Image = ImgDesc
  type Alloc Image = Device
  type Ctx Image = Device
  type Flags Image = Layout

  createResource (ImgDesc n) dev = do
    logHere dev ("create " <> n)
    Image <$> newIORef "undefined" <*> pure False

  destroyResource (ImgDesc n) dev _ = logHere dev ("destroy " <> n)

  preRead _ d f dev im = barrier d f dev im
  preWrite _ d f dev im = barrier d f dev im

  isShared (Image _ sh) = sh

  preAcquire _ (ImgDesc n) f _peer dev _ = logHere dev ("acquire-hook " <> n <> " ->" <> layoutName f)
  preRelease _ (ImgDesc n) f _peer dev _ = logHere dev ("release-hook " <> n <> " ->" <> layoutName f)

barrier :: ImgDesc -> Layout -> Device -> Image -> IO ()
barrier (ImgDesc n) f dev (Image stateRef _) = do
  cur <- readIORef stateRef
  let target = layoutName f
  when (cur /= target) do
    logHere dev ("barrier " <> n <> " " <> cur <> "->" <> target)
    writeIORef stateRef target

img :: Text -> ImgDesc
img = ImgDesc

{- | A 'QueueBackend' that records the schedule onto the mock device.

Hand-offs are ground-truthed ('devHandoffs'): a release arms (handle, dst)
exactly once, the acquire consumes it exactly once from the right peer, and
only after a wait covering the releasing pass's signal — the invariants a
real QFOT pair needs to not deadlock or corrupt. 'assertHandoffsDrained'
closes the frame: every release found its acquire.
-}
mkBackend :: Device -> QueueBackend
mkBackend dev =
  QueueBackend
    { invoke = \_ body -> body
    , beforePass = \ps -> do
        writeIORef dev.devQueue ps.queue
        let QueueId our = ps.queue
        for_ ps.waits \w -> do
          let QueueId q = w.queue
          tl <- readIORef dev.devTimeline
          let reached = IntMap.findWithDefault 0 q tl
          when (reached < w.value) $
            assertFailure ("wait on q" <> show q <> " for " <> show w.value <> " but only " <> show reached <> " reached")
          modifyIORef' dev.devWaited (Map.insertWith max (our, q) w.value)
          logQ dev ps.queue ("wait q" <> tshow q <> ">=" <> tshow w.value)
        for_ ps.waitEvents \se -> do
          let EventId e = se.event
          evs <- readIORef dev.devEvents
          unless (Set.member e evs) $
            assertFailure ("waited on unsignaled event " <> show e)
          logQ dev ps.queue ("waitEvent " <> tshow e)
        waited <- readIORef dev.devWaited
        for_ ps.acquires \(Transfer h peer _flags) -> do
          let QueueId src = peer
          Map.lookup (tshow h, our) <$> readIORef dev.devHandoffs >>= \case
            Nothing ->
              assertFailure ("acquire of " <> show h <> " on q" <> show our <> " without a pending release")
            Just (src', sig) -> do
              when (src' /= src) $
                assertFailure ("acquire of " <> show h <> " names q" <> show src <> " but q" <> show src' <> " released it")
              let covered = Map.findWithDefault 0 (our, src) waited
              when (covered < sig) $
                assertFailure ("acquire of " <> show h <> " precedes a wait covering its release (waited " <> show covered <> ", released at " <> show sig <> ")")
              modifyIORef' dev.devHandoffs (Map.delete (tshow h, our))
          logQ dev ps.queue ("acquire " <> tshow h <> " from q" <> tshow src)
    , afterPass = \ps -> do
        let QueueId our = ps.queue
        for_ ps.releases \(Transfer h peer _flags) -> do
          let QueueId dst = peer
          pend <- readIORef dev.devHandoffs
          when (Map.member (tshow h, dst) pend) $
            assertFailure ("double release of " <> show h <> " to q" <> show dst)
          modifyIORef' dev.devHandoffs (Map.insert (tshow h, dst) (our, ps.signal))
          logQ dev ps.queue ("release " <> tshow h <> " to q" <> tshow dst)
        for_ ps.signalEvents \se -> do
          let EventId e = se.event
          modifyIORef' dev.devEvents (Set.insert e)
          logQ dev ps.queue ("signalEvent " <> tshow e)
        let QueueId q = ps.queue
        modifyIORef' dev.devTimeline (IntMap.insert q ps.signal)
        logQ dev ps.queue ("signal " <> tshow ps.signal)
    , completed = do
        tl <- readIORef dev.devTimeline
        pure do
          (q, v) <- IntMap.toList tl
          pure (QueueId q, v)
    }

-- | Every release found its acquire; call after 'FG.executeQueued'.
assertHandoffsDrained :: Device -> IO ()
assertHandoffsDrained dev = do
  pend <- readIORef dev.devHandoffs
  unless (Map.null pend) $
    assertFailure ("released but never acquired: " <> show (Map.keys pend))

-- | Execute on a fresh mock device and close the frame's hand-off ledger.
runQueued :: FrameGraph Device Device -> IO Device
runQueued g = do
  dev <- newDevice
  rq <- FG.newRecycleQueue
  FG.executeQueued g (mkBackend dev) (Just rq) dev dev
  assertHandoffsDrained dev
  pure dev

runDevice :: Device -> Text -> IO ()
runDevice dev n = logHere dev ("run " <> n)

-- | 'runDevice' from inside an execution callback.
runHere :: Text -> FG.Exec Device Device ()
runHere n = do
  dev <- FG.askCtx
  liftIO (runDevice dev n)

syncOf :: Snapshot -> Text -> PassSync
syncOf s n = fromJust do
  p <- find (\p -> p.name == n) s.passes
  p.sync

tshow :: (Show a) => a -> Text
tshow = Text.pack . show