Compile.hs

{-|
The compile phase.

Reference counting, culling, lifetime computation and the derivation of
the multi-queue synchronization schedule.
-}
module Fragr.Compile
  ( compile
  , compileWith
  , canExecute
  , validateAliasGroups
  ) where

import Control.Exception (throwIO)
import Control.Monad (guard, when)
import Control.Monad.IO.Class (MonadIO (..))
import Data.Either (partitionEithers)
import Data.Foldable (for_, toList)
import Data.IORef (readIORef, writeIORef)
import Data.IntMap.Strict (IntMap)
import Data.IntMap.Strict qualified as IntMap
import Data.IntSet qualified as IntSet
import Data.List (foldl')
import Data.Map.Strict qualified as Map
import Data.Maybe (listToMaybe, maybeToList)
import Data.Sequence (Seq)
import Data.Sequence qualified as Seq
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Text (Text)
import Data.Word (Word64)

import Fragr.Error (FragrError (..))
import Fragr.Graph (FrameGraph (..), PassNode (..), ResourceEntry (..), ResourceNode (..), producedNodes, requireCompiled, touchedNodes)
import Fragr.Resource (Access (..), accessId)
import Fragr.Sync (Compiled (..), PassSync (..), SyncEvent (..), Transfer (..), Wait (..))
import Fragr.Types (EventId (..), FamilyId, Handle (..), QueueId (..))

{- |
Reference-count, cull dead passes and resources, and compute resource
lifetimes. Must run after all passes are registered and before 'Fragr.Execute.execute'.

Cross-queue transfers come out one per consuming queue; 'compileWith'
derives them per consuming /family/ instead.
-}
{-# INLINE compile #-}
compile :: (MonadIO m) => FrameGraph ctx alloc -> m ()
compile = compileWith []

{- |
'compile' under a queue-family partition: ownership of a cross-queue
output moves once per consuming family — the family's first-registered
consumer acquires, its sibling queues ride the schedule's waits — and a
version consumed on two distinct families is 'ReleasedToTwoFamilies' (a
single owner cannot be released twice; 'Fragr.Graph.markShared' resources
are exempt). Queues left out of the partition (e.g. a host queue) keep
their per-queue transfers.
-}
{-# INLINEABLE compileWith #-}
compileWith :: (MonadIO m) => [(QueueId, FamilyId)] -> FrameGraph ctx alloc -> m ()
compileWith families g = liftIO do
  passes <- readIORef g.passesRef
  nodes <- readIORef g.nodesRef
  entries <- readIORef g.entriesRef
  let
    familyMap = IntMap.fromList do
      (QueueId q, f) <- families
      pure (q, f)
    (compiled, violations) = compilePure passes nodes entries familyMap
  for_ (listToMaybe violations) \(res, consumers) ->
    throwIO $ ReleasedToTwoFamilies res consumers
  writeIORef g.compiledRef $ Just compiled

-- | Whether a pass survives culling, given the pass refcounts.
{-# INLINE canExecute #-}
canExecute :: IntMap Int -> PassNode ctx alloc -> Bool
canExecute refs p =
  p.sideEffect || IntMap.findWithDefault 0 p.passId refs > 0

{- |
Validate a plan-time aliasing decision: each member entry of every group
must be written before it is read, since over aliased memory a first read
observes whatever the previous occupant left behind. A renaming write
counts as a read — it consumes the prior version. Group membership is by
entry id (see 'Fragr.Snapshot.EntryInfo'); throws 'AliasedReadBeforeWrite'
naming the resource and the reading pass. Fatal before 'compile'.

Note this is the generic half of the rule: a backend still owes the
full-overwrite discipline (clear, discard, or copy-over) that flags alone
can express.
-}
validateAliasGroups :: (MonadIO m) => FrameGraph ctx alloc -> [[Int]] -> m ()
validateAliasGroups g groups = liftIO do
  compiled <- requireCompiled g "validateAliasGroups"
  passes <- readIORef g.passesRef
  nodes <- readIORef g.nodesRef
  entries <- readIORef g.entriesRef
  let
    members = IntSet.fromList (concat groups)
    entryOfNode h = (nodes `Seq.index` h).resourceId
    -- The first executing pass that reads or writes each member entry, and
    -- whether it reads it. Creates alone are materialization, not access.
    firstTouch = IntMap.fromListWith (\_later earlier -> earlier) do
      p <- filter (canExecute compiled.passRefs) (toList passes)
      (e, readsIt) <-
        IntMap.toList . IntMap.fromListWith (||) $
          [(entryOfNode (accessId a), True) | a <- p.reads]
            <> [(entryOfNode (accessId a), False) | a <- p.writes]
      guard (e `IntSet.member` members)
      pure (e, (p, readsIt))
  for_ (concat groups) \e -> do
    when (e < 0 || e >= Seq.length entries) $
      throwIO (ResourceIdOutOfRange e)
    for_ (IntMap.lookup e firstTouch) \(p, readsIt) ->
      when readsIt $
        throwIO (AliasedReadBeforeWrite (entries `Seq.index` e).name p.name)

{- | Compile-internal ordering edge, with the accesses behind each
endpoint: the producing pass's on 'srcAccess', the consuming pass's on
'dstAccess'.
-}
data SyncEdge = SyncEdge
  { from :: Int
  , to :: Int
  , srcAccess :: Set Access
  , dstAccess :: Set Access
  }

-- Alongside the schedule: the 'ReleasedToTwoFamilies' violations, each a
-- resource name with (first consumer, family) per consuming family.
compilePure
  :: Seq (PassNode ctx alloc)
  -> Seq ResourceNode
  -> Seq (ResourceEntry ctx alloc)
  -> IntMap FamilyId
  -> (Compiled, [(Text, [(Text, FamilyId)])])
compilePure passes nodes entries families =
  ( Compiled
      { passRefs = finalPassRefs
      , nodeRefs = finalNodeRefs
      , passSync = syncMap
      , passLevel = levelMap
      , passAnti = antiAfterMap
      , entryRetire = entryRetireMap
      , retireAfter = retireAfterMap
      }
  , ownershipViolations
  )
  where
    passList = toList passes

    -- 4.1 initial reference counting
    -- A pass's refcount is the number of distinct resource /nodes/ it
    -- produces — creates and writes — that are still referenced. Count
    -- distinct handles: writing one created handle several times (with
    -- different flags) still names a single node, which culling can
    -- decrement only once.
    passRefs0 = IntMap.fromList do
      p <- passList
      pure (p.passId, IntSet.size (IntSet.fromList (producedNodes p)))
    nodeRefs0 =
      IntMap.unionWith
        (+)
        ( IntMap.fromList do
            n <- toList nodes
            pure (n.nodeId, 0)
        )
        ( IntMap.fromListWith (+) do
            p <- passList
            h <- map accessId p.reads
            pure (h, 1 :: Int)
        )
    nodeProducer = IntMap.fromList do
      p <- passList
      h <- producedNodes p
      pure (h, p.passId)
    passSideEffect = IntMap.fromList do
      p <- passList
      pure (p.passId, p.sideEffect)
    passReads = IntMap.fromList do
      p <- passList
      pure (p.passId, map accessId p.reads)

    -- 4.2 culling by flood fill from unreferenced resource nodes
    (finalPassRefs, finalNodeRefs) =
      cull passRefs0 nodeRefs0 do
        (n, 0) <- IntMap.toList nodeRefs0
        pure n

    cull pr nr = \case
      [] -> (pr, nr)
      n : rest ->
        case IntMap.lookup n nodeProducer of
          Nothing -> cull pr nr rest
          Just pid
            | IntMap.findWithDefault False pid passSideEffect -> cull pr nr rest
            | otherwise ->
                let
                  refs = pr IntMap.! pid - 1
                  pr' = IntMap.insert pid refs pr
                in
                  if refs > 0 then
                    cull pr' nr rest
                  else
                    let
                      dec (m, acc) h =
                        let v = m IntMap.! h - 1
                        in (IntMap.insert h v m, if v == 0 then h : acc else acc)
                      (nr', newlyDead) = foldl' dec (nr, []) (IntMap.findWithDefault [] pid passReads)
                    in
                      cull pr' nr' (newlyDead <> rest)

    entryOfNode h = (nodes `Seq.index` h).resourceId
    entryIds f = Set.fromList do
      e <- toList entries
      guard (f e)
      pure e.entryId
    importedEntries = entryIds (.imported)
    sharedEntries = entryIds (.shared)

    -- Sync schedule (see the module-level notes). Everything below runs
    -- over the /executing/ passes only.
    executing = filter (canExecute finalPassRefs) passList
    queueOf = IntMap.fromList do
      p <- executing
      pure (p.passId, p.queue)
    qOf i = queueOf IntMap.! i

    -- Per-queue registration-order position, 1-based: a pass at position i
    -- signals value i on its queue's timeline.
    posOf :: IntMap Word64
    posOf = snd (foldl' step (IntMap.empty, IntMap.empty) executing)
      where
        step (counters, acc) p =
          let
            QueueId q = p.queue
            n = IntMap.findWithDefault 0 q counters + 1 :: Word64
          in
            (IntMap.insert q n counters, IntMap.insert p.passId n acc)

    -- The declared write accesses behind each written node; a node has
    -- exactly one producer, so the node id alone keys them. Accumulated as
    -- sets — deduplicated and ascending by construction, here and in
    -- everything downstream that unions them.
    writesByNode :: IntMap (Set Access)
    writesByNode = IntMap.fromListWith (<>) do
      p <- executing
      a@Access{handle = Handle i} <- p.writes
      pure (i, Set.singleton a)

    -- Producer -> consumer edges, one per declared consumer read, carrying
    -- the accesses on both endpoints.
    dataEdges :: [SyncEdge]
    dataEdges = do
      c <- executing
      a@Access{handle = Handle i} <- c.reads
      pid <- maybeToList (IntMap.lookup i nodeProducer)
      guard (IntMap.member pid queueOf)
      guard (pid /= c.passId)
      pure
        SyncEdge
          { from = pid
          , to = c.passId
          , srcAccess = IntMap.findWithDefault Set.empty i writesByNode
          , dstAccess = Set.singleton a
          }

    -- Reader -> renamer anti-edges (write-after-read): the pass that renames
    -- a node must run after every executing reader of that node. Readers
    -- always register before the renamer (a rename makes the handle stale),
    -- so anti-edges point backward like data edges.
    nodeByVer = Map.fromList do
      n <- toList nodes
      pure ((n.resourceId, n.version), n.nodeId)
    -- Old node id -> (renaming pass, the new version's node id).
    renamerOf :: IntMap (Int, Int)
    renamerOf = IntMap.fromList do
      n <- toList nodes
      next <- maybeToList (Map.lookup (n.resourceId, n.version + 1) nodeByVer)
      w <- maybeToList (IntMap.lookup next nodeProducer)
      guard (IntMap.member w queueOf)
      pure (n.nodeId, (w, next))
    antiEdges :: [SyncEdge]
    antiEdges = do
      r <- executing
      a@Access{handle = Handle i} <- r.reads
      (w, next) <- maybeToList (IntMap.lookup i renamerOf)
      guard (w /= r.passId)
      pure
        SyncEdge
          { from = r.passId
          , to = w
          , srcAccess = Set.singleton a
          , dstAccess = IntMap.findWithDefault Set.empty next writesByNode
          }

    -- Ordering edges: data edges and anti-edges alike need waits (cross
    -- queue) or events (same queue, gapped); only data edges transfer
    -- ownership.
    syncEdges :: [SyncEdge]
    syncEdges = dataEdges <> antiEdges

    -- Longest-path depth over the ordering edges (see 'Fragr.Snapshot.level').
    -- Edges point backward, so one sweep in registration order sees every
    -- predecessor's level already settled.
    predecessors :: IntMap [Int]
    predecessors = IntMap.fromListWith (<>) do
      e <- syncEdges
      pure (e.to, [e.from])

    levelMap :: IntMap Int
    levelMap = foldl' step IntMap.empty executing
      where
        step acc p = IntMap.insert p.passId (maximum (0 : map depth preds)) acc
          where
            preds = IntMap.findWithDefault [] p.passId predecessors
            depth q = 1 + IntMap.findWithDefault 0 q acc

    -- The anti-edges as a per-reader adjacency, for writers that want the
    -- ordering the dataflow does not draw.
    antiAfterMap :: IntMap [Int]
    antiAfterMap =
      IntSet.toList <$> IntMap.fromListWith (<>) do
        e <- antiEdges
        pure (e.from, IntSet.singleton e.to)

    -- Cross-queue waits, deduplicated by a per-(consumer,producer)-queue
    -- watermark: because passes on a queue run in order and each waits
    -- before running, a value already awaited by an earlier same-queue pass
    -- is implied and dropped (transitive reduction). Only the wait is
    -- implied, not its scope: a backend derives stage masks from covers, so
    -- a dropped wait's accesses migrate to the kept wait behind the
    -- watermark.
    foreignNeed :: IntMap (IntMap (Word64, Set Access))
    foreignNeed =
      IntMap.fromListWith (IntMap.unionWith merge) (edgeNeed <> siblingNeed)
      where
        edgeNeed = do
          e <- syncEdges
          let
            QueueId pq = qOf e.from
            QueueId cq = qOf e.to
          guard (pq /= cq)
          pure (e.to, IntMap.singleton pq (posOf IntMap.! e.from, e.dstAccess))
        merge (v1, a1) (v2, a2) = (max v1 v2, a1 <> a2)

    waitsMap :: IntMap [Wait]
    waitsMap = snd (foldl' go (Map.empty, IntMap.empty) executing)
      where
        go acc@(wm, out) c =
          case IntMap.lookup c.passId foreignNeed of
            Nothing -> acc
            Just needs ->
              let
                QueueId cq = c.queue
                ((wm', out'), ws) = foldl' (perPq cq c.passId) ((wm, out), []) (IntMap.toAscList needs)
              in
                (wm', IntMap.insert c.passId (reverse ws) out')
        perPq cq pid ((wm, ws'), ws) (pq, (v, accs)) =
          case Map.lookup (cq, pq) wm of
            Just (cur, keeper)
              | v <= cur ->
                  ((wm, IntMap.adjust (map (widen pq accs)) keeper ws'), ws)
            _ ->
              ( (Map.insert (cq, pq) (v, pid) wm, ws')
              , Wait{queue = QueueId pq, value = v, covers = Set.toAscList accs} : ws
              )
        widen pq accs w
          | w.queue == QueueId pq =
              Wait{queue = w.queue, value = w.value, covers = Set.toAscList (Set.fromList w.covers <> accs)}
          | otherwise = w

    -- Same-queue dependencies with at least one pass in between get a
    -- split-barrier event pair; adjacent ones rely on a plain barrier. Each
    -- endpoint covers its own pass's accesses (barrier scopes).
    gappedByPair :: Map.Map (Int, Int) (Set Access, Set Access)
    gappedByPair = Map.fromListWith (<>) do
      e <- syncEdges
      guard (qOf e.from == qOf e.to)
      guard (posOf IntMap.! e.to - posOf IntMap.! e.from > 1)
      pure ((e.from, e.to), (e.srcAccess, e.dstAccess))
    eventOf = zip (Map.toAscList gappedByPair) (map EventId [0 ..])
    eventsBy pick = IntMap.fromListWith (<>) do
      ((pair, ends), eid) <- eventOf
      let (pass, accs) = pick pair ends
      pure (pass, Set.singleton SyncEvent{event = eid, covers = Set.toAscList accs})
    signalEventsMap = eventsBy \(pid, _) (src, _) -> (pid, src)
    waitEventsMap = eventsBy \(_, c) (_, dst) -> (c, dst)

    -- Cross-queue ownership transfer: the producer releases the output to
    -- the consumer's queue, the consumer acquires it from the producer's,
    -- both under the consuming access's flags. Data edges only — anti-edges
    -- carry no payload.
    --
    -- A rename consumes the old version through its implicit flagless read;
    -- its transfer carries the renaming write's access instead — the state
    -- the consumer actually uses the resource in, and the only flags the
    -- release / acquire hooks could ever see.
    renameAccess c = \case
      Access{handle = Handle i, flags = Nothing}
        | Just (w, next) <- IntMap.lookup i renamerOf
        , w == c
        , a' : _ <- Set.toAscList (IntMap.findWithDefault Set.empty next writesByNode) ->
            a'
      a -> a
    famOf :: QueueId -> Maybe FamilyId
    famOf (QueueId q) = IntMap.lookup q families

    -- Ownership moves per family where the partition says so: consumers
    -- whose queue has a family are grouped per (node, family) and only the
    -- group's first-registered consumer — the primary acquirer — carries
    -- the transfer, under its own accesses. Consumers outside the
    -- partition (e.g. a host queue) keep one transfer per consuming queue.
    (legacyRows, familyRows) = partitionEithers do
      e <- dataEdges
      guard (qOf e.from /= qOf e.to)
      a <- Set.toAscList e.dstAccess
      pure case famOf (qOf e.to) of
        Nothing -> Left ((accessId a, qOf e.to), IntMap.singleton e.to (Set.singleton a))
        Just fam -> Right ((accessId a, fam), IntMap.singleton e.to (Set.singleton a))
    -- Per (node, family) / (node, queue) outside the partition: each
    -- consuming pass's accesses.
    familyMembers :: Map.Map (Int, FamilyId) (IntMap (Set Access))
    familyMembers = Map.fromListWith (IntMap.unionWith (<>)) familyRows
    legacyMembers :: Map.Map (Int, QueueId) (IntMap (Set Access))
    legacyMembers = Map.fromListWith (IntMap.unionWith (<>)) legacyRows
    -- The producer and the group's first-registered consumer — the primary
    -- acquirer — with its accesses; a version released once must be
    -- acquired once.
    primaryOf node members = case IntMap.findMin members of
      (c, accs) -> (nodeProducer IntMap.! node, c, accs)
    familyGroups :: Map.Map (Int, FamilyId) (Int, Int, Set Access)
    familyGroups = Map.mapWithKey (\(node, _fam) -> primaryOf node) familyMembers
    legacyGroups :: Map.Map (Int, QueueId) (Int, Int, Set Access)
    legacyGroups = Map.mapWithKey (\(node, _q) -> primaryOf node) legacyMembers
    -- The sibling-queue acquire guard: the acquire barrier is recorded on
    -- the primary acquirer's queue, so a same-family consumer on another
    -- queue must also wait for the acquiring pass, not just the producer.
    -- Only where a real acquire happens — a hand-off within the producer's
    -- own family or of a shared resource transitions producer-side, which
    -- the data-edge waits already order.
    siblingNeed :: [(Int, IntMap (Word64, Set Access))]
    siblingNeed = do
      ((node, fam), members) <- Map.toAscList familyMembers
      guard (entryOfNode node `Set.notMember` sharedEntries)
      guard (famOf (qOf (nodeProducer IntMap.! node)) /= Just fam)
      let
        (primary, _) = IntMap.findMin members
        QueueId pq = qOf primary
      (c, accs) <- IntMap.toAscList (IntMap.delete primary members)
      guard (qOf c /= qOf primary)
      pure (c, IntMap.singleton pq (posOf IntMap.! primary, accs))
    transferPairs = do
      (pid, primary, accs) <- Map.elems legacyGroups <> Map.elems familyGroups
      a <- Set.toAscList accs
      pure (pid, primary, renameAccess primary a)
    releasesMap =
      IntMap.fromListWith (<>) do
        (pid, c, Access h f) <- transferPairs
        pure (pid, Set.singleton Transfer{handle = h, peer = qOf c, flags = f})
    acquiresMap =
      IntMap.fromListWith (<>) do
        (pid, c, Access h f) <- transferPairs
        pure (c, Set.singleton Transfer{handle = h, peer = qOf pid, flags = f})

    -- One version consumed on two distinct families: a single owner cannot
    -- be released twice. Shared entries carry no ownership and are exempt.
    nodeFamilies :: IntMap (Map.Map FamilyId Int)
    nodeFamilies = IntMap.fromListWith (<>) do
      ((node, fam), (_pid, primary, _accs)) <- Map.toAscList familyGroups
      pure (node, Map.singleton fam primary)
    ownershipViolations = do
      (node, fams) <- IntMap.toAscList nodeFamilies
      guard (Map.size fams > 1)
      guard (entryOfNode node `Set.notMember` sharedEntries)
      pure
        ( (entries `Seq.index` entryOfNode node).name
        , do
            (fam, c) <- Map.toAscList fams
            pure ((passes `Seq.index` c).name, fam)
        )

    syncMap = IntMap.fromList do
      p <- executing
      let
        setAt :: IntMap (Set a) -> [a]
        setAt m = Set.toAscList (IntMap.findWithDefault Set.empty p.passId m)
      pure
        ( p.passId
        , PassSync
            { passId = p.passId
            , name = p.name
            , queue = p.queue
            , waits = IntMap.findWithDefault [] p.passId waitsMap
            , signal = posOf IntMap.! p.passId
            , waitEvents = setAt waitEventsMap
            , signalEvents = setAt signalEventsMap
            , acquires = setAt acquiresMap
            , releases = setAt releasesMap
            }
        )

    -- Per-entry deferred-reclamation requirements (transients only — the
    -- graph never reclaims imports): the last executing user position on
    -- each queue that touches the entry.
    entryUses = do
      p <- executing
      h <- touchedNodes p
      let e = entryOfNode h
      guard (e `Set.notMember` importedEntries)
      pure (e, p)
    entryRetireMap =
      IntMap.map
        ( \m -> do
            (q, v) <- Map.toAscList m
            pure (QueueId q, v)
        )
        ( IntMap.fromListWith (Map.unionWith max) do
            (e, p) <- entryUses
            let QueueId q = p.queue
            pure (e, Map.singleton q (posOf IntMap.! p.passId))
        )
    entryLastMap = IntMap.fromListWith max do
      (e, p) <- entryUses
      pure (e, p.passId)
    retireAfterMap = IntMap.fromListWith (flip (<>)) do
      (e, pid) <- IntMap.toList entryLastMap
      pure (pid, [e])