leb128: no allocations in decoding (#941)

Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>

internal/leb128/leb128.go

@@ -93,37 +93,27 @@ func EncodeUint64(value uint64) (buf []byte) {
 	}
 }
 
-type nextByte interface {
-	next(i int) (byte, error)
-}
-
-type byteSliceNext []byte
-
-func (n byteSliceNext) next(i int) (byte, error) {
-	if i >= len(n) {
-		return 0, io.EOF
-	}
-	return n[i], nil
-}
-
-type byteReaderNext struct{ io.ByteReader }
-
-func (n byteReaderNext) next(_ int) (byte, error) { return n.ReadByte() }
+type nextByte func(i int) (byte, error)
 
 func DecodeUint32(r io.ByteReader) (ret uint32, bytesRead uint64, err error) {
-	return decodeUint32(byteReaderNext{r})
+	return decodeUint32(func(_ int) (byte, error) { return r.ReadByte() })
 }
 
 func LoadUint32(buf []byte) (ret uint32, bytesRead uint64, err error) {
-	return decodeUint32(byteSliceNext(buf))
+	return decodeUint32(func(i int) (byte, error) {
+		if i >= len(buf) {
+			return 0, io.EOF
+		}
+		return buf[i], nil
+	})
 }
 
-func decodeUint32(buf nextByte) (ret uint32, bytesRead uint64, err error) {
+func decodeUint32(next nextByte) (ret uint32, bytesRead uint64, err error) {
 	// Derived from https://github.com/golang/go/blob/aafad20b617ee63d58fcd4f6e0d98fe27760678c/src/encoding/binary/varint.go
 	// with the modification on the overflow handling tailored for 32-bits.
 	var s uint32
 	for i := 0; i < maxVarintLen32; i++ {
-		b, err := buf.next(i)
+		b, err := next(i)
 		if err != nil {
 			return 0, 0, err
 		}
@@ -167,18 +157,23 @@ func LoadUint64(buf []byte) (ret uint64, bytesRead uint64, err error) {
 }
 
 func DecodeInt32(r io.ByteReader) (ret int32, bytesRead uint64, err error) {
-	return decodeInt32(byteReaderNext{r})
+	return decodeInt32(func(_ int) (byte, error) { return r.ReadByte() })
 }
 
 func LoadInt32(buf []byte) (ret int32, bytesRead uint64, err error) {
-	return decodeInt32(byteSliceNext(buf))
+	return decodeInt32(func(i int) (byte, error) {
+		if i >= len(buf) {
+			return 0, io.EOF
+		}
+		return buf[i], nil
+	})
 }
 
-func decodeInt32(buf nextByte) (ret int32, bytesRead uint64, err error) {
+func decodeInt32(next nextByte) (ret int32, bytesRead uint64, err error) {
 	var shift int
 	var b byte
 	for {
-		b, err = buf.next(int(bytesRead))
+		b, err = next(int(bytesRead))
 		if err != nil {
 			return 0, 0, fmt.Errorf("readByte failed: %w", err)
 		}
@@ -248,18 +243,23 @@ func DecodeInt33AsInt64(r io.ByteReader) (ret int64, bytesRead uint64, err error
 }
 
 func DecodeInt64(r io.ByteReader) (ret int64, bytesRead uint64, err error) {
-	return decodeInt64(byteReaderNext{r})
+	return decodeInt64(func(_ int) (byte, error) { return r.ReadByte() })
 }
 
 func LoadInt64(buf []byte) (ret int64, bytesRead uint64, err error) {
-	return decodeInt64(byteSliceNext(buf))
+	return decodeInt64(func(i int) (byte, error) {
+		if i >= len(buf) {
+			return 0, io.EOF
+		}
+		return buf[i], nil
+	})
 }
 
-func decodeInt64(buf nextByte) (ret int64, bytesRead uint64, err error) {
+func decodeInt64(next nextByte) (ret int64, bytesRead uint64, err error) {
 	var shift int
 	var b byte
 	for {
-		b, err = buf.next(int(bytesRead))
+		b, err = next(int(bytesRead))
 		if err != nil {
 			return 0, 0, fmt.Errorf("readByte failed: %w", err)
 		}

internal/leb128/leb128_alloc_test.go

@@ -0,0 +1,119 @@
+package leb128
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/tetratelabs/wazero/internal/testing/require"
+)
+
+// TestLeb128NoAlloc ensures no allocation required in the leb128 package.
+func TestLeb128NoAlloc(t *testing.T) {
+	t.Run("LoadUint32", func(t *testing.T) {
+		result := testing.Benchmark(BenchmarkLoadUint32)
+		require.Zero(t, result.AllocsPerOp())
+	})
+	t.Run("LoadUint64", func(t *testing.T) {
+		result := testing.Benchmark(BenchmarkLoadUint64)
+		require.Zero(t, result.AllocsPerOp())
+	})
+	t.Run("LoadInt32", func(t *testing.T) {
+		result := testing.Benchmark(BenchmarkLoadInt32)
+		require.Zero(t, result.AllocsPerOp())
+	})
+	t.Run("LoadInt64", func(t *testing.T) {
+		result := testing.Benchmark(BenchmarkLoadInt64)
+		require.Zero(t, result.AllocsPerOp())
+	})
+	t.Run("DecodeUint32", func(t *testing.T) {
+		result := testing.Benchmark(BenchmarkDecodeUint32)
+		require.Zero(t, result.AllocsPerOp())
+	})
+	t.Run("DecodeInt32", func(t *testing.T) {
+		result := testing.Benchmark(BenchmarkDecodeInt32)
+		require.Zero(t, result.AllocsPerOp())
+	})
+	t.Run("DecodeInt64", func(t *testing.T) {
+		result := testing.Benchmark(BenchmarkDecodeInt64)
+		require.Zero(t, result.AllocsPerOp())
+	})
+}
+
+func BenchmarkLoadUint32(b *testing.B) {
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		_, _, err := LoadUint32([]byte{0x80, 0x80, 0x80, 0x4f})
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func BenchmarkLoadUint64(b *testing.B) {
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		_, _, err := LoadUint64([]byte{0x80, 0x80, 0x80, 0x4f})
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func BenchmarkLoadInt32(b *testing.B) {
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		_, _, err := LoadInt32([]byte{0x80, 0x80, 0x80, 0x4f})
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func BenchmarkLoadInt64(b *testing.B) {
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		_, _, err := LoadInt64([]byte{0x80, 0x80, 0x80, 0x4f})
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func BenchmarkDecodeUint32(b *testing.B) {
+	b.ReportAllocs()
+	data := []byte{0x80, 0x80, 0x80, 0x4f}
+	r := bytes.NewReader(data)
+	for i := 0; i < b.N; i++ {
+		_, _, err := DecodeUint32(r)
+		if err != nil {
+			b.Fatal(err)
+		}
+		r.Reset(data)
+	}
+}
+
+func BenchmarkDecodeInt32(b *testing.B) {
+	b.ReportAllocs()
+	data := []byte{0x80, 0x80, 0x80, 0x4f}
+	r := bytes.NewReader(data)
+	for i := 0; i < b.N; i++ {
+		_, _, err := DecodeInt32(r)
+		if err != nil {
+			b.Fatal(err)
+		}
+		r.Reset(data)
+	}
+}
+
+func BenchmarkDecodeInt64(b *testing.B) {
+	b.ReportAllocs()
+	data := []byte{0x80, 0x80, 0x80, 0x4f}
+	r := bytes.NewReader(data)
+	for i := 0; i < b.N; i++ {
+		_, _, err := DecodeInt64(r)
+		if err != nil {
+			b.Fatal(err)
+		}
+		r.Reset(data)
+	}
+}