module Resource.Mesh.Codec where

import RIO

import Codec.Compression.Zstd qualified as Zstd
import Codec.Serialise qualified as CBOR
import Crypto.Hash.MD5 qualified as MD5
import Data.Binary.Get (runGet)
import Data.Binary.Get qualified as Get
import Data.Binary.Put (runPut)
import Data.Binary.Put qualified as Put
import Data.ByteString.Internal qualified as BSI
import Data.ByteString.Unsafe (unsafePackCStringLen)
import Data.Typeable (typeRep, typeRepTyCon)
import Data.Vector qualified as Vector
import Data.Vector.Generic qualified as Generic
import Data.Vector.Storable qualified as Storable
import Foreign qualified
import Geomancy.Vec3 qualified as Vec3
import RIO.ByteString qualified as ByteString
import RIO.ByteString.Lazy qualified as BSL
import System.IO.Unsafe (unsafePerformIO)
import UnliftIO.Resource (MonadResource)
import UnliftIO.Resource qualified as Resource
import Vulkan.Core10 qualified as Vk

import Engine.Vulkan.Types (MonadVulkan, Queues)
import Resource.Buffer qualified as Buffer
import Resource.Model qualified as Model

-- * Format meta

pattern VER_BREAKS :: Word8
pattern $bVER_BREAKS :: Word8
$mVER_BREAKS :: forall {r}. Word8 -> (Void# -> r) -> (Void# -> r) -> r
VER_BREAKS = 3

pattern VER_TWEAKS :: Word8
pattern $bVER_TWEAKS :: Word8
$mVER_TWEAKS :: forall {r}. Word8 -> (Void# -> r) -> (Void# -> r) -> r
VER_TWEAKS = 0

-- * Encoding

encodeFile
  :: forall vp vi va vn attrs nodes meta env
  .  ( Generic.Vector vp Vec3.Packed
     , Generic.Vector vi Word32
     , Generic.Vector va attrs
     , Generic.Vector vn nodes
     , Storable attrs
     , Storable nodes
     , CBOR.Serialise meta
     , HasLogFunc env
     )
  => FilePath
  -> vp Vec3.Packed
  -> vi Word32
  -> va attrs
  -> vn nodes
  -> meta
  -> RIO env ()
encodeFile :: forall (vp :: * -> *) (vi :: * -> *) (va :: * -> *) (vn :: * -> *)
       attrs nodes meta env.
(Vector vp Packed, Vector vi Word32, Vector va attrs,
 Vector vn nodes, Storable attrs, Storable nodes, Serialise meta,
 HasLogFunc env) =>
FilePath
-> vp Packed
-> vi Word32
-> va attrs
-> vn nodes
-> meta
-> RIO env ()
encodeFile FilePath
fp vp Packed
positions vi Word32
indices va attrs
attrs vn nodes
nodes meta
meta = do
  (ByteString
posDigest, ByteString
posCompressed) <- vp Packed -> RIO env (ByteString, ByteString)
forall a (v :: * -> *) (m :: * -> *).
(Storable a, Vector v a, MonadIO m) =>
v a -> m (ByteString, ByteString)
encodeItems vp Packed
positions
  Utf8Builder -> RIO env ()
forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) =>
Utf8Builder -> m ()
logDebug (Utf8Builder -> RIO env ()) -> Utf8Builder -> RIO env ()
forall a b. (a -> b) -> a -> b
$ Utf8Builder
"Position digest: " Utf8Builder -> Utf8Builder -> Utf8Builder
forall a. Semigroup a => a -> a -> a
<> ByteString -> Utf8Builder
forall a. Show a => a -> Utf8Builder
displayShow ByteString
posDigest

  (ByteString
indDigest, ByteString
indCompressed) <- vi Word32 -> RIO env (ByteString, ByteString)
forall a (v :: * -> *) (m :: * -> *).
(Storable a, Vector v a, MonadIO m) =>
v a -> m (ByteString, ByteString)
encodeItems vi Word32
indices
  Utf8Builder -> RIO env ()
forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) =>
Utf8Builder -> m ()
logDebug (Utf8Builder -> RIO env ()) -> Utf8Builder -> RIO env ()
forall a b. (a -> b) -> a -> b
$ Utf8Builder
"Index digest: " Utf8Builder -> Utf8Builder -> Utf8Builder
forall a. Semigroup a => a -> a -> a
<> ByteString -> Utf8Builder
forall a. Show a => a -> Utf8Builder
displayShow ByteString
indDigest

  (ByteString
attDigest, ByteString
attCompressed) <- va attrs -> RIO env (ByteString, ByteString)
forall a (v :: * -> *) (m :: * -> *).
(Storable a, Vector v a, MonadIO m) =>
v a -> m (ByteString, ByteString)
encodeItems va attrs
attrs
  Utf8Builder -> RIO env ()
forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) =>
Utf8Builder -> m ()
logDebug (Utf8Builder -> RIO env ()) -> Utf8Builder -> RIO env ()
forall a b. (a -> b) -> a -> b
$ Utf8Builder
"Attribute digest: " Utf8Builder -> Utf8Builder -> Utf8Builder
forall a. Semigroup a => a -> a -> a
<> ByteString -> Utf8Builder
forall a. Show a => a -> Utf8Builder
displayShow ByteString
attDigest

  (ByteString
nodDigest, ByteString
nodCompressed) <- vn nodes -> RIO env (ByteString, ByteString)
forall a (v :: * -> *) (m :: * -> *).
(Storable a, Vector v a, MonadIO m) =>
v a -> m (ByteString, ByteString)
encodeItems vn nodes
nodes
  Utf8Builder -> RIO env ()
forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) =>
Utf8Builder -> m ()
logDebug (Utf8Builder -> RIO env ()) -> Utf8Builder -> RIO env ()
forall a b. (a -> b) -> a -> b
$ Utf8Builder
"Node digest: " Utf8Builder -> Utf8Builder -> Utf8Builder
forall a. Semigroup a => a -> a -> a
<> ByteString -> Utf8Builder
forall a. Show a => a -> Utf8Builder
displayShow ByteString
nodDigest

  let (Int
metSize, ByteString
metDigest, ByteString
metCompressed) = meta -> (Int, ByteString, ByteString)
forall a. Serialise a => a -> (Int, ByteString, ByteString)
encodeCBOR meta
meta
  Utf8Builder -> RIO env ()
forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) =>
Utf8Builder -> m ()
logDebug (Utf8Builder -> RIO env ()) -> Utf8Builder -> RIO env ()
forall a b. (a -> b) -> a -> b
$ Utf8Builder
"Meta digest: " Utf8Builder -> Utf8Builder -> Utf8Builder
forall a. Semigroup a => a -> a -> a
<> ByteString -> Utf8Builder
forall a. Show a => a -> Utf8Builder
displayShow ByteString
metDigest

  FilePath -> IOMode -> (Handle -> RIO env ()) -> RIO env ()
forall (m :: * -> *) a.
MonadUnliftIO m =>
FilePath -> IOMode -> (Handle -> m a) -> m a
withFile FilePath
fp IOMode
WriteMode \Handle
out -> do
    Handle -> LByteString -> RIO env ()
forall (m :: * -> *). MonadIO m => Handle -> LByteString -> m ()
BSL.hPut Handle
out (LByteString -> RIO env ()) -> LByteString -> RIO env ()
forall a b. (a -> b) -> a -> b
$ Put -> LByteString
runPut do
      -- 0x00 + 4 + 4
      FilePath -> Put
Put.putStringUtf8 FilePath
"🌋📦"

      -- 0x08 + 1 + 1
      Word8 -> Put
Put.putWord8 Word8
VER_BREAKS
      Word8 -> Put
Put.putWord8 Word8
VER_TWEAKS

      -- 0x0A + 2 + 4
      -- XXX: reserved
      Word16 -> Put
Put.putWord16le Word16
forall a. Bounded a => a
maxBound
      Word32 -> Put
Put.putWord32le Word32
forall a. Bounded a => a
maxBound

      -- 0x10 + 4 + 4
      Word32 -> Put
Put.putWord32le (Word32 -> Put) -> (Int -> Word32) -> Int -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Put) -> Int -> Put
forall a b. (a -> b) -> a -> b
$ vp Packed -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
Generic.length vp Packed
positions
      Word32 -> Put
Put.putWord32le (Word32 -> Put) -> (Int -> Word32) -> Int -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Put) -> Int -> Put
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
ByteString.length ByteString
posCompressed

      -- 0x18 + 4 + 4
      Word32 -> Put
Put.putWord32le (Word32 -> Put) -> (Int -> Word32) -> Int -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Put) -> Int -> Put
forall a b. (a -> b) -> a -> b
$ vi Word32 -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
Generic.length vi Word32
indices
      Word32 -> Put
Put.putWord32le (Word32 -> Put) -> (Int -> Word32) -> Int -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Put) -> Int -> Put
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
ByteString.length ByteString
indCompressed

      -- 0x20 + 4 + 4
      Word32 -> Put
Put.putWord32le (Word32 -> Put) -> (Int -> Word32) -> Int -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Put) -> Int -> Put
forall a b. (a -> b) -> a -> b
$ attrs -> Int
forall a. Storable a => a -> Int
Foreign.sizeOf (FilePath -> attrs
forall a. HasCallStack => FilePath -> a
error FilePath
"sizeOf" :: attrs)
      Word32 -> Put
Put.putWord32le (Word32 -> Put) -> (Int -> Word32) -> Int -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Put) -> Int -> Put
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
ByteString.length ByteString
attCompressed

      -- 0x28 + 4 + 4
      Word32 -> Put
Put.putWord32le (Word32 -> Put) -> (Int -> Word32) -> Int -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Put) -> Int -> Put
forall a b. (a -> b) -> a -> b
$ nodes -> Int
forall a. Storable a => a -> Int
Foreign.sizeOf (FilePath -> nodes
forall a. HasCallStack => FilePath -> a
error FilePath
"sizeOf" :: nodes)
      Word32 -> Put
Put.putWord32le (Word32 -> Put) -> (Int -> Word32) -> Int -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Put) -> Int -> Put
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
ByteString.length ByteString
nodCompressed

      -- 0x30 + 4 + 4
      Word32 -> Put
Put.putWord32le (Word32 -> Put) -> Word32 -> Put
forall a b. (a -> b) -> a -> b
$ Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
metSize
      Word32 -> Put
Put.putWord32le (Word32 -> Put) -> (Int -> Word32) -> Int -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Put) -> Int -> Put
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
ByteString.length ByteString
metCompressed

      -- 0x38 + 4 + 4
      -- XXX: reserved
      Word32 -> Put
Put.putWord32le Word32
forall a. Bounded a => a
maxBound
      Word32 -> Put
Put.putWord32le Word32
forall a. Bounded a => a
maxBound

      -- 0x40 + 16 + 16 + 16 + 16 + 16
      ByteString -> Put
Put.putByteString ByteString
posDigest
      ByteString -> Put
Put.putByteString ByteString
indDigest
      ByteString -> Put
Put.putByteString ByteString
attDigest
      ByteString -> Put
Put.putByteString ByteString
nodDigest
      ByteString -> Put
Put.putByteString ByteString
metDigest

      -- 0x90 + posCompressed
      ByteString -> Put
Put.putByteString ByteString
posCompressed

      -- 0x90 + posCompressed + indCompressed
      ByteString -> Put
Put.putByteString ByteString
indCompressed

      -- 0x90 + posCompressed + indCompressed + attCompressed
      ByteString -> Put
Put.putByteString ByteString
attCompressed

      -- 0x90 + posCompressed + indCompressed + attCompressed + nodCompressed
      ByteString -> Put
Put.putByteString ByteString
nodCompressed

      -- 0x90 + posCompressed + indCompressed + attCompressed + nodCompressed + metCompressed
      ByteString -> Put
Put.putByteString ByteString
metCompressed

      -- 0x90 + posCompressed + indCompressed + attCompressed + nodCompressed + metCompressed + 4 + 4
      FilePath -> Put
Put.putStringUtf8 FilePath
"📦🌋"

encodeItems
  :: ( Storable a
     , Generic.Vector v a
     , MonadIO m
     )
  => v a
  -> m (ByteString, ByteString)
encodeItems :: forall a (v :: * -> *) (m :: * -> *).
(Storable a, Vector v a, MonadIO m) =>
v a -> m (ByteString, ByteString)
encodeItems v a
items = do
  let bytes :: Vector Word8
bytes = forall a b. (Storable a, Storable b) => Vector a -> Vector b
Storable.unsafeCast @_ @Word8 (Vector a -> Vector Word8) -> Vector a -> Vector Word8
forall a b. (a -> b) -> a -> b
$ v a -> Vector a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
Vector.convert v a
items
  IO (ByteString, ByteString) -> m (ByteString, ByteString)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (ByteString, ByteString) -> m (ByteString, ByteString))
-> IO (ByteString, ByteString) -> m (ByteString, ByteString)
forall a b. (a -> b) -> a -> b
$ Vector Word8
-> (Ptr Word8 -> IO (ByteString, ByteString))
-> IO (ByteString, ByteString)
forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
Storable.unsafeWith Vector Word8
bytes \Ptr Word8
ptr -> do
    ByteString
buf <- CStringLen -> IO ByteString
unsafePackCStringLen
      ( Ptr Word8 -> Ptr CChar
forall a b. Ptr a -> Ptr b
Foreign.castPtr Ptr Word8
ptr
      , Vector Word8 -> Int
forall a. Storable a => Vector a -> Int
Storable.length Vector Word8
bytes
      )
    let
      -- XXX: Process buffer before bytes/ptr go out of scope!
      !bufHash :: ByteString
bufHash = ByteString -> ByteString
MD5.hash ByteString
buf
      !zBuf :: ByteString
zBuf = Int -> ByteString -> ByteString
Zstd.compress Int
Zstd.maxCLevel ByteString
buf
    (ByteString, ByteString) -> IO (ByteString, ByteString)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString
bufHash, ByteString
zBuf)

encodeCBOR :: CBOR.Serialise a => a -> (Int, ByteString, ByteString)
encodeCBOR :: forall a. Serialise a => a -> (Int, ByteString, ByteString)
encodeCBOR a
stuff =
  ( ByteString -> Int
ByteString.length ByteString
buf
  , ByteString -> ByteString
MD5.hash ByteString
buf
  , Int -> ByteString -> ByteString
Zstd.compress Int
Zstd.maxCLevel ByteString
buf
  )
  where
    buf :: ByteString
buf = LByteString -> ByteString
BSL.toStrict (a -> LByteString
forall a. Serialise a => a -> LByteString
CBOR.serialise a
stuff)

-- * Decoding

loadIndexed
  :: ( Storable attrs
     , Storable nodes
     , CBOR.Serialise meta
     , Show meta
     , Typeable nodes
     , HasLogFunc env
     , MonadResource m
     , MonadVulkan env m
     )
  => Queues Vk.CommandPool
  -> FilePath
  -> m
    ( Resource.ReleaseKey
    , (meta, Storable.Vector nodes, Model.Indexed 'Buffer.Staged Vec3.Packed attrs)
    )
loadIndexed :: forall attrs nodes meta env (m :: * -> *).
(Storable attrs, Storable nodes, Serialise meta, Show meta,
 Typeable nodes, HasLogFunc env, MonadResource m,
 MonadVulkan env m) =>
Queues CommandPool
-> FilePath
-> m (ReleaseKey,
      (meta, Vector nodes, Indexed 'Staged Packed attrs))
loadIndexed Queues CommandPool
pools FilePath
fp = do
  Utf8Builder -> m ()
forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) =>
Utf8Builder -> m ()
logInfo (Utf8Builder -> m ()) -> Utf8Builder -> m ()
forall a b. (a -> b) -> a -> b
$ Utf8Builder
"Loading " Utf8Builder -> Utf8Builder -> Utf8Builder
forall a. Semigroup a => a -> a -> a
<> FilePath -> Utf8Builder
forall a. IsString a => FilePath -> a
fromString FilePath
fp
  (meta
meta, Vector nodes
nodes, (Vector Packed
positions, Vector Word32
indices, Vector attrs
attrs)) <- FilePath
-> m (meta, Vector nodes,
      (Vector Packed, Vector Word32, Vector attrs))
forall attrs env nodes meta (m :: * -> *).
(Storable attrs, Serialise meta, Storable nodes, Typeable nodes,
 HasLogFunc env, MonadReader env m, MonadIO m) =>
FilePath
-> m (meta, Vector nodes,
      (Vector Packed, Vector Word32, Vector attrs))
loadBlobs FilePath
fp
  Utf8Builder -> m ()
forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) =>
Utf8Builder -> m ()
logDebug (Utf8Builder -> m ()) -> Utf8Builder -> m ()
forall a b. (a -> b) -> a -> b
$ meta -> Utf8Builder
forall a. Show a => a -> Utf8Builder
displayShow meta
meta

  Utf8Builder -> m ()
forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) =>
Utf8Builder -> m ()
logDebug (Utf8Builder -> m ()) -> Utf8Builder -> m ()
forall a b. (a -> b) -> a -> b
$ Utf8Builder
"Staging " Utf8Builder -> Utf8Builder -> Utf8Builder
forall a. Semigroup a => a -> a -> a
<> FilePath -> Utf8Builder
forall a. IsString a => FilePath -> a
fromString FilePath
fp
  env
context <- m env
forall r (m :: * -> *). MonadReader r m => m r
ask
  Indexed 'Staged Packed attrs
indexed <- env
-> Queues CommandPool
-> Vector Packed
-> Vector attrs
-> Vector Word32
-> m (Indexed 'Staged Packed attrs)
forall context pos attrs (io :: * -> *).
(HasVulkan context, Storable pos, Storable attrs,
 MonadUnliftIO io) =>
context
-> Queues CommandPool
-> Vector pos
-> Vector attrs
-> Vector Word32
-> io (Indexed 'Staged pos attrs)
Model.createStaged env
context Queues CommandPool
pools Vector Packed
positions Vector attrs
attrs Vector Word32
indices
  ReleaseKey
key <- IO () -> m ReleaseKey
forall (m :: * -> *). MonadResource m => IO () -> m ReleaseKey
Resource.register (IO () -> m ReleaseKey) -> IO () -> m ReleaseKey
forall a b. (a -> b) -> a -> b
$ env -> Indexed 'Staged Packed attrs -> IO ()
forall context (io :: * -> *) (storage :: Store) pos attrs.
(HasVulkan context, MonadUnliftIO io) =>
context -> Indexed storage pos attrs -> io ()
Model.destroyIndexed env
context Indexed 'Staged Packed attrs
indexed
  pure (ReleaseKey
key, (meta
meta, Vector nodes
nodes, Indexed 'Staged Packed attrs
indexed))

loadBlobs
  :: forall attrs env nodes meta m
  .  ( Storable attrs
     , CBOR.Serialise meta
     , Storable nodes
     , Typeable nodes
     , HasLogFunc env
     , MonadReader env m
     , MonadIO m
     )
  => FilePath
  -> m
    ( meta
    , Storable.Vector nodes
    , ( Storable.Vector Vec3.Packed
      , Storable.Vector Word32
      , Storable.Vector attrs
      )
    )
loadBlobs :: forall attrs env nodes meta (m :: * -> *).
(Storable attrs, Serialise meta, Storable nodes, Typeable nodes,
 HasLogFunc env, MonadReader env m, MonadIO m) =>
FilePath
-> m (meta, Vector nodes,
      (Vector Packed, Vector Word32, Vector attrs))
loadBlobs FilePath
fp = do
  LByteString
blob <- FilePath -> m LByteString
forall (m :: * -> *). MonadIO m => FilePath -> m LByteString
BSL.readFile FilePath
fp
  let
    getter :: Get
  (Word8, meta, Vector nodes,
   (Vector Packed, Vector Word32, Vector attrs))
getter = do
      ByteString
magicStart <- Int -> Get ByteString
Get.getByteString (Int
4Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
4)

      Word8
verBreaks <- Get Word8
Get.getWord8
      Word8
verTweaks <- Get Word8
Get.getWord8
      FilePath -> Word8 -> Word8 -> Get ()
forall (m :: * -> *) a.
(MonadFail m, Show a, Eq a) =>
FilePath -> a -> a -> m ()
guardEq FilePath
"Codec version" Word8
VER_BREAKS Word8
verBreaks

      Word16
_reserved16 <- Get Word16
Get.getWord16le
      Word32
_reserved32 <- Get Word32
Get.getWord32le

      Int
numPositions <- (Word32 -> Int) -> Get Word32 -> Get Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Get Word32
Get.getWord32le
      Int
lenPositions <- (Word32 -> Int) -> Get Word32 -> Get Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Get Word32
Get.getWord32le

      Int
numIndices <- (Word32 -> Int) -> Get Word32 -> Get Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Get Word32
Get.getWord32le
      Int
lenIndices <- (Word32 -> Int) -> Get Word32 -> Get Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Get Word32
Get.getWord32le

      Int
sizeOfAttr <- (Word32 -> Int) -> Get Word32 -> Get Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Get Word32
Get.getWord32le
      Int
lenAttrs   <- (Word32 -> Int) -> Get Word32 -> Get Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Get Word32
Get.getWord32le
      FilePath -> Int -> Int -> Get ()
forall (m :: * -> *) a.
(MonadFail m, Show a, Eq a) =>
FilePath -> a -> a -> m ()
guardEq FilePath
"Attribute size" (attrs -> Int
forall a. Storable a => a -> Int
Foreign.sizeOf (FilePath -> attrs
forall a. HasCallStack => FilePath -> a
error FilePath
"sizeOfAttr" :: attrs)) Int
sizeOfAttr

      Int
sizeOfNode <- (Word32 -> Int) -> Get Word32 -> Get Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Get Word32
Get.getWord32le
      Int
lenNodes   <- (Word32 -> Int) -> Get Word32 -> Get Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Get Word32
Get.getWord32le
      FilePath -> Int -> Int -> Get ()
forall (m :: * -> *) a.
(MonadFail m, Show a, Eq a) =>
FilePath -> a -> a -> m ()
guardEq
        (FilePath
"Node size for " FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> TyCon -> FilePath
forall a. Show a => a -> FilePath
show (TypeRep -> TyCon
typeRepTyCon (TypeRep -> TyCon)
-> (Proxy nodes -> TypeRep) -> Proxy nodes -> TyCon
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Proxy nodes -> TypeRep
forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (Proxy nodes -> TyCon) -> Proxy nodes -> TyCon
forall a b. (a -> b) -> a -> b
$ forall {t}. Proxy t
forall {k} (t :: k). Proxy t
Proxy @nodes))
        (nodes -> Int
forall a. Storable a => a -> Int
Foreign.sizeOf (FilePath -> nodes
forall a. HasCallStack => FilePath -> a
error FilePath
"sizeOfNode" :: nodes))
        Int
sizeOfNode

      Int
sizeOfMeta <- (Word32 -> Int) -> Get Word32 -> Get Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Get Word32
Get.getWord32le
      Int
lenMeta    <- (Word32 -> Int) -> Get Word32 -> Get Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Get Word32
Get.getWord32le
      -- guardEq "Meta size" (Foreign.sizeOf (error "sizeOf" :: meta)) sizeOfMeta

      Word32
_reserved32 <- Get Word32
Get.getWord32le
      Word32
_reserved32 <- Get Word32
Get.getWord32le

      ByteString
posDigest <- Int -> Get ByteString
Get.getByteString Int
16
      ByteString
indDigest <- Int -> Get ByteString
Get.getByteString Int
16
      ByteString
attDigest <- Int -> Get ByteString
Get.getByteString Int
16
      ByteString
nodDigest <- Int -> Get ByteString
Get.getByteString Int
16
      ByteString
metDigest <- Int -> Get ByteString
Get.getByteString Int
16

      -- XXX: end for static part for VER_BREAKS
      Int64
staticDone <- Get Int64
Get.bytesRead
      FilePath -> Int64 -> Int64 -> Get ()
forall (m :: * -> *) a.
(MonadFail m, Show a, Eq a) =>
FilePath -> a -> a -> m ()
guardEq (FilePath
"End of static part for v" FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> Word8 -> FilePath
forall a. Show a => a -> FilePath
show Word8
VER_BREAKS) Int64
0x90 Int64
staticDone

      let payloadSize :: Int64
payloadSize = Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
lenPositions Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lenIndices Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lenAttrs Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lenNodes Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lenMeta)
      FilePath -> Int64 -> Int64 -> Get ()
forall (m :: * -> *) a.
(MonadFail m, Show a, Eq a) =>
FilePath -> a -> a -> m ()
guardEq FilePath
"Blob size" (LByteString -> Int64
BSL.length LByteString
blob) (Int64 -> Get ()) -> Int64 -> Get ()
forall a b. (a -> b) -> a -> b
$
        Int64
staticDone Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
+ Int64
payloadSize Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
+ Int64
4 Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
+ Int64
4

      ByteString
zPositions <- Int -> Get ByteString
Get.getByteString Int
lenPositions
      ByteString
zIndices   <- Int -> Get ByteString
Get.getByteString Int
lenIndices
      ByteString
zAttrs     <- Int -> Get ByteString
Get.getByteString Int
lenAttrs
      ByteString
zNodes     <- Int -> Get ByteString
Get.getByteString Int
lenNodes
      ByteString
zMetas     <- Int -> Get ByteString
Get.getByteString Int
lenMeta

      ByteString
magicFinish <- Int -> Get ByteString
Get.getByteString (Int
4Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
4)

      let magicReverse :: ByteString
magicReverse = Int -> ByteString -> ByteString
ByteString.drop Int
4 ByteString
magicFinish ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Int -> ByteString -> ByteString
ByteString.take Int
4 ByteString
magicFinish
      FilePath -> ByteString -> ByteString -> Get ()
forall (m :: * -> *) a.
(MonadFail m, Show a, Eq a) =>
FilePath -> a -> a -> m ()
guardEq FilePath
"Magic final" ByteString
magicStart ByteString
magicReverse

      Vector Packed
positions <- FilePath
-> ByteString -> Maybe Int -> ByteString -> Get (Vector Packed)
forall item (m :: * -> *).
(Storable item, MonadFail m) =>
FilePath
-> ByteString -> Maybe Int -> ByteString -> m (Vector item)
decodeItems FilePath
"Positions"  ByteString
posDigest (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
numPositions) ByteString
zPositions
      Vector Word32
indices   <- FilePath
-> ByteString -> Maybe Int -> ByteString -> Get (Vector Word32)
forall item (m :: * -> *).
(Storable item, MonadFail m) =>
FilePath
-> ByteString -> Maybe Int -> ByteString -> m (Vector item)
decodeItems FilePath
"Indices"    ByteString
indDigest (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
numIndices)   ByteString
zIndices
      Vector attrs
attrs     <- FilePath
-> ByteString -> Maybe Int -> ByteString -> Get (Vector attrs)
forall item (m :: * -> *).
(Storable item, MonadFail m) =>
FilePath
-> ByteString -> Maybe Int -> ByteString -> m (Vector item)
decodeItems FilePath
"Attributes" ByteString
attDigest (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
numPositions) ByteString
zAttrs
      Vector nodes
nodes     <- FilePath
-> ByteString -> Maybe Int -> ByteString -> Get (Vector nodes)
forall item (m :: * -> *).
(Storable item, MonadFail m) =>
FilePath
-> ByteString -> Maybe Int -> ByteString -> m (Vector item)
decodeItems FilePath
"Nodes"      ByteString
nodDigest Maybe Int
forall a. Maybe a
Nothing             ByteString
zNodes

      meta
meta      <- FilePath -> ByteString -> Int -> ByteString -> Get meta
forall a (m :: * -> *).
(Serialise a, MonadFail m) =>
FilePath -> ByteString -> Int -> ByteString -> m a
decodeCBOR  FilePath
"Metadata"   ByteString
metDigest Int
sizeOfMeta          ByteString
zMetas

      pure (Word8
verTweaks, meta
meta, Vector nodes
nodes, (Vector Packed
positions, Vector Word32
indices, Vector attrs
attrs))

  let (Word8
verTweaks, meta
meta, Vector nodes
nodes, (Vector Packed, Vector Word32, Vector attrs)
blobs) = Get
  (Word8, meta, Vector nodes,
   (Vector Packed, Vector Word32, Vector attrs))
-> LByteString
-> (Word8, meta, Vector nodes,
    (Vector Packed, Vector Word32, Vector attrs))
forall a. Get a -> LByteString -> a
runGet Get
  (Word8, meta, Vector nodes,
   (Vector Packed, Vector Word32, Vector attrs))
getter LByteString
blob

  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Word8
verTweaks Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
VER_TWEAKS) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
    Utf8Builder -> m ()
forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) =>
Utf8Builder -> m ()
logWarn (Utf8Builder -> m ()) -> Utf8Builder -> m ()
forall a b. (a -> b) -> a -> b
$ [Utf8Builder] -> Utf8Builder
forall a. Monoid a => [a] -> a
mconcat
      [ Utf8Builder
"Format tweak version mismatch: "
      , Word8 -> Utf8Builder
forall a. Display a => a -> Utf8Builder
display Word8
verTweaks
      , Utf8Builder
" /= "
      , Word8 -> Utf8Builder
forall a. Display a => a -> Utf8Builder
display Word8
VER_TWEAKS
      ]

  pure (meta
meta, Vector nodes
nodes, (Vector Packed, Vector Word32, Vector attrs)
blobs)

decodeItems
  :: forall item m
  . (Storable item, MonadFail m)
  => String
  -> ByteString
  -> Maybe Int
  -> ByteString
  -> m (Storable.Vector item)
decodeItems :: forall item (m :: * -> *).
(Storable item, MonadFail m) =>
FilePath
-> ByteString -> Maybe Int -> ByteString -> m (Vector item)
decodeItems FilePath
label ByteString
digest Maybe Int
expectedSize ByteString
zBytes =
  case ByteString -> Decompress
Zstd.decompress ByteString
zBytes of
    Zstd.Error FilePath
err ->
      FilePath -> m (Vector item)
forall (m :: * -> *) a. MonadFail m => FilePath -> m a
fail FilePath
err
    Decompress
Zstd.Skip -> do
      FilePath -> Int -> Int -> m ()
forall (m :: * -> *) a.
(MonadFail m, Show a, Eq a) =>
FilePath -> a -> a -> m ()
guardEq (FilePath
label FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
" size") Int
0 Int
itemSize
      case Maybe Int
expectedSize of
        Just Int
size ->
          Vector item -> m (Vector item)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vector item -> m (Vector item))
-> (item -> Vector item) -> item -> m (Vector item)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> item -> Vector item
forall a. Storable a => Int -> a -> Vector a
Storable.replicate Int
size (item -> m (Vector item)) -> item -> m (Vector item)
forall a b. (a -> b) -> a -> b
$
            IO item -> item
forall a. IO a -> a
unsafePerformIO (Ptr item -> IO item
forall a. Storable a => Ptr a -> IO a
Foreign.peek Ptr item
forall a. Ptr a
Foreign.nullPtr)
        Maybe Int
Nothing ->
          Vector item -> m (Vector item)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vector item
forall a. Storable a => Vector a
Storable.empty
    Zstd.Decompress ByteString
bytes -> do
      let (ForeignPtr Word8
buf, Int
bufOff, Int
bufLen) = ByteString -> (ForeignPtr Word8, Int, Int)
BSI.toForeignPtr ByteString
bytes
      FilePath -> Int -> Int -> m ()
forall (m :: * -> *) a.
(MonadFail m, Show a, Eq a) =>
FilePath -> a -> a -> m ()
guardEq (FilePath
label FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
" buffer offset") Int
0 Int
bufOff
      case Maybe Int
expectedSize of
        Maybe Int
Nothing ->
          () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        Just Int
size ->
          FilePath -> Int -> Int -> m ()
forall (m :: * -> *) a.
(MonadFail m, Show a, Eq a) =>
FilePath -> a -> a -> m ()
guardEq (FilePath
label FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
" buffer size") (Int
itemSize Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
size) Int
bufLen
      FilePath -> ByteString -> ByteString -> m ()
forall (m :: * -> *) a.
(MonadFail m, Show a, Eq a) =>
FilePath -> a -> a -> m ()
guardEq (FilePath
label FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
" hash") ByteString
digest (ByteString -> ByteString
MD5.hash ByteString
bytes)
      let
        !items :: Vector item
items =
          forall a b. (Storable a, Storable b) => Vector a -> Vector b
Storable.unsafeCast @Word8 @item (Vector Word8 -> Vector item) -> Vector Word8 -> Vector item
forall a b. (a -> b) -> a -> b
$
            ForeignPtr Word8 -> Int -> Vector Word8
forall a. Storable a => ForeignPtr a -> Int -> Vector a
Storable.unsafeFromForeignPtr0 ForeignPtr Word8
buf Int
bufLen
      case Maybe Int
expectedSize of
        Maybe Int
Nothing ->
          () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        Just Int
size ->
          FilePath -> Int -> Int -> m ()
forall (m :: * -> *) a.
(MonadFail m, Show a, Eq a) =>
FilePath -> a -> a -> m ()
guardEq (FilePath
label FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
" size") Int
size (Vector item -> Int
forall a. Storable a => Vector a -> Int
Storable.length Vector item
items)
      pure Vector item
items
  where
    itemSize :: Int
itemSize = forall a. Storable a => a -> Int
Foreign.sizeOf @item item
forall a. HasCallStack => a
undefined

decodeCBOR
  :: ( CBOR.Serialise a
     , MonadFail m
     )
  => String
  -> ByteString
  -> Int
  -> ByteString
  -> m a
decodeCBOR :: forall a (m :: * -> *).
(Serialise a, MonadFail m) =>
FilePath -> ByteString -> Int -> ByteString -> m a
decodeCBOR FilePath
label ByteString
digest Int
expectedSize ByteString
zBytes =
  case ByteString -> Decompress
Zstd.decompress ByteString
zBytes of
    Zstd.Error FilePath
err ->
      FilePath -> m a
forall (m :: * -> *) a. MonadFail m => FilePath -> m a
fail (FilePath -> m a) -> FilePath -> m a
forall a b. (a -> b) -> a -> b
$ FilePath
label FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
": zstd error (" FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
err FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
")"
    Decompress
Zstd.Skip ->
      FilePath -> m a
forall (m :: * -> *) a. MonadFail m => FilePath -> m a
fail (FilePath -> m a) -> FilePath -> m a
forall a b. (a -> b) -> a -> b
$ FilePath
label FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
": empty zstd"
    Zstd.Decompress ByteString
bytes -> do
      FilePath -> Int -> Int -> m ()
forall (m :: * -> *) a.
(MonadFail m, Show a, Eq a) =>
FilePath -> a -> a -> m ()
guardEq (FilePath
label FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
" size") Int
expectedSize (ByteString -> Int
ByteString.length ByteString
bytes)
      FilePath -> ByteString -> ByteString -> m ()
forall (m :: * -> *) a.
(MonadFail m, Show a, Eq a) =>
FilePath -> a -> a -> m ()
guardEq (FilePath
label FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
" hash") ByteString
digest (ByteString -> ByteString
MD5.hash ByteString
bytes)
      case LByteString -> Either DeserialiseFailure a
forall a. Serialise a => LByteString -> Either DeserialiseFailure a
CBOR.deserialiseOrFail (ByteString -> LByteString
BSL.fromStrict ByteString
bytes) of
        Right a
value ->
          a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
value
        Left (CBOR.DeserialiseFailure Int64
_off FilePath
err) ->
          FilePath -> m a
forall (m :: * -> *) a. MonadFail m => FilePath -> m a
fail (FilePath -> m a) -> FilePath -> m a
forall a b. (a -> b) -> a -> b
$ [FilePath] -> FilePath
unlines
            [ FilePath
label FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
" deserialise failure:"
            , FilePath
err
            ]

-- * Utils

guardEq :: (MonadFail m, Show a, Eq a) => String -> a -> a -> m ()
guardEq :: forall (m :: * -> *) a.
(MonadFail m, Show a, Eq a) =>
FilePath -> a -> a -> m ()
guardEq FilePath
label a
expected a
got = do
  if a
expected a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
got then
    -- traceM . fromString $ label <> " match: " <> show got
    () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  else
    FilePath -> m ()
forall (m :: * -> *) a. MonadFail m => FilePath -> m a
fail (FilePath -> m ()) -> FilePath -> m ()
forall a b. (a -> b) -> a -> b
$ [FilePath] -> FilePath
unlines
      [ FilePath
label FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
" mismatch"
      , FilePath
"\tExpected: " FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> a -> FilePath
forall a. Show a => a -> FilePath
show a
expected
      , FilePath
"\tGot     : " FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> a -> FilePath
forall a. Show a => a -> FilePath
show a
got
      ]