quicvarint: speed up parsing of 1, 2 and 4 byte varints (#5229)

quicvarint/varint.go

@@ -1,6 +1,7 @@
 package quicvarint
 
 import (
+	"encoding/binary"
 	"fmt"
 	"io"
 )
@@ -74,23 +75,30 @@ func Parse(b []byte) (uint64 /* value */, int /* bytes consumed */, error) {
 	if len(b) == 0 {
 		return 0, 0, io.EOF
 	}
-	firstByte := b[0]
-	// the first two bits of the first byte encode the length
-	l := 1 << ((firstByte & 0xc0) >> 6)
-	if len(b) < l {
-		return 0, 0, io.ErrUnexpectedEOF
+
+	first := b[0]
+	switch first >> 6 {
+	case 0: // 1-byte encoding: 00xxxxxx
+		return uint64(first & 0b00111111), 1, nil
+	case 1: // 2-byte encoding: 01xxxxxx
+		if len(b) < 2 {
+			return 0, 0, io.ErrUnexpectedEOF
+		}
+		return uint64(b[1]) | uint64(first&0b00111111)<<8, 2, nil
+	case 2: // 4-byte encoding: 10xxxxxx
+		if len(b) < 4 {
+			return 0, 0, io.ErrUnexpectedEOF
+		}
+		return uint64(b[3]) | uint64(b[2])<<8 | uint64(b[1])<<16 | uint64(first&0b00111111)<<24, 4, nil
+	case 3: // 8-byte encoding: 00xxxxxx
+		if len(b) < 8 {
+			return 0, 0, io.ErrUnexpectedEOF
+		}
+		// binary.BigEndian.Uint64 only reads the first 8 bytes. Passing the full slice avoids slicing overhead.
+		return binary.BigEndian.Uint64(b) & 0x3fffffffffffffff, 8, nil
 	}
-	b0 := firstByte & (0xff - 0xc0)
-	if l == 1 {
-		return uint64(b0), 1, nil
-	}
-	if l == 2 {
-		return uint64(b[1]) + uint64(b0)<<8, 2, nil
-	}
-	if l == 4 {
-		return uint64(b[3]) + uint64(b[2])<<8 + uint64(b[1])<<16 + uint64(b0)<<24, 4, nil
-	}
-	return uint64(b[7]) + uint64(b[6])<<8 + uint64(b[5])<<16 + uint64(b[4])<<24 + uint64(b[3])<<32 + uint64(b[2])<<40 + uint64(b[1])<<48 + uint64(b0)<<56, 8, nil
+
+	panic("unreachable")
 }
 
 // Append appends i in the QUIC varint format.

quicvarint/varint_test.go

@@ -14,7 +14,7 @@ func TestLimits(t *testing.T) {
 	require.Equal(t, uint64(1<<62-1), uint64(Max))
 }
 
-func TestParsing(t *testing.T) {
+func TestRead(t *testing.T) {
 	tests := []struct {
 		name     string
 		input    []byte
@@ -38,6 +38,29 @@ func TestParsing(t *testing.T) {
 	}
 }
 
+func TestParse(t *testing.T) {
+	tests := []struct {
+		name          string
+		input         []byte
+		expectedValue uint64
+		expectedLen   int
+	}{
+		{"1 byte", []byte{0b00011001}, 25, 1},
+		{"2 byte", []byte{0b01111011, 0xbd}, 15293, 2},
+		{"4 byte", []byte{0b10011101, 0x7f, 0x3e, 0x7d}, 494878333, 4},
+		{"8 byte", []byte{0b11000010, 0x19, 0x7c, 0x5e, 0xff, 0x14, 0xe8, 0x8c}, 151288809941952652, 8},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			value, l, err := Parse(tt.input)
+			require.Equal(t, tt.expectedValue, value)
+			require.Equal(t, tt.expectedLen, l)
+			require.Nil(t, err)
+		})
+	}
+}
+
 func TestParsingFailures(t *testing.T) {
 	tests := []struct {
 		name        string
@@ -50,15 +73,27 @@ func TestParsingFailures(t *testing.T) {
 			expectedErr: io.EOF,
 		},
 		{
-			name:        "slice too short",
-			input:       Append(nil, maxVarInt2*10)[:3],
+			name:        "2-byte encoding: not enough bytes",
+			input:       []byte{0b01000001},
+			expectedErr: io.ErrUnexpectedEOF,
+		},
+		{
+			name:        "4-byte encoding: not enough bytes",
+			input:       []byte{0b10000000, 0x0, 0x0},
+			expectedErr: io.ErrUnexpectedEOF,
+		},
+		{
+			name:        "8-byte encoding: not enough bytes",
+			input:       []byte{0b11000000, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0},
 			expectedErr: io.ErrUnexpectedEOF,
 		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			_, _, err := Parse(tt.input)
+			value, l, err := Parse(tt.input)
+			require.Equal(t, uint64(0), value)
+			require.Equal(t, 0, l)
 			require.Equal(t, tt.expectedErr, err)
 		})
 	}