Add ByteBuffer.ShrinkTo/ShrinkBy

byte_buffer.go

@@ -361,3 +361,22 @@ func (b *ByteBuffer) ClaimFixed(n int) (claimed []byte) {
 	}
 	return
 }
+
+// ShrinkBy shrinks the write area by at most `n` bytes.
+func (b *ByteBuffer) ShrinkBy(n int) int {
+	if n <= 0 {
+		return 0
+	}
+
+	if length := b.WriteLen(); n > length {
+		n = length
+	}
+	b.wi -= n
+	b.data = b.data[:b.wi]
+	return n
+}
+
+// ShrinkTo shrinks the write to contain min(n, WriteLen()) bytes.
+func (b *ByteBuffer) ShrinkTo(n int) (shrunkBy int) {
+	return b.ShrinkBy(b.WriteLen() - n)
+}

byte_buffer_test.go

@@ -4,6 +4,7 @@ import (
 	"bufio"
 	"bytes"
 	"errors"
+	"fmt"
 	"math/rand"
 	"testing"
 	"time"
@@ -568,6 +569,75 @@ func TestByteBufferClaimFixed(t *testing.T) {
 	}
 }
 
+func TestByteBufferShrinkBy(t *testing.T) {
+	{
+		b := NewByteBuffer()
+		b.Write([]byte("hello"))
+		if b.WriteLen() != 5 {
+			t.Fatal("wrong write length")
+		}
+
+		if n := b.ShrinkBy(1); n != 1 && b.WriteLen() != 4 {
+			t.Fatal("wrong shrink by")
+		}
+
+		if n := b.ShrinkBy(100); n != 4 && b.WriteLen() != 0 {
+			t.Fatal("wrong shrink by")
+		}
+
+		if n := b.ShrinkBy(-100); n != 0 && b.WriteLen() != 0 {
+			t.Fatal("wrong shrink by")
+		}
+	}
+	{
+		b := NewByteBuffer()
+		b.Write([]byte("hello"))
+		if n := b.ShrinkBy(2); n != 2 && b.WriteLen() != 3 {
+			t.Fatal("wrong shrink to")
+		}
+		b.Commit(100)
+		if string(b.Data()) != "hel" {
+			t.Fatal("wrong shrink by")
+		}
+	}
+}
+
+func TestByteBufferShrinkTo(t *testing.T) {
+	{
+		b := NewByteBuffer()
+		b.Write([]byte("hello"))
+		if b.WriteLen() != 5 {
+			t.Fatal("wrong write length")
+		}
+
+		if n := b.ShrinkTo(4); n != 1 && b.WriteLen() != 4 {
+			t.Fatal("wrong shrink to")
+		}
+
+		if n := b.ShrinkTo(0); n != 4 && b.WriteLen() != 0 {
+			fmt.Println(n, b.WriteLen())
+			t.Fatal("wrong shrink to")
+		}
+	}
+	{
+		b := NewByteBuffer()
+		if n := b.ShrinkTo(1); n != 0 {
+			t.Fatal("wrong shrink to")
+		}
+	}
+	{
+		b := NewByteBuffer()
+		b.Write([]byte("hello"))
+		if n := b.ShrinkTo(3); n != 2 && b.WriteLen() != 3 {
+			t.Fatal("wrong shrink to")
+		}
+		b.Commit(100)
+		if string(b.Data()) != "hel" {
+			t.Fatal("wrong shrink to")
+		}
+	}
+}
+
 func BenchmarkByteBuffer(b *testing.B) {
 	var letters = []byte("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")