ktx-codec-0.0.2.0: test/Spec.hs

module Main where

import Codec.Compression.Zstd.FFI qualified as Zstd
import Control.Exception (bracket)
import Control.Monad (guard, unless)
import Data.ByteString qualified as BS
import Data.ByteString.Builder (toLazyByteString)
import Data.ByteString.Lazy qualified as BSL
import Data.List (isSuffixOf)
import Data.Vector qualified as Vector
import Debug.Trace (traceM)
import Foreign qualified
import Shower qualified
import System.Directory (listDirectory)
import System.Environment (lookupEnv)
import System.FilePath ((</>))
import Test.Tasty
import Test.Tasty.HUnit

import Codec.Ktx qualified as Ktx1
import Codec.Ktx.KeyValue qualified as KeyValue
import Codec.Ktx2.DFD (DFD(..))
import Codec.Ktx2.DFD.Khronos.BasicV2 qualified as BasicV2
import Codec.Ktx2.Header qualified as Header
import Codec.Ktx2.Level qualified as Level
import Codec.Ktx2.Read qualified as Read

assetsPath :: FilePath
assetsPath = ".." </> "assets"

main :: IO ()
main = do
  assets <- listDirectory assetsPath
  defaultMain $
    testGroup "ktx-codec"
      [ testGroup "KTX 1" do
          fp <- assets
          guard $ ".ktx" `isSuffixOf` fp
          pure . testCase fp $
            testKtx1 (assetsPath </> fp)
      , testGroup "KTX 2 (files)" do
          fp <- assets
          guard $ ".ktx2" `isSuffixOf` fp
          pure . testCase fp $
            testKtx2File (assetsPath </> fp)
      , testGroup "KTX 2 (bytes)" do
          fp <- assets
          guard $ ".ktx2" `isSuffixOf` fp
          pure . testCase fp $
            testKtx2Bytes (assetsPath </> fp)
      ]

testKtx1 :: FilePath -> IO ()
testKtx1 source =
  Ktx1.fromFile source >>= \case
    Left (offset, err) -> do
      traceM err
      fail $ unwords
        [ "KTX-1 load error in"
        , source
        , "at"
        , show offset
        ]
    Right ktx -> do
      printer ktx
      let encodedBytes = toLazyByteString (Ktx1.toBuilder ktx)
      originalBytes <- BSL.readFile source
      unless (encodedBytes == originalBytes) $
        fail $ "Encoded KTX does not match original for " <> source

testKtx2File :: FilePath -> IO ()
testKtx2File source =
  bracket (Read.open source) Read.close $
    testKtx2Context source

testKtx2Bytes :: FilePath -> IO ()
testKtx2Bytes source = do
  contents <- BS.readFile source
  Read.bytes contents >>= testKtx2Context source

testKtx2Context
  :: ( Read.ReadChunk a
     , Read.ReadLevel a
     , Show (Read.Context a)
     )
  => FilePath
  -> Read.Context a
  -> IO ()
testKtx2Context source ctx = do
  printer ctx

  levels <- Read.levels ctx
  printer levels

  DFD{dfdBlocks} <- Read.dataFormatDescriptor ctx
  block0 <- Vector.headM dfdBlocks
  case BasicV2.fromBlock block0 of
    Nothing -> do
      printer block0
      fail $ "Unexpected block type in test suite in " <> source
    Just basic -> do
      -- printer basic
      let back = BasicV2.toBlock basic
      unless (back == block0) do
        printer (back, block0)
        fail "BasicV2 tripping failed"

  fmap KeyValue.textual (Read.keyValueData ctx) >>= Shower.printer
  Read.supercompressionGlobalData ctx >>= printer
  Vector.forM_ levels \level -> do
    bytes <- Read.levelData ctx level
    unless (BS.length bytes == fromIntegral (Level.byteLength level)) $
      fail "Uncompressed level data does not match with index!"

    Foreign.allocaBytesAligned (fromIntegral $ Level.uncompressedByteLength level) 16 \dst -> do
      case Header.supercompressionScheme (Read.header ctx) of
        Header.SC_NONE -> do
          -- XXX: Write to destination directly
          ok <- Read.levelToPtr ctx level dst
          unless ok $
            fail $ "Data read failed for " <> show level

        Header.SC_ZSTANDARD ->
          -- XXX: Prepare an intermediate buffer for decompression
          Foreign.allocaBytesAligned (fromIntegral $ Level.byteLength level) 16 \stage -> do
            ok <- Read.levelToPtr ctx level stage
            unless ok $
              fail $ "Data read failed for " <> show level

            let expected = fromIntegral $ Level.uncompressedByteLength level
            res <-
              Zstd.checkError $
                Zstd.decompress
                  dst
                  expected
                  stage
                  (fromIntegral $ Level.byteLength level)
            case res of
              Right decoded | decoded == expected ->
                printer (decoded, level)
              Right mismatch ->
                fail $ "Zstd decoded unexpected number of bytes (" <> show mismatch <> ") for level " <> show level
              Left err ->
                fail $ "Zstd error: " <> err

        huh ->
          fail $ "Unexpected supercompression: " <> show huh

printer :: Show a => a -> IO ()
printer x =
  lookupEnv "TEST_DUMP" >>= \case
    Nothing ->
      pure ()
    Just _ ->
      Shower.printer x