Ktx2.hs

{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE NoForeignFunctionInterface #-}

module Resource.Texture.Ktx2
  ( load
  , loadBytes
  , loadKtx2
  ) where

import RIO

import Codec.Compression.Zstd.FFI qualified as Zstd
import Codec.Ktx.KeyValue qualified as KeyValue
import Codec.Ktx2.Header qualified as Header
import Codec.Ktx2.Level qualified as Level
import Codec.Ktx2.Read qualified as Read
import Data.Kind (Type)
import Data.Vector qualified as Vector
import Foreign qualified
import GHC.Stack (withFrozenCallStack)
import UnliftIO.Resource (MonadResource)
import Vulkan.Core10 qualified as Vk
import VulkanMemoryAllocator qualified as VMA

import Engine.Vulkan.Types (HasVulkan(..), MonadVulkan, Queues)
import Resource.Image qualified as Image
import Resource.Source (Source(..))
import Resource.Source qualified as Source
import Resource.Texture (Texture(..), TextureLayers(..))
import Resource.Texture qualified as Texture

load
  :: forall (a :: Type) env m
  .  ( TextureLayers a
     , MonadVulkan env m
     , MonadResource m
     , MonadThrow m
     , HasLogFunc env
     , Typeable a
     , HasCallStack
     )
  => Queues Vk.CommandPool
  -> Source
  -> m (Texture a)
load pool source =
  withFrozenCallStack $
    case source of
      Source.File label path ->
        -- XXX: the codec has a more efficient loader for files
        loadFile label pool path
      _bytes ->
        Source.load (loadBytes source.label pool) source

loadFile
  :: ( TextureLayers a
     , MonadVulkan env m
     , MonadResource m
     , MonadThrow m
     , HasLogFunc env
     )
  => Maybe Text
  -> Queues Vk.CommandPool
  -> FilePath
  -> m (Texture a)
loadFile label pool path =
  bracket (Read.open path) Read.close $
    loadKtx2 (label <|> Just (fromString path)) pool

loadBytes
  :: ( TextureLayers a
     , MonadVulkan env m
     , MonadResource m
     , MonadThrow m
     , HasLogFunc env
     )
  => Maybe Text
  -> Queues Vk.CommandPool
  -> ByteString
  -> m (Texture a)
loadBytes label pool bytes =
  Read.bytes bytes >>= loadKtx2 label pool

loadKtx2
  :: forall a m env src
  .  ( TextureLayers a
     , MonadVulkan env m
     , MonadResource m
     , MonadThrow m
     , HasLogFunc env
     , Read.ReadChunk src
     , Read.ReadLevel src
     )
  => Maybe Text
  -> Queues Vk.CommandPool
  -> Read.Context src
  -> m (Texture a)
loadKtx2 label pool ktx2@(Read.Context _src header) = do
  logDebug $ displayShow (label, header)

  kvd <- Read.keyValueData ktx2
  logDebug $ displayShow (label, format, extent, numLayers, KeyValue.textual kvd)

  levels <- Read.levels ktx2

  unless (Vector.length levels > 0) $
    throwString "Ktx2 contains no image levels"

  unless (Vector.length levels == fromIntegral header.levelCount) $
    throwString $ "Ktx2 level count mismatch " <> show (Vector.length levels, header.levelCount)

  unless (numLayers == textureLayers @a) $
    throwM $ Texture.ArrayError (textureLayers @a) numLayers

  let
    levelSizes = fmap (fromIntegral . (.uncompressedByteLength)) levels
    totalSize = Vector.sum levelSizes
    offsets = Vector.init $ Vector.scanl' (+) 0 levelSizes

  dst <- Image.allocateDst
    pool
    label
    extent
    (fromIntegral $ Vector.length levels)
    numLayers
    format

  vma <- asks getAllocator
  VMA.withBuffer vma (Texture.stageBufferCI totalSize) Texture.stageAllocationCI bracket \(staging, stage, stageInfo) -> do
    liftIO case header.supercompressionScheme of
      Header.SC_NONE ->
        Vector.forM_ (Vector.zip offsets levels) \(offset, level) ->
          liftIO . Read.levelToPtr ktx2 level $
            Foreign.plusPtr (VMA.mappedData stageInfo) offset

      Header.SC_ZSTANDARD -> do
        let maxSize = Vector.maximum levelSizes
        Foreign.allocaBytesAligned maxSize 16 \src ->
          Vector.forM_ (Vector.zip offsets levels) \(offset, level) -> do
            let expected = fromIntegral level.uncompressedByteLength
            Read.levelToPtr ktx2 level src

            res <-
              Zstd.checkError $
                Zstd.decompress
                  (Foreign.plusPtr (VMA.mappedData stageInfo) offset)
                  expected
                  src
                  (fromIntegral level.byteLength)
            case res of
              Right size | size == expected ->
                pure ()
              Right unexpected ->
                throwString $
                  "Zstd decompressed unexpected amount of bytes: " <> show (unexpected, expected)
              Left err ->
                throwString err

      huh ->
        error $ "Unexpected supercompression scheme: " ++ show huh

    VMA.flushAllocation vma stage 0 Vk.WHOLE_SIZE

    final <- Image.copyBufferToDst
      pool
      staging
      dst
      offsets

    pure Texture
      { tFormat         = format
      , tMipLevels      = header.levelCount
      , tLayers         = numLayers
      , tAllocatedImage = final
      }
  where
    format = Vk.Format (fromIntegral header.vkFormat)

    extent = Vk.Extent3D
      { width = header.pixelWidth
      , height = header.pixelHeight
      , depth = max 1 header.pixelDepth
      }

    -- XXX: can be flat array or a cubemap
    numLayers = max header.faceCount header.layerCount