perf: vendor and improve binary pkg

binary/compact-u16.go

@@ -0,0 +1,138 @@
+// Copyright 2021 github.com/gagliardetto
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package bin
+
+import (
+	"fmt"
+	"io"
+	"math"
+)
+
+// EncodeCompactU16Length encodes a "Compact-u16" length into the provided slice pointer.
+// See https://docs.solana.com/developing/programming-model/transactions#compact-u16-format
+// See https://github.com/solana-labs/solana/blob/2ef2b6daa05a7cff057e9d3ef95134cee3e4045d/web3.js/src/util/shortvec-encoding.ts
+func EncodeCompactU16Length(buf *[]byte, ln int) error {
+	if ln < 0 || ln > math.MaxUint16 {
+		return fmt.Errorf("length %d out of range", ln)
+	}
+	u := uint(ln)
+	switch {
+	case u < 0x80:
+		*buf = append(*buf, byte(u))
+	case u < 0x4000:
+		*buf = append(*buf, byte(u)|0x80, byte(u>>7))
+	default:
+		*buf = append(*buf, byte(u)|0x80, byte(u>>7)|0x80, byte(u>>14))
+	}
+	return nil
+}
+
+// PutCompactU16Length writes a "Compact-u16" length into dst and returns the
+// number of bytes written (1, 2, or 3). dst must be at least 3 bytes long.
+// This is the allocation-free variant of EncodeCompactU16Length, used by the
+// Encoder's scratch-buffer hot path.
+func PutCompactU16Length(dst []byte, ln int) (int, error) {
+	if ln < 0 || ln > math.MaxUint16 {
+		return 0, fmt.Errorf("length %d out of range", ln)
+	}
+	u := uint(ln)
+	switch {
+	case u < 0x80:
+		dst[0] = byte(u)
+		return 1, nil
+	case u < 0x4000:
+		dst[0] = byte(u) | 0x80
+		dst[1] = byte(u >> 7)
+		return 2, nil
+	default:
+		dst[0] = byte(u) | 0x80
+		dst[1] = byte(u>>7) | 0x80
+		dst[2] = byte(u >> 14)
+		return 3, nil
+	}
+}
+
+const _MAX_COMPACTU16_ENCODING_LENGTH = 3
+
+// DecodeCompactU16 decodes a Solana "Compact-u16" length from bytes and returns
+// (value, bytes_consumed, error). Hand-unrolled for the max 3-byte encoding to
+// avoid a per-iteration loop overhead.
+func DecodeCompactU16(bytes []byte) (int, int, error) {
+	if len(bytes) == 0 {
+		return 0, 0, io.ErrUnexpectedEOF
+	}
+	b0 := int(bytes[0])
+	if b0&0x80 == 0 {
+		return b0, 1, nil
+	}
+	if len(bytes) < 2 {
+		return 0, 0, io.ErrUnexpectedEOF
+	}
+	b1 := int(bytes[1])
+	if b1&0x80 == 0 {
+		if b1 == 0 {
+			return 0, 0, fmt.Errorf("compact-u16: non-canonical 2-byte encoding (trailing zero byte)")
+		}
+		return (b0 & 0x7f) | (b1 << 7), 2, nil
+	}
+	if len(bytes) < 3 {
+		return 0, 0, io.ErrUnexpectedEOF
+	}
+	b2 := int(bytes[2])
+	if b2 == 0 {
+		return 0, 0, fmt.Errorf("compact-u16: non-canonical 3-byte encoding (trailing zero byte)")
+	}
+	if b2&0x80 != 0 {
+		return 0, 0, fmt.Errorf("byte three continues")
+	}
+	ln := (b0 & 0x7f) | ((b1 & 0x7f) << 7) | (b2 << 14)
+	if ln > math.MaxUint16 {
+		return 0, 0, fmt.Errorf("invalid length: %d", ln)
+	}
+	return ln, 3, nil
+}
+
+// DecodeCompactU16LengthFromByteReader decodes a "Compact-u16" length from the provided io.ByteReader.
+func DecodeCompactU16LengthFromByteReader(reader io.ByteReader) (int, error) {
+	ln := 0
+	size := 0
+	for nthByte := range _MAX_COMPACTU16_ENCODING_LENGTH {
+		elemByte, err := reader.ReadByte()
+		if err != nil {
+			return 0, err
+		}
+		elem := int(elemByte)
+		if elem == 0 && nthByte != 0 {
+			return 0, fmt.Errorf("compact-u16: non-canonical encoding (trailing zero byte at position %d)", nthByte)
+		}
+		if nthByte == _MAX_COMPACTU16_ENCODING_LENGTH-1 && (elem&0x80) != 0 {
+			return 0, fmt.Errorf("compact-u16: byte three has continuation bit set")
+		}
+		ln |= (elem & 0x7f) << (size * 7)
+		size += 1
+		if (elem & 0x80) == 0 {
+			break
+		}
+	}
+	// check for non-valid sizes
+	if size == 0 || size > _MAX_COMPACTU16_ENCODING_LENGTH {
+		return 0, fmt.Errorf("compact-u16: invalid size: %d", size)
+	}
+	// check for non-valid lengths
+	if ln < 0 || ln > math.MaxUint16 {
+		return 0, fmt.Errorf("compact-u16: invalid length: %d", ln)
+	}
+	return ln, nil
+}

binary/compact-u16_test.go

@@ -0,0 +1,221 @@
+// Copyright 2021 github.com/gagliardetto
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package bin
+
+import (
+	"bytes"
+	"math"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestCompactU16(t *testing.T) {
+	candidates := []int{0, 1, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 100, 1000, 10000, math.MaxUint16 - 1, math.MaxUint16}
+	for _, val := range candidates {
+		if val < 0 || val > math.MaxUint16 {
+			panic("value too large")
+		}
+		buf := make([]byte, 0)
+		require.NoError(t, EncodeCompactU16Length(&buf, val))
+
+		buf = append(buf, []byte("hello world")...)
+		decoded, _, err := DecodeCompactU16(buf)
+		require.NoError(t, err)
+
+		require.Equal(t, val, decoded)
+	}
+	for _, val := range candidates {
+		buf := make([]byte, 0)
+		EncodeCompactU16Length(&buf, val)
+
+		buf = append(buf, []byte("hello world")...)
+		{
+			decoded, err := DecodeCompactU16LengthFromByteReader(bytes.NewReader(buf))
+			require.NoError(t, err)
+			require.Equal(t, val, decoded)
+		}
+		{
+			decoded, _, err := DecodeCompactU16(buf)
+			require.NoError(t, err)
+			require.Equal(t, val, decoded)
+		}
+	}
+	{
+		// now test all from 0 to 0xffff
+		for i := 0; i < math.MaxUint16; i++ {
+			buf := make([]byte, 0)
+			EncodeCompactU16Length(&buf, i)
+
+			buf = append(buf, []byte("hello world")...)
+			{
+				decoded, err := DecodeCompactU16LengthFromByteReader(bytes.NewReader(buf))
+				require.NoError(t, err)
+				require.Equal(t, i, decoded)
+			}
+			{
+				decoded, _, err := DecodeCompactU16(buf)
+				require.NoError(t, err)
+				require.Equal(t, i, decoded)
+			}
+		}
+	}
+}
+
+func BenchmarkCompactU16(b *testing.B) {
+	// generate 1000 random values
+	candidates := make([]int, 1000)
+	for i := 0; i < 1000; i++ {
+		candidates[i] = i
+	}
+
+	buf := make([]byte, 0)
+	EncodeCompactU16Length(&buf, math.MaxUint16)
+
+	b.ResetTimer()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		_, _, _ = DecodeCompactU16(buf)
+	}
+}
+
+func BenchmarkCompactU16Encode(b *testing.B) {
+	b.ResetTimer()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		buf := make([]byte, 0)
+		EncodeCompactU16Length(&buf, math.MaxUint16)
+	}
+}
+
+func BenchmarkCompactU16Reader(b *testing.B) {
+	// generate 1000 random values
+	candidates := make([]int, 1000)
+	for i := 0; i < 1000; i++ {
+		candidates[i] = i
+	}
+
+	buf := make([]byte, 0)
+	EncodeCompactU16Length(&buf, math.MaxUint16)
+
+	reader := NewBorshDecoder(buf)
+	b.ResetTimer()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		out, _ := reader.ReadCompactU16()
+		if out != math.MaxUint16 {
+			panic("not equal")
+		}
+		reader.SetPosition(0)
+	}
+}
+
+func encode_len(len uint16) []byte {
+	buf := make([]byte, 0)
+	err := EncodeCompactU16Length(&buf, int(len))
+	if err != nil {
+		panic(err)
+	}
+	return buf
+}
+
+func assert_len_encoding(t *testing.T, len uint16, buf []byte) {
+	require.Equal(t, encode_len(len), buf, "unexpected usize encoding")
+	decoded, _, err := DecodeCompactU16(buf)
+	require.NoError(t, err)
+	require.Equal(t, int(len), decoded)
+	{
+		// now try with a reader
+		reader := bytes.NewReader(buf)
+		out, _ := DecodeCompactU16LengthFromByteReader(reader)
+		require.Equal(t, int(len), out)
+	}
+}
+
+func TestShortVecEncodeLen(t *testing.T) {
+	assert_len_encoding(t, 0x0, []byte{0x0})
+	assert_len_encoding(t, 0x7f, []byte{0x7f})
+	assert_len_encoding(t, 0x80, []byte{0x80, 0x01})
+	assert_len_encoding(t, 0xff, []byte{0xff, 0x01})
+	assert_len_encoding(t, 0x100, []byte{0x80, 0x02})
+	assert_len_encoding(t, 0x7fff, []byte{0xff, 0xff, 0x01})
+	assert_len_encoding(t, 0xffff, []byte{0xff, 0xff, 0x03})
+}
+
+func assert_good_deserialized_value(t *testing.T, value uint16, buf []byte) {
+	decoded, _, err := DecodeCompactU16(buf)
+	require.NoError(t, err)
+	require.Equal(t, int(value), decoded)
+	{
+		// now try with a reader
+		reader := bytes.NewReader(buf)
+		out, _ := DecodeCompactU16LengthFromByteReader(reader)
+		require.Equal(t, int(value), out)
+	}
+}
+
+func assert_bad_deserialized_value(t *testing.T, buf []byte) {
+	_, _, err := DecodeCompactU16(buf)
+	require.Error(t, err, "expected an error for bytes: %v", buf)
+	{
+		// now try with a reader
+		reader := bytes.NewReader(buf)
+		_, err := DecodeCompactU16LengthFromByteReader(reader)
+		require.Error(t, err, "expected an error for bytes: %v", buf)
+	}
+}
+
+func TestDeserialize(t *testing.T) {
+	assert_good_deserialized_value(t, 0x0000, []byte{0x00})
+	assert_good_deserialized_value(t, 0x007f, []byte{0x7f})
+	assert_good_deserialized_value(t, 0x0080, []byte{0x80, 0x01})
+	assert_good_deserialized_value(t, 0x00ff, []byte{0xff, 0x01})
+	assert_good_deserialized_value(t, 0x0100, []byte{0x80, 0x02})
+	assert_good_deserialized_value(t, 0x07ff, []byte{0xff, 0x0f})
+	assert_good_deserialized_value(t, 0x3fff, []byte{0xff, 0x7f})
+	assert_good_deserialized_value(t, 0x4000, []byte{0x80, 0x80, 0x01})
+	assert_good_deserialized_value(t, 0xffff, []byte{0xff, 0xff, 0x03})
+
+	// aliases
+	// 0x0000
+	assert_bad_deserialized_value(t, []byte{0x80, 0x00})
+	assert_bad_deserialized_value(t, []byte{0x80, 0x80, 0x00})
+	// 0x007f
+	assert_bad_deserialized_value(t, []byte{0xff, 0x00})
+	assert_bad_deserialized_value(t, []byte{0xff, 0x80, 0x00})
+	// 0x0080
+	assert_bad_deserialized_value(t, []byte{0x80, 0x81, 0x00})
+	// 0x00ff
+	assert_bad_deserialized_value(t, []byte{0xff, 0x81, 0x00})
+	// 0x0100
+	assert_bad_deserialized_value(t, []byte{0x80, 0x82, 0x00})
+	// 0x07ff
+	assert_bad_deserialized_value(t, []byte{0xff, 0x8f, 0x00})
+	// 0x3fff
+	assert_bad_deserialized_value(t, []byte{0xff, 0xff, 0x00})
+
+	// too short
+	assert_bad_deserialized_value(t, []byte{})
+	assert_bad_deserialized_value(t, []byte{0x80})
+
+	// too long
+	assert_bad_deserialized_value(t, []byte{0x80, 0x80, 0x80, 0x00})
+
+	// too large
+	// 0x0001_0000
+	assert_bad_deserialized_value(t, []byte{0x80, 0x80, 0x04})
+	// 0x0001_8000
+	assert_bad_deserialized_value(t, []byte{0x80, 0x80, 0x06})
+}