chore: Improve VarInt write

layou233 · Sep 1, 2024 · 1f44918 · 1f44918
1 parent 0bb3c52
commit 1f44918
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 21 deletions.
diff --git a/common/mcprotocol/varint.go b/common/mcprotocol/varint.go
@@ -1,6 +1,7 @@
 package mcprotocol
 
 import (
+	"encoding/binary"
 	"errors"
 	"io"
 
@@ -37,17 +38,32 @@ func (i VarInt) WriteToBuffer(buffer *buf.Buffer) {
 // PutVarInt encodes a Minecraft variable-length format int32 into bs and returns the number of bytes written.
 // If the buffer is too small, PutVarInt will panic.
 func PutVarInt(bs []byte, n int32) (numWrite int) {
+	// https://steinborn.me/posts/performance/how-fast-can-you-write-a-varint/
 	num := uint32(n)
-	for num != 0 {
-		b := num & 0x7F
-		num >>= 7
-		if num != 0 {
-			b |= 0x80
-		}
-		bs[numWrite] = byte(b)
-		numWrite++
+	if num&0xFFFFFF80 == 0 {
+		bs[0] = byte(num)
+		return 1
+	} else if num&0xFFFFC000 == 0 {
+		result := uint16((num&0x7F|0x80)<<8 | (num >> 7))
+		binary.BigEndian.PutUint16(bs, result)
+		return 2
+	} else if num&0xFFE00000 == 0 {
+		bs[2] = byte(num >> 14)
+		startingBytes := uint16((num&0x7F|0x80)<<8 | ((num>>7)&0x7F | 0x80))
+		binary.BigEndian.PutUint16(bs, startingBytes)
+		return 3
+	} else if num&0xF0000000 == 0 {
+		result := (num&0x7F|0x80)<<24 | (((num>>7)&0x7F | 0x80) << 16) |
+			((num>>14)&0x7F|0x80)<<8 | (num >> 21)
+		binary.BigEndian.PutUint32(bs, result)
+		return 4
+	} else {
+		bs[4] = byte(num >> 28)
+		startingBytes := (num&0x7F|0x80)<<24 | ((num>>7)&0x7F|0x80)<<16 |
+			((num>>14)&0x7F|0x80)<<8 | ((num>>21)&0x7F | 0x80)
+		binary.BigEndian.PutUint32(bs, startingBytes)
+		return 5
 	}
-	return
 }
 
 func VarIntLen(n int32) int {

diff --git a/common/mcprotocol/varint_test.go b/common/mcprotocol/varint_test.go
@@ -10,14 +10,17 @@ import (
 
 // samples from https://wiki.vg/Protocol#VarInt_and_VarLong
 
-func checkWrite(t *testing.T, n int32, result [MaxVarIntLen]byte) {
+func checkWrite(t *testing.T, n int32, result [MaxVarIntLen]byte, expectedLen int) {
 	buffer := buf.NewSize(MaxVarIntLen + 1)
 	defer buffer.Release()
 	_, err := VarInt(n).WriteTo(buffer)
 	if err != nil {
 		return
 	}
 	t.Log("VarInt", n, "WriteTo", buffer.Bytes())
+	if buffer.Len() != expectedLen {
+		t.Errorf("VarInt WriteTo length error: got %d, expect %d", buffer.Len(), expectedLen)
+	}
 	buffer.Truncate(MaxVarIntLen)
 	if !bytes.Equal(buffer.Bytes(), result[:]) {
 		t.Fatalf("VarInt WriteTo error: got %v, expect %v", buffer.Bytes(), result)
@@ -38,17 +41,17 @@ func checkRead(t *testing.T, n int32, result [MaxVarIntLen]byte) {
 }
 
 func TestVarInt_WriteTo(t *testing.T) {
-	checkWrite(t, 0, [MaxVarIntLen]byte{0})
-	checkWrite(t, 1, [MaxVarIntLen]byte{1})
-	checkWrite(t, 2, [MaxVarIntLen]byte{2})
-	checkWrite(t, 127, [MaxVarIntLen]byte{127})
-	checkWrite(t, 128, [MaxVarIntLen]byte{128, 1})
-	checkWrite(t, 255, [MaxVarIntLen]byte{255, 1})
-	checkWrite(t, 25565, [MaxVarIntLen]byte{221, 199, 1})
-	checkWrite(t, 2097151, [MaxVarIntLen]byte{255, 255, 127})
-	checkWrite(t, 2147483647, [MaxVarIntLen]byte{255, 255, 255, 255, 7})
-	checkWrite(t, -1, [MaxVarIntLen]byte{255, 255, 255, 255, 15})
-	checkWrite(t, -2147483648, [MaxVarIntLen]byte{128, 128, 128, 128, 8})
+	checkWrite(t, 0, [MaxVarIntLen]byte{0}, 1)
+	checkWrite(t, 1, [MaxVarIntLen]byte{1}, 1)
+	checkWrite(t, 2, [MaxVarIntLen]byte{2}, 1)
+	checkWrite(t, 127, [MaxVarIntLen]byte{127}, 1)
+	checkWrite(t, 128, [MaxVarIntLen]byte{128, 1}, 2)
+	checkWrite(t, 255, [MaxVarIntLen]byte{255, 1}, 2)
+	checkWrite(t, 25565, [MaxVarIntLen]byte{221, 199, 1}, 3)
+	checkWrite(t, 2097151, [MaxVarIntLen]byte{255, 255, 127}, 3)
+	checkWrite(t, 2147483647, [MaxVarIntLen]byte{255, 255, 255, 255, 7}, 5)
+	checkWrite(t, -1, [MaxVarIntLen]byte{255, 255, 255, 255, 15}, 5)
+	checkWrite(t, -2147483648, [MaxVarIntLen]byte{128, 128, 128, 128, 8}, 5)
 }
 
 func TestReadFrom(t *testing.T) {