{-# 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.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
  :: ( TextureLayers a
     , MonadVulkan env m
     , MonadResource m
     , MonadThrow m
     , HasLogFunc env
     , Typeable a
     , HasCallStack
     )
  => Queues Vk.CommandPool
  -> Source
  -> m (Texture a)
load :: forall a env (m :: * -> *).
(TextureLayers a, MonadVulkan env m, MonadResource m, MonadThrow m,
 HasLogFunc env, Typeable a, HasCallStack) =>
Queues CommandPool -> Source -> m (Texture a)
load Queues CommandPool
pool Source
source =
  forall a. HasCallStack => (HasCallStack => a) -> a
withFrozenCallStack forall a b. (a -> b) -> a -> b
$
    case Source
source of
      Source.File Maybe Text
label FilePath
path ->
        -- XXX: the codec has a more efficient loader for files
        forall a env (m :: * -> *).
(TextureLayers a, MonadVulkan env m, MonadResource m, MonadThrow m,
 HasLogFunc env) =>
Maybe Text -> Queues CommandPool -> FilePath -> m (Texture a)
loadFile Maybe Text
label Queues CommandPool
pool FilePath
path
      Source
_bytes ->
        forall a (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, Typeable a,
 HasCallStack) =>
(ByteString -> m a) -> Source -> m a
Source.load (forall a env (m :: * -> *).
(TextureLayers a, MonadVulkan env m, MonadResource m, MonadThrow m,
 HasLogFunc env) =>
Maybe Text -> Queues CommandPool -> ByteString -> m (Texture a)
loadBytes Source
source.label Queues CommandPool
pool) Source
source

loadFile
  :: ( TextureLayers a
     , MonadVulkan env m
     , MonadResource m
     , MonadThrow m
     , HasLogFunc env
     )
  => Maybe Text
  -> Queues Vk.CommandPool
  -> FilePath
  -> m (Texture a)
loadFile :: forall a env (m :: * -> *).
(TextureLayers a, MonadVulkan env m, MonadResource m, MonadThrow m,
 HasLogFunc env) =>
Maybe Text -> Queues CommandPool -> FilePath -> m (Texture a)
loadFile Maybe Text
label Queues CommandPool
pool FilePath
path =
  forall (m :: * -> *) a b c.
MonadUnliftIO m =>
m a -> (a -> m b) -> (a -> m c) -> m c
bracket (forall (io :: * -> *). MonadIO io => FilePath -> io FileContext
Read.open FilePath
path) forall (io :: * -> *). MonadIO io => FileContext -> io ()
Read.close forall a b. (a -> b) -> a -> b
$
    forall a (m :: * -> *) env src.
(TextureLayers a, MonadVulkan env m, MonadResource m, MonadThrow m,
 HasLogFunc env, ReadChunk src, ReadLevel src) =>
Maybe Text -> Queues CommandPool -> Context src -> m (Texture a)
loadKtx2 (Maybe Text
label forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> forall a. a -> Maybe a
Just (forall a. IsString a => FilePath -> a
fromString FilePath
path)) Queues CommandPool
pool

loadBytes
  :: ( TextureLayers a
     , MonadVulkan env m
     , MonadResource m
     , MonadThrow m
     , HasLogFunc env
     )
  => Maybe Text
  -> Queues Vk.CommandPool
  -> ByteString
  -> m (Texture a)
loadBytes :: forall a env (m :: * -> *).
(TextureLayers a, MonadVulkan env m, MonadResource m, MonadThrow m,
 HasLogFunc env) =>
Maybe Text -> Queues CommandPool -> ByteString -> m (Texture a)
loadBytes Maybe Text
label Queues CommandPool
pool ByteString
bytes =
  forall (io :: * -> *). MonadIO io => ByteString -> io BytesContext
Read.bytes ByteString
bytes forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall a (m :: * -> *) env src.
(TextureLayers a, MonadVulkan env m, MonadResource m, MonadThrow m,
 HasLogFunc env, ReadChunk src, ReadLevel src) =>
Maybe Text -> Queues CommandPool -> Context src -> m (Texture a)
loadKtx2 Maybe Text
label Queues CommandPool
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 :: forall a (m :: * -> *) env src.
(TextureLayers a, MonadVulkan env m, MonadResource m, MonadThrow m,
 HasLogFunc env, ReadChunk src, ReadLevel src) =>
Maybe Text -> Queues CommandPool -> Context src -> m (Texture a)
loadKtx2 Maybe Text
label Queues CommandPool
pool ktx2 :: Context src
ktx2@(Read.Context src
_src Header
header) = do
  forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) =>
Utf8Builder -> m ()
logDebug forall a b. (a -> b) -> a -> b
$ forall a. Show a => a -> Utf8Builder
displayShow (Maybe Text
label, Header
header)

  KeyValueData
kvd <- forall src (io :: * -> *).
(ReadChunk src, MonadIO io) =>
Context src -> io KeyValueData
Read.keyValueData Context src
ktx2
  forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) =>
Utf8Builder -> m ()
logDebug forall a b. (a -> b) -> a -> b
$ forall a. Show a => a -> Utf8Builder
displayShow (Maybe Text
label, Format
format, Extent3D
extent, Word32
numLayers, KeyValueData -> Map Text Text
KeyValue.textual KeyValueData
kvd)

  Vector Level
levels <- forall src (io :: * -> *).
(ReadChunk src, MonadIO io) =>
Context src -> io (Vector Level)
Read.levels Context src
ktx2

  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (forall a. Vector a -> Int
Vector.length Vector Level
levels forall a. Ord a => a -> a -> Bool
> Int
0) forall a b. (a -> b) -> a -> b
$
    forall (m :: * -> *) a.
(MonadIO m, HasCallStack) =>
FilePath -> m a
throwString FilePath
"Ktx2 contains no image levels"

  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (forall a. Vector a -> Int
Vector.length Vector Level
levels forall a. Eq a => a -> a -> Bool
== forall a b. (Integral a, Num b) => a -> b
fromIntegral Header
header.levelCount) forall a b. (a -> b) -> a -> b
$
    forall (m :: * -> *) a.
(MonadIO m, HasCallStack) =>
FilePath -> m a
throwString forall a b. (a -> b) -> a -> b
$ FilePath
"Ktx2 level count mismatch " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> FilePath
show (forall a. Vector a -> Int
Vector.length Vector Level
levels, Header
header.levelCount)

  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Word32
numLayers forall a. Eq a => a -> a -> Bool
== forall a. TextureLayers a => Word32
textureLayers @a) forall a b. (a -> b) -> a -> b
$
    forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ Word32 -> Word32 -> TextureError
Texture.ArrayError (forall a. TextureLayers a => Word32
textureLayers @a) Word32
numLayers

  let
    levelSizes :: Vector Int
levelSizes = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall b c a. (b -> c) -> (a -> b) -> a -> c
. (.uncompressedByteLength)) Vector Level
levels
    totalSize :: Int
totalSize = forall a. Num a => Vector a -> a
Vector.sum Vector Int
levelSizes
    offsets :: Vector Int
offsets = forall a. Vector a -> Vector a
Vector.init forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b -> a) -> a -> Vector b -> Vector a
Vector.scanl' forall a. Num a => a -> a -> a
(+) Int
0 Vector Int
levelSizes

  DstImage
dst <- forall env (m :: * -> *).
(MonadVulkan env m, MonadResource m) =>
Queues CommandPool
-> Maybe Text
-> Extent3D
-> Word32
-> Word32
-> Format
-> m DstImage
Image.allocateDst
    Queues CommandPool
pool
    Maybe Text
label
    Extent3D
extent
    (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ forall a. Vector a -> Int
Vector.length Vector Level
levels)
    Word32
numLayers
    Format
format

  Allocator
vma <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a. HasVulkan a => a -> Allocator
getAllocator
  forall (a :: [*]) (io :: * -> *) r.
(Extendss BufferCreateInfo a, PokeChain a, MonadIO io) =>
Allocator
-> BufferCreateInfo a
-> AllocationCreateInfo
-> (io (Buffer, Allocation, AllocationInfo)
    -> ((Buffer, Allocation, AllocationInfo) -> io ()) -> r)
-> r
VMA.withBuffer Allocator
vma (forall a. Integral a => a -> BufferCreateInfo '[]
Texture.stageBufferCI Int
totalSize) AllocationCreateInfo
Texture.stageAllocationCI forall (m :: * -> *) a b c.
MonadUnliftIO m =>
m a -> (a -> m b) -> (a -> m c) -> m c
bracket \(Buffer
staging, Allocation
stage, AllocationInfo
stageInfo) -> do
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO case Header
header.supercompressionScheme of
      Word32
Header.SC_NONE ->
        forall (m :: * -> *) a b. Monad m => Vector a -> (a -> m b) -> m ()
Vector.forM_ (forall a b. Vector a -> Vector b -> Vector (a, b)
Vector.zip Vector Int
offsets Vector Level
levels) \(Int
offset, Level
level) ->
          forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall src (io :: * -> *).
(ReadLevel src, MonadIO io) =>
Context src -> Level -> Ptr () -> io Bool
Read.levelToPtr Context src
ktx2 Level
level forall a b. (a -> b) -> a -> b
$
            forall a b. Ptr a -> Int -> Ptr b
Foreign.plusPtr (AllocationInfo -> Ptr ()
VMA.mappedData AllocationInfo
stageInfo) Int
offset

      Word32
Header.SC_ZSTANDARD -> do
        let maxSize :: Int
maxSize = forall a. Ord a => Vector a -> a
Vector.maximum Vector Int
levelSizes
        forall a b. Int -> Int -> (Ptr a -> IO b) -> IO b
Foreign.allocaBytesAligned Int
maxSize Int
16 \Ptr ()
src ->
          forall (m :: * -> *) a b. Monad m => Vector a -> (a -> m b) -> m ()
Vector.forM_ (forall a b. Vector a -> Vector b -> Vector (a, b)
Vector.zip Vector Int
offsets Vector Level
levels) \(Int
offset, Level
level) -> do
            let expected :: CSize
expected = forall a b. (Integral a, Num b) => a -> b
fromIntegral Level
level.uncompressedByteLength
            forall src (io :: * -> *).
(ReadLevel src, MonadIO io) =>
Context src -> Level -> Ptr () -> io Bool
Read.levelToPtr Context src
ktx2 Level
level Ptr ()
src

            Either FilePath CSize
res <-
              IO CSize -> IO (Either FilePath CSize)
Zstd.checkError forall a b. (a -> b) -> a -> b
$
                forall dst src. Ptr dst -> CSize -> Ptr src -> CSize -> IO CSize
Zstd.decompress
                  (forall a b. Ptr a -> Int -> Ptr b
Foreign.plusPtr (AllocationInfo -> Ptr ()
VMA.mappedData AllocationInfo
stageInfo) Int
offset)
                  CSize
expected
                  Ptr ()
src
                  (forall a b. (Integral a, Num b) => a -> b
fromIntegral Level
level.byteLength)
            case Either FilePath CSize
res of
              Right CSize
size | CSize
size forall a. Eq a => a -> a -> Bool
== CSize
expected ->
                forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
              Right CSize
unexpected ->
                forall (m :: * -> *) a.
(MonadIO m, HasCallStack) =>
FilePath -> m a
throwString forall a b. (a -> b) -> a -> b
$
                  FilePath
"Zstd decompressed unexpected amount of bytes: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> FilePath
show (CSize
unexpected, CSize
expected)
              Left FilePath
err ->
                forall (m :: * -> *) a.
(MonadIO m, HasCallStack) =>
FilePath -> m a
throwString FilePath
err

      Word32
huh ->
        forall a. HasCallStack => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Unexpected supercompression scheme: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> FilePath
show Word32
huh

    forall (io :: * -> *).
MonadIO io =>
Allocator -> Allocation -> Word64 -> Word64 -> io ()
VMA.flushAllocation Allocator
vma Allocation
stage Word64
0 Word64
Vk.WHOLE_SIZE

    AllocatedImage
final <- forall env (m :: * -> *) deviceSize (t :: * -> *).
(MonadVulkan env m, Integral deviceSize, Foldable t) =>
Queues CommandPool
-> Buffer
-> DstImage
-> ("mip offsets" ::: t deviceSize)
-> m AllocatedImage
Image.copyBufferToDst
      Queues CommandPool
pool
      Buffer
staging
      DstImage
dst
      Vector Int
offsets

    pure Texture
      { $sel:tFormat:Texture :: Format
tFormat         = Format
format
      , $sel:tMipLevels:Texture :: Word32
tMipLevels      = Header
header.levelCount
      , $sel:tLayers:Texture :: Word32
tLayers         = Word32
numLayers
      , $sel:tAllocatedImage:Texture :: AllocatedImage
tAllocatedImage = AllocatedImage
final
      }
  where
    format :: Format
format = Int32 -> Format
Vk.Format (forall a b. (Integral a, Num b) => a -> b
fromIntegral Header
header.vkFormat)

    extent :: Extent3D
extent = Vk.Extent3D
      { $sel:width:Extent3D :: Word32
width = Header
header.pixelWidth
      , $sel:height:Extent3D :: Word32
height = Header
header.pixelHeight
      , $sel:depth:Extent3D :: Word32
depth = forall a. Ord a => a -> a -> a
max Word32
1 Header
header.pixelDepth
      }

    -- XXX: can be flat array or a cubemap
    numLayers :: Word32
numLayers = forall a. Ord a => a -> a -> a
max Header
header.faceCount Header
header.layerCount