{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell #-}
module Minecraft.Anvil where
import Control.Monad (when, replicateM)
import Control.Monad.State.Strict (execStateT, modify')
import Codec.Compression.Zlib (decompress, compress)
import Data.Bits (shiftR, shiftL, (.&.), (.|.))
import Data.ByteString.Lazy (ByteString)
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as LB
import Data.Data (Data, Typeable)
import Data.List (mapAccumL)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.NBT (NBT)
import Data.Serialize (Serialize(..), Get, Put, decodeLazy, encodeLazy, getLazyByteString, getWord8, getWord32be, putLazyByteString, putWord8, putWord32be)
import Data.Time.Clock.POSIX (POSIXTime)
import Data.Vector (Vector)
import qualified Data.Vector as Vector
import Data.Word (Word8, Word32)
import GHC.Generics
import Pipes.ByteString (fromHandle, toHandle)
import Pipes.Cereal (decodeGet, encodePut)
import Pipes.Parse (runStateT)
import Pipes (Pipe, runEffect, (>->), await, each, yield)
import System.IO (Handle, hSeek, SeekMode(AbsoluteSeek))
{-
regionX :: Double -- ^ chunkX
-> Int
regionX x = floor (x / 32.0)
regionY :: Double -- ^ chunkY
-> Int
regionY y = foor (y / 32.0)
-}
type ChunkX = Int
type ChunkZ = Int
type RegionX = Int
type RegionZ = Int
type ChunkMap = Map (ChunkX, ChunkZ) (ChunkData, POSIXTime)
regionX :: ChunkX
-> RegionX
regionX x = x `shiftR` 5
regionZ :: ChunkZ
-> RegionZ
regionZ z = z `shiftR` 5
chunkIndex :: (ChunkX, ChunkZ)
-> Int
chunkIndex (x, z) = (x `pmod` 32) + ((z `pmod` 32) * 32)
where
pmod n m =
let n' = n `mod` m
in if (n' >= 0)
then n'
else n' + m
type Word24 = Word32
-- | ChunkLocation
--
-- ChunkLocation 0 0 means chunk is not present in the file
data ChunkLocation = ChunkLocation
{ chunkOffset :: Word24 -- ^ number of 4KiB sectors from start of file
, chunkLength :: Word8 -- ^ length of chunk -- units are 4KiB sectors, rounded up
}
deriving (Eq, Ord, Read, Show, Data, Typeable, Generic)
-- | ChunkLocation 0 0 means chunk is not present in the file
emptyChunkLocation :: ChunkLocation
emptyChunkLocation = ChunkLocation 0 0
putChunkLocation :: ChunkLocation -> Put
putChunkLocation (ChunkLocation offset len) =
do putWord8 (fromIntegral ((offset `shiftR` 16) .&. 0xFF))
putWord8 (fromIntegral ((offset `shiftR` 8) .&. 0xFF))
putWord8 (fromIntegral (offset .&. 0xFF))
putWord8 len
getChunkLocation :: Get ChunkLocation
getChunkLocation =
do b2 <- fromIntegral <$> getWord8
b1 <- fromIntegral <$> getWord8
b0 <- fromIntegral <$> getWord8
len <- fromIntegral <$> getWord8
pure (ChunkLocation ((b2 `shiftL` 16) .|. (b1 `shiftL` 8) .|. b0) len)
instance Serialize ChunkLocation where
put = putChunkLocation
get = getChunkLocation
data AnvilHeader = AnvilHeader
{ locations :: Vector ChunkLocation
, timestamps :: Vector POSIXTime
}
deriving (Eq, Ord, Show, Data, Typeable, Generic)
showAnvilHeader :: AnvilHeader -> String
showAnvilHeader ah = show $ AnvilHeader { locations = Vector.filter (/= emptyChunkLocation) (locations ah)
, timestamps = Vector.filter (/= 0) (timestamps ah)
}
-- guessing on the timestamp format a bit
putTimestamp :: POSIXTime -> Put
putTimestamp t = putWord32be (round t)
getTimestamp :: Get POSIXTime
getTimestamp =
do l <- getWord32be
pure $ (realToFrac l)
putAnvilHeader :: AnvilHeader -> Put
putAnvilHeader (AnvilHeader locations timestamps)
| (Vector.length locations /= 1024) || (Vector.length timestamps /= 1024) =
error "putAnvilHeader: locations and timestamps fields must be exactly 1024 entries long."
| otherwise =
do mapM_ putChunkLocation (Vector.toList locations)
mapM_ putTimestamp (Vector.toList timestamps)
getAnvilHeader :: Get AnvilHeader
getAnvilHeader =
do locations <- replicateM 1024 getChunkLocation
timestamps <- replicateM 1024 getTimestamp
pure $ AnvilHeader (Vector.fromList locations) (Vector.fromList timestamps)
instance Serialize AnvilHeader where
put = putAnvilHeader
get = getAnvilHeader
emptyAnvilHeader :: AnvilHeader
emptyAnvilHeader =
AnvilHeader { locations = Vector.replicate 1024 emptyChunkLocation
, timestamps = Vector.replicate 1024 0
}
-- | make an 'AnvilHeader'
--
-- assumes the chunks will be written in the same order as they appear in the list
mkAnvilHeader :: [((ChunkX, ChunkZ), (ChunkData, POSIXTime))]
-> AnvilHeader
mkAnvilHeader chunks =
let timestamps' = map (\(i, (_, t)) -> (chunkIndex i, t)) chunks
locations' = snd $ mapAccumL mkLocation 0x2 chunks -- ^ first chunk is at sector 0x2, after the AnvilHeader
in AnvilHeader { locations = (locations emptyAnvilHeader) Vector.// locations'
, timestamps = (timestamps emptyAnvilHeader) Vector.// timestamps'
}
where
mkLocation :: Word24 -> ((ChunkX, ChunkZ), (ChunkData, POSIXTime)) -> (Word24, (Int, ChunkLocation))
mkLocation offset (chunkPos, (chunkData, _)) =
let paddedSectorLength = ((chunkDataLength chunkData) + 4 + 4095) `div` 4096
offset' = offset + paddedSectorLength
in (offset', (chunkIndex chunkPos, ChunkLocation offset (fromIntegral paddedSectorLength)))
data CompressionType
= GZip -- ^ unused in practice
| Zlib
deriving (Eq, Ord, Read, Show, Data, Typeable, Generic)
-- | ChunkData
--
-- Compressed chunk data
data ChunkData = ChunkData
{ chunkDataLength :: Word32 -- ^ length of chunkData + 1
, chunkDataCompression :: CompressionType -- ^ compression type
, chunkData :: ByteString -- ^ compressed data (length - 1)
}
deriving (Eq, Ord, Read, Show, Data, Typeable, Generic)
putChunkData :: ChunkData -> Put
putChunkData cd =
do putWord32be (chunkDataLength cd)
case chunkDataCompression cd of
GZip -> putWord8 1
Zlib -> putWord8 2
putLazyByteString (chunkData cd)
getChunkData :: Get ChunkData
getChunkData =
do len <- getWord32be
comp <- do w <- getWord8
case w of
1 -> pure GZip
2 -> pure Zlib
_ -> error $ "Unknown compression code in getChunkData: " ++ show (len, w)
bs <- getLazyByteString (fromIntegral (len - 1))
pure $ ChunkData len comp bs
instance Serialize ChunkData where
put = putChunkData
get = getChunkData
-- | read 'ChunkData' from a Seekable 'Handle'
readChunkData :: Handle -> ChunkLocation -> IO (Maybe ChunkData)
readChunkData h chunkLocation
| chunkLocation == emptyChunkLocation = pure Nothing
| otherwise =
do hSeek h AbsoluteSeek (fromIntegral $ ((chunkOffset chunkLocation) * 4096))
(r, _) <- runStateT (decodeGet getChunkData) (fromHandle h)
case r of
(Left err) -> error err
(Right cd) -> pure (Just cd)
-- | write 'ChunkData' to a Seekable 'Handle'
--
-- FIXME: chunks are supposed to be guaranteed less than 1MB
{-
writeChunkData :: Handle -> ChunkData -> IO ()
writeChunkData h chunkData =
do runEffect $ (encodePut (putChunkData chunkData)) >-> (toHandle h)
let padding = 4096 - ((chunkDataLength chunkData + 4) `mod` 4096)
when (padding > 0) (B.hPutStr h (B.replicate (fromIntegral padding) 0))
-}
writeChunkData :: (Monad m) => Pipe ChunkData B.ByteString m ()
writeChunkData = do
chunkData <- await
encodePut (putChunkData chunkData)
let padding = 4096 - ((chunkDataLength chunkData + 4) `mod` 4096)
when (padding > 0) (yield (B.replicate (fromIntegral padding) 0))
{-
do bytes <- execStateT (runEffect $ (encodePut (putChunkData chunkData)) >-> counter >-> (toHandle h)) 0
pure bytes
where
counter =
do bs <- await
modify' $ \i -> i + B.length bs
yield bs
-}
{-
readChunkData :: Handle -> ChunkLocation -> IO (Maybe ChunkData)
readChunkData h chunkLocation
| chunkLocation == emptyChunkLocation = pure Nothing
| otherwise =
do hSeek h AbsoluteSeek (fromIntegral $ ((chunkOffset chunkLocation) * 4096))
(r, _) <- runStateT (decodeGet getChunkData) (fromHandle h)
case r of
(Left err) -> error err
(Right cd) -> pure (Just cd)
-}
decompressChunkData :: ChunkData -> Either String NBT
decompressChunkData cd
| chunkDataCompression cd == Zlib =
decodeLazy (decompress (chunkData cd))
| otherwise = error $ "decompressChunkData not implemented for " ++ show (chunkDataCompression cd)
-- | NBT needs to be a Chunk
compressChunkData :: NBT -> ChunkData
compressChunkData nbt =
let d = compress (encodeLazy nbt)
in ChunkData { chunkDataLength = 1 + fromIntegral (LB.length d)
, chunkDataCompression = Zlib
, chunkData = d
}
writeChunkMap :: Handle -> ChunkMap -> IO ()
writeChunkMap h chunkMap =
do hSeek h AbsoluteSeek 0
let chunks = Map.toAscList chunkMap
runEffect $ (encodePut (putAnvilHeader (mkAnvilHeader chunks))) >-> (toHandle h)
runEffect $ (each $ map (fst . snd) chunks) >-> writeChunkData >-> (toHandle h)
-- runEffect $ (encodePut (putAnvilHeader emptyAnvilHeader)) >-> (toHandle h)
-- [(ChunkPos, ChunkData)] -> [(ChunkPos, Length)]
-- chunks' <- mapM (\(chunkPos, (cd, modTime)) -> do i <- writeChunkData h cd ; pure (chunkPos, (i, modTime))) chunks
pure ()