{-
Module : Data.Nbt
Copyright : (c) David R. Garland 2022
License : MIT
Maintainer : davidrgarland@me.com
Functions related to parsing & serializing the NBT file format.
-}
module Data.Nbt
( Type (..)
, Nbt (..)
, Nbt'
, MapNbt
, Tag (..)
, Tag'
, MapTag
, Cmpnd (..)
, Cmpnd'
, MapCmpnd
, pattern Compound'
, lookupNbt
, getNbt
, lookupTag
, getTag
, typeOf
, readUncompressed
, writeUncompressed
, readCompressed
, writeCompressed
) where
import Codec.Compression.GZip as GZ
import Control.Monad
import Data.ByteString qualified as B
import Data.ByteString.Lazy qualified as BL
import Data.Foldable
import Data.Foldable.WithIndex
import Data.Int
import Data.Map.Strict qualified as M
import Data.Map.Strict (Map)
import Data.RRBVector qualified as R
import Data.RRBVector (Vector, (<|))
import Data.Serialize hiding (label)
import Data.Text (Text)
import Data.Text.Encoding
import Data.Word
-- | An NBT tag type, in order so that 'toEnum' and 'fromEnum' are useful for serialization and
-- deserialization.
data Type
= EndType -- ^ 0x00 NUL. Denotes the end of a file.
| ByteType -- ^ 0x01 SOH. Corresponds to `Data.Nbt.Byte`.
| ShortType -- ^ 0x02 STX. Corresponds to `Data.Nbt.Short`.
| IntType -- ^ 0x03 ETX. Corresponds to `Data.Nbt.Int`.
| LongType -- ^ 0x04 EOT. Corresponds to `Data.Nbt.Long`.
| FloatType -- ^ 0x05 ENQ. Corresponds to `Data.Nbt.Float`.
| DoubleType -- ^ 0x06 ACK. Corresponds to `Data.Nbt.Double`.
| ByteArrayType -- ^ 0x07 BEL. Corresponds to `ByteArray`.
| StringType -- ^ 0x08 BS. Corresponds to `Data.Nbt.String`.
| ListType -- ^ 0x09 HT. Corresponds to `List`.
| CompoundType -- ^ 0x0a LF. Corresponds to `Compound`.
| IntArrayType -- ^ 0x0b VT. Corresponds to `IntArray`.
| LongArrayType -- ^ 0x0c FF. Corresponds to `LongArray`.
deriving (Show, Eq, Ord, Enum)
-- | The main NBT type, with support for serialization and deserialization. It couples a 'Text' label with a 'Tag'.
--
-- When serialized: 1-byte unsigned 'Type' + 2-byte unsigned integer length + N bytes utf-8 text + 'Tag'
data Nbt b = Nbt
{ label :: Text
, tag :: Tag b
}
deriving (Show, Eq, Ord, Functor)
-- | The "simple" form of `Nbt`, which has no label map on `Compound`s.
type Nbt' = Nbt ()
-- | A version of `Nbt` which has a map from labels to element indices for fast lookup.
type MapNbt = Nbt (Map Text Int)
-- | An NBT tag, responsible for storing the actual data.
data Tag b
= Byte Int8 -- ^ 1-byte signed integer
| Short Int16 -- ^ 2-byte signed integer
| Int Int32 -- ^ 4-byte signed integer
| Long Int64 -- ^ 8-byte signed integer
| Float Float -- ^ 4-byte float
| Double Double -- ^ 8-byte double
| ByteArray (Vector Int8) -- ^ 4-byte signed integer length + N 1-byte signed integers
| String Text -- ^ 2-byte unsigned integer length + N bytes utf-8 text
| List (Vector (Tag b)) -- ^ 1-byte unsigned 'Type' + 4-byte signed integer length + N homogenous tags
| Compound (Cmpnd b) -- ^ N heterogenous NBT elements punctuated by an 'EndType'
| IntArray (Vector Int32) -- ^ 4-byte signed integer length + N 4-byte signed integers
| LongArray (Vector Int64) -- ^ 4-byte signed integer length + N 8-byte signed integers
deriving (Show, Eq, Ord, Functor)
-- | A pattern synonym for "Compound (Cmpnd () v)".
pattern Compound' :: Vector Nbt' -> Tag'
pattern Compound' v <- Compound (Cmpnd () v)
where Compound' v = Compound (Cmpnd () v)
{-# COMPLETE Byte, Short, Int, Long, Float, Double, ByteArray, String, List, Compound', IntArray, LongArray #-}
-- | The "simple" form of `Tag`, which has no label map on `Compound`s.
type Tag' = Tag ()
-- | A version of `Tag` which has a map from labels to element indices for fast lookup.
type MapTag = Tag (Map Text Int)
-- | The payload of the `Compound` constructor.
data Cmpnd b = Cmpnd b (Vector (Nbt b))
deriving (Show, Eq, Ord, Functor)
-- | The "simple" form of `Cmpnd`, which has no label map.
type Cmpnd' = Cmpnd ()
-- | A version of `Cmpnd` which has a map from labels to element indices for fast lookup.
type MapCmpnd = Cmpnd (Map Text Int)
-- | Use a `MapCmpnd` to quickly look up a `MapNbt` by label.
lookupNbt :: Text -> MapCmpnd -> Maybe MapNbt
lookupNbt t (Cmpnd m v) = (v R.!) <$> M.lookup t m
-- | Use a `MapCmpnd` to quickly get a `MapNbt` by label.
getNbt :: Text -> MapCmpnd -> MapNbt
getNbt t (Cmpnd m v) = v R.! (m M.! t)
-- | Runs `lookupNbt`, then retrieves the `tag`.
lookupTag :: Text -> MapCmpnd -> Maybe MapTag
lookupTag t c = tag <$> lookupNbt t c
-- | Runs `getNbt`, then retrieves the `tag`.
getTag :: Text -> MapCmpnd -> MapTag
getTag t c = tag $ getNbt t c
instance Serialize (Nbt ()) where
get = get >>= getNbtByType
put (Nbt n d) = put (typeOf d) >> putString n >> putTag d
instance Serialize (Nbt (Map Text Int)) where
get = get >>= getNbtByType
put (Nbt n d) = put (typeOf d) >> putString n >> putTag d
instance Serialize Type where
get = toEnum . fromIntegral <$> getWord8
put = putWord8 . fromIntegral . fromEnum
class GenMap b where
genMap :: Vector (Nbt b) -> b
instance GenMap () where
genMap = const ()
{-# INLINE CONLIKE genMap #-}
instance GenMap (Map Text Int) where
genMap = ifoldr (\i (Nbt l _) a -> M.insert l i a) M.empty
{-# INLINE genMap #-}
-- | Get the type of an NBT tag.
typeOf :: Tag b -> Type
typeOf = \case
Byte _ -> ByteType
Short _ -> ShortType
Int _ -> IntType
Long _ -> LongType
Float _ -> FloatType
Double _ -> DoubleType
ByteArray _ -> ByteArrayType
String _ -> StringType
List _ -> ListType
Compound _ -> CompoundType
IntArray _ -> IntArrayType
LongArray _ -> LongArrayType
putTag :: Serialize (Nbt b) => Tag b -> Put
putTag = \case
Byte b -> put b
Short s -> put s
Int i -> put i
Long l -> put l
Float f -> putFloat32be f
Double d -> putFloat64be d
ByteArray bs -> putArray put bs
String str -> putString str
List ts -> do
ty <-
if null ts then do
pure EndType
else do
let xt = typeOf (ts R.! 0)
let xs = R.drop 1 ts
if all (\e -> typeOf e == xt) xs then
pure xt
else
error "attempted to write heterogenous list"
put ty
putArray putTag ts
Compound (Cmpnd _ ts) -> traverse_ put ts >> put EndType
IntArray is -> putArray put is
LongArray is -> putArray put is
where
putArray :: (e -> Put) -> Vector e -> Put
putArray putter a = do
put (fromIntegral (length a) :: Int32)
traverse_ putter a
{-# SPECIALIZE putTag :: Tag () -> Put #-}
{-# SPECIALIZE putTag :: Tag (Map Text Int) -> Put #-}
getNbtByType :: GenMap b => Type -> Get (Nbt b)
getNbtByType ty = Nbt <$> getString <*> getByType ty
{-# SPECIALIZE getNbtByType :: Type -> Get (Nbt ()) #-}
{-# SPECIALIZE getNbtByType :: Type -> Get (Nbt (Map Text Int)) #-}
getByType :: GenMap b => Type -> Get (Tag b)
getByType = \case
EndType -> fail "cannot get a value of EndType"
ByteType -> Byte <$> get
ShortType -> Short <$> get
IntType -> Int <$> get
LongType -> Long <$> get
FloatType -> Float <$> getFloat32be
DoubleType -> Double <$> getFloat64be
ByteArrayType -> ByteArray <$> getArray get
StringType -> String <$> getString
ListType -> List <$> (getArray . getByType =<< get)
CompoundType -> do
x <- getCompound
pure $ Compound (Cmpnd (genMap x) x)
IntArrayType -> IntArray <$> getArray get
LongArrayType -> LongArray <$> getArray get
where
getCompound :: GenMap b => Get (Vector (Nbt b))
getCompound = get >>= \case
EndType -> pure R.empty
ty -> (<|) <$> getNbtByType ty <*> getCompound
{-# SPECIALIZE getByType :: Type -> Get (Tag ()) #-}
{-# SPECIALIZE getByType :: Type -> Get (Tag (Map Text Int)) #-}
{-
-- File IO
-}
-- | Read an uncompressed NBT file. Exceptions are thrown as usual if the internal call to 'Data.ByteString.readFile' fails,
-- and the 'Either' message is that returned by 'Data.Serialize.decode'.
readUncompressed :: Serialize (Nbt b) => FilePath -> IO (Either String (Nbt b))
readUncompressed fp = decode <$> B.readFile fp
-- | Write an uncompressed NBT file.
writeUncompressed :: Serialize (Nbt b) => FilePath -> Nbt b -> IO ()
writeUncompressed fp nbt = B.writeFile fp $ runPut (put nbt)
-- | Read a GZip-compressed NBT file. Exceptions are thrown as usual if the internal call to 'Data.ByteString.Lazy.readFile' fails,
-- and the 'Either' message is that returned by 'Data.Serialize.decode'.
readCompressed :: Serialize (Nbt b) => FilePath -> IO (Either String (Nbt b))
readCompressed fp = do
bs <- GZ.decompress <$> BL.readFile fp
pure . decode $ B.toStrict bs
-- | Write a GZip-compressed NBT file.
writeCompressed :: Serialize (Nbt b) => FilePath -> Nbt b -> IO ()
writeCompressed fp nbt = do
let bs = GZ.compress . BL.fromStrict $ runPut (put nbt)
B.writeFile fp (B.toStrict bs)
{-
-- Internal Utility Functions
-}
getArray :: Get e -> Get (Vector e)
getArray getter = do
len <- get :: Get Int32
elts <- replicateM (fromIntegral len) getter
pure . R.fromList $ elts
getString :: Get Text
getString = do
len <- get :: Get Word16
decodeUtf8 <$> getByteString (fromIntegral len)
putString :: Text -> Put
putString t = do
let b = encodeUtf8 t
put (fromIntegral (B.length b) :: Word16)
putByteString b