From 1cef82d54dda8726dcb5708b69fa51b04c215439 Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Wed, 5 May 2021 22:13:54 -0400
Subject: [PATCH] Unroll the VarInt writing loop

This is about as optimized as it can get. Thanks to @Leymooo for the idea, I simply expanded on it. We optimize for the common 1-3 byte cases, and punt more "complicated" cases to the original VarInt writing function we had before.
---
 .../proxy/protocol/ProtocolUtils.java         | 31 ++++++++++++++-----
 .../netty/MinecraftVarintLengthEncoder.java   |  3 +-
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/proxy/src/main/java/com/velocitypowered/proxy/protocol/ProtocolUtils.java b/proxy/src/main/java/com/velocitypowered/proxy/protocol/ProtocolUtils.java
index 9ac62ff69..1f696bb7a 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/protocol/ProtocolUtils.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/protocol/ProtocolUtils.java
@@ -120,15 +120,32 @@ public enum ProtocolUtils {
    * @param value the integer to write
    */
   public static void writeVarInt(ByteBuf buf, int value) {
-    // Inspired by https://richardstartin.github.io/posts/dont-use-protobuf-for-telemetry
-    // This has been slightly modified in that we reduce the length to 32-bit only, since Velocity
-    // doesn't look at any part of the Minecraft protocol that requires us to look at VarLongs.
-    int continuationBytes = (31 - Integer.numberOfLeadingZeros(value)) / 7;
-    for (int i = 0; i < continuationBytes; ++i) {
-      buf.writeByte(((byte) ((value & 0x7F) | 0x80)));
+    int bytes = varIntBytes(value);
+    // Optimization: focus on 1-3 byte VarInts as they are the most common
+    if (bytes == 1) {
+      buf.writeByte(value & 0x7f);
+    } else if (bytes == 2) {
+      int w = (value & 0x7F | 0x80) << 8 | (value >>> 7);
+      buf.writeShort(w);
+    } else if (bytes == 3) {
+      int w = (value & 0x7F | 0x80) << 16 | ((value >>> 7) & 0x7F | 0x80) << 8 | (value >>> 14);
+      buf.writeMedium(w);
+    } else {
+      // 4 and 5 byte VarInts aren't common so split those cases off
+      writeVarIntUncommon(buf, value);
+    }
+  }
+
+  private static void writeVarIntUncommon(ByteBuf buf, int value) {
+    while (true) {
+      if ((value & 0xFFFFFF80) == 0) {
+        buf.writeByte(value);
+        return;
+      }
+
+      buf.writeByte(value & 0x7F | 0x80);
       value >>>= 7;
     }
-    buf.writeByte((byte) value);
   }
 
   public static String readString(ByteBuf buf) {
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftVarintLengthEncoder.java b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftVarintLengthEncoder.java
index 867ea9b7f..94bfa84d0 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftVarintLengthEncoder.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftVarintLengthEncoder.java
@@ -43,7 +43,8 @@ public class MinecraftVarintLengthEncoder extends MessageToByteEncoder<ByteBuf>
   @Override
   protected ByteBuf allocateBuffer(ChannelHandlerContext ctx, ByteBuf msg, boolean preferDirect)
       throws Exception {
-    int anticipatedRequiredCapacity = 5 + msg.readableBytes();
+    int anticipatedRequiredCapacity = ProtocolUtils.varIntBytes(msg.readableBytes())
+        + msg.readableBytes();
     return IS_JAVA_CIPHER
         ? ctx.alloc().heapBuffer(anticipatedRequiredCapacity)
         : ctx.alloc().directBuffer(anticipatedRequiredCapacity);