Demo.hs

{-# LANGUAGE TypeFamilies #-}

{-|
The demo frame graph and its staging-buffer pool.

A tiny streaming frame: two upload passes each borrow a pooled staging
buffer, fill it, and copy their chunk into an imported GPU mesh buffer; a
draw pass consumes the synced mesh, and a dead pass gets culled. A staging
transient retires the moment its upload ends, so chunk B reuses chunk A's
slot within the frame and a second frame allocates nothing at all.
-}
module App.Demo
  ( buildDemo
  , StagingPool
  , newStagingPool
  ) where

import Control.Monad.IO.Class (liftIO)
import Data.IORef (IORef, modifyIORef', newIORef, readIORef, writeIORef)
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Text.IO qualified as Text.IO

import Fragr (FrameGraph, Handle)
import Fragr qualified as FG

-- | A mapped buffer stand-in: the name of the pool slot backing it.
newtype Buffer = Buffer Text

data BufferDesc = BufferDesc
  { label :: Text
  , size :: Int
  }

-- | Size-keyed free lists of retained slots, plus a slot-name counter.
newtype StagingPool = StagingPool (IORef PoolState)

data PoolState = PoolState
  { nextSlot :: Int
  , free :: Map Int [Buffer]
  }

newStagingPool :: IO StagingPool
newStagingPool = StagingPool <$> newIORef PoolState{nextSlot = 0, free = Map.empty}

instance FG.Resource Buffer where
  type Desc Buffer = BufferDesc
  type Alloc Buffer = StagingPool
  type Ctx Buffer = ()

  createResource desc (StagingPool ref) = do
    pool <- readIORef ref
    case Map.findWithDefault [] desc.size pool.free of
      slot@(Buffer name) : rest -> do
        writeIORef ref pool{free = Map.insert desc.size rest pool.free}
        Text.IO.putStrLn ("  reuse " <> name <> " for " <> desc.label)
        pure slot
      [] -> do
        let name = "buf" <> Text.pack (show pool.nextSlot)
        writeIORef ref pool{nextSlot = pool.nextSlot + 1}
        Text.IO.putStrLn ("  alloc " <> name <> " for " <> desc.label)
        pure (Buffer name)

  destroyResource desc (StagingPool ref) slot@(Buffer name) = do
    Text.IO.putStrLn ("  recycle " <> name <> " after " <> desc.label)
    modifyIORef' ref \pool -> pool{free = Map.insertWith (<>) desc.size [slot] pool.free}

  describeDesc desc =
    Text.pack (show desc.size) <> "B"

-- | One upload's pass data: the borrowed staging slot and the synced mesh.
data Upload = Upload
  { staging :: Handle Buffer
  , mesh :: Handle Buffer
  }

-- | Build and compile one frame's demo graph.
buildDemo :: IO (FrameGraph () StagingPool)
buildDemo = do
  g <- FG.newFrameGraph

  mesh <- FG.importResource g "mesh" BufferDesc{label = "mesh", size = 1024 * 1024} (Buffer "mesh")

  meshA <- uploadChunk g "A" mesh
  meshB <- uploadChunk g "B" meshA

  FG.addPass_
    g
    "Draw"
    do
      FG.read meshB
      FG.setSideEffect
    (liftIO (putStrLn "  run Draw (mesh in sync)"))

  -- This one is dead code and gets culled.
  FG.addPass_
    g
    "Orphan"
    do
      h <- FG.create @Buffer "staging.orphan" BufferDesc{label = "staging.orphan", size = 64 * 1024}
      FG.write_ h
    (liftIO (putStrLn "  run Orphan (should never happen)"))

  FG.compile g
  pure g

{- | Stream one chunk: borrow a staging slot, fill it, copy it into the
mesh. The staging transient's last user is this very pass, so the slot
returns to the pool as soon as the upload ends.
-}
uploadChunk :: FrameGraph () StagingPool -> Text -> Handle Buffer -> IO (Handle Buffer)
uploadChunk g chunk mesh = do
  up <-
    FG.addPass
      g
      ("Upload " <> chunk)
      do
        staging <- FG.create @Buffer ("staging." <> chunk) BufferDesc{label = "staging." <> chunk, size = 64 * 1024}
        staging' <- FG.write staging
        mesh' <- FG.write mesh
        pure Upload{staging = staging', mesh = mesh'}
      \up -> do
        Buffer slot <- FG.get @Buffer up.staging
        liftIO (Text.IO.putStrLn ("  run Upload " <> chunk <> " (fill " <> slot <> ", copy to mesh)"))
  pure up.mesh