diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +1.4.0 + - Improve decoding performance + - Remove internal toMap function + 1.3.0 - Support GHC 9.2 - Prevent inlining for GHCJS diff --git a/bench/Main.hs b/bench/Main.hs new file mode 100644 --- /dev/null +++ b/bench/Main.hs @@ -0,0 +1,175 @@ +{-# LANGUAGE ApplicativeDo #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + + +module Main where + +import qualified Data.ByteString as B +import qualified Data.ByteString.Lazy as BL +import qualified Proto3.Wire.Decode as De +import qualified Proto3.Wire.Encode as En +import Proto3.Wire + +import Control.Applicative (liftA2, liftA3) +import Control.Monad (forM) +import Data.Maybe +import Data.Word +import Data.IORef + +import Criterion (bench) +import qualified Criterion as C +import Criterion.Main (defaultMain) + +data Tree a = Leaf | Branch a (Tree a) (Tree a) + deriving (Eq, Functor) + +data Rose a = Bud | Rose a [Rose a] + +instance Foldable Tree where + foldr _ z Leaf = z + foldr f z (Branch a t1 t2) = foldr f (f a (foldr f z t2)) t1 + + sum Leaf = 0 + sum (Branch a t1 t2) = + let !a1 = sum t1 + !a2 = sum t2 + in a + a1 + a2 + +instance Foldable Rose where + foldMap f Bud = mempty + foldMap f (Rose x rs) = f x <> ((foldMap.foldMap) f rs) + +intTreeParser :: De.Parser De.RawMessage (Tree Word64) +intTreeParser = liftA3 combine + (De.at (De.repeated De.fixed64) (FieldNumber 0)) + (De.at (De.one (De.embedded' intTreeParser) Leaf) (FieldNumber 1)) + (De.at (De.one (De.embedded' intTreeParser) Leaf) (FieldNumber 2)) + where + combine xs y z = Branch (sum xs) y z + +intRoseParser :: De.Parser De.RawMessage (Rose Word64) +intRoseParser = liftA2 (Rose @Word64) + (De.at (De.one De.fixed64 0) (FieldNumber 0)) + (De.at (De.repeated (De.embedded' intRoseParser)) (FieldNumber 1)) + +detRandom :: [Word64] +detRandom = concat . replicate 10 $ + [ 227, 133, 16, 164, 43, + 159, 207, 87, 180, 236, + 245, 128, 249, 170, 216, + 181, 164, 162, 239, 249, + 76, 237, 197, 246, 209, + 231, 124, 154, 55, 64, + 4, 114, 79, 199, 252, + 163, 116, 237, 209, 138, + 240, 148, 212, 224, 88, + 131, 122, 114, 158, 97, + 186, 3, 223, 230, 223, + 207, 93, 168, 48, 130, + 77, 122, 30, 222, 221, + 224, 243, 19, 175, 61, + 112, 246, 201, 57, 185, + 19, 128, 129, 138, 209, + 4, 153, 196, 238, 72, + 254, 157, 233, 81, 30, + 106, 249, 57, 214, 104, + 171, 146, 175, 185, 192, + 159, 207, 87, 180, 236, + 227, 133, 16, 164, 43, + 245, 128, 249, 170, 216, + 181, 164, 162, 239, 249, + 76, 237, 197, 246, 209, + 231, 124, 154, 55, 64, + 4, 114, 79, 199, 252, + 163, 116, 237, 209, 138, + 240, 148, 212, 224, 88, + 131, 122, 114, 158, 97, + 186, 3, 223, 230, 223, + 207, 93, 168, 48, 130, + 77, 122, 30, 222, 221, + 224, 243, 19, 175, 61, + 112, 246, 201, 57, 185, + 19, 128, 129, 138, 209, + 4, 153, 196, 238, 72, + 254, 157, 233, 81, 30, + 106, 249, 57, 214, 104, + 171, 146, 175, 185, 192, + 159, 207, 87, 180, 236, + 227, 133, 16, 164, 43, + 245, 128, 249, 170, 216, + 181, 164, 162, 239, 249, + 76, 237, 197, 246, 209, + 231, 124, 154, 55, 64, + 4, 114, 79, 199, 252, + 163, 116, 237, 209, 138, + 240, 148, 212, 224, 88, + 131, 122, 114, 158, 97, + 186, 3, 223, 230, 223, + 207, 93, 168, 48, 130, + 77, 122, 30, 222, 221, + 224, 243, 19, 175, 61, + 112, 246, 201, 57, 185, + 19, 128, 129, 138, 209, + 4, 153, 196, 238, 72, + 254, 157, 233, 81, 30, + 106, 249, 57, 214, 104, + 171, 146, 175, 185, 192 + ] + +pullInt :: IORef [Word64] -> IO Word64 +pullInt xs = do + xs' <- readIORef xs + case xs' of + [] -> pure (-1) + x : xs' -> do + writeIORef xs xs' + pure x + +mkTree0 :: IO Word64 -> IO En.MessageBuilder +mkTree0 ints = do + shouldFork <- (\(i :: Word64) -> (i `mod` 8) < 6) <$> ints + if shouldFork + then do + i <- En.fixed64 (FieldNumber 0) <$> ints + left <- En.embedded (FieldNumber 1) <$> mkTree0 ints + right <- En.embedded (FieldNumber 2) <$> mkTree0 ints + pure (i <> left <> right) + else pure mempty + +mkRose0 :: IO Word64 -> IO En.MessageBuilder +mkRose0 ints = do + next <- fromIntegral <$> ints + if next == -1 then pure mempty else do + let nBranches = next `mod` 9 + if nBranches == 0 then pure mempty else do + loc <- (\i -> (i `mod` nBranches)) . fromIntegral <$> ints + i <- En.fixed64 (FieldNumber 0) <$> ints + rs1 <- forM (replicate loc ()) $ \() -> + En.embedded (FieldNumber 1) <$> mkTree0 ints + rs2 <- forM (replicate (nBranches - loc) ()) $ \() -> + En.embedded (FieldNumber 1) <$> mkTree0 ints + pure (mconcat rs1 <> i <> mconcat rs2) + +mkTree :: IO B.ByteString +mkTree = BL.toStrict . En.toLazyByteString <$> (mkTree0 . pullInt =<< newIORef detRandom) + +mkRose :: IO B.ByteString +mkRose = BL.toStrict . En.toLazyByteString <$> (mkRose0 . pullInt =<< newIORef detRandom) + +decode :: Foldable f => De.Parser De.RawMessage (f Word64) -> B.ByteString -> IO (Maybe Word64) +decode p = pure . fmap sum . toMaybe . De.parse p + where + toMaybe (Left _) = Nothing + toMaybe (Right x) = Just x + +unwrap :: (Functor m, Foldable f) => m (f a) -> m a +unwrap = fmap (foldr1 const) + +main :: IO () +main = + defaultMain + [ bench "Parse int tree" $ C.perRunEnv mkTree (unwrap . decode intTreeParser) + , bench "Parse int rose tree" $ C.perRunEnv mkRose (unwrap . decode intRoseParser)] diff --git a/proto3-wire.cabal b/proto3-wire.cabal --- a/proto3-wire.cabal +++ b/proto3-wire.cabal @@ -1,5 +1,5 @@ name: proto3-wire -version: 1.3.0 +version: 1.4.0 synopsis: A low-level implementation of the Protocol Buffers (version 3) wire format license: Apache-2.0 license-file: LICENSE @@ -65,3 +65,11 @@ text >= 0.2 && <1.3, transformers >=0.5.6.2 && <0.6, vector >=0.12.0.2 && <0.13 + +benchmark bench + type: exitcode-stdio-1.0 + main-is: Main.hs + build-depends: base >= 4 && < 5, bytestring, random, criterion, proto3-wire + hs-source-dirs: bench + ghc-options: -O2 -Wall + default-language: Haskell2010 diff --git a/src/Proto3/Wire/Decode.hs b/src/Proto3/Wire/Decode.hs --- a/src/Proto3/Wire/Decode.hs +++ b/src/Proto3/Wire/Decode.hs @@ -74,8 +74,6 @@ , embedded' -- * ZigZag codec , zigZagDecode - -- * Exported For Doctest Only - , toMap ) where import Control.Applicative @@ -123,29 +121,48 @@ -- | Convert key-value pairs to a map of keys to a sequence of values with that -- key, in their reverse occurrence order. -- --- >>> toMap ([(FieldNumber 1, 3),(FieldNumber 2, 4),(FieldNumber 1, 6)] :: [(FieldNumber,Int)]) --- fromList [(1,[6,3]),(2,[4])] -- -toMap :: [(FieldNumber, v)] -> M.IntMap [v] -toMap kvs0 = M.fromListWith (<>) . map (fmap (:[])) . map (first (fromIntegral . getFieldNumber)) $ kvs0 +decodeWireMessage :: B.ByteString -> Either String RawMessage +decodeWireMessage = decodeWire0 combineSeen' Nothing close + where + close Nothing = M.empty + close (Just (m, k, v)) = M.insertWith (++) k v m --- | Parses data in the raw wire format into an untyped 'Map' representation. + combineSeen' :: Maybe (M.IntMap [v], Int, [v]) -> FieldNumber -> v -> Maybe (M.IntMap [v], Int, [v]) + combineSeen' b (FieldNumber fn) v = combineSeen b (fromIntegral fn) v + + combineSeen :: Maybe (M.IntMap [v], Int, [v]) -> Int -> v -> Maybe (M.IntMap [v], Int, [v]) + combineSeen Nothing k1 a1 = Just (M.empty, k1, [a1]) + combineSeen (Just (m, k2, as)) k1 a1 = + if k1 == k2 + then Just (m, k1, a1 : as) + -- It might seem that we want to use DList but we don't because: + -- - alter has worse performance than insertWith, and there's no upsert + -- - We're building up a list of elements in a recursive way + -- that will be opaque to GHC + -- - DList would add another dependency + else let !m' = M.insertWith (++) k2 as m + in Just (m', k1, [a1]) + decodeWire :: B.ByteString -> Either String [(FieldNumber, ParsedField)] -decodeWire bstr = drloop bstr [] +decodeWire = decodeWire0 (\xs k v -> (k,v):xs) [] reverse + +-- | Parses data in the raw wire format into an untyped 'Map' representation. +decodeWire0 :: (b -> FieldNumber -> ParsedField -> b) -> b -> (b -> r) -> B.ByteString -> Either String r +decodeWire0 cl z finish bstr = drloop bstr z where - drloop !bs xs | B.null bs = Right $ reverse xs + drloop !bs xs | B.null bs = Right $ finish xs drloop !bs xs | otherwise = do (w, rest) <- takeVarInt bs wt <- gwireType $ fromIntegral (w .&. 7) let fn = w `shiftR` 3 (res, rest2) <- takeWT wt rest - drloop rest2 ((FieldNumber fn,res):xs) - + drloop rest2 (cl xs (FieldNumber fn) res) +{-# INLINE decodeWire0 #-} eitherUncons :: B.ByteString -> Either String (Word8, B.ByteString) eitherUncons = maybe (Left "failed to parse varint128") Right . B.uncons - takeVarInt :: B.ByteString -> Either String (Word64, B.ByteString) takeVarInt !bs = case B.uncons bs of @@ -190,6 +207,7 @@ if w10 < 128 then return (val9 + (fromIntegral w10 `shiftL` 63), r10) else do Left ("failed to parse varint128: too big; " ++ show val6) +{-# INLINE takeVarInt #-} gwireType :: Word8 -> Either String WireType @@ -198,6 +216,7 @@ gwireType 1 = return Fixed64 gwireType 2 = return LengthDelimited gwireType wt = Left $ "wireType got unknown wire type: " ++ show wt +{-# INLINE gwireType #-} safeSplit :: Int -> B.ByteString -> Either String (B.ByteString, B.ByteString) safeSplit !i !b | B.length b < i = Left "failed to parse varint128: not enough bytes" @@ -210,6 +229,7 @@ takeWT LengthDelimited b = do (!len, rest) <- takeVarInt b fmap (first LengthDelimitedField) $ safeSplit (fromIntegral len) rest +{-# INLINE takeWT #-} -- * Parser Interface @@ -289,9 +309,10 @@ -- | Parse a message (encoded in the raw wire format) using the specified -- `Parser`. parse :: Parser RawMessage a -> B.ByteString -> Either ParseError a -parse parser bs = case decodeWire bs of +parse parser bs = case decodeWireMessage bs of Left err -> Left (BinaryError (pack err)) - Right res -> runParser parser (toMap res) + Right res -> runParser parser res +{-# INLINE parse #-} -- | To comply with the protobuf spec, if there are multiple fields with the same -- field number, this will always return the last one. @@ -300,22 +321,31 @@ [] -> Nothing (x:_) -> Just x -throwWireTypeError :: Show input - => String - -> input - -> Either ParseError expected -throwWireTypeError expected wrong = - Left (WireTypeError (pack msg)) +wireTypeError :: String + -> RawPrimitive + -> ParseError +wireTypeError expected wrong = WireTypeError (pack msg) where msg = "Wrong wiretype. Expected " ++ expected ++ " but got " ++ show wrong -throwCerealError :: String -> String -> Either ParseError a -throwCerealError expected cerealErr = - Left (BinaryError (pack msg)) +throwWireTypeError :: String + -> RawPrimitive + -> Either ParseError expected +throwWireTypeError expected wrong = + Left (wireTypeError expected wrong) +{-# INLINE throwWireTypeError #-} + +cerealTypeError :: String -> String -> ParseError +cerealTypeError expected cerealErr = (BinaryError (pack msg)) where msg = "Failed to parse contents of " ++ expected ++ " field. " ++ "Error from cereal was: " ++ cerealErr +throwCerealError :: String -> String -> Either ParseError a +throwCerealError expected cerealErr = + Left (cerealTypeError expected cerealErr) +{-# INLINE throwCerealError #-} + parseVarInt :: Integral a => Parser RawPrimitive a parseVarInt = Parser $ \case @@ -486,6 +516,7 @@ -- > one float `at` fieldNumber 1 :: Parser RawMessage (Maybe Float) at :: Parser RawField a -> FieldNumber -> Parser RawMessage a at parser fn = Parser $ runParser parser . fromMaybe mempty . M.lookup (fromIntegral . getFieldNumber $ fn) +{-# INLINE at #-} -- | Try to parse different field numbers with their respective parsers. This is -- used to express alternative between possible fields of a oneof. @@ -538,15 +569,20 @@ repeated :: Parser RawPrimitive a -> Parser RawField [a] repeated parser = Parser $ fmap reverse . mapM (runParser parser) + +embeddedParseError :: String -> ParseError +embeddedParseError err = EmbeddedError msg Nothing + where + msg = "Failed to parse embedded message: " <> (pack err) +{-# NOINLINE embeddedParseError #-} + -- | For a field containing an embedded message, parse as far as getting the -- wire-level fields out of the message. embeddedToParsedFields :: RawPrimitive -> Either ParseError RawMessage embeddedToParsedFields (LengthDelimitedField bs) = - case decodeWire bs of - Left err -> Left (EmbeddedError ("Failed to parse embedded message: " - <> (pack err)) - Nothing) - Right result -> return (toMap result) + case decodeWireMessage bs of + Left err -> Left (embeddedParseError err) + Right result -> Right result embeddedToParsedFields wrong = throwWireTypeError "embedded" wrong @@ -568,6 +604,7 @@ let combinedMap = foldl' (M.unionWith (<>)) M.empty innerMaps parsed <- runParser p combinedMap return $ Just parsed +{-# INLINE embedded #-} -- | Create a primitive parser for an embedded message from a message parser. -- @@ -578,10 +615,11 @@ \case LengthDelimitedField bs -> case parse parser bs of - Left err -> Left (EmbeddedError "Failed to parse embedded message." - (Just err)) + Left err -> Left (EmbeddedError "Failed to parse embedded message." (Just err)) Right result -> return result wrong -> throwWireTypeError "embedded" wrong +{-# INLINE embedded' #-} + -- TODO test repeated and embedded better for reverse logic... diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -16,6 +16,7 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE LambdaCase #-} +{-# LANGUAGE ScopedTypeVariables #-} {-# OPTIONS_GHC -Wno-warnings-deprecations #-} @@ -70,6 +71,7 @@ , decodeNonsense , varIntHeavyTests , packedLargeTests + , decodeWireRoundTrip ] data StringOrInt64 = TString T.Text | TInt64 Int64 @@ -215,7 +217,6 @@ , (fieldNumber 3, Just . TString <$> one Decode.text mempty) ] ) - ] roundTrip :: (Show a, Eq a, Arbitrary a) @@ -230,6 +231,16 @@ case Decode.parse decode (BL.toStrict bytes) of Left _ -> error "Could not decode encoded message" Right x' -> x === x' + + +decodeWireRoundTrip :: TestTree +decodeWireRoundTrip = QC.testProperty "decodeWire round trips" $ + \(inp :: [(FieldNumber, Int32)]) -> + let bytes = Encode.toLazyByteString (foldMap (\(k, v) -> Encode.int32 k v) inp) + x = map (second $ Decode.VarintField . fromIntegral) inp + in case Decode.decodeWire (BL.toStrict bytes) of + Left _ -> error "decodeWire failed" + Right x' -> x === x' buildSingleChunk :: TestTree buildSingleChunk = HU.testCase "Legacy Builder creates a single chunk" $ do