An NBT parser written in Haskell

NBT.hs

module Minecraft.NBT (findTag, indexTag, mapTag, parseNBT, Tag(..)) where

import qualified Data.ByteString.Char8 as BC
import qualified Data.ByteString.Lazy as BL
import qualified Codec.Compression.Zlib as ZL
import qualified Codec.Compression.GZip as GZ
import qualified Data.Map as M
import Data.Attoparsec.Lazy as AP
import Data.Attoparsec.Binary
import Data.Word
import Control.Monad (void, mapM)
import Control.Applicative
import Data.List (lookup)
import Data.Maybe (fromJust)

data Tag = TagByte      String Int
         | TagShort     String Int
         | TagInt       String Int
         | TagLong      String Int
         | TagFloat     String Float
         | TagDouble    String Double
         | TagByteArray String [Int]
         | TagString    String String   
         | TagList      String [Tag]
         | TagCompound  String (M.Map String Tag)
         | TagIntArray  String [Int]
         deriving (Show, Eq)

tagName (TagByte n _)      = n   
tagName (TagShort n _)     = n   
tagName (TagInt n _)       = n   
tagName (TagLong n _)      = n   
tagName (TagFloat n _)     = n   
tagName (TagDouble n _)    = n   
tagName (TagByteArray n _) = n
tagName (TagString n _)    = n   
tagName (TagList n _)      = n   
tagName (TagCompound n _)  = n   
tagName (TagIntArray n _)  = n

parseTagByteArrayNoHeader :: Parser [Int]
parseTagByteArrayNoHeader = do
  len <- fromIntegral <$> anyWord32be
  mapM (const (fromIntegral <$> anyWord8)) [1 .. len]

parseTagListNoHeader :: Parser [Tag]
parseTagListNoHeader = do
  id <- fromIntegral <$> anyWord8
  len <- fromIntegral <$> anyWord32be
  let f = fromJust $ lookup id idToFunc
  mapM (const f) [1 .. len]

parseTagCompoundNoHeader :: Parser (M.Map String Tag)
parseTagCompoundNoHeader = do
  tags <- many parseTag
  parseTagEnd
  let map = foldl (\acc t -> M.insert (tagName t) t acc) M.empty tags
  return map 

parseTagIntArrayNoHeader :: Parser [Int]
parseTagIntArrayNoHeader = do
  len <- fromIntegral <$> anyWord32be
  mapM (const (fromIntegral <$> anyWord32be)) [1 .. len]

idToFunc :: [(Int, Parser Tag)]
idToFunc = [(1, TagByte <$> pure "" <*> (fromIntegral <$> anyWord8)),
            (2, TagShort <$> pure "" <*> (fromIntegral <$> anyWord16be)),
            (3, TagInt <$> pure "" <*> (fromIntegral <$> anyWord32be)),
            (4, TagLong <$> pure "" <*> (fromIntegral <$> anyWord64be)),
            (5, TagFloat <$> pure "" <*> (fromIntegral <$> anyWord32be)),
            (6, TagDouble <$> pure "" <*> (fromIntegral <$> anyWord64be)),
            (7, TagByteArray <$> pure "" <*> parseTagByteArrayNoHeader),
            (8, TagString <$> pure "" <*> (fmap fromIntegral anyWord16be >>= AP.take 
                                            >>= return . BC.unpack)),
            (9, TagList <$> pure "" <*> parseTagListNoHeader),
            (10, TagCompound <$> pure "" <*> parseTagCompoundNoHeader),
            (11, TagIntArray <$> pure "" <*> parseTagIntArrayNoHeader)]

parseTagEnd :: Parser ()
parseTagEnd = void (word8 0)

parseTagHeader :: Word8 -> Parser String
parseTagHeader id = (word8 id >> anyWord16be >>= fmap BC.unpack . AP.take . fromIntegral)
                    <?> ("parseTagHeader " ++ show id)

parseTag :: Parser Tag
parseTag = parseTagByte <|> parseTagShort <|> parseTagInt <|> parseTagLong <|> parseTagFloat
                        <|> parseTagDouble <|> parseTagByteArray <|> parseTagString 
                        <|> parseTagList <|> parseTagCompound <|> parseTagIntArray

parseTagByte :: Parser Tag
parseTagByte = TagByte <$> parseTagHeader 1 <*> (fromIntegral <$> anyWord8)

parseTagShort :: Parser Tag
parseTagShort = TagShort <$> parseTagHeader 2 <*> (fromIntegral <$> anyWord16be)

parseTagInt :: Parser Tag
parseTagInt = TagInt <$> parseTagHeader 3 <*> (fromIntegral <$> anyWord32be)

parseTagLong :: Parser Tag
parseTagLong = TagLong <$> parseTagHeader 4 <*> (fromIntegral <$> anyWord64be)

parseTagFloat :: Parser Tag
parseTagFloat = TagFloat <$> parseTagHeader 5 <*> (fromIntegral <$> anyWord32be)

parseTagDouble :: Parser Tag
parseTagDouble = TagDouble <$> parseTagHeader 6 <*> (fromIntegral <$> anyWord64be)

parseTagByteArray :: Parser Tag
parseTagByteArray = TagByteArray <$> parseTagHeader 7 <*> parseTagByteArrayNoHeader

parseTagString :: Parser Tag
parseTagString = TagString <$> parseTagHeader 8 <*> (fmap fromIntegral anyWord16be >>= AP.take 
                                                                                   >>= return . BC.unpack)

parseTagList :: Parser Tag
parseTagList = TagList <$> parseTagHeader 9 <*> parseTagListNoHeader

parseTagCompound :: Parser Tag
parseTagCompound = TagCompound <$> parseTagHeader 10 <*> parseTagCompoundNoHeader

parseTagIntArray :: Parser Tag
parseTagIntArray = TagIntArray <$> parseTagHeader 11 <*> parseTagIntArrayNoHeader

findTag :: String -> Tag -> Maybe Tag
findTag xs (TagCompound _ map) = M.lookup xs map
findTag xs t = error (show t)

indexTag :: Int -> Tag -> Maybe Tag
indexTag i (TagList _ xs)
  | i >= 0 && i < length xs = Just (xs !! i)
  | otherwise               = Nothing

mapTag :: (Tag -> Maybe a) -> Tag -> Maybe [a]
mapTag f (TagList _ tags) = mapM f tags

parseNBT :: BL.ByteString -> Int -> Result Tag
parseNBT xs compression = parse parseTag nbt
  where nbt = if compression == 0 
                then GZ.decompress xs
                else ZL.decompress xs

main = do
  contents <- BL.readFile "World/level.dat"
  print $ parseNBT contents 0