Repository: tetratelabs/wazero
Pull request #941 (2022-12-19 02:09:10)

leb128: no allocations in decoding
This is a hot pass during compilation. Found while doing #936, and raised this as a separate PR.
```
benchstat leb128_old.txt leb128_new.txt
name             old time/op    new time/op    delta
LoadUint32-10      41.8ns ± 1%    12.0ns ± 0%   -71.35%  (p=0.000 n=10+10)
LoadUint64-10      4.96ns ± 0%    4.95ns ± 0%      ~     (p=0.813 n=9+9)
LoadInt32-10       42.1ns ± 2%    12.4ns ± 0%   -70.56%  (p=0.000 n=10+10)
LoadInt64-10       41.7ns ± 1%    12.4ns ± 0%   -70.27%  (p=0.000 n=10+10)
DecodeUint32-10    38.4ns ± 1%    17.3ns ± 1%   -55.03%  (p=0.000 n=10+10)
DecodeInt32-10     39.3ns ± 1%    17.8ns ± 0%   -54.59%  (p=0.000 n=10+9)
DecodeInt64-10     39.3ns ± 0%    17.6ns ± 0%   -55.28%  (p=0.000 n=10+10)

name             old alloc/op   new alloc/op   delta
LoadUint32-10       28.0B ± 0%      0.0B       -100.00%  (p=0.000 n=10+10)
LoadUint64-10       0.00B          0.00B           ~     (all equal)
LoadInt32-10        28.0B ± 0%      0.0B       -100.00%  (p=0.000 n=10+10)
LoadInt64-10        28.0B ± 0%      0.0B       -100.00%  (p=0.000 n=10+10)
DecodeUint32-10     16.0B ± 0%      0.0B       -100.00%  (p=0.000 n=10+10)
DecodeInt32-10      16.0B ± 0%      0.0B       -100.00%  (p=0.000 n=10+10)
DecodeInt64-10      16.0B ± 0%      0.0B       -100.00%  (p=0.000 n=10+10)

name             old allocs/op  new allocs/op  delta
LoadUint32-10        2.00 ± 0%      0.00       -100.00%  (p=0.000 n=10+10)
LoadUint64-10        0.00           0.00           ~     (all equal)
LoadInt32-10         2.00 ± 0%      0.00       -100.00%  (p=0.000 n=10+10)
LoadInt64-10         2.00 ± 0%      0.00       -100.00%  (p=0.000 n=10+10)
DecodeUint32-10      1.00 ± 0%      0.00       -100.00%  (p=0.000 n=10+10)
DecodeInt32-10       1.00 ± 0%      0.00       -100.00%  (p=0.000 n=10+10)
DecodeInt64-10       1.00 ± 0%      0.00       -100.00%  (p=0.000 n=10+10)

```


Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>



Patch:
diff --git a/internal/leb128/leb128.go b/internal/leb128/leb128.go
index 0e35585d..52eccfe4 100644
--- a/internal/leb128/leb128.go
+++ b/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)
 		}


Test patch:
diff --git a/internal/leb128/leb128_alloc_test.go b/internal/leb128/leb128_alloc_test.go
new file mode 100644
index 00000000..ccc9d0eb
--- /dev/null
+++ b/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)
+	}
+}