Decoder.hs

{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE BangPatterns #-}

import Control.Arrow
import Foreign.Ptr
import Data.Word
import Data.Bits
import Foreign.Storable
import Data.ByteString (ByteString)
import qualified Data.ByteString.Internal as ByteString

data Decoding s a = Done !(Ptr Word8) !a
                  | Yield !(Ptr Word8) !Int !s

-- | A decoder specifies how a thing should be parsed.
--   If a Decoder exceeds its initial guess of consumed bytes
--   it yields so the driver can make sure to buffer enough bytes.
data Decoder a = forall s. D !Int !s (Ptr Word8 -> s -> IO (Decoding s a))

occDecoder :: Decoder (Word64, (Word64, Word64))
occDecoder = varint >*< varint >*< varint
{-# INLINE occDecoder #-}

termDecoder :: Decoder (Word64, (ByteString, (Word64, Word64)))
termDecoder =
  varint >*< varintPrefixedByteString >*< varint >*< varint
{-# INLINE termDecoder #-}

data VPB = VPB0
         | VPB1 !Int

varintPrefixedByteString :: Decoder ByteString
varintPrefixedByteString = case varint of
    D szVint vi0 decVarint -> D szVint VPB0 step where

      step op VPB0 = do
        decVi <- decVarint op vi0
        case decVi of
          Done op' n -> return $ Yield op' (fromIntegral n) (VPB1 (fromIntegral n))
          Yield op vi1 _ -> error "The impossible happened"

      step op (VPB1 n) = do
        let bs = ByteString.unsafeCreate n $ \p -> ByteString.memcpy p op n
        return $ Done (op `plusPtr` n) bs

data PD a b c = PD0 !a !b
              | PD1 !b !c

pairDecoder :: Decoder a -> Decoder b -> Decoder (a, b)
pairDecoder (D n1 s1' dec1) (D n2 s2' dec2) =
  D (n1 + n2) (PD0 s1' s2') step
  where
    step op (PD0 s1 s2) = do
      sa <- dec1 op s1
      case sa of
        Done op' a -> do
          sb <- dec2 op s2
          case sb of
            Done op'' b -> return $ Done op'' (a, b)
            Yield op'' n s2' -> return $ Yield op'' n (PD1 s2 a)
        Yield op' n s1' -> return $ Yield op' n (PD0 s1' s2)
    step op (PD1 s2 a) = do
      sb <- dec2 op s2
      case sb of
        Done op' b -> return $ Done op' (a, b)
        Yield op' n s2' -> return $ Yield op' n (PD1 s2' a)
{-# INLINE pairDecoder #-}

(>*<) = pairDecoder
infixr 5 >*<

word64be :: Decoder Word64
word64be = D 8 () $ \op _ -> do x <- peek (castPtr op)
                                return $ Done (op `plusPtr` 8) x
{-# INLINE word64be #-}

varint :: Decoder Word64
varint = D 9 () $ \op _ -> do
  x <- peek op
  loop (op `plusPtr` 1) x (fromIntegral x)
  where loop :: Ptr Word8 -> Word8 -> Word64 -> IO (Decoding () Word64)
        loop !op !n !acc
          | testBit n 7 = do
            x <- peek op
            loop (op `plusPtr` 1) x (fromIntegral x `unsafeShiftL` 7 .|. clearBit acc 7)
          | otherwise = return $ Done op acc
{-# INLINE varint #-}

main :: IO ()
main = undefined