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

@@ -265,21 +265,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 +287,69 @@ 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_SliceBoundsCheck(t *testing.T) {
+	buf := newBuffer([]byte("abcd"), mem.NopBufferPool{})
+	defer buf.Free()
+	defer func() {
+		r := recover()
+		if r == nil {
+			t.Fatalf("Out of bounds slice did not panic")
+		}
+	}()
+	buf.Slice(0, 5)
+}
+
+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()