module Resource.Texture.Ktx1
  ( load
  , loadBytes
  , loadKtx1
  ) where

import RIO

import Codec.Ktx qualified as Ktx1
import Data.ByteString.Unsafe (unsafeUseAsCStringLen)
import Data.Text qualified as Text
import Data.Vector qualified as Vector
import Foreign qualified
import GHC.Stack (withFrozenCallStack)
import Vulkan.Core10 qualified as Vk
import Vulkan.Utils.FromGL qualified as FromGL
import VulkanMemoryAllocator qualified as VMA

import Engine.Vulkan.Types (HasVulkan(..), MonadVulkan, Queues)
import Resource.Image (AllocatedImage(..))
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
     , MonadThrow m
     , HasLogFunc env
     , Typeable a
     , HasCallStack
     )
  => Queues Vk.CommandPool
  -> Source
  -> m (Texture a)
load :: forall a env (m :: * -> *).
(TextureLayers a, MonadVulkan env 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
$
    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, MonadThrow m,
 HasLogFunc env) =>
Queues CommandPool -> ByteString -> m (Texture a)
loadBytes Queues CommandPool
pool) Source
source

loadBytes
  :: ( TextureLayers a
     , MonadVulkan env m
     , MonadThrow m
     , HasLogFunc env
     )
  => Queues Vk.CommandPool
  -> ByteString
  -> m (Texture a)
loadBytes :: forall a env (m :: * -> *).
(TextureLayers a, MonadVulkan env m, MonadThrow m,
 HasLogFunc env) =>
Queues CommandPool -> ByteString -> m (Texture a)
loadBytes Queues CommandPool
pool ByteString
bytes = do
  case ByteString -> Either (ByteOffset, String) Ktx
Ktx1.fromByteString ByteString
bytes of
    Left (ByteOffset
offset, String
err) -> do
      forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) =>
Utf8Builder -> m ()
logError forall a b. (a -> b) -> a -> b
$ Utf8Builder
"Texture load error: " forall a. Semigroup a => a -> a -> a
<> forall a. IsString a => String -> a
fromString String
err
      forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ ByteOffset -> Text -> TextureError
Texture.LoadError ByteOffset
offset (String -> Text
Text.pack String
err)
    Right Ktx
ktx1 ->
      forall a (m :: * -> *) env.
(TextureLayers a, MonadVulkan env m, MonadThrow m,
 HasLogFunc env) =>
Queues CommandPool -> Ktx -> m (Texture a)
loadKtx1 Queues CommandPool
pool Ktx
ktx1

loadKtx1
  :: forall a m env
  .  ( TextureLayers a
     , MonadVulkan env m
     , MonadThrow m
     , HasLogFunc env
     )
  => Queues Vk.CommandPool
  -> Ktx1.Ktx
  -> m (Texture a)
loadKtx1 :: forall a (m :: * -> *) env.
(TextureLayers a, MonadVulkan env m, MonadThrow m,
 HasLogFunc env) =>
Queues CommandPool -> Ktx -> m (Texture a)
loadKtx1 Queues CommandPool
pool Ktx1.Ktx{Header
$sel:header:Ktx :: Ktx -> Header
header :: Header
header, $sel:images:Ktx :: Ktx -> MipLevels
images=MipLevels
ktxImages} = do
  env
context <- forall r (m :: * -> *). MonadReader r m => m r
ask
  let vma :: Allocator
vma = forall a. HasVulkan a => a -> Allocator
getAllocator env
context

  let
    skipMips :: Int
skipMips = Int
0 -- DEBUG: Vector.length ktxImages `div` 2
    mipsSkipped :: Int
mipsSkipped = forall a. Ord a => a -> a -> a
min (forall a. Vector a -> Int
Vector.length MipLevels
ktxImages forall a. Num a => a -> a -> a
- Int
1) Int
skipMips
    images :: MipLevels
images = forall a. Int -> Vector a -> Vector a
Vector.drop Int
mipsSkipped MipLevels
ktxImages
    mipLevels :: Word32
mipLevels = Header -> Word32
Ktx1.numberOfMipmapLevels Header
header forall a. Num a => a -> a -> a
- forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
mipsSkipped

  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall a. Vector a -> Bool
Vector.null MipLevels
images) 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
$ ByteOffset -> Text -> TextureError
Texture.LoadError ByteOffset
0 Text
"At least one image must be present in KTX"

  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
mipLevels forall a. Eq a => a -> a -> Bool
== forall a. Vector a -> Int
Vector.length MipLevels
images) 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 -> Int -> TextureError
Texture.MipLevelsError Word32
mipLevels (forall a. Vector a -> Int
Vector.length MipLevels
images)

  -- XXX: https://github.com/KhronosGroup/KTX-Software/blob/bf849b7f/lib/vk_format.h#L676
  Format
format <- case Header -> Word32
Ktx1.glInternalFormat Header
header of
    -- XXX: Force all BC7s to SRGB
    Word32
36492 ->
      forall (f :: * -> *) a. Applicative f => a -> f a
pure Format
Vk.FORMAT_BC7_SRGB_BLOCK
    Word32
other ->
      case forall a. (Eq a, Num a) => a -> Maybe Format
FromGL.internalFormat Word32
other of
        Maybe Format
Nothing ->
          forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$ String
"Unexpected glInternalFormat: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show Word32
other -- TODO: throwIo
        Just Format
fmt ->
          -- XXX: going in blind
          forall (f :: * -> *) a. Applicative f => a -> f a
pure Format
fmt
  forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) =>
Utf8Builder -> m ()
logDebug forall a b. (a -> b) -> a -> b
$ forall a. Monoid a => [a] -> a
mconcat
    [ Utf8Builder
"Loading format "
    , forall a. Display a => a -> Utf8Builder
display forall a b. (a -> b) -> a -> b
$ Header -> Word32
Ktx1.glInternalFormat Header
header
    , Utf8Builder
" as "
    , forall a. Show a => a -> Utf8Builder
displayShow Format
format
    ]

  -- XXX: https://github.com/KhronosGroup/KTX-Software/blob/bf849b7f/lib/vkloader.c#L552
  let
    extent :: Extent3D
extent = Vk.Extent3D
      { $sel:width:Extent3D :: Word32
Vk.width  = Header -> Word32
Ktx1.pixelWidth Header
header forall a. Bits a => a -> Int -> a
`Foreign.shiftR` Int
mipsSkipped
      , $sel:height:Extent3D :: Word32
Vk.height = Header -> Word32
Ktx1.pixelHeight Header
header forall a. Bits a => a -> Int -> a
`Foreign.shiftR` Int
mipsSkipped
      , $sel:depth:Extent3D :: Word32
Vk.depth  = forall a. Ord a => a -> a -> a
max Word32
1 forall a b. (a -> b) -> a -> b
$ Header -> Word32
Ktx1.pixelDepth Header
header
      }
    arrayLayers :: Word32
arrayLayers = forall a. Ord a => a -> a -> a
max Word32
1 forall a b. (a -> b) -> a -> b
$ Header -> Word32
Ktx1.numberOfArrayElements Header
header

  -- TODO: basisu can encode movies as arrays, this could be handy
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Word32
arrayLayers forall a. Eq a => a -> a -> Bool
== Word32
1) do
    forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) =>
Utf8Builder -> m ()
logError Utf8Builder
"TODO: arrayLayers > 1"
    forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ Word32 -> Word32 -> TextureError
Texture.ArrayError Word32
1 Word32
arrayLayers

  let
    numLayers :: Word32
numLayers = Header -> Word32
Ktx1.numberOfFaces Header
header
    mipSizes :: Vector Word32
mipSizes = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. Num a => a -> a -> a
(*) Word32
numLayers forall b c a. (b -> c) -> (a -> b) -> a -> c
. MipLevel -> Word32
Ktx1.imageSize) MipLevels
images

    offsets' :: Vector Word32
offsets' = forall a b. (a -> b -> a) -> a -> Vector b -> Vector a
Vector.scanl' forall a. Num a => a -> a -> a
(+) Word32
0 Vector Word32
mipSizes
    totalSize :: Word32
totalSize = forall a. Vector a -> a
Vector.last Vector Word32
offsets'
    offsets :: Vector Word32
offsets = forall a. Vector a -> Vector a
Vector.init Vector Word32
offsets'

  forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) =>
Utf8Builder -> m ()
logDebug forall a b. (a -> b) -> a -> b
$ Utf8Builder
"mipSizes: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> Utf8Builder
displayShow Vector Word32
mipSizes
  forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) =>
Utf8Builder -> m ()
logDebug forall a b. (a -> b) -> a -> b
$ Utf8Builder
"offsets: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> Utf8Builder
displayShow Vector Word32
offsets

  {- XXX:
    Image created before staging buffer due to `Image.copyBufferToImage`
    issued inside `VMA.withBuffer` bracket.
  -}
  (Image
image, Allocation
allocation, AllocationInfo
_info) <- forall (a :: [*]) (io :: * -> *).
(Extendss ImageCreateInfo a, PokeChain a, MonadIO io) =>
Allocator
-> ImageCreateInfo a
-> AllocationCreateInfo
-> io (Image, Allocation, AllocationInfo)
VMA.createImage
    Allocator
vma
    (Format -> Extent3D -> Word32 -> Word32 -> ImageCreateInfo '[]
Texture.imageCI Format
format Extent3D
extent Word32
mipLevels Word32
numLayers)
    AllocationCreateInfo
Texture.imageAllocationCI

  forall context (m :: * -> *).
(HasVulkan context, MonadUnliftIO m) =>
context
-> Queues CommandPool
-> Image
-> Word32
-> Word32
-> Format
-> ("old" ::: ImageLayout)
-> ("old" ::: ImageLayout)
-> m ()
Image.transitionLayout
    env
context
    Queues CommandPool
pool
    Image
image
    Word32
mipLevels
    Word32
numLayers -- XXX: arrayLayers is always 0 for now
    Format
format
    "old" ::: ImageLayout
Vk.IMAGE_LAYOUT_UNDEFINED
    "old" ::: ImageLayout
Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL

  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

  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 Word32
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
    let mipImages :: Vector (Word32, MipLevel)
mipImages = forall a b. Vector a -> Vector b -> Vector (a, b)
Vector.zip Vector Word32
offsets MipLevels
images
    forall (m :: * -> *) a b.
Monad m =>
Vector a -> (Int -> a -> m b) -> m ()
Vector.iforM_ Vector (Word32, MipLevel)
mipImages \Int
mipIx (Word32
offset, Ktx1.MipLevel{Word32
imageSize :: Word32
$sel:imageSize:MipLevel :: MipLevel -> Word32
imageSize, Vector ArrayElement
$sel:arrayElements:MipLevel :: MipLevel -> Vector ArrayElement
arrayElements :: Vector ArrayElement
arrayElements}) -> do
      forall (m :: * -> *) a b.
Monad m =>
Vector a -> (Int -> a -> m b) -> m ()
Vector.iforM_ Vector ArrayElement
arrayElements \Int
arrayIx Ktx1.ArrayElement{Vector Face
$sel:faces:ArrayElement :: ArrayElement -> Vector Face
faces :: Vector Face
faces} -> do
        forall (m :: * -> *) a b.
Monad m =>
Vector a -> (Int -> a -> m b) -> m ()
Vector.iforM_ Vector Face
faces \Int
faceIx Ktx1.Face{Vector ZSlice
$sel:zSlices:Face :: Face -> Vector ZSlice
zSlices :: Vector ZSlice
zSlices} -> do
          forall (m :: * -> *) a b.
Monad m =>
Vector a -> (Int -> a -> m b) -> m ()
Vector.iforM_ Vector ZSlice
zSlices \Int
sliceIx Ktx1.ZSlice{ByteString
$sel:block:ZSlice :: ZSlice -> ByteString
block :: ByteString
block} -> do
            let
              indices :: Utf8Builder
indices = forall a. Monoid a => [a] -> a
mconcat
                [ Utf8Builder
"["
                , Utf8Builder
" mip:" forall a. Semigroup a => a -> a -> a
<> forall a. Display a => a -> Utf8Builder
display Int
mipIx
                , Utf8Builder
" arr:" forall a. Semigroup a => a -> a -> a
<> forall a. Display a => a -> Utf8Builder
display Int
arrayIx
                , Utf8Builder
" fac:" forall a. Semigroup a => a -> a -> a
<> forall a. Display a => a -> Utf8Builder
display Int
faceIx
                , Utf8Builder
" slc:" forall a. Semigroup a => a -> a -> a
<> forall a. Display a => a -> Utf8Builder
display Int
sliceIx
                , Utf8Builder
" ]"
                ]
            let blockOffset :: Word32
blockOffset = Word32
offset forall a. Num a => a -> a -> a
+ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
faceIx forall a. Num a => a -> a -> a
* Word32
imageSize
            let sectionPtr :: Ptr Any
sectionPtr = forall a b. Ptr a -> Int -> Ptr b
Foreign.plusPtr (AllocationInfo -> Ptr ()
VMA.mappedData AllocationInfo
stageInfo) (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
blockOffset)
            forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) =>
Utf8Builder -> m ()
logDebug forall a b. (a -> b) -> a -> b
$ forall a. Monoid a => [a] -> a
mconcat
              [ Utf8Builder
indices
              , Utf8Builder
" base offset = "
              , forall a. Display a => a -> Utf8Builder
display Word32
offset
              , Utf8Builder
" image offset = "
              , forall a. Display a => a -> Utf8Builder
display forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
faceIx forall a. Num a => a -> a -> a
* Word32
imageSize
              , Utf8Builder
" image size = "
              , forall a. Display a => a -> Utf8Builder
display Word32
imageSize
              ]
            forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. ByteString -> (CStringLen -> IO a) -> IO a
unsafeUseAsCStringLen ByteString
block \(Ptr CChar
pixelsPtr, Int
pixelBytes) -> do
              if Int
pixelBytes forall a. Eq a => a -> a -> Bool
/= forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
imageSize then
                forall a. HasCallStack => String -> a
error String
"assert: MipLevel.imageSize matches block.pixelBytes"
              else
                -- traceShowM (sectionPtr, pixelBytes)
                forall a. Ptr a -> Ptr a -> Int -> IO ()
Foreign.copyBytes Ptr Any
sectionPtr (forall a b. Ptr a -> Ptr b
Foreign.castPtr Ptr CChar
pixelsPtr) Int
pixelBytes

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

    -- XXX: copying to image while the staging buffer is still alive
    forall context (t :: * -> *) deviceSize (m :: * -> *).
(HasVulkan context, Foldable t, Integral deviceSize,
 MonadUnliftIO m) =>
context
-> Queues CommandPool
-> Buffer
-> Image
-> Extent3D
-> ("mip offsets" ::: t deviceSize)
-> Word32
-> m ()
Image.copyBufferToImage
      env
context
      Queues CommandPool
pool
      Buffer
staging
      Image
image
      Extent3D
extent
      Vector Word32
offsets
      Word32
numLayers

  -- XXX: staging buffer is gone

  forall context (m :: * -> *).
(HasVulkan context, MonadUnliftIO m) =>
context
-> Queues CommandPool
-> Image
-> Word32
-> Word32
-> Format
-> ("old" ::: ImageLayout)
-> ("old" ::: ImageLayout)
-> m ()
Image.transitionLayout
    env
context
    Queues CommandPool
pool
    Image
image
    Word32
mipLevels
    Word32
numLayers -- XXX: arrayLayers is always 0 for now
    Format
format
    "old" ::: ImageLayout
Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL
    "old" ::: ImageLayout
Vk.IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL

  ImageView
imageView <- forall (io :: * -> *) context.
(MonadIO io, HasVulkan context) =>
context -> Image -> Format -> Word32 -> Word32 -> io ImageView
Texture.createImageView
    env
context
    Image
image
    Format
format
    Word32
mipLevels
    Word32
numLayers

  let
    extent2d :: Extent2D
extent2d =
      let
        Vk.Extent3D{Word32
depth :: Word32
height :: Word32
width :: Word32
$sel:depth:Extent3D :: Extent3D -> Word32
$sel:height:Extent3D :: Extent3D -> Word32
$sel:width:Extent3D :: Extent3D -> Word32
..} = Extent3D
extent
      in
        Vk.Extent2D{Word32
$sel:width:Extent2D :: Word32
$sel:height:Extent2D :: Word32
height :: Word32
width :: Word32
..}

    allocatedImage :: AllocatedImage
allocatedImage = AllocatedImage
      { $sel:aiAllocation:AllocatedImage :: Allocation
aiAllocation = Allocation
allocation
      , $sel:aiExtent:AllocatedImage :: Extent2D
aiExtent     = Extent2D
extent2d
      , $sel:aiFormat:AllocatedImage :: Format
aiFormat     = Format
format
      , $sel:aiImage:AllocatedImage :: Image
aiImage      = Image
image
      , $sel:aiImageView:AllocatedImage :: ImageView
aiImageView  = ImageView
imageView
      , $sel:aiImageRange:AllocatedImage :: ImageSubresourceRange
aiImageRange = ImageAspectFlags -> Word32 -> Word32 -> ImageSubresourceRange
Image.subresource ImageAspectFlags
Vk.IMAGE_ASPECT_COLOR_BIT Word32
mipLevels Word32
numLayers
      }
  pure Texture
    { $sel:tFormat:Texture :: Format
tFormat         = Format
format
    , $sel:tMipLevels:Texture :: Word32
tMipLevels      = Word32
mipLevels
    , $sel:tLayers:Texture :: Word32
tLayers         = Word32
numLayers
    , $sel:tAllocatedImage:Texture :: AllocatedImage
tAllocatedImage = AllocatedImage
allocatedImage
    }