From 150fd9a9cf239b769fe46f607c95a85a6c197377 Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Sat, 8 May 2021 23:26:43 -0400
Subject: [PATCH] Add highly-optimized VarInt writing method

---
 .../proxy/protocol/ProtocolUtils.java         | 30 ++++++++-----------
 .../MinecraftCompressorAndLengthEncoder.java  |  4 +--
 .../proxy/protocol/ProtocolUtilsTest.java     |  2 +-
 3 files changed, 16 insertions(+), 20 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 a40598e43..ee3482c9d 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/protocol/ProtocolUtils.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/protocol/ProtocolUtils.java
@@ -120,7 +120,7 @@ public enum ProtocolUtils {
    * @param value the integer to write
    */
   public static void writeVarInt(ByteBuf buf, int value) {
-    // Optimization: focus on 1-3 byte VarInts as they are the most common
+    // See https://steinborn.me/posts/performance/how-fast-can-you-write-a-varint/
     if ((value & (0xFFFFFFFF << 7)) == 0) {
       buf.writeByte(value);
     } else if ((value & (0xFFFFFFFF << 14)) == 0) {
@@ -129,30 +129,26 @@ public enum ProtocolUtils {
     } else if ((value & (0xFFFFFFFF << 21)) == 0) {
       int w = (value & 0x7F | 0x80) << 16 | ((value >>> 7) & 0x7F | 0x80) << 8 | (value >>> 14);
       buf.writeMedium(w);
+    } else if ((value & (0xFFFFFFFF << 28)) == 0) {
+      int w = (value & 0x7F | 0x80) << 24 | (((value >>> 7) & 0x7F | 0x80) << 16)
+          | ((value >>> 14) & 0x7F | 0x80) << 8 | (value >>> 21);
+      buf.writeInt(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;
+      int w = (value & 0x7F | 0x80) << 24 | ((value >>> 7) & 0x7F | 0x80) << 16
+          | ((value >>> 14) & 0x7F | 0x80) << 8 | ((value >>> 21) & 0x7F | 0x80);
+      buf.writeInt(w);
+      buf.writeByte(value >>> 28);
     }
   }
 
   /**
-   * Writes all integers as 3 bytes Minecraft-style VarInt to the specified {@code buf}.
+   * Writes the specified {@code value} as a 21-bit Minecraft VarInt to the specified {@code buf}.
+   * The upper 11 bits will be discarded.
    * @param buf the buffer to read from
    * @param value the integer to write
    */
-  public static void writeVarIntAs3Bytes(ByteBuf buf, int value) {
+  public static void write21BitVarInt(ByteBuf buf, int value) {
+    // See https://steinborn.me/posts/performance/how-fast-can-you-write-a-varint/
     int w = (value & 0x7F | 0x80) << 16 | ((value >>> 7) & 0x7F | 0x80) << 8 | (value >>> 14);
     buf.writeMedium(w);
   }
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressorAndLengthEncoder.java b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressorAndLengthEncoder.java
index 34d247bc7..28d005161 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressorAndLengthEncoder.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressorAndLengthEncoder.java
@@ -61,7 +61,7 @@ public class MinecraftCompressorAndLengthEncoder extends MessageToByteEncoder<By
       throws DataFormatException {
     int uncompressed = msg.readableBytes();
 
-    ProtocolUtils.writeVarIntAs3Bytes(out, 0); // Dummy packet length
+    ProtocolUtils.write21BitVarInt(out, 0); // Dummy packet length
     ProtocolUtils.writeVarInt(out, uncompressed);
     ByteBuf compatibleIn = MoreByteBufUtils.ensureCompatible(ctx.alloc(), compressor, msg);
 
@@ -81,7 +81,7 @@ public class MinecraftCompressorAndLengthEncoder extends MessageToByteEncoder<By
     int writerIndex = out.writerIndex();
     int packetLength = out.readableBytes() - 3;
     out.writerIndex(0);
-    ProtocolUtils.writeVarIntAs3Bytes(out, packetLength); // Rewrite packet length
+    ProtocolUtils.write21BitVarInt(out, packetLength); // Rewrite packet length
     out.writerIndex(writerIndex);
   }
 
diff --git a/proxy/src/test/java/com/velocitypowered/proxy/protocol/ProtocolUtilsTest.java b/proxy/src/test/java/com/velocitypowered/proxy/protocol/ProtocolUtilsTest.java
index 58530a4a3..04b2f3193 100644
--- a/proxy/src/test/java/com/velocitypowered/proxy/protocol/ProtocolUtilsTest.java
+++ b/proxy/src/test/java/com/velocitypowered/proxy/protocol/ProtocolUtilsTest.java
@@ -81,7 +81,7 @@ public class ProtocolUtilsTest {
 
   private void writeReadTest3Bytes(ByteBuf buf, int test) {
     buf.clear();
-    ProtocolUtils.writeVarIntAs3Bytes(buf, test);
+    ProtocolUtils.write21BitVarInt(buf, test);
     assertEquals(test, ProtocolUtils.readVarInt(buf));
   }