content: --
  -- Licensed to the Apache Software Foundation (ASF) under one
  -- or more contributor license agreements. See the NOTICE file
  -- distributed with this work for additional information
  -- regarding copyright ownership. The ASF licenses this file
  -- to you under the Apache License, Version 2.0 (the
  -- "License"); you may not use this file except in compliance
  -- with the License. You may obtain a copy of the License at
  --
  --   http://www.apache.org/licenses/LICENSE-2.0
  --
  -- Unless required by applicable law or agreed to in writing,
  -- software distributed under the License is distributed on an
  -- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  -- KIND, either express or implied. See the License for the
  -- specific language governing permissions and limitations
  -- under the License.
  --
  
  {-# LANGUAGE CPP #-}
  {-# LANGUAGE ExistentialQuantification #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  
  module Thrift.Protocol.Compact
      ( module Thrift.Protocol
      , CompactProtocol(..)
      , parseVarint
      , buildVarint
      ) where
  
  import Control.Applicative
  import Control.Monad
  import Data.Attoparsec.ByteString as P
  import Data.Attoparsec.ByteString.Lazy as LP
  import Data.Bits
  import Data.ByteString.Lazy.Builder as B
  import Data.Int
  import Data.List as List
  import Data.Monoid
  import Data.Word
  import Data.Text.Lazy.Encoding ( decodeUtf8, encodeUtf8 )
  
  import Thrift.Protocol
  import Thrift.Transport
  import Thrift.Types
  
  import qualified Data.ByteString as BS
  import qualified Data.ByteString.Lazy as LBS
  import qualified Data.HashMap.Strict as Map
  import qualified Data.Text.Lazy as LT
  
  -- | the Compact Protocol implements the standard Thrift 'TCompactProcotol'
  -- which is similar to the 'TBinaryProtocol', but takes less space on the wire.
  -- Integral types are encoded using as varints.
  data CompactProtocol a = CompactProtocol a
                           -- ^ Constuct a 'CompactProtocol' with a 'Transport'
  
  protocolID, version, versionMask, typeMask, typeBits :: Word8
  protocolID  = 0x82 -- 1000 0010
  version     = 0x01
  versionMask = 0x1f -- 0001 1111
  typeMask    = 0xe0 -- 1110 0000
  typeBits    = 0x07 -- 0000 0111
  typeShiftAmount :: Int
  typeShiftAmount = 5
  
  getTransport :: Transport t => CompactProtocol t -> t
  getTransport (CompactProtocol t) = t
  
  instance Transport t => Protocol (CompactProtocol t) where
      readByte p = tReadAll (getTransport p) 1
      writeMessage p (n, t, s) f = do
        tWrite (getTransport p) messageBegin
        f
        tFlush $ getTransport p
        where
          messageBegin = toLazyByteString $
            B.word8 protocolID <>
            B.word8 ((version .&. versionMask) .|.
                    (((fromIntegral $ fromEnum t) `shiftL`
                    typeShiftAmount) .&. typeMask)) <>
            buildVarint (i32ToZigZag s) <>
            buildCompactValue (TString $ encodeUtf8 n)
  
      readMessage p f = readMessageBegin >>= f
        where
          readMessageBegin = runParser p $ do
            pid <- fromIntegral <$> P.anyWord8
            when (pid /= protocolID) $ error "Bad Protocol ID"
            w <- fromIntegral <$> P.anyWord8
            let ver = w .&. versionMask
            when (ver /= version) $ error "Bad Protocol version"
            let typ = (w `shiftR` typeShiftAmount) .&. typeBits
            seqId <- parseVarint zigZagToI32
            TString name <- parseCompactValue T_STRING
            return (decodeUtf8 name, toEnum $ fromIntegral $ typ, seqId)
  
      writeVal p = tWrite (getTransport p) . toLazyByteString . buildCompactValue
      readVal p ty = runParser p $ parseCompactValue ty
  
  instance Transport t => StatelessProtocol (CompactProtocol t) where
      serializeVal _ = toLazyByteString . buildCompactValue
      deserializeVal _ ty bs =
        case LP.eitherResult $ LP.parse (parseCompactValue ty) bs of
          Left s -> error s
          Right val -> val
  
  -- | Writing Functions
  buildCompactValue :: ThriftVal -> Builder
  buildCompactValue (TStruct fields) = buildCompactStruct fields
  buildCompactValue (TMap kt vt entries) =
    let len = fromIntegral $ length entries :: Word32 in
    if len == 0
    then B.word8 0x00
    else buildVarint len <>
         B.word8 (fromTType kt `shiftL` 4 .|. fromTType vt) <>
         buildCompactMap entries
  buildCompactValue (TList ty entries) =
    let len = length entries in
    (if len < 15
     then B.word8 $ (fromIntegral len `shiftL` 4) .|. fromTType ty
     else B.word8 (0xF0 .|. fromTType ty) <>
          buildVarint (fromIntegral len :: Word32)) <>
    buildCompactList entries
  buildCompactValue (TSet ty entries) = buildCompactValue (TList ty entries)
  buildCompactValue (TBool b) =
    B.word8 $ toEnum $ if b then 1 else 0
  buildCompactValue (TByte b) = int8 b
  buildCompactValue (TI16 i) = buildVarint $ i16ToZigZag i
  buildCompactValue (TI32 i) = buildVarint $ i32ToZigZag i
  buildCompactValue (TI64 i) = buildVarint $ i64ToZigZag i
  buildCompactValue (TDouble d) = doubleLE d
  buildCompactValue (TString s) = buildVarint len <> lazyByteString s
    where
      len = fromIntegral (LBS.length s) :: Word32
  buildCompactValue (TBinary s) = buildCompactValue (TString s)
  
  buildCompactStruct :: Map.HashMap Int16 (LT.Text, ThriftVal) -> Builder
  buildCompactStruct = flip (loop 0) mempty . Map.toList
    where
      loop _ [] acc = acc <> B.word8 (fromTType T_STOP)
      loop lastId ((fid, (_,val)) : fields) acc = loop fid fields $ acc <>
        (if fid > lastId && fid - lastId <= 15
         then B.word8 $ fromIntegral ((fid - lastId) `shiftL` 4) .|. typeOf val
         else B.word8 (typeOf val) <> buildVarint (i16ToZigZag fid)) <>
        (if typeOf val > 0x02 -- Not a T_BOOL
         then buildCompactValue val
         else mempty) -- T_BOOLs are encoded in the type
  buildCompactMap :: [(ThriftVal, ThriftVal)] -> Builder
  buildCompactMap = foldl combine mempty
    where
      combine s (key, val) = buildCompactValue key <> buildCompactValue val <> s
  
  buildCompactList :: [ThriftVal] -> Builder
  buildCompactList = foldr (mappend . buildCompactValue) mempty
  
  -- | Reading Functions
  parseCompactValue :: ThriftType -> Parser ThriftVal
  parseCompactValue (T_STRUCT tmap) = TStruct <$> parseCompactStruct tmap
  parseCompactValue (T_MAP kt' vt') = do
    n <- parseVarint id
    if n == 0
      then return $ TMap kt' vt' []
      else do
      w <- P.anyWord8
      let kt = typeFrom $ w `shiftR` 4
          vt = typeFrom $ w .&. 0x0F
      TMap kt vt <$> parseCompactMap kt vt n
  parseCompactValue (T_LIST ty) = TList ty <$> parseCompactList
  parseCompactValue (T_SET ty) = TSet ty <$> parseCompactList
  parseCompactValue T_BOOL = TBool . (/=0) <$> P.anyWord8
  parseCompactValue T_BYTE = TByte . fromIntegral <$> P.anyWord8
  parseCompactValue T_I16 = TI16 <$> parseVarint zigZagToI16
  parseCompactValue T_I32 = TI32 <$> parseVarint zigZagToI32
  parseCompactValue T_I64 = TI64 <$> parseVarint zigZagToI64
  parseCompactValue T_DOUBLE = TDouble . bsToDoubleLE <$> P.take 8
  parseCompactValue T_STRING = parseCompactString TString
  parseCompactValue T_BINARY = parseCompactString TBinary
  parseCompactValue ty = error $ "Cannot read value of type " ++ show ty
  
  parseCompactString ty = do
    len :: Word32 <- parseVarint id
    ty . LBS.fromStrict <$> P.take (fromIntegral len)
  
  parseCompactStruct :: TypeMap -> Parser (Map.HashMap Int16 (LT.Text, ThriftVal))
  parseCompactStruct tmap = Map.fromList <$> parseFields 0
    where
      parseFields :: Int16 -> Parser [(Int16, (LT.Text, ThriftVal))]
      parseFields lastId = do
        w <- P.anyWord8
        if w == 0x00
          then return []
          else do
            let ty = typeFrom (w .&. 0x0F)
                modifier = (w .&. 0xF0) `shiftR` 4
            fid <- if modifier /= 0
                   then return (lastId + fromIntegral modifier)
                   else parseVarint zigZagToI16
            val <- if ty == T_BOOL
                   then return (TBool $ (w .&. 0x0F) == 0x01)
                   else case (ty, Map.lookup fid tmap) of
                          (T_STRING, Just (_, T_BINARY)) -> parseCompactValue T_BINARY
                          _ -> parseCompactValue ty
            ((fid, (LT.empty, val)) : ) <$> parseFields fid
  
  parseCompactMap :: ThriftType -> ThriftType -> Int32 ->
                     Parser [(ThriftVal, ThriftVal)]
  parseCompactMap kt vt n | n <= 0 = return []
                          | otherwise = do
    k <- parseCompactValue kt
    v <- parseCompactValue vt
    ((k,v) :) <$> parseCompactMap kt vt (n-1)
  
  parseCompactList :: Parser [ThriftVal]
  parseCompactList = do
    w <- P.anyWord8
    let ty = typeFrom $ w .&. 0x0F
        lsize = w `shiftR` 4
    size <- if lsize == 0xF
            then parseVarint id
            else return $ fromIntegral lsize
    loop ty size
    where
      loop :: ThriftType -> Int32 -> Parser [ThriftVal]
      loop ty n | n <= 0 = return []
                | otherwise = liftM2 (:) (parseCompactValue ty)
                              (loop ty (n-1))
  
  -- Signed numbers must be converted to "Zig Zag" format before they can be
  -- serialized in the Varint format
  i16ToZigZag :: Int16 -> Word16
  i16ToZigZag n = fromIntegral $ (n `shiftL` 1) `xor` (n `shiftR` 15)
  
  zigZagToI16 :: Word16 -> Int16
  zigZagToI16 n = fromIntegral $ (n `shiftR` 1) `xor` negate (n .&. 0x1)
  
  i32ToZigZag :: Int32 -> Word32
  i32ToZigZag n = fromIntegral $ (n `shiftL` 1) `xor` (n `shiftR` 31)
  
  zigZagToI32 :: Word32 -> Int32
  zigZagToI32 n = fromIntegral $ (n `shiftR` 1) `xor` negate (n .&. 0x1)
  
  i64ToZigZag :: Int64 -> Word64
  i64ToZigZag n = fromIntegral $ (n `shiftL` 1) `xor` (n `shiftR` 63)
  
  zigZagToI64 :: Word64 -> Int64
  zigZagToI64 n = fromIntegral $ (n `shiftR` 1) `xor` negate (n .&. 0x1)
  
  buildVarint :: (Bits a, Integral a)  => a -> Builder
  buildVarint n | n .&. complement 0x7F == 0 = B.word8 $ fromIntegral n
                | otherwise = B.word8 (0x80 .|. (fromIntegral n .&. 0x7F)) <>
                              buildVarint (n `shiftR` 7)
  
  parseVarint :: (Bits a, Integral a, Ord a) => (a -> b) -> Parser b
  parseVarint fromZigZag = do
    bytestemp <- BS.unpack <$> P.takeTill (not . flip testBit 7)
    lsb <- P.anyWord8
    let bytes = lsb : List.reverse bytestemp
    return $ fromZigZag $ List.foldl' combine 0x00 bytes
    where combine a b = (a `shiftL` 7) .|. (fromIntegral b .&. 0x7f)
  
  -- | Compute the Compact Type
  fromTType :: ThriftType -> Word8
  fromTType ty = case ty of
    T_STOP -> 0x00
    T_BOOL -> 0x01
    T_BYTE -> 0x03
    T_I16 -> 0x04
    T_I32 -> 0x05
    T_I64 -> 0x06
    T_DOUBLE -> 0x07
    T_STRING -> 0x08
    T_BINARY -> 0x08
    T_LIST{} -> 0x09
    T_SET{} -> 0x0A
    T_MAP{} -> 0x0B
    T_STRUCT{} -> 0x0C
    T_VOID -> error "No Compact type for T_VOID"
  
  typeOf :: ThriftVal -> Word8
  typeOf v = case v of
    TBool True -> 0x01
    TBool False -> 0x02
    TByte _ -> 0x03
    TI16 _ -> 0x04
    TI32 _ -> 0x05
    TI64 _ -> 0x06
    TDouble _ -> 0x07
    TString _ -> 0x08
    TBinary _ -> 0x08
    TList{} -> 0x09
    TSet{} -> 0x0A
    TMap{} -> 0x0B
    TStruct{} -> 0x0C
  
  typeFrom :: Word8 -> ThriftType
  typeFrom w = case w of
    0x01 -> T_BOOL
    0x02 -> T_BOOL
    0x03 -> T_BYTE
    0x04 -> T_I16
    0x05 -> T_I32
    0x06 -> T_I64
    0x07 -> T_DOUBLE
    0x08 -> T_STRING
    0x09 -> T_LIST T_VOID
    0x0A -> T_SET T_VOID
    0x0B -> T_MAP T_VOID T_VOID
    0x0C -> T_STRUCT Map.empty
    n -> error $ "typeFrom: " ++ show n ++ " is not a compact type"
  