Repository: tokio-rs/prost (language: rust)

Issue:
varint encoding/decoding/len functions are not optimized
Protobuf performance tends to be dominated by varint operations.  The [varint functions](https://github.com/danburkert/prost/blob/v0.1.0/src/encoding.rs#L38-L98) in `prost` have not been optimized at all. There are some benchmarks, but I think they've bit-rotted at this point.


Discussion / hints:

Fix patch:
diff --git a/.travis.yml b/.travis.yml
index 4c754928e..3d84e2d03 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -23,3 +23,4 @@ rust:
 script:
   - cargo build --verbose --all
   - cargo test --verbose --all
+  - cargo bench --benches --verbose --all
diff --git a/benches/varint.rs b/benches/varint.rs
index a3219fda6..bf1c372fc 100644
--- a/benches/varint.rs
+++ b/benches/varint.rs
@@ -1,45 +1,41 @@
 #![feature(test)]
 
-extern crate proto;
+extern crate prost;
 extern crate test;
 extern crate bytes;
 
 use test::Bencher;
 
-use bytes::{
-    BytesMut,
-    IntoBuf,
-};
+use bytes::IntoBuf;
 
-use proto::field::{
-    encode_varint,
+use prost::encoding::{
     decode_varint,
-    varint_width
+    encode_varint,
+    encoded_len_varint,
 };
 
 macro_rules! varint_bench {
-    ($encode_name:ident, $decode_name:ident, $expected_bytes:expr, $encode:expr) => {
+    ($encode_name:ident, $decode_name:ident, $encoded_len_name: ident, $encode:expr) => {
         #[bench]
         fn $encode_name(b: &mut Bencher) {
-            let mut buf = BytesMut::with_capacity(4096);
+            let mut buf = Vec::with_capacity(100 * 10);
             b.iter(|| {
                 buf.clear();
                 $encode(&mut buf);
                 test::black_box(&buf[..]);
             });
-            assert_eq!(buf.len(), $expected_bytes);
             b.bytes = 100 * 8;
         }
         #[bench]
         fn $decode_name(b: &mut Bencher) {
-            let mut buf = BytesMut::with_capacity(4096);
+            let mut buf = Vec::with_capacity(100 * 10);
             $encode(&mut buf);
-            let buf = buf.freeze();
+            let buf = &buf[..];
 
             let mut values = [0u64; 100];
 
             b.iter(|| {
-                let mut buf = buf.clone().into_buf();
+                let mut buf = buf.into_buf();
                 for i in 0..100 {
                     values[i] = decode_varint(&mut buf).unwrap();
                 }
@@ -47,11 +43,32 @@ macro_rules! varint_bench {
             });
             b.bytes = 100 * 8;
         }
+        #[bench]
+        fn $encoded_len_name(b: &mut Bencher) {
+            let mut values = [0u64; 100];
+            {
+                let mut buf = Vec::with_capacity(100 * 10);
+                $encode(&mut buf);
+                let mut buf = (&buf[..]).into_buf();
+                for i in 0..100 {
+                    values[i] = decode_varint(&mut buf).unwrap();
+                }
+            }
+
+            b.iter(|| {
+                let mut sum = 0;
+                for &value in values.iter() {
+                    sum += encoded_len_varint(value);
+                }
+                test::black_box(sum);
+            });
+            b.bytes = 100 * 8;
+        }
     }
 }
 
 /// Benchmark encoding and decoding 100 varints of mixed width (average 5.5 bytes).
-varint_bench!(encode_varint_mixed, decode_varint_mixed, 550, |ref mut buf| {
+varint_bench!(encode_varint_mixed, decode_varint_mixed, encoded_len_varint_mixed, |ref mut buf| {
     for width in 0..10 {
         let exponent = width * 7;
         for offset in 0..10 {
@@ -61,14 +78,14 @@ varint_bench!(encode_varint_mixed, decode_varint_mixed, 550, |ref mut buf| {
 });
 
 /// Benchmark encoding and decoding 100 small (1 byte) varints.
-varint_bench!(encode_varint_small, decode_varint_small, 100, |ref mut buf| {
+varint_bench!(encode_varint_small, decode_varint_small, encoded_len_varint_small, |ref mut buf| {
     for value in 0..100 {
         encode_varint(value, buf);
     }
 });
 
 /// Benchmark encoding and decoding 100 medium (5 byte) varints.
-varint_bench!(encode_varint_medium, decode_varint_medium, 500, |ref mut buf| {
+varint_bench!(encode_varint_medium, decode_varint_medium, encoded_len_varint_medium, |ref mut buf| {
     let start = 1 << 28;
     for value in start..start + 100 {
         encode_varint(value, buf);
@@ -76,7 +93,7 @@ varint_bench!(encode_varint_medium, decode_varint_medium, 500, |ref mut buf| {
 });
 
 /// Benchmark encoding and decoding 100 large (10 byte) varints.
-varint_bench!(encode_varint_large, decode_varint_large, 1000, |ref mut buf| {
+varint_bench!(encode_varint_large, decode_varint_large, encoded_len_varint_large, |ref mut buf| {
     let start = 1 << 63;
     for value in start..start + 100 {
         encode_varint(value, buf);
diff --git a/src/encoding.rs b/src/encoding.rs
index 6cfee8adb..96c49a77d 100644
--- a/src/encoding.rs
+++ b/src/encoding.rs
@@ -20,12 +20,18 @@ use Message;
 /// The buffer must have enough remaining space (maximum 10 bytes).
 #[inline]
 pub fn encode_varint<B>(mut value: u64, buf: &mut B) where B: BufMut {
+    // Safety notes:
+    //
+    // - bytes_mut is unsafe because it may return an uninitialized slice.
+    //   The use here is safe because the slice is only written to, never read from.
+    //
+    // - advance_mut is unsafe because it could cause uninitialized memory to be
+    //   advanced over. The use here is safe since each byte which is advanced over
+    //   has been written to in the previous loop iteration.
     let mut i;
     'outer: loop {
         i = 0;
 
-        // bytes_mut is unsafe because it may return an uninitialized slice.
-        // This use is safe because the slice is only written to, not read from.
         for byte in unsafe { buf.bytes_mut() } {
             i += 1;
             if value < 0x80 {
@@ -38,18 +44,86 @@ pub fn encode_varint<B>(mut value: u64, buf: &mut B) where B: BufMut {
         }
 
         unsafe { buf.advance_mut(i); }
-        assert!(buf.has_remaining_mut());
+        debug_assert!(buf.has_remaining_mut());
     }
 
-    // advance_mut is unsafe because it could cause uninitialized memory to be
-    // advanced over. This use is safe since each byte which is advanced over
-    // has been written to in the previous loop.
     unsafe { buf.advance_mut(i); }
 }
 
 /// Decodes a LEB128-encoded variable length integer from the buffer.
-#[inline]
 pub fn decode_varint<B>(buf: &mut B) -> Result<u64, DecodeError> where B: Buf {
+    // NLL hack.
+    'slow: loop {
+        // Another NLL hack.
+        let (value, advance) = {
+            let bytes = buf.bytes();
+            let len = bytes.len();
+            if len == 0 {
+                return Err(DecodeError::new("invalid varint"));
+            }
+
+            let byte = bytes[0];
+            if byte < 0x80 {
+                (byte as u64, 1)
+            } else if len > 10 || bytes[len - 1] < 0x80 {
+                decode_varint_slice(bytes)?
+            } else {
+                break 'slow;
+            }
+        };
+
+        buf.advance(advance);
+        return Ok(value);
+    }
+    decode_varint_slow(buf)
+}
+
+/// Decodes a LEB128-encoded variable length integer from the slice, returning the value and the
+/// number of bytes read.
+///
+/// Based loosely on [ReadVarint64FromArray][1].
+///
+/// [1]: https://github.com/google/protobuf/blob/3.3.x/src/google/protobuf/io/coded_stream.cc#L365-L406
+#[inline]
+fn decode_varint_slice(bytes: &[u8]) -> Result<(u64, usize), DecodeError> {
+    // Fully unrolled varint decoding loop. Splitting into 32-bit pieces gives better performance.
+
+    let mut b: u8;
+    let mut part0: u32;
+    b = bytes[0]; part0  =  b as u32       ; if b < 0x80 { return Ok((part0 as u64, 1)) };
+    part0 -= 0x80;
+    b = bytes[1]; part0 += (b as u32) <<  7; if b < 0x80 { return Ok((part0 as u64, 2)) };
+    part0 -= 0x80 << 7;
+    b = bytes[2]; part0 += (b as u32) << 14; if b < 0x80 { return Ok((part0 as u64, 3)) };
+    part0 -= 0x80 << 14;
+    b = bytes[3]; part0 += (b as u32) << 21; if b < 0x80 { return Ok((part0 as u64, 4)) };
+    part0 -= 0x80 << 21;
+    let value = part0 as u64;
+
+    let mut part1: u32;
+    b = bytes[4]; part1  =  b as u32       ; if b < 0x80 { return Ok((value + ((part1 as u64) << 28), 5)) };
+    part1 -= 0x80;
+    b = bytes[5]; part1 += (b as u32) <<  7; if b < 0x80 { return Ok((value + ((part1 as u64) << 28), 6)) };
+    part1 -= 0x80 << 7;
+    b = bytes[6]; part1 += (b as u32) << 14; if b < 0x80 { return Ok((value + ((part1 as u64) << 28), 7)) };
+    part1 -= 0x80 << 14;
+    b = bytes[7]; part1 += (b as u32) << 21; if b < 0x80 { return Ok((value + ((part1 as u64) << 28), 8)) };
+    part1 -= 0x80 << 21;
+    let value = value + ((part1 as u64) << 28);
+
+    let mut part2: u32;
+    b = bytes[8]; part2  =  b as u32       ; if b < 0x80 { return Ok((value + ((part2 as u64) << 56), 9)) };
+    part2 -= 0x80;
+    b = bytes[9]; part2 += (b as u32) <<  7; if b < 0x80 { return Ok((value + ((part2 as u64) << 56), 10)) };
+
+    // We have overrun the maximum size of a varint (10 bytes). Assume the data is corrupt.
+    return Err(DecodeError::new("invalid varint"));
+}
+
+/// Decodes a LEB128-encoded variable length integer from the buffer, advancing the buffer as
+/// necessary.
+#[inline(never)]
+fn decode_varint_slow<B>(buf: &mut B) -> Result<u64, DecodeError> where B: Buf {
     let mut value = 0;
     for count in 0..min(10, buf.remaining()) {
         let byte = buf.get_u8();
@@ -66,16 +140,9 @@ pub fn decode_varint<B>(buf: &mut B) -> Result<u64, DecodeError> where B: Buf {
 /// The returned value will be between 1 and 10, inclusive.
 #[inline]
 pub fn encoded_len_varint(value: u64) -> usize {
-         if value < 1 <<  7 { 1 }
-    else if value < 1 << 14 { 2 }
-    else if value < 1 << 21 { 3 }
-    else if value < 1 << 28 { 4 }
-    else if value < 1 << 35 { 5 }
-    else if value < 1 << 42 { 6 }
-    else if value < 1 << 49 { 7 }
-    else if value < 1 << 56 { 8 }
-    else if value < 1 << 63 { 9 }
-    else { 10 }
+    // Based on [VarintSize64][1].
+    // [1]: https://github.com/google/protobuf/blob/3.3.x/src/google/protobuf/io/coded_stream.h#L1301-L1309
+    ((((value | 1).leading_zeros() ^ 63) * 9 + 73) / 64) as usize
 }
 
 #[derive(Clone, Copy, Debug, PartialEq)]
@@ -768,6 +835,7 @@ pub mod btree_map {
 mod test {
     use std::fmt::Debug;
     use std::io::Cursor;
+    use std::u64;
 
     use bytes::{Bytes, BytesMut, IntoBuf};
     use quickcheck::TestResult;
@@ -915,26 +983,57 @@ mod test {
     #[test]
     fn varint() {
         fn check(value: u64, encoded: &[u8]) {
-            let mut buf = Vec::new();
-
+            // Small buffer.
+            let mut buf = Vec::with_capacity(1);
             encode_varint(value, &mut buf);
+            assert_eq!(buf, encoded);
 
+            // Large buffer.
+            let mut buf = Vec::with_capacity(100);
+            encode_varint(value, &mut buf);
             assert_eq!(buf, encoded);
 
-            let roundtrip_value = decode_varint(&mut Bytes::from(encoded).into_buf()).expect("decoding failed");
+            assert_eq!(encoded_len_varint(value), encoded.len());
+
+            let roundtrip_value = decode_varint(&mut encoded.into_buf()).expect("decoding failed");
+            assert_eq!(value, roundtrip_value);
+
+            let roundtrip_value = decode_varint_slow(&mut encoded.into_buf()).expect("slow decoding failed");
             assert_eq!(value, roundtrip_value);
         }
 
-        check(0, &[0b0000_0000]);
-        check(1, &[0b0000_0001]);
+        check(2u64.pow(0) - 1, &[0x00]);
+        check(2u64.pow(0), &[0x01]);
+
+        check(2u64.pow(7) - 1, &[0x7F]);
+        check(2u64.pow(7), &[0x80, 0x01]);
+        check(300, &[0xAC, 0x02]);
+
+        check(2u64.pow(14) - 1, &[0xFF, 0x7F]);
+        check(2u64.pow(14), &[0x80, 0x80, 0x01]);
+
+        check(2u64.pow(21) - 1, &[0xFF, 0xFF, 0x7F]);
+        check(2u64.pow(21), &[0x80, 0x80, 0x80, 0x01]);
+
+        check(2u64.pow(28) - 1, &[0xFF, 0xFF, 0xFF, 0x7F]);
+        check(2u64.pow(28), &[0x80, 0x80, 0x80, 0x80, 0x01]);
+
+        check(2u64.pow(35) - 1, &[0xFF, 0xFF, 0xFF, 0xFF, 0x7F]);
+        check(2u64.pow(35), &[0x80, 0x80, 0x80, 0x80, 0x80, 0x01]);
+
+        check(2u64.pow(42) - 1, &[0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x7F]);
+        check(2u64.pow(42), &[0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x01]);
+
+        check(2u64.pow(49) - 1, &[0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x7F]);
+        check(2u64.pow(49), &[0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x01]);
 
-        check(127, &[0b0111_1111]);
-        check(128, &[0b1000_0000, 0b0000_0001]);
+        check(2u64.pow(56) - 1, &[0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x7F]);
+        check(2u64.pow(56), &[0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x01]);
 
-        check(300, &[0b1010_1100, 0b0000_0010]);
+        check(2u64.pow(63) - 1, &[0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x7F]);
+        check(2u64.pow(63), &[0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x01]);
 
-        check(16_383, &[0b1111_1111, 0b0111_1111]);
-        check(16_384, &[0b1000_0000, 0b1000_0000, 0b0000_0001]);
+        check(u64::MAX, &[0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x01]);
     }
 
     /// This big bowl o' macro soup generates a quickcheck encoding test for each


Test patch: