Eliminate interface-boxing allocations in blocks.go

go.mod

@@ -1,3 +1,3 @@
 module github.com/pierrec/lz4/v4
 
-go 1.14
+go 1.17

internal/lz4block/blocks.go

@@ -1,7 +1,10 @@
 // Package lz4block provides LZ4 BlockSize types and pools of buffers.
 package lz4block
 
-import "sync"
+import (
+	"fmt"
+	"sync"
+)
 
 const (
 	Block64Kb uint32 = 1 << (16 + iota*2)
@@ -12,11 +15,11 @@ const (
 )
 
 var (
-	BlockPool64K  = sync.Pool{New: func() interface{} { return make([]byte, Block64Kb) }}
-	BlockPool256K = sync.Pool{New: func() interface{} { return make([]byte, Block256Kb) }}
-	BlockPool1M   = sync.Pool{New: func() interface{} { return make([]byte, Block1Mb) }}
-	BlockPool4M   = sync.Pool{New: func() interface{} { return make([]byte, Block4Mb) }}
-	BlockPool8M   = sync.Pool{New: func() interface{} { return make([]byte, Block8Mb) }}
+	blockPool64K  = sync.Pool{New: func() interface{} { return &[Block64Kb]byte{} }}
+	blockPool256K = sync.Pool{New: func() interface{} { return &[Block256Kb]byte{} }}
+	blockPool1M   = sync.Pool{New: func() interface{} { return &[Block1Mb]byte{} }}
+	blockPool4M   = sync.Pool{New: func() interface{} { return &[Block4Mb]byte{} }}
+	blockPool8M   = sync.Pool{New: func() interface{} { return &[Block8Mb]byte{} }}
 )
 
 func Index(b uint32) BlockSizeIndex {
@@ -50,35 +53,37 @@ func (b BlockSizeIndex) IsValid() bool {
 }
 
 func (b BlockSizeIndex) Get() []byte {
-	var buf interface{}
 	switch b {
 	case 4:
-		buf = BlockPool64K.Get()
+		return blockPool64K.Get().(*[Block64Kb]byte)[:]
 	case 5:
-		buf = BlockPool256K.Get()
+		return blockPool256K.Get().(*[Block256Kb]byte)[:]
 	case 6:
-		buf = BlockPool1M.Get()
+		return blockPool1M.Get().(*[Block1Mb]byte)[:]
 	case 7:
-		buf = BlockPool4M.Get()
+		return blockPool4M.Get().(*[Block4Mb]byte)[:]
 	case 3:
-		buf = BlockPool8M.Get()
+		return blockPool8M.Get().(*[Block8Mb]byte)[:]
+	default:
+		panic(fmt.Errorf("invalid block index %d", b))
 	}
-	return buf.([]byte)
 }
 
 func Put(buf []byte) {
 	// Safeguard: do not allow invalid buffers.
 	switch c := cap(buf); uint32(c) {
 	case Block64Kb:
-		BlockPool64K.Put(buf[:c])
+		blockPool64K.Put((*[Block64Kb]byte)(buf[:c]))
 	case Block256Kb:
-		BlockPool256K.Put(buf[:c])
+		blockPool256K.Put((*[Block256Kb]byte)(buf[:c]))
 	case Block1Mb:
-		BlockPool1M.Put(buf[:c])
+		blockPool1M.Put((*[Block1Mb]byte)(buf[:c]))
 	case Block4Mb:
-		BlockPool4M.Put(buf[:c])
+		blockPool4M.Put((*[Block4Mb]byte)(buf[:c]))
 	case Block8Mb:
-		BlockPool8M.Put(buf[:c])
+		blockPool8M.Put((*[Block8Mb]byte)(buf[:c]))
+	default:
+		panic(fmt.Errorf("invalid block size %d", c))
 	}
 }
 

internal/lz4block/blocks_test.go

@@ -0,0 +1,42 @@
+package lz4block
+
+import (
+	"testing"
+)
+
+func TestBlockGetPut(t *testing.T) {
+	expect := map[BlockSizeIndex]uint32{ // block index -> expected size
+		4: Block64Kb,
+		5: Block256Kb,
+		6: Block1Mb,
+		7: Block4Mb,
+		3: Block8Mb,
+	}
+	for idx, size := range expect {
+		buf := idx.Get()
+		if uint32(cap(buf)) != size {
+			t.Errorf("expected size %d for index %d, got %d", size, idx, cap(buf))
+		}
+		Put(buf) // ensure no panic
+	}
+}
+
+func TestBlockGetInvalid(t *testing.T) {
+	defer func() { recover() }() // swallow panic
+	_ = BlockSizeIndex(123).Get()
+	t.Fatalf("expected panic on bad Get")
+}
+
+func TestBlockPutInvalid(t *testing.T) {
+	defer func() { recover() }() // swallow panic
+	Put(make([]byte, 123))
+	t.Fatalf("expected panic on bad Put")
+}
+
+func BenchmarkGetPut(b *testing.B) {
+	const idx = BlockSizeIndex(4)
+	for i := 0; i < b.N; i++ {
+		buf := idx.Get()
+		Put(buf)
+	}
+}