Add buf2 package

Author: nekohasekai

common/buf2/buffer.go

@@ -0,0 +1,166 @@
+package buf2
+
+import (
+	"strconv"
+	"sync"
+)
+
+var bufferPool = sync.Pool{
+	New: func() any {
+		return &Buffer{}
+	},
+}
+
+type Buffer struct {
+	start int
+	end   int
+	chunk *Chunk
+}
+
+func New(capacity int) *Buffer {
+	chunk := NewChunk(capacity)
+	buffer := bufferPool.Get().(*Buffer)
+	*buffer = Buffer{chunk: chunk}
+	return buffer
+}
+
+func From(data []byte) *Buffer {
+	buffer := New(len(data))
+	buffer.Write(data)
+	return buffer
+}
+
+func (b *Buffer) Cut(start int, end int) *Buffer {
+	if b == nil {
+		panic("cannot cut a nil view")
+	}
+	b.chunk.IncRef()
+	newBuffer := bufferPool.Get().(*Buffer)
+	newBuffer.chunk = b.chunk
+	newBuffer.start = start
+	newBuffer.end = end
+	return newBuffer
+}
+
+func (b *Buffer) Release() {
+	if b == nil {
+		panic("cannot release a nil view")
+	}
+	b.chunk.DecRef()
+	*b = Buffer{}
+	bufferPool.Put(b)
+}
+
+func (b *Buffer) Reset() {
+	if b == nil {
+		panic("cannot reset a nil view")
+	}
+	b.start = 0
+	b.end = 0
+}
+
+func (b *Buffer) Full() bool {
+	return b == nil || b.end == len(b.chunk.data)
+}
+
+func (b *Buffer) Clone() *Buffer {
+	if b == nil {
+		panic("cannot clone a nil view")
+	}
+	b.chunk.IncRef()
+	newBuffer := bufferPool.Get().(*Buffer)
+	newBuffer.chunk = b.chunk
+	newBuffer.start = b.start
+	newBuffer.end = b.end
+	return newBuffer
+}
+
+func (b *Buffer) Capacity() int {
+	if b == nil {
+		return 0
+	}
+	return len(b.chunk.data)
+}
+
+func (b *Buffer) Size() int {
+	if b == nil {
+		return 0
+	}
+	return b.end - b.start
+}
+
+func (b *Buffer) AsSlice() []byte {
+	if b.Size() == 0 {
+		return nil
+	}
+	return b.chunk.data[b.start:b.end]
+}
+
+func (b *Buffer) ToSlice() []byte {
+	if b.Size() == 0 {
+		return nil
+	}
+	s := make([]byte, b.Size())
+	copy(s, b.AsSlice())
+	return s
+}
+
+func (b *Buffer) AvailableSize() int {
+	if b == nil {
+		return 0
+	}
+	return len(b.chunk.data) - b.end
+}
+
+func (b *Buffer) AvailableSlice() []byte {
+	if b == nil {
+		panic("cannot get available slice from a nil view")
+	}
+	return b.availableSlice()
+}
+
+func (b *Buffer) ExtendHeader(n int) []byte {
+	if b == nil {
+		panic("cannot extend a nil view")
+	}
+	if b.start < n {
+		panic("buffer overflow: cap " + strconv.Itoa(b.Capacity()) + ",start " + strconv.Itoa(b.start) + ", need " + strconv.Itoa(n))
+	}
+	if b.sharesChunk() {
+		defer b.chunk.DecRef()
+		b.chunk = b.chunk.Clone()
+	}
+	b.start -= n
+	return b.chunk.data[b.start : b.start+n]
+}
+
+func (b *Buffer) Extend(n int) []byte {
+	if b == nil {
+		panic("cannot extend a nil view")
+	}
+	newEnd := b.end + n
+	if newEnd > b.Capacity() {
+		panic("buffer overflow: cap " + strconv.Itoa(b.Capacity()) + ",end " + strconv.Itoa(b.end) + ", need " + strconv.Itoa(n))
+	}
+	if b.sharesChunk() {
+		defer b.chunk.DecRef()
+		b.chunk = b.chunk.Clone()
+	}
+	newSlice := b.chunk.data[b.end : b.end+n]
+	b.end = newEnd
+	return newSlice
+}
+
+func (b *Buffer) Advance(from int) {
+	if b == nil {
+		panic("cannot advance a nil view")
+	}
+	b.start += from
+}
+
+func (b *Buffer) Truncate(to int) {
+	if b == nil {
+		panic("cannot truncate a nil view")
+	}
+	b.end = b.start + to
+}

common/buf2/buffer_rw.go

@@ -0,0 +1,140 @@
+package buf2
+
+import (
+	"crypto/rand"
+	"fmt"
+	"io"
+
+	"github.com/sagernet/sing/common"
+)
+
+func (b *Buffer) Read(p []byte) (int, error) {
+	if len(p) == 0 {
+		return 0, nil
+	}
+	if b.Size() == 0 {
+		return 0, io.EOF
+	}
+	n := copy(p, b.AsSlice())
+	b.Advance(n)
+	return n, nil
+}
+
+func (b *Buffer) Write(p []byte) (n int, err error) {
+	if b == nil {
+		panic("cannot write to a nil view")
+	}
+	if len(p) == 0 {
+		return
+	}
+	if b.Full() {
+		return 0, io.ErrShortBuffer
+	}
+	if b.sharesChunk() {
+		defer b.chunk.DecRef()
+		b.chunk = b.chunk.Clone()
+	}
+	n = copy(b.chunk.data[b.end:], p)
+	b.end += n
+	return
+}
+
+func (b *Buffer) ReadAt(p []byte, off int) (int, error) {
+	if off < 0 || off > b.Size() {
+		return 0, fmt.Errorf("ReadAt(): offset out of bounds: want 0 < off < %d, got off=%d", b.Size(), off)
+	}
+	n := copy(p, b.AsSlice()[off:])
+	return n, nil
+}
+
+func (b *Buffer) WriteAt(p []byte, off int) (int, error) {
+	if b == nil {
+		panic("cannot write to a nil view")
+	}
+	if off < 0 || off > b.Size() {
+		return 0, fmt.Errorf("write offset out of bounds: want 0 < off < %d, got off=%d", b.Size(), off)
+	}
+	if b.sharesChunk() {
+		defer b.chunk.DecRef()
+		b.chunk = b.chunk.Clone()
+	}
+	n := copy(b.AsSlice()[off:], p)
+	if n < len(p) {
+		return n, io.ErrShortWrite
+	}
+	return n, nil
+}
+
+func (b *Buffer) ReadByte() (byte, error) {
+	if b.Size() == 0 {
+		return 0, io.EOF
+	}
+	p := b.AsSlice()[0]
+	b.start++
+	return p, nil
+}
+
+func (b *Buffer) WriteByte(d byte) error {
+	if b == nil {
+		panic("cannot write to a nil view")
+	} else if b.Full() {
+		return io.ErrShortBuffer
+	}
+	if b.sharesChunk() {
+		defer b.chunk.DecRef()
+		b.chunk = b.chunk.Clone()
+	}
+	b.chunk.data[b.end] = d
+	b.end++
+	return nil
+}
+
+func (b *Buffer) WriteTo(w io.Writer) (n int64, err error) {
+	if b.Size() > 0 {
+		sz := b.Size()
+		m, e := w.Write(b.AsSlice())
+		b.Advance(m)
+		n = int64(m)
+		if e != nil {
+			return n, e
+		}
+		if m != sz {
+			return n, io.ErrShortWrite
+		}
+	}
+	return n, nil
+}
+
+func (b *Buffer) ReadFrom0(r io.Reader) (n int64, err error) {
+	if b == nil {
+		panic("cannot write to a nil view")
+	}
+	if b.AvailableSize() == 0 {
+		return 0, io.ErrShortBuffer
+	}
+	nInt, err := r.Read(b.availableSlice())
+	b.end += nInt
+	n += int64(nInt)
+	if err == io.EOF {
+		return n, nil
+	}
+	return
+}
+
+func (b *Buffer) WriteRandom(size int) []byte {
+	buffer := b.Extend(size)
+	common.Must1(io.ReadFull(rand.Reader, buffer))
+	return buffer
+}
+
+func (b *Buffer) sharesChunk() bool {
+	return b.chunk.references.Load() > 1
+}
+
+func (b *Buffer) availableSlice() []byte {
+	if b.sharesChunk() {
+		defer b.chunk.DecRef()
+		b.chunk = b.chunk.Clone()
+	}
+	return b.chunk.data[b.end:]
+}

common/buf2/chunk.go

@@ -0,0 +1,67 @@
+package buf2
+
+import (
+	"github.com/sagernet/sing/common/atomic"
+	F "github.com/sagernet/sing/common/format"
+)
+
+type Chunk struct {
+	data       []byte
+	references atomic.Int64
+}
+
+func NewChunk(size int) *Chunk {
+	var chunk *Chunk
+	if size > MaxChunkSize {
+		chunk = &Chunk{data: make([]byte, size)}
+	} else {
+		chunk = getChunkPool(size).Get().(*Chunk)
+		for i := range chunk.data {
+			chunk.data[i] = 0
+		}
+	}
+	chunk.references.Store(1)
+	return chunk
+}
+
+//go:nosplit
+func (c *Chunk) IncRef() {
+	v := c.references.Add(1)
+	if v <= 1 {
+		panic(F.ToString("Incrementing non-positive count ", v, " on buffer"))
+	}
+}
+
+//go:nosplit
+func (c *Chunk) TryIncRef() bool {
+	const speculativeRef = 1 << 32
+	if v := c.references.Add(speculativeRef); int32(v) == 0 {
+		c.references.Add(-speculativeRef)
+		return false
+	}
+	c.references.Add(-speculativeRef + 1)
+	return true
+}
+
+func (c *Chunk) ReadRefs() int64 {
+	return c.references.Load()
+}
+
+//go:nosplit
+func (c *Chunk) DecRef() {
+	v := c.references.Add(-1)
+	switch {
+	case v < 0:
+		panic(F.ToString("Decrementing non-positive ref count ", v, " on buffer"))
+	case v == 0:
+		if len(c.data) < MaxChunkSize {
+			getChunkPool(len(c.data)).Put(c)
+		}
+	}
+}
+
+func (c *Chunk) Clone() *Chunk {
+	newChunk := NewChunk(len(c.data))
+	copy(newChunk.data, c.data)
+	return newChunk
+}