BenchmarkEncoding.hs

{-# 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

Bytes.hs

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)

encoding.proto

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;
}

varint.c

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;
}