Add nil receiver and io methods to bufferv2.

Allowing for nil receivers in read-like operations lets views be treated more like slices, which is easier for the programmer generally. io methods are useful for copy operations to avoid unnecessary allocations. PiperOrigin-RevId: 457086370
2026-05-22 17:12:49 -07:00 · 2022-06-24 13:49:22 -07:00
parent ffabadf010
commit efdc289a96
4 changed files with 216 additions and 41 deletions
@@ -294,13 +294,13 @@ func (b *Buffer) prependOwned(v *View) {
 }

 // PullUp makes the specified range contiguous and returns the backing memory.
-func (b *Buffer) PullUp(offset, length int) (*View, bool) {
+func (b *Buffer) PullUp(offset, length int) (View, bool) {
 	if length == 0 {
-		return nil, true
+		return View{}, true
 	}
 	tgt := Range{begin: offset, end: offset + length}
 	if tgt.Intersect(Range{end: int(b.size)}).Len() != length {
-		return nil, false
+		return View{}, false
 	}

 	curr := Range{}
@@ -312,12 +312,13 @@ func (b *Buffer) PullUp(offset, length int) (*View, bool) {
 		if x := curr.Intersect(tgt); x.Len() == tgt.Len() {
 			// buf covers the whole requested target range.
 			sub := x.Offset(-curr.begin)
-			new := viewPool.Get().(*View)
-			new.read = sub.begin
-			new.write = sub.end
 			// Don't increment the reference count of the underlying chunk. Views
 			// returned by PullUp are explicitly unowned and read only
-			new.chunk = v.chunk
+			new := View{
+				read:  v.read + sub.begin,
+				write: v.read + sub.end,
+				chunk: v.chunk,
+			}
 			return new, true
 		} else if x.Len() > 0 {
 			// buf is pointing at the starting buffer we want to merge.
@@ -357,10 +358,11 @@ func (b *Buffer) PullUp(offset, length int) (*View, bool) {
 	b.removeView(v)

 	r := tgt.Offset(-curr.begin)
-	pulled := viewPool.Get().(*View)
-	pulled.read = r.begin
-	pulled.write = r.end
-	pulled.chunk = merged.chunk
+	pulled := View{
+		read:  r.begin,
+		write: r.end,
+		chunk: merged.chunk,
+	}
 	return pulled, true
 }

@@ -418,11 +420,11 @@ func (b *Buffer) Apply(fn func(*View)) {
 // outside of b is ignored.
 func (b *Buffer) SubApply(offset, length int, fn func(*View)) {
 	for v := b.data.Front(); length > 0 && v != nil; v = v.Next() {
-		d := v.Clone()
-		if offset >= d.Size() {
-			offset -= d.Size()
+		if offset >= v.Size() {
+			offset -= v.Size()
 			continue
 		}
+		d := v.Clone()
 		if offset > 0 {
 			d.TrimFront(offset)
 			offset = 0
@@ -612,7 +612,7 @@ func TestBufferPullUp(t *testing.T) {

 			got, gotOk := b.PullUp(tc.offset, tc.length)
 			want, wantOk := []byte(tc.output), !tc.failed
-			if gotOk == wantOk && got == nil && len(want) == 0 {
+			if gotOk == wantOk && got.Size() == 0 && len(want) == 0 {
 				return
 			}
 			if gotOk != wantOk || !bytes.Equal(got.AsSlice(), want) {
@@ -630,6 +630,25 @@ func TestBufferPullUp(t *testing.T) {
 	}
 }

+func TestPullUpModifiedViews(t *testing.T) {
+	var b Buffer
+	defer b.Release()
+	for _, s := range []string{"abcdef", "123456", "ghijkl"} {
+		v := NewViewWithData([]byte(s))
+		v.TrimFront(3)
+		b.appendOwned(v)
+	}
+
+	v, ok := b.PullUp(3, 3)
+	if !ok {
+		t.Errorf("PullUp failed: want ok=true, got ok=false")
+	}
+	want := []byte("456")
+	if !bytes.Equal(v.AsSlice(), want) {
+		t.Errorf("PullUp failed: want %v, got %v", want, v.AsSlice())
+	}
+}
+
 func TestBufferClone(t *testing.T) {
 	const (
 		originalSize  = 90
@@ -21,6 +21,10 @@ import (
 	"gvisor.dev/gvisor/pkg/sync"
 )

+// ReadSize is the default amount that a View's size is increased by when an
+// io.Reader has more data than a View can hold during calls to ReadFrom.
+const ReadSize = 512
+
 var viewPool = sync.Pool{
 	New: func() interface{} {
 		return &View{}
@@ -42,8 +46,6 @@ var viewPool = sync.Pool{
 // must use Write/WriteAt/CopyIn to modify the underlying View. This preserves
 // the safety guarantees of copy-on-write.
 type View struct {
-	sync.NoCopy
-
 	viewEntry
 	read  int
 	write int
@@ -84,6 +86,9 @@ func NewViewWithData(data []byte) *View {
 // The caller must own the View to call Clone. It is not safe to call Clone
 // on a borrowed or shared View because it can race with other View methods.
 func (v *View) Clone() *View {
+	if v == nil {
+		panic("cannot clone a nil view")
+	}
 	v.chunk.IncRef()
 	newV := viewPool.Get().(*View)
 	newV.chunk = v.chunk
@@ -94,6 +99,9 @@ func (v *View) Clone() *View {

 // Release releases the chunk held by v and returns v to the pool.
 func (v *View) Release() {
+	if v == nil {
+		panic("cannot release a nil view")
+	}
 	v.chunk.DecRef()
 	*v = View{}
 	viewPool.Put(v)
@@ -112,16 +120,25 @@ func (v *View) Full() bool {

 // Capacity returns the total size of this view's chunk.
 func (v *View) Capacity() int {
+	if v == nil {
+		return 0
+	}
 	return len(v.chunk.data)
 }

 // Size returns the size of data written to the view.
 func (v *View) Size() int {
+	if v == nil {
+		return 0
+	}
 	return v.write - v.read
 }

 // TrimFront advances the read index by the given amount.
 func (v *View) TrimFront(n int) {
+	if v.read+n > v.write {
+		panic("cannot trim past the end of a view")
+	}
 	v.read += n
 }

@@ -135,6 +152,9 @@ func (v *View) AsSlice() []byte {

 // AvailableSize returns the number of bytes available for writing.
 func (v *View) AvailableSize() int {
+	if v == nil {
+		return 0
+	}
 	return len(v.chunk.data) - v.write
 }

@@ -153,6 +173,26 @@ func (v *View) Read(p []byte) (int, error) {
 	return n, nil
 }

+// WriteTo writes data to w until the view is empty or an error occurs. The
+// return value n is the number of bytes written.
+//
+// WriteTo implements the io.WriterTo interface.
+func (v *View) WriteTo(w io.Writer) (n int64, err error) {
+	if v.Size() > 0 {
+		sz := v.Size()
+		m, e := w.Write(v.AsSlice())
+		v.TrimFront(m)
+		n = int64(m)
+		if e != nil {
+			return n, e
+		}
+		if m != sz {
+			return n, io.ErrShortWrite
+		}
+	}
+	return n, nil
+}
+
 // ReadAt reads data to the p starting at offset.
 //
 // Implements the io.ReaderAt interface.
@@ -170,24 +210,62 @@ func (v *View) ReadAt(p []byte, off int) (int, error) {
 //
 // Implements the io.Writer interface.
 func (v *View) Write(p []byte) (int, error) {
-	if v.sharesChunk() {
+	if v == nil {
+		panic("cannot write to a nil view")
+	}
+	if v.AvailableSize() < len(p) {
+		v.growCap(len(p) - v.AvailableSize())
+	} else if v.sharesChunk() {
 		defer v.chunk.DecRef()
 		v.chunk = v.chunk.Clone()
 	}
 	n := copy(v.chunk.data[v.write:], p)
 	v.write += n
 	if n < len(p) {
-		return n, fmt.Errorf("could not finish write: want len(p) <= v.AvailableSize(), got len(p)=%d, v.AvailableSize()=%d", len(p), v.AvailableSize())
+		return n, io.ErrShortWrite
 	}
 	return n, nil
 }

+// ReadFrom reads data from r until EOF and appends it to the buffer, growing
+// the buffer as needed. The return value n is the number of bytes read. Any
+// error except io.EOF encountered during the read is also returned.
+//
+// ReadFrom implements the io.ReaderFrom interface.
+func (v *View) ReadFrom(r io.Reader) (n int64, err error) {
+	if v == nil {
+		panic("cannot write to a nil view")
+	}
+	if v.sharesChunk() {
+		defer v.chunk.DecRef()
+		v.chunk = v.chunk.Clone()
+	}
+	for {
+		if v.AvailableSize() == 0 {
+			v.growCap(ReadSize)
+		}
+		m, e := r.Read(v.availableSlice())
+		v.write += m
+		n += int64(m)
+
+		if e == io.EOF {
+			return n, nil
+		}
+		if e != nil {
+			return n, e
+		}
+	}
+}
+
 // WriteAt writes data to the views's chunk starting at start. If the
 // view's chunk has a reference count greater than 1, the chunk is copied first
 // and then written to.
 //
 // Implements the io.WriterAt interface.
 func (v *View) WriteAt(p []byte, off int) (int, error) {
+	if v == nil {
+		panic("cannot write to a nil view")
+	}
 	if off < 0 || off > v.Size() {
 		return 0, fmt.Errorf("write offset out of bounds: want 0 < off < %d, got off=%d", v.Size(), off)
 	}
@@ -197,22 +275,43 @@ func (v *View) WriteAt(p []byte, off int) (int, error) {
 	}
 	n := copy(v.AsSlice()[off:], p)
 	if n < len(p) {
-		return n, fmt.Errorf("could not finish write: want off + len(p) < v.Capacity(), got off=%d, len(p)=%d ,v.Size() = %d", off, len(p), v.Size())
+		return n, io.ErrShortWrite
 	}
 	return n, nil
 }

-// Grow advances the write index by the given amount.
+// Grow increases the size of the view. If the new size is greater than the
+// view's current capacity, Grow will reallocate the view with an increased
+// capacity.
 func (v *View) Grow(n int) {
-	if n+v.write > v.Capacity() {
-		panic("cannot grow view past capacity")
+	if v == nil {
+		panic("cannot grow a nil view")
+	}
+	if v.write+n > v.Capacity() {
+		v.growCap(n)
 	}
 	v.write += n
 }

+// growCap increases the capacity of the view by at least n.
+func (v *View) growCap(n int) {
+	if v == nil {
+		panic("cannot grow a nil view")
+	}
+	defer v.chunk.DecRef()
+	old := v.AsSlice()
+	v.chunk = newChunk(v.Capacity() + n)
+	copy(v.chunk.data, old)
+	v.read = 0
+	v.write = len(old)
+}
+
 // CapLength caps the length of the view's read slice to n. If n > v.Size(),
 // the function is a no-op.
 func (v *View) CapLength(n int) {
+	if v == nil {
+		panic("cannot resize a nil view")
+	}
 	if n < 0 {
 		panic("n must be >= 0")
 	}
@@ -15,6 +15,7 @@
 package buffer

 import (
+	"bytes"
 	"math/rand"
 	"testing"

@@ -80,7 +81,7 @@ func TestWrite(t *testing.T) {
 	for _, tc := range []struct {
 		name      string
 		view      *View
-		initData  []byte
+		initSize  int
 		writeSize int
 	}{
 		{
@@ -91,45 +92,41 @@ func TestWrite(t *testing.T) {
 		{
 			name:      "full view",
 			view:      NewView(100),
-			initData:  make([]byte, 100),
+			initSize:  100,
 			writeSize: 50,
 		},
 		{
 			name:      "full write to partially full view",
 			view:      NewView(100),
-			initData:  make([]byte, 20),
+			initSize:  20,
 			writeSize: 50,
 		},
 		{
 			name:      "partial write to partially full view",
 			view:      NewView(100),
-			initData:  make([]byte, 80),
+			initSize:  80,
 			writeSize: 50,
 		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
-			tc.view.Write(tc.initData)
+			tc.view.Grow(tc.initSize)
 			defer tc.view.Release()
-			origWriteSize := tc.view.AvailableSize()
-
-			var orig []byte
-			orig = append(orig, tc.view.AsSlice()...)
+			orig := append([]byte(nil), tc.view.AsSlice()...)
 			toWrite := make([]byte, tc.writeSize)
 			rand.Read(toWrite)

-			n, _ := tc.view.Write(toWrite)
-
-			if n > origWriteSize {
-				t.Errorf("got tc.view.Write() = %d, want <=%d", n, origWriteSize)
+			n, err := tc.view.Write(toWrite)
+			if err != nil {
+				t.Errorf("Write failed: %s", err)
 			}
-			if tc.writeSize > origWriteSize {
-				toWrite = toWrite[:origWriteSize]
+			if n != tc.writeSize {
+				t.Errorf("got n=%d, want %d", n, tc.writeSize)
 			}
-			if tc.view.AvailableSize() != tc.view.Capacity()-(len(toWrite)+len(orig)) {
-				t.Errorf("got tc.view.WriteSize() = %d, want %d", tc.view.AvailableSize(), tc.view.Capacity()-(len(toWrite)+len(orig)))
+			if tc.view.Size() != len(orig)+tc.writeSize {
+				t.Errorf("got Size()=%d, want %d", tc.view.Size(), len(orig)+tc.writeSize)
 			}
 			if !cmp.Equal(tc.view.AsSlice(), append(orig, toWrite...)) {
-				t.Errorf("got tc.view.ReadSlice() = %d, want %d", tc.view.AsSlice(), toWrite)
+				t.Errorf("got tc.view.AsSlice() = %d, want %d", tc.view.AsSlice(), toWrite)
 			}
 		})
 	}
@@ -172,3 +169,61 @@ func TestWriteAt(t *testing.T) {
 		t.Errorf("got v.AsSlice()[:off] = %v, want %v", v.AsSlice()[:off], orig.AsSlice()[:off])
 	}
 }
+
+func TestWriteTo(t *testing.T) {
+	writeToSize := 100
+	v := NewViewSize(writeToSize)
+	defer v.Release()
+
+	w := bytes.NewBuffer(make([]byte, 100))
+
+	n, err := v.WriteTo(w)
+	if err != nil {
+		t.Errorf("WriteTo failed: %s", err)
+	}
+	if n != int64(writeToSize) {
+		t.Errorf("got n=%d, want 100", n)
+	}
+	if v.Size() != 0 {
+		t.Errorf("got v.Size()=%d, want 0", v.Size())
+	}
+}
+
+func TestReadFrom(t *testing.T) {
+	for _, tc := range []struct {
+		name string
+		data []byte
+		view *View
+	}{
+		{
+			name: "basic",
+			data: []byte{1, 2, 3},
+			view: NewView(10),
+		},
+		{
+			name: "requires grow",
+			data: []byte{4, 5, 6},
+			view: NewViewSize(63),
+		},
+	} {
+		defer tc.view.Release()
+		clone := tc.view.Clone()
+		defer clone.Release()
+		r := bytes.NewReader(tc.data)
+
+		n, err := tc.view.ReadFrom(r)
+
+		if err != nil {
+			t.Errorf("v.ReadFrom failed: %s", err)
+		}
+		if int(n) != len(tc.data) {
+			t.Errorf("v.ReadFrom failed: want n=%d, got %d", len(tc.data), n)
+		}
+		if tc.view.Size() == clone.Size() {
+			t.Errorf("expected clone.Size() != v.Size(), got match")
+		}
+		if !bytes.Equal(tc.view.AsSlice(), append(clone.AsSlice(), tc.data...)) {
+			t.Errorf("v.ReadFrom failed: want %v, got %v", tc.data, tc.view.AsSlice())
+		}
+	}
+}