From ed42c9e3abde49c2b18feefaef316c20cb7472e7 Mon Sep 17 00:00:00 2001
From: Levi Morrison <levi.morrison@datadoghq.com>
Date: Tue, 30 Jan 2024 20:25:37 -0700
Subject: [PATCH 1/2] perf: fix encode_varint{32,64}

The previous code was full of jumps that aren't easily predictable.
This version is less code, generates less code, and performs better
(well, on x86_64 anyway, I am not an aarch64 expert). All while using
safe Rust.
---
 protobuf/src/varint/encode.rs | 117 ++++++++++++----------------------
 1 file changed, 42 insertions(+), 75 deletions(-)

diff --git a/protobuf/src/varint/encode.rs b/protobuf/src/varint/encode.rs
index eb9fc2d03..4f751f4f2 100644
--- a/protobuf/src/varint/encode.rs
+++ b/protobuf/src/varint/encode.rs
@@ -1,92 +1,59 @@
 use std::mem::MaybeUninit;
 
+use crate::varint::MAX_VARINT32_ENCODED_LEN;
 use crate::varint::MAX_VARINT_ENCODED_LEN;
 
-/// Encode u64 as varint.
-/// Panics if buffer length is less than 10.
-#[inline]
-pub(crate) fn encode_varint64(mut value: u64, buf: &mut [MaybeUninit<u8>]) -> usize {
-    assert!(buf.len() >= MAX_VARINT_ENCODED_LEN);
+struct VarInt64Iterator {
+    num: u64,
+    cont: bool,
+}
 
-    fn iter(value: &mut u64, byte: &mut MaybeUninit<u8>) -> bool {
-        if (*value & !0x7F) > 0 {
-            byte.write(((*value & 0x7F) | 0x80) as u8);
-            *value >>= 7;
-            true
+impl VarInt64Iterator {
+    fn new(num: u64) -> Self {
+        Self { num, cont: true }
+    }
+}
+
+impl Iterator for VarInt64Iterator {
+    type Item = u8;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.cont {
+            self.cont = (self.num & !0x7F) != 0;
+            let num = self.num;
+            let val = if self.cont { (num & 0x7F) | 0x80 } else { num } as u8;
+            self.num >>= 7;
+            Some(val)
         } else {
-            byte.write(*value as u8);
-            false
+            None
         }
     }
+}
 
-    // Explicitly unroll loop to avoid either
-    // unsafe code or bound checking when writing to `buf`
-
-    if !iter(&mut value, &mut buf[0]) {
-        return 1;
-    };
-    if !iter(&mut value, &mut buf[1]) {
-        return 2;
-    };
-    if !iter(&mut value, &mut buf[2]) {
-        return 3;
-    };
-    if !iter(&mut value, &mut buf[3]) {
-        return 4;
-    };
-    if !iter(&mut value, &mut buf[4]) {
-        return 5;
-    };
-    if !iter(&mut value, &mut buf[5]) {
-        return 6;
-    };
-    if !iter(&mut value, &mut buf[6]) {
-        return 7;
-    };
-    if !iter(&mut value, &mut buf[7]) {
-        return 8;
-    };
-    if !iter(&mut value, &mut buf[8]) {
-        return 9;
-    };
-    buf[9].write(value as u8);
-    10
+fn encode_varint(value: u64, buf: &mut [MaybeUninit<u8>]) -> usize {
+    let iter = VarInt64Iterator::new(value);
+    let mut bytes_written = 0;
+    for (val, slot) in iter.zip(buf.iter_mut()) {
+        slot.write(val);
+        bytes_written += 1;
+    }
+    bytes_written
+}
+
+/// Encode u64 as varint.
+/// Panics if buffer length is less than 10.
+#[inline]
+pub(crate) fn encode_varint64(value: u64, buf: &mut [MaybeUninit<u8>]) -> usize {
+    assert!(buf.len() >= MAX_VARINT_ENCODED_LEN);
+    encode_varint(value, buf)
 }
 
 /// Encode u32 value as varint.
 /// Panics if buffer length is less than 5.
 #[inline]
-pub(crate) fn encode_varint32(mut value: u32, buf: &mut [MaybeUninit<u8>]) -> usize {
-    assert!(buf.len() >= 5);
-
-    fn iter(value: &mut u32, byte: &mut MaybeUninit<u8>) -> bool {
-        if (*value & !0x7F) > 0 {
-            byte.write(((*value & 0x7F) | 0x80) as u8);
-            *value >>= 7;
-            true
-        } else {
-            byte.write(*value as u8);
-            false
-        }
-    }
-
-    // Explicitly unroll loop to avoid either
-    // unsafe code or bound checking when writing to `buf`
-
-    if !iter(&mut value, &mut buf[0]) {
-        return 1;
-    };
-    if !iter(&mut value, &mut buf[1]) {
-        return 2;
-    };
-    if !iter(&mut value, &mut buf[2]) {
-        return 3;
-    };
-    if !iter(&mut value, &mut buf[3]) {
-        return 4;
-    };
-    buf[4].write(value as u8);
-    5
+pub(crate) fn encode_varint32(value: u32, buf: &mut [MaybeUninit<u8>]) -> usize {
+    assert!(buf.len() >= MAX_VARINT32_ENCODED_LEN);
+    encode_varint(value as u64, buf)
 }
 
 /// Encoded size of u64 value.

From c6f18ef1080c1ab5cac1010f1b7d5737c48c75d3 Mon Sep 17 00:00:00 2001
From: Levi Morrison <levi.morrison@datadoghq.com>
Date: Fri, 2 Feb 2024 14:08:47 -0700
Subject: [PATCH 2/2] refactor: simplify encoded_varint64_len

This greatly cleans up the assembly as well, but this function is not
as likely to matter from a perf perspective.
---
 protobuf/src/varint/encode.rs | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/protobuf/src/varint/encode.rs b/protobuf/src/varint/encode.rs
index 4f751f4f2..c8d984d2a 100644
--- a/protobuf/src/varint/encode.rs
+++ b/protobuf/src/varint/encode.rs
@@ -59,12 +59,10 @@ pub(crate) fn encode_varint32(value: u32, buf: &mut [MaybeUninit<u8>]) -> usize
 /// Encoded size of u64 value.
 #[inline]
 pub(crate) fn encoded_varint64_len(value: u64) -> usize {
-    if value == 0 {
-        1
-    } else {
-        let significant_bits = 64 - value.leading_zeros();
-        (significant_bits + 6) as usize / 7
-    }
+    // Bitwise-or'ing by 1 allows the `value = zero` case to work without
+    // affecting other cases.
+    let significant_bits = 64 - (value | 1).leading_zeros();
+    (significant_bits + 6) as usize / 7
 }
 
 #[cfg(test)]