{-# LANGUAGE OverloadedLists #-}

{-| Shared graphics-pipeline construction behind the two rendering paths,
"Vulkan.Utils.RenderPass" and "Vulkan.Utils.DynamicRendering". Not meant for
direct use — import one of those modules instead.

The only difference between the paths is whether the pipeline references a
'Vk.RenderPass' or carries a @PipelineRenderingCreateInfo@ in its pNext chain,
so everything else (the vanilla rasterizer/blend/dynamic-state config, the
transient empty layout, the shader-module lifetime) lives here once.
-}
module Vulkan.Utils.Pipeline.Internal
  ( basePipelineCreateInfo
  , buildColorPipeline
  , withCompiledStages
  ) where

import Control.Monad.IO.Unlift (MonadUnliftIO)
import Control.Monad.Trans.Resource (MonadResource, ReleaseKey, allocate, release)
import Data.Bits ((.|.))
import Data.ByteString (ByteString)
import Data.Foldable (traverse_)
import Data.Maybe (fromMaybe)
import Data.Vector (Vector)
import qualified Data.Vector as V
import Vulkan.CStruct.Extends (SomeStruct (..))
import qualified Vulkan.Core10 as Vk
import Vulkan.Utils.Pipeline.Specialization (Specialization, withSpecialization)
import Vulkan.Utils.Shader (shaderModuleStage)
import Vulkan.Zero (zero)

{- | The shared body of the vanilla graphics pipeline: the given @dynamicStates@,
@colorAttachmentCount@ identical non-blended color attachments, an optional
depth-stencil state, and empty vertex input. The static values left in the
create-info (cull mode, topology, …) are the baked defaults for any state not
listed dynamic; states that are listed dynamic ignore them, so callers MUST emit
the matching @cmdSet*@ before drawing.

The colour and depth shape MUST match the attachments the pipeline renders to —
the render pass (render-pass path) or the @PipelineRenderingCreateInfo@ formats
(dynamic-rendering path):

  * @colorAttachmentCount == 0@ omits @colorBlendState@ entirely (a depth-only
    pipeline); otherwise one RGBA, non-blended attachment per colour target.
  * @depth@ adds a zeroed @depthStencilState@ — present (non-NULL) is required
    whenever a depth attachment is used; the actual test config is dynamic, so a
    zeroed struct is correct.

Pass @Just@ the target render pass (render-pass path), or @Nothing@ and attach
a @PipelineRenderingCreateInfo@ to the returned struct's pNext chain
(dynamic-rendering path).
-}
basePipelineCreateInfo
  :: Vk.PipelineLayout
  -> Maybe Vk.RenderPass
  -> Int
  -- ^ Colour attachment count (blend attachments); @0@ for depth-only.
  -> Bool
  -- ^ Whether a depth attachment is present.
  -> Vk.PipelineVertexInputStateCreateInfo '[]
  -- ^ Vertex input (bindings + attributes); @zero@ for none.
  -> Vector Vk.DynamicState
  -> Vector (SomeStruct Vk.PipelineShaderStageCreateInfo)
  -> Vk.GraphicsPipelineCreateInfo '[]
basePipelineCreateInfo :: PipelineLayout
-> Maybe RenderPass
-> Int
-> Bool
-> PipelineVertexInputStateCreateInfo '[]
-> Vector DynamicState
-> Vector (SomeStruct PipelineShaderStageCreateInfo)
-> GraphicsPipelineCreateInfo '[]
basePipelineCreateInfo PipelineLayout
pipelineLayout Maybe RenderPass
renderPass Int
colorAttachmentCount Bool
depth PipelineVertexInputStateCreateInfo '[]
vertexInput Vector DynamicState
dynamicStates' Vector (SomeStruct PipelineShaderStageCreateInfo)
stages =
  GraphicsPipelineCreateInfo '[]
forall a. Zero a => a
zero
    { Vk.stages = stages
    , Vk.vertexInputState = Just (SomeStruct vertexInput)
    , Vk.inputAssemblyState =
        Just
          zero
            { Vk.topology = Vk.PRIMITIVE_TOPOLOGY_TRIANGLE_LIST
            , Vk.primitiveRestartEnable = False
            }
    , Vk.viewportState =
        Just $
          SomeStruct
            zero
              { -- The counts MUST be zero when the matching @*_WITH_COUNT@ state
                -- is dynamic (set then via @cmdSetViewportWithCount@); otherwise
                -- the static count stands (plain @VIEWPORT@/@SCISSOR@ only swap
                -- the values). VUID-VkGraphicsPipelineCreateInfo-pDynamicStates-03379/03380.
                Vk.viewportCount =
                  if Vk.DYNAMIC_STATE_VIEWPORT_WITH_COUNT `V.elem` dynamicStates' then 0 else 1
              , Vk.scissorCount =
                  if Vk.DYNAMIC_STATE_SCISSOR_WITH_COUNT `V.elem` dynamicStates' then 0 else 1
              }
    , Vk.rasterizationState =
        Just $
          SomeStruct
            zero
              { Vk.depthClampEnable = False
              , Vk.rasterizerDiscardEnable = False
              , Vk.lineWidth = 1
              , Vk.polygonMode = Vk.POLYGON_MODE_FILL
              , Vk.cullMode = Vk.CULL_MODE_NONE
              , Vk.frontFace = Vk.FRONT_FACE_COUNTER_CLOCKWISE
              , Vk.depthBiasEnable = False
              }
    , Vk.multisampleState =
        Just $
          SomeStruct
            zero
              { Vk.sampleShadingEnable = False
              , Vk.rasterizationSamples = Vk.SAMPLE_COUNT_1_BIT
              , Vk.minSampleShading = 1
              , Vk.sampleMask = [maxBound]
              }
    , Vk.depthStencilState =
        if depth then Just zero else Nothing
    , Vk.colorBlendState =
        if colorAttachmentCount == 0
          then Nothing
          else
            Just $
              SomeStruct
                zero
                  { Vk.logicOpEnable = False
                  , Vk.attachments = V.replicate colorAttachmentCount colorBlendAttachment
                  }
    , Vk.dynamicState = Just zero{Vk.dynamicStates = dynamicStates'}
    , Vk.layout = pipelineLayout
    , Vk.renderPass = fromMaybe Vk.NULL_HANDLE renderPass
    , Vk.subpass = 0
    , Vk.basePipelineHandle = zero
    }
  where
    colorBlendAttachment :: Vk.PipelineColorBlendAttachmentState
    colorBlendAttachment :: PipelineColorBlendAttachmentState
colorBlendAttachment =
      PipelineColorBlendAttachmentState
forall a. Zero a => a
zero
        { Vk.colorWriteMask =
            Vk.COLOR_COMPONENT_R_BIT
              .|. Vk.COLOR_COMPONENT_G_BIT
              .|. Vk.COLOR_COMPONENT_B_BIT
              .|. Vk.COLOR_COMPONENT_A_BIT
        , Vk.blendEnable = False
        }

{- | Build a single graphics pipeline from the given create-info builder. With
'Nothing', a transient empty pipeline layout is allocated and freed after the
build (the historical behaviour: no descriptor sets, no push constants); with
@Just layout@, the caller's layout is used and remains owned (and kept alive)
by the caller. The returned 'ReleaseKey' frees the pipeline.
-}
buildColorPipeline
  :: (MonadResource m, MonadFail m)
  => Vk.Device
  -> Maybe Vk.PipelineLayout
  -> (Vk.PipelineLayout -> SomeStruct Vk.GraphicsPipelineCreateInfo)
  -> m (ReleaseKey, Vk.Pipeline)
buildColorPipeline :: forall (m :: * -> *).
(MonadResource m, MonadFail m) =>
Device
-> Maybe PipelineLayout
-> (PipelineLayout -> SomeStruct GraphicsPipelineCreateInfo)
-> m (ReleaseKey, Pipeline)
buildColorPipeline Device
dev Maybe PipelineLayout
layout PipelineLayout -> SomeStruct GraphicsPipelineCreateInfo
mkCreateInfo = case Maybe PipelineLayout
layout of
  Just PipelineLayout
pipelineLayout -> PipelineLayout -> m (ReleaseKey, Pipeline)
build PipelineLayout
pipelineLayout
  Maybe PipelineLayout
Nothing -> do
    (layoutKey, pipelineLayout) <- Device
-> PipelineLayoutCreateInfo
-> Maybe AllocationCallbacks
-> (IO PipelineLayout
    -> (PipelineLayout -> IO ()) -> m (ReleaseKey, PipelineLayout))
-> m (ReleaseKey, PipelineLayout)
forall (io :: * -> *) r.
MonadIO io =>
Device
-> PipelineLayoutCreateInfo
-> Maybe AllocationCallbacks
-> (io PipelineLayout -> (PipelineLayout -> io ()) -> r)
-> r
Vk.withPipelineLayout Device
dev PipelineLayoutCreateInfo
forall a. Zero a => a
zero Maybe AllocationCallbacks
forall a. Maybe a
Nothing IO PipelineLayout
-> (PipelineLayout -> IO ()) -> m (ReleaseKey, PipelineLayout)
forall (m :: * -> *) a.
MonadResource m =>
IO a -> (a -> IO ()) -> m (ReleaseKey, a)
allocate
    built <- build pipelineLayout
    release layoutKey
    pure built
  where
    build :: PipelineLayout -> m (ReleaseKey, Pipeline)
build PipelineLayout
pipelineLayout = do
      (key, (_, [pipeline])) <-
        Device
-> PipelineCache
-> Vector (SomeStruct GraphicsPipelineCreateInfo)
-> Maybe AllocationCallbacks
-> (IO (Result, Vector Pipeline)
    -> ((Result, Vector Pipeline) -> IO ())
    -> m (ReleaseKey, (Result, Vector Pipeline)))
-> m (ReleaseKey, (Result, Vector Pipeline))
forall (io :: * -> *) r.
MonadIO io =>
Device
-> PipelineCache
-> Vector (SomeStruct GraphicsPipelineCreateInfo)
-> Maybe AllocationCallbacks
-> (io (Result, Vector Pipeline)
    -> ((Result, Vector Pipeline) -> io ()) -> r)
-> r
Vk.withGraphicsPipelines Device
dev PipelineCache
forall a. Zero a => a
zero [PipelineLayout -> SomeStruct GraphicsPipelineCreateInfo
mkCreateInfo PipelineLayout
pipelineLayout] Maybe AllocationCallbacks
forall a. Maybe a
Nothing IO (Result, Vector Pipeline)
-> ((Result, Vector Pipeline) -> IO ())
-> m (ReleaseKey, (Result, Vector Pipeline))
forall (m :: * -> *) a.
MonadResource m =>
IO a -> (a -> IO ()) -> m (ReleaseKey, a)
allocate
      pure (key, pipeline)

{- | Compile each @(stage, SPIR-V)@ pair into a shader module, run the
continuation with the resulting stages, then release the now-redundant
module handles. Shader modules are only needed during pipeline creation, so
the continuation typically returns the built pipeline.
-}
withCompiledStages
  :: (MonadResource m, MonadUnliftIO m, Specialization spec)
  => Vk.Device
  -> spec
  -> [(Vk.ShaderStageFlagBits, ByteString)]
  -> (Vector (SomeStruct Vk.PipelineShaderStageCreateInfo) -> m a)
  -> m a
withCompiledStages :: forall (m :: * -> *) spec a.
(MonadResource m, MonadUnliftIO m, Specialization spec) =>
Device
-> spec
-> [(ShaderStageFlagBits, ByteString)]
-> (Vector (SomeStruct PipelineShaderStageCreateInfo) -> m a)
-> m a
withCompiledStages Device
dev spec
spec [(ShaderStageFlagBits, ByteString)]
shaders Vector (SomeStruct PipelineShaderStageCreateInfo) -> m a
k =
  spec -> (Maybe SpecializationInfo -> m a) -> m a
forall spec (m :: * -> *) a.
(Specialization spec, MonadUnliftIO m) =>
spec -> (Maybe SpecializationInfo -> m a) -> m a
withSpecialization spec
spec ((Maybe SpecializationInfo -> m a) -> m a)
-> (Maybe SpecializationInfo -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \Maybe SpecializationInfo
specializationInfo -> do
    compiled <-
      ((ShaderStageFlagBits, ByteString)
 -> m (ReleaseKey, SomeStruct PipelineShaderStageCreateInfo))
-> [(ShaderStageFlagBits, ByteString)]
-> m [(ReleaseKey, SomeStruct PipelineShaderStageCreateInfo)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse
        (\(ShaderStageFlagBits
stage, ByteString
code) -> Device
-> ShaderStageFlagBits
-> Maybe SpecializationInfo
-> ByteString
-> m (ReleaseKey, SomeStruct PipelineShaderStageCreateInfo)
forall (m :: * -> *).
MonadResource m =>
Device
-> ShaderStageFlagBits
-> Maybe SpecializationInfo
-> ByteString
-> m (ReleaseKey, SomeStruct PipelineShaderStageCreateInfo)
shaderModuleStage Device
dev ShaderStageFlagBits
stage Maybe SpecializationInfo
specializationInfo ByteString
code)
        [(ShaderStageFlagBits, ByteString)]
shaders
    let (keys, stages) = unzip compiled
    result <- k (V.fromList stages)
    traverse_ release keys
    pure result