mem: add Buffer.Slice()

Works just like Go's slice slicing.

Author: ash2k

mem/buffers.go

@@ -53,6 +53,10 @@ type Buffer interface {
 	Free()
 	// Len returns the Buffer's size.
 	Len() int
+	// Slice returns a new Buffer that is a view into this buffer's data
+	// from [start:end). The buffer is not modified. Panics if the buffer
+	// has been freed or if start/end are out of bounds.
+	Slice(start, end int) Buffer
 
 	split(n int) (left, right Buffer)
 	read(buf []byte) (int, Buffer)
@@ -180,6 +184,22 @@ func (b *buffer) Len() int {
 	return len(b.ReadOnlyData())
 }
 
+func (b *buffer) Slice(start, end int) Buffer {
+	if b.rootBuf == nil || b.rootBuf.refs.Add(1) <= 1 {
+		panic("Cannot slice freed buffer")
+	}
+
+	if start < 0 || end > len(b.data) || start > end {
+		panic(fmt.Sprintf("slice bounds out of range [%d:%d] with length %d", start, end, len(b.data)))
+	}
+
+	s := newBuffer()
+	s.data = b.data[start:end]
+	s.rootBuf = b.rootBuf
+	s.refs.Store(1)
+	return s
+}
+
 func (b *buffer) split(n int) (Buffer, Buffer) {
 	if b.rootBuf == nil || b.rootBuf.refs.Add(1) <= 1 {
 		panic("Cannot split freed buffer")
@@ -240,6 +260,13 @@ func (e emptyBuffer) Len() int {
 	return 0
 }
 
+func (e emptyBuffer) Slice(start, end int) Buffer {
+	if start != 0 || end != 0 {
+		panic(fmt.Sprintf("slice bounds out of range [%d:%d] with length 0", start, end))
+	}
+	return e
+}
+
 func (e emptyBuffer) split(int) (left, right Buffer) {
 	return e, e
 }
@@ -264,6 +291,9 @@ func (s SliceBuffer) Free() {}
 // Len is a noop implementation of Len.
 func (s SliceBuffer) Len() int { return len(s) }
 
+// Slice returns a new SliceBuffer that is a view into the receiver from [start:end).
+func (s SliceBuffer) Slice(start, end int) Buffer { return s[start:end] }
+
 func (s SliceBuffer) split(n int) (left, right Buffer) {
 	return s[:n], s[n:]
 }

mem/buffers_test.go

@@ -20,6 +20,7 @@ package mem_test
 
 import (
 	"bytes"
+	"fmt"
 	"testing"
 
 	"google.golang.org/grpc/internal"
@@ -265,21 +266,15 @@ func (s) TestBuffer_Split(t *testing.T) {
 		}
 		freed = true
 	}))
-	checkBufData := func(b mem.Buffer, expected []byte) {
-		t.Helper()
-		if !bytes.Equal(b.ReadOnlyData(), expected) {
-			t.Fatalf("Buffer did not contain expected data %v, got %v", expected, b.ReadOnlyData())
-		}
-	}
 
 	buf, split1 := mem.SplitUnsafe(buf, 2)
-	checkBufData(buf, data[:2])
-	checkBufData(split1, data[2:])
+	checkBufData(t, buf, data[:2])
+	checkBufData(t, split1, data[2:])
 
 	// Check that splitting the buffer more than once works as intended.
 	split1, split2 := mem.SplitUnsafe(split1, 1)
-	checkBufData(split1, data[2:3])
-	checkBufData(split2, data[3:])
+	checkBufData(t, split1, data[2:3])
+	checkBufData(t, split2, data[3:])
 
 	// If any of the following frees actually free the buffer, the test will fail.
 	buf.Free()
@@ -293,6 +288,158 @@ func (s) TestBuffer_Split(t *testing.T) {
 	}
 }
 
+func (s) TestBuffer_Slice(t *testing.T) {
+	ready := false
+	freed := false
+	data := []byte{1, 2, 3, 4}
+	buf := mem.NewBuffer(&data, poolFunc(func(*[]byte) {
+		if !ready {
+			t.Fatalf("Freed too early")
+		}
+		freed = true
+	}))
+
+	// Slice the buffer and verify the data.
+	slice1 := buf.Slice(1, 3)
+	checkBufData(t, slice1, data[1:3])
+
+	// Verify the original buffer is not modified.
+	checkBufData(t, buf, data)
+
+	// Slice the slice.
+	slice2 := slice1.Slice(0, 1)
+	checkBufData(t, slice2, data[1:2])
+
+	// Free original and first slice — root should not be freed yet.
+	buf.Free()
+	slice1.Free()
+
+	// The last slice keeps the root alive.
+	checkBufData(t, slice2, data[1:2])
+
+	ready = true
+	slice2.Free()
+
+	if !freed {
+		t.Fatalf("Buffer never freed")
+	}
+}
+
+func (s) TestBuffer_SliceAfterFree(t *testing.T) {
+	buf := newBuffer([]byte("abcd"), mem.NopBufferPool{})
+	buf.Free()
+	defer checkForPanic(t, "Cannot slice freed buffer")
+	buf.Slice(0, 2)
+}
+
+func (s) TestBuffer_SliceBasic(t *testing.T) {
+	type sliceCase struct {
+		start, end int
+		want       []byte
+	}
+	cases := []sliceCase{
+		{1, 3, []byte{2, 3}},
+		{0, 4, []byte{1, 2, 3, 4}},
+		{0, 0, []byte{}},
+		{4, 4, []byte{}},
+	}
+	ctors := map[string]func() mem.Buffer{
+		"pooled": newPooledBuffer,
+		"slice":  newSliceBuf,
+	}
+	for name, newBuf := range ctors {
+		for _, tc := range cases {
+			t.Run(fmt.Sprintf("%s[%d:%d]", name, tc.start, tc.end), func(t *testing.T) {
+				buf := newBuf()
+				got := buf.Slice(tc.start, tc.end)
+				checkBufData(t, got, tc.want)
+			})
+		}
+		t.Run(name+" subslice", func(t *testing.T) {
+			buf := newBuf()
+			slice := buf.Slice(1, 3)
+			slice2 := slice.Slice(0, 1)
+			checkBufData(t, slice2, []byte{2})
+		})
+	}
+	t.Run("empty buffer", func(t *testing.T) {
+		buf := newEmptyBuf()
+		got := buf.Slice(0, 0)
+		checkBufData(t, got, nil)
+	})
+}
+
+func (s) TestBuffer_SliceBoundsCheck(t *testing.T) {
+	type panicCase struct {
+		name       string
+		start, end int
+	}
+	nonEmptyCases := []panicCase{
+		{"end_out_of_bounds", 0, 5},
+		{"start_negative", -1, 3},
+		{"start_greater_than_end", 3, 0},
+	}
+	tests := []struct {
+		name       string
+		buf        func() mem.Buffer
+		panicCases []panicCase
+	}{
+		{
+			name:       "buffer",
+			buf:        newPooledBuffer,
+			panicCases: nonEmptyCases,
+		},
+		{
+			name:       "SliceBuffer",
+			buf:        newSliceBuf,
+			panicCases: nonEmptyCases,
+		},
+		{
+			name: "emptyBuffer",
+			buf:  newEmptyBuf,
+			panicCases: []panicCase{
+				{"end_out_of_bounds", 0, 1},
+				{"start_negative", -1, 0},
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			for _, pc := range tt.panicCases {
+				t.Run(pc.name, func(t *testing.T) {
+					buf := tt.buf()
+					defer func() {
+						if recover() == nil {
+							t.Fatalf("Slice(%d, %d) did not panic", pc.start, pc.end)
+						}
+					}()
+					buf.Slice(pc.start, pc.end)
+				})
+			}
+		})
+	}
+}
+
+func newPooledBuffer() mem.Buffer {
+	return newBuffer([]byte{1, 2, 3, 4}, mem.NopBufferPool{})
+}
+
+func newSliceBuf() mem.Buffer {
+	return mem.SliceBuffer{1, 2, 3, 4}
+}
+
+func newEmptyBuf() mem.Buffer {
+	var bs mem.BufferSlice
+	return bs.MaterializeToBuffer(mem.NopBufferPool{})
+}
+
+func checkBufData(t *testing.T, b mem.Buffer, expected []byte) {
+	t.Helper()
+	if !bytes.Equal(b.ReadOnlyData(), expected) {
+		t.Fatalf("Buffer did not contain expected data %v, got %v", expected, b.ReadOnlyData())
+	}
+}
+
 func checkForPanic(t *testing.T, wantErr string) {
 	t.Helper()
 	r := recover()