Graph.hs

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

{-|
The mutable graph object and its bookkeeping.

Pass / resource records and the graph-level operations that need no pass
context. Everything public here is re-exported through "Fragr"; the
bookkeeping internals carry no stability guarantees when imported
directly.
-}
module Fragr.Graph
  ( FrameGraph (..)
  , newFrameGraph
  , PassNode (..)
  , ResourceNode (..)
  , ResourceEntry (..)
  , SomeEntry (..)
  , Resources (..)
  , importResource
  , importScratch
  , markShared
  , markObserved
  , addPreExec
  , addPostExec
  , isValid
  , getDescriptor
  , appendEntry
  , appendNode
  , nodeAt
  , entryAt
  , entryOf
  , producedNodes
  , touchedNodes
  , requireCompiled
  , assertValid
  , castEntry
  , describeEntry
  ) where

import Type.Reflection

import Control.Exception (throwIO)
import Control.Monad (unless)
import Control.Monad.IO.Class (MonadIO (..))
import Data.IORef
import Data.IntSet (IntSet)
import Data.Sequence (Seq, (|>))
import Data.Sequence qualified as Seq
import Data.Text (Text)

import Fragr.Error (FragrError (..))
import Fragr.Resource (Access, Resource (..), accessId)
import Fragr.Sync (Compiled)
import Fragr.Types (Handle (..), QueueId, SomeHandle (..), someHandleId)

{- |
The frame graph. @ctx@ and @alloc@ are the opaque user values accepted by
'execute'; every resource used with the graph must agree on them (via its
'Ctx' and 'Alloc' associated types).
-}
data FrameGraph ctx alloc = FrameGraph
  { passesRef :: IORef (Seq (PassNode ctx alloc))
  , nodesRef :: IORef (Seq ResourceNode)
  , entriesRef :: IORef (Seq (ResourceEntry ctx alloc))
  , compiledRef :: IORef (Maybe Compiled)
  , preExecRef :: IORef (Seq (ctx -> IO ()))
  , postExecRef :: IORef (Seq (ctx -> IO ()))
  }

{-# INLINEABLE newFrameGraph #-}
-- XXX: The @forall@ keeps @ctx@ and @alloc@ first for type applications;
-- the other type-applied entry points ('create', 'importResource', 'get',
-- 'getDesc', 'getDescriptor') fix @r@ first the same way.
newFrameGraph :: forall ctx alloc m. (MonadIO m) => m (FrameGraph ctx alloc)
newFrameGraph = liftIO do
  FrameGraph
    <$> newIORef mempty
    <*> newIORef mempty
    <*> newIORef mempty
    <*> newIORef Nothing
    <*> newIORef mempty
    <*> newIORef mempty

-- | One registered pass.
data PassNode ctx alloc = PassNode
  { passId :: Int
  , name :: Text
  , creates :: [SomeHandle]
  , reads :: [Access]
  , writes :: [Access]
  , declared :: IntSet
  {- ^ every declared handle (creates, reads, writes): the membership
  index behind execution-time access checks
  -}
  , sideEffect :: Bool
  , queue :: QueueId
  , run :: Resources ctx alloc -> IO ()
  }

-- | One version of a resource. Immutable after creation.
data ResourceNode = ResourceNode
  { nodeId :: Int
  , resourceId :: Int
  , version :: Int
  }

-- | One distinct resource: identity, current version, type-erased payload.
data ResourceEntry ctx alloc = ResourceEntry
  { entryId :: Int
  , name :: Text
  , imported :: Bool
  , observed :: Bool
  {- ^ Writes are externally visible ('importResource'): writers become
  side effects, immune to culling. 'importScratch' and created
  transients are only read through the graph, so demand decides.
  -}
  , shared :: Bool
  {- ^ Exempt from single-owner validation ('markShared'): the allocation
  tolerates concurrent access from several queue families at once, so
  consuming one version on two families is not an error for it. Ownership
  itself is never stored — each version's owner is implied by its
  producing pass, and transfers move it along the data edges.
  -}
  , versionRef :: IORef Int
  , payload :: SomeEntry ctx alloc
  }

{- | Type-erased resource object with its descriptor. The object slot is
'Nothing' for a transient that has not been materialized (yet, or
anymore).
-}
data SomeEntry ctx alloc where
  SomeEntry
    :: (Resource r, Alloc r ~ alloc, Ctx r ~ ctx)
    => TypeRep r
    -> Desc r
    -> IORef (Maybe r)
    -> SomeEntry ctx alloc

initialVersion :: Int
initialVersion = 1

{- |
The environment 'Exec' reads: the executing pass, its graph, and the
frame context. Resource access is restricted to the handles the pass
declared via create, read or write.
-}
data Resources ctx alloc = Resources
  { graph :: FrameGraph ctx alloc
  , pass :: PassNode ctx alloc
  , ctx :: ctx
  }

{- |
Import an externally-owned, already-constructed resource. The graph never
'createResource's nor 'destroyResource's it. Not tied to any pass.

Its contents are externally observable, so every pass writing it becomes a
side effect, immune to culling; import scratch-like targets with
'importScratch' instead.
-}
{-# INLINEABLE importResource #-}
importResource
  :: forall r ctx alloc m
   . (Resource r, Alloc r ~ alloc, Ctx r ~ ctx, MonadIO m)
  => FrameGraph ctx alloc
  -> Text
  -> Desc r
  -> r
  -> m (Handle r)
importResource g resName desc obj = liftIO do
  objRef <- newIORef (Just obj)
  appendEntry g resName True True desc objRef

{- |
'importResource' for a target nothing observes from outside the graph: the
object and its memory are still externally owned, but its contents matter
only to passes of this graph, so writers stay subject to demand culling
like any transient's. Mark the passes that feed a between-graphs consumer
'setSideEffect' explicitly.
-}
{-# INLINEABLE importScratch #-}
importScratch
  :: forall r ctx alloc m
   . (Resource r, Alloc r ~ alloc, Ctx r ~ ctx, MonadIO m)
  => FrameGraph ctx alloc
  -> Text
  -> Desc r
  -> r
  -> m (Handle r)
importScratch g resName desc obj = liftIO do
  objRef <- newIORef (Just obj)
  appendEntry g resName True False desc objRef

{- |
Mark the resource behind the handle as tolerating concurrent access from
several queue families at once — Vulkan @CONCURRENT@ sharing — so
'Fragr.Compile.compileWith' skips its single-owner check
('Fragr.Error.ReleasedToTwoFamilies'). Its transfers are still derived:
a backend melts them into plain state transitions, but the hooks must
fire.

Imports read this off the object ('Fragr.Resource.isShared'); marking by
hand is for created transients, whose object does not exist yet.
-}
{-# INLINEABLE markShared #-}
markShared :: (MonadIO m) => FrameGraph ctx alloc -> Handle r -> m ()
markShared g h = liftIO do
  node <- nodeAt g h
  modifyIORef' g.entriesRef (Seq.adjust' (\e -> e{shared = True}) node.resourceId)
  writeIORef g.compiledRef Nothing

{- |
Make writes to the resource behind the handle externally observable, as if
imported through 'importResource': later writers become side effects.
Not exported through "Fragr" — 'Fragr.Builder.importOwned' flips its entry
/after/ the synthetic pass, keeping that pass itself cullable.
-}
markObserved :: (MonadIO m) => FrameGraph ctx alloc -> Handle r -> m ()
markObserved g h = liftIO do
  node <- nodeAt g h
  modifyIORef' g.entriesRef (Seq.adjust' (\e -> e{observed = True}) node.resourceId)

{- |
Install a per-pass flush point: an action fired by 'execute' and
'executeQueued' after every executing pass's hooks ('preAcquire',
'preRead', 'preWrite') and before its execution callback. Hooks that
accumulate work into @ctx@ (e.g. image barriers to batch into one command)
emit it here. Every installed action fires for every executing pass, in
installation order — a library adapter and the application can hook the
same graph. Installation is append-only: on a graph kept across frames,
install during setup, not per frame.
-}
{-# INLINE addPreExec #-}
addPreExec :: (MonadIO m) => FrameGraph ctx alloc -> (ctx -> IO ()) -> m ()
addPreExec g f = liftIO $ modifyIORef' g.preExecRef (|> f)

{- |
The post-pass counterpart of 'addPreExec': an action fired after every
executing pass's execution callback and its 'preRelease' hooks, so a pass
releasing several resources to another queue can batch the release
barriers its hooks accumulated into one command.
-}
{-# INLINE addPostExec #-}
addPostExec :: (MonadIO m) => FrameGraph ctx alloc -> (ctx -> IO ()) -> m ()
addPostExec g f = liftIO $ modifyIORef' g.postExecRef (|> f)

{- |
True iff the handle names the /latest/ version of its resource. A handle
out of range is a fatal error, not 'False'.
-}
{-# INLINEABLE isValid #-}
isValid :: (MonadIO m) => FrameGraph ctx alloc -> Handle r -> m Bool
isValid g h = liftIO do
  node <- nodeAt g h
  entry <- entryAt g node.resourceId
  current <- readIORef entry.versionRef
  pure (node.version == current)

-- | The descriptor of the resource behind the handle.
{-# INLINEABLE getDescriptor #-}
getDescriptor
  :: forall r ctx alloc m
   . (Resource r, MonadIO m)
  => FrameGraph ctx alloc
  -> Handle r
  -> m (Desc r)
getDescriptor g h = liftIO do
  entry <- entryOf g h
  (desc, _) <- castEntry @r "getDescriptor" entry
  pure desc

appendEntry
  :: forall r ctx alloc
   . (Resource r, Alloc r ~ alloc, Ctx r ~ ctx)
  => FrameGraph ctx alloc
  -> Text
  -> Bool
  -> Bool
  -> Desc r
  -> IORef (Maybe r)
  -> IO (Handle r)
appendEntry g resName isImported isObserved desc objRef = do
  entries <- readIORef g.entriesRef
  versionRef <- newIORef initialVersion
  -- An import's object is already here and knows its sharing; a created
  -- transient has none yet ('markShared' is its channel).
  obj <- readIORef objRef
  let entry =
        ResourceEntry
          { entryId = Seq.length entries
          , name = resName
          , imported = isImported
          , observed = isObserved
          , shared = maybe False isShared obj
          , versionRef
          , payload = SomeEntry (typeRep @r) desc objRef
          }
  writeIORef g.entriesRef (entries |> entry)
  appendNode g entry.entryId initialVersion

-- Polymorphic in @r@: the caller pins the minted handle's resource type
-- (the entry's for 'appendEntry', the renamed handle's for a write).
appendNode :: FrameGraph ctx alloc -> Int -> Int -> IO (Handle r)
appendNode g rid ver = do
  nodes <- readIORef g.nodesRef
  let node =
        ResourceNode
          { nodeId = Seq.length nodes
          , resourceId = rid
          , version = ver
          }
  writeIORef g.nodesRef (nodes |> node)
  -- Like registering a pass ('addPass') or 'markShared', appending a node
  -- invalidates any prior compilation.
  writeIORef g.compiledRef Nothing
  pure (Handle node.nodeId)

{-# INLINEABLE nodeAt #-}
nodeAt :: FrameGraph ctx alloc -> Handle r -> IO ResourceNode
nodeAt g (Handle i) = do
  nodes <- readIORef g.nodesRef
  case Seq.lookup i nodes of
    Nothing -> throwIO $ HandleOutOfRange i
    Just node -> pure node

{-# INLINEABLE entryAt #-}
entryAt :: FrameGraph ctx alloc -> Int -> IO (ResourceEntry ctx alloc)
entryAt g i = do
  entries <- readIORef g.entriesRef
  case Seq.lookup i entries of
    Nothing -> throwIO $ ResourceIdOutOfRange i
    Just entry -> pure entry

-- | The compiled schedule, or 'NotCompiled' naming the caller.
{-# INLINEABLE requireCompiled #-}
requireCompiled :: FrameGraph ctx alloc -> Text -> IO Compiled
requireCompiled g who =
  readIORef g.compiledRef >>= \case
    Nothing -> throwIO $ NotCompiled who
    Just c -> pure c

{-# INLINEABLE entryOf #-}
entryOf :: FrameGraph ctx alloc -> Handle r -> IO (ResourceEntry ctx alloc)
entryOf g h = do
  node <- nodeAt g h
  entryAt g node.resourceId

{- | The node ids the pass produces: one per create and per write.

Culling refcounts the pass by these, and 'Fragr.Builder.finalize' finds a
handle's producer through them.
-}
producedNodes :: PassNode ctx alloc -> [Int]
producedNodes p = map someHandleId p.creates <> map accessId p.writes

-- | The node ids the pass touches at all: 'producedNodes' plus its reads.
touchedNodes :: PassNode ctx alloc -> [Int]
touchedNodes p = producedNodes p <> map accessId p.reads

{-# INLINEABLE assertValid #-}
assertValid :: Text -> FrameGraph ctx alloc -> Handle r -> IO ()
assertValid who g h = do
  ok <- isValid g h
  unless ok $
    throwIO $
      StaleHandle who (SomeHandle h)

{-# INLINEABLE castEntry #-}
castEntry
  :: forall r ctx alloc
   . (Resource r)
  => Text
  -> ResourceEntry ctx alloc
  -> IO (Desc r, IORef (Maybe r))
castEntry who entry = case entry.payload of
  SomeEntry rep desc objRef ->
    case eqTypeRep rep (typeRep @r) of
      Just HRefl -> pure (desc, objRef)
      Nothing ->
        throwIO $ TypeMismatch who (SomeTypeRep rep) (SomeTypeRep (typeRep @r))

describeEntry :: SomeEntry ctx alloc -> Text
describeEntry (SomeEntry (_ :: TypeRep r) desc _) = describeDesc @r desc