Read.hs

{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}

{- | 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(..))
import Data.Typeable (Typeable, Proxy(..), typeRep)

-- * Context

-- | A bundle of source data and header information used by reader functions.
data Context a = Context
  { context :: a
  , 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 handle offset size = liftIO do
    IO.hSeek handle IO.AbsoluteSeek (fromIntegral offset)
    BS.hGet handle size

instance ReadLevel IO.Handle where
  readLevelTo handle Level{..} ptr = liftIO do
    IO.hSeek
      handle
      IO.AbsoluteSeek
      (fromIntegral byteOffset)
    got <- IO.hGetBuf handle ptr (fromIntegral byteLength)
    pure $ fromIntegral got == byteLength

instance Show (Context IO.Handle) where
  show (Context handle header) = mconcat
    [ "Context ("
    , show handle
    , ") "
    , show header
    ]

open :: MonadIO io => FilePath -> io FileContext
open path = do
  handle <- liftIO $ IO.openBinaryFile path IO.ReadMode
  header <- decodeAt handle 0 80 get
  pure $ Context handle header

close :: MonadIO io => FileContext -> io ()
close (Context handle _header) = liftIO $ IO.hClose 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 bs offset size =
    if offset + size > BS.length bs then
      liftIO . throwIO . ChunkError $ mconcat
        [ "Offset " <> fromString (show offset)
        , " and size " <> fromString (show size)
        , " is beyond the size of the buffer: "
        , fromString (show $ BS.length bs)
        ]
    else
      pure $ BS.take size (BS.drop offset bs)

instance ReadLevel ByteString where
  readLevelTo buf Level{..} dst =
    liftIO $ BSU.unsafeUseAsCStringLen buf \(src, size) ->
      if size < fromIntegral (byteOffset + byteLength) then
        pure False
      else do
        Foreign.copyBytes
          dst
          (plusPtr src $ fromIntegral byteOffset)
          (fromIntegral byteLength)
        pure True

instance Show (Context ByteString) where
  show (Context buf header) = mconcat
    [ "Context ["
    , show $ BS.length buf
    , "] "
    , show header
    ]

bytes :: MonadIO io => ByteString -> io BytesContext
bytes src = do
  header <- decodeAt src 0 80 get
  pure $ Context src header

-- * File contents

-- | Read the level index.
levels
  :: ( ReadChunk src
     , MonadIO io
     )
  => Context src
  -> io (Vector Level)
levels (Context handle Header{levelCount}) =
  decodeAt handle 80 (numLevels * 8 * 3) $
    Vector.replicateM numLevels get
  where
    numLevels = max 1 $ fromIntegral 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 (Context handle _header) = readLevelTo handle

-- | Copy level data into a managed buffer.
levelData
  :: ( ReadChunk src
     , MonadIO io
     )
  => Context src
  -> Level
  -> io ByteString
levelData (Context handle _header) Level{..} = do
  readChunkAt
    handle
    (fromIntegral byteOffset)
    (fromIntegral 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 (Context handle Header{..}) =
  if dfdByteLength == 0 then
    pure DFD
      { dfdTotalSize = 0
      , dfdBlocks = mempty
      }
  else
    decodeAt
      handle
      (fromIntegral dfdByteOffset)
      (fromIntegral dfdByteLength)
      get

{- | Read and parse Key-Value Data block.
-}
keyValueData
  :: ( ReadChunk src
     , MonadIO io
     )
  => Context src
  -> io KeyValueData
keyValueData (Context handle Header{..}) =
  if kvdByteLength == 0 then
    pure mempty
  else
    decodeAt
      handle
      (fromIntegral kvdByteOffset)
      (fromIntegral kvdByteLength)
      (KeyValue.getDataLe $ fromIntegral kvdByteLength)

-- | Get a copy of global supercompression data.
supercompressionGlobalData
  :: ( ReadChunk src
     , MonadIO io
     )
  => Context src
  -> io ByteString
supercompressionGlobalData (Context handle Header{..}) =
  decodeAt
    handle
    (fromIntegral sgdByteOffset)
    (fromIntegral sgdByteLength)
    (getByteString $ fromIntegral 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 (Eq, Show)

instance Exception ChunkError

-- | Get a chunk of data and run a decoder on it.
decodeAt
  :: forall a src io
  .  ( ReadChunk src
     , Show a
     , Typeable a
     , MonadIO io
     )
  => src
  -> Int
  -> Int
  -> Get a
  -> io a
decodeAt src offset size action = do
  chunk <- readChunkAt src offset size
  case runGetOrFail action (BSL.fromChunks [chunk]) of
    Right ("", used, ok) | fromIntegral used == size ->
      pure ok
    Right (remains, finalOffset, _okBut) ->
      liftIO . throwIO $ DecodeError
        (fromIntegral offset + finalOffset)
        ( "BUG: unused data " <>
          fromString (show remains)
        )
    Left (_remains, errorOffset, message) ->
      liftIO . throwIO $ DecodeError
        (fromIntegral offset + errorOffset)
        ( fromString $ unwords
            [ typeName
            , "-"
            , message
            ]
        )
  where
    typeName = show $ typeRep (Proxy :: Proxy a)

data DecodeError = DecodeError ByteOffset Text
  deriving (Eq, Show)

instance Exception DecodeError

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