Downsample.hs

{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE OverloadedRecordDot #-}

module Resource.Image.Downsample
  ( mipPyramid
    -- * Internals
  ,
  ) where

import RIO

import Data.Bits (shiftR)
import Vulkan.CStruct.Extends (SomeStruct(..))
import Vulkan.Core10 qualified as ImageBlit (ImageBlit(..))
import Vulkan.Core10 qualified as ImageSubresourceRange (ImageSubresourceRange(..))
import Vulkan.Core10 qualified as Vk
import Vulkan.NamedType ((:::))
import Vulkan.Zero (zero)

import Engine.Vulkan.Types (MonadVulkan)
import Resource.Image (AllocatedImage, DstImage(..))
import Resource.Image qualified as Image

-- https://docs.vulkan.org/samples/latest/samples/api/texture_mipmap_generation/README.html
mipPyramid
  :: MonadVulkan env m
  => Vk.CommandBuffer -- XXX: tag to require qGraphics pool
  -> AllocatedImage
  -> Vk.ImageLayout
  -> DstImage
  -> m AllocatedImage
mipPyramid cmd src srcLayout dst = do
  -- dst is undefined everywhere, src is srcLayout
  copyImageBaseToDst cmd src srcLayout dst
  -- src copied and reverted, dst[0] is dst-optimal, the rest are undefined
  transferTransitions cmd dst Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL Vk.IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL 0
  -- dst[0] is src-optimal, dst[1+] remain undefined
  for_ @[] [1 .. dst.aiImageRange.levelCount - 1] \toMip -> do
    -- dst[toMip-1] is src-optimal, dst[toMip] is undefined
    transferTransitions cmd dst Vk.IMAGE_LAYOUT_UNDEFINED Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL toMip
    -- dst[toMip] is dst-optimal
    blitMipDown cmd dst (toMip - 1) toMip
    transferTransitions cmd dst Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL Vk.IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL toMip
    -- dst[toMip] is src-optimal, will become dst[toMip - 1]
  -- dst[*] are all src-optimal
  finalizeDst cmd dst srcLayout
  -- dst[*] are all srcLayout

{- | Transition src image to src-optimal, copy the data, and transition it back.

The dst image asssumed to be in the dst-optimal layout and not transitioned.
-}
copyImageBaseToDst
  :: MonadIO m
  => Vk.CommandBuffer
  -> "src image"  ::: AllocatedImage
  -> "src layout" ::: Vk.ImageLayout
  -> "dst image"  ::: DstImage
  -> m ()
copyImageBaseToDst cmd src srcLayout (DstImage dst) = liftIO do
  Vk.cmdPipelineBarrier
    cmd
    Vk.PIPELINE_STAGE_TOP_OF_PIPE_BIT
    Vk.PIPELINE_STAGE_TRANSFER_BIT
    zero mempty mempty
    [prepareDst, prepareSrc]

  {-
  VUID-vkCmdCopyImage-srcImageLayout-01917
    srcImageLayout must be VK_IMAGE_LAYOUT_SHARED_PRESENT_KHR, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, or VK_IMAGE_LAYOUT_GENERAL
  -}
  Vk.cmdCopyImage
    cmd
    src.aiImage
    Vk.IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL
    dst.aiImage
    Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL
    [copyBase]

  Vk.cmdPipelineBarrier
    cmd
    Vk.PIPELINE_STAGE_TRANSFER_BIT
    Vk.PIPELINE_STAGE_FRAGMENT_SHADER_BIT -- hm...
    zero mempty mempty
    [revertSrc]
  where
    prepareDst = SomeStruct zero
      { Vk.srcAccessMask       = zero
      , Vk.dstAccessMask       = Vk.ACCESS_TRANSFER_WRITE_BIT
      , Vk.oldLayout           = Vk.IMAGE_LAYOUT_UNDEFINED
      , Vk.newLayout           = Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL
      , Vk.srcQueueFamilyIndex = Vk.QUEUE_FAMILY_IGNORED
      , Vk.dstQueueFamilyIndex = Vk.QUEUE_FAMILY_IGNORED
      , Vk.image               = dst.aiImage
      , Vk.subresourceRange    = dst.aiImageRange
      }

    onlyBaseLevel = src.aiImageRange
      { ImageSubresourceRange.levelCount = 1
      }

    prepareSrc = SomeStruct zero
      { Vk.srcAccessMask       = Vk.ACCESS_SHADER_WRITE_BIT
      , Vk.dstAccessMask       = Vk.ACCESS_TRANSFER_WRITE_BIT
      , Vk.oldLayout           = srcLayout
      , Vk.newLayout           = Vk.IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL
      , Vk.srcQueueFamilyIndex = Vk.QUEUE_FAMILY_IGNORED
      , Vk.dstQueueFamilyIndex = Vk.QUEUE_FAMILY_IGNORED
      , Vk.image               = src.aiImage
      , Vk.subresourceRange    = onlyBaseLevel
      }

    copyBase = Vk.ImageCopy
      { srcSubresource = baseISL
      , srcOffset = zero
      , dstSubresource = baseISL
      , dstOffset = zero
      , extent = src.aiExtent
      }
      where
        baseISL = Vk.ImageSubresourceLayers
          { aspectMask     = Vk.IMAGE_ASPECT_COLOR_BIT
          , mipLevel       = 0
          , baseArrayLayer = 0
          , layerCount     = Vk.REMAINING_ARRAY_LAYERS
          }

    revertSrc = SomeStruct zero
      { Vk.srcAccessMask       = Vk.ACCESS_TRANSFER_WRITE_BIT
      , Vk.dstAccessMask       = Vk.ACCESS_SHADER_READ_BIT
      , Vk.oldLayout           = Vk.IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL
      , Vk.newLayout           = srcLayout
      , Vk.srcQueueFamilyIndex = Vk.QUEUE_FAMILY_IGNORED
      , Vk.dstQueueFamilyIndex = Vk.QUEUE_FAMILY_IGNORED
      , Vk.image               = src.aiImage
      , Vk.subresourceRange    = onlyBaseLevel
      }

transferTransitions
  :: MonadIO m
  => Vk.CommandBuffer
  -> DstImage
  -> Vk.ImageLayout
  -> Vk.ImageLayout
  -> Word32
  -> m ()
transferTransitions cmd (DstImage image) old new mip =
  Vk.cmdPipelineBarrier cmd srcStage Vk.PIPELINE_STAGE_TRANSFER_BIT zero mempty mempty [barrier]
  where
    barrier = SomeStruct zero
      { Vk.srcAccessMask       = srcAccess
      , Vk.dstAccessMask       = dstAccess
      , Vk.oldLayout           = old
      , Vk.newLayout           = new
      , Vk.srcQueueFamilyIndex = Vk.QUEUE_FAMILY_IGNORED
      , Vk.dstQueueFamilyIndex = Vk.QUEUE_FAMILY_IGNORED
      , Vk.image               = image.aiImage
      , Vk.subresourceRange    = subr
      }
    srcAccess = case old of
      Vk.IMAGE_LAYOUT_UNDEFINED -> zero
      Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL -> Vk.ACCESS_TRANSFER_WRITE_BIT
      Vk.IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL -> Vk.ACCESS_TRANSFER_READ_BIT
      _ -> error $ "unexpected old layout: " <> show old
    dstAccess = case new of
      Vk.IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL -> Vk.ACCESS_TRANSFER_READ_BIT
      Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL -> Vk.ACCESS_TRANSFER_WRITE_BIT
      _ -> error $ "unexpected new layout: " <> show new
    srcStage = case old of
      Vk.IMAGE_LAYOUT_UNDEFINED -> Vk.PIPELINE_STAGE_TOP_OF_PIPE_BIT
      _ -> Vk.PIPELINE_STAGE_TRANSFER_BIT
    subr = Vk.ImageSubresourceRange
      { aspectMask     = image.aiImageRange.aspectMask
      , baseMipLevel   = mip
      , levelCount     = 1
      , baseArrayLayer = 0
      , layerCount     = Vk.REMAINING_ARRAY_LAYERS
      }

blitMipDown
  :: MonadIO m
  => Vk.CommandBuffer -- XXX: tag to require qGraphics pool
  -> DstImage
  -> Word32
  -> Word32
  -> m ()
blitMipDown cmd (DstImage image) fromMip toMip =
  Vk.cmdBlitImage
    cmd
    image.aiImage
    Vk.IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL
    image.aiImage
    Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL
    [region]
    Vk.FILTER_LINEAR
  where
    region = zero
      { ImageBlit.srcSubresource = prevLayers
      , ImageBlit.srcOffsets =
          ( zero
          , Vk.Offset3D
              (max 1 $ fromIntegral image.aiExtent.width `shiftR` fromIntegral fromMip)
              (max 1 $ fromIntegral image.aiExtent.height `shiftR` fromIntegral fromMip)
              1
          )
      , ImageBlit.dstSubresource = currLayers
      , ImageBlit.dstOffsets =
          ( zero
          , Vk.Offset3D
              (max 1 $ fromIntegral image.aiExtent.width `shiftR` fromIntegral toMip)
              (max 1 $ fromIntegral image.aiExtent.height `shiftR` fromIntegral toMip)
              1
          )
      }
      where
        prevLayers = Vk.ImageSubresourceLayers
          { aspectMask     = image.aiImageRange.aspectMask
          , mipLevel       = fromMip
          , baseArrayLayer = 0
          , layerCount     = Vk.REMAINING_ARRAY_LAYERS
          }
        currLayers = Vk.ImageSubresourceLayers
          { aspectMask     = image.aiImageRange.aspectMask
          , mipLevel       = toMip
          , baseArrayLayer = 0
          , layerCount     = Vk.REMAINING_ARRAY_LAYERS
          }

finalizeDst
  :: MonadIO m
  => Vk.CommandBuffer
  -> DstImage
  -> Vk.ImageLayout
  -> m AllocatedImage
finalizeDst cmd (DstImage image) toLayout = do
  Vk.cmdPipelineBarrier
    cmd
    Vk.PIPELINE_STAGE_TRANSFER_BIT
    Vk.PIPELINE_STAGE_FRAGMENT_SHADER_BIT
    zero
    mempty
    mempty
    [barrier]
  pure image
  where
    barrier = SomeStruct zero
      { Vk.srcAccessMask       = Vk.ACCESS_TRANSFER_READ_BIT
      , Vk.dstAccessMask       = Vk.ACCESS_SHADER_READ_BIT
      , Vk.oldLayout           = Vk.IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL
      , Vk.newLayout           = toLayout
      , Vk.srcQueueFamilyIndex = Vk.QUEUE_FAMILY_IGNORED
      , Vk.dstQueueFamilyIndex = Vk.QUEUE_FAMILY_IGNORED
      , Vk.image               = image.aiImage
      , Vk.subresourceRange    = image.aiImageRange
      }