test.hs

module Main where

import Control.Applicative
import Control.Monad
import Data.Binary.Get
import Data.Bits
import Data.Word
import qualified Data.ByteString.Lazy as BS

import Debug.Trace

data ItemType =
      BT_STOP
    | BT_STOP_BASE
    | BT_BOOL
    | BT_UINT8
    | BT_UINT16
    | BT_UINT32
    | BT_UINT64
    | BT_FLOAT
    | BT_DOUBLE
    | BT_STRING
    | BT_STRUCT
    | BT_LIST
    | BT_SET
    | BT_MAP
    | BT_INT8
    | BT_INT16
    | BT_INT32
    | BT_INT64
    | BT_WSTRING
    deriving (Show, Enum, Eq)

getVarInt :: Get Int
getVarInt = step 0
    where
    step :: Int -> Get Int
    step n | n > 4 = fail "VarInt: sequence too long"
    step n = do
        b <- fromIntegral <$> getWord8
        rest <- if b `testBit` 7 then step (n + 1)  else return (0 :: Int)
        return $ (b `clearBit` 7) .|. (rest `shiftL` 7)

startStruct :: Get ()
startStruct = do
    b <- bytesRead
    traceShowM ("begin struct", b)

endStruct :: Get ()
endStruct = do
    b <- bytesRead
    traceShowM ("end struct", b)

endBase :: Get ()
endBase = do
    b <- bytesRead
    traceShowM ("end base", b)

item :: ItemType -> Get ()
item BT_DOUBLE = void getWord64le
item BT_FLOAT = void getWord32le
item BT_STRING = do
    n <- getVarInt
    skip n
item BT_WSTRING = do
    n <- getVarInt
    skip (n * 2)
item BT_UINT64 = void getWord64le
item BT_UINT16 = void getWord16le
item BT_UINT32 = void getWord32le
item BT_UINT8 = void getWord8
item BT_INT8 = void getWord8
item BT_INT16 = void getWord16le
item BT_INT32 = void getWord32le
item BT_INT64 = void getWord64le
item BT_BOOL = void getWord8
item BT_STRUCT = getStruct
item BT_LIST = do
    tag <- getWord8
    let itemtype = toEnum (fromIntegral tag)
    n <- getVarInt
    forM_ [0 .. n - 1] $ \i -> do
        traceShowM ("item", i)
        item itemtype
item BT_SET = do
    tag <- getWord8
    let itemtype = toEnum (fromIntegral tag)
    n <- getVarInt
    forM_ [0 .. n - 1] $ \i -> do
        traceShowM ("item", i)
        item itemtype
item BT_MAP = do
    ktag <- getWord8
    vtag <- getWord8
    let ktype = toEnum (fromIntegral ktag)
    let vtype = toEnum (fromIntegral vtag)
    n <- getVarInt
    forM_ [0 .. n - 1] $ \i -> do
        traceShowM ("item", i)
        item ktype
        item vtype

field :: ItemType -> Word16 -> Get ()
field t n = do
    traceShowM ("field", t, n)
    void $ item t

getStruct :: Get ()
getStruct = do
    startStruct
    loop
    where
    loop = do
        tag <- getWord8
        case toEnum (fromIntegral tag) of
            BT_STOP -> do
                endStruct
                return ()
            BT_STOP_BASE -> do
                endBase
                loop
            t -> do
                n <- getWord16le
                field t n
                loop

main :: IO()
main = do
    b <- BS.readFile "/home/blaze/bond/test/compat/data/compat.fast.dat"
    print $ runGet getStruct b