Builder.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE TypeFamilies #-}

-- | The pass-declaration monad: registering passes and declaring accesses.
module Fragr.Builder
  ( addPass
  , addPass_
  , Build
  , Builder (..)
  , create
  , read
  , readWith
  , write
  , write_
  , writeWith
  , writeWith_
  , setSideEffect
  , setQueue
  , finalize
  , importOwned
  ) where

import Prelude hiding (read)

import Control.Exception (throwIO)
import Control.Monad (unless, void, when)
import Control.Monad.IO.Class (MonadIO (..))
import Control.Monad.Trans.Reader (ReaderT (..))
import Data.Foldable (find)
import Data.IORef
import Data.IntSet qualified as IntSet
import Data.Sequence ((|>))
import Data.Sequence qualified as Seq
import Data.Text (Text)

import Fragr.Error (FragrError (..))
import Fragr.Exec (Exec)
import Fragr.Graph (FrameGraph (..), PassNode (..), ResourceEntry (..), ResourceNode (..), appendEntry, appendNode, assertValid, entryAt, entryOf, importScratch, markObserved, nodeAt, producedNodes)
import Fragr.Resource (Access (..), Resource (..), accessId)
import Fragr.Types (Handle, QueueId, SomeHandle (..), defaultQueue, handleId, someHandleId)

{- |
Register a pass.

The setup block runs immediately, exactly once, before 'addPass' returns;
it declares the pass's accesses in the 'Build' monad and returns the pass
data (typically a record of handles). The execution callback is stored and
runs later, during 'Fragr.Execute.execute', only if the pass survives
culling; it receives the pass data and runs in 'Exec', which carries this
pass's resource accessor ('Fragr.Exec.get', 'Fragr.Exec.getDesc') and the
context given to 'Fragr.Execute.execute' ('Fragr.Exec.askCtx').

Registration order is execution order; no reordering is performed.
Returns the pass data produced by the setup block.
-}
{-# INLINEABLE addPass #-}
addPass
  :: (MonadIO m)
  => FrameGraph ctx alloc
  -> Text
  -> Build ctx alloc d
  -> (d -> Exec ctx alloc ())
  -> m d
addPass g passName setup exec = liftIO do
  draftRef <-
    newIORef
      PassDraft
        { draftCreates = []
        , draftReads = []
        , draftWrites = []
        , draftSideEffect = False
        , draftQueue = defaultQueue
        }
  dat <-
    runReaderT
      setup
      Builder
        { graph = g
        , draftRef
        }
  draft <- readIORef draftRef
  passes <- readIORef g.passesRef
  let node =
        PassNode
          { passId = Seq.length passes
          , name = passName
          , creates = reverse draft.draftCreates
          , reads = reverse draft.draftReads
          , writes = reverse draft.draftWrites
          , declared =
              IntSet.fromList $
                map someHandleId draft.draftCreates
                  <> map accessId (draft.draftReads <> draft.draftWrites)
          , sideEffect = draft.draftSideEffect
          , queue = draft.draftQueue
          , run = runReaderT (exec dat)
          }
  writeIORef g.passesRef (passes |> node)
  -- Registering a pass invalidates any prior compilation:
  -- 'Fragr.Execute.execute' would otherwise run against a stale schedule
  -- (silently skipping the new pass, or failing to find its 'passSync').
  writeIORef g.compiledRef Nothing
  pure dat

{- |
'addPass' for sink passes (present, readback, metering): the setup returns
no pass data, so the execution callback drops the then-vestigial data
argument. Together with 'write_' / 'writeWith_' a sink registers with zero
discarded binds.
-}
{-# INLINE addPass_ #-}
addPass_
  :: (MonadIO m)
  => FrameGraph ctx alloc
  -> Text
  -> Build ctx alloc ()
  -> Exec ctx alloc ()
  -> m ()
addPass_ g passName setup exec = addPass g passName setup (const exec)

{- |
The pass-declaration monad: the setup argument of 'addPass'. Declarations
('create', 'read', 'write', 'setQueue', 'setSideEffect') implicitly target
the pass being set up.

Setup is a declaration, not an effect: values a pass needs from the
outside world (device caps, per-frame state, scratch refs) are produced
before 'addPass' and captured, in line with the build-once-per-frame
model. The full reader surface is inherited for whoever insists.
-}
type Build ctx alloc = ReaderT (Builder ctx alloc) IO

-- | The environment 'Build' reads: the graph plus the open pass draft.
data Builder ctx alloc = Builder
  { graph :: FrameGraph ctx alloc
  , draftRef :: IORef PassDraft
  }

data PassDraft = PassDraft
  { draftCreates :: [SomeHandle]
  , draftReads :: [Access]
  , draftWrites :: [Access]
  , draftSideEffect :: Bool
  , draftQueue :: QueueId
  }

{- |
Declare a new transient resource, created and destroyed by the graph. The
resource object is /not/ materialized here; 'createResource' runs during
'Fragr.Execute.execute', just before the first pass that needs it.
-}
{-# INLINEABLE create #-}
create
  :: forall r ctx alloc
   . (Resource r, Alloc r ~ alloc, Ctx r ~ ctx)
  => Text
  -> Desc r
  -> Build ctx alloc (Handle r)
create resName desc = ReaderT \b -> do
  objRef <- newIORef (Nothing @r)
  h <- appendEntry b.graph resName False False desc objRef
  modifyIORef' b.draftRef \d -> d{draftCreates = SomeHandle h : d.draftCreates}
  pure h

-- | Like 'readWith', without flags: no hook fires for the access.
{-# INLINEABLE read #-}
read :: (Resource r) => Handle r -> Build ctx alloc ()
read h = ReaderT \b -> declareRead b h Nothing

{- |
Declare that the pass reads the resource version named by the handle; the
flags reach the resource's 'preRead' hook. Reads never rename, so no new
handle is minted. Exact duplicates (same handle /and/ same flags) are
recorded once.

Fatal errors: stale handle; handle created or written by this same pass
(within one pass a resource is either input or output).
-}
{-# INLINEABLE readWith #-}
readWith :: (Resource r) => Handle r -> Flags r -> Build ctx alloc ()
readWith h flags = ReaderT \b -> declareRead b h (Just flags)

declareRead :: (Resource r) => Builder ctx alloc -> Handle r -> Maybe (Flags r) -> IO ()
declareRead b h flags = do
  assertValid "read" b.graph h
  d <- readIORef b.draftRef
  when (SomeHandle h `elem` d.draftCreates || handleId h `elem` map accessId d.draftWrites) $
    throwIO $
      ReadsOwnOutput (SomeHandle h)
  unless (Access{handle = h, flags} `elem` d.draftReads) $
    writeIORef b.draftRef d{draftReads = Access{handle = h, flags} : d.draftReads}

-- | Like 'writeWith', without flags: no hook fires for the access.
{-# INLINEABLE write #-}
write :: (Resource r) => Handle r -> Build ctx alloc (Handle r)
write h = ReaderT \b -> declareWrite b h Nothing

{- | 'write' for a terminal write: the minted handle is dead by design and
dropped. Intentional discards stay visible and greppable; accidental ones
stay type errors.
-}
{-# INLINE write_ #-}
write_ :: (Resource r) => Handle r -> Build ctx alloc ()
write_ h = void (write h)

{- |
Declare that the pass writes the resource version named by the handle; the
flags reach the resource's 'preWrite' hook.

If this pass created the handle, the same handle is returned. Otherwise
the resource is /renamed/: the pass implicitly also reads the old version
(without flags), a new version (and node) is minted, and the /new/ handle
is returned — the old one becomes stale. Always keep the returned handle:
@h' <- FG.write h@.

Writing an observed import ('importResource') automatically marks the pass
as having a side effect (it must never be culled); writes to an
'importScratch' resource stay cullable.
-}
{-# INLINEABLE writeWith #-}
writeWith :: (Resource r) => Handle r -> Flags r -> Build ctx alloc (Handle r)
writeWith h flags = ReaderT \b -> declareWrite b h (Just flags)

-- | 'writeWith' for a terminal write, discarding the handle like 'write_'.
{-# INLINE writeWith_ #-}
writeWith_ :: (Resource r) => Handle r -> Flags r -> Build ctx alloc ()
writeWith_ h flags = void (writeWith h flags)

declareWrite :: (Resource r) => Builder ctx alloc -> Handle r -> Maybe (Flags r) -> IO (Handle r)
declareWrite b h flags = do
  assertValid "write" b.graph h
  node <- nodeAt b.graph h
  entry <- entryAt b.graph node.resourceId
  when entry.observed $ markSideEffect b
  d <- readIORef b.draftRef
  if SomeHandle h `elem` d.draftCreates then do
    unless (Access{handle = h, flags} `elem` d.draftWrites) $
      writeIORef b.draftRef d{draftWrites = Access{handle = h, flags} : d.draftWrites}
    pure h
  else do
    unless (Access{handle = h, flags = Nothing} `elem` d.draftReads) $
      modifyIORef' b.draftRef \d' -> d'{draftReads = Access{handle = h, flags = Nothing} : d'.draftReads}
    v <- readIORef entry.versionRef
    let version' = v + 1
    writeIORef entry.versionRef version'
    h' <- appendNode b.graph entry.entryId version'
    modifyIORef' b.draftRef \d' -> d'{draftWrites = Access{handle = h', flags} : d'.draftWrites}
    pure h'

{- | Mark the pass as having an observable output besides graph resources
(presenting to screen, CPU readback, ...): it is immune to culling.
-}
{-# INLINE setSideEffect #-}
setSideEffect :: Build ctx alloc ()
setSideEffect = ReaderT markSideEffect

markSideEffect :: Builder ctx alloc -> IO ()
markSideEffect b = modifyIORef' b.draftRef \d -> d{draftSideEffect = True}

{- | Assign the pass to a submission queue (default 'defaultQueue'). The
schedule computed by 'Fragr.Compile.compile' derives cross-queue timeline waits,
same-queue split-barrier events and ownership transfers from these
assignments; see 'Fragr.Sync.PassSync' and 'Fragr.Execute.executeQueued'.
-}
{-# INLINE setQueue #-}
setQueue :: QueueId -> Build ctx alloc ()
setQueue q = ReaderT \b -> modifyIORef' b.draftRef \d -> d{draftQueue = q}

{- |
'Fragr.Graph.importResource' plus the queue whose family currently owns the
contents: registers a synthetic pass on it, standing in for the work that
left them there, so a first touch on another family this frame has a
producer edge to derive the release / acquire pair from
('Fragr.Compile.compileWith') — the release records on the owning queue even
when nothing else runs there. Returns the handle the synthetic rename
minted.

A first touch within the owner's family melts like any same-family
hand-off, and on an entirely untouched import the synthetic pass culls
with the rest of the chain.
-}
{-# INLINEABLE importOwned #-}
importOwned
  :: forall r ctx alloc m
   . (Resource r, Alloc r ~ alloc, Ctx r ~ ctx, MonadIO m)
  => FrameGraph ctx alloc
  -> Text
  -> Desc r
  -> r
  -> QueueId
  -> m (Handle r)
importOwned g resName desc obj owner = liftIO do
  h0 <- importScratch g resName desc obj
  -- Scratch first, observed after: the synthetic write must not become a
  -- side effect, or an unused import would defeat culling.
  h <- addPass g ("import " <> resName) (setQueue owner *> write h0) (\_data -> pure ())
  markObserved g h
  pure h

{- |
Declare the terminal state of a resource: registers a synthetic
side-effecting pass, named after the resource, that writes the handle with
the given flags. The producing chain thus survives culling, and the
'preWrite' hook runs as the resource's last access, leaving it in the
required state (e.g. presentable). The handle minted by the underlying
rename is discarded — the resource is final.

The synthetic pass runs on the queue of the pass that produced the handle
— created or wrote it; 'defaultQueue' for an unwritten import — so
finalizing manufactures no cross-queue edge; a resource must be finalized
elsewhere through an explicit pass.
-}
{-# INLINEABLE finalize #-}
finalize :: (MonadIO m, Resource r) => FrameGraph ctx alloc -> Handle r -> Flags r -> m ()
finalize g h flags = liftIO do
  entry <- entryOf g h
  passes <- readIORef g.passesRef
  let
    produces p = handleId h `elem` producedNodes p
    producerQueue = maybe defaultQueue (.queue) (find produces passes)
  addPass_
    g
    ("finalize " <> entry.name)
    do
      setQueue producerQueue
      writeWith_ h flags
      setSideEffect
    (pure ())