Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
{-# LANGUAGE OverloadedLabels #-}
{-# LANGUAGE OverloadedStrings #-}
module Main (main) where
import Criterion (bench, bgroup, whnf)
import Criterion.Main (Benchmark)
import qualified Data.ByteString as BS
import Data.ByteString.Builder (toLazyByteString)
import Data.ByteString.Lazy (toStrict)
import Data.ProtoLens (encodeMessage)
import Data.ProtoLens.BenchmarkUtil (benchmarkMain)
import Data.ProtoLens.Encoding.Bytes
import Data.ProtoLens.Message (defMessage)
import Data.Word (Word64)
import Lens.Family ((&), (.~))
import Proto.Encoding
import Proto.Encoding_Fields
numbers :: [Word64]
numbers = [0, 19, 300, 70000, 1566433440, 9223372036854, 922337203685477, 9223372036854775803]
benchmaker :: Int -> [Benchmark]
benchmaker _ = [bgroup "putVarInt" [
bgroup "one" $ one putVarInt <$> numbers,
bgroup "five" $ nInts putVarInt 5 <$> numbers
],
bgroup "proto" [
bgroup "one msg" $ oneMsg . fromIntegral <$> numbers,
bgroup "five msg" $ fiveMsg . fromIntegral <$> numbers
]
]
where
one f n = bench name $ whnf build1 n
where name = show n <> " (" <> (show . BS.length . build1) n <> "B)"
build1 :: Word64 -> BS.ByteString
build1 = toStrict . toLazyByteString . f
nInts f n x = bench name $ whnf buildN x
where name = show x <> " (" <> (show . BS.length . buildN) x <> "B)"
buildN :: Word64 -> BS.ByteString
buildN = toStrict . toLazyByteString . foldMap f . replicate n
oneMsg n = bench name $ whnf encodeMessage proto
where
name = show n <> " (" <> (show . BS.length . encodeMessage) proto <> "B)"
proto :: OneInt64
proto = defMessage & oneInt64 .~ n
fiveMsg n = bench name $ whnf encodeMessage proto
where
name = show n <> " (" <> (show . BS.length . encodeMessage) proto <> "B)"
proto :: FiveInt64s
proto = defMessage & firstInt64 .~ n
& secondInt64 .~ n
& thirdInt64 .~ n
& fourthInt64 .~ n
& fifthInt64 .~ n
main :: IO ()
main = benchmarkMain 0 benchmaker
foreign import ccall unsafe "static _hs_protobuf_put_varint" c_varint
:: CULLong -> Ptr Word8 -> IO (Ptr Word8)
-- I was swapping out the 'otherwise' to test different variants.
putVarInt :: Word64 -> Builder
putVarInt n
| n < 128 = Builder.word8 (fromIntegral n)
| otherwise = Prim.primBounded cvar64 (coerce n)
| otherwise = Builder.word8 (fromIntegral $ n .&. 127 .|. 128)
<> putVarInt (n `shiftR` 7)
| otherwise = Prim.primBounded var64 n
where
cvar64 :: PrimI.BoundedPrim CULLong
cvar64 = PrimI.boudedPrim 10 c_varint
var64 :: PrimI.BoundedPrim Word64
var64 = PrimI.boudedPrim 10 enc
where
enc n' ptr
| n' < 128 = poke1 ptr n'
| otherwise = poke1 ptr (n' .&. 127 .|. 128) >>= enc (n' `shiftR` 7)
poke1 ptr v = poke ptr (fromIntegral v :: Word8) >> pure (ptr `plusPtr` 1)
syntax = "proto3";
package encoding;
message OneInt64 {
int64 one_int64 = 1;
}
message FiveInt64s {
int64 first_int64 = 1;
int64 second_int64 = 2;
int64 third_int64 = 3;
int64 fourth_int64 = 4;
int64 fifth_int64 = 5;
}
char* _hs_protobuf_put_varint(unsigned long long value, char* buf) {
if (value < 0x80) {
buf[0] = value & 0xff;
return buf + 1;
}
buf[0] = (value | 0x80) & 0xff;
value >>= 7;
if (value < 0x80) {
buf[1] = value & 0xff;
return buf + 2;
}
buf++;
do {
*buf = (value | 0x80) & 0xff;
value >>= 7;
++buf;
} while (value >= 0x80);
*buf++ = value & 0xff;
return buf;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment