wireform-thrift (empty) → 0.1.0.0
raw patch · 26 files changed
+6283/−0 lines, 26 filesdep +aesondep +basedep +base16-bytestring
Dependencies added: aeson, base, base16-bytestring, base64-bytestring, bytestring, containers, criterion, deepseq, hashable, hedgehog, megaparsec, prettyprinter, scientific, stm, sydtest, sydtest-hedgehog, template-haskell, text, unordered-containers, vector, wireform-core, wireform-derive, wireform-thrift
Files
- LICENSE +26/−0
- bench/Bench.hs +83/−0
- scripts/protocol_golden.py +129/−0
- src/Thrift.hs +43/−0
- src/Thrift/Class.hs +835/−0
- src/Thrift/CodeGen.hs +534/−0
- src/Thrift/Decode.hs +366/−0
- src/Thrift/Derive.hs +462/−0
- src/Thrift/Encode.hs +489/−0
- src/Thrift/Encoding.hs +42/−0
- src/Thrift/JSON.hs +536/−0
- src/Thrift/Message.hs +266/−0
- src/Thrift/Parser.hs +438/−0
- src/Thrift/QQ.hs +39/−0
- src/Thrift/Registry.hs +78/−0
- src/Thrift/Schema.hs +200/−0
- src/Thrift/Transport.hs +75/−0
- src/Thrift/Value.hs +58/−0
- src/Thrift/Wire.hs +882/−0
- test/Main.hs +16/−0
- test/Test/Thrift/CompactBool.hs +343/−0
- test/Test/Thrift/Derive.hs +118/−0
- test/Test/Thrift/Derive/Instances.hs +13/−0
- test/Test/Thrift/Derive/Types.hs +85/−0
- test/fixtures/protocol-encodings.tsv +10/−0
- wireform-thrift.cabal +117/−0
+ LICENSE view
@@ -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.
+ bench/Bench.hs view
@@ -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+ ]+ ]+ ]
+ scripts/protocol_golden.py view
@@ -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()
+ src/Thrift.hs view
@@ -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
+ src/Thrift/Class.hs view
@@ -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)
+ src/Thrift/CodeGen.hs view
@@ -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
+ src/Thrift/Decode.hs view
@@ -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 #-}
+ src/Thrift/Derive.hs view
@@ -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
+ src/Thrift/Encode.hs view
@@ -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)
+ src/Thrift/Encoding.hs view
@@ -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
+ src/Thrift/JSON.hs view
@@ -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)+ _ -> "<complex>"+++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
+ src/Thrift/Message.hs view
@@ -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 #-}
+ src/Thrift/Parser.hs view
@@ -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) "<thrift>" 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)
+ src/Thrift/QQ.hs view
@@ -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
+ src/Thrift/Registry.hs view
@@ -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)}
+ src/Thrift/Schema.hs view
@@ -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)
+ src/Thrift/Transport.hs view
@@ -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 #-}
+ src/Thrift/Value.hs view
@@ -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
+ src/Thrift/Wire.hs view
@@ -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 #-}
+ test/Main.hs view
@@ -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+ ]+ )
+ test/Test/Thrift/CompactBool.hs view
@@ -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<bool> 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
+ test/Test/Thrift/Derive.hs view
@@ -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
+ test/Test/Thrift/Derive/Instances.hs view
@@ -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
+ test/Test/Thrift/Derive/Types.hs view
@@ -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 = "<unknown>"+++{-# 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)) #-}
+ test/fixtures/protocol-encodings.tsv view
@@ -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
+ wireform-thrift.cabal view
@@ -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