diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,26 @@ +Copyright 2026 Ian Duncan + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, + this list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its contributors + may be used to endorse or promote products derived from this software + without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/bench/Bench.hs b/bench/Bench.hs new file mode 100644 --- /dev/null +++ b/bench/Bench.hs @@ -0,0 +1,83 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE OverloadedStrings #-} + +{- | Microbench for wireform-thrift's encode + decode hot paths +across the binary and compact wire protocols. + +Test fixtures: a small @Person@ record and a 100-element vector of +the same. Same shape as every other per-package bench in the +monorepo so cross-format comparisons in a unified dashboard stay +comparable. +-} +module Main (main) where + +import Control.DeepSeq (NFData) +import Criterion.Main +import Data.ByteString (ByteString) +import Data.Text (Text) +import Data.Text qualified as T +import Data.Vector (Vector) +import Data.Vector qualified as V +import GHC.Generics (Generic) +import Thrift.Class + + +data Person = Person + { personName :: !Text + , personAge :: !Int + , personEmail :: !Text + } + deriving stock (Show, Eq, Generic) + deriving anyclass (ToThrift, FromThrift, NFData) + + +small :: Person +small = Person "Alice" 30 "alice@example.com" + + +medium :: Vector Person +medium = + V.fromList + [ Person + (T.pack ("user-" <> show i)) + (20 + i `mod` 50) + (T.pack ("user" <> show i <> "@example.com")) + | i <- [1 .. 100 :: Int] + ] + + +main :: IO () +main = + defaultMain + [ bgroup + "binary" + [ bgroup + "encode" + [ bench "Person" $ nf encodeThriftBinary small + , bench "[Person] x 100" $ nf encodeThriftBinary medium + ] + , bgroup + "decode" + [ env (pure (encodeThriftBinary small)) $ \bs -> + bench "Person" $ nf (decodeThriftBinary :: ByteString -> Either String Person) bs + , env (pure (encodeThriftBinary medium)) $ \bs -> + bench "[Person] x 100" $ nf (decodeThriftBinary :: ByteString -> Either String (Vector Person)) bs + ] + ] + , bgroup + "compact" + [ bgroup + "encode" + [ bench "Person" $ nf encodeThriftCompact small + , bench "[Person] x 100" $ nf encodeThriftCompact medium + ] + , bgroup + "decode" + [ env (pure (encodeThriftCompact small)) $ \bs -> + bench "Person" $ nf (decodeThriftCompact :: ByteString -> Either String Person) bs + , env (pure (encodeThriftCompact medium)) $ \bs -> + bench "[Person] x 100" $ nf (decodeThriftCompact :: ByteString -> Either String (Vector Person)) bs + ] + ] + ] diff --git a/scripts/protocol_golden.py b/scripts/protocol_golden.py new file mode 100644 --- /dev/null +++ b/scripts/protocol_golden.py @@ -0,0 +1,129 @@ +#!/usr/bin/env python3 +"""Regenerate/check wire goldens using Apache's Python thrift package (0.25.0). + +Run from the package directory: python scripts/protocol_golden.py --check +Use --write after intentionally changing fixtures. The Haskell suite consumes +these independent bytes without requiring Python at test time. +""" +import argparse +from importlib.metadata import version +from pathlib import Path +from uuid import UUID + +from thrift.Thrift import TType +from thrift.protocol import TBinaryProtocol, TCompactProtocol +from thrift.transport import TTransport + + +def struct(fields): + return TType.STRUCT, fields + + +def sequence(kind, element_type, values): + return kind, (element_type, values) + + +def fixtures(): + yield "scalars", struct([ + (1, (TType.BOOL, False)), (2, (TType.BOOL, True)), + (3, (TType.BYTE, -128)), (4, (TType.I16, -32768)), + (5, (TType.I32, -2147483648)), (6, (TType.I64, -9223372036854775808)), + (7, (TType.DOUBLE, -123.5)), (8, (TType.STRING, "é雪\0")), + (9, (TType.STRING, bytes([0, 255, 128, 1]))), + (10, (TType.UUID, UUID(bytes=bytes(range(16))))), + ]) + yield "field-ids", struct([ + (32767, (TType.I32, 1)), (-32768, (TType.I32, 2)), (0, (TType.BOOL, False)), + ]) + yield "container-booleans", struct([ + (1, sequence(TType.LIST, TType.BOOL, [False, True, False])), + (2, sequence(TType.SET, TType.BOOL, [False, True])), + (3, (TType.MAP, (TType.BOOL, TType.BOOL, [(False, True), (True, False)]))), + ]) + for count in (0, 14, 15, 128): + yield f"list-{count}", struct([ + (1, sequence(TType.LIST, TType.I32, list(range(-64, count - 64)))) + ]) + yield "nested", struct([ + (1, sequence(TType.LIST, TType.STRUCT, [ + [(1, (TType.STRING, "first")), (2, (TType.BOOL, False))], + [(1, (TType.STRING, "雪")), (2, (TType.BOOL, True))], + ])), + (2, (TType.MAP, (TType.STRING, TType.LIST, [ + ("a", (TType.BOOL, [False, True])), ("é", (TType.BOOL, [])), + ]))), + (3, (TType.MAP, (TType.STOP, TType.STOP, []))), + ]) + + +def write(protocol, typ, value): + if typ == TType.STRUCT: + protocol.writeStructBegin("fixture") + for fid, (field_type, field_value) in sorted(value): + protocol.writeFieldBegin("field", field_type, fid) + write(protocol, field_type, field_value) + protocol.writeFieldEnd() + protocol.writeFieldStop() + protocol.writeStructEnd() + elif typ in (TType.LIST, TType.SET): + element_type, values = value + begin = protocol.writeListBegin if typ == TType.LIST else protocol.writeSetBegin + end = protocol.writeListEnd if typ == TType.LIST else protocol.writeSetEnd + begin(element_type, len(values)) + for element in values: + write(protocol, element_type, element) + end() + elif typ == TType.MAP: + kt, vt, entries = value + protocol.writeMapBegin(kt, vt, len(entries)) + for key, element in entries: + write(protocol, kt, key) + write(protocol, vt, element) + protocol.writeMapEnd() + elif typ == TType.STRING: + if isinstance(value, str): + protocol.writeString(value) + else: + protocol.writeBinary(value) + else: + method = { + TType.BOOL: protocol.writeBool, TType.BYTE: protocol.writeByte, + TType.I16: protocol.writeI16, TType.I32: protocol.writeI32, + TType.I64: protocol.writeI64, TType.DOUBLE: protocol.writeDouble, + TType.UUID: protocol.writeUuid, + }[typ] + method(value) + + +def rendered(): + lines = ["# Apache Python thrift 0.25.0; regenerate with scripts/protocol_golden.py --write", + "# fixture\tbinary hex\tcompact hex"] + for name, (typ, value) in fixtures(): + encodings = [] + for implementation in (TBinaryProtocol.TBinaryProtocol, TCompactProtocol.TCompactProtocol): + transport = TTransport.TMemoryBuffer() + write(implementation(transport), typ, value) + encodings.append(transport.getvalue().hex()) + lines.append("\t".join([name, *encodings])) + return "\n".join(lines) + "\n" + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + mode = parser.add_mutually_exclusive_group(required=True) + mode.add_argument("--check", action="store_true") + mode.add_argument("--write", action="store_true") + args = parser.parse_args() + path = Path(__file__).resolve().parents[1] / "test/fixtures/protocol-encodings.tsv" + if version("thrift") != "0.25.0": + raise SystemExit("Install Apache Python thrift==0.25.0 to reproduce these fixtures") + expected = rendered() + if args.write: + path.write_text(expected) + elif path.read_text() != expected: + raise SystemExit("Apache Thrift protocol fixtures differ; inspect before regenerating") + print("PASS: 8 fixtures x Binary/Compact match Apache Thrift") + + +if __name__ == "__main__": + main() diff --git a/src/Thrift.hs b/src/Thrift.hs new file mode 100644 --- /dev/null +++ b/src/Thrift.hs @@ -0,0 +1,43 @@ +{- | Convenience re-export module for the common wireform-thrift surface. + +Import this module to get the Thrift binary codec — encoding, decoding, the +wire types, the dynamic 'Value' model, the 'ToThrift'/'FromThrift' typeclasses, +and the deriver — in one go. + +@ +import Thrift +@ + +Specialized surfaces that are deliberately *not* re-exported here — opt into +them directly when you need them: "Thrift.Schema" (IDL schema model), +"Thrift.Parser" (IDL parser), "Thrift.CodeGen" (IDL code generation), +"Thrift.Message" (message framing), "Thrift.Transport" (transport layer), +"Thrift.Registry" (type registry), "Thrift.JSON" (Thrift ↔ JSON mapping), +and "Thrift.QQ" (quasiquoter). +-} +module Thrift ( + -- * Encoding + module Thrift.Encode, + module Thrift.Encoding, + + -- * Decoding + module Thrift.Decode, + + -- * Wire format + module Thrift.Wire, + + -- * Dynamic values + module Thrift.Value, + + -- * Typeclass-based codec + deriving + module Thrift.Class, + module Thrift.Derive, +) where + +import Thrift.Class +import Thrift.Decode +import Thrift.Derive +import Thrift.Encode +import Thrift.Encoding +import Thrift.Value +import Thrift.Wire diff --git a/src/Thrift/Class.hs b/src/Thrift/Class.hs new file mode 100644 --- /dev/null +++ b/src/Thrift/Class.hs @@ -0,0 +1,835 @@ +{-# LANGUAGE DefaultSignatures #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +{- | Typeclass-based Thrift serialization with GHC Generics support. + +Provides 'ToThrift' and 'FromThrift' typeclasses for converting Haskell +values to\/from 'Thrift.Value.Value'. Records are encoded as Thrift +structs with field IDs assigned sequentially (1, 2, 3, ...). +Derive instances via @DeriveGeneric@. + +@ +{\-\# LANGUAGE DeriveGeneric \#-\} +import GHC.Generics (Generic) +import Thrift.Class + +data LogEntry = LogEntry { level :: Text, message :: Text, code :: Int } + deriving (Generic) +instance ToThrift LogEntry +instance FromThrift LogEntry + +let bytes = encodeThriftBinary (LogEntry \"ERROR\" \"disk full\" 507) +let Right entry = decodeThriftBinary bytes :: Either String LogEntry +@ +-} +module Thrift.Class ( + ToThrift (..), + FromThrift (..), + encodeThriftBinary, + encodeThriftBinaryDirect, + decodeThriftBinary, + encodeThriftCompact, + encodeThriftCompactDirect, + decodeThriftCompact, + genericToEncoding, + GToThrift (..), + GFromThrift (..), +) where + +import Data.ByteString (ByteString) +import Data.ByteString.Lazy qualified as BSL +import Data.Functor.Compose (Compose (..)) +import Data.Functor.Const (Const (..)) +import Data.Functor.Identity (Identity (..)) +import Data.Functor.Product qualified as FProduct +import Data.Functor.Sum qualified as FSum +import Data.HashMap.Strict (HashMap) +import Data.HashMap.Strict qualified as HM +import Data.HashSet (HashSet) +import Data.HashSet qualified as HS +import Data.Hashable (Hashable) +import Data.Int (Int16, Int32, Int64, Int8) +import Data.IntMap.Strict (IntMap) +import Data.IntMap.Strict qualified as IntMap +import Data.IntSet (IntSet) +import Data.IntSet qualified as IntSet +import Data.List.NonEmpty (NonEmpty (..)) +import Data.List.NonEmpty qualified as NE +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.Monoid qualified as Mon +import Data.Ord (Down (..)) +import Data.Ratio (Ratio, denominator, numerator, (%)) +import Data.Semigroup qualified as Semi +import Data.Sequence (Seq) +import Data.Sequence qualified as Seq +import Data.Set (Set) +import Data.Set qualified as Set +import Data.Text (Text) +import Data.Text qualified as T +import Data.Text.Lazy qualified as TL +import Data.Vector (Vector) +import Data.Vector qualified as V +import Data.Version (Version, makeVersion, versionBranch) +import Data.Word (Word16, Word32, Word64, Word8) +import GHC.Generics +import Numeric.Natural (Natural) +import Thrift.Decode qualified as TD +import Thrift.Encode qualified as TE +import Thrift.Encoding (Encoding) +import Thrift.Encoding qualified as Enc +import Thrift.Value qualified as TV +import Thrift.Wire (ThriftType (..)) + + +class ToThrift a where + toThrift :: a -> TV.Value + default toThrift :: (Generic a, GToThrift (Rep a)) => a -> TV.Value + toThrift = gToThrift . from + + + {- | aeson-style direct encoder. Thrift's binary and compact wire + formats need protocol commitment before bytes flow, so + 'Encoding' wraps a fully-built 'TV.Value' and routes through + 'TE.encodeBinary' / 'TE.encodeCompact' at run time. + -} + toEncoding :: a -> Encoding + toEncoding = Enc.value . toThrift + + +class FromThrift a where + fromThrift :: TV.Value -> Either String a + default fromThrift :: (Generic a, GFromThrift (Rep a)) => TV.Value -> Either String a + fromThrift v = to <$> gFromThrift v + + +encodeThriftBinary :: ToThrift a => a -> ByteString +encodeThriftBinary = TE.encodeBinary . toThrift + + +-- | Encode binary protocol via 'toEncoding'. +encodeThriftBinaryDirect :: ToThrift a => a -> ByteString +encodeThriftBinaryDirect = Enc.encodingToBinaryByteString . toEncoding + + +decodeThriftBinary :: FromThrift a => ByteString -> Either String a +decodeThriftBinary bs = TD.decodeBinary bs >>= fromThrift + + +encodeThriftCompact :: ToThrift a => a -> ByteString +encodeThriftCompact = TE.encodeCompact . toThrift + + +-- | Encode compact protocol via 'toEncoding'. +encodeThriftCompactDirect :: ToThrift a => a -> ByteString +encodeThriftCompactDirect = Enc.encodingToCompactByteString . toEncoding + + +genericToEncoding :: (Generic a, GToThrift (Rep a)) => a -> Encoding +genericToEncoding = Enc.value . gToThrift . from + + +decodeThriftCompact :: FromThrift a => ByteString -> Either String a +decodeThriftCompact bs = TD.decodeCompact bs >>= fromThrift + + +instance ToThrift Bool where + toThrift = TV.Bool + + +instance FromThrift Bool where + fromThrift (TV.Bool b) = Right b + fromThrift _ = Left "FromThrift Bool: expected Bool" + + +instance ToThrift Int where + toThrift n = TV.I64 (fromIntegral n) + + +instance FromThrift Int where + fromThrift (TV.I64 n) = Right (fromIntegral n) + fromThrift (TV.I32 n) = Right (fromIntegral n) + fromThrift (TV.I16 n) = Right (fromIntegral n) + fromThrift (TV.Byte n) = Right (fromIntegral n) + fromThrift _ = Left "FromThrift Int: expected integer type" + + +instance ToThrift Int8 where + toThrift = TV.Byte + + +instance FromThrift Int8 where + fromThrift (TV.Byte n) = Right n + fromThrift (TV.I16 n) = Right (fromIntegral n) + fromThrift (TV.I32 n) = Right (fromIntegral n) + fromThrift (TV.I64 n) = Right (fromIntegral n) + fromThrift _ = Left "FromThrift Int8: expected integer type" + + +instance ToThrift Int16 where + toThrift = TV.I16 + + +instance FromThrift Int16 where + fromThrift (TV.I16 n) = Right n + fromThrift (TV.Byte n) = Right (fromIntegral n) + fromThrift (TV.I32 n) = Right (fromIntegral n) + fromThrift (TV.I64 n) = Right (fromIntegral n) + fromThrift _ = Left "FromThrift Int16: expected integer type" + + +instance ToThrift Int32 where + toThrift = TV.I32 + + +instance FromThrift Int32 where + fromThrift (TV.I32 n) = Right n + fromThrift (TV.Byte n) = Right (fromIntegral n) + fromThrift (TV.I16 n) = Right (fromIntegral n) + fromThrift (TV.I64 n) = Right (fromIntegral n) + fromThrift _ = Left "FromThrift Int32: expected integer type" + + +instance ToThrift Int64 where + toThrift = TV.I64 + + +instance FromThrift Int64 where + fromThrift (TV.I64 n) = Right n + fromThrift (TV.I32 n) = Right (fromIntegral n) + fromThrift (TV.I16 n) = Right (fromIntegral n) + fromThrift (TV.Byte n) = Right (fromIntegral n) + fromThrift _ = Left "FromThrift Int64: expected integer type" + + +instance ToThrift Word where + toThrift n = TV.I64 (fromIntegral n) + + +instance FromThrift Word where + fromThrift (TV.I64 n) = Right (fromIntegral n) + fromThrift (TV.I32 n) = Right (fromIntegral n) + fromThrift (TV.I16 n) = Right (fromIntegral n) + fromThrift (TV.Byte n) = Right (fromIntegral n) + fromThrift _ = Left "FromThrift Word: expected integer type" + + +instance ToThrift Word8 where + toThrift n = TV.Byte (fromIntegral n) + + +instance FromThrift Word8 where + fromThrift (TV.Byte n) = Right (fromIntegral n) + fromThrift (TV.I16 n) = Right (fromIntegral n) + fromThrift (TV.I32 n) = Right (fromIntegral n) + fromThrift (TV.I64 n) = Right (fromIntegral n) + fromThrift _ = Left "FromThrift Word8: expected integer type" + + +instance ToThrift Word16 where + toThrift n = TV.I16 (fromIntegral n) + + +instance FromThrift Word16 where + fromThrift (TV.I16 n) = Right (fromIntegral n) + fromThrift (TV.Byte n) = Right (fromIntegral n) + fromThrift (TV.I32 n) = Right (fromIntegral n) + fromThrift (TV.I64 n) = Right (fromIntegral n) + fromThrift _ = Left "FromThrift Word16: expected integer type" + + +instance ToThrift Word32 where + toThrift n = TV.I32 (fromIntegral n) + + +instance FromThrift Word32 where + fromThrift (TV.I32 n) = Right (fromIntegral n) + fromThrift (TV.Byte n) = Right (fromIntegral n) + fromThrift (TV.I16 n) = Right (fromIntegral n) + fromThrift (TV.I64 n) = Right (fromIntegral n) + fromThrift _ = Left "FromThrift Word32: expected integer type" + + +instance ToThrift Word64 where + toThrift n = TV.I64 (fromIntegral n) + + +instance FromThrift Word64 where + fromThrift (TV.I64 n) = Right (fromIntegral n) + fromThrift (TV.I32 n) = Right (fromIntegral n) + fromThrift (TV.I16 n) = Right (fromIntegral n) + fromThrift (TV.Byte n) = Right (fromIntegral n) + fromThrift _ = Left "FromThrift Word64: expected integer type" + + +instance ToThrift Float where + toThrift f = TV.Double (realToFrac f) + + +instance FromThrift Float where + fromThrift (TV.Double d) = Right (realToFrac d) + fromThrift _ = Left "FromThrift Float: expected Double" + + +instance ToThrift Double where + toThrift = TV.Double + + +instance FromThrift Double where + fromThrift (TV.Double d) = Right d + fromThrift _ = Left "FromThrift Double: expected Double" + + +instance ToThrift Text where + toThrift = TV.String + + +instance FromThrift Text where + fromThrift (TV.String t) = Right t + fromThrift _ = Left "FromThrift Text: expected String" + + +instance ToThrift ByteString where + toThrift = TV.Binary + + +instance FromThrift ByteString where + fromThrift (TV.Binary bs) = Right bs + fromThrift (TV.UUID bs) = Right bs + fromThrift _ = Left "FromThrift ByteString: expected Binary" + + +instance ToThrift a => ToThrift [a] where + toThrift xs = + let vals = V.fromList (map toThrift xs) + tt = if V.null vals then TT_I64 else TV.thriftTypeOf (V.head vals) + in TV.List tt vals + + +instance FromThrift a => FromThrift [a] where + fromThrift (TV.List _ vs) = traverse fromThrift (V.toList vs) + fromThrift (TV.Set _ vs) = traverse fromThrift (V.toList vs) + fromThrift _ = Left "FromThrift [a]: expected List or Set" + + +instance ToThrift a => ToThrift (Vector a) where + toThrift xs = + let vals = V.map toThrift xs + tt = if V.null vals then TT_I64 else TV.thriftTypeOf (V.head vals) + in TV.List tt vals + + +instance FromThrift a => FromThrift (Vector a) where + fromThrift (TV.List _ vs) = V.mapM fromThrift vs + fromThrift (TV.Set _ vs) = V.mapM fromThrift vs + fromThrift _ = Left "FromThrift Vector: expected List or Set" + + +instance (ToThrift a, ToThrift b) => ToThrift (a, b) where + toThrift (a, b) = + let vals = V.fromList [toThrift a, toThrift b] + tt = TV.thriftTypeOf (toThrift a) + in TV.List tt vals + + +instance (FromThrift a, FromThrift b) => FromThrift (a, b) where + fromThrift (TV.List _ vs) + | V.length vs == 2 = (,) <$> fromThrift (vs V.! 0) <*> fromThrift (vs V.! 1) + fromThrift _ = Left "FromThrift (a,b): expected List of length 2" + + +-- Aeson-parity instances --------------------------------------------------- + +thriftListOf :: ToThrift a => [a] -> TV.Value +thriftListOf xs = + let vals = V.fromList (fmap toThrift xs) + tt = if V.null vals then TT_I64 else TV.thriftTypeOf (V.head vals) + in TV.List tt vals + + +instance ToThrift Char where + toThrift c = TV.String (T.singleton c) + + +instance FromThrift Char where + fromThrift (TV.String t) | T.length t == 1 = Right (T.head t) + fromThrift _ = Left "FromThrift Char: expected single-character String" + + +instance ToThrift Integer where + toThrift = TV.I64 . fromInteger + + +instance FromThrift Integer where + fromThrift (TV.I64 n) = Right (toInteger n) + fromThrift (TV.I32 n) = Right (toInteger n) + fromThrift (TV.I16 n) = Right (toInteger n) + fromThrift (TV.Byte n) = Right (toInteger n) + fromThrift _ = Left "FromThrift Integer: expected integer type" + + +instance ToThrift Natural where + toThrift = TV.I64 . fromIntegral + + +instance FromThrift Natural where + fromThrift v = do + n <- fromThrift v :: Either String Integer + if n < 0 + then Left "FromThrift Natural: negative integer" + else Right (fromInteger n) + + +instance ToThrift TL.Text where + toThrift = TV.String . TL.toStrict + + +instance FromThrift TL.Text where + fromThrift v = TL.fromStrict <$> fromThrift v + + +instance ToThrift BSL.ByteString where + toThrift = TV.Binary . BSL.toStrict + + +instance FromThrift BSL.ByteString where + fromThrift v = BSL.fromStrict <$> fromThrift v + + +instance ToThrift a => ToThrift (NonEmpty a) where + toThrift = thriftListOf . NE.toList + + +instance FromThrift a => FromThrift (NonEmpty a) where + fromThrift v = do + xs <- fromThrift v + case xs of + [] -> Left "FromThrift NonEmpty: empty list" + (y : ys) -> Right (y :| ys) + + +-- | 'Either' encodes as a single-field struct: field 1 = Left, field 2 = Right. +instance (ToThrift a, ToThrift b) => ToThrift (Either a b) where + toThrift (Left x) = TV.Struct (V.singleton (1, toThrift x)) + toThrift (Right x) = TV.Struct (V.singleton (2, toThrift x)) + + +instance (FromThrift a, FromThrift b) => FromThrift (Either a b) where + fromThrift (TV.Struct kvs) + | V.length kvs == 1 = case V.head kvs of + (1, v) -> Left <$> fromThrift v + (2, v) -> Right <$> fromThrift v + _ -> Left "FromThrift Either: expected field id 1 or 2" + fromThrift _ = Left "FromThrift Either: expected single-field Struct" + + +instance (Ord a, ToThrift a) => ToThrift (Set a) where + toThrift xs = + let vals = V.fromList (fmap toThrift (Set.toList xs)) + tt = if V.null vals then TT_I64 else TV.thriftTypeOf (V.head vals) + in TV.Set tt vals + + +instance (Ord a, FromThrift a) => FromThrift (Set a) where + fromThrift v = Set.fromList <$> fromThrift v + + +instance ToThrift a => ToThrift (Seq a) where + toThrift s = thriftListOf (foldr (:) [] s) + + +instance FromThrift a => FromThrift (Seq a) where + fromThrift v = Seq.fromList <$> fromThrift v + + +instance (ToThrift k, ToThrift v) => ToThrift (Map k v) where + toThrift m = + let pairs = V.fromList [(toThrift k, toThrift v) | (k, v) <- Map.toList m] + kt = if V.null pairs then TT_STRING else TV.thriftTypeOf (fst (V.head pairs)) + vt = if V.null pairs then TT_STRING else TV.thriftTypeOf (snd (V.head pairs)) + in TV.Map kt vt pairs + + +instance (Ord k, FromThrift k, FromThrift v) => FromThrift (Map k v) where + fromThrift (TV.Map _ _ kvs) = do + pairs <- traverse (\(k, v) -> (,) <$> fromThrift k <*> fromThrift v) (V.toList kvs) + Right (Map.fromList pairs) + fromThrift _ = Left "FromThrift Map: expected Map" + + +instance (Hashable k, ToThrift k, ToThrift v) => ToThrift (HashMap k v) where + toThrift m = + let pairs = V.fromList [(toThrift k, toThrift v) | (k, v) <- HM.toList m] + kt = if V.null pairs then TT_STRING else TV.thriftTypeOf (fst (V.head pairs)) + vt = if V.null pairs then TT_STRING else TV.thriftTypeOf (snd (V.head pairs)) + in TV.Map kt vt pairs + + +instance (Eq k, Hashable k, FromThrift k, FromThrift v) => FromThrift (HashMap k v) where + fromThrift (TV.Map _ _ kvs) = do + pairs <- traverse (\(k, v) -> (,) <$> fromThrift k <*> fromThrift v) (V.toList kvs) + Right (HM.fromList pairs) + fromThrift _ = Left "FromThrift HashMap: expected Map" + + +instance (Hashable a, ToThrift a) => ToThrift (HashSet a) where + toThrift xs = + let vals = V.fromList (fmap toThrift (HS.toList xs)) + tt = if V.null vals then TT_I64 else TV.thriftTypeOf (V.head vals) + in TV.Set tt vals + + +instance (Eq a, Hashable a, FromThrift a) => FromThrift (HashSet a) where + fromThrift v = HS.fromList <$> fromThrift v + + +instance ToThrift v => ToThrift (IntMap v) where + toThrift m = + let pairs = V.fromList [(toThrift k, toThrift v) | (k, v) <- IntMap.toList m] + vt = if V.null pairs then TT_STRING else TV.thriftTypeOf (snd (V.head pairs)) + in TV.Map TT_I64 vt pairs + + +instance FromThrift v => FromThrift (IntMap v) where + fromThrift (TV.Map _ _ kvs) = do + pairs <- traverse (\(k, v) -> (,) <$> fromThrift k <*> fromThrift v) (V.toList kvs) + Right (IntMap.fromList pairs) + fromThrift _ = Left "FromThrift IntMap: expected Map" + + +instance ToThrift IntSet where + toThrift = thriftListOf . IntSet.toList + + +instance FromThrift IntSet where + fromThrift v = IntSet.fromList <$> fromThrift v + + +instance (ToThrift a, ToThrift b, ToThrift c) => ToThrift (a, b, c) where + toThrift (a, b, c) = TV.List TT_STRUCT (V.fromList [toThrift a, toThrift b, toThrift c]) + + +instance (FromThrift a, FromThrift b, FromThrift c) => FromThrift (a, b, c) where + fromThrift (TV.List _ vs) + | V.length vs == 3 = + (,,) <$> fromThrift (vs V.! 0) <*> fromThrift (vs V.! 1) <*> fromThrift (vs V.! 2) + fromThrift _ = Left "FromThrift (a,b,c): expected List of length 3" + + +instance (ToThrift a, ToThrift b, ToThrift c, ToThrift d) => ToThrift (a, b, c, d) where + toThrift (a, b, c, d) = TV.List TT_STRUCT (V.fromList [toThrift a, toThrift b, toThrift c, toThrift d]) + + +instance (FromThrift a, FromThrift b, FromThrift c, FromThrift d) => FromThrift (a, b, c, d) where + fromThrift (TV.List _ vs) + | V.length vs == 4 = + (,,,) + <$> fromThrift (vs V.! 0) + <*> fromThrift (vs V.! 1) + <*> fromThrift (vs V.! 2) + <*> fromThrift (vs V.! 3) + fromThrift _ = Left "FromThrift (a,b,c,d): expected List of length 4" + + +instance ToThrift a => ToThrift (Identity a) where + toThrift (Identity x) = toThrift x + + +instance FromThrift a => FromThrift (Identity a) where + fromThrift v = Identity <$> fromThrift v + + +instance ToThrift a => ToThrift (Const a b) where + toThrift (Const x) = toThrift x + + +instance FromThrift a => FromThrift (Const a b) where + fromThrift v = Const <$> fromThrift v + + +instance ToThrift a => ToThrift (Down a) where + toThrift (Down x) = toThrift x + + +instance FromThrift a => FromThrift (Down a) where + fromThrift v = Down <$> fromThrift v + + +instance ToThrift Version where + toThrift = toThrift . versionBranch + + +instance FromThrift Version where + fromThrift v = makeVersion <$> fromThrift v + + +instance ToThrift () where + toThrift () = TV.Struct V.empty + + +instance FromThrift () where + fromThrift (TV.Struct kvs) | V.null kvs = Right () + fromThrift _ = Left "FromThrift (): expected empty Struct" + + +instance ToThrift a => ToThrift (Maybe a) where + toThrift Nothing = TV.Struct V.empty + toThrift (Just x) = TV.Struct (V.singleton (1, toThrift x)) + + +instance FromThrift a => FromThrift (Maybe a) where + fromThrift (TV.Struct kvs) + | V.null kvs = Right Nothing + | V.length kvs == 1, (1, v) <- V.head kvs = Just <$> fromThrift v + fromThrift v = Just <$> fromThrift v + + +instance (Integral a, ToThrift a) => ToThrift (Ratio a) where + toThrift r = TV.List TT_STRUCT (V.fromList [toThrift (numerator r), toThrift (denominator r)]) + + +instance (Integral a, FromThrift a) => FromThrift (Ratio a) where + fromThrift (TV.List _ vs) + | V.length vs == 2 = do + n <- fromThrift (vs V.! 0) + d <- fromThrift (vs V.! 1) + if d == 0 + then Left "FromThrift Ratio: zero denominator" + else Right (n % d) + fromThrift _ = Left "FromThrift Ratio: expected List of length 2" + + +-- Functor / monoid newtype instances -------------------------------------- + +instance ToThrift a => ToThrift (Mon.Sum a) where + toThrift = toThrift . Mon.getSum + + +instance FromThrift a => FromThrift (Mon.Sum a) where + fromThrift v = Mon.Sum <$> fromThrift v + + +instance ToThrift a => ToThrift (Mon.Product a) where + toThrift = toThrift . Mon.getProduct + + +instance FromThrift a => FromThrift (Mon.Product a) where + fromThrift v = Mon.Product <$> fromThrift v + + +instance ToThrift a => ToThrift (Mon.Dual a) where + toThrift = toThrift . Mon.getDual + + +instance FromThrift a => FromThrift (Mon.Dual a) where + fromThrift v = Mon.Dual <$> fromThrift v + + +instance ToThrift Mon.All where + toThrift = toThrift . Mon.getAll + + +instance FromThrift Mon.All where + fromThrift v = Mon.All <$> fromThrift v + + +instance ToThrift Mon.Any where + toThrift = toThrift . Mon.getAny + + +instance FromThrift Mon.Any where + fromThrift v = Mon.Any <$> fromThrift v + + +instance ToThrift a => ToThrift (Mon.First a) where + toThrift = toThrift . Mon.getFirst + + +instance FromThrift a => FromThrift (Mon.First a) where + fromThrift v = Mon.First <$> fromThrift v + + +instance ToThrift a => ToThrift (Mon.Last a) where + toThrift = toThrift . Mon.getLast + + +instance FromThrift a => FromThrift (Mon.Last a) where + fromThrift v = Mon.Last <$> fromThrift v + + +instance ToThrift a => ToThrift (Semi.Min a) where + toThrift = toThrift . Semi.getMin + + +instance FromThrift a => FromThrift (Semi.Min a) where + fromThrift v = Semi.Min <$> fromThrift v + + +instance ToThrift a => ToThrift (Semi.Max a) where + toThrift = toThrift . Semi.getMax + + +instance FromThrift a => FromThrift (Semi.Max a) where + fromThrift v = Semi.Max <$> fromThrift v + + +instance ToThrift a => ToThrift (Semi.First a) where + toThrift = toThrift . Semi.getFirst + + +instance FromThrift a => FromThrift (Semi.First a) where + fromThrift v = Semi.First <$> fromThrift v + + +instance ToThrift a => ToThrift (Semi.Last a) where + toThrift = toThrift . Semi.getLast + + +instance FromThrift a => FromThrift (Semi.Last a) where + fromThrift v = Semi.Last <$> fromThrift v + + +instance ToThrift a => ToThrift (Semi.WrappedMonoid a) where + toThrift = toThrift . Semi.unwrapMonoid + + +instance FromThrift a => FromThrift (Semi.WrappedMonoid a) where + fromThrift v = Semi.WrapMonoid <$> fromThrift v + + +instance (ToThrift a, ToThrift b) => ToThrift (Semi.Arg a b) where + toThrift (Semi.Arg a b) = TV.List TT_STRUCT (V.fromList [toThrift a, toThrift b]) + + +instance (FromThrift a, FromThrift b) => FromThrift (Semi.Arg a b) where + fromThrift (TV.List _ vs) + | V.length vs == 2 = Semi.Arg <$> fromThrift (vs V.! 0) <*> fromThrift (vs V.! 1) + fromThrift _ = Left "FromThrift Arg: expected List of length 2" + + +instance ToThrift (f (g a)) => ToThrift (Compose f g a) where + toThrift = toThrift . getCompose + + +instance FromThrift (f (g a)) => FromThrift (Compose f g a) where + fromThrift v = Compose <$> fromThrift v + + +instance (ToThrift (f a), ToThrift (g a)) => ToThrift (FProduct.Product f g a) where + toThrift (FProduct.Pair x y) = TV.List TT_STRUCT (V.fromList [toThrift x, toThrift y]) + + +instance (FromThrift (f a), FromThrift (g a)) => FromThrift (FProduct.Product f g a) where + fromThrift (TV.List _ vs) + | V.length vs == 2 = FProduct.Pair <$> fromThrift (vs V.! 0) <*> fromThrift (vs V.! 1) + fromThrift _ = Left "FromThrift Functor.Product: expected List of length 2" + + +instance (ToThrift (f a), ToThrift (g a)) => ToThrift (FSum.Sum f g a) where + toThrift (FSum.InL x) = TV.Struct (V.singleton (1, toThrift x)) + toThrift (FSum.InR x) = TV.Struct (V.singleton (2, toThrift x)) + + +instance (FromThrift (f a), FromThrift (g a)) => FromThrift (FSum.Sum f g a) where + fromThrift (TV.Struct kvs) + | V.length kvs == 1 = case V.head kvs of + (1, v) -> FSum.InL <$> fromThrift v + (2, v) -> FSum.InR <$> fromThrift v + _ -> Left "FromThrift Functor.Sum: expected field id 1 or 2" + fromThrift _ = Left "FromThrift Functor.Sum: expected single-field Struct" + + +instance ToThrift TV.Value where + toThrift = id + + +instance FromThrift TV.Value where + fromThrift = Right + + +-- GHC.Generics support + +class GToThrift f where + gToThrift :: f p -> TV.Value + + +class GFromThrift f where + gFromThrift :: TV.Value -> Either String (f p) + + +instance GToThrift f => GToThrift (M1 D c f) where + gToThrift (M1 x) = gToThrift x + + +instance GFromThrift f => GFromThrift (M1 D c f) where + gFromThrift v = M1 <$> gFromThrift v + + +instance (GToThriftFields f) => GToThrift (M1 C c f) where + gToThrift (M1 x) = + let fields = gToThriftFields 1 x + in TV.Struct (V.fromList fields) + + +instance (GFromThriftFields f) => GFromThrift (M1 C c f) where + gFromThrift (TV.Struct kvs) = + let lkup fid = lookupField fid kvs + in M1 <$> gFromThriftFields 1 lkup + gFromThrift _ = Left "GFromThrift: expected Struct for record type" + + +lookupField :: Int16 -> Vector (Int16, TV.Value) -> Maybe TV.Value +lookupField fid kvs = go 0 + where + len = V.length kvs + go i + | i >= len = Nothing + | (k, v) <- kvs V.! i, k == fid = Just v + | otherwise = go (i + 1) + + +class GToThriftFields f where + gToThriftFields :: Int16 -> f p -> [(Int16, TV.Value)] + + +class GFromThriftFields f where + gFromThriftFields :: Int16 -> (Int16 -> Maybe TV.Value) -> Either String (f p) + + +instance (GToThriftFields a, GToThriftFields b, GFieldCount a) => GToThriftFields (a :*: b) where + gToThriftFields fid (a :*: b) = + gToThriftFields fid a ++ gToThriftFields (fid + gFieldCount (undefined :: a p)) b + + +instance (GFromThriftFields a, GFromThriftFields b, GFieldCount a) => GFromThriftFields (a :*: b) where + gFromThriftFields fid lkup = + (:*:) + <$> gFromThriftFields fid lkup + <*> gFromThriftFields (fid + gFieldCount (undefined :: a p)) lkup + + +instance (ToThrift a) => GToThriftFields (M1 S s (K1 i a)) where + gToThriftFields fid (M1 (K1 x)) = [(fid, toThrift x)] + + +instance (FromThrift a) => GFromThriftFields (M1 S s (K1 i a)) where + gFromThriftFields fid lkup = + case lkup fid of + Nothing -> Left $ "GFromThrift: missing field " ++ show fid + Just v -> M1 . K1 <$> fromThrift v + + +class GFieldCount f where + gFieldCount :: f p -> Int16 + + +instance GFieldCount (M1 S s (K1 i a)) where + gFieldCount _ = 1 + + +instance (GFieldCount a, GFieldCount b) => GFieldCount (a :*: b) where + gFieldCount _ = gFieldCount (undefined :: a p) + gFieldCount (undefined :: b p) diff --git a/src/Thrift/CodeGen.hs b/src/Thrift/CodeGen.hs new file mode 100644 --- /dev/null +++ b/src/Thrift/CodeGen.hs @@ -0,0 +1,534 @@ +{-# LANGUAGE TemplateHaskell #-} + +{- | Thrift code generation — generates Haskell data types and +ToThrift\/FromThrift instances from Thrift schemas. + +== Text generation + +'generateThriftTypes' takes a 'ThriftSchema' and produces Haskell source +as 'Text'. The output includes @data@ declarations for structs, enums, +and service method descriptors. + +== Template Haskell + +'deriveThrift' generates declarations at compile time from a 'ThriftSchema'. +-} +module Thrift.CodeGen ( + generateThriftTypes, + generateThriftTypesWithRegistry, + deriveThrift, +) where + +import Data.ByteString (ByteString) +import Data.Char (toLower, toUpper) +import Data.Int (Int16, Int32, Int64, Int8) +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.Text (Text) +import Data.Text qualified as T +import Data.Vector qualified as V +import Language.Haskell.TH +import Thrift.Registry +import Thrift.Schema + + +-- --------------------------------------------------------------------------- +-- Text-based code generation +-- --------------------------------------------------------------------------- + +{- | Generate Haskell source code (as 'Text') for all types in a 'ThriftSchema'. +Uses 'defaultThriftRegistry'. +-} +generateThriftTypes :: ThriftSchema -> Text +generateThriftTypes = generateThriftTypesWithRegistry defaultThriftRegistry + + +{- | Generate Haskell source code using a custom 'ThriftRegistry'. +When a field has an annotation matching a handler, the type transformation +is applied and extra code is emitted. +-} +generateThriftTypesWithRegistry :: ThriftRegistry -> ThriftSchema -> Text +generateThriftTypesWithRegistry reg schema = + let typedefDecls = map genThriftTypedef (tsTypedefs schema) + constDecls = map genThriftConst (tsConsts schema) + enumDecls = map genThriftEnum (tsEnums schema) + structDecls = map (genThriftStructWithRegistry reg) (tsStructs schema) + serviceDecls = map genThriftService (tsServices schema) + in T.intercalate "\n\n" (typedefDecls <> constDecls <> enumDecls <> structDecls <> serviceDecls) + + +-- --------------------------------------------------------------------------- +-- Typedef generation (text) +-- --------------------------------------------------------------------------- + +genThriftTypedef :: ThriftTypedef -> Text +genThriftTypedef td = + "type " <> ttName td <> " = " <> thriftInnerHsType (ttType td) + + +-- --------------------------------------------------------------------------- +-- Const generation (text) +-- --------------------------------------------------------------------------- + +genThriftConst :: ThriftConst -> Text +genThriftConst tc = + let name = lowerFirst (snakeToCamel (T.toLower (tcName tc))) + hsType = thriftInnerHsType (tcType tc) + hsVal = constValueToHs (tcType tc) (tcValue tc) + in name <> " :: " <> hsType <> "\n" <> name <> " = " <> hsVal + + +constValueToHs :: ThriftType -> ThriftConstValue -> Text +constValueToHs _ (TCVInt n) = T.pack (show n) +constValueToHs _ (TCVDouble d) = T.pack (show d) +constValueToHs _ (TCVString s) = "\"" <> s <> "\"" +constValueToHs _ (TCVBool True) = "True" +constValueToHs _ (TCVBool False) = "False" +constValueToHs _ (TCVIdent i) = i +constValueToHs _ (TCVList vs) = + "[" <> T.intercalate ", " (map (constValueToHs TString) vs) <> "]" +constValueToHs _ (TCVMap _) = "mempty" + + +-- --------------------------------------------------------------------------- +-- Enum generation (text) +-- --------------------------------------------------------------------------- + +genThriftEnum :: ThriftEnum -> Text +genThriftEnum te = + let name = teName te + vals = teValues te + in T.unlines $ + case vals of + [] -> + [ "data " <> name + , " deriving stock (Show, Eq, Ord, Enum, Bounded, Generic)" + , " deriving anyclass NFData" + ] + ((sym, _) : rest) -> + ["data " <> name] + <> [" = " <> thriftEnumConName name sym] + <> map (\(s, _) -> " | " <> thriftEnumConName name s) rest + <> [ " deriving stock (Show, Eq, Ord, Enum, Bounded, Generic)" + , " deriving anyclass NFData" + ] + <> genThriftEnumToThrift name vals + <> genThriftEnumFromThrift name vals + + +thriftEnumConName :: Text -> Text -> Text +thriftEnumConName enumName valName = + enumName <> upperFirst (snakeToCamel (T.toLower valName)) + + +genThriftEnumToThrift :: Text -> [(Text, Int32)] -> [Text] +genThriftEnumToThrift name vals = + [ "" + , "instance ToThrift " <> name <> " where" + ] + <> zipWith + ( \i (sym, _) -> + " toThrift " <> thriftEnumConName name sym <> " = Thrift.Value.I32 " <> T.pack (show i) + ) + [0 :: Int ..] + vals + + +genThriftEnumFromThrift :: Text -> [(Text, Int32)] -> [Text] +genThriftEnumFromThrift name vals = + [ "" + , "instance FromThrift " <> name <> " where" + ] + <> zipWith + ( \i (sym, _) -> + " fromThrift (Thrift.Value.I32 " <> T.pack (show i) <> ") = Right " <> thriftEnumConName name sym + ) + [0 :: Int ..] + vals + <> [" fromThrift _ = Left \"FromThrift " <> name <> ": expected I32 enum value\""] + + +-- --------------------------------------------------------------------------- +-- Struct generation (text) +-- --------------------------------------------------------------------------- + +genThriftStructWithRegistry :: ThriftRegistry -> ThriftStruct -> Text +genThriftStructWithRegistry reg ts = + let name = tsName ts + fields = tsFields ts + structAnns = V.toList (tsAnnotations ts) + extraStructCode = + concatMap + ( \(k, v) -> + case Map.lookup k (trStructAnnotations reg) of + Just handler -> sahExtraCode handler k v + Nothing -> [] + ) + structAnns + extraDerivs = + concatMap + ( \(k, v) -> + case Map.lookup k (trStructAnnotations reg) of + Just handler -> sahExtraDerivations handler v + Nothing -> [] + ) + structAnns + fieldExtraCode = + concatMap + ( \fld -> + concatMap + ( \(k, v) -> + case Map.lookup k (trFieldAnnotations reg) of + Just handler -> fahExtraCode handler (tfFieldName fld) v + Nothing -> [] + ) + (V.toList (tfAnnotations fld)) + ) + fields + in T.unlines $ + genStructDataDeclWithRegistry reg name fields + <> (if null extraDerivs then [] else map (" deriving " <>) extraDerivs) + <> genStructToThrift name fields + <> genStructFromThrift name fields + <> extraStructCode + <> fieldExtraCode + + +genStructDataDeclWithRegistry :: ThriftRegistry -> Text -> [ThriftField] -> [Text] +genStructDataDeclWithRegistry reg name fields = + ["data " <> name <> " = " <> name] + <> case fields of + [] -> + [ " deriving stock (Show, Eq, Generic)" + , " deriving anyclass NFData" + ] + (f : fs) -> + [" { " <> genThriftFieldDeclWithRegistry reg name f] + <> map (\fld -> " , " <> genThriftFieldDeclWithRegistry reg name fld) fs + <> [ " } deriving stock (Show, Eq, Generic)" + , " deriving anyclass NFData" + ] + + +genThriftFieldDeclWithRegistry :: ThriftRegistry -> Text -> ThriftField -> Text +genThriftFieldDeclWithRegistry reg recName fld = + let accessor = thriftFieldAccessorName recName (tfFieldName fld) + baseType = thriftFieldHsType (tfFieldType fld) (tfRequiredness fld) + hsType = applyFieldAnnotationTransforms reg fld baseType + in accessor <> " :: " <> hsType + + +applyFieldAnnotationTransforms :: ThriftRegistry -> ThriftField -> Text -> Text +applyFieldAnnotationTransforms reg fld ty = + foldl + ( \acc (k, _v) -> + case Map.lookup k (trFieldAnnotations reg) of + Just handler -> fahTransformType handler acc + Nothing -> acc + ) + ty + (V.toList (tfAnnotations fld)) + + +thriftFieldAccessorName :: Text -> Text -> Text +thriftFieldAccessorName recName fieldName = + lowerFirst recName <> upperFirst (snakeToCamel fieldName) + + +thriftFieldHsType :: ThriftType -> Requiredness -> Text +thriftFieldHsType ty req = case req of + Optional -> "!(Maybe " <> thriftInnerHsType ty <> ")" + Required -> thriftStrictHsType ty + Default -> thriftStrictHsType ty + + +thriftStrictHsType :: ThriftType -> Text +thriftStrictHsType = \case + TBool -> "!Bool" + TByte -> "{-# UNPACK #-} !Int8" + TI16 -> "{-# UNPACK #-} !Int16" + TI32 -> "{-# UNPACK #-} !Int32" + TI64 -> "{-# UNPACK #-} !Int64" + TDouble -> "{-# UNPACK #-} !Double" + TString -> "!Text" + TBinary -> "!ByteString" + TUUID -> "!ByteString" + TStruct n -> "!" <> n + TEnum n -> "!" <> n + TTypedef n -> "!" <> n + TList elemTy -> "!(Vector " <> thriftInnerHsType elemTy <> ")" + TSet elemTy -> "!(Vector " <> thriftInnerHsType elemTy <> ")" + TMap keyTy valTy -> "!(Map " <> thriftInnerHsType keyTy <> " " <> thriftInnerHsType valTy <> ")" + + +thriftInnerHsType :: ThriftType -> Text +thriftInnerHsType = \case + TBool -> "Bool" + TByte -> "Int8" + TI16 -> "Int16" + TI32 -> "Int32" + TI64 -> "Int64" + TDouble -> "Double" + TString -> "Text" + TBinary -> "ByteString" + TUUID -> "ByteString" + TStruct n -> n + TEnum n -> n + TTypedef n -> n + TList elemTy -> "(Vector " <> thriftInnerHsType elemTy <> ")" + TSet elemTy -> "(Vector " <> thriftInnerHsType elemTy <> ")" + TMap keyTy valTy -> "(Map " <> thriftInnerHsType keyTy <> " " <> thriftInnerHsType valTy <> ")" + + +-- --------------------------------------------------------------------------- +-- Struct ToThrift instance (text) +-- --------------------------------------------------------------------------- + +genStructToThrift :: Text -> [ThriftField] -> [Text] +genStructToThrift name fields = + [ "" + , "instance ToThrift " <> name <> " where" + , " toThrift msg = Thrift.Value.Struct $ V.fromList" + ] + <> case fields of + [] -> [" []"] + _ -> + [" [ " <> toThriftFieldExpr name (head fields)] + <> map (\f -> " , " <> toThriftFieldExpr name f) (tail fields) + <> [" ]"] + + +toThriftFieldExpr :: Text -> ThriftField -> Text +toThriftFieldExpr recName fld = + let accessor = "msg." <> thriftFieldAccessorName recName (tfFieldName fld) + fid = T.pack (show (tfFieldId fld)) + in case tfRequiredness fld of + Optional -> + "maybe (0, Thrift.Value.Bool False) (\\v -> (" <> fid <> ", " <> toThriftConvert (tfFieldType fld) "v" <> ")) " <> accessor + _ -> + "(" <> fid <> ", " <> toThriftConvert (tfFieldType fld) accessor <> ")" + + +toThriftConvert :: ThriftType -> Text -> Text +toThriftConvert ty accessor = case ty of + TBool -> "Thrift.Value.Bool " <> accessor + TByte -> "Thrift.Value.Byte " <> accessor + TI16 -> "Thrift.Value.I16 " <> accessor + TI32 -> "Thrift.Value.I32 " <> accessor + TI64 -> "Thrift.Value.I64 " <> accessor + TDouble -> "Thrift.Value.Double " <> accessor + TString -> "Thrift.Value.String " <> accessor + TBinary -> "Thrift.Value.Binary " <> accessor + TUUID -> "Thrift.Value.UUID " <> accessor + TStruct _ -> "toThrift " <> accessor + TEnum _ -> "toThrift " <> accessor + TTypedef _ -> "toThrift " <> accessor + TList _ -> "toThrift " <> accessor + TSet _ -> "toThrift " <> accessor + TMap _ _ -> "toThrift " <> accessor + + +-- --------------------------------------------------------------------------- +-- Struct FromThrift instance (text) +-- --------------------------------------------------------------------------- + +genStructFromThrift :: Text -> [ThriftField] -> [Text] +genStructFromThrift name fields = + [ "" + , "instance FromThrift " <> name <> " where" + , " fromThrift (Thrift.Value.Struct fieldVec) = do" + , " let fieldMap = Map.fromList (V.toList fieldVec)" + ] + <> map (fromThriftFieldBind name) fields + <> [ " pure " <> name + , " { " + <> T.intercalate + "\n , " + ( map + ( \fld -> + let accessor = thriftFieldAccessorName name (tfFieldName fld) + in accessor <> " = " <> accessor <> "'" + ) + fields + ) + , " }" + , " fromThrift _ = Left \"FromThrift " <> name <> ": expected Struct\"" + ] + + +fromThriftFieldBind :: Text -> ThriftField -> Text +fromThriftFieldBind recName fld = + let accessor = thriftFieldAccessorName recName (tfFieldName fld) + fid = T.pack (show (tfFieldId fld)) + in case tfRequiredness fld of + Optional -> + " let " <> accessor <> "' = case Map.lookup " <> fid <> " fieldMap of { Just v -> case fromThrift v of { Right x -> Just x; _ -> Nothing }; Nothing -> Nothing }" + _ -> + " " <> accessor <> "' <- case Map.lookup " <> fid <> " fieldMap of { Just v -> fromThrift v; Nothing -> Left \"missing field " <> tfFieldName fld <> "\" }" + + +-- --------------------------------------------------------------------------- +-- Service generation (text) — type-level method descriptors +-- --------------------------------------------------------------------------- + +genThriftService :: ThriftService -> Text +genThriftService svc = + let name = tsvName svc + methods = tsvMethods svc + in T.unlines $ + [ "-- | Service @" <> name <> "@ method descriptors." + , "data " <> name <> "Method" + ] + <> case methods of + [] -> + [" deriving stock (Show, Eq, Ord, Enum, Bounded)"] + (m : ms) -> + [" = " <> name <> upperFirst (snakeToCamel (tmName m)) <> "Method"] + <> map + ( \method -> + " | " <> name <> upperFirst (snakeToCamel (tmName method)) <> "Method" + ) + ms + <> [" deriving stock (Show, Eq, Ord, Enum, Bounded)"] + <> [ "" + , name <> "ServiceName :: Text" + , name <> "ServiceName = \"" <> name <> "\"" + , "" + , name <> "MethodName :: " <> name <> "Method -> Text" + ] + <> map + ( \method -> + name <> "MethodName " <> name <> upperFirst (snakeToCamel (tmName method)) <> "Method = \"" <> tmName method <> "\"" + ) + methods + + +-- --------------------------------------------------------------------------- +-- Template Haskell +-- --------------------------------------------------------------------------- + +-- | Generate Haskell declarations from a 'ThriftSchema' at compile time. +deriveThrift :: ThriftSchema -> Q [Dec] +deriveThrift schema = do + enumDecs <- concat <$> mapM deriveThriftEnumTH (tsEnums schema) + structDecs <- concat <$> mapM deriveThriftStructTH (tsStructs schema) + pure (enumDecs <> structDecs) + + +-- --------------------------------------------------------------------------- +-- TH: Enum +-- --------------------------------------------------------------------------- + +deriveThriftEnumTH :: ThriftEnum -> Q [Dec] +deriveThriftEnumTH te = do + let name = teName te + tyName = mkName (T.unpack name) + vals = teValues te + cons = map (\(sym, _) -> NormalC (mkName (T.unpack (thriftEnumConName name sym))) []) vals + dataDec = + DataD + [] + tyName + [] + Nothing + cons + [ DerivClause + (Just StockStrategy) + [ConT ''Show, ConT ''Eq, ConT ''Ord, ConT ''Enum, ConT ''Bounded] + ] + pure [dataDec] + + +-- --------------------------------------------------------------------------- +-- TH: Struct +-- --------------------------------------------------------------------------- + +deriveThriftStructTH :: ThriftStruct -> Q [Dec] +deriveThriftStructTH ts = do + let name = tsName ts + tyName = mkName (T.unpack name) + conName = mkName (T.unpack name) + fields = tsFields ts + fieldDecs <- mapM (mkThriftRecordField name) fields + let dataDec = + DataD + [] + tyName + [] + Nothing + [RecC conName fieldDecs] + [DerivClause (Just StockStrategy) [ConT ''Show, ConT ''Eq]] + pure [dataDec] + + +mkThriftRecordField :: Text -> ThriftField -> Q VarBangType +mkThriftRecordField recName fld = do + let accessor = thriftFieldAccessorName recName (tfFieldName fld) + accName = mkName (T.unpack accessor) + hsTy <- thriftFieldToTHType (tfFieldType fld) (tfRequiredness fld) + let bang = Bang NoSourceUnpackedness SourceStrict + pure (accName, bang, hsTy) + + +thriftFieldToTHType :: ThriftType -> Requiredness -> Q Type +thriftFieldToTHType ty req = case req of + Optional -> do + inner <- thriftTypeToTH ty + pure (AppT (ConT ''Maybe) inner) + _ -> thriftTypeToTH ty + + +thriftTypeToTH :: ThriftType -> Q Type +thriftTypeToTH = \case + TBool -> [t|Bool|] + TByte -> [t|Int8|] + TI16 -> [t|Int16|] + TI32 -> [t|Int32|] + TI64 -> [t|Int64|] + TDouble -> [t|Double|] + TString -> [t|Text|] + TBinary -> [t|ByteString|] + TUUID -> [t|ByteString|] + TStruct n -> pure (ConT (mkName (T.unpack n))) + TEnum n -> pure (ConT (mkName (T.unpack n))) + TTypedef n -> pure (ConT (mkName (T.unpack n))) + TList elemTy -> do + inner <- thriftTypeToTH elemTy + pure (AppT (ConT ''V.Vector) inner) + TSet elemTy -> do + inner <- thriftTypeToTH elemTy + pure (AppT (ConT ''V.Vector) inner) + TMap keyTy valTy -> do + k <- thriftTypeToTH keyTy + v <- thriftTypeToTH valTy + pure (AppT (AppT (ConT ''Map) k) v) + + +-- --------------------------------------------------------------------------- +-- Name helpers +-- --------------------------------------------------------------------------- + +lowerFirst :: Text -> Text +lowerFirst s = case T.uncons s of + Just (c, rest) -> T.cons (toLower c) rest + Nothing -> s + + +upperFirst :: Text -> Text +upperFirst s = case T.uncons s of + Just (c, rest) -> T.cons (toUpper c) rest + Nothing -> s + + +snakeToCamel :: Text -> Text +snakeToCamel t = + let parts = T.splitOn "_" t + in case parts of + [] -> t + (p : ps) -> T.concat (lowerFirst p : map titleCase ps) + + +titleCase :: Text -> Text +titleCase s = case T.uncons s of + Just (c, rest) -> T.cons (toUpper c) (T.toLower rest) + Nothing -> s diff --git a/src/Thrift/Decode.hs b/src/Thrift/Decode.hs new file mode 100644 --- /dev/null +++ b/src/Thrift/Decode.hs @@ -0,0 +1,366 @@ +{-# LANGUAGE BangPatterns #-} + +{- | High-level Thrift decoding for Binary and Compact protocols. + +Reads a wire-format 'ByteString' and produces a 'Thrift.Value.Value' tree. +Container elements use pre-allocated mutable vectors. + +@ +import Thrift.Decode (decodeBinary, decodeCompact) + +case decodeBinary bytes of + Right val -> print val + Left err -> putStrLn err +@ +-} +module Thrift.Decode ( + decodeBinary, + decodeCompact, + decodeCompactFrom, + decodeCompactBinaryStrings, +) where + +import Control.Monad.ST (ST, runST) +import Data.ByteString (ByteString) +import Data.ByteString qualified as BS +import Data.Int (Int16) +import Data.Vector qualified as V +import Data.Vector.Mutable qualified as MV +import Thrift.Value qualified as TV +import Thrift.Wire +import Wireform.FFI (decodeTextFast) + + +-------------------------------------------------------------------------------- +-- Binary Protocol +-------------------------------------------------------------------------------- + +decodeBinary :: ByteString -> Either String TV.Value +decodeBinary bs = case decodeBinStruct bs 0 of + Nothing -> Left "decodeBinary: failed to decode struct" + Just (v, _off) -> Right v + + +decodeBinStruct :: ByteString -> Int -> Maybe (TV.Value, Int) +decodeBinStruct bs off = go off [] + where + go !o !acc = case tBinDecodeFieldBegin bs o of + Nothing -> Nothing + Just (TT_STOP, _, o') -> Just (TV.Struct (V.fromList (reverse acc)), o') + Just (tt, fid, o') -> case decodeBinValue tt bs o' of + Nothing -> Nothing + Just (v, o'') -> go o'' ((fid, v) : acc) +{-# INLINE decodeBinStruct #-} + + +decodeBinValue :: ThriftType -> ByteString -> Int -> Maybe (TV.Value, Int) +decodeBinValue !tt !bs !off = case tt of + TT_BOOL -> case tBinDecodeBool bs off of + Just (b, o) -> Just (TV.Bool b, o) + Nothing -> Nothing + TT_BYTE -> case tBinDecodeI8 bs off of + Just (v, o) -> Just (TV.Byte v, o) + Nothing -> Nothing + TT_I16 -> case tBinDecodeI16 bs off of + Just (v, o) -> Just (TV.I16 v, o) + Nothing -> Nothing + TT_I32 -> case tBinDecodeI32 bs off of + Just (v, o) -> Just (TV.I32 v, o) + Nothing -> Nothing + TT_I64 -> case tBinDecodeI64 bs off of + Just (v, o) -> Just (TV.I64 v, o) + Nothing -> Nothing + TT_DOUBLE -> case tBinDecodeDouble bs off of + Just (d, o) -> Just (TV.Double d, o) + Nothing -> Nothing + TT_STRING -> case tBinDecodeString bs off of + Just (b, o) -> case decodeTextFast b of + Right t -> Just (TV.String t, o) + Left _ -> Just (TV.Binary b, o) + Nothing -> Nothing + TT_STRUCT -> decodeBinStruct bs off + TT_MAP -> case tBinDecodeMapBegin bs off of + Nothing -> Nothing + Just (kt, vt, sz, o) -> + let !count = fromIntegral sz + in decodeBinMapEntries kt vt count bs o + TT_LIST -> case tBinDecodeListBegin bs off of + Nothing -> Nothing + Just (et, sz, o) -> + let !count = fromIntegral sz + in decodeBinListEntries et count bs o + TT_SET -> case tBinDecodeSetBegin bs off of + Nothing -> Nothing + Just (et, sz, o) -> + let !count = fromIntegral sz + in decodeBinSetEntries et count bs o + TT_UUID + | off + 16 > BS.length bs -> Nothing + | otherwise -> Just (TV.UUID (BS.take 16 (BS.drop off bs)), off + 16) + TT_STOP -> Nothing +{-# INLINE decodeBinValue #-} + + +decodeBinMapEntries + :: ThriftType + -> ThriftType + -> Int + -> ByteString + -> Int + -> Maybe (TV.Value, Int) +decodeBinMapEntries kt vt !count bs !off + | count < 0 || off < 0 || off > BS.length bs || count > (BS.length bs - off) `quot` 2 = Nothing + | count == 0 = Just (TV.Map kt vt V.empty, off) + | otherwise = runST $ do + mv <- MV.new count + go mv 0 off + where + go :: MV.MVector s (TV.Value, TV.Value) -> Int -> Int -> ST s (Maybe (TV.Value, Int)) + go !mv !i !o + | i >= count = do + vec <- V.unsafeFreeze mv + pure $! Just (TV.Map kt vt vec, o) + | otherwise = case decodeBinValue kt bs o of + Nothing -> pure Nothing + Just (k, o1) -> case decodeBinValue vt bs o1 of + Nothing -> pure Nothing + Just (v, o2) -> do + MV.unsafeWrite mv i (k, v) + go mv (i + 1) o2 +{-# INLINE decodeBinMapEntries #-} + + +decodeBinListEntries + :: ThriftType + -> Int + -> ByteString + -> Int + -> Maybe (TV.Value, Int) +decodeBinListEntries et !count bs !off + | count < 0 || off < 0 || off > BS.length bs || count > BS.length bs - off = Nothing + | count == 0 = Just (TV.List et V.empty, off) + | otherwise = runST $ do + mv <- MV.new count + go mv 0 off + where + go :: MV.MVector s TV.Value -> Int -> Int -> ST s (Maybe (TV.Value, Int)) + go !mv !i !o + | i >= count = do + vec <- V.unsafeFreeze mv + pure $! Just (TV.List et vec, o) + | otherwise = case decodeBinValue et bs o of + Nothing -> pure Nothing + Just (v, o') -> do + MV.unsafeWrite mv i v + go mv (i + 1) o' +{-# INLINE decodeBinListEntries #-} + + +decodeBinSetEntries + :: ThriftType + -> Int + -> ByteString + -> Int + -> Maybe (TV.Value, Int) +decodeBinSetEntries et !count bs !off + | count < 0 || off < 0 || off > BS.length bs || count > BS.length bs - off = Nothing + | count == 0 = Just (TV.Set et V.empty, off) + | otherwise = runST $ do + mv <- MV.new count + go mv 0 off + where + go :: MV.MVector s TV.Value -> Int -> Int -> ST s (Maybe (TV.Value, Int)) + go !mv !i !o + | i >= count = do + vec <- V.unsafeFreeze mv + pure $! Just (TV.Set et vec, o) + | otherwise = case decodeBinValue et bs o of + Nothing -> pure Nothing + Just (v, o') -> do + MV.unsafeWrite mv i v + go mv (i + 1) o' +{-# INLINE decodeBinSetEntries #-} + + +-------------------------------------------------------------------------------- +-- Compact Protocol +-------------------------------------------------------------------------------- + +decodeCompact :: ByteString -> Either String TV.Value +decodeCompact = decodeCompactWith True + + +{- | Preserve every Compact @STRING@ payload as 'TV.Binary'. Thrift uses the +same wire type for UTF-8 strings and opaque bytes. Schema-aware readers can +decode text fields later without copying binary statistics or index bounds +through Text. Nested structs and containers retain the same policy. +Binary payloads are slices sharing the input buffer's storage. +-} +decodeCompactBinaryStrings :: ByteString -> Either String TV.Value +decodeCompactBinaryStrings = decodeCompactWith False + + +decodeCompactWith :: Bool -> ByteString -> Either String TV.Value +decodeCompactWith decodeStrings bs = case decodeCompStruct decodeStrings bs 0 of + Nothing -> Left "decodeCompact: failed to decode struct" + Just (v, _off) -> Right v + + +{- | Decode one Thrift Compact struct starting at an offset; return the value +and the next byte offset after the struct. +-} +decodeCompactFrom :: ByteString -> Int -> Either String (TV.Value, Int) +decodeCompactFrom bs off = case decodeCompStruct True bs off of + Nothing -> Left "decodeCompactFrom: failed to decode struct" + Just (v, o) -> Right (v, o) + + +decodeCompStruct :: Bool -> ByteString -> Int -> Maybe (TV.Value, Int) +decodeCompStruct decodeStrings bs off = go off 0 [] + where + go :: Int -> Int16 -> [(Int16, TV.Value)] -> Maybe (TV.Value, Int) + go !o !lastFid !acc = case tCompDecodeFieldBegin bs o lastFid of + Nothing -> Nothing + Just (TT_STOP, _, o', _) -> Just (TV.Struct (V.fromList (reverse acc)), o') + Just (tt, fid, o', boolVal) -> case tt of + TT_BOOL -> go o' fid ((fid, TV.Bool boolVal) : acc) + _ -> case decodeCompValue decodeStrings tt bs o' of + Nothing -> Nothing + Just (v, o'') -> go o'' fid ((fid, v) : acc) +{-# INLINE decodeCompStruct #-} + + +decodeCompValue :: Bool -> ThriftType -> ByteString -> Int -> Maybe (TV.Value, Int) +decodeCompValue decodeStrings !tt !bs !off = case tt of + TT_BOOL -> case tCompDecodeBool bs off of + Just (b, o) -> Just (TV.Bool b, o) + Nothing -> Nothing + TT_BYTE -> case tCompDecodeI8 bs off of + Just (v, o) -> Just (TV.Byte v, o) + Nothing -> Nothing + TT_I16 -> case tCompDecodeI16 bs off of + Just (v, o) -> Just (TV.I16 v, o) + Nothing -> Nothing + TT_I32 -> case tCompDecodeI32 bs off of + Just (v, o) -> Just (TV.I32 v, o) + Nothing -> Nothing + TT_I64 -> case tCompDecodeI64 bs off of + Just (v, o) -> Just (TV.I64 v, o) + Nothing -> Nothing + TT_DOUBLE -> case tCompDecodeDouble bs off of + Just (d, o) -> Just (TV.Double d, o) + Nothing -> Nothing + TT_STRING -> case tCompDecodeString bs off of + Just (b, o) + | not decodeStrings -> Just (TV.Binary b, o) + | otherwise -> case decodeTextFast b of + Right t -> Just (TV.String t, o) + Left _ -> Just (TV.Binary b, o) + Nothing -> Nothing + TT_STRUCT -> decodeCompStruct decodeStrings bs off + TT_MAP -> case tCompDecodeMapBegin bs off of + Nothing -> Nothing + Just (kt, vt, sz, o) + | sz == 0 -> Just (TV.Map kt vt V.empty, o) + | otherwise -> + let !count = fromIntegral sz + in decodeCompMapEntries decodeStrings kt vt count bs o + TT_LIST -> case tCompDecodeListBegin bs off of + Nothing -> Nothing + Just (et, sz, o) -> + let !count = fromIntegral sz + in decodeCompListEntries decodeStrings et count bs o + TT_SET -> case tCompDecodeSetBegin bs off of + Nothing -> Nothing + Just (et, sz, o) -> + let !count = fromIntegral sz + in decodeCompSetEntries decodeStrings et count bs o + TT_UUID + | off + 16 > BS.length bs -> Nothing + | otherwise -> Just (TV.UUID (BS.take 16 (BS.drop off bs)), off + 16) + TT_STOP -> Nothing +{-# INLINE decodeCompValue #-} + + +decodeCompMapEntries + :: Bool + -> ThriftType + -> ThriftType + -> Int + -> ByteString + -> Int + -> Maybe (TV.Value, Int) +decodeCompMapEntries decodeStrings kt vt !count bs !off + | count < 0 || off < 0 || off > BS.length bs || count > (BS.length bs - off) `quot` 2 = Nothing + | count == 0 = Just (TV.Map kt vt V.empty, off) + | otherwise = runST $ do + mv <- MV.new count + go mv 0 off + where + go :: MV.MVector s (TV.Value, TV.Value) -> Int -> Int -> ST s (Maybe (TV.Value, Int)) + go !mv !i !o + | i >= count = do + vec <- V.unsafeFreeze mv + pure $! Just (TV.Map kt vt vec, o) + | otherwise = case decodeCompValue decodeStrings kt bs o of + Nothing -> pure Nothing + Just (k, o1) -> case decodeCompValue decodeStrings vt bs o1 of + Nothing -> pure Nothing + Just (v, o2) -> do + MV.unsafeWrite mv i (k, v) + go mv (i + 1) o2 +{-# INLINE decodeCompMapEntries #-} + + +decodeCompListEntries + :: Bool + -> ThriftType + -> Int + -> ByteString + -> Int + -> Maybe (TV.Value, Int) +decodeCompListEntries decodeStrings et !count bs !off + | count < 0 || off < 0 || off > BS.length bs || count > BS.length bs - off = Nothing + | count == 0 = Just (TV.List et V.empty, off) + | otherwise = runST $ do + mv <- MV.new count + go mv 0 off + where + go :: MV.MVector s TV.Value -> Int -> Int -> ST s (Maybe (TV.Value, Int)) + go !mv !i !o + | i >= count = do + vec <- V.unsafeFreeze mv + pure $! Just (TV.List et vec, o) + | otherwise = case decodeCompValue decodeStrings et bs o of + Nothing -> pure Nothing + Just (v, o') -> do + MV.unsafeWrite mv i v + go mv (i + 1) o' +{-# INLINE decodeCompListEntries #-} + + +decodeCompSetEntries + :: Bool + -> ThriftType + -> Int + -> ByteString + -> Int + -> Maybe (TV.Value, Int) +decodeCompSetEntries decodeStrings et !count bs !off + | count < 0 || off < 0 || off > BS.length bs || count > BS.length bs - off = Nothing + | count == 0 = Just (TV.Set et V.empty, off) + | otherwise = runST $ do + mv <- MV.new count + go mv 0 off + where + go :: MV.MVector s TV.Value -> Int -> Int -> ST s (Maybe (TV.Value, Int)) + go !mv !i !o + | i >= count = do + vec <- V.unsafeFreeze mv + pure $! Just (TV.Set et vec, o) + | otherwise = case decodeCompValue decodeStrings et bs o of + Nothing -> pure Nothing + Just (v, o') -> do + MV.unsafeWrite mv i v + go mv (i + 1) o' +{-# INLINE decodeCompSetEntries #-} diff --git a/src/Thrift/Derive.hs b/src/Thrift/Derive.hs new file mode 100644 --- /dev/null +++ b/src/Thrift/Derive.hs @@ -0,0 +1,462 @@ +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TemplateHaskell #-} + +{- | Annotation-driven Template Haskell deriver for Thrift +'Thrift.Class.ToThrift' / 'Thrift.Class.FromThrift' instances. + +Encoding shape: + +* 'TypeShapeNewtype' — pass-through. +* 'TypeShapeRecord' — Thrift @Struct@ with one entry per record + field. Field IDs default to the field's positional index + starting at @1@; an explicit @tag N@ modifier overrides it. +* 'TypeShapeEnum' — encoded as a Thrift @I32@. Each constructor + maps to its zero-based positional index unless overridden by + @tag N@. +* 'TypeShapeSum' — encoded as a Thrift @Struct@ where the + single set field's ID identifies the constructor (Thrift's + union convention). Constructor field IDs are positional unless + overridden by @tag N@. + +Modifiers honoured: 'tag' (field id / enum value override), +'rename' (ignored — Thrift is structural), 'skip', 'defaults', +'optional', 'required', 'coerced'. Modifiers irrelevant to Thrift +(e.g. 'renameStyle') are silently ignored on Thrift-only paths. +-} +module Thrift.Derive ( + deriveThrift, + deriveToThrift, + deriveFromThrift, +) where + +import Data.Coerce (coerce) +import Data.Foldable (foldlM) +import Data.Int (Int16, Int32) +import Data.Vector qualified as Vector +import Language.Haskell.TH +import Thrift.Class qualified as TC +import Thrift.Value qualified as TV +import Wireform.Derive.Backend +import Wireform.Derive.ModifierInfo +import Wireform.Derive.TypeInfo + + +-- --------------------------------------------------------------------------- +-- Public entry points +-- --------------------------------------------------------------------------- + +deriveThrift :: Name -> Q [Dec] +deriveThrift nm = (++) <$> deriveToThrift nm <*> deriveFromThrift nm + + +deriveToThrift :: Name -> Q [Dec] +deriveToThrift nm = do + ti <- reifyTypeInfo nm + body <- toThriftBody ti + let typ = applyTypeArgs (ConT (typeInfoName ti)) (typeInfoVarTypes ti) + decl = + InstanceD + Nothing + [] + (AppT (ConT ''TC.ToThrift) typ) + [FunD 'TC.toThrift [Clause [] (NormalB body) []]] + pure [decl] + + +deriveFromThrift :: Name -> Q [Dec] +deriveFromThrift nm = do + ti <- reifyTypeInfo nm + body <- fromThriftBody ti + let typ = applyTypeArgs (ConT (typeInfoName ti)) (typeInfoVarTypes ti) + decl = + InstanceD + Nothing + [] + (AppT (ConT ''TC.FromThrift) typ) + [FunD 'TC.fromThrift [Clause [] (NormalB body) []]] + pure [decl] + + +-- --------------------------------------------------------------------------- +-- ToThrift +-- --------------------------------------------------------------------------- + +toThriftBody :: TypeInfo -> Q Exp +toThriftBody ti = case typeInfoShape ti of + TypeShapeNewtype c -> toThriftNewtype c + TypeShapeRecord c -> toThriftRecord c + TypeShapeEnum cs -> toThriftEnum cs + TypeShapeSum cs -> toThriftSum cs + + +toThriftNewtype :: ConInfo -> Q Exp +toThriftNewtype c = case conInfoFields c of + [FieldInfo (Just sel) _] -> do + x <- newName "x" + lamE [varP x] [|TC.toThrift ($(varE sel) $(varE x))|] + [FieldInfo Nothing _] -> do + x <- newName "x" + lamE [conP (conInfoName c) [varP x]] [|TC.toThrift $(varE x)|] + _ -> fail "Thrift.Derive: newtype must have exactly one field" + + +toThriftRecord :: ConInfo -> Q Exp +toThriftRecord c = do + x <- newName "x" + pairs <- recordToThriftPairs (varE x) (conInfoFields c) + lamE + [varP x] + [|TV.Struct (Vector.fromList $(pure pairs))|] + + +recordToThriftPairs :: Q Exp -> [FieldInfo] -> Q Exp +recordToThriftPairs varExp fields = do + pairExpss <- mapM go (zip [1 ..] fields) + pure (ListE (concat pairExpss)) + where + go :: (Int16, FieldInfo) -> Q [Exp] + go (defaultId, FieldInfo mSel _) = do + selName <- requireSelector mSel + mi <- reifyModifierInfoFor backendThrift selName + if miSkip mi + then pure [] + else do + let fid = case miTag mi of + Just n -> fromIntegral n :: Int16 + Nothing -> defaultId + getter = appE (varE selName) varExp + encoded = case miCoerce mi of + Nothing -> [|TC.toThrift $getter|] + Just _ -> [|TC.toThrift (coerce $getter)|] + pair <- [|($(litE (integerL (fromIntegral fid))) :: Int16, $encoded)|] + pure [pair] + + +{- | Enums encoded as I32. Default values are zero-based positional +indices unless overridden by 'tag'. +-} +toThriftEnum :: [ConInfo] -> Q Exp +toThriftEnum cs = do + v <- newName "v" + matches <- mapM enumMatchTo (zip [0 ..] cs) + body <- caseE (varE v) (map pure matches) + lamE [varP v] (pure body) + where + enumMatchTo :: (Int32, ConInfo) -> Q Match + enumMatchTo (defaultIdx, c) = do + mi <- reifyModifierInfoFor backendThrift (conInfoName c) + let n = case miTag mi of + Just t -> fromIntegral t :: Int32 + Nothing -> defaultIdx + bodyE <- [|TV.I32 $(litE (integerL (fromIntegral n)))|] + pure (Match (ConP (conInfoName c) [] []) (NormalB bodyE) []) + + +{- | Sums encoded as Thrift unions: a Struct with exactly one entry +whose field ID identifies the constructor. +-} +toThriftSum :: [ConInfo] -> Q Exp +toThriftSum cs = do + v <- newName "v" + matches <- mapM sumMatchTo (zip [1 ..] cs) + body <- caseE (varE v) (map pure matches) + lamE [varP v] (pure body) + where + sumMatchTo :: (Int16, ConInfo) -> Q Match + sumMatchTo (defaultId, c) = do + mi <- reifyModifierInfoFor backendThrift (conInfoName c) + let fid = case miTag mi of + Just t -> fromIntegral t :: Int16 + Nothing -> defaultId + fieldNames <- mapM (\_ -> newName "f") (conInfoFields c) + let pat = ConP (conInfoName c) [] (map VarP fieldNames) + payload <- case fieldNames of + [] -> [|TV.Bool True|] + [n] -> [|TC.toThrift $(varE n)|] + ns -> + [| + TC.toThrift + ( $(pure (ListE (map (AppE (VarE 'TC.toThrift) . VarE) ns))) + :: [TV.Value] + ) + |] + bodyE <- + [| + TV.Struct + ( Vector.singleton + ($(litE (integerL (fromIntegral fid))) :: Int16, $(pure payload)) + ) + |] + pure (Match pat (NormalB bodyE) []) + + +-- --------------------------------------------------------------------------- +-- FromThrift +-- --------------------------------------------------------------------------- + +fromThriftBody :: TypeInfo -> Q Exp +fromThriftBody ti = case typeInfoShape ti of + TypeShapeNewtype c -> fromThriftNewtype c + TypeShapeRecord c -> fromThriftRecord c + TypeShapeEnum cs -> fromThriftEnum cs + TypeShapeSum cs -> fromThriftSum cs + + +fromThriftNewtype :: ConInfo -> Q Exp +fromThriftNewtype c = case conInfoFields c of + [FieldInfo _ _] -> [|fmap $(conE (conInfoName c)) . TC.fromThrift|] + _ -> fail "Thrift.Derive: newtype must have exactly one field" + + +fromThriftRecord :: ConInfo -> Q Exp +fromThriftRecord c = do + v <- newName "v" + kvs <- newName "kvs" + bodyE <- recordParser kvs (conInfoFields c) (conInfoName c) + lamE + [varP v] + ( caseE + (varE v) + [ match + (conP 'TV.Struct [varP kvs]) + (normalB (pure bodyE)) + [] + , match + wildP + ( normalB + [|Left "Thrift.Derive: expected Struct for record type"|] + ) + [] + ] + ) + + +recordParser :: Name -> [FieldInfo] -> Name -> Q Exp +recordParser kvs fields conName = case zip [1 ..] fields of + [] -> [|Right $(conE conName)|] + (p0 : ps) -> do + e0 <- fieldParser kvs p0 + hd <- [|$(conE conName) <$> $(pure e0)|] + foldlM + ( \acc fp -> do + ef <- fieldParser kvs fp + [|$(pure acc) <*> $(pure ef)|] + ) + hd + ps + + +fieldParser :: Name -> (Int16, FieldInfo) -> Q Exp +fieldParser kvs (defaultId, FieldInfo mSel _) = do + selName <- requireSelector mSel + mi <- reifyModifierInfoFor backendThrift selName + if miSkip mi + then case miDefaults mi of + Just defNm -> [|Right $(varE defNm)|] + Nothing -> + [| + Left + ( "Thrift.Derive: missing 'defaults' for skipped field " + ++ $(litE (stringL (nameBase selName))) + ) + |] + else do + let fid = case miTag mi of + Just t -> fromIntegral t :: Int16 + Nothing -> defaultId + isOptional = miRequired mi == Just False + base <- + if isOptional + then + [| + case lookupThriftField $(litE (integerL (fromIntegral fid))) $(varE kvs) of + Nothing -> Right Nothing + Just v -> fmap Just (TC.fromThrift v) + |] + else + [| + case lookupThriftField $(litE (integerL (fromIntegral fid))) $(varE kvs) of + Nothing -> + Left + ( "Thrift.Derive: missing field id " + ++ show ($(litE (integerL (fromIntegral fid))) :: Int16) + ) + Just v -> TC.fromThrift v + |] + case miCoerce mi of + Nothing -> pure base + Just _ -> [|fmap coerce $(pure base)|] + + +lookupThriftField :: Int16 -> Vector.Vector (Int16, TV.Value) -> Maybe TV.Value +lookupThriftField fid kvs = Vector.foldr step Nothing kvs + where + step (k, v) acc + | k == fid = Just v + | otherwise = acc + + +fromThriftEnum :: [ConInfo] -> Q Exp +fromThriftEnum cs = do + v <- newName "v" + i <- newName "i" + branches <- mapM (enumMatchFrom i) (zip [0 ..] cs) + let multi = + MultiIfE + ( map (\(g, e) -> (g, AppE (ConE 'Right) e)) branches + ++ [ + ( NormalG (ConE 'True) + , AppE + (ConE 'Left) + ( AppE + ( AppE + (VarE 'mappend) + (LitE (StringL "Thrift.Derive: unknown enum value ")) + ) + (AppE (VarE 'show) (VarE i)) + ) + ) + ] + ) + lamE + [varP v] + ( caseE + (varE v) + [ match + (conP 'TV.I32 [varP i]) + (normalB (pure multi)) + [] + , match + wildP + (normalB [|Left "Thrift.Derive: enum expected I32"|]) + [] + ] + ) + where + enumMatchFrom :: Name -> (Int32, ConInfo) -> Q (Guard, Exp) + enumMatchFrom iVar (defaultIdx, c) = do + mi <- reifyModifierInfoFor backendThrift (conInfoName c) + let n = case miTag mi of + Just t -> fromIntegral t :: Int32 + Nothing -> defaultIdx + guardE = + InfixE + (Just (VarE iVar)) + (VarE '(==)) + (Just (LitE (IntegerL (fromIntegral n)))) + pure (NormalG guardE, ConE (conInfoName c)) + + +fromThriftSum :: [ConInfo] -> Q Exp +fromThriftSum cs = do + v <- newName "v" + kvs <- newName "kvs" + fid <- newName "fid" + pay <- newName "payload" + branches <- mapM (sumMatchFrom fid pay) (zip [1 ..] cs) + let multi = + MultiIfE + ( branches + ++ [ + ( NormalG (ConE 'True) + , AppE + (ConE 'Left) + ( AppE + ( AppE + (VarE 'mappend) + (LitE (StringL "Thrift.Derive: unknown sum field id ")) + ) + (AppE (VarE 'show) (VarE fid)) + ) + ) + ] + ) + lamE + [varP v] + ( caseE + (varE v) + [ match + (conP 'TV.Struct [varP kvs]) + ( normalB + [| + case Vector.toList $(varE kvs) of + [(fidLocal, payloadLocal)] -> + let $(varP fid) = fidLocal + $(varP pay) = payloadLocal + in $(pure multi) + _ -> Left "Thrift.Derive: sum struct must have exactly one field" + |] + ) + [] + , match + wildP + (normalB [|Left "Thrift.Derive: sum expected Struct"|]) + [] + ] + ) + where + sumMatchFrom :: Name -> Name -> (Int16, ConInfo) -> Q (Guard, Exp) + sumMatchFrom fidVar payVar (defaultId, c) = do + mi <- reifyModifierInfoFor backendThrift (conInfoName c) + let fid = case miTag mi of + Just t -> fromIntegral t :: Int16 + Nothing -> defaultId + guardE = + InfixE + (Just (VarE fidVar)) + (VarE '(==)) + (Just (LitE (IntegerL (fromIntegral fid)))) + bodyE <- case conInfoFields c of + [] -> [|Right $(conE (conInfoName c))|] + [_one] -> [|fmap $(conE (conInfoName c)) (TC.fromThrift $(varE payVar))|] + many -> sumNAry payVar (conInfoName c) (length many) + pure (NormalG guardE, bodyE) + + +sumNAry :: Name -> Name -> Int -> Q Exp +sumNAry payVar conName arity = do + arr <- newName "arr" + let parseI :: Int -> Q Exp + parseI i = [|TC.fromThrift ($(varE arr) !! $(litE (integerL (fromIntegral i))))|] + hd <- do + e0 <- parseI 0 + [|$(conE conName) <$> $(pure e0)|] + body <- + foldlM + ( \acc i -> do + ei <- parseI i + [|$(pure acc) <*> $(pure ei)|] + ) + hd + [1 .. arity - 1] + let conNameStr = nameBase conName + [| + do + $(varP arr) <- TC.fromThrift $(varE payVar) :: Either String [TV.Value] + if length $(varE arr) /= $(litE (integerL (fromIntegral arity))) + then + Left + ( "Thrift.Derive: " + ++ conNameStr + ++ " expected " + ++ show ($(litE (integerL (fromIntegral arity))) :: Int) + ++ " contents, got " + ++ show (length $(varE arr)) + ) + else $(pure body) + |] + + +-- --------------------------------------------------------------------------- +-- Helpers +-- --------------------------------------------------------------------------- + +requireSelector :: Maybe Name -> Q Name +requireSelector (Just n) = pure n +requireSelector Nothing = + fail "Thrift.Derive: cannot derive Thrift for non-record positional field" + + +applyTypeArgs :: Type -> [Type] -> Type +applyTypeArgs = foldl AppT diff --git a/src/Thrift/Encode.hs b/src/Thrift/Encode.hs new file mode 100644 --- /dev/null +++ b/src/Thrift/Encode.hs @@ -0,0 +1,489 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE UnboxedSums #-} +{-# LANGUAGE UnboxedTuples #-} + +{- | High-level Thrift encoding for Binary and Compact protocols. + +Converts a 'Thrift.Value.Value' tree into a wire-format 'ByteString'. +Uses direct buffer writes via 'Proto.Encode.Direct.directEncode' to +avoid Builder allocation overhead. + +@ +import Thrift.Encode (encodeBinary, encodeCompact) +import qualified Thrift.Value as T +import qualified Data.Vector as V + +let person = T.Struct (V.fromList [(1, T.String \"Alice\")]) +let binBytes = encodeBinary person +let compactBytes = encodeCompact person +@ +-} +module Thrift.Encode ( + encodeBinary, + encodeBinaryBounded, + encodeCompact, + encodeCompactBounded, + binaryEncodedSizeBounded, + compactEncodedSizeBounded, +) where + +import Data.Bits (countLeadingZeros, shiftL, shiftR, xor, (.&.), (.|.)) +import Data.ByteString (ByteString) +import Data.ByteString qualified as BS +import Data.ByteString.Internal qualified as BSI +import Data.Int (Int16, Int32, Int64) +import Data.List (sortBy) +import Data.Ord (comparing) +import Data.Text.Encoding qualified as TE +import Data.Text.Foreign qualified as TF +import Data.Vector qualified as V +import Data.Word (Word16, Word32, Word64, Word8, byteSwap64) +import Foreign.ForeignPtr (withForeignPtr) +import Foreign.Ptr (Ptr, plusPtr) +import Foreign.Storable (pokeByteOff) +import GHC.ByteOrder (ByteOrder (..), targetByteOrder) +import GHC.Exts (Int (I#), Int#, writeWord8OffAddrAsWord64#, (+#), (-#), (<#), (==#), (>#)) +import GHC.Float (castDoubleToWord64) +import GHC.IO (IO (..)) +import GHC.Ptr (Ptr (Ptr)) +import GHC.Word (Word64 (W64#)) +import Thrift.Value qualified as TV +import Thrift.Wire (ThriftType (..), thriftTypeToBin, thriftTypeToCompact) +import Wireform.Encode.Direct (directEncode) + + +-------------------------------------------------------------------------------- +-- Binary Protocol +-------------------------------------------------------------------------------- + +encodeBinary :: TV.Value -> ByteString +encodeBinary = either error id . encodeBinaryBounded maxBound +{-# INLINE encodeBinary #-} + + +-- | Validate the complete value and byte budget before allocating output. +encodeBinaryBounded :: Int -> TV.Value -> Either String ByteString +encodeBinaryBounded limit value = do + size <- binaryEncodedSizeBounded limit value + Right (directEncode size (writeBin value)) + + +-- | Measure exactly without encoding strings or allocating output buffers. +binaryEncodedSizeBounded :: Int -> TV.Value -> Either String Int +binaryEncodedSizeBounded = checkedSize measureBinary# + + +writeBin :: TV.Value -> Ptr Word8 -> Int -> IO Int +writeBin = writeBinValue +{-# INLINE writeBin #-} + + +writeBinValue :: TV.Value -> Ptr Word8 -> Int -> IO Int +writeBinValue val p off = case val of + TV.Bool b -> do + pokeByteOff p off (if b then 1 :: Word8 else 0) + pure $! off + 1 + TV.Byte v -> do + pokeByteOff p off (fromIntegral v :: Word8) + pure $! off + 1 + TV.I16 v -> writeBE16 p off (fromIntegral v) + TV.I32 v -> writeBE32 p off (fromIntegral v) + TV.I64 v -> writeBE64 p off (fromIntegral v) + TV.Double d -> writeBE64 p off (castDoubleToWord64 d) + TV.String t -> do + let !bs = TE.encodeUtf8 t + writeBinBytes p off bs + TV.Binary b -> writeBinBytes p off b + TV.UUID b -> writeRaw p off b + TV.Struct fields -> do + let sorted = sortBy (comparing fst) (V.toList fields) + off1 <- writeBinStructFields p off sorted + pokeByteOff p off1 (0x00 :: Word8) + pure $! off1 + 1 + TV.Map kt vt entries -> do + pokeByteOff p off (thriftTypeToBin kt) + pokeByteOff p (off + 1) (thriftTypeToBin vt) + off1 <- writeBE32 p (off + 2) (fromIntegral (V.length entries) :: Word32) + V.foldM' (\o (k, v) -> do o1 <- writeBinValue k p o; writeBinValue v p o1) off1 entries + TV.List et elems -> do + pokeByteOff p off (thriftTypeToBin et) + off1 <- writeBE32 p (off + 1) (fromIntegral (V.length elems) :: Word32) + V.foldM' (\o v -> writeBinValue v p o) off1 elems + TV.Set et elems -> do + pokeByteOff p off (thriftTypeToBin et) + off1 <- writeBE32 p (off + 1) (fromIntegral (V.length elems) :: Word32) + V.foldM' (\o v -> writeBinValue v p o) off1 elems + + +writeBinStructFields :: Ptr Word8 -> Int -> [(Int16, TV.Value)] -> IO Int +writeBinStructFields _ off [] = pure off +writeBinStructFields p off ((fid, v) : rest) = do + pokeByteOff p off (thriftTypeToBin (TV.thriftTypeOf v)) + off1 <- writeBE16 p (off + 1) (fromIntegral fid :: Word16) + off2 <- writeBinValue v p off1 + writeBinStructFields p off2 rest + + +writeBinBytes :: Ptr Word8 -> Int -> ByteString -> IO Int +writeBinBytes p off bs = do + off1 <- writeBE32 p off (fromIntegral (BS.length bs) :: Word32) + writeRaw p off1 bs +{-# INLINE writeBinBytes #-} + + +-------------------------------------------------------------------------------- +-- Compact Protocol +-------------------------------------------------------------------------------- + +encodeCompact :: TV.Value -> ByteString +encodeCompact = either error id . encodeCompactBounded maxBound +{-# INLINE encodeCompact #-} + + +-- | Encode only after exact sizing and validation fit the requested budget. +encodeCompactBounded :: Int -> TV.Value -> Either String ByteString +encodeCompactBounded limit value = do + size <- compactEncodedSizeBounded limit value + Right (directEncode size (writeComp value)) + + +compactEncodedSizeBounded :: Int -> TV.Value -> Either String Int +compactEncodedSizeBounded = checkedSize measureCompact# + + +type Measure = Int# -> TV.Value -> (# String | Int# #) + + +-- Subtract from a remaining budget instead of adding sizes that can overflow. +-- Recursive workers return raw counters/results, never a per-value Either. +checkedSize :: Measure -> Int -> TV.Value -> Either String Int +checkedSize measure (I# limit) value = case limit <# 0# of + 1# -> Left "Thrift.Encode: negative size limit" + _ -> case measure limit value of + (# message | #) -> Left message + (# | remaining #) -> Right (I# (limit -# remaining)) +{-# INLINE checkedSize #-} + + +spend# :: Int -> Int# -> (# String | Int# #) +spend# (I# bytes) remaining = case bytes <# 0# of + 1# -> (# "Thrift.Encode: negative size" | #) + _ -> case bytes ># remaining of + 1# -> (# "Thrift.Encode: encoded value exceeds size limit" | #) + _ -> (# | remaining -# bytes #) +{-# INLINE spend# #-} + + +countFits :: Word64 -> Int -> Bool +countFits maximumCount count = count >= 0 && fromIntegral count <= maximumCount +{-# INLINE countFits #-} + + +measurePayload# :: Word64 -> Int -> Int -> Int# -> (# String | Int# #) +measurePayload# maximumCount prefixSize size remaining + | not (countFits maximumCount size) = (# "Thrift.Encode: byte length exceeds protocol limit" | #) + | otherwise = case spend# prefixSize remaining of + (# message | #) -> (# message | #) + (# | afterPrefix #) -> spend# size afterPrefix +{-# INLINE measurePayload# #-} + + +measureUUID# :: ByteString -> Int# -> (# String | Int# #) +measureUUID# bytes remaining + | BS.length bytes /= 16 = (# "Thrift.Encode: UUID must contain exactly 16 bytes" | #) + | otherwise = spend# 16 remaining +{-# INLINE measureUUID# #-} + + +measureBinary# :: Measure +measureBinary# remaining value = case value of + TV.Bool _ -> spend# 1 remaining + TV.Byte _ -> spend# 1 remaining + TV.I16 _ -> spend# 2 remaining + TV.I32 _ -> spend# 4 remaining + TV.I64 _ -> spend# 8 remaining + TV.Double _ -> spend# 8 remaining + TV.String text -> measurePayload# 0x7fffffff 4 (TF.lengthWord8 text) remaining + TV.Binary bytes -> measurePayload# 0x7fffffff 4 (BS.length bytes) remaining + TV.UUID bytes -> measureUUID# bytes remaining + TV.Struct fields -> case spend# 1 remaining of + (# message | #) -> (# message | #) + (# | rest #) -> + if V.length fields > I# rest `quot` 3 + then (# "Thrift.Encode: encoded value exceeds size limit" | #) + else measureFields# False 0 (sortBy (comparing fst) (V.toList fields)) rest + TV.Map keyType valueType entries -> measureMap# False keyType valueType entries remaining + TV.List elementType elements -> measureSequence# False elementType elements remaining + TV.Set elementType elements -> measureSequence# False elementType elements remaining + + +measureCompact# :: Measure +measureCompact# remaining value = case value of + TV.Bool _ -> spend# 1 remaining + TV.Byte _ -> spend# 1 remaining + TV.I16 n -> spend# (compVarintSize (fromIntegral (zigZagEncode32 (fromIntegral n)))) remaining + TV.I32 n -> spend# (compVarintSize (fromIntegral (zigZagEncode32 n))) remaining + TV.I64 n -> spend# (compVarintSize (zigZagEncode64 n)) remaining + TV.Double _ -> spend# 8 remaining + TV.String text -> bytes (TF.lengthWord8 text) + TV.Binary valueBytes -> bytes (BS.length valueBytes) + TV.UUID valueBytes -> measureUUID# valueBytes remaining + TV.Struct fields -> case spend# 1 remaining of + (# message | #) -> (# message | #) + (# | rest #) -> + if V.length fields > I# rest + then (# "Thrift.Encode: encoded value exceeds size limit" | #) + else measureFields# True 0 (sortBy (comparing fst) (V.toList fields)) rest + TV.Map keyType valueType entries -> measureMap# True keyType valueType entries remaining + TV.List elementType elements -> measureSequence# True elementType elements remaining + TV.Set elementType elements -> measureSequence# True elementType elements remaining + where + bytes size = measurePayload# 0xffffffff (compVarintSize (fromIntegral size)) size remaining + + +measureFields# :: Bool -> Int16 -> [(Int16, TV.Value)] -> Int# -> (# String | Int# #) +measureFields# _ _ [] remaining = (# | remaining #) +measureFields# compact lastFid ((fid, value) : fields) remaining = + let !delta = fromIntegral fid - fromIntegral lastFid :: Int + !headerSize = if not compact then 3 else if delta > 0 && delta <= 15 then 1 else 1 + compVarintSize (fromIntegral (zigZagEncode32 (fromIntegral fid))) + in case spend# headerSize remaining of + (# message | #) -> (# message | #) + (# | rest #) -> case value of + TV.Bool _ | compact -> measureFields# compact fid fields rest + _ -> case (if compact then measureCompact# else measureBinary#) rest value of + (# message | #) -> (# message | #) + (# | afterValue #) -> measureFields# compact fid fields afterValue + + +measureSequence# :: Bool -> ThriftType -> V.Vector TV.Value -> Int# -> (# String | Int# #) +measureSequence# compact elementType elements remaining + | not (countFits (if compact then 0xffffffff else 0x7fffffff) count) = (# "Thrift.Encode: container length exceeds protocol limit" | #) + | otherwise = case spend# headerSize remaining of + (# message | #) -> (# message | #) + (# | rest #) -> + if count > I# rest + then (# "Thrift.Encode: encoded value exceeds size limit" | #) + else + let go index budget = case index ==# count# of + 1# -> (# | budget #) + _ -> + let value = V.unsafeIndex elements (I# index) + in if TV.thriftTypeOf value /= elementType + then (# "Thrift.Encode: container element type mismatch" | #) + else case (if compact then measureCompact# else measureBinary#) budget value of + (# message | #) -> (# message | #) + (# | next #) -> go (index +# 1#) next + in go 0# rest + where + !count@(I# count#) = V.length elements + !headerSize = if not compact then 5 else if count < 15 then 1 else 1 + compVarintSize (fromIntegral count) + + +measureMap# :: Bool -> ThriftType -> ThriftType -> V.Vector (TV.Value, TV.Value) -> Int# -> (# String | Int# #) +measureMap# compact keyType valueType entries remaining + | not (countFits (if compact then 0xffffffff else 0x7fffffff) count) = (# "Thrift.Encode: container length exceeds protocol limit" | #) + | otherwise = case spend# headerSize remaining of + (# message | #) -> (# message | #) + (# | rest #) -> + if count > I# rest `quot` 2 + then (# "Thrift.Encode: encoded value exceeds size limit" | #) + else + let go index budget = case index ==# count# of + 1# -> (# | budget #) + _ -> + let (key, value) = V.unsafeIndex entries (I# index) + in if TV.thriftTypeOf key /= keyType || TV.thriftTypeOf value /= valueType + then (# "Thrift.Encode: map key/value type mismatch" | #) + else case (if compact then measureCompact# else measureBinary#) budget key of + (# message | #) -> (# message | #) + (# | afterKey #) -> case (if compact then measureCompact# else measureBinary#) afterKey value of + (# message | #) -> (# message | #) + (# | next #) -> go (index +# 1#) next + in go 0# rest + where + !count@(I# count#) = V.length entries + !headerSize = if not compact then 6 else if count == 0 then 1 else 1 + compVarintSize (fromIntegral count) + + +compVarintSize :: Word64 -> Int +compVarintSize !n = + let !bits = 64 - countLeadingZeros (n .|. 1) + in (bits + 6) `quot` 7 +{-# INLINE compVarintSize #-} + + +zigZagEncode32 :: Int32 -> Word32 +zigZagEncode32 !n = fromIntegral ((n `shiftL` 1) `xor` (n `shiftR` 31)) +{-# INLINE zigZagEncode32 #-} + + +zigZagEncode64 :: Int64 -> Word64 +zigZagEncode64 !n = fromIntegral ((n `shiftL` 1) `xor` (n `shiftR` 63)) +{-# INLINE zigZagEncode64 #-} + + +writeComp :: TV.Value -> Ptr Word8 -> Int -> IO Int +writeComp val p off = case val of + TV.Struct fields -> writeCompStruct p off fields + _ -> writeCompValue val p off +{-# INLINE writeComp #-} + + +writeCompStruct :: Ptr Word8 -> Int -> V.Vector (Int16, TV.Value) -> IO Int +writeCompStruct p off fields = do + let sorted = sortBy (comparing fst) (V.toList fields) + off1 <- goWriteCompStruct p off 0 sorted + pokeByteOff p off1 (0x00 :: Word8) + pure $! off1 + 1 + + +goWriteCompStruct :: Ptr Word8 -> Int -> Int16 -> [(Int16, TV.Value)] -> IO Int +goWriteCompStruct _ off _ [] = pure off +goWriteCompStruct p off !lastFid ((fid, v) : rest) = do + let !delta = fid - lastFid + !ctype = case v of + TV.Bool b -> if b then 1 else 2 + _ -> thriftTypeToCompact (TV.thriftTypeOf v) + off1 <- + if delta > 0 && delta <= 15 + then do + pokeByteOff p off (fromIntegral delta `shiftL` 4 .|. ctype :: Word8) + pure $! off + 1 + else do + pokeByteOff p off ctype + writeVarint p (off + 1) (fromIntegral (zigZagEncode32 (fromIntegral fid))) + off2 <- case v of + TV.Bool _ -> pure off1 + _ -> writeCompValue v p off1 + goWriteCompStruct p off2 fid rest + + +writeCompValue :: TV.Value -> Ptr Word8 -> Int -> IO Int +writeCompValue val p off = case val of + TV.Bool b -> do + pokeByteOff p off (if b then 1 :: Word8 else 2) + pure $! off + 1 + TV.Byte v -> do + pokeByteOff p off (fromIntegral v :: Word8) + pure $! off + 1 + TV.I16 v -> writeVarint p off (fromIntegral (zigZagEncode32 (fromIntegral v))) + TV.I32 v -> writeVarint p off (fromIntegral (zigZagEncode32 v)) + TV.I64 v -> writeVarint p off (zigZagEncode64 v) + TV.Double d -> writeLE64 p off (castDoubleToWord64 d) + TV.String t -> do + let !bs = TE.encodeUtf8 t + writeCompBytes p off bs + TV.Binary b -> writeCompBytes p off b + TV.UUID b -> writeRaw p off b + TV.Struct fields -> writeCompStruct p off fields + TV.Map kt vt entries -> + if V.null entries + then do pokeByteOff p off (0x00 :: Word8); pure $! off + 1 + else do + off1 <- writeVarint p off (fromIntegral (V.length entries)) + pokeByteOff p off1 (thriftTypeToCompact kt `shiftL` 4 .|. thriftTypeToCompact vt :: Word8) + V.foldM' (\o (k, v) -> do o1 <- writeCompValue k p o; writeCompValue v p o1) (off1 + 1) entries + TV.List et elems -> do + off1 <- writeCompListHeader p off et (V.length elems) + V.foldM' (\o v -> writeCompValue v p o) off1 elems + TV.Set et elems -> do + off1 <- writeCompListHeader p off et (V.length elems) + V.foldM' (\o v -> writeCompValue v p o) off1 elems + + +writeCompBytes :: Ptr Word8 -> Int -> ByteString -> IO Int +writeCompBytes p off bs = do + off1 <- writeVarint p off (fromIntegral (BS.length bs)) + writeRaw p off1 bs +{-# INLINE writeCompBytes #-} + + +writeCompListHeader :: Ptr Word8 -> Int -> ThriftType -> Int -> IO Int +writeCompListHeader p off et sz = do + let !ctype = thriftTypeToCompact et + if sz < 15 + then do pokeByteOff p off (fromIntegral sz `shiftL` 4 .|. ctype :: Word8); pure $! off + 1 + else do pokeByteOff p off (0xF0 .|. ctype :: Word8); writeVarint p (off + 1) (fromIntegral sz) +{-# INLINE writeCompListHeader #-} + + +-------------------------------------------------------------------------------- +-- Shared write helpers +-------------------------------------------------------------------------------- + +writeBE16 :: Ptr Word8 -> Int -> Word16 -> IO Int +writeBE16 p off w = do + pokeByteOff p off (fromIntegral (w `shiftR` 8) :: Word8) + pokeByteOff p (off + 1) (fromIntegral (w .&. 0xFF) :: Word8) + pure $! off + 2 +{-# INLINE writeBE16 #-} + + +writeBE32 :: Ptr Word8 -> Int -> Word32 -> IO Int +writeBE32 p off w = do + pokeByteOff p off (fromIntegral (w `shiftR` 24) :: Word8) + pokeByteOff p (off + 1) (fromIntegral ((w `shiftR` 16) .&. 0xFF) :: Word8) + pokeByteOff p (off + 2) (fromIntegral ((w `shiftR` 8) .&. 0xFF) :: Word8) + pokeByteOff p (off + 3) (fromIntegral (w .&. 0xFF) :: Word8) + pure $! off + 4 +{-# INLINE writeBE32 #-} + + +writeBE64 :: Ptr Word8 -> Int -> Word64 -> IO Int +writeBE64 p off w = do + pokeByteOff p off (fromIntegral (w `shiftR` 56) :: Word8) + pokeByteOff p (off + 1) (fromIntegral ((w `shiftR` 48) .&. 0xFF) :: Word8) + pokeByteOff p (off + 2) (fromIntegral ((w `shiftR` 40) .&. 0xFF) :: Word8) + pokeByteOff p (off + 3) (fromIntegral ((w `shiftR` 32) .&. 0xFF) :: Word8) + pokeByteOff p (off + 4) (fromIntegral ((w `shiftR` 24) .&. 0xFF) :: Word8) + pokeByteOff p (off + 5) (fromIntegral ((w `shiftR` 16) .&. 0xFF) :: Word8) + pokeByteOff p (off + 6) (fromIntegral ((w `shiftR` 8) .&. 0xFF) :: Word8) + pokeByteOff p (off + 7) (fromIntegral (w .&. 0xFF) :: Word8) + pure $! off + 8 +{-# INLINE writeBE64 #-} + + +writeLE64 :: Ptr Word8 -> Int -> Word64 -> IO Int +writeLE64 (Ptr address) (I# offset) value = + let !little = case targetByteOrder of LittleEndian -> value; BigEndian -> byteSwap64 value + in case little of + W64# word -> IO $ \state -> + case writeWord8OffAddrAsWord64# address offset word state of + next -> (# next, I# (offset +# 8#) #) +{-# INLINE writeLE64 #-} + + +writeRaw :: Ptr Word8 -> Int -> ByteString -> IO Int +writeRaw p off (BSI.BS fp len) = do + withForeignPtr fp $ \src -> BSI.memcpy (p `plusPtr` off) src len + pure $! off + len +{-# INLINE writeRaw #-} + + +writeVarint :: Ptr Word8 -> Int -> Word64 -> IO Int +writeVarint !p !off !n + | n < 0x80 = do + pokeByteOff p off (fromIntegral n :: Word8) + pure $! off + 1 + | n < 0x4000 = do + pokeByteOff p off (fromIntegral (n .&. 0x7F .|. 0x80) :: Word8) + pokeByteOff p (off + 1) (fromIntegral (n `shiftR` 7) :: Word8) + pure $! off + 2 + | n < 0x200000 = do + pokeByteOff p off (fromIntegral (n .&. 0x7F .|. 0x80) :: Word8) + pokeByteOff p (off + 1) (fromIntegral ((n `shiftR` 7) .&. 0x7F .|. 0x80) :: Word8) + pokeByteOff p (off + 2) (fromIntegral (n `shiftR` 14) :: Word8) + pure $! off + 3 + | otherwise = writeVarintSlow p off n +{-# INLINE writeVarint #-} + + +writeVarintSlow :: Ptr Word8 -> Int -> Word64 -> IO Int +writeVarintSlow !p !off !n + | n < 0x80 = do + pokeByteOff p off (fromIntegral n :: Word8) + pure $! off + 1 + | otherwise = do + pokeByteOff p off (fromIntegral (n .&. 0x7F .|. 0x80) :: Word8) + writeVarintSlow p (off + 1) (n `shiftR` 7) diff --git a/src/Thrift/Encoding.hs b/src/Thrift/Encoding.hs new file mode 100644 --- /dev/null +++ b/src/Thrift/Encoding.hs @@ -0,0 +1,42 @@ +{- | Direct-to-AST Thrift encoding, mirroring aeson's @toEncoding@ +approach. + +Thrift's wire formats (binary and compact) require both a +container's field id and the inferred 'ThriftType' of every field +before any bytes can be emitted; building an 'Encoding' that +streams bytes directly would force the caller to pre-commit to a +protocol. We therefore expose 'Encoding' as a thin wrapper around +'Thrift.Value.Value' so 'Thrift.Class.toEncoding' is uniform with +the other formats and a future direct-write path can be slotted +in without an API break. +-} +module Thrift.Encoding ( + Encoding (..), + encodingToValue, + encodingToBinaryByteString, + encodingToCompactByteString, + value, +) where + +import Data.ByteString (ByteString) +import Thrift.Encode qualified as TE +import Thrift.Value qualified as TV + + +newtype Encoding = Encoding {runEncoding :: TV.Value} + + +encodingToValue :: Encoding -> TV.Value +encodingToValue = runEncoding + + +encodingToBinaryByteString :: Encoding -> ByteString +encodingToBinaryByteString = TE.encodeBinary . runEncoding + + +encodingToCompactByteString :: Encoding -> ByteString +encodingToCompactByteString = TE.encodeCompact . runEncoding + + +value :: TV.Value -> Encoding +value = Encoding diff --git a/src/Thrift/JSON.hs b/src/Thrift/JSON.hs new file mode 100644 --- /dev/null +++ b/src/Thrift/JSON.hs @@ -0,0 +1,536 @@ +{-# OPTIONS_GHC -Wno-orphans #-} + +{- | Thrift JSON protocol encoding/decoding. + +Provides two JSON representations: + +* /Simple JSON/ ('thriftToJSON' \/ 'thriftFromJSON'): primitives as direct + JSON values, structs as objects keyed by field ID, containers as arrays. + Decoding is template-guided (a sample 'TV.Value' provides the type + structure). + +* /Typed JSON/ ('thriftToTypedJSON' \/ 'thriftFromTypedJSON'): each value is + wrapped in a type-tagged object (TJSONProtocol-like), enabling decoding + without a separate schema. +-} +module Thrift.JSON ( + -- * Simple JSON + thriftToJSON, + thriftFromJSON, + + -- * Typed JSON (TJSONProtocol-like) + thriftToTypedJSON, + thriftFromTypedJSON, +) where + +import Data.Aeson qualified as Aeson +import Data.Aeson.Key qualified as Key +import Data.Aeson.KeyMap qualified as KM +import Data.ByteString qualified as BS +import Data.ByteString.Base16 qualified as Base16 +import Data.ByteString.Base64 qualified as Base64 +import Data.Int (Int16, Int64, Int8) +import Data.List (sortBy) +import Data.Ord (comparing) +import Data.Scientific qualified as Sci +import Data.Text (Text) +import Data.Text qualified as Text +import Data.Text.Encoding qualified as TE +import Data.Text.Read qualified as TR +import Data.Vector qualified as V +import Thrift.Value qualified as TV +import Thrift.Wire (ThriftType (..)) + + +-------------------------------------------------------------------------------- +-- Simple JSON encoding +-------------------------------------------------------------------------------- + +-- | Encode a 'TV.Value' as simple JSON. +thriftToJSON :: TV.Value -> Aeson.Value +thriftToJSON = \case + TV.Bool b -> Aeson.Bool b + TV.Byte v -> Aeson.Number (fromIntegral v) + TV.I16 v -> Aeson.Number (fromIntegral v) + TV.I32 v -> Aeson.Number (fromIntegral v) + TV.I64 v -> Aeson.String (Text.pack (show v)) + TV.Double d -> Aeson.Number (Sci.fromFloatDigits d) + TV.String t -> Aeson.String t + TV.Binary bs -> Aeson.String (TE.decodeUtf8 (Base64.encode bs)) + TV.UUID bs -> Aeson.String (TE.decodeUtf8 (Base16.encode bs)) + TV.Struct fields -> + Aeson.object + [ (Key.fromText (Text.pack (show fid)), thriftToJSON v) + | (fid, v) <- V.toList fields + ] + TV.List _ elems -> + Aeson.Array (V.map thriftToJSON elems) + TV.Set _ elems -> + Aeson.Array (V.map thriftToJSON elems) + TV.Map _ _ entries -> + Aeson.Array + ( V.map + (\(k, v) -> Aeson.Array (V.fromList [thriftToJSON k, thriftToJSON v])) + entries + ) + + +-------------------------------------------------------------------------------- +-- Simple JSON decoding (template-guided) +-------------------------------------------------------------------------------- + +{- | Decode a JSON value into a 'TV.Value', using a template value that +supplies the expected type structure. +-} +thriftFromJSON :: TV.Value -> Aeson.Value -> Either String TV.Value +thriftFromJSON template json = case template of + TV.Bool _ -> case json of + Aeson.Bool b -> Right (TV.Bool b) + _ -> Left "Expected JSON bool" + TV.Byte _ -> case json of + Aeson.Number n -> case Sci.toBoundedInteger n of + Just v -> Right (TV.Byte (v :: Int8)) + Nothing -> Left "Number out of range for byte" + _ -> Left "Expected JSON number for byte" + TV.I16 _ -> case json of + Aeson.Number n -> case Sci.toBoundedInteger n of + Just v -> Right (TV.I16 v) + Nothing -> Left "Number out of range for i16" + _ -> Left "Expected JSON number for i16" + TV.I32 _ -> case json of + Aeson.Number n -> case Sci.toBoundedInteger n of + Just v -> Right (TV.I32 v) + Nothing -> Left "Number out of range for i32" + _ -> Left "Expected JSON number for i32" + TV.I64 _ -> case json of + Aeson.String s -> case (TR.signed TR.decimal s :: Either String (Int64, Text)) of + Right (v, "") -> Right (TV.I64 v) + _ -> Left ("Invalid i64 string: " <> Text.unpack s) + Aeson.Number n -> case Sci.toBoundedInteger n of + Just v -> Right (TV.I64 (v :: Int64)) + Nothing -> Left "Number out of range for i64" + _ -> Left "Expected JSON string or number for i64" + TV.Double _ -> case json of + Aeson.Number n -> Right (TV.Double (Sci.toRealFloat n)) + _ -> Left "Expected JSON number for double" + TV.String _ -> case json of + Aeson.String t -> Right (TV.String t) + _ -> Left "Expected JSON string" + TV.Binary _ -> case json of + Aeson.String s -> case Base64.decode (TE.encodeUtf8 s) of + Right bs -> Right (TV.Binary bs) + Left err -> Left ("Invalid base64: " <> err) + _ -> Left "Expected JSON string (base64) for binary" + TV.UUID _ -> case json of + Aeson.String s -> case Base16.decode (TE.encodeUtf8 s) of + Right bs + | BS.length bs == 16 -> Right (TV.UUID bs) + | otherwise -> Left "UUID hex must decode to 16 bytes" + Left err -> Left ("Invalid hex for UUID: " <> err) + _ -> Left "Expected JSON string (hex) for UUID" + TV.Struct templateFields -> case json of + Aeson.Object obj -> do + let tmplMap = + [ (Text.pack (show fid), tmpl) + | (fid, tmpl) <- V.toList templateFields + ] + fields <- + mapM + ( \(key, val) -> do + let kt = Key.toText key + fid <- parseFieldId kt + case lookup kt tmplMap of + Just tmpl -> do + v <- thriftFromJSON tmpl val + Right (fid, v) + Nothing -> Left ("Unknown field: " <> Text.unpack kt) + ) + (KM.toList obj) + Right (TV.Struct (V.fromList (sortBy (comparing fst) fields))) + _ -> Left "Expected JSON object for struct" + TV.List et tmplElems -> case json of + Aeson.Array arr -> do + let etmpl = + if V.null tmplElems + then defaultForWireType et + else V.head tmplElems + elems <- V.mapM (thriftFromJSON etmpl) arr + Right (TV.List et elems) + _ -> Left "Expected JSON array for list" + TV.Set et tmplElems -> case json of + Aeson.Array arr -> do + let etmpl = + if V.null tmplElems + then defaultForWireType et + else V.head tmplElems + elems <- V.mapM (thriftFromJSON etmpl) arr + Right (TV.Set et elems) + _ -> Left "Expected JSON array for set" + TV.Map kt vt tmplEntries -> case json of + Aeson.Array arr -> do + let (kTmpl, vTmpl) = + if V.null tmplEntries + then (defaultForWireType kt, defaultForWireType vt) + else let (k, v) = V.head tmplEntries in (k, v) + entries <- + V.mapM + ( \pairJson -> case pairJson of + Aeson.Array pairArr + | V.length pairArr == 2 -> do + k <- thriftFromJSON kTmpl (pairArr V.! 0) + v <- thriftFromJSON vTmpl (pairArr V.! 1) + Right (k, v) + _ -> Left "Expected [key, value] pair in map" + ) + arr + Right (TV.Map kt vt entries) + _ -> Left "Expected JSON array of pairs for map" + + +-------------------------------------------------------------------------------- +-- Typed JSON encoding (TJSONProtocol-like) +-------------------------------------------------------------------------------- + +-- | Encode a 'TV.Value' as typed JSON (TJSONProtocol format). +thriftToTypedJSON :: TV.Value -> Aeson.Value +thriftToTypedJSON = \case + TV.Struct fields -> encodeTypedStruct fields + v -> wrapTyped v + + +encodeTypedStruct :: V.Vector (Int16, TV.Value) -> Aeson.Value +encodeTypedStruct fields = + Aeson.object + [ (Key.fromText (Text.pack (show fid)), wrapTyped v) + | (fid, v) <- V.toList fields + ] + + +wrapTyped :: TV.Value -> Aeson.Value +wrapTyped v = Aeson.object [(Key.fromText (typeAbbrev v), rawTyped v)] + + +rawTyped :: TV.Value -> Aeson.Value +rawTyped = \case + TV.Bool b -> Aeson.Bool b + TV.Byte v -> Aeson.Number (fromIntegral v) + TV.I16 v -> Aeson.Number (fromIntegral v) + TV.I32 v -> Aeson.Number (fromIntegral v) + TV.I64 v -> Aeson.String (Text.pack (show v)) + TV.Double d -> Aeson.Number (Sci.fromFloatDigits d) + TV.String t -> Aeson.String t + TV.Binary bs -> Aeson.String (TE.decodeUtf8 (Base64.encode bs)) + TV.UUID bs -> Aeson.String (TE.decodeUtf8 (Base16.encode bs)) + TV.Struct fields -> encodeTypedStruct fields + TV.List et elems -> + let etTag = containerElemTag et elems + in Aeson.Array + ( V.fromList + ( Aeson.String etTag + : Aeson.Number (fromIntegral (V.length elems)) + : V.toList (V.map rawTyped elems) + ) + ) + TV.Set et elems -> + let etTag = containerElemTag et elems + in Aeson.Array + ( V.fromList + ( Aeson.String etTag + : Aeson.Number (fromIntegral (V.length elems)) + : V.toList (V.map rawTyped elems) + ) + ) + TV.Map kt vt entries -> + let ktTag = if V.null entries then wireTypeTag kt else typeAbbrev (fst (V.head entries)) + vtTag = if V.null entries then wireTypeTag vt else typeAbbrev (snd (V.head entries)) + in Aeson.Array + ( V.fromList + [ Aeson.String ktTag + , Aeson.String vtTag + , Aeson.Number (fromIntegral (V.length entries)) + , Aeson.object + [ (Key.fromText (stringifyKey k), rawTyped v) + | (k, v) <- V.toList entries + ] + ] + ) + + +-------------------------------------------------------------------------------- +-- Typed JSON decoding +-------------------------------------------------------------------------------- + +-- | Decode a typed JSON value (TJSONProtocol format) into a 'TV.Value'. +thriftFromTypedJSON :: Aeson.Value -> Either String TV.Value +thriftFromTypedJSON json = case json of + Aeson.Object obj + | KM.null obj -> Right (TV.Struct V.empty) + | [(key, val)] <- KM.toList obj + , isTypeTag (Key.toText key) -> + decodeTypedWrapped (Key.toText key) val + | otherwise -> decodeTypedStruct obj + _ -> Left "Expected JSON object for typed Thrift value" + + +decodeTypedStruct :: KM.KeyMap Aeson.Value -> Either String TV.Value +decodeTypedStruct obj = do + fields <- + mapM + ( \(key, val) -> do + fid <- parseFieldId (Key.toText key) + v <- decodeWrappedField val + Right (fid, v) + ) + (KM.toList obj) + Right (TV.Struct (V.fromList (sortBy (comparing fst) fields))) + + +decodeWrappedField :: Aeson.Value -> Either String TV.Value +decodeWrappedField json = case json of + Aeson.Object obj + | [(key, val)] <- KM.toList obj -> + decodeTypedWrapped (Key.toText key) val + | otherwise -> Left "Expected single type-tag in wrapped field" + _ -> Left "Expected type-wrapped object for field" + + +decodeTypedWrapped :: Text -> Aeson.Value -> Either String TV.Value +decodeTypedWrapped tag val = case tag of + "tf" -> case val of + Aeson.Bool b -> Right (TV.Bool b) + _ -> Left "Expected bool for 'tf'" + "i8" -> case val of + Aeson.Number n -> case Sci.toBoundedInteger n of + Just v -> Right (TV.Byte v) + Nothing -> Left "Out of range for i8" + _ -> Left "Expected number for 'i8'" + "i16" -> case val of + Aeson.Number n -> case Sci.toBoundedInteger n of + Just v -> Right (TV.I16 v) + Nothing -> Left "Out of range for i16" + _ -> Left "Expected number for 'i16'" + "i32" -> case val of + Aeson.Number n -> case Sci.toBoundedInteger n of + Just v -> Right (TV.I32 v) + Nothing -> Left "Out of range for i32" + _ -> Left "Expected number for 'i32'" + "i64" -> case val of + Aeson.String s -> case (TR.signed TR.decimal s :: Either String (Int64, Text)) of + Right (v, "") -> Right (TV.I64 v) + _ -> Left ("Invalid i64: " <> Text.unpack s) + Aeson.Number n -> case Sci.toBoundedInteger n of + Just v -> Right (TV.I64 (v :: Int64)) + Nothing -> Left "Out of range for i64" + _ -> Left "Expected string or number for 'i64'" + "dbl" -> case val of + Aeson.Number n -> Right (TV.Double (Sci.toRealFloat n)) + _ -> Left "Expected number for 'dbl'" + "str" -> case val of + Aeson.String t -> Right (TV.String t) + _ -> Left "Expected string for 'str'" + "bin" -> case val of + Aeson.String s -> case Base64.decode (TE.encodeUtf8 s) of + Right bs -> Right (TV.Binary bs) + Left err -> Left ("Invalid base64: " <> err) + _ -> Left "Expected string for 'bin'" + "uuid" -> case val of + Aeson.String s -> case Base16.decode (TE.encodeUtf8 s) of + Right bs + | BS.length bs == 16 -> Right (TV.UUID bs) + | otherwise -> Left "UUID hex must be 16 bytes" + Left err -> Left ("Invalid hex: " <> err) + _ -> Left "Expected string for 'uuid'" + "rec" -> thriftFromTypedJSON val + "lst" -> decodeTypedList val + "set" -> decodeTypedSet val + "map" -> decodeTypedMap val + _ -> Left ("Unknown type tag: " <> Text.unpack tag) + + +decodeTypedList :: Aeson.Value -> Either String TV.Value +decodeTypedList json = case json of + Aeson.Array arr + | V.length arr >= 2 -> do + etTag <- parseTypeTagStr (arr V.! 0) + et <- wireTypeFromTag etTag + let elems = V.toList (V.drop 2 arr) + decoded <- mapM (decodeTypedWrapped etTag) elems + Right (TV.List et (V.fromList decoded)) + _ -> Left "Expected array (length >= 2) for typed list" + + +decodeTypedSet :: Aeson.Value -> Either String TV.Value +decodeTypedSet json = case json of + Aeson.Array arr + | V.length arr >= 2 -> do + etTag <- parseTypeTagStr (arr V.! 0) + et <- wireTypeFromTag etTag + let elems = V.toList (V.drop 2 arr) + decoded <- mapM (decodeTypedWrapped etTag) elems + Right (TV.Set et (V.fromList decoded)) + _ -> Left "Expected array (length >= 2) for typed set" + + +decodeTypedMap :: Aeson.Value -> Either String TV.Value +decodeTypedMap json = case json of + Aeson.Array arr + | V.length arr >= 3 -> do + ktTag <- parseTypeTagStr (arr V.! 0) + vtTag <- parseTypeTagStr (arr V.! 1) + kt <- wireTypeFromTag ktTag + vt <- wireTypeFromTag vtTag + if V.length arr < 4 + then Right (TV.Map kt vt V.empty) + else case arr V.! 3 of + Aeson.Object obj -> do + entries <- + mapM + ( \(k, v) -> do + key <- parseTypedKeyFromTag ktTag (Key.toText k) + val' <- decodeTypedWrapped vtTag v + Right (key, val') + ) + (KM.toList obj) + Right (TV.Map kt vt (V.fromList entries)) + _ -> Left "Expected object for map entries" + _ -> Left "Expected array (length >= 3) for typed map" + + +-------------------------------------------------------------------------------- +-- Internal helpers +-------------------------------------------------------------------------------- + +parseFieldId :: Text -> Either String Int16 +parseFieldId t = case (TR.signed TR.decimal t :: Either String (Integer, Text)) of + Right (n, "") -> Right (fromIntegral n) + _ -> Left ("Invalid field ID: " <> Text.unpack t) + + +parseTypeTagStr :: Aeson.Value -> Either String Text +parseTypeTagStr (Aeson.String t) = Right t +parseTypeTagStr _ = Left "Expected string type tag" + + +wireTypeFromTag :: Text -> Either String ThriftType +wireTypeFromTag = \case + "tf" -> Right TT_BOOL + "i8" -> Right TT_BYTE + "i16" -> Right TT_I16 + "i32" -> Right TT_I32 + "i64" -> Right TT_I64 + "dbl" -> Right TT_DOUBLE + "str" -> Right TT_STRING + "bin" -> Right TT_STRING + "rec" -> Right TT_STRUCT + "map" -> Right TT_MAP + "lst" -> Right TT_LIST + "set" -> Right TT_SET + "uuid" -> Right TT_UUID + t -> Left ("Unknown type tag: " <> Text.unpack t) + + +wireTypeTag :: ThriftType -> Text +wireTypeTag = \case + TT_BOOL -> "tf" + TT_BYTE -> "i8" + TT_I16 -> "i16" + TT_I32 -> "i32" + TT_I64 -> "i64" + TT_DOUBLE -> "dbl" + TT_STRING -> "str" + TT_STRUCT -> "rec" + TT_MAP -> "map" + TT_LIST -> "lst" + TT_SET -> "set" + TT_UUID -> "uuid" + TT_STOP -> "stop" + + +typeAbbrev :: TV.Value -> Text +typeAbbrev = \case + TV.Bool {} -> "tf" + TV.Byte {} -> "i8" + TV.I16 {} -> "i16" + TV.I32 {} -> "i32" + TV.I64 {} -> "i64" + TV.Double {} -> "dbl" + TV.String {} -> "str" + TV.Binary {} -> "bin" + TV.Struct {} -> "rec" + TV.Map {} -> "map" + TV.List {} -> "lst" + TV.Set {} -> "set" + TV.UUID {} -> "uuid" + + +isTypeTag :: Text -> Bool +isTypeTag t = + t + `elem` ["tf", "i8", "i16", "i32", "i64", "dbl", "str", "bin", "rec", "map", "lst", "set", "uuid"] + + +containerElemTag :: ThriftType -> V.Vector TV.Value -> Text +containerElemTag et elems + | V.null elems = wireTypeTag et + | otherwise = typeAbbrev (V.head elems) + + +stringifyKey :: TV.Value -> Text +stringifyKey = \case + TV.String t -> t + TV.Bool b -> if b then "true" else "false" + TV.Byte v -> Text.pack (show v) + TV.I16 v -> Text.pack (show v) + TV.I32 v -> Text.pack (show v) + TV.I64 v -> Text.pack (show v) + TV.Double d -> Text.pack (show d) + TV.Binary bs -> TE.decodeUtf8 (Base64.encode bs) + TV.UUID bs -> TE.decodeUtf8 (Base16.encode bs) + _ -> "" + + +parseTypedKeyFromTag :: Text -> Text -> Either String TV.Value +parseTypedKeyFromTag tag keyStr = case tag of + "str" -> Right (TV.String keyStr) + "tf" -> case keyStr of + "true" -> Right (TV.Bool True) + "false" -> Right (TV.Bool False) + _ -> Left "Invalid bool key" + "i8" -> readIntKey (\n -> TV.Byte (fromIntegral n)) keyStr + "i16" -> readIntKey (\n -> TV.I16 (fromIntegral n)) keyStr + "i32" -> readIntKey (\n -> TV.I32 (fromIntegral n)) keyStr + "i64" -> readIntKey TV.I64 keyStr + _ -> Left ("Unsupported map key type: " <> Text.unpack tag) + + +readIntKey :: (Int64 -> TV.Value) -> Text -> Either String TV.Value +readIntKey f t = case (TR.signed TR.decimal t :: Either String (Int64, Text)) of + Right (n, "") -> Right (f n) + _ -> Left ("Invalid numeric key: " <> Text.unpack t) + + +defaultForWireType :: ThriftType -> TV.Value +defaultForWireType = \case + TT_BOOL -> TV.Bool False + TT_BYTE -> TV.Byte 0 + TT_I16 -> TV.I16 0 + TT_I32 -> TV.I32 0 + TT_I64 -> TV.I64 0 + TT_DOUBLE -> TV.Double 0 + TT_STRING -> TV.String "" + TT_STRUCT -> TV.Struct V.empty + TT_MAP -> TV.Map TT_STOP TT_STOP V.empty + TT_LIST -> TV.List TT_STOP V.empty + TT_SET -> TV.Set TT_STOP V.empty + TT_UUID -> TV.UUID BS.empty + TT_STOP -> TV.Bool False + + +instance Aeson.ToJSON TV.Value where + toJSON = thriftToJSON + + +instance Aeson.FromJSON TV.Value where + parseJSON v = case thriftFromTypedJSON v of + Right val -> pure val + Left err -> fail err diff --git a/src/Thrift/Message.hs b/src/Thrift/Message.hs new file mode 100644 --- /dev/null +++ b/src/Thrift/Message.hs @@ -0,0 +1,266 @@ +{-# LANGUAGE BangPatterns #-} + +{- | Thrift RPC message header encoding\/decoding for Binary and Compact protocols. + +A Thrift RPC message wraps a struct payload with a header containing the +method name, message type (call\/reply\/exception\/oneway), and a sequence ID. +-} +module Thrift.Message ( + ThriftMessageType (..), + ThriftMessage (..), + encodeMessageBinary, + decodeMessageBinary, + encodeMessageCompact, + decodeMessageCompact, +) where + +import Data.Bits (shiftL, shiftR, (.&.), (.|.)) +import Data.ByteString (ByteString) +import Data.ByteString qualified as BS +import Data.ByteString.Lazy qualified as BL +import Data.ByteString.Unsafe qualified as BSU +import Data.Int (Int32) +import Data.Text (Text) +import Data.Text.Encoding qualified as TE +import Data.Word (Word32, Word64, Word8) +import Thrift.Decode (decodeBinary, decodeCompact) +import Thrift.Encode (encodeBinary, encodeCompact) +import Thrift.Value qualified as TV +import Wireform.Builder qualified as B + + +-- | Thrift RPC message types. +data ThriftMessageType + = TMsgCall + | TMsgReply + | TMsgException + | TMsgOneway + deriving stock (Show, Eq, Ord, Enum, Bounded) + + +-- | A complete Thrift RPC message: header + payload. +data ThriftMessage = ThriftMessage + { tmsgName :: !Text + , tmsgType :: !ThriftMessageType + , tmsgSeqId :: {-# UNPACK #-} !Int32 + , tmsgPayload :: !TV.Value + } + deriving stock (Show, Eq) + + +msgTypeToWord8 :: ThriftMessageType -> Word8 +msgTypeToWord8 TMsgCall = 1 +msgTypeToWord8 TMsgReply = 2 +msgTypeToWord8 TMsgException = 3 +msgTypeToWord8 TMsgOneway = 4 +{-# INLINE msgTypeToWord8 #-} + + +msgTypeFromWord8 :: Word8 -> Maybe ThriftMessageType +msgTypeFromWord8 1 = Just TMsgCall +msgTypeFromWord8 2 = Just TMsgReply +msgTypeFromWord8 3 = Just TMsgException +msgTypeFromWord8 4 = Just TMsgOneway +msgTypeFromWord8 _ = Nothing +{-# INLINE msgTypeFromWord8 #-} + + +-------------------------------------------------------------------------------- +-- Binary Protocol +-------------------------------------------------------------------------------- + +binaryVersion :: Word32 +binaryVersion = 0x80010000 + + +encodeMessageBinary :: ThriftMessage -> ByteString +encodeMessageBinary (ThriftMessage name mtype seqid payload) = + BL.toStrict $ + B.toLazyByteString $ + let !nameBytes = TE.encodeUtf8 name + !version = binaryVersion .|. fromIntegral (msgTypeToWord8 mtype) + in putBE32 version + <> putBE32 (fromIntegral (BS.length nameBytes)) + <> B.byteString nameBytes + <> putBE32i (fromIntegral seqid) + <> B.byteString (encodeBinary payload) + + +decodeMessageBinary :: ByteString -> Either String ThriftMessage +decodeMessageBinary !bs + | BS.length bs < 4 = Left "decodeMessageBinary: insufficient data" + | otherwise = + let !w = getBE32 bs 0 + in if w .&. 0x80000000 /= 0 + then decodeStrictBinary bs w + else decodeOldBinary bs + + +decodeStrictBinary :: ByteString -> Word32 -> Either String ThriftMessage +decodeStrictBinary !bs !versionWord = + let !mtypeByte = fromIntegral (versionWord .&. 0xFF) :: Word8 + in case msgTypeFromWord8 mtypeByte of + Nothing -> Left $ "decodeMessageBinary: invalid message type: " ++ show mtypeByte + Just !mtype -> do + (nameBytes, off1) <- getStr bs 4 + case TE.decodeUtf8' nameBytes of + Left _ -> Left "decodeMessageBinary: invalid UTF-8 in method name" + Right name -> do + seqid <- getI32At bs off1 + let !payloadBs = BS.drop (off1 + 4) bs + payload <- decodeBinary payloadBs + Right (ThriftMessage name mtype seqid payload) + + +decodeOldBinary :: ByteString -> Either String ThriftMessage +decodeOldBinary !bs = do + (nameBytes, off1) <- getStr bs 0 + case TE.decodeUtf8' nameBytes of + Left _ -> Left "decodeMessageBinary: invalid UTF-8 in method name" + Right name -> do + if off1 >= BS.length bs + then Left "decodeMessageBinary: missing message type byte" + else do + let !mtypeByte = BSU.unsafeIndex bs off1 + case msgTypeFromWord8 mtypeByte of + Nothing -> Left $ "decodeMessageBinary: invalid message type: " ++ show mtypeByte + Just !mtype -> do + seqid <- getI32At bs (off1 + 1) + let !payloadBs = BS.drop (off1 + 5) bs + payload <- decodeBinary payloadBs + Right (ThriftMessage name mtype seqid payload) + + +-------------------------------------------------------------------------------- +-- Compact Protocol +-------------------------------------------------------------------------------- + +compactProtocolId :: Word8 +compactProtocolId = 0x82 + + +compactVersion :: Word8 +compactVersion = 1 + + +encodeMessageCompact :: ThriftMessage -> ByteString +encodeMessageCompact (ThriftMessage name mtype seqid payload) = + BL.toStrict $ + B.toLazyByteString $ + let !nameBytes = TE.encodeUtf8 name + !vtByte = (compactVersion .&. 0x1F) .|. (msgTypeToWord8 mtype `shiftL` 5) + in B.word8 compactProtocolId + <> B.word8 vtByte + <> putVarint (fromIntegral (fromIntegral seqid :: Word32)) + <> putVarint (fromIntegral (BS.length nameBytes)) + <> B.byteString nameBytes + <> B.byteString (encodeCompact payload) + + +decodeMessageCompact :: ByteString -> Either String ThriftMessage +decodeMessageCompact !bs + | BS.length bs < 2 = Left "decodeMessageCompact: insufficient data" + | BSU.unsafeIndex bs 0 /= compactProtocolId = + Left $ "decodeMessageCompact: invalid protocol ID: " ++ show (BSU.unsafeIndex bs 0) + | otherwise = + let !vtByte = BSU.unsafeIndex bs 1 + !ver = vtByte .&. 0x1F + !mtypeBits = (vtByte `shiftR` 5) .&. 0x07 + in if ver /= compactVersion + then Left $ "decodeMessageCompact: unsupported version: " ++ show ver + else case msgTypeFromWord8 mtypeBits of + Nothing -> Left $ "decodeMessageCompact: invalid message type: " ++ show mtypeBits + Just !mtype -> do + (seqidW, off1) <- + maybeToEither + "decodeMessageCompact: failed to read seqid varint" + (getVarint bs 2) + let !seqid = fromIntegral seqidW :: Int32 + (nameLen, off2) <- + maybeToEither + "decodeMessageCompact: failed to read name length varint" + (getVarint bs off1) + let !nLen = fromIntegral nameLen :: Int + if off2 + nLen > BS.length bs + then Left "decodeMessageCompact: name extends past end of data" + else do + let !nameBytes = BS.take nLen (BS.drop off2 bs) + case TE.decodeUtf8' nameBytes of + Left _ -> Left "decodeMessageCompact: invalid UTF-8 in method name" + Right name -> do + let !payloadBs = BS.drop (off2 + nLen) bs + payload <- decodeCompact payloadBs + Right (ThriftMessage name mtype seqid payload) + + +-------------------------------------------------------------------------------- +-- Internal helpers +-------------------------------------------------------------------------------- + +maybeToEither :: String -> Maybe a -> Either String a +maybeToEither msg Nothing = Left msg +maybeToEither _ (Just a) = Right a +{-# INLINE maybeToEither #-} + + +putBE32 :: Word32 -> B.Builder +putBE32 !w = + B.word8 (fromIntegral (w `shiftR` 24)) + <> B.word8 (fromIntegral ((w `shiftR` 16) .&. 0xFF)) + <> B.word8 (fromIntegral ((w `shiftR` 8) .&. 0xFF)) + <> B.word8 (fromIntegral (w .&. 0xFF)) +{-# INLINE putBE32 #-} + + +putBE32i :: Word32 -> B.Builder +putBE32i = putBE32 +{-# INLINE putBE32i #-} + + +getBE32 :: ByteString -> Int -> Word32 +getBE32 !bs !off = + let !b0 = fromIntegral (BSU.unsafeIndex bs off) :: Word32 + !b1 = fromIntegral (BSU.unsafeIndex bs (off + 1)) :: Word32 + !b2 = fromIntegral (BSU.unsafeIndex bs (off + 2)) :: Word32 + !b3 = fromIntegral (BSU.unsafeIndex bs (off + 3)) :: Word32 + in (b0 `shiftL` 24) .|. (b1 `shiftL` 16) .|. (b2 `shiftL` 8) .|. b3 +{-# INLINE getBE32 #-} + + +getStr :: ByteString -> Int -> Either String (ByteString, Int) +getStr !bs !off + | off + 4 > BS.length bs = Left "unexpected end of data reading string length" + | otherwise = + let !len = fromIntegral (getBE32 bs off) :: Int + !start = off + 4 + in if len < 0 || start + len > BS.length bs + then Left "string length exceeds available data" + else Right (BS.take len (BS.drop start bs), start + len) + + +getI32At :: ByteString -> Int -> Either String Int32 +getI32At !bs !off + | off + 4 > BS.length bs = Left "unexpected end of data reading i32" + | otherwise = Right (fromIntegral (getBE32 bs off)) + + +putVarint :: Word64 -> B.Builder +putVarint !n + | n < 0x80 = B.word8 (fromIntegral n) + | otherwise = B.word8 (fromIntegral (n .&. 0x7F) .|. 0x80) <> putVarint (n `shiftR` 7) +{-# INLINE putVarint #-} + + +getVarint :: ByteString -> Int -> Maybe (Word64, Int) +getVarint !bs !off = go off 0 0 + where + go !i !acc !shift + | i >= BS.length bs = Nothing + | shift >= 64 = Nothing + | otherwise = + let !b = fromIntegral (BSU.unsafeIndex bs i) :: Word64 + !acc' = acc .|. ((b .&. 0x7F) `shiftL` shift) + in if b < 0x80 + then Just (acc', i + 1) + else go (i + 1) acc' (shift + 7) +{-# INLINE getVarint #-} diff --git a/src/Thrift/Parser.hs b/src/Thrift/Parser.hs new file mode 100644 --- /dev/null +++ b/src/Thrift/Parser.hs @@ -0,0 +1,438 @@ +{- | Thrift IDL parser. + +Parses Thrift IDL definitions (@.thrift@ files) into a +'Thrift.Schema.ThriftSchema' AST using Megaparsec. Supports structs, +unions, exceptions, enums, services, typedefs, constants, includes, +namespaces, and all Thrift type annotations. + +@ +struct Person { + 1: required string name; + 2: optional i32 age; +} +@ + +@ +import Thrift.Parser (parseThrift) +let Right schema = parseThrift input +@ +-} +module Thrift.Parser ( + parseThrift, +) where + +import Control.Monad (void) +import Data.Char (isAlpha, isAlphaNum) +import Data.Int (Int32, Int64) +import Data.Maybe (fromMaybe) +import Data.Text (Text) +import Data.Text qualified as T +import Data.Vector qualified as V +import Data.Void (Void) +import Text.Megaparsec +import Text.Megaparsec.Char +import Text.Megaparsec.Char.Lexer qualified as L +import Thrift.Schema + + +type Parser = Parsec Void Text + + +-- | Parse Thrift IDL text into a 'ThriftSchema'. +parseThrift :: Text -> Either String ThriftSchema +parseThrift input = + case parse (sc *> documentP <* eof) "" input of + Left err -> Left (errorBundlePretty err) + Right s -> Right s + + +-- Whitespace and comments +sc :: Parser () +sc = + L.space + space1 + (L.skipLineComment "//") + (L.skipBlockComment "/*" "*/") + + +lexeme :: Parser a -> Parser a +lexeme = L.lexeme sc + + +symbol :: Text -> Parser Text +symbol = L.symbol sc + + +reserved :: Text -> Parser () +reserved w = lexeme (string w *> notFollowedBy alphaNumChar) + + +braces :: Parser a -> Parser a +braces = between (symbol "{") (symbol "}") + + +parens :: Parser a -> Parser a +parens = between (symbol "(") (symbol ")") + + +angles :: Parser a -> Parser a +angles = between (symbol "<") (symbol ">") + + +comma :: Parser () +comma = void (symbol ",") + + +semicolon :: Parser () +semicolon = void (symbol ";") + + +listSep :: Parser () +listSep = comma <|> semicolon + + +optListSep :: Parser () +optListSep = void (optional listSep) + + +identifier :: Parser Text +identifier = lexeme $ do + c <- satisfy (\ch -> isAlpha ch || ch == '_') + rest <- takeWhileP Nothing (\ch -> isAlphaNum ch || ch == '_' || ch == '.') + pure (T.cons c rest) + + +stringLiteral :: Parser Text +stringLiteral = lexeme $ do + delim <- char '"' <|> char '\'' + content <- takeWhileP Nothing (/= delim) + void (char delim) + pure content + + +integerLiteral :: Parser Int64 +integerLiteral = lexeme $ do + sign <- optional (char '-') + n <- try (string "0x" *> L.hexadecimal) <|> L.decimal + let val = case sign of + Just _ -> negate n + Nothing -> n + pure val + + +-- Document +documentP :: Parser ThriftSchema +documentP = do + _ <- many headerP + defs <- many definitionP + let structs = [s | DefStruct s <- defs] + enums = [e | DefEnum e <- defs] + typedefs = [t | DefTypedef t <- defs] + consts = [c | DefConst c <- defs] + services = [s | DefService s <- defs] + pure + ThriftSchema + { tsStructs = structs + , tsEnums = enums + , tsTypedefs = typedefs + , tsConsts = consts + , tsServices = services + } + + +data Definition + = DefStruct ThriftStruct + | DefEnum ThriftEnum + | DefTypedef ThriftTypedef + | DefConst ThriftConst + | DefService ThriftService + + +headerP :: Parser () +headerP = includeP <|> namespaceP <|> hashCommentP <|> cppIncludeP + + +hashCommentP :: Parser () +hashCommentP = lexeme $ do + void (char '#') + void (takeWhileP Nothing (/= '\n')) + + +includeP :: Parser () +includeP = do + reserved "include" + _ <- stringLiteral + pure () + + +cppIncludeP :: Parser () +cppIncludeP = do + reserved "cpp_include" + _ <- stringLiteral + pure () + + +namespaceP :: Parser () +namespaceP = do + reserved "namespace" + _ <- identifier + _ <- identifier <|> stringLiteral + optional (parens (many annotationPairP)) + pure () + + +annotationPairP :: Parser (Text, Text) +annotationPairP = do + key <- identifier + val <- option "" (symbol "=" *> (stringLiteral <|> identifier)) + optListSep + pure (key, val) + + +definitionP :: Parser Definition +definitionP = + choice + [ DefStruct <$> structP + , DefStruct <$> unionP + , DefStruct <$> exceptionP + , DefEnum <$> enumP + , DefService <$> serviceP + , DefTypedef <$> typedefP + , DefConst <$> constP + ] + + +structP :: Parser ThriftStruct +structP = do + reserved "struct" + name <- identifier + fields <- braces (many fieldP) + anns <- annotationsP + pure + ThriftStruct + { tsName = name + , tsKind = StructNormal + , tsFields = fields + , tsAnnotations = anns + } + + +unionP :: Parser ThriftStruct +unionP = do + reserved "union" + name <- identifier + fields <- braces (many fieldP) + anns <- annotationsP + pure + ThriftStruct + { tsName = name + , tsKind = StructUnion + , tsFields = fields + , tsAnnotations = anns + } + + +exceptionP :: Parser ThriftStruct +exceptionP = do + reserved "exception" + name <- identifier + fields <- braces (many fieldP) + anns <- annotationsP + pure + ThriftStruct + { tsName = name + , tsKind = StructException + , tsFields = fields + , tsAnnotations = anns + } + + +fieldP :: Parser ThriftField +fieldP = do + fid <- integerLiteral + void (symbol ":") + req <- optional (requirednessP) + ftype <- typeP + fname <- identifier + dflt <- optional (symbol "=" *> constValueP) + optListSep + anns <- annotationsP + pure + ThriftField + { tfFieldId = fromIntegral fid + , tfFieldName = fname + , tfFieldType = ftype + , tfRequiredness = fromMaybe Default req + , tfDefault = dflt + , tfAnnotations = anns + } + + +requirednessP :: Parser Requiredness +requirednessP = (Required <$ reserved "required") <|> (Optional <$ reserved "optional") + + +typeP :: Parser ThriftType +typeP = + choice + [ TBool <$ reserved "bool" + , TByte <$ (reserved "byte" <|> reserved "i8") + , TI16 <$ reserved "i16" + , TI32 <$ reserved "i32" + , TI64 <$ reserved "i64" + , TDouble <$ reserved "double" + , TString <$ reserved "string" + , TBinary <$ reserved "binary" + , TUUID <$ reserved "uuid" + , listTypeP + , setTypeP + , mapTypeP + , TStruct <$> identifier + ] + + +listTypeP :: Parser ThriftType +listTypeP = do + reserved "list" + TList <$> angles typeP + + +setTypeP :: Parser ThriftType +setTypeP = do + reserved "set" + TSet <$> angles typeP + + +mapTypeP :: Parser ThriftType +mapTypeP = do + reserved "map" + angles $ do + kt <- typeP + void (symbol ",") + vt <- typeP + pure (TMap kt vt) + + +enumP :: Parser ThriftEnum +enumP = do + reserved "enum" + name <- identifier + vals <- braces (many enumValueP) + _ <- annotationsP + let numbered = assignEnumValues vals 0 + pure ThriftEnum {teName = name, teValues = numbered} + + +enumValueP :: Parser (Text, Maybe Int32) +enumValueP = do + name <- identifier + val <- optional (symbol "=" *> (fromIntegral <$> integerLiteral)) + optListSep + _ <- annotationsP + pure (name, val) + + +assignEnumValues :: [(Text, Maybe Int32)] -> Int32 -> [(Text, Int32)] +assignEnumValues [] _ = [] +assignEnumValues ((name, mval) : rest) nextVal = + let val = fromMaybe nextVal mval + in (name, val) : assignEnumValues rest (val + 1) + + +serviceP :: Parser ThriftService +serviceP = do + reserved "service" + name <- identifier + extends <- optional (reserved "extends" *> identifier) + methods <- braces (many methodP) + _ <- annotationsP + pure + ThriftService + { tsvName = name + , tsvExtends = extends + , tsvMethods = methods + } + + +methodP :: Parser ThriftMethod +methodP = do + ow <- option False (True <$ reserved "oneway") + retType <- (Nothing <$ reserved "void") <|> (Just <$> typeP) + name <- identifier + params <- parens (many fieldP) + throws <- option [] (reserved "throws" *> parens (many fieldP)) + optListSep + _ <- annotationsP + pure + ThriftMethod + { tmName = name + , tmReturnType = retType + , tmParams = params + , tmThrows = throws + , tmOneway = ow + } + + +typedefP :: Parser ThriftTypedef +typedefP = do + reserved "typedef" + ty <- typeP + name <- identifier + _ <- annotationsP + optListSep + pure ThriftTypedef {ttName = name, ttType = ty} + + +constP :: Parser ThriftConst +constP = do + reserved "const" + ty <- typeP + name <- identifier + void (symbol "=") + val <- constValueP + optListSep + pure ThriftConst {tcName = name, tcType = ty, tcValue = val} + + +constValueP :: Parser ThriftConstValue +constValueP = + choice + [ TCVBool True <$ reserved "true" + , TCVBool False <$ reserved "false" + , TCVString <$> stringLiteral + , try constListP + , try constMapP + , try (TCVDouble <$> lexeme (L.signed (pure ()) L.float)) + , TCVInt <$> integerLiteral + , TCVIdent <$> identifier + ] + + +constListP :: Parser ThriftConstValue +constListP = do + void (symbol "[") + vals <- many (constValueP <* optListSep) + void (symbol "]") + pure (TCVList vals) + + +constMapP :: Parser ThriftConstValue +constMapP = do + void (symbol "{") + entries <- many (constMapEntry <* optListSep) + void (symbol "}") + pure (TCVMap entries) + + +constMapEntry :: Parser (ThriftConstValue, ThriftConstValue) +constMapEntry = do + k <- constValueP + void (symbol ":") + v <- constValueP + pure (k, v) + + +annotationsP :: Parser (V.Vector (Text, Text)) +annotationsP = do + mAnns <- optional $ parens (many annotationPairP) + pure $ V.fromList (fromMaybe [] mAnns) diff --git a/src/Thrift/QQ.hs b/src/Thrift/QQ.hs new file mode 100644 --- /dev/null +++ b/src/Thrift/QQ.hs @@ -0,0 +1,39 @@ +{- | QuasiQuoter for inline Thrift IDL definitions. + +@ +{\-\# LANGUAGE QuasiQuotes \#-\} +{\-\# LANGUAGE TemplateHaskell \#-\} +import Thrift.QQ + +[thrift| + struct Person { 1: string name; 2: i32 age; } +|] +@ +-} +module Thrift.QQ ( + thrift, +) where + +import Data.Text qualified as T +import Language.Haskell.TH +import Language.Haskell.TH.Quote +import Thrift.CodeGen (deriveThrift) +import Thrift.Parser (parseThrift) + + +-- | QuasiQuoter for Thrift IDL. +thrift :: QuasiQuoter +thrift = + QuasiQuoter + { quoteExp = \_ -> fail "thrift: not an expression quoter" + , quotePat = \_ -> fail "thrift: not a pattern quoter" + , quoteType = \_ -> fail "thrift: not a type quoter" + , quoteDec = thriftDec + } + + +thriftDec :: String -> Q [Dec] +thriftDec src = do + case parseThrift (T.pack src) of + Left err -> fail ("thrift parse error: " ++ err) + Right schema -> deriveThrift schema diff --git a/src/Thrift/Registry.hs b/src/Thrift/Registry.hs new file mode 100644 --- /dev/null +++ b/src/Thrift/Registry.hs @@ -0,0 +1,78 @@ +{- | User-extensible annotation handler registration for Thrift. + +Thrift's extension mechanism is annotations on fields and structs. +This module allows users to register custom handlers that transform +types and emit extra code during Thrift code generation. +-} +module Thrift.Registry ( + -- * Registry + ThriftRegistry (..), + defaultThriftRegistry, + + -- * Field annotation handlers + FieldAnnotationHandler (..), + registerFieldAnnotation, + + -- * Struct annotation handlers + StructAnnotationHandler (..), + registerStructAnnotation, +) where + +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.Text (Text) + + +{- | Handler for a Thrift field annotation. When a field carries an +annotation matching a registered handler, the handler can transform +the field's Haskell type and emit extra lines of code. +-} +data FieldAnnotationHandler = FieldAnnotationHandler + { fahTransformType :: !(Text -> Text) + , fahExtraCode :: !(Text -> Text -> [Text]) + } + + +{- | Handler for a Thrift struct-level annotation. Can add extra +deriving clauses and emit extra lines of code. +-} +data StructAnnotationHandler = StructAnnotationHandler + { sahExtraDerivations :: !(Text -> [Text]) + , sahExtraCode :: !(Text -> Text -> [Text]) + } + + +-- | Registry of custom annotation handlers for Thrift code generation. +data ThriftRegistry = ThriftRegistry + { trFieldAnnotations :: !(Map Text FieldAnnotationHandler) + , trStructAnnotations :: !(Map Text StructAnnotationHandler) + } + + +instance Semigroup ThriftRegistry where + a <> b = + ThriftRegistry + { trFieldAnnotations = trFieldAnnotations a <> trFieldAnnotations b + , trStructAnnotations = trStructAnnotations a <> trStructAnnotations b + } + + +instance Monoid ThriftRegistry where + mempty = ThriftRegistry Map.empty Map.empty + + +-- | Default (empty) Thrift registry with no custom annotation handlers. +defaultThriftRegistry :: ThriftRegistry +defaultThriftRegistry = mempty + + +-- | Register a field annotation handler. +registerFieldAnnotation :: Text -> FieldAnnotationHandler -> ThriftRegistry -> ThriftRegistry +registerFieldAnnotation name handler reg = + reg {trFieldAnnotations = Map.insert name handler (trFieldAnnotations reg)} + + +-- | Register a struct annotation handler. +registerStructAnnotation :: Text -> StructAnnotationHandler -> ThriftRegistry -> ThriftRegistry +registerStructAnnotation name handler reg = + reg {trStructAnnotations = Map.insert name handler (trStructAnnotations reg)} diff --git a/src/Thrift/Schema.hs b/src/Thrift/Schema.hs new file mode 100644 --- /dev/null +++ b/src/Thrift/Schema.hs @@ -0,0 +1,200 @@ +{-# LANGUAGE BangPatterns #-} + +{- | Thrift schema AST — a complete representation of Thrift IDL constructs. + +Covers base types, containers, structs, unions, exceptions, enums, +typedefs, constants, and services. +-} +module Thrift.Schema ( + -- * Top-level schema + ThriftSchema (..), + + -- * Types + ThriftType (..), + + -- * Struct / Union / Exception + ThriftStruct (..), + StructKind (..), + ThriftField (..), + Requiredness (..), + + -- * Enum + ThriftEnum (..), + + -- * Service + ThriftService (..), + ThriftMethod (..), + + -- * Const + ThriftConst (..), + ThriftConstValue (..), + + -- * Typedef + ThriftTypedef (..), +) where + +import Data.Int (Int32, Int64) +import Data.Text (Text) +import Data.Vector (Vector) + + +-------------------------------------------------------------------------------- +-- Requiredness +-------------------------------------------------------------------------------- + +-- | Field requiredness as specified in the Thrift IDL. +data Requiredness + = -- | The field must be present. + Required + | -- | The field may be absent. + Optional + | -- | Requiredness not specified (Thrift default behavior). + Default + deriving stock (Show, Eq, Ord, Enum, Bounded) + + +-------------------------------------------------------------------------------- +-- Thrift types +-------------------------------------------------------------------------------- + +-- | All types expressible in Thrift IDL. +data ThriftType + = TBool + | TByte + | TI16 + | TI32 + | TI64 + | TDouble + | TString + | TBinary + | TUUID + | TStruct !Text + | TEnum !Text + | TTypedef !Text + | TList !ThriftType + | TSet !ThriftType + | TMap !ThriftType !ThriftType + deriving stock (Show, Eq, Ord) + + +-------------------------------------------------------------------------------- +-- Fields +-------------------------------------------------------------------------------- + +-- | A single field in a struct, union, or exception. +data ThriftField = ThriftField + { tfFieldId :: {-# UNPACK #-} !Int32 + , tfFieldName :: !Text + , tfFieldType :: !ThriftType + , tfRequiredness :: !Requiredness + , tfDefault :: !(Maybe ThriftConstValue) + , tfAnnotations :: !(Vector (Text, Text)) + } + deriving stock (Show, Eq) + + +-------------------------------------------------------------------------------- +-- Struct / Union / Exception +-------------------------------------------------------------------------------- + +-- | Struct kind tag. +data StructKind = StructNormal | StructUnion | StructException + deriving stock (Show, Eq, Ord, Enum, Bounded) + + +-- | A Thrift struct, union, or exception. +data ThriftStruct = ThriftStruct + { tsName :: !Text + , tsKind :: !StructKind + , tsFields :: ![ThriftField] + , tsAnnotations :: !(Vector (Text, Text)) + } + deriving stock (Show, Eq) + + +-------------------------------------------------------------------------------- +-- Enum +-------------------------------------------------------------------------------- + +-- | A Thrift enum definition. +data ThriftEnum = ThriftEnum + { teName :: !Text + , teValues :: ![(Text, Int32)] + } + deriving stock (Show, Eq) + + +-------------------------------------------------------------------------------- +-- Typedef +-------------------------------------------------------------------------------- + +-- | A Thrift typedef. +data ThriftTypedef = ThriftTypedef + { ttName :: !Text + , ttType :: !ThriftType + } + deriving stock (Show, Eq) + + +-------------------------------------------------------------------------------- +-- Constants +-------------------------------------------------------------------------------- + +-- | Constant values in the Thrift IDL. +data ThriftConstValue + = TCVInt {-# UNPACK #-} !Int64 + | TCVDouble {-# UNPACK #-} !Double + | TCVString !Text + | TCVBool !Bool + | TCVList ![ThriftConstValue] + | TCVMap ![(ThriftConstValue, ThriftConstValue)] + | TCVIdent !Text + deriving stock (Show, Eq) + + +-- | A top-level constant definition. +data ThriftConst = ThriftConst + { tcName :: !Text + , tcType :: !ThriftType + , tcValue :: !ThriftConstValue + } + deriving stock (Show, Eq) + + +-------------------------------------------------------------------------------- +-- Service +-------------------------------------------------------------------------------- + +-- | A single method in a Thrift service. +data ThriftMethod = ThriftMethod + { tmName :: !Text + , tmReturnType :: !(Maybe ThriftType) + , tmParams :: ![ThriftField] + , tmThrows :: ![ThriftField] + , tmOneway :: !Bool + } + deriving stock (Show, Eq) + + +-- | A Thrift service definition. +data ThriftService = ThriftService + { tsvName :: !Text + , tsvExtends :: !(Maybe Text) + , tsvMethods :: ![ThriftMethod] + } + deriving stock (Show, Eq) + + +-------------------------------------------------------------------------------- +-- Top-level schema +-------------------------------------------------------------------------------- + +-- | A complete Thrift schema (one .thrift file). +data ThriftSchema = ThriftSchema + { tsStructs :: ![ThriftStruct] + , tsEnums :: ![ThriftEnum] + , tsTypedefs :: ![ThriftTypedef] + , tsConsts :: ![ThriftConst] + , tsServices :: ![ThriftService] + } + deriving stock (Show, Eq) diff --git a/src/Thrift/Transport.hs b/src/Thrift/Transport.hs new file mode 100644 --- /dev/null +++ b/src/Thrift/Transport.hs @@ -0,0 +1,75 @@ +{-# LANGUAGE BangPatterns #-} + +{- | Thrift framed transport: 4-byte big-endian length prefix framing. + +Many Thrift servers (e.g. TNonblockingServer, TThreadedSelectorServer) +require framed transport where each message is preceded by a 4-byte +big-endian length prefix. +-} +module Thrift.Transport ( + frameMessage, + unframeMessage, + unframeMessages, +) where + +import Data.Bits (shiftL, shiftR, (.&.), (.|.)) +import Data.ByteString (ByteString) +import Data.ByteString qualified as BS +import Data.ByteString.Unsafe qualified as BSU +import Data.Word (Word32, Word8) + + +-- | Frame a message: 4-byte big-endian length prefix + payload. +frameMessage :: ByteString -> ByteString +frameMessage !payload = + let !len = BS.length payload + !b0 = fromIntegral ((len `shiftR` 24) .&. 0xFF) :: Word8 + !b1 = fromIntegral ((len `shiftR` 16) .&. 0xFF) :: Word8 + !b2 = fromIntegral ((len `shiftR` 8) .&. 0xFF) :: Word8 + !b3 = fromIntegral (len .&. 0xFF) :: Word8 + in BS.pack [b0, b1, b2, b3] <> payload + + +-- | Unframe: read 4-byte length, extract payload. +unframeMessage :: ByteString -> Either String ByteString +unframeMessage !bs + | BS.length bs < 4 = Left "Thrift.Transport: input too short for frame header" + | otherwise = + let !len = readBE32 bs 0 + in if fromIntegral len + 4 > BS.length bs + then + Left $ + "Thrift.Transport: frame claims " + ++ show len + ++ " bytes but only " + ++ show (BS.length bs - 4) + ++ " available" + else Right (BSU.unsafeTake (fromIntegral len) (BSU.unsafeDrop 4 bs)) + + +-- | Unframe multiple messages from a stream. +unframeMessages :: ByteString -> Either String [ByteString] +unframeMessages !bs = go 0 [] + where + !bsLen = BS.length bs + go !off !acc + | off >= bsLen = Right (reverse acc) + | off + 4 > bsLen = Left "Thrift.Transport: truncated frame header" + | otherwise = + let !len = fromIntegral (readBE32 bs off) :: Int + !payOff = off + 4 + in if payOff + len > bsLen + then Left $ "Thrift.Transport: truncated frame payload" + else + let !payload = BSU.unsafeTake len (BSU.unsafeDrop payOff bs) + in go (payOff + len) (payload : acc) + + +readBE32 :: ByteString -> Int -> Word32 +readBE32 !bs !off = + let !b0 = fromIntegral (BSU.unsafeIndex bs off) :: Word32 + !b1 = fromIntegral (BSU.unsafeIndex bs (off + 1)) :: Word32 + !b2 = fromIntegral (BSU.unsafeIndex bs (off + 2)) :: Word32 + !b3 = fromIntegral (BSU.unsafeIndex bs (off + 3)) :: Word32 + in (b0 `shiftL` 24) .|. (b1 `shiftL` 16) .|. (b2 `shiftL` 8) .|. b3 +{-# INLINE readBE32 #-} diff --git a/src/Thrift/Value.hs b/src/Thrift/Value.hs new file mode 100644 --- /dev/null +++ b/src/Thrift/Value.hs @@ -0,0 +1,58 @@ +{- | Shared Value type for high-level Thrift encode\/decode. + +Each variant carries its wire data directly. Struct fields are +tagged by their Int16 field ID so both protocols can encode them +with the correct field headers. Covers bool, byte, i16, i32, i64, +double, string, binary, list, set, map, and struct. + +@ +import qualified Thrift.Value as T +import qualified Data.Vector as V + +let person = T.Struct (V.fromList [(1, T.String \"Alice\"), (2, T.I32 30)]) +@ +-} +module Thrift.Value ( + Value (..), + thriftTypeOf, +) where + +import Data.ByteString (ByteString) +import Data.Int (Int16, Int32, Int64, Int8) +import Data.Text (Text) +import Data.Vector (Vector) +import Thrift.Wire (ThriftType (..)) + + +data Value + = Bool Bool + | Byte Int8 + | I16 Int16 + | I32 Int32 + | I64 Int64 + | Double Double + | String Text + | Binary ByteString + | Struct (Vector (Int16, Value)) + | Map ThriftType ThriftType (Vector (Value, Value)) + | List ThriftType (Vector Value) + | Set ThriftType (Vector Value) + | UUID ByteString + deriving stock (Show, Eq) + + +thriftTypeOf :: Value -> ThriftType +thriftTypeOf = \case + Bool {} -> TT_BOOL + Byte {} -> TT_BYTE + I16 {} -> TT_I16 + I32 {} -> TT_I32 + I64 {} -> TT_I64 + Double {} -> TT_DOUBLE + String {} -> TT_STRING + Binary {} -> TT_STRING + Struct {} -> TT_STRUCT + Map {} -> TT_MAP + List {} -> TT_LIST + Set {} -> TT_SET + UUID {} -> TT_UUID diff --git a/src/Thrift/Wire.hs b/src/Thrift/Wire.hs new file mode 100644 --- /dev/null +++ b/src/Thrift/Wire.hs @@ -0,0 +1,882 @@ +{-# LANGUAGE BangPatterns #-} + +{- | Low-level Thrift wire encoding primitives for both the Binary Protocol +and the Compact Protocol. + +Binary Protocol: big-endian integers, IEEE 754 doubles, i32-prefixed strings. +Compact Protocol: zigzag+varint integers, delta field headers, little-endian doubles. + +Performance: direct 'pokeByteOff' \/ 'peekByteOff', INLINE on everything. +-} +module Thrift.Wire ( + -- * Thrift type codes + ThriftType (..), + thriftTypeToBin, + thriftTypeFromBin, + thriftTypeToCompact, + thriftTypeFromCompact, + + -- * Binary Protocol — Encode + tBinEncodeI8, + tBinEncodeI16, + tBinEncodeI32, + tBinEncodeI64, + tBinEncodeBool, + tBinEncodeDouble, + tBinEncodeString, + tBinEncodeBinary, + tBinEncodeFieldBegin, + tBinEncodeFieldStop, + tBinEncodeListBegin, + tBinEncodeSetBegin, + tBinEncodeMapBegin, + + -- * Binary Protocol — Decode + tBinDecodeI8, + tBinDecodeI16, + tBinDecodeI32, + tBinDecodeI64, + tBinDecodeBool, + tBinDecodeDouble, + tBinDecodeString, + tBinDecodeBinary, + tBinDecodeFieldBegin, + tBinDecodeListBegin, + tBinDecodeSetBegin, + tBinDecodeMapBegin, + + -- * Compact Protocol — Encode + tCompEncodeI8, + tCompEncodeI16, + tCompEncodeI32, + tCompEncodeI64, + tCompEncodeBool, + tCompEncodeDouble, + tCompEncodeString, + tCompEncodeBinary, + tCompEncodeFieldBegin, + tCompEncodeFieldStop, + tCompEncodeListBegin, + tCompEncodeSetBegin, + tCompEncodeMapBegin, + tCompEncodeVarint, + tCompEncodeZigZag, + + -- * Compact Protocol — Decode + tCompDecodeI8, + tCompDecodeI16, + tCompDecodeI32, + tCompDecodeI64, + tCompDecodeBool, + tCompDecodeDouble, + tCompDecodeString, + tCompDecodeBinary, + tCompDecodeFieldBegin, + tCompDecodeListBegin, + tCompDecodeSetBegin, + tCompDecodeMapBegin, + tCompDecodeVarint, + tCompDecodeZigZag, +) where + +import Data.Bits (shiftL, shiftR, xor, (.&.), (.|.)) +import Data.ByteString (ByteString) +import Data.ByteString qualified as BS +import Data.ByteString.Unsafe qualified as BSU +import Data.Int (Int16, Int32, Int64, Int8) +import Data.Word (Word16, Word32, Word64, Word8) +import GHC.Float (castDoubleToWord64, castWord64ToDouble) +import Wireform.Builder (Builder) +import Wireform.Builder qualified as B + + +-------------------------------------------------------------------------------- +-- Thrift type codes +-------------------------------------------------------------------------------- + +-- | Thrift logical types used in both protocols. +data ThriftType + = TT_STOP + | TT_BOOL + | TT_BYTE + | TT_I16 + | TT_I32 + | TT_I64 + | TT_DOUBLE + | TT_STRING + | TT_STRUCT + | TT_MAP + | TT_LIST + | TT_SET + | TT_UUID + deriving stock (Show, Eq, Ord, Enum, Bounded) + + +-- | Binary protocol type codes. +thriftTypeToBin :: ThriftType -> Word8 +thriftTypeToBin !t = case t of + TT_STOP -> 0 + TT_BOOL -> 2 + TT_BYTE -> 3 + TT_I16 -> 6 + TT_I32 -> 8 + TT_I64 -> 10 + TT_DOUBLE -> 4 + TT_STRING -> 11 + TT_STRUCT -> 12 + TT_MAP -> 13 + TT_LIST -> 15 + TT_SET -> 14 + TT_UUID -> 16 +{-# INLINE thriftTypeToBin #-} + + +-- | Decode a binary protocol type code. +thriftTypeFromBin :: Word8 -> Maybe ThriftType +thriftTypeFromBin !w = case w of + 0 -> Just TT_STOP + 2 -> Just TT_BOOL + 3 -> Just TT_BYTE + 6 -> Just TT_I16 + 8 -> Just TT_I32 + 10 -> Just TT_I64 + 4 -> Just TT_DOUBLE + 11 -> Just TT_STRING + 12 -> Just TT_STRUCT + 13 -> Just TT_MAP + 15 -> Just TT_LIST + 14 -> Just TT_SET + 16 -> Just TT_UUID + _ -> Nothing +{-# INLINE thriftTypeFromBin #-} + + +{- | Compact protocol type codes. + +BOOLEAN_TRUE=1, BOOLEAN_FALSE=2, BYTE=3, I16=4, I32=5, I64=6, +DOUBLE=7, BINARY=8, LIST=9, SET=10, MAP=11, STRUCT=12, UUID=13 +-} +thriftTypeToCompact :: ThriftType -> Word8 +thriftTypeToCompact !t = case t of + TT_STOP -> 0 + TT_BOOL -> 1 -- BOOLEAN_TRUE by default; caller adjusts for field headers + TT_BYTE -> 3 + TT_I16 -> 4 + TT_I32 -> 5 + TT_I64 -> 6 + TT_DOUBLE -> 7 + TT_STRING -> 8 + TT_STRUCT -> 12 + TT_MAP -> 11 + TT_LIST -> 9 + TT_SET -> 10 + TT_UUID -> 13 +{-# INLINE thriftTypeToCompact #-} + + +-- | Decode a compact protocol type code. +thriftTypeFromCompact :: Word8 -> Maybe ThriftType +thriftTypeFromCompact !w = case w of + 0 -> Just TT_STOP + 1 -> Just TT_BOOL + 2 -> Just TT_BOOL + 3 -> Just TT_BYTE + 4 -> Just TT_I16 + 5 -> Just TT_I32 + 6 -> Just TT_I64 + 7 -> Just TT_DOUBLE + 8 -> Just TT_STRING + 9 -> Just TT_LIST + 10 -> Just TT_SET + 11 -> Just TT_MAP + 12 -> Just TT_STRUCT + 13 -> Just TT_UUID + _ -> Nothing +{-# INLINE thriftTypeFromCompact #-} + + +-------------------------------------------------------------------------------- +-- Internal helpers +-------------------------------------------------------------------------------- + +-- Big-endian encode/decode helpers using pokeByteOff/peekByteOff + +putBE16 :: Word16 -> Builder +putBE16 !w = + B.word8 (fromIntegral (w `shiftR` 8)) + <> B.word8 (fromIntegral (w .&. 0xFF)) +{-# INLINE putBE16 #-} + + +putBE32 :: Word32 -> Builder +putBE32 !w = + B.word8 (fromIntegral (w `shiftR` 24)) + <> B.word8 (fromIntegral ((w `shiftR` 16) .&. 0xFF)) + <> B.word8 (fromIntegral ((w `shiftR` 8) .&. 0xFF)) + <> B.word8 (fromIntegral (w .&. 0xFF)) +{-# INLINE putBE32 #-} + + +putBE64 :: Word64 -> Builder +putBE64 !w = + B.word8 (fromIntegral (w `shiftR` 56)) + <> B.word8 (fromIntegral ((w `shiftR` 48) .&. 0xFF)) + <> B.word8 (fromIntegral ((w `shiftR` 40) .&. 0xFF)) + <> B.word8 (fromIntegral ((w `shiftR` 32) .&. 0xFF)) + <> B.word8 (fromIntegral ((w `shiftR` 24) .&. 0xFF)) + <> B.word8 (fromIntegral ((w `shiftR` 16) .&. 0xFF)) + <> B.word8 (fromIntegral ((w `shiftR` 8) .&. 0xFF)) + <> B.word8 (fromIntegral (w .&. 0xFF)) +{-# INLINE putBE64 #-} + + +getBE16 :: ByteString -> Int -> Maybe (Word16, Int) +getBE16 !bs !off + | off + 2 > BS.length bs = Nothing + | otherwise = + let !b0 = fromIntegral (BSU.unsafeIndex bs off) + !b1 = fromIntegral (BSU.unsafeIndex bs (off + 1)) + !v = (b0 `shiftL` 8) .|. b1 + in Just (v, off + 2) +{-# INLINE getBE16 #-} + + +getBE32 :: ByteString -> Int -> Maybe (Word32, Int) +getBE32 !bs !off + | off + 4 > BS.length bs = Nothing + | otherwise = + let !b0 = fromIntegral (BSU.unsafeIndex bs off) + !b1 = fromIntegral (BSU.unsafeIndex bs (off + 1)) + !b2 = fromIntegral (BSU.unsafeIndex bs (off + 2)) + !b3 = fromIntegral (BSU.unsafeIndex bs (off + 3)) + !v = (b0 `shiftL` 24) .|. (b1 `shiftL` 16) .|. (b2 `shiftL` 8) .|. b3 + in Just (v, off + 4) +{-# INLINE getBE32 #-} + + +getBE64 :: ByteString -> Int -> Maybe (Word64, Int) +getBE64 !bs !off + | off + 8 > BS.length bs = Nothing + | otherwise = + let !b0 = fromIntegral (BSU.unsafeIndex bs off) :: Word64 + !b1 = fromIntegral (BSU.unsafeIndex bs (off + 1)) :: Word64 + !b2 = fromIntegral (BSU.unsafeIndex bs (off + 2)) :: Word64 + !b3 = fromIntegral (BSU.unsafeIndex bs (off + 3)) :: Word64 + !b4 = fromIntegral (BSU.unsafeIndex bs (off + 4)) :: Word64 + !b5 = fromIntegral (BSU.unsafeIndex bs (off + 5)) :: Word64 + !b6 = fromIntegral (BSU.unsafeIndex bs (off + 6)) :: Word64 + !b7 = fromIntegral (BSU.unsafeIndex bs (off + 7)) :: Word64 + !v = + (b0 `shiftL` 56) + .|. (b1 `shiftL` 48) + .|. (b2 `shiftL` 40) + .|. (b3 `shiftL` 32) + .|. (b4 `shiftL` 24) + .|. (b5 `shiftL` 16) + .|. (b6 `shiftL` 8) + .|. b7 + in Just (v, off + 8) +{-# INLINE getBE64 #-} + + +-- Little-endian 64-bit (for compact protocol doubles) +putLE64 :: Word64 -> Builder +putLE64 !w = + B.word8 (fromIntegral (w .&. 0xFF)) + <> B.word8 (fromIntegral ((w `shiftR` 8) .&. 0xFF)) + <> B.word8 (fromIntegral ((w `shiftR` 16) .&. 0xFF)) + <> B.word8 (fromIntegral ((w `shiftR` 24) .&. 0xFF)) + <> B.word8 (fromIntegral ((w `shiftR` 32) .&. 0xFF)) + <> B.word8 (fromIntegral ((w `shiftR` 40) .&. 0xFF)) + <> B.word8 (fromIntegral ((w `shiftR` 48) .&. 0xFF)) + <> B.word8 (fromIntegral ((w `shiftR` 56) .&. 0xFF)) +{-# INLINE putLE64 #-} + + +getLE64 :: ByteString -> Int -> Maybe (Word64, Int) +getLE64 !bs !off + | off + 8 > BS.length bs = Nothing + | otherwise = + let !b0 = fromIntegral (BSU.unsafeIndex bs off) :: Word64 + !b1 = fromIntegral (BSU.unsafeIndex bs (off + 1)) :: Word64 + !b2 = fromIntegral (BSU.unsafeIndex bs (off + 2)) :: Word64 + !b3 = fromIntegral (BSU.unsafeIndex bs (off + 3)) :: Word64 + !b4 = fromIntegral (BSU.unsafeIndex bs (off + 4)) :: Word64 + !b5 = fromIntegral (BSU.unsafeIndex bs (off + 5)) :: Word64 + !b6 = fromIntegral (BSU.unsafeIndex bs (off + 6)) :: Word64 + !b7 = fromIntegral (BSU.unsafeIndex bs (off + 7)) :: Word64 + !v = + b0 + .|. (b1 `shiftL` 8) + .|. (b2 `shiftL` 16) + .|. (b3 `shiftL` 24) + .|. (b4 `shiftL` 32) + .|. (b5 `shiftL` 40) + .|. (b6 `shiftL` 48) + .|. (b7 `shiftL` 56) + in Just (v, off + 8) +{-# INLINE getLE64 #-} + + +-- Varint encode/decode for compact protocol +putVarint :: Word64 -> Builder +putVarint !n + | n < 0x80 = B.word8 (fromIntegral n) + | n < 0x4000 = + B.word8 (fromIntegral (n .&. 0x7F) .|. 0x80) + <> B.word8 (fromIntegral (n `shiftR` 7)) + | n < 0x200000 = + B.word8 (fromIntegral (n .&. 0x7F) .|. 0x80) + <> B.word8 (fromIntegral ((n `shiftR` 7) .&. 0x7F) .|. 0x80) + <> B.word8 (fromIntegral (n `shiftR` 14)) + | otherwise = putVarintSlow n +{-# INLINE putVarint #-} + + +putVarintSlow :: Word64 -> Builder +putVarintSlow = go + where + go !n + | n < 0x80 = B.word8 (fromIntegral n) + | otherwise = B.word8 (fromIntegral (n .&. 0x7F) .|. 0x80) <> go (n `shiftR` 7) +{-# INLINE putVarintSlow #-} + + +getVarint :: ByteString -> Int -> Maybe (Word64, Int) +getVarint !bs !off + | off < 0 = Nothing + | otherwise = go off 0 0 + where + go !i !acc !shift + | i >= BS.length bs = Nothing + | shift >= 64 = Nothing + | otherwise = + let !b = fromIntegral (BSU.unsafeIndex bs i) :: Word64 + !acc' = acc .|. ((b .&. 0x7F) `shiftL` shift) + in if shift == 63 && b > 1 + then Nothing + else + if b < 0x80 + then Just (acc', i + 1) + else go (i + 1) acc' (shift + 7) +{-# INLINE getVarint #-} + + +-- ZigZag encoding (same as protobuf) +zigZagEncode64 :: Int64 -> Word64 +zigZagEncode64 !n = fromIntegral ((n `shiftL` 1) `xor` (n `shiftR` 63)) +{-# INLINE zigZagEncode64 #-} + + +zigZagDecode64 :: Word64 -> Int64 +zigZagDecode64 !n = fromIntegral ((n `shiftR` 1) `xor` negate (n .&. 1)) +{-# INLINE zigZagDecode64 #-} + + +zigZagEncode32 :: Int32 -> Word32 +zigZagEncode32 !n = fromIntegral ((n `shiftL` 1) `xor` (n `shiftR` 31)) +{-# INLINE zigZagEncode32 #-} + + +zigZagDecode32 :: Word32 -> Int32 +zigZagDecode32 !n = fromIntegral ((n `shiftR` 1) `xor` negate (n .&. 1)) +{-# INLINE zigZagDecode32 #-} + + +-------------------------------------------------------------------------------- +-- Binary Protocol — Encode +-------------------------------------------------------------------------------- + +tBinEncodeI8 :: Int8 -> Builder +tBinEncodeI8 !v = B.word8 (fromIntegral v) +{-# INLINE tBinEncodeI8 #-} + + +tBinEncodeI16 :: Int16 -> Builder +tBinEncodeI16 !v = putBE16 (fromIntegral v) +{-# INLINE tBinEncodeI16 #-} + + +tBinEncodeI32 :: Int32 -> Builder +tBinEncodeI32 !v = putBE32 (fromIntegral v) +{-# INLINE tBinEncodeI32 #-} + + +tBinEncodeI64 :: Int64 -> Builder +tBinEncodeI64 !v = putBE64 (fromIntegral v) +{-# INLINE tBinEncodeI64 #-} + + +tBinEncodeBool :: Bool -> Builder +tBinEncodeBool !b = B.word8 (if b then 1 else 0) +{-# INLINE tBinEncodeBool #-} + + +tBinEncodeDouble :: Double -> Builder +tBinEncodeDouble !d = putBE64 (castDoubleToWord64 d) +{-# INLINE tBinEncodeDouble #-} + + +-- | Encode a string (i32 length prefix + bytes). +tBinEncodeString :: ByteString -> Builder +tBinEncodeString !bs = + putBE32 (fromIntegral (BS.length bs)) <> B.byteString bs +{-# INLINE tBinEncodeString #-} + + +-- | Encode raw binary (same wire format as string). +tBinEncodeBinary :: ByteString -> Builder +tBinEncodeBinary = tBinEncodeString +{-# INLINE tBinEncodeBinary #-} + + +-- | Encode a struct field header: type byte + field id. +tBinEncodeFieldBegin :: ThriftType -> Int16 -> Builder +tBinEncodeFieldBegin !tt !fid = + B.word8 (thriftTypeToBin tt) <> putBE16 (fromIntegral fid) +{-# INLINE tBinEncodeFieldBegin #-} + + +-- | Encode the struct stop byte. +tBinEncodeFieldStop :: Builder +tBinEncodeFieldStop = B.word8 0x00 +{-# INLINE tBinEncodeFieldStop #-} + + +-- | Encode a list header: element type + i32 size. +tBinEncodeListBegin :: ThriftType -> Int32 -> Builder +tBinEncodeListBegin !elemType !sz = + B.word8 (thriftTypeToBin elemType) <> putBE32 (fromIntegral sz) +{-# INLINE tBinEncodeListBegin #-} + + +-- | Encode a set header (same as list header). +tBinEncodeSetBegin :: ThriftType -> Int32 -> Builder +tBinEncodeSetBegin = tBinEncodeListBegin +{-# INLINE tBinEncodeSetBegin #-} + + +-- | Encode a map header: key type + value type + i32 size. +tBinEncodeMapBegin :: ThriftType -> ThriftType -> Int32 -> Builder +tBinEncodeMapBegin !keyType !valType !sz = + B.word8 (thriftTypeToBin keyType) + <> B.word8 (thriftTypeToBin valType) + <> putBE32 (fromIntegral sz) +{-# INLINE tBinEncodeMapBegin #-} + + +-------------------------------------------------------------------------------- +-- Binary Protocol — Decode +-------------------------------------------------------------------------------- + +-- | Decode result: value + new offset, or Nothing on error. +tBinDecodeI8 :: ByteString -> Int -> Maybe (Int8, Int) +tBinDecodeI8 !bs !off + | off >= BS.length bs = Nothing + | otherwise = Just (fromIntegral (BSU.unsafeIndex bs off), off + 1) +{-# INLINE tBinDecodeI8 #-} + + +tBinDecodeI16 :: ByteString -> Int -> Maybe (Int16, Int) +tBinDecodeI16 !bs !off = case getBE16 bs off of + Just (!w, !off') -> Just (fromIntegral w, off') + Nothing -> Nothing +{-# INLINE tBinDecodeI16 #-} + + +tBinDecodeI32 :: ByteString -> Int -> Maybe (Int32, Int) +tBinDecodeI32 !bs !off = case getBE32 bs off of + Just (!w, !off') -> Just (fromIntegral w, off') + Nothing -> Nothing +{-# INLINE tBinDecodeI32 #-} + + +tBinDecodeI64 :: ByteString -> Int -> Maybe (Int64, Int) +tBinDecodeI64 !bs !off = case getBE64 bs off of + Just (!w, !off') -> Just (fromIntegral w, off') + Nothing -> Nothing +{-# INLINE tBinDecodeI64 #-} + + +tBinDecodeBool :: ByteString -> Int -> Maybe (Bool, Int) +tBinDecodeBool !bs !off + | off >= BS.length bs = Nothing + | otherwise = Just (BSU.unsafeIndex bs off /= 0, off + 1) +{-# INLINE tBinDecodeBool #-} + + +tBinDecodeDouble :: ByteString -> Int -> Maybe (Double, Int) +tBinDecodeDouble !bs !off = case getBE64 bs off of + Just (!w, !off') -> Just (castWord64ToDouble w, off') + Nothing -> Nothing +{-# INLINE tBinDecodeDouble #-} + + +-- | Decode a string: i32 length prefix + bytes. +tBinDecodeString :: ByteString -> Int -> Maybe (ByteString, Int) +tBinDecodeString !bs !off = do + (!lenW, !off1) <- getBE32 bs off + let !len = fromIntegral lenW :: Int + if len < 0 || off1 + len > BS.length bs + then Nothing + else Just (BS.take len (BS.drop off1 bs), off1 + len) +{-# INLINE tBinDecodeString #-} + + +-- | Decode raw binary (same wire format as string). +tBinDecodeBinary :: ByteString -> Int -> Maybe (ByteString, Int) +tBinDecodeBinary = tBinDecodeString +{-# INLINE tBinDecodeBinary #-} + + +{- | Decode a struct field header. Returns Nothing if at end of bytes, +returns (TT_STOP, 0) on stop byte, or (type, fieldId) otherwise. +-} +tBinDecodeFieldBegin :: ByteString -> Int -> Maybe (ThriftType, Int16, Int) +tBinDecodeFieldBegin !bs !off + | off >= BS.length bs = Nothing + | otherwise = + let !typeByte = BSU.unsafeIndex bs off + in if typeByte == 0x00 + then Just (TT_STOP, 0, off + 1) + else case thriftTypeFromBin typeByte of + Nothing -> Nothing + Just !tt -> case getBE16 bs (off + 1) of + Nothing -> Nothing + Just (!fid, !off') -> Just (tt, fromIntegral fid, off') +{-# INLINE tBinDecodeFieldBegin #-} + + +-- | Decode a list header: element type + i32 size. +tBinDecodeListBegin :: ByteString -> Int -> Maybe (ThriftType, Int32, Int) +tBinDecodeListBegin !bs !off + | off >= BS.length bs = Nothing + | otherwise = + let !typeByte = BSU.unsafeIndex bs off + in case thriftTypeFromBin typeByte of + Nothing -> Nothing + Just !tt -> case getBE32 bs (off + 1) of + Nothing -> Nothing + Just (!sz, !off') -> Just (tt, fromIntegral sz, off') +{-# INLINE tBinDecodeListBegin #-} + + +-- | Decode a set header (same as list). +tBinDecodeSetBegin :: ByteString -> Int -> Maybe (ThriftType, Int32, Int) +tBinDecodeSetBegin = tBinDecodeListBegin +{-# INLINE tBinDecodeSetBegin #-} + + +-- | Decode a map header: key type + value type + i32 size. +tBinDecodeMapBegin :: ByteString -> Int -> Maybe (ThriftType, ThriftType, Int32, Int) +tBinDecodeMapBegin !bs !off + | off + 2 > BS.length bs = Nothing + | otherwise = + let !kByte = BSU.unsafeIndex bs off + !vByte = BSU.unsafeIndex bs (off + 1) + in case (thriftTypeFromBin kByte, thriftTypeFromBin vByte) of + (Just !kt, Just !vt) -> case getBE32 bs (off + 2) of + Nothing -> Nothing + Just (!sz, !off') -> Just (kt, vt, fromIntegral sz, off') + _ -> Nothing +{-# INLINE tBinDecodeMapBegin #-} + + +-------------------------------------------------------------------------------- +-- Compact Protocol — Encode +-------------------------------------------------------------------------------- + +-- | Encode a byte (compact protocol). +tCompEncodeI8 :: Int8 -> Builder +tCompEncodeI8 !v = B.word8 (fromIntegral v) +{-# INLINE tCompEncodeI8 #-} + + +-- | Encode i16 as zigzag + varint. +tCompEncodeI16 :: Int16 -> Builder +tCompEncodeI16 !v = putVarint (fromIntegral (zigZagEncode32 (fromIntegral v))) +{-# INLINE tCompEncodeI16 #-} + + +-- | Encode i32 as zigzag + varint. +tCompEncodeI32 :: Int32 -> Builder +tCompEncodeI32 !v = putVarint (fromIntegral (zigZagEncode32 v)) +{-# INLINE tCompEncodeI32 #-} + + +-- | Encode i64 as zigzag + varint. +tCompEncodeI64 :: Int64 -> Builder +tCompEncodeI64 !v = putVarint (zigZagEncode64 v) +{-# INLINE tCompEncodeI64 #-} + + +{- | Encode a bool. In compact protocol, booleans are typically encoded +in the field header. This standalone version uses a single byte. +-} +tCompEncodeBool :: Bool -> Builder +tCompEncodeBool !b = B.word8 (if b then 1 else 2) +{-# INLINE tCompEncodeBool #-} + + +-- | Encode a double as 8 bytes LITTLE-endian IEEE 754. +tCompEncodeDouble :: Double -> Builder +tCompEncodeDouble !d = putLE64 (castDoubleToWord64 d) +{-# INLINE tCompEncodeDouble #-} + + +-- | Encode a string: varint length + bytes. +tCompEncodeString :: ByteString -> Builder +tCompEncodeString !bs = + putVarint (fromIntegral (BS.length bs)) <> B.byteString bs +{-# INLINE tCompEncodeString #-} + + +-- | Encode raw binary (same wire format as string in compact). +tCompEncodeBinary :: ByteString -> Builder +tCompEncodeBinary = tCompEncodeString +{-# INLINE tCompEncodeBinary #-} + + +{- | Encode a compact field header using delta encoding. + +If the delta from the previous field ID fits in 4 bits (1-15), +use the short form: (delta << 4 | compactType). +Otherwise, use the long form: (0x00 | compactType), then zigzag varint field ID. + +For bool fields, the compact type encodes the value: 1=true, 2=false. +-} +tCompEncodeFieldBegin :: ThriftType -> Int16 -> Int16 -> Bool -> Builder +tCompEncodeFieldBegin !tt !fid !lastFid !boolVal = + let !delta = fid - lastFid + !ctype = case tt of + TT_BOOL -> if boolVal then 1 else 2 + _ -> thriftTypeToCompact tt + in if delta > 0 && delta <= 15 + then B.word8 (fromIntegral delta `shiftL` 4 .|. ctype) + else B.word8 ctype <> putVarint (fromIntegral (zigZagEncode32 (fromIntegral fid))) +{-# INLINE tCompEncodeFieldBegin #-} + + +-- | Encode the struct stop byte (compact). +tCompEncodeFieldStop :: Builder +tCompEncodeFieldStop = B.word8 0x00 +{-# INLINE tCompEncodeFieldStop #-} + + +{- | Encode a compact list header. + +If size < 15: single byte (size << 4 | elemType). +Else: (0xF0 | elemType), varint size. +-} +tCompEncodeListBegin :: ThriftType -> Int32 -> Builder +tCompEncodeListBegin !elemType !sz = + let !ctype = thriftTypeToCompact elemType + in if sz < 15 + then B.word8 (fromIntegral sz `shiftL` 4 .|. ctype) + else B.word8 (0xF0 .|. ctype) <> putVarint (fromIntegral sz) +{-# INLINE tCompEncodeListBegin #-} + + +-- | Encode a compact set header (same as list). +tCompEncodeSetBegin :: ThriftType -> Int32 -> Builder +tCompEncodeSetBegin = tCompEncodeListBegin +{-# INLINE tCompEncodeSetBegin #-} + + +{- | Encode a compact map header. + +If empty: single 0x00 byte. +Else: varint size, then (keyType << 4 | valType). +-} +tCompEncodeMapBegin :: ThriftType -> ThriftType -> Int32 -> Builder +tCompEncodeMapBegin !keyType !valType !sz + | sz == 0 = B.word8 0x00 + | otherwise = + putVarint (fromIntegral sz) + <> B.word8 (thriftTypeToCompact keyType `shiftL` 4 .|. thriftTypeToCompact valType) +{-# INLINE tCompEncodeMapBegin #-} + + +-- | Raw varint encode (exposed for external use). +tCompEncodeVarint :: Word64 -> Builder +tCompEncodeVarint = putVarint +{-# INLINE tCompEncodeVarint #-} + + +-- | ZigZag encode (exposed for external use). +tCompEncodeZigZag :: Int64 -> Word64 +tCompEncodeZigZag = zigZagEncode64 +{-# INLINE tCompEncodeZigZag #-} + + +-------------------------------------------------------------------------------- +-- Compact Protocol — Decode +-------------------------------------------------------------------------------- + +-- | Decode a byte (compact protocol). +tCompDecodeI8 :: ByteString -> Int -> Maybe (Int8, Int) +tCompDecodeI8 !bs !off + | off >= BS.length bs = Nothing + | otherwise = Just (fromIntegral (BSU.unsafeIndex bs off), off + 1) +{-# INLINE tCompDecodeI8 #-} + + +-- | Decode i16 from zigzag + varint. +tCompDecodeI16 :: ByteString -> Int -> Maybe (Int16, Int) +tCompDecodeI16 !bs !off = case getVarint bs off of + Just (!w, !off') -> Just (fromIntegral (zigZagDecode32 (fromIntegral w)), off') + Nothing -> Nothing +{-# INLINE tCompDecodeI16 #-} + + +-- | Decode i32 from zigzag + varint. +tCompDecodeI32 :: ByteString -> Int -> Maybe (Int32, Int) +tCompDecodeI32 !bs !off = case getVarint bs off of + Just (!w, !off') -> Just (zigZagDecode32 (fromIntegral w), off') + Nothing -> Nothing +{-# INLINE tCompDecodeI32 #-} + + +-- | Decode i64 from zigzag + varint. +tCompDecodeI64 :: ByteString -> Int -> Maybe (Int64, Int) +tCompDecodeI64 !bs !off = case getVarint bs off of + Just (!w, !off') -> Just (zigZagDecode64 w, off') + Nothing -> Nothing +{-# INLINE tCompDecodeI64 #-} + + +-- | Decode a bool (standalone, not from field header). +tCompDecodeBool :: ByteString -> Int -> Maybe (Bool, Int) +tCompDecodeBool !bs !off + | off < 0 || off >= BS.length bs = Nothing + | otherwise = case BSU.unsafeIndex bs off of + 1 -> Just (True, off + 1) + 2 -> Just (False, off + 1) + -- Read older wireform files whose element booleans used binary's 0. + -- New compact writers always emit the standard false marker, 2. + 0 -> Just (False, off + 1) + _ -> Nothing +{-# INLINE tCompDecodeBool #-} + + +-- | Decode a double as 8 bytes LITTLE-endian. +tCompDecodeDouble :: ByteString -> Int -> Maybe (Double, Int) +tCompDecodeDouble !bs !off = case getLE64 bs off of + Just (!w, !off') -> Just (castWord64ToDouble w, off') + Nothing -> Nothing +{-# INLINE tCompDecodeDouble #-} + + +-- | Decode a string: varint length + bytes. +tCompDecodeString :: ByteString -> Int -> Maybe (ByteString, Int) +tCompDecodeString !bs !off = do + (!lenW, !off1) <- getVarint bs off + if lenW > fromIntegral (maxBound :: Int32) || lenW > fromIntegral (BS.length bs - off1) + then Nothing + else + let !len = fromIntegral lenW :: Int + in Just (BS.take len (BS.drop off1 bs), off1 + len) +{-# INLINE tCompDecodeString #-} + + +-- | Decode raw binary (same wire format as string in compact). +tCompDecodeBinary :: ByteString -> Int -> Maybe (ByteString, Int) +tCompDecodeBinary = tCompDecodeString +{-# INLINE tCompDecodeBinary #-} + + +{- | Decode a compact field header. + +Returns (ThriftType, fieldId, newOffset, boolValue). +The boolValue is meaningful only when the type is TT_BOOL. +Returns (TT_STOP, 0, off+1, False) on stop byte. + +Takes the previous field ID for delta decoding. +-} +tCompDecodeFieldBegin :: ByteString -> Int -> Int16 -> Maybe (ThriftType, Int16, Int, Bool) +tCompDecodeFieldBegin !bs !off !lastFid + | off < 0 || off >= BS.length bs = Nothing + | otherwise = + let !hdr = BSU.unsafeIndex bs off + in if hdr == 0x00 + then Just (TT_STOP, 0, off + 1, False) + else + let !ctype = hdr .&. 0x0F + !delta = hdr `shiftR` 4 + !boolVal = ctype == 1 -- BOOLEAN_TRUE + in case thriftTypeFromCompact ctype of + Nothing -> Nothing + Just !tt -> + if delta /= 0 + then Just (tt, lastFid + fromIntegral delta, off + 1, boolVal) + else case getVarint bs (off + 1) of + Nothing -> Nothing + Just (!zz, !off') -> + let !fid = fromIntegral (zigZagDecode32 (fromIntegral zz)) + in Just (tt, fid, off', boolVal) +{-# INLINE tCompDecodeFieldBegin #-} + + +{- | Decode a compact list header. +Returns (elemType, size, newOffset). +-} +tCompDecodeListBegin :: ByteString -> Int -> Maybe (ThriftType, Int32, Int) +tCompDecodeListBegin !bs !off + | off < 0 || off >= BS.length bs = Nothing + | otherwise = + let !hdr = BSU.unsafeIndex bs off + !szHi = hdr `shiftR` 4 + !ctype = hdr .&. 0x0F + in case thriftTypeFromCompact ctype of + Nothing -> Nothing + Just !tt -> + if szHi /= 0x0F + then Just (tt, fromIntegral szHi, off + 1) + else case getVarint bs (off + 1) of + Nothing -> Nothing + Just (!sz, !off') + | sz > fromIntegral (maxBound :: Int32) -> Nothing + | otherwise -> Just (tt, fromIntegral sz, off') +{-# INLINE tCompDecodeListBegin #-} + + +-- | Decode a compact set header (same as list). +tCompDecodeSetBegin :: ByteString -> Int -> Maybe (ThriftType, Int32, Int) +tCompDecodeSetBegin = tCompDecodeListBegin +{-# INLINE tCompDecodeSetBegin #-} + + +{- | Decode a compact map header. +Returns (keyType, valType, size, newOffset). +For empty maps: a single 0x00 byte (size=0). +-} +tCompDecodeMapBegin :: ByteString -> Int -> Maybe (ThriftType, ThriftType, Int32, Int) +tCompDecodeMapBegin !bs !off = do + (!szW, !off1) <- getVarint bs off + if szW > fromIntegral (maxBound :: Int32) then Nothing else Just () + let !sz = fromIntegral szW :: Int32 + if sz == 0 + then Just (TT_STOP, TT_STOP, 0, off1) + else do + if off1 >= BS.length bs + then Nothing + else do + let !kvByte = BSU.unsafeIndex bs off1 + !kCode = kvByte `shiftR` 4 + !vCode = kvByte .&. 0x0F + kt <- thriftTypeFromCompact kCode + vt <- thriftTypeFromCompact vCode + Just (kt, vt, sz, off1 + 1) +{-# INLINE tCompDecodeMapBegin #-} + + +-- | Raw varint decode (exposed for external use). +tCompDecodeVarint :: ByteString -> Int -> Maybe (Word64, Int) +tCompDecodeVarint = getVarint +{-# INLINE tCompDecodeVarint #-} + + +-- | ZigZag decode (exposed for external use). +tCompDecodeZigZag :: Word64 -> Int64 +tCompDecodeZigZag = zigZagDecode64 +{-# INLINE tCompDecodeZigZag #-} diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,16 @@ +module Main (main) where + +import Test.Syd +import Test.Thrift.CompactBool qualified +import Test.Thrift.Derive qualified + + +main :: IO () +main = + sydTest + ( describe "wireform-thrift" $ + sequence_ + [ Test.Thrift.Derive.tests + , Test.Thrift.CompactBool.tests + ] + ) diff --git a/test/Test/Thrift/CompactBool.hs b/test/Test/Thrift/CompactBool.hs new file mode 100644 --- /dev/null +++ b/test/Test/Thrift/CompactBool.hs @@ -0,0 +1,343 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Test.Thrift.CompactBool (tests) where + +import Control.Exception (evaluate) +import Control.Monad (forM_, when) +import Data.ByteString qualified as BS +import Data.ByteString.Internal qualified as BSI +import Data.ByteString.Lazy qualified as BL +import Data.Either (isLeft) +import Data.Text qualified as T +import Data.Text.Encoding qualified as Text +import Data.Vector qualified as V +import Data.Vector.Generic.Base qualified as VG +import Data.Word (Word32) +import Foreign.ForeignPtr (mallocForeignPtrBytes) +import GHC.Conc (getAllocationCounter) +import GHC.Float (castWord64ToDouble) +import Hedgehog +import Hedgehog.Gen qualified as Gen +import Hedgehog.Range qualified as Range +import Numeric (readHex) +import Test.Syd +import Test.Syd.Hedgehog () +import Thrift.Decode (decodeBinary, decodeCompact, decodeCompactBinaryStrings, decodeCompactFrom) +import Thrift.Encode (binaryEncodedSizeBounded, compactEncodedSizeBounded, encodeBinaryBounded, encodeCompact, encodeCompactBounded) +import Thrift.Value qualified as Value +import Thrift.Wire qualified as Wire +import Wireform.Builder qualified as Builder + + +tests :: Spec +tests = describe "Compact container booleans" $ do + encodeBoundsTests + protocolGoldenTests + let list = Value.Struct (V.singleton (1, Value.List Wire.TT_BOOL (V.fromList [Value.Bool False, Value.Bool True, Value.Bool False]))) + listBytes = BS.pack [0x19, 0x31, 2, 1, 2, 0] + set = Value.Struct (V.singleton (1, Value.Set Wire.TT_BOOL (V.fromList [Value.Bool False, Value.Bool True]))) + setBytes = BS.pack [0x1a, 0x21, 2, 1, 0] + mapValue = + Value.Struct + ( V.singleton + ( 1 + , Value.Map + Wire.TT_BOOL + Wire.TT_BOOL + (V.fromList [(Value.Bool False, Value.Bool True), (Value.Bool True, Value.Bool False)]) + ) + ) + mapBytes = BS.pack [0x1b, 2, 0x11, 2, 1, 1, 2, 0] + it "writes the published list byte sequence" $ encodeCompact list `shouldBe` listBytes + it "reads standard false/true/false list elements" $ decodeCompact listBytes `shouldBe` Right list + it "writes standard boolean set elements" $ encodeCompact set `shouldBe` setBytes + it "reads standard boolean set elements" $ decodeCompact setBytes `shouldBe` Right set + it "writes standard boolean map keys and values" $ encodeCompact mapValue `shouldBe` mapBytes + it "reads standard boolean map keys and values" $ decodeCompact mapBytes `shouldBe` Right mapValue + it "uses marker2 in the builder API as well" $ + BL.toStrict (Builder.toLazyByteString (Wire.tCompEncodeBool False)) `shouldBe` BS.singleton 2 + it "retains binary protocol's separate false marker" $ + BL.toStrict (Builder.toLazyByteString (Wire.tBinEncodeBool False)) `shouldBe` BS.singleton 0 + it "reads old wireform zero markers but rewrites standard markers" $ do + decodeCompact (BS.pack [0x19, 0x31, 0, 1, 0, 0]) `shouldBe` Right list + encodeCompact list `shouldBe` listBytes + it "rejects invalid markers and offsets" $ do + Wire.tCompDecodeBool (BS.singleton 3) 0 `shouldBe` Nothing + Wire.tCompDecodeBool (BS.singleton 1) (-1) `shouldBe` Nothing + Wire.tCompDecodeBool (BS.singleton 1) 1 `shouldBe` Nothing + it "preserves arbitrary mixed boolean lists" $ property $ do + values <- forAll (Gen.list (Range.linear 0 100) Gen.bool) + let value = Value.Struct (V.singleton (1, Value.List Wire.TT_BOOL (V.fromList (Value.Bool <$> values)))) + decodeCompact (encodeCompact value) === Right value + describe "opaque Compact strings" $ do + let nested payloads = + let binaries = Value.Binary <$> V.fromList payloads + bools = Value.List Wire.TT_BOOL (V.fromList [Value.Bool False, Value.Bool True]) + -- Compact omits type tags for an empty map. + opaqueMap key value entries + | V.null entries = Value.Map Wire.TT_STOP Wire.TT_STOP entries + | otherwise = Value.Map key value entries + in Value.Struct + ( V.fromList + [ (1, Value.List Wire.TT_STRING binaries) + , (2, Value.Set Wire.TT_STRING binaries) + , (3, opaqueMap Wire.TT_STRING Wire.TT_LIST (V.map (\b -> (b, bools)) binaries)) + , (4, Value.Struct (V.singleton (1, Value.List Wire.TT_STRING binaries))) + , (5, opaqueMap Wire.TT_LIST Wire.TT_STRING (V.map (\b -> (bools, b)) binaries)) + ] + ) + it "preserves UTF-8, NULs, empty and non-UTF-8 bytes in every container" $ do + let value = nested [BS.empty, "plain", Text.encodeUtf8 "é雪", BS.pack [0, 1, 0, 127], BS.pack [255, 192, 128]] + decodeCompactBinaryStrings (encodeCompact value) `shouldBe` Right value + it "keeps the default decoder's text interpretation" $ do + let value = Value.Struct (V.singleton (1, Value.String "plain-é")) + binary = Value.Struct (V.singleton (1, Value.Binary (Text.encodeUtf8 "plain-é"))) + bytes = encodeCompact value + decodeCompact bytes `shouldBe` Right value + decodeCompactBinaryStrings bytes `shouldBe` Right binary + it "preserves arbitrary sliced binary payloads in nested containers" $ property $ do + payloads <- forAll (Gen.list (Range.linear 0 40) (Gen.bytes (Range.linear 0 256))) + wirePrefix <- forAll (Gen.bytes (Range.linear 1 16)) + let value = nested payloads + wire = BS.drop (BS.length wirePrefix) (wirePrefix <> encodeCompact value) + decodeCompactBinaryStrings wire === Right value + it "rejects impossible compact container counts before allocating" $ + mapM_ + (\bytes -> isLeft (decodeCompact (BS.pack bytes)) `shouldBe` True) + [ [0x19, 0xf1, 0xff, 0xff, 0xff, 0xff, 0x07, 0] + , [0x1a, 0xf1, 0xff, 0xff, 0xff, 0xff, 0x07, 0] + , [0x1b, 0xff, 0xff, 0xff, 0xff, 0x07, 0x11, 0] + , [0x19, 0xf1, 0x80, 0x80, 0x80, 0x80, 0x08, 0] + , [0x1a, 0xf1, 0x80, 0x80, 0x80, 0x80, 0x08, 0] + , [0x1b, 0x80, 0x80, 0x80, 0x80, 0x08, 0x11, 0] + , [0x19, 0xf1, 0x80, 0x80, 0x80, 0x80, 0x10, 0] + , [0x1a, 0xf1, 0x80, 0x80, 0x80, 0x80, 0x10, 0] + , [0x1b, 0x80, 0x80, 0x80, 0x80, 0x10, 0] + , [0x19, 0xf1, 0x81, 0x80, 0x80, 0x80, 0x10, 1, 0] + , [0x1a, 0xf1, 0x81, 0x80, 0x80, 0x80, 0x10, 1, 0] + , [0x1b, 0x81, 0x80, 0x80, 0x80, 0x10, 0x11, 1, 1, 0] + ] + it "rejects overflowing varints, wrapped string lengths and negative root offsets" $ do + Wire.tCompDecodeVarint (BS.pack (replicate 9 0x80 ++ [2])) 0 `shouldBe` Nothing + Wire.tCompDecodeVarint (BS.singleton 0) (-1) `shouldBe` Nothing + Wire.tCompDecodeString (BS.pack [0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 1]) 0 `shouldBe` Nothing + Wire.tCompDecodeListBegin (BS.singleton 0) (-1) `shouldBe` Nothing + isLeft (decodeCompactFrom (BS.singleton 0) (-1)) `shouldBe` True + it "rejects negative and impossible binary container counts" $ + mapM_ + (\bytes -> isLeft (decodeBinary (BS.pack bytes)) `shouldBe` True) + [ [15, 0, 1, 2, 0xff, 0xff, 0xff, 0xff, 0] + , [14, 0, 1, 2, 0xff, 0xff, 0xff, 0xff, 0] + , [13, 0, 1, 2, 2, 0xff, 0xff, 0xff, 0xff, 0] + , [15, 0, 1, 2, 0x7f, 0xff, 0xff, 0xff, 0] + , [14, 0, 1, 2, 0x7f, 0xff, 0xff, 0xff, 0] + , [13, 0, 1, 2, 2, 0x7f, 0xff, 0xff, 0xff, 0] + ] + + +data EncoderCase = EncoderCase + { encoderLabel :: !String + , boundedEncode :: !(Int -> Value.Value -> Either String BS.ByteString) + , boundedSize :: !(Int -> Value.Value -> Either String Int) + } + + +encoderCases :: [EncoderCase] +encoderCases = [EncoderCase "Binary" encodeBinaryBounded binaryEncodedSizeBounded, EncoderCase "Compact" encodeCompactBounded compactEncodedSizeBounded] + + +encodeBoundsTests :: Spec +encodeBoundsTests = describe "Bounded Binary/Compact encoders" $ do + forM_ encoderCases $ \encoder -> describe (encoderLabel encoder) $ do + it "validates UUID width before writing a fixed-size buffer" $ + forM_ [0, 1, 15, 17, 32] $ \n -> + boundedEncode encoder maxBound (Value.UUID (BS.replicate n 0)) + `shouldBe` Left "Thrift.Encode: UUID must contain exactly 16 bytes" + it "writes exactly sixteen sliced UUID bytes" $ do + let bytes = BS.pack [0 .. 15] + sliced = BS.take 16 (BS.drop 3 ("pad" <> bytes <> "tail")) + boundedEncode encoder 16 (Value.UUID sliced) `shouldBe` Right bytes + isLeft (boundedEncode encoder 15 (Value.UUID sliced)) `shouldBe` True + it "rejects negative budgets and mismatched container tags" $ do + isLeft (boundedEncode encoder (-1) (Value.Bool True)) `shouldBe` True + forM_ + [ Value.List Wire.TT_I32 (V.singleton (Value.Bool True)) + , Value.Set Wire.TT_STRING (V.singleton (Value.I64 1)) + , Value.Map Wire.TT_BOOL Wire.TT_I32 (V.singleton (Value.I32 1, Value.I32 2)) + , Value.Map Wire.TT_BOOL Wire.TT_I32 (V.singleton (Value.Bool True, Value.String "x")) + ] + (\value -> isLeft (boundedEncode encoder maxBound value) `shouldBe` True) + it "measures UTF8 bytes without allocating the encoded text" $ do + let text = T.replicate (1024 * 1024) "é" + _ <- evaluate (T.length text) + allocationBefore <- getAllocationCounter + rejected <- evaluate (isLeft (boundedEncode encoder 32 (Value.String text))) + allocationAfter <- getAllocationCounter + rejected `shouldBe` True + (allocationBefore - allocationAfter < 65536) `shouldBe` True + it "matches actual encoded lengths and exact byte budgets for nested values" $ property $ do + value <- forAll (genValue 3) + bytes <- evalEither (boundedEncode encoder maxBound value) + let size = BS.length bytes + boundedSize encoder maxBound value === Right size + boundedEncode encoder size value === Right bytes + assert (isLeft (boundedEncode encoder (size - 1) value)) + it "checks impossible container views before traversing backing storage" $ do + let elements = VG.basicUnsafeSlice 0 maxBound (V.singleton (Value.Bool True)) + entries = VG.basicUnsafeSlice 0 maxBound (V.singleton (Value.Bool True, Value.Bool False)) + fields = VG.basicUnsafeSlice 0 maxBound (V.singleton (1, Value.Bool True)) + forM_ [Value.List Wire.TT_BOOL elements, Value.Set Wire.TT_BOOL elements, Value.Map Wire.TT_BOOL Wire.TT_BOOL entries, Value.Struct fields] $ \value -> + isLeft (boundedEncode encoder 16 value) `shouldBe` True + it "preserves compact double bits at an unaligned byte offset" $ do + let number = castWord64ToDouble 0x7ff8000000000042 + value = Value.Struct (V.fromList [(1, Value.Byte 0), (2, Value.Double number)]) + bytes <- either fail pure (encodeCompactBounded maxBound value) + BS.take 8 (BS.drop 3 bytes) `shouldBe` BS.pack [66, 0, 0, 0, 0, 0, 248, 127] + it "uses full field headers across the signed16 field-id range" $ do + let value = Value.Struct (V.fromList [(maxBound, Value.I32 1), (minBound, Value.I32 2), (0, Value.Bool False)]) + bytes <- either fail pure (encodeCompactBounded maxBound value) + compactEncodedSizeBounded maxBound value `shouldBe` Right (BS.length bytes) + decodeCompact bytes `shouldBe` Right (Value.Struct (V.fromList [(minBound, Value.I32 2), (0, Value.Bool False), (maxBound, Value.I32 1)])) + it "checks signed Binary and unsigned Compact length fields before copying" $ + when (toInteger (maxBound :: Int) > toInteger (maxBound :: Word32)) $ do + storage <- mallocForeignPtrBytes 1 + let bytes n = Value.Binary (BSI.fromForeignPtr storage 0 n) + signed = 2147483647 + unsigned = 4294967295 + binaryEncodedSizeBounded maxBound (bytes signed) `shouldBe` Right (signed + 4) + isLeft (encodeBinaryBounded maxBound (bytes (signed + 1))) `shouldBe` True + compactEncodedSizeBounded maxBound (bytes unsigned) `shouldBe` Right (unsigned + 5) + isLeft (encodeCompactBounded maxBound (bytes (unsigned + 1))) `shouldBe` True + isLeft (encodeCompactBounded 1024 (bytes unsigned)) `shouldBe` True + it "rejects cumulative sizes without materializing repeated large payloads" $ do + let bytes = BS.replicate (1024 * 1024) 0 + value = Value.List Wire.TT_STRING (V.replicate 4096 (Value.Binary bytes)) + _ <- evaluate (BS.length bytes) + forM_ encoderCases $ \encoder -> do + allocationBefore <- getAllocationCounter + rejected <- evaluate (isLeft (boundedEncode encoder 0xffffffff value)) + allocationAfter <- getAllocationCounter + rejected `shouldBe` True + (allocationBefore - allocationAfter < 2 * 1024 * 1024) `shouldBe` True + + +genValue :: Int -> Gen Value.Value +genValue depth = do + typ <- Gen.element (scalarTypes ++ if depth <= 0 then [] else [Wire.TT_STRUCT, Wire.TT_LIST, Wire.TT_SET, Wire.TT_MAP]) + genTyped depth typ + where + scalarTypes = [Wire.TT_BOOL, Wire.TT_BYTE, Wire.TT_I16, Wire.TT_I32, Wire.TT_I64, Wire.TT_DOUBLE, Wire.TT_STRING, Wire.TT_UUID] + + +genTyped :: Int -> Wire.ThriftType -> Gen Value.Value +genTyped depth typ = case typ of + Wire.TT_BOOL -> Value.Bool <$> Gen.bool + Wire.TT_BYTE -> Value.Byte <$> Gen.int8 Range.linearBounded + Wire.TT_I16 -> Value.I16 <$> Gen.int16 Range.linearBounded + Wire.TT_I32 -> Value.I32 <$> Gen.int32 Range.linearBounded + Wire.TT_I64 -> Value.I64 <$> Gen.int64 Range.linearBounded + Wire.TT_DOUBLE -> Value.Double <$> Gen.double (Range.linearFrac (-10000) 10000) + Wire.TT_STRING -> Gen.choice [Value.Binary <$> Gen.bytes (Range.linear 0 64), Value.String <$> Gen.text (Range.linear 0 32) Gen.unicode] + Wire.TT_UUID -> Value.UUID <$> Gen.bytes (Range.singleton 16) + Wire.TT_STRUCT -> do + fields <- Gen.list (Range.linear 0 6) $ do + fid <- Gen.int16 Range.linearBounded + value <- genValue (depth - 1) + pure (fid, value) + pure (Value.Struct (V.fromList fields)) + Wire.TT_LIST -> sequenceOf Value.List + Wire.TT_SET -> sequenceOf Value.Set + Wire.TT_MAP -> do + key <- genValue (depth - 1) + value <- genValue (depth - 1) + let kt = Value.thriftTypeOf key + vt = Value.thriftTypeOf value + entries <- Gen.list (Range.linear 0 6) ((,) <$> genTyped (depth - 1) kt <*> genTyped (depth - 1) vt) + pure (Value.Map kt vt (V.fromList entries)) + Wire.TT_STOP -> pure (Value.Struct V.empty) + where + sequenceOf constructor = do + sample <- genValue (depth - 1) + let elementType = Value.thriftTypeOf sample + elements <- V.fromList <$> Gen.list (Range.linear 0 8) (genTyped (depth - 1) elementType) + pure (constructor elementType elements) + + +-- Independent expected bytes are generated by Apache's Python implementation. +-- Keep these values separate from the Python fixture definitions so that a +-- writer regression cannot update its own expected wire representation. +data ProtocolFixture = ProtocolFixture + { protocolFixtureName :: !String + , protocolFixtureValue :: !Value.Value + } + + +protocolGoldenTests :: Spec +protocolGoldenTests = describe "Apache Thrift protocol goldens" $ + it "matches both reference protocols for scalars, field ids and nested containers" $ do + contents <- readFile "test/fixtures/protocol-encodings.tsv" + let dataRow [] = False + dataRow ('#' : _) = False + dataRow _ = True + rows = filter dataRow (lines contents) + length rows `shouldBe` length protocolFixtures + forM_ rows $ \line -> case words line of + [name, binaryHex, compactHex] -> case filter ((== name) . protocolFixtureName) protocolFixtures of + [fixture] -> do + binary <- parseHex binaryHex + compact <- parseHex compactHex + encodeBinaryBounded maxBound (protocolFixtureValue fixture) `shouldBe` Right binary + encodeCompactBounded maxBound (protocolFixtureValue fixture) `shouldBe` Right compact + _ -> fail ("unknown or duplicate protocol fixture: " ++ name) + _ -> fail "malformed protocol fixture row" + where + parseHex [] = pure BS.empty + parseHex (a : b : rest) = case readHex [a, b] of + [(byte, "")] -> BS.cons byte <$> parseHex rest + _ -> fail "invalid protocol fixture hex" + parseHex _ = fail "odd protocol fixture hex length" + + +protocolFixtures :: [ProtocolFixture] +protocolFixtures = + [ ProtocolFixture "scalars" $ + struct + [ (1, Value.Bool False) + , (2, Value.Bool True) + , (3, Value.Byte minBound) + , (4, Value.I16 minBound) + , (5, Value.I32 minBound) + , (6, Value.I64 minBound) + , (7, Value.Double (-123.5)) + , (8, Value.String "é雪\0") + , (9, Value.Binary (BS.pack [0, 255, 128, 1])) + , (10, Value.UUID (BS.pack [0 .. 15])) + ] + , ProtocolFixture "field-ids" $ + struct [(maxBound, Value.I32 1), (minBound, Value.I32 2), (0, Value.Bool False)] + , ProtocolFixture "container-booleans" $ + struct + [ (1, Value.List Wire.TT_BOOL (V.fromList [Value.Bool False, Value.Bool True, Value.Bool False])) + , (2, Value.Set Wire.TT_BOOL (V.fromList [Value.Bool False, Value.Bool True])) + , (3, Value.Map Wire.TT_BOOL Wire.TT_BOOL (V.fromList [(Value.Bool False, Value.Bool True), (Value.Bool True, Value.Bool False)])) + ] + ] + ++ map + (\count -> ProtocolFixture ("list-" ++ show count) (struct [(1, Value.List Wire.TT_I32 (V.fromList (map (Value.I32 . fromIntegral) [-64 .. count - 65])))])) + ([0, 14, 15, 128] :: [Int]) + ++ [ ProtocolFixture "nested" $ + struct + [ (1, Value.List Wire.TT_STRUCT (V.fromList [struct [(1, Value.String "first"), (2, Value.Bool False)], struct [(1, Value.String "雪"), (2, Value.Bool True)]])) + , + ( 2 + , Value.Map Wire.TT_STRING Wire.TT_LIST $ + V.fromList + [ (Value.String "a", Value.List Wire.TT_BOOL (V.fromList [Value.Bool False, Value.Bool True])) + , (Value.String "é", Value.List Wire.TT_BOOL V.empty) + ] + ) + , (3, Value.Map Wire.TT_STOP Wire.TT_STOP V.empty) + ] + ] + where + struct = Value.Struct . V.fromList diff --git a/test/Test/Thrift/Derive.hs b/test/Test/Thrift/Derive.hs new file mode 100644 --- /dev/null +++ b/test/Test/Thrift/Derive.hs @@ -0,0 +1,118 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Test.Thrift.Derive (tests) where + +import Data.Vector qualified as V +import Test.Syd +import Test.Thrift.Derive.Instances () +import Test.Thrift.Derive.Types +import Thrift.Class qualified as TC +import Thrift.Value qualified as TV + + +tests :: Spec +tests = + describe "Thrift.Derive" $ + sequence_ + [ recordTests + , newtypeTests + , enumTests + , unionTests + ] + + +-- --------------------------------------------------------------------------- + +recordTests :: Spec +recordTests = + describe "record" $ + sequence_ + [ it "field IDs default to positional, tag overrides, skip drops" $ do + let e = LogEntry 1700000000 "boot" 42 "abc" + case TC.toThrift e of + TV.Struct kvs -> do + (V.any (idIs 1) kvs) `shouldBe` True + (V.any (idIs 2) kvs) `shouldBe` True + (V.any (idIs 7) kvs) `shouldBe` True + -- Note: positional default for logCode (id 3) is shadowed by + -- the explicit `tag 7`, so id 3 must NOT be present. + (not (V.any (idIs 3) kvs)) `shouldBe` True + (not (V.any (idIs 4) kvs)) `shouldBe` True + v -> expectationFailure ("expected Struct, got " ++ show v) + , it "round-trip fills skipped from defaults" $ do + let e = LogEntry 1700000000 "boot" 42 "abc" + case TC.fromThrift (TC.toThrift e) of + Right e' -> do + logTimestamp e' `shouldBe` logTimestamp e + logMessage e' `shouldBe` logMessage e + logCode e' `shouldBe` logCode e + logRequestId e' `shouldBe` defaultRequestId + Left err -> expectationFailure err + ] + where + idIs n (k, _) = k == n + + +-- --------------------------------------------------------------------------- + +newtypeTests :: Spec +newtypeTests = + describe "newtype" $ + sequence_ + [ it "round-trip" $ + TC.fromThrift (TC.toThrift (RequestId 7)) `shouldBe` Right (RequestId 7) + ] + + +-- --------------------------------------------------------------------------- + +enumTests :: Spec +enumTests = + describe "enum (I32-encoded)" $ + sequence_ + [ it "Debug -> 0 (default positional)" $ + TC.toThrift Debug `shouldBe` TV.I32 0 + , it "Info -> 1" $ + TC.toThrift Info `shouldBe` TV.I32 1 + , it "Critical -> 99 (explicit tag)" $ + TC.toThrift Critical `shouldBe` TV.I32 99 + , it "round-trip" $ mapM_ rt [Debug, Info, Warn, Critical] + , it "unknown value fails" $ + case TC.fromThrift (TV.I32 17) :: Either String Severity of + Left _ -> pure () + Right s -> expectationFailure ("unexpected " ++ show s) + ] + where + rt :: Severity -> IO () + rt s = TC.fromThrift (TC.toThrift s) `shouldBe` Right s + + +-- --------------------------------------------------------------------------- + +unionTests :: Spec +unionTests = + describe "sum (Thrift union)" $ + sequence_ + [ it "EvHeartbeat -> field id 1, payload Bool True" $ + TC.toThrift EvHeartbeat + `shouldBe` TV.Struct (V.singleton (1, TV.Bool True)) + , it "EvData 5 -> field id 2, payload I64 5" $ + TC.toThrift (EvData 5) + `shouldBe` TV.Struct (V.singleton (2, TV.I64 5)) + , it "EvAlert -> tag-overridden field id 10" $ + case TC.toThrift (EvAlert "fire" 7) of + TV.Struct kvs -> + (V.any (\(k, _) -> k == 10) kvs) `shouldBe` True + v -> expectationFailure ("expected Struct, got " ++ show v) + , it "round-trip EvHeartbeat" $ rt EvHeartbeat + , it "round-trip EvData" $ rt (EvData 99) + , it "round-trip EvAlert" $ rt (EvAlert "x" 1) + , it "unknown field id fails" $ do + let bad = TV.Struct (V.singleton (255, TV.Bool True)) + case TC.fromThrift bad :: Either String Event of + Left _ -> pure () + Right e -> expectationFailure ("unexpected " ++ show e) + ] + where + rt :: Event -> IO () + rt e = TC.fromThrift (TC.toThrift e) `shouldBe` Right e diff --git a/test/Test/Thrift/Derive/Instances.hs b/test/Test/Thrift/Derive/Instances.hs new file mode 100644 --- /dev/null +++ b/test/Test/Thrift/Derive/Instances.hs @@ -0,0 +1,13 @@ +{-# LANGUAGE TemplateHaskell #-} +{-# OPTIONS_GHC -Wno-orphans #-} + +module Test.Thrift.Derive.Instances () where + +import Test.Thrift.Derive.Types +import Thrift.Derive + + +deriveThrift ''LogEntry +deriveThrift ''RequestId +deriveThrift ''Severity +deriveThrift ''Event diff --git a/test/Test/Thrift/Derive/Types.hs b/test/Test/Thrift/Derive/Types.hs new file mode 100644 --- /dev/null +++ b/test/Test/Thrift/Derive/Types.hs @@ -0,0 +1,85 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TemplateHaskell #-} + +{- | Annotated fixture types for the Thrift deriver round-trip tests. + +Thrift differs from CBOR / MsgPack: fields are identified by +numeric IDs (Int16) rather than text-string keys, so 'rename' and +'renameStyle' are ignored on Thrift wire keys. The 'tag' modifier +overrides the otherwise-positional default. +-} +module Test.Thrift.Derive.Types ( + -- * Record (mixed default + tag-overridden field IDs) + LogEntry (..), + defaultRequestId, + + -- * Newtype + RequestId (..), + + -- * Enum (with explicit tags) + Severity (..), + + -- * Sum (Thrift union, sequential field IDs) + Event (..), +) where + +import Data.Text (Text) +import Wireform.Derive.Backend +import Wireform.Derive.Modifier + + +-- --------------------------------------------------------------------------- +-- Record +-- --------------------------------------------------------------------------- + +data LogEntry = LogEntry + { logTimestamp :: !Int -- field id 1 by default + , logMessage :: !Text -- field id 2 by default + , logCode :: !Int -- explicit tag 7 + , logRequestId :: !Text -- skipped under Thrift, default value + } + deriving (Eq, Show) + + +defaultRequestId :: Text +defaultRequestId = "" + + +{-# ANN logCode (forBackend backendThrift (tag 7)) #-} + + +{-# ANN logRequestId (forBackend backendThrift skip) #-} +{-# ANN logRequestId (forBackend backendThrift (defaults 'defaultRequestId)) #-} + + +-- --------------------------------------------------------------------------- +-- Newtype +-- --------------------------------------------------------------------------- + +newtype RequestId = RequestId {unRequestId :: Int} + deriving (Eq, Show) + + +-- --------------------------------------------------------------------------- +-- Enum +-- --------------------------------------------------------------------------- + +data Severity = Debug | Info | Warn | Critical + deriving (Eq, Show) + + +{-# ANN Critical (forBackend backendThrift (tag 99)) #-} + + +-- --------------------------------------------------------------------------- +-- Sum (encoded as a Thrift union) +-- --------------------------------------------------------------------------- + +data Event + = EvHeartbeat + | EvData !Int + | EvAlert !Text !Int + deriving (Eq, Show) + + +{-# ANN EvAlert (forBackend backendThrift (tag 10)) #-} diff --git a/test/fixtures/protocol-encodings.tsv b/test/fixtures/protocol-encodings.tsv new file mode 100644 --- /dev/null +++ b/test/fixtures/protocol-encodings.tsv @@ -0,0 +1,10 @@ +# Apache Python thrift 0.25.0; regenerate with scripts/protocol_golden.py --write +# fixture binary hex compact hex +scalars 0200010002000201030003800600048000080005800000000a00068000000000000000040007c05ee000000000000b000800000006c3a9e99baa000b00090000000400ff800110000a000102030405060708090a0b0c0d0e0f00 1211138014ffff0315ffffffff0f16ffffffffffffffffff01170000000000e05ec01806c3a9e99baa00180400ff80011d000102030405060708090a0b0c0d0e0f00 +field-ids 0880000000000202000000087fff0000000100 05ffff0304020005feff030200 +container-booleans 0f000102000000030001000e0002020000000200010d00030202000000020001010000 19310201021a2102011b02110201010200 +list-0 0f0001080000000000 190500 +list-14 0f0001080000000effffffc0ffffffc1ffffffc2ffffffc3ffffffc4ffffffc5ffffffc6ffffffc7ffffffc8ffffffc9ffffffcaffffffcbffffffccffffffcd00 19e57f7d7b79777573716f6d6b69676500 +list-15 0f0001080000000fffffffc0ffffffc1ffffffc2ffffffc3ffffffc4ffffffc5ffffffc6ffffffc7ffffffc8ffffffc9ffffffcaffffffcbffffffccffffffcdffffffce00 19f50f7f7d7b79777573716f6d6b6967656300 +list-128 0f00010800000080ffffffc0ffffffc1ffffffc2ffffffc3ffffffc4ffffffc5ffffffc6ffffffc7ffffffc8ffffffc9ffffffcaffffffcbffffffccffffffcdffffffceffffffcfffffffd0ffffffd1ffffffd2ffffffd3ffffffd4ffffffd5ffffffd6ffffffd7ffffffd8ffffffd9ffffffdaffffffdbffffffdcffffffddffffffdeffffffdfffffffe0ffffffe1ffffffe2ffffffe3ffffffe4ffffffe5ffffffe6ffffffe7ffffffe8ffffffe9ffffffeaffffffebffffffecffffffedffffffeeffffffeffffffff0fffffff1fffffff2fffffff3fffffff4fffffff5fffffff6fffffff7fffffff8fffffff9fffffffafffffffbfffffffcfffffffdfffffffeffffffff000000000000000100000002000000030000000400000005000000060000000700000008000000090000000a0000000b0000000c0000000d0000000e0000000f000000100000001100000012000000130000001400000015000000160000001700000018000000190000001a0000001b0000001c0000001d0000001e0000001f000000200000002100000022000000230000002400000025000000260000002700000028000000290000002a0000002b0000002c0000002d0000002e0000002f000000300000003100000032000000330000003400000035000000360000003700000038000000390000003a0000003b0000003c0000003d0000003e0000003f00 19f580017f7d7b79777573716f6d6b69676563615f5d5b59575553514f4d4b49474543413f3d3b39373533312f2d2b29272523211f1d1b19171513110f0d0b090705030100020406080a0c0e10121416181a1c1e20222426282a2c2e30323436383a3c3e40424446484a4c4e50525456585a5c5e60626466686a6c6e70727476787a7c7e00 +nested 0f00010c000000020b000100000005666972737402000200000b000100000003e99baa02000201000d00020b0f0000000200000001610200000002000100000002c3a902000000000d000300000000000000 192c1805666972737412001803e99baa11001b0289016121020102c3a9011b0000 diff --git a/wireform-thrift.cabal b/wireform-thrift.cabal new file mode 100644 --- /dev/null +++ b/wireform-thrift.cabal @@ -0,0 +1,117 @@ +cabal-version: 3.0 +name: wireform-thrift +version: 0.1.0.0 +synopsis: Apache Thrift binary protocol, schema, and JSON mapping +description: + Binary and Compact protocol codecs, dynamic values, IDL parsing and code + generation, schema-driven Haskell derivation, and JSON mapping. +category: Data, Codec +license: BSD-3-Clause +license-file: LICENSE +author: Ian Duncan +maintainer: ian@iankduncan.com +copyright: 2026 Ian Duncan +build-type: Simple +extra-source-files: + scripts/protocol_golden.py + test/fixtures/protocol-encodings.tsv + +common defaults + ghc-options: + -Wall + -Wcompat + -Widentities + -Wincomplete-record-updates + -Wincomplete-uni-patterns + -Wpartial-fields + -Wredundant-constraints + -O2 + default-language: GHC2021 + default-extensions: + StrictData + DerivingStrategies + DeriveGeneric + DeriveAnyClass + OverloadedStrings + LambdaCase + +library + import: defaults + hs-source-dirs: src + exposed-modules: + Thrift + Thrift.Class + Thrift.CodeGen + Thrift.Decode + Thrift.Derive + Thrift.Encode + Thrift.Encoding + Thrift.JSON + Thrift.Message + Thrift.Parser + Thrift.QQ + Thrift.Registry + Thrift.Schema + Thrift.Transport + Thrift.Value + Thrift.Wire + build-depends: + base >= 4.20 && < 5, + wireform-core >= 0.3 && < 0.4, + wireform-derive == 0.2.*, + bytestring >= 0.11 && < 0.13, + base16-bytestring >= 1.0 && < 1.1, + base64-bytestring >= 1.2 && < 1.3, + containers >= 0.6 && < 0.9, + text >= 2.0 && < 2.2, + vector >= 0.13 && < 0.14, + megaparsec >= 9.0 && < 10, + template-haskell >= 2.18 && < 2.25, + prettyprinter >= 1.7 && < 1.8, + hashable >= 1.3 && < 1.6, + unordered-containers >= 0.2.16 && < 0.3, + deepseq >= 1.4 && < 1.6, + aeson >= 2.1 && < 2.4, + scientific >= 0.3 && < 0.4, + stm >= 2.5 && < 2.6 + +benchmark wireform-thrift-bench + import: defaults + type: exitcode-stdio-1.0 + hs-source-dirs: bench + main-is: Bench.hs + ghc-options: -O2 + build-depends: + base, + wireform-thrift, + bytestring, + deepseq, + text, + vector, + criterion >= 1.5 && < 1.7 + +test-suite wireform-thrift-derive-test + import: defaults + ghc-options: -rtsopts + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Main.hs + other-modules: + Test.Thrift.CompactBool + Test.Thrift.Derive + Test.Thrift.Derive.Types + Test.Thrift.Derive.Instances + default-extensions: + ScopedTypeVariables + TemplateHaskell + build-depends: + base, + wireform-thrift, + wireform-derive, + wireform-core, + bytestring, + hedgehog, + sydtest-hedgehog, + text, + vector, + sydtest