From 68a1f18b65a923c9174f13d2cad74d22edb0906e Mon Sep 17 00:00:00 2001
From: Matthew Kim <38759997+friendlymatthew@users.noreply.github.com>
Date: Sun, 17 Mar 2024 14:28:59 -0400
Subject: [PATCH] feat: add sizevarint + sizeuvarint

---
 src/encoding/binary/varint.go      | 17 +++++++++++++++++
 src/encoding/binary/varint_test.go | 10 ++++++++++
 2 files changed, 27 insertions(+)

diff --git a/src/encoding/binary/varint.go b/src/encoding/binary/varint.go
index 64dd9d61b44..eb969b8b301 100644
--- a/src/encoding/binary/varint.go
+++ b/src/encoding/binary/varint.go
@@ -27,6 +27,7 @@ package binary
 import (
 	"errors"
 	"io"
+	"math/bits"
 )
 
 // MaxVarintLenN is the maximum length of a varint-encoded N-bit integer.
@@ -164,3 +165,19 @@ func ReadVarint(r io.ByteReader) (int64, error) {
 	}
 	return x, err
 }
+
+// SizeUvarint returns the encoded size of a unsigned varint.
+// The size is guaranteed to be within 1 and 10, inclusive.
+func SizeUvarint(x uint64) int {
+	return int(9*uint32(bits.Len64(x))+64) / 64
+}
+
+// SizeVarint returns the encoded size of a signed varint.
+// The size is guaranteed to be within 1 and 10, inclusive.
+func SizeVarint(x int64) int {
+	ux := uint64(x) << 1
+	if x < 0 {
+		ux = ^ux
+	}
+	return SizeUvarint(ux)
+}
diff --git a/src/encoding/binary/varint_test.go b/src/encoding/binary/varint_test.go
index 5c3ea318c39..ab6698459cd 100644
--- a/src/encoding/binary/varint_test.go
+++ b/src/encoding/binary/varint_test.go
@@ -36,6 +36,11 @@ func testVarint(t *testing.T, x int64) {
 		t.Errorf("Varint(%d): got n = %d; want %d", x, m, n)
 	}
 
+	l := SizeVarint(x)
+	if n != l {
+		t.Errorf("SizeVarint(%d): got n = %d, want %d", x, m, l)
+	}
+
 	buf2 := []byte("prefix")
 	buf2 = AppendVarint(buf2, x)
 	if string(buf2) != "prefix"+string(buf[:n]) {
@@ -62,6 +67,11 @@ func testUvarint(t *testing.T, x uint64) {
 		t.Errorf("Uvarint(%d): got n = %d; want %d", x, m, n)
 	}
 
+	l := SizeUvarint(x)
+	if n != l {
+		t.Errorf("SizeUvarint(%d): got n = %d, want %d", x, m, l)
+	}
+
 	buf2 := []byte("prefix")
 	buf2 = AppendUvarint(buf2, x)
 	if string(buf2) != "prefix"+string(buf[:n]) {