From dd67478593c700bb90a781199acc302c09f0952c Mon Sep 17 00:00:00 2001
From: Tnze <cjd001113@outlook.com>
Date: Fri, 31 Mar 2023 13:12:45 +0800
Subject: [PATCH] Add Len method for VarInt and VarLong.

---
 net/packet/packet_test.go | 43 +++++++++++++++++++++++++++++++++++++++
 net/packet/types.go       | 42 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 85 insertions(+)

diff --git a/net/packet/packet_test.go b/net/packet/packet_test.go
index 9c1559b..6982a87 100644
--- a/net/packet/packet_test.go
+++ b/net/packet/packet_test.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	_ "embed"
 	"fmt"
+	"io"
 	"testing"
 
 	pk "github.com/Tnze/go-mc/net/packet"
@@ -60,6 +61,22 @@ func TestVarInt_ReadFrom_tooLongData(t *testing.T) {
 	}
 }
 
+func FuzzVarInt_Len(f *testing.F) {
+	for _, v := range VarInts {
+		f.Add(int32(v))
+	}
+	var buf bytes.Buffer
+	f.Fuzz(func(t *testing.T, v int32) {
+		defer buf.Reset()
+		if _, err := pk.VarInt(v).WriteTo(&buf); err != nil {
+			t.Fatal(err)
+		}
+		if a, b := buf.Len(), pk.VarInt(v).Len(); a != b {
+			t.Errorf("%#v Length calculation error: calculated to be %d, actually %d", pk.VarInt(v), b, a)
+		}
+	})
+}
+
 var VarLongs = []pk.VarLong{0, 1, 2, 127, 128, 255, 2147483647, 9223372036854775807, -1, -2147483648, -9223372036854775808}
 
 var PackedVarLongs = [][]byte{
@@ -101,6 +118,22 @@ func TestVarLong_ReadFrom(t *testing.T) {
 	}
 }
 
+func FuzzVarLong_Len(f *testing.F) {
+	for _, v := range VarLongs {
+		f.Add(int64(v))
+	}
+	var buf bytes.Buffer
+	f.Fuzz(func(t *testing.T, v int64) {
+		defer buf.Reset()
+		if _, err := pk.VarLong(v).WriteTo(&buf); err != nil {
+			t.Fatal(err)
+		}
+		if a, b := buf.Len(), pk.VarLong(v).Len(); a != b {
+			t.Errorf("%#v Length calculation error: calculated to be %d, actually %d", pk.VarLong(v), b, a)
+		}
+	})
+}
+
 //go:embed joingame_test.bin
 var testJoinGameData []byte
 
@@ -171,3 +204,13 @@ func ExampleMarshal_setSlot() {
 	// 15 00 00 05 01 01 01 00
 	// 15 00 00 05 01 01 01 03 00 00 12 34 56 78
 }
+
+func BenchmarkPacket_Pack(b *testing.B) {
+	p := pk.Packet{ID: 0, Data: make([]byte, 64)}
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		if err := p.Pack(io.Discard, -1); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
diff --git a/net/packet/types.go b/net/packet/types.go
index f601a20..050fc84 100644
--- a/net/packet/types.go
+++ b/net/packet/types.go
@@ -290,6 +290,23 @@ func (v *VarInt) ReadFrom(r io.Reader) (n int64, err error) {
 	return
 }
 
+func (v VarInt) Len() int {
+	switch {
+	case v < 0:
+		return MaxVarIntLen
+	case v < 1<<(7*1):
+		return 1
+	case v < 1<<(7*2):
+		return 2
+	case v < 1<<(7*3):
+		return 3
+	case v < 1<<(7*4):
+		return 4
+	default:
+		return 5
+	}
+}
+
 func (v VarLong) WriteTo(w io.Writer) (n int64, err error) {
 	vi := make([]byte, 0, MaxVarLongLen)
 	num := uint64(v)
@@ -328,6 +345,31 @@ func (v *VarLong) ReadFrom(r io.Reader) (n int64, err error) {
 	return
 }
 
+func (v VarLong) Len() int {
+	switch {
+	case v < 0:
+		return MaxVarLongLen
+	case v < 1<<(7*1):
+		return 1
+	case v < 1<<(7*2):
+		return 2
+	case v < 1<<(7*3):
+		return 3
+	case v < 1<<(7*4):
+		return 4
+	case v < 1<<(7*5):
+		return 5
+	case v < 1<<(7*6):
+		return 6
+	case v < 1<<(7*7):
+		return 7
+	case v < 1<<(7*8):
+		return 8
+	default:
+		return 9
+	}
+}
+
 func (p Position) WriteTo(w io.Writer) (n int64, err error) {
 	var b [8]byte
 	position := uint64(p.X&0x3FFFFFF)<<38 | uint64((p.Z&0x3FFFFFF)<<12) | uint64(p.Y&0xFFF)