{-# LANGUAGE FlexibleInstances #-}

{- | Block-by-block extraction of data from a KTX2 file.

* Acquire a "Context". Header data is available without reading the rest of the source.
* Read 'levels' index. An memory allocation information is available.
* Consult 'Codec.Ktx2.Header.supercompressionScheme' and copy level data to decompression staging buffer or GPU memory directly.

Extra information is available when needed:

* Image metadata.
* Data Format Descriptor. Khronos Basic descriptor block is usually present, but a file may contain more.
* Supercompression data shared between all the levels.
-}

module Codec.Ktx2.Read
  ( Context(..)

  , FileContext
  , open
  , close

  , BytesContext
  , bytes

  -- * Reading blocks

  -- ** Image data
  , levels
  , levelToPtr
  , levelData

  -- ** Supplemental information
  , dataFormatDescriptor
  , keyValueData
  , supercompressionGlobalData

  -- * Decoding internals
  , ReadChunk(..)
  , ChunkError(..)
  , decodeAt
  , DecodeError(..)
  , ReadLevel(..)
  ) where

import Control.Exception (Exception, throwIO)
import Control.Monad.IO.Class (MonadIO(..))
import Data.Binary (Binary(..))
import Data.Binary.Get (Get, ByteOffset, getByteString, runGetOrFail)
import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.ByteString.Lazy qualified as BSL
import Data.ByteString.Unsafe qualified as BSU
import Data.String (fromString)
import Data.Text (Text)
import Data.Vector (Vector)
import Data.Vector qualified as Vector
import Foreign (Ptr, plusPtr)
import Foreign qualified
import System.IO qualified as IO

import Codec.Ktx.KeyValue (KeyValueData)
import Codec.Ktx.KeyValue qualified as KeyValue
import Codec.Ktx2.Header (Header(..))
import Codec.Ktx2.Level (Level(..))
import Codec.Ktx2.DFD (DFD)

-- * Context

-- | A bundle of source data and header information used by reader functions.
data Context a = Context
  { forall a. Context a -> a
context :: a
  , forall a. Context a -> Header
header :: Header
  }

-- ** Reading from files

-- | Context for reading from a file. The file has to be seekable.
type FileContext = Context IO.Handle

instance ReadChunk IO.Handle where
  readChunkAt :: forall (io :: * -> *).
MonadIO io =>
Handle -> Int -> Int -> io ByteString
readChunkAt Handle
handle Int
offset Int
size = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
    Handle -> SeekMode -> Integer -> IO ()
IO.hSeek Handle
handle SeekMode
IO.AbsoluteSeek (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
offset)
    Handle -> Int -> IO ByteString
BS.hGet Handle
handle Int
size

instance ReadLevel IO.Handle where
  readLevelTo :: forall (io :: * -> *).
MonadIO io =>
Handle -> Level -> Ptr () -> io Bool
readLevelTo Handle
handle Level{Word64
uncompressedByteLength :: Level -> Word64
byteLength :: Level -> Word64
byteOffset :: Level -> Word64
uncompressedByteLength :: Word64
byteLength :: Word64
byteOffset :: Word64
..} Ptr ()
ptr = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
    Handle -> SeekMode -> Integer -> IO ()
IO.hSeek
      Handle
handle
      SeekMode
IO.AbsoluteSeek
      (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
byteOffset)
    Int
got <- forall a. Handle -> Ptr a -> Int -> IO Int
IO.hGetBuf Handle
handle Ptr ()
ptr (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
byteLength)
    pure $ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
got forall a. Eq a => a -> a -> Bool
== Word64
byteLength

instance Show (Context IO.Handle) where
  show :: Context Handle -> String
show (Context Handle
handle Header
header) = forall a. Monoid a => [a] -> a
mconcat
    [ String
"Context ("
    , forall a. Show a => a -> String
show Handle
handle
    , String
") "
    , forall a. Show a => a -> String
show Header
header
    ]

open :: MonadIO io => FilePath -> io FileContext
open :: forall (io :: * -> *). MonadIO io => String -> io (Context Handle)
open String
path = do
  Handle
handle <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ String -> IOMode -> IO Handle
IO.openBinaryFile String
path IOMode
IO.ReadMode
  Header
header <- forall src a (io :: * -> *).
(ReadChunk src, Show a, MonadIO io) =>
src -> Int -> Int -> Get a -> io a
decodeAt Handle
handle Int
0 Int
80 forall t. Binary t => Get t
get
  pure $ forall a. a -> Header -> Context a
Context Handle
handle Header
header

close :: MonadIO io => FileContext -> io ()
close :: forall (io :: * -> *). MonadIO io => Context Handle -> io ()
close (Context Handle
handle Header
_header) = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ Handle -> IO ()
IO.hClose Handle
handle

-- ** Reading from memory

-- | Context for reading from memory. Useful when the data is embedded in a module or otherwise already available in full.
type BytesContext = Context ByteString

instance ReadChunk ByteString where
  readChunkAt :: forall (io :: * -> *).
MonadIO io =>
ByteString -> Int -> Int -> io ByteString
readChunkAt ByteString
bs Int
offset Int
size =
    if Int
offset forall a. Num a => a -> a -> a
+ Int
size forall a. Ord a => a -> a -> Bool
> ByteString -> Int
BS.length ByteString
bs then
      forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall e a. Exception e => e -> IO a
throwIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ChunkError
ChunkError forall a b. (a -> b) -> a -> b
$ forall a. Monoid a => [a] -> a
mconcat
        [ Text
"Offset " forall a. Semigroup a => a -> a -> a
<> forall a. IsString a => String -> a
fromString (forall a. Show a => a -> String
show Int
offset)
        , Text
" and size " forall a. Semigroup a => a -> a -> a
<> forall a. IsString a => String -> a
fromString (forall a. Show a => a -> String
show Int
size)
        , Text
" is beyond the size of the buffer: "
        , forall a. IsString a => String -> a
fromString (forall a. Show a => a -> String
show forall a b. (a -> b) -> a -> b
$ ByteString -> Int
BS.length ByteString
bs)
        ]
    else
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Int -> ByteString -> ByteString
BS.take Int
size (Int -> ByteString -> ByteString
BS.drop Int
offset ByteString
bs)

instance ReadLevel ByteString where
  readLevelTo :: forall (io :: * -> *).
MonadIO io =>
ByteString -> Level -> Ptr () -> io Bool
readLevelTo ByteString
buf Level{Word64
uncompressedByteLength :: Word64
byteLength :: Word64
byteOffset :: Word64
uncompressedByteLength :: Level -> Word64
byteLength :: Level -> Word64
byteOffset :: Level -> Word64
..} Ptr ()
dst =
    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
BSU.unsafeUseAsCStringLen ByteString
buf \(Ptr CChar
src, Int
size) ->
      if Int
size forall a. Ord a => a -> a -> Bool
< forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64
byteOffset forall a. Num a => a -> a -> a
+ Word64
byteLength) then
        forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
      else do
        forall a. Ptr a -> Ptr a -> Int -> IO ()
Foreign.copyBytes
          Ptr ()
dst
          (forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr CChar
src forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
byteOffset)
          (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
byteLength)
        pure Bool
True

instance Show (Context ByteString) where
  show :: Context ByteString -> String
show (Context ByteString
buf Header
header) = forall a. Monoid a => [a] -> a
mconcat
    [ String
"Context ["
    , forall a. Show a => a -> String
show forall a b. (a -> b) -> a -> b
$ ByteString -> Int
BS.length ByteString
buf
    , String
"] "
    , forall a. Show a => a -> String
show Header
header
    ]

bytes :: MonadIO io => ByteString -> io BytesContext
bytes :: forall (io :: * -> *).
MonadIO io =>
ByteString -> io (Context ByteString)
bytes ByteString
src = do
  Header
header <- forall src a (io :: * -> *).
(ReadChunk src, Show a, MonadIO io) =>
src -> Int -> Int -> Get a -> io a
decodeAt ByteString
src Int
0 Int
80 forall t. Binary t => Get t
get
  pure $ forall a. a -> Header -> Context a
Context ByteString
src Header
header

-- * File contents

-- | Read the level index.
levels
  :: ( ReadChunk src
     , MonadIO io
     )
  => Context src
  -> io (Vector Level)
levels :: forall src (io :: * -> *).
(ReadChunk src, MonadIO io) =>
Context src -> io (Vector Level)
levels (Context src
handle Header{Word32
levelCount :: Header -> Word32
levelCount :: Word32
levelCount}) =
  forall src a (io :: * -> *).
(ReadChunk src, Show a, MonadIO io) =>
src -> Int -> Int -> Get a -> io a
decodeAt src
handle Int
80 (Int
numLevels forall a. Num a => a -> a -> a
* Int
8 forall a. Num a => a -> a -> a
* Int
3) forall a b. (a -> b) -> a -> b
$
    forall (m :: * -> *) a. Monad m => Int -> m a -> m (Vector a)
Vector.replicateM Int
numLevels forall t. Binary t => Get t
get
  where
    numLevels :: Int
numLevels = forall a. Ord a => a -> a -> a
max Int
1 forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
levelCount

{- | Copy level data into a provided pointer.

The buffer must be large enough for the 'byteLength' of the "Level" being accessed.
-}
{-# INLINE levelToPtr #-}
levelToPtr
  :: ( ReadLevel src
     , MonadIO io
     )
  => Context src
  -> Level
  -> Ptr ()
  -> io Bool
levelToPtr :: forall src (io :: * -> *).
(ReadLevel src, MonadIO io) =>
Context src -> Level -> Ptr () -> io Bool
levelToPtr (Context src
handle Header
_header) = forall a (io :: * -> *).
(ReadLevel a, MonadIO io) =>
a -> Level -> Ptr () -> io Bool
readLevelTo src
handle

-- | Copy level data into a managed buffer.
levelData
  :: ( ReadChunk src
     , MonadIO io
     )
  => Context src
  -> Level
  -> io ByteString
levelData :: forall src (io :: * -> *).
(ReadChunk src, MonadIO io) =>
Context src -> Level -> io ByteString
levelData (Context src
handle Header
_header) Level{Word64
uncompressedByteLength :: Word64
byteLength :: Word64
byteOffset :: Word64
uncompressedByteLength :: Level -> Word64
byteLength :: Level -> Word64
byteOffset :: Level -> Word64
..} = do
  forall a (io :: * -> *).
(ReadChunk a, MonadIO io) =>
a -> Int -> Int -> io ByteString
readChunkAt
    src
handle
    (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
byteOffset)
    (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
byteLength)

{- | Read DFD block data.

Further processing is performed according to descriptor vendor/type/version.
E.g. "Codec.Ktx2.DFD.Khronos.BasicV2".
-}
dataFormatDescriptor
  :: ( ReadChunk src
     , MonadIO io
     )
  => Context src
  -> io DFD
dataFormatDescriptor :: forall src (io :: * -> *).
(ReadChunk src, MonadIO io) =>
Context src -> io DFD
dataFormatDescriptor (Context src
handle Header{Word32
Word64
sgdByteLength :: Header -> Word64
sgdByteOffset :: Header -> Word64
kvdByteLength :: Header -> Word32
kvdByteOffset :: Header -> Word32
dfdByteLength :: Header -> Word32
dfdByteOffset :: Header -> Word32
supercompressionScheme :: Header -> Word32
faceCount :: Header -> Word32
layerCount :: Header -> Word32
pixelDepth :: Header -> Word32
pixelHeight :: Header -> Word32
pixelWidth :: Header -> Word32
typeSize :: Header -> Word32
vkFormat :: Header -> Word32
sgdByteLength :: Word64
sgdByteOffset :: Word64
kvdByteLength :: Word32
kvdByteOffset :: Word32
dfdByteLength :: Word32
dfdByteOffset :: Word32
supercompressionScheme :: Word32
levelCount :: Word32
faceCount :: Word32
layerCount :: Word32
pixelDepth :: Word32
pixelHeight :: Word32
pixelWidth :: Word32
typeSize :: Word32
vkFormat :: Word32
levelCount :: Header -> Word32
..}) =
  forall src a (io :: * -> *).
(ReadChunk src, Show a, MonadIO io) =>
src -> Int -> Int -> Get a -> io a
decodeAt
    src
handle
    (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
dfdByteOffset)
    (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
dfdByteLength)
    forall t. Binary t => Get t
get

{- | Read and parse Key-Value Data block.
-}
keyValueData
  :: ( ReadChunk src
     , MonadIO io
     )
  => Context src
  -> io KeyValueData
keyValueData :: forall src (io :: * -> *).
(ReadChunk src, MonadIO io) =>
Context src -> io KeyValueData
keyValueData (Context src
handle Header{Word32
Word64
sgdByteLength :: Word64
sgdByteOffset :: Word64
kvdByteLength :: Word32
kvdByteOffset :: Word32
dfdByteLength :: Word32
dfdByteOffset :: Word32
supercompressionScheme :: Word32
levelCount :: Word32
faceCount :: Word32
layerCount :: Word32
pixelDepth :: Word32
pixelHeight :: Word32
pixelWidth :: Word32
typeSize :: Word32
vkFormat :: Word32
sgdByteLength :: Header -> Word64
sgdByteOffset :: Header -> Word64
kvdByteLength :: Header -> Word32
kvdByteOffset :: Header -> Word32
dfdByteLength :: Header -> Word32
dfdByteOffset :: Header -> Word32
supercompressionScheme :: Header -> Word32
faceCount :: Header -> Word32
layerCount :: Header -> Word32
pixelDepth :: Header -> Word32
pixelHeight :: Header -> Word32
pixelWidth :: Header -> Word32
typeSize :: Header -> Word32
vkFormat :: Header -> Word32
levelCount :: Header -> Word32
..}) =
  forall src a (io :: * -> *).
(ReadChunk src, Show a, MonadIO io) =>
src -> Int -> Int -> Get a -> io a
decodeAt
    src
handle
    (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
kvdByteOffset)
    (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
kvdByteLength)
    (Int -> Get KeyValueData
KeyValue.getDataLe forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
kvdByteLength)

-- | Get a copy of global supercompression data.
supercompressionGlobalData
  :: ( ReadChunk src
     , MonadIO io
     )
  => Context src
  -> io ByteString
supercompressionGlobalData :: forall src (io :: * -> *).
(ReadChunk src, MonadIO io) =>
Context src -> io ByteString
supercompressionGlobalData (Context src
handle Header{Word32
Word64
sgdByteLength :: Word64
sgdByteOffset :: Word64
kvdByteLength :: Word32
kvdByteOffset :: Word32
dfdByteLength :: Word32
dfdByteOffset :: Word32
supercompressionScheme :: Word32
levelCount :: Word32
faceCount :: Word32
layerCount :: Word32
pixelDepth :: Word32
pixelHeight :: Word32
pixelWidth :: Word32
typeSize :: Word32
vkFormat :: Word32
sgdByteLength :: Header -> Word64
sgdByteOffset :: Header -> Word64
kvdByteLength :: Header -> Word32
kvdByteOffset :: Header -> Word32
dfdByteLength :: Header -> Word32
dfdByteOffset :: Header -> Word32
supercompressionScheme :: Header -> Word32
faceCount :: Header -> Word32
layerCount :: Header -> Word32
pixelDepth :: Header -> Word32
pixelHeight :: Header -> Word32
pixelWidth :: Header -> Word32
typeSize :: Header -> Word32
vkFormat :: Header -> Word32
levelCount :: Header -> Word32
..}) =
  forall src a (io :: * -> *).
(ReadChunk src, Show a, MonadIO io) =>
src -> Int -> Int -> Get a -> io a
decodeAt
    src
handle
    (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
sgdByteOffset)
    (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
sgdByteLength)
    (Int -> Get ByteString
getByteString forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
sgdByteLength)

-- * IO helpers

class ReadChunk a where
  {- | Get a chunk of data.

  The context handle must have enough information to check whether requested region is safe to access.
  Throw "ChunkError" when it isn't possible to fullfill the request.
  -}
  readChunkAt :: MonadIO io => a -> Int -> Int -> io ByteString

newtype ChunkError = ChunkError Text
  deriving (ChunkError -> ChunkError -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ChunkError -> ChunkError -> Bool
$c/= :: ChunkError -> ChunkError -> Bool
== :: ChunkError -> ChunkError -> Bool
$c== :: ChunkError -> ChunkError -> Bool
Eq, Int -> ChunkError -> ShowS
[ChunkError] -> ShowS
ChunkError -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ChunkError] -> ShowS
$cshowList :: [ChunkError] -> ShowS
show :: ChunkError -> String
$cshow :: ChunkError -> String
showsPrec :: Int -> ChunkError -> ShowS
$cshowsPrec :: Int -> ChunkError -> ShowS
Show)

instance Exception ChunkError

-- | Get a chunk of data and run a decoder on it.
decodeAt
  :: ( ReadChunk src
     , Show a
     , MonadIO io
     )
  => src
  -> Int
  -> Int
  -> Get a
  -> io a
decodeAt :: forall src a (io :: * -> *).
(ReadChunk src, Show a, MonadIO io) =>
src -> Int -> Int -> Get a -> io a
decodeAt src
src Int
offset Int
size Get a
action = do
  ByteString
chunk <- forall a (io :: * -> *).
(ReadChunk a, MonadIO io) =>
a -> Int -> Int -> io ByteString
readChunkAt src
src Int
offset Int
size
  case forall a.
Get a
-> ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, a)
runGetOrFail Get a
action ([ByteString] -> ByteString
BSL.fromChunks [ByteString
chunk]) of
    Right (ByteString
"", ByteOffset
used, a
ok) | forall a b. (Integral a, Num b) => a -> b
fromIntegral ByteOffset
used forall a. Eq a => a -> a -> Bool
== Int
size ->
      forall (f :: * -> *) a. Applicative f => a -> f a
pure a
ok
    Right (ByteString
remains, ByteOffset
finalOffset, a
_okBut) ->
      forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall e a. Exception e => e -> IO a
throwIO forall a b. (a -> b) -> a -> b
$ ByteOffset -> Text -> DecodeError
DecodeError ByteOffset
finalOffset forall a b. (a -> b) -> a -> b
$
        Text
"BUG: unused data" forall a. Semigroup a => a -> a -> a
<> forall a. IsString a => String -> a
fromString (forall a. Show a => a -> String
show ByteString
remains)
    Left (ByteString
_remains, ByteOffset
errorOffset, String
message) ->
      forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall e a. Exception e => e -> IO a
throwIO forall a b. (a -> b) -> a -> b
$ ByteOffset -> Text -> DecodeError
DecodeError ByteOffset
errorOffset (forall a. IsString a => String -> a
fromString String
message)

data DecodeError = DecodeError ByteOffset Text
  deriving (DecodeError -> DecodeError -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: DecodeError -> DecodeError -> Bool
$c/= :: DecodeError -> DecodeError -> Bool
== :: DecodeError -> DecodeError -> Bool
$c== :: DecodeError -> DecodeError -> Bool
Eq, Int -> DecodeError -> ShowS
[DecodeError] -> ShowS
DecodeError -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [DecodeError] -> ShowS
$cshowList :: [DecodeError] -> ShowS
show :: DecodeError -> String
$cshow :: DecodeError -> String
showsPrec :: Int -> DecodeError -> ShowS
$cshowsPrec :: Int -> DecodeError -> ShowS
Show)

instance Exception DecodeError

class ReadLevel a where
  readLevelTo :: MonadIO io => a -> Level -> Ptr () -> io Bool