{-# LANGUAGE OverloadedLists #-}

module Render.ShadowMap.RenderPass
  ( ShadowMap(..)
  , allocate
  , usePass
  ) where

import RIO

import Control.Monad.Trans.Resource qualified as Resource
import Data.Bits ((.|.))
import Data.Vector qualified as Vector
import Vulkan.Core10 qualified as Vk
import Vulkan.Core11.Promoted_From_VK_KHR_multiview qualified as Khr
import Vulkan.NamedType ((:::))
import Vulkan.Utils.Debug qualified as Debug
import Vulkan.Zero (zero)
import Vulkan.CStruct.Extends (pattern (:&), pattern (::&))

import Engine.Types.RefCounted (RefCounted, resourceTRefCount)
import Engine.Types.RefCounted qualified as RefCounted
import Engine.Vulkan.Types (HasVulkan(..), HasSwapchain(..), HasRenderPass(..), RenderPass(..), MonadVulkan)
import Render.Pass (usePass)
import Resource.Image (AllocatedImage)
import Resource.Image qualified as Image
import Resource.Region qualified as Region
import Resource.Vulkan.Named qualified as Named

-- * Depth-only pass for shadowmapping pipelines

data ShadowMap = ShadowMap
  { ShadowMap -> RenderPass
smRenderPass  :: Vk.RenderPass
  , ShadowMap -> AllocatedImage
smDepthImage  :: AllocatedImage
  , ShadowMap -> Framebuffer
smFrameBuffer :: Vk.Framebuffer
  , ShadowMap -> Rect2D
smRenderArea  :: Vk.Rect2D
  , ShadowMap -> Extent2D
smExtent      :: Vk.Extent2D
  , ShadowMap -> Word32
smLayerCount  :: Word32
  , ShadowMap -> Vector ClearValue
smClear       :: Vector Vk.ClearValue
  , ShadowMap -> RefCounted
smRelease     :: RefCounted
  }

instance HasRenderPass ShadowMap where
  getRenderPass :: ShadowMap -> RenderPass
getRenderPass   = ShadowMap -> RenderPass
smRenderPass
  getFramebuffers :: ShadowMap -> Vector Framebuffer
getFramebuffers = Int -> Framebuffer -> Vector Framebuffer
forall a. Int -> a -> Vector a
Vector.replicate Int
10 (Framebuffer -> Vector Framebuffer)
-> (ShadowMap -> Framebuffer) -> ShadowMap -> Vector Framebuffer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShadowMap -> Framebuffer
smFrameBuffer
  getClearValues :: ShadowMap -> Vector ClearValue
getClearValues  = ShadowMap -> Vector ClearValue
smClear
  getRenderArea :: ShadowMap -> Rect2D
getRenderArea   = ShadowMap -> Rect2D
smRenderArea

instance RenderPass ShadowMap where
  refcountRenderpass :: forall env. MonadResource (RIO env) => ShadowMap -> RIO env ()
refcountRenderpass  = RefCounted -> RIO env ()
forall (f :: * -> *). MonadResource f => RefCounted -> f ()
resourceTRefCount (RefCounted -> RIO env ())
-> (ShadowMap -> RefCounted) -> ShadowMap -> RIO env ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShadowMap -> RefCounted
smRelease

allocate
  :: ( Resource.MonadResource m
     , MonadVulkan env m
     , HasLogFunc env
     , HasSwapchain context
     )
  => context
  -> Word32
  -> "light count" ::: Word32
  -> m ShadowMap
allocate :: forall (m :: * -> *) env context.
(MonadResource m, MonadVulkan env m, HasLogFunc env,
 HasSwapchain context) =>
context -> Word32 -> Word32 -> m ShadowMap
allocate context
context Word32
mapSize Word32
layerCount = do
  (ReleaseKey
_rpKey, RenderPass
renderPass) <- context -> Word32 -> Word32 -> m (ReleaseKey, RenderPass)
forall env (m :: * -> *) swapchain.
(MonadVulkan env m, MonadResource m, HasSwapchain swapchain) =>
swapchain -> Word32 -> Word32 -> m (ReleaseKey, RenderPass)
allocateRenderPass context
context (Word32
2 Word32 -> Word32 -> Word32
forall a b. (Num a, Integral b) => a -> b -> a
^ Word32
layerCount Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
- Word32
1) Word32
0
  (RefCounted
release, Framebuffers
resources) <- m (ReleaseKey, Framebuffers) -> m (RefCounted, Framebuffers)
forall (m :: * -> *) a.
MonadResource m =>
m (ReleaseKey, a) -> m (RefCounted, a)
RefCounted.wrapped (m (ReleaseKey, Framebuffers) -> m (RefCounted, Framebuffers))
-> m (ReleaseKey, Framebuffers) -> m (RefCounted, Framebuffers)
forall a b. (a -> b) -> a -> b
$ ResourceT m Framebuffers -> m (ReleaseKey, Framebuffers)
forall (m :: * -> *) a.
MonadResource m =>
ResourceT m a -> m (ReleaseKey, a)
Region.run do
    Utf8Builder -> Utf8Builder -> ResourceT m ()
forall (m :: * -> *) env.
(MonadUnliftIO m, MonadReader env m, HasLogFunc env,
 HasCallStack) =>
Utf8Builder -> Utf8Builder -> ResourceT m ()
Region.logDebug
      Utf8Builder
"Allocating ShadowMap resources"
      Utf8Builder
"Releasing ShadowMap resources"
    context
-> Extent2D -> Word32 -> RenderPass -> ResourceT m Framebuffers
forall (m :: * -> *) env swapchain.
(MonadResource m, MonadVulkan env m, HasSwapchain swapchain) =>
swapchain -> Extent2D -> Word32 -> RenderPass -> m Framebuffers
allocateFramebuffer context
context Extent2D
extent Word32
layerCount RenderPass
renderPass
  let (AllocatedImage
depthImage, Framebuffer
framebuffer) = Framebuffers
resources

  ShadowMap -> m ShadowMap
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ShadowMap
    { $sel:smRenderPass:ShadowMap :: RenderPass
smRenderPass  = RenderPass
renderPass
    , $sel:smExtent:ShadowMap :: Extent2D
smExtent      = Extent2D
extent
    , $sel:smLayerCount:ShadowMap :: Word32
smLayerCount  = Word32
layerCount
    , $sel:smRenderArea:ShadowMap :: Rect2D
smRenderArea  = Rect2D
fullSurface
    , $sel:smClear:ShadowMap :: Vector ClearValue
smClear       = Vector ClearValue
clear
    , $sel:smDepthImage:ShadowMap :: AllocatedImage
smDepthImage  = AllocatedImage
depthImage
    , $sel:smFrameBuffer:ShadowMap :: Framebuffer
smFrameBuffer = Framebuffer
framebuffer
    , $sel:smRelease:ShadowMap :: RefCounted
smRelease     = RefCounted
release
    }
  where
    extent :: Extent2D
extent = Vk.Extent2D{$sel:width:Extent2D :: Word32
width=Word32
mapSize, $sel:height:Extent2D :: Word32
height=Word32
mapSize}

    fullSurface :: Rect2D
fullSurface = Vk.Rect2D
      { $sel:offset:Rect2D :: Offset2D
Vk.offset = Offset2D
forall a. Zero a => a
zero
      , $sel:extent:Rect2D :: Extent2D
Vk.extent = Extent2D
extent
      }

    clear :: Vector ClearValue
clear = [ClearValue] -> Vector ClearValue
forall a. [a] -> Vector a
Vector.fromList
      [ ClearDepthStencilValue -> ClearValue
Vk.DepthStencil (Float -> Word32 -> ClearDepthStencilValue
Vk.ClearDepthStencilValue Float
1.0 Word32
0)
      ]

-- ** Render pass

allocateRenderPass
  :: ( MonadVulkan env m
     , Resource.MonadResource m
     , HasSwapchain swapchain
     )
  => swapchain
  -> Word32
  -> Word32
  -> m (Resource.ReleaseKey, Vk.RenderPass)
allocateRenderPass :: forall env (m :: * -> *) swapchain.
(MonadVulkan env m, MonadResource m, HasSwapchain swapchain) =>
swapchain -> Word32 -> Word32 -> m (ReleaseKey, RenderPass)
allocateRenderPass swapchain
swapchain Word32
viewMask Word32
correlationMask = do
  Device
device <- (env -> Device) -> m Device
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks env -> Device
forall a. HasVulkan a => a -> Device
getDevice

  res :: (ReleaseKey, RenderPass)
res@(ReleaseKey
_key, RenderPass
object) <- Device
-> RenderPassCreateInfo '[RenderPassMultiviewCreateInfo]
-> Maybe AllocationCallbacks
-> (IO RenderPass
    -> (RenderPass -> IO ()) -> m (ReleaseKey, RenderPass))
-> m (ReleaseKey, RenderPass)
forall (a :: [*]) (io :: * -> *) r.
(Extendss RenderPassCreateInfo a, PokeChain a, MonadIO io) =>
Device
-> RenderPassCreateInfo a
-> Maybe AllocationCallbacks
-> (io RenderPass -> (RenderPass -> io ()) -> r)
-> r
Vk.withRenderPass Device
device RenderPassCreateInfo '[RenderPassMultiviewCreateInfo]
createInfo Maybe AllocationCallbacks
forall a. Maybe a
Nothing IO RenderPass
-> (RenderPass -> IO ()) -> m (ReleaseKey, RenderPass)
forall (m :: * -> *) a.
MonadResource m =>
IO a -> (a -> IO ()) -> m (ReleaseKey, a)
Resource.allocate
  Device -> RenderPass -> ByteString -> m ()
forall a (m :: * -> *).
(HasObjectType a, MonadIO m) =>
Device -> a -> ByteString -> m ()
Debug.nameObject Device
device RenderPass
object ByteString
"ShadowMap.RP"
  pure (ReleaseKey, RenderPass)
res
  where
    depthFormat :: Format
depthFormat = swapchain -> Format
forall a. HasSwapchain a => a -> Format
getDepthFormat swapchain
swapchain

    createInfo :: RenderPassCreateInfo '[RenderPassMultiviewCreateInfo]
createInfo =
      RenderPassCreateInfo '[]
forall a. Zero a => a
zero
        { $sel:attachments:RenderPassCreateInfo :: Vector AttachmentDescription
Vk.attachments  = [AttachmentDescription] -> Vector AttachmentDescription
forall a. [a] -> Vector a
Vector.fromList [Item [AttachmentDescription]
AttachmentDescription
depth]
        , $sel:subpasses:RenderPassCreateInfo :: Vector SubpassDescription
Vk.subpasses    = [SubpassDescription] -> Vector SubpassDescription
forall a. [a] -> Vector a
Vector.fromList [Item [SubpassDescription]
SubpassDescription
subpass]
        , $sel:dependencies:RenderPassCreateInfo :: Vector SubpassDependency
Vk.dependencies = [SubpassDependency] -> Vector SubpassDependency
forall a. [a] -> Vector a
Vector.fromList [Item [SubpassDependency]
SubpassDependency
pre, Item [SubpassDependency]
SubpassDependency
post]
        }
      RenderPassCreateInfo '[]
-> Chain '[RenderPassMultiviewCreateInfo]
-> RenderPassCreateInfo '[RenderPassMultiviewCreateInfo]
forall (a :: [*] -> *) (es :: [*]) (es' :: [*]).
Extensible a =>
a es' -> Chain es -> a es
::& Khr.RenderPassMultiviewCreateInfo
        { $sel:viewMasks:RenderPassMultiviewCreateInfo :: Vector Word32
Khr.viewMasks        = [Word32
Item (Vector Word32)
viewMask]
        , $sel:viewOffsets:RenderPassMultiviewCreateInfo :: Vector Int32
Khr.viewOffsets      = []
        , $sel:correlationMasks:RenderPassMultiviewCreateInfo :: Vector Word32
Khr.correlationMasks = [Word32
Item (Vector Word32)
correlationMask]
        }
      RenderPassMultiviewCreateInfo
-> Chain '[] -> Chain '[RenderPassMultiviewCreateInfo]
forall e (es :: [*]). e -> Chain es -> Chain (e : es)
:& ()

    depth :: AttachmentDescription
depth = AttachmentDescription
forall a. Zero a => a
zero
      { $sel:format:AttachmentDescription :: Format
Vk.format         = Format
depthFormat
      , $sel:samples:AttachmentDescription :: SampleCountFlagBits
Vk.samples        = SampleCountFlagBits
Vk.SAMPLE_COUNT_1_BIT
      , $sel:loadOp:AttachmentDescription :: AttachmentLoadOp
Vk.loadOp         = AttachmentLoadOp
Vk.ATTACHMENT_LOAD_OP_CLEAR
      , $sel:storeOp:AttachmentDescription :: AttachmentStoreOp
Vk.storeOp        = AttachmentStoreOp
Vk.ATTACHMENT_STORE_OP_STORE
      , $sel:stencilLoadOp:AttachmentDescription :: AttachmentLoadOp
Vk.stencilLoadOp  = AttachmentLoadOp
Vk.ATTACHMENT_LOAD_OP_DONT_CARE
      , $sel:stencilStoreOp:AttachmentDescription :: AttachmentStoreOp
Vk.stencilStoreOp = AttachmentStoreOp
Vk.ATTACHMENT_STORE_OP_DONT_CARE
      , $sel:initialLayout:AttachmentDescription :: ImageLayout
Vk.initialLayout  = ImageLayout
Vk.IMAGE_LAYOUT_UNDEFINED
      , $sel:finalLayout:AttachmentDescription :: ImageLayout
Vk.finalLayout    = ImageLayout
Vk.IMAGE_LAYOUT_DEPTH_STENCIL_READ_ONLY_OPTIMAL
      }

    subpass :: SubpassDescription
subpass = SubpassDescription
forall a. Zero a => a
zero
      { $sel:pipelineBindPoint:SubpassDescription :: PipelineBindPoint
Vk.pipelineBindPoint = PipelineBindPoint
Vk.PIPELINE_BIND_POINT_GRAPHICS
      , $sel:depthStencilAttachment:SubpassDescription :: Maybe AttachmentReference
Vk.depthStencilAttachment = AttachmentReference -> Maybe AttachmentReference
forall a. a -> Maybe a
Just AttachmentReference
forall a. Zero a => a
zero
          { $sel:attachment:AttachmentReference :: Word32
Vk.attachment = Word32
0
          , $sel:layout:AttachmentReference :: ImageLayout
Vk.layout     = ImageLayout
Vk.IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL
          }
      }

    pre :: SubpassDependency
pre = SubpassDependency
forall a. Zero a => a
zero
      { $sel:dependencyFlags:SubpassDependency :: DependencyFlags
Vk.dependencyFlags = DependencyFlags
Vk.DEPENDENCY_BY_REGION_BIT

      , $sel:srcSubpass:SubpassDependency :: Word32
Vk.srcSubpass      = Word32
Vk.SUBPASS_EXTERNAL
      , $sel:srcStageMask:SubpassDependency :: PipelineStageFlags
Vk.srcStageMask    = PipelineStageFlags
Vk.PIPELINE_STAGE_FRAGMENT_SHADER_BIT
      , $sel:srcAccessMask:SubpassDependency :: AccessFlags
Vk.srcAccessMask   = AccessFlags
Vk.ACCESS_SHADER_READ_BIT

      , $sel:dstSubpass:SubpassDependency :: Word32
Vk.dstSubpass      = Word32
0
      , $sel:dstStageMask:SubpassDependency :: PipelineStageFlags
Vk.dstStageMask    = PipelineStageFlags
Vk.PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT
      , $sel:dstAccessMask:SubpassDependency :: AccessFlags
Vk.dstAccessMask   = AccessFlags
Vk.ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT
      }

    post :: SubpassDependency
post = SubpassDependency
forall a. Zero a => a
zero
      { $sel:dependencyFlags:SubpassDependency :: DependencyFlags
Vk.dependencyFlags = DependencyFlags
Vk.DEPENDENCY_BY_REGION_BIT

      , $sel:srcSubpass:SubpassDependency :: Word32
Vk.srcSubpass      = Word32
0
      , $sel:srcStageMask:SubpassDependency :: PipelineStageFlags
Vk.srcStageMask    = PipelineStageFlags
Vk.PIPELINE_STAGE_LATE_FRAGMENT_TESTS_BIT
      , $sel:srcAccessMask:SubpassDependency :: AccessFlags
Vk.srcAccessMask   = AccessFlags
Vk.ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT

      , $sel:dstSubpass:SubpassDependency :: Word32
Vk.dstSubpass      = Word32
Vk.SUBPASS_EXTERNAL
      , $sel:dstStageMask:SubpassDependency :: PipelineStageFlags
Vk.dstStageMask    = PipelineStageFlags
Vk.PIPELINE_STAGE_FRAGMENT_SHADER_BIT
      , $sel:dstAccessMask:SubpassDependency :: AccessFlags
Vk.dstAccessMask   = AccessFlags
Vk.ACCESS_SHADER_READ_BIT
      }

-- ** Framebuffer

type Framebuffers =
  ( Image.AllocatedImage
  , Vk.Framebuffer
  )

allocateFramebuffer
  :: ( Resource.MonadResource m
     , MonadVulkan env m
     , HasSwapchain swapchain
     )
  => swapchain
  -> Vk.Extent2D
  -> Word32
  -> Vk.RenderPass
  -> m Framebuffers
allocateFramebuffer :: forall (m :: * -> *) env swapchain.
(MonadResource m, MonadVulkan env m, HasSwapchain swapchain) =>
swapchain -> Extent2D -> Word32 -> RenderPass -> m Framebuffers
allocateFramebuffer swapchain
swapchain Extent2D
extent Word32
layerCount RenderPass
renderPass = do
  AllocatedImage
depth <- Maybe Text
-> ImageAspectFlags
-> ("image dimensions" ::: Extent3D)
-> Word32
-> Word32
-> SampleCountFlagBits
-> Format
-> ImageUsageFlags
-> m AllocatedImage
forall env (io :: * -> *).
(MonadVulkan env io, MonadResource io) =>
Maybe Text
-> ImageAspectFlags
-> ("image dimensions" ::: Extent3D)
-> Word32
-> Word32
-> SampleCountFlagBits
-> Format
-> ImageUsageFlags
-> io AllocatedImage
Image.allocate
    (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"ShadowMap.depth")
    ImageAspectFlags
Vk.IMAGE_ASPECT_DEPTH_BIT
    (Extent2D -> Word32 -> "image dimensions" ::: Extent3D
Image.inflateExtent Extent2D
extent Word32
1)
    Word32
1
    Word32
layerCount
    SampleCountFlagBits
Vk.SAMPLE_COUNT_1_BIT
    (swapchain -> Format
forall a. HasSwapchain a => a -> Format
getDepthFormat swapchain
swapchain)
    (ImageUsageFlags
Vk.IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT ImageUsageFlags -> ImageUsageFlags -> ImageUsageFlags
forall a. Bits a => a -> a -> a
.|. ImageUsageFlags
Vk.IMAGE_USAGE_SAMPLED_BIT)

  let
    attachments :: Vector ImageView
attachments = [ImageView] -> Vector ImageView
forall a. [a] -> Vector a
Vector.fromList
      [ AllocatedImage -> ImageView
Image.aiImageView AllocatedImage
depth
      ]

    {- XXX:
      If the render pass uses multiview, then layers must be one and each attachment
      requires a number of layers that is greater than the maximum bit index set in
      the view mask in the subpasses in which it is used.
    -}
    fbNumLayers :: Word32
fbNumLayers = Word32
1

    Vk.Extent2D{Word32
$sel:width:Extent2D :: Extent2D -> Word32
width :: Word32
width, Word32
$sel:height:Extent2D :: Extent2D -> Word32
height :: Word32
height} = Extent2D
extent
    fbCI :: FramebufferCreateInfo '[]
fbCI = FramebufferCreateInfo '[]
forall a. Zero a => a
zero
      { $sel:renderPass:FramebufferCreateInfo :: RenderPass
Vk.renderPass  = RenderPass
renderPass
      , $sel:width:FramebufferCreateInfo :: Word32
Vk.width       = Word32
width
      , $sel:height:FramebufferCreateInfo :: Word32
Vk.height      = Word32
height
      , $sel:attachments:FramebufferCreateInfo :: Vector ImageView
Vk.attachments = Vector ImageView
attachments
      , $sel:layers:FramebufferCreateInfo :: Word32
Vk.layers      = Word32
fbNumLayers
      }

  Device
device <- (env -> Device) -> m Device
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks env -> Device
forall a. HasVulkan a => a -> Device
getDevice
  (ReleaseKey
_framebufferKey, Framebuffer
framebuffer) <- Device
-> FramebufferCreateInfo '[]
-> Maybe AllocationCallbacks
-> (IO Framebuffer
    -> (Framebuffer -> IO ()) -> m (ReleaseKey, Framebuffer))
-> m (ReleaseKey, Framebuffer)
forall (a :: [*]) (io :: * -> *) r.
(Extendss FramebufferCreateInfo a, PokeChain a, MonadIO io) =>
Device
-> FramebufferCreateInfo a
-> Maybe AllocationCallbacks
-> (io Framebuffer -> (Framebuffer -> io ()) -> r)
-> r
Vk.withFramebuffer Device
device FramebufferCreateInfo '[]
fbCI Maybe AllocationCallbacks
forall a. Maybe a
Nothing IO Framebuffer
-> (Framebuffer -> IO ()) -> m (ReleaseKey, Framebuffer)
forall (m :: * -> *) a.
MonadResource m =>
IO a -> (a -> IO ()) -> m (ReleaseKey, a)
Resource.allocate
  Framebuffer -> Text -> m ()
forall env (m :: * -> *) a.
(MonadVulkan env m, HasObjectType a) =>
a -> Text -> m ()
Named.object Framebuffer
framebuffer Text
"ShadowMap.FB"

  pure (AllocatedImage
depth, Framebuffer
framebuffer)