/*
 *
 * Copyright 2024 gRPC authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */


// Package mem provides utilities that facilitate memory reuse in byte slices
// that are used as buffers.
//
// # Experimental
//
// Notice: All APIs in this package are EXPERIMENTAL and may be changed or
// removed in a later release.
package mem

import (
	
	
	
)

// A Buffer represents a reference counted piece of data (in bytes) that can be
// acquired by a call to NewBuffer() or Copy(). A reference to a Buffer may be
// released by calling Free(), which invokes the free function given at creation
// only after all references are released.
//
// Note that a Buffer is not safe for concurrent access and instead each
// goroutine should use its own reference to the data, which can be acquired via
// a call to Ref().
//
// Attempts to access the underlying data after releasing the reference to the
// Buffer will panic.
type Buffer interface {
	// ReadOnlyData returns the underlying byte slice. Note that it is undefined
	// behavior to modify the contents of this slice in any way.
	ReadOnlyData() []byte
	// Ref increases the reference counter for this Buffer.
	Ref()
	// Free decrements this Buffer's reference counter and frees the underlying
	// byte slice if the counter reaches 0 as a result of this call.
	Free()
	// Len returns the Buffer's size.
	Len() int

	split(n int) (left, right Buffer)
	read(buf []byte) (int, Buffer)
}

var (
	bufferPoolingThreshold = 1 << 10

	bufferObjectPool = sync.Pool{New: func() any { return new(buffer) }}
)

// IsBelowBufferPoolingThreshold returns true if the given size is less than or
// equal to the threshold for buffer pooling. This is used to determine whether
// to pool buffers or allocate them directly.
func ( int) bool {
	return  <= bufferPoolingThreshold
}

type buffer struct {
	refs atomic.Int32
	data []byte

	// rootBuf is the buffer responsible for returning origData to the pool
	// once the reference count drops to 0.
	//
	// When a buffer is split, the new buffer inherits the rootBuf of the
	// original and increments the root's reference count. For the
	// initial buffer (the root), this field points to itself.
	rootBuf *buffer

	// The following fields are only set for root buffers.
	origData *[]byte
	pool     BufferPool
}

func newBuffer() *buffer {
	return bufferObjectPool.Get().(*buffer)
}

// NewBuffer creates a new Buffer from the given data, initializing the reference
// counter to 1. The data will then be returned to the given pool when all
// references to the returned Buffer are released. As a special case to avoid
// additional allocations, if the given buffer pool is nil, the returned buffer
// will be a "no-op" Buffer where invoking Buffer.Free() does nothing and the
// underlying data is never freed.
//
// Note that the backing array of the given data is not copied.
func ( *[]byte,  BufferPool) Buffer {
	// Use the buffer's capacity instead of the length, otherwise buffers may
	// not be reused under certain conditions. For example, if a large buffer
	// is acquired from the pool, but fewer bytes than the buffering threshold
	// are written to it, the buffer will not be returned to the pool.
	if  == nil || IsBelowBufferPoolingThreshold(cap(*)) {
		return (SliceBuffer)(*)
	}
	 := newBuffer()
	.origData = 
	.data = *
	.pool = 
	.rootBuf = 
	.refs.Store(1)
	return 
}

// Copy creates a new Buffer from the given data, initializing the reference
// counter to 1.
//
// It acquires a []byte from the given pool and copies over the backing array
// of the given data. The []byte acquired from the pool is returned to the
// pool when all references to the returned Buffer are released.
func ( []byte,  BufferPool) Buffer {
	if IsBelowBufferPoolingThreshold(len()) {
		 := make(SliceBuffer, len())
		copy(, )
		return 
	}

	 := .Get(len())
	copy(*, )
	return NewBuffer(, )
}

func ( *buffer) () []byte {
	if .rootBuf == nil {
		panic("Cannot read freed buffer")
	}
	return .data
}

func ( *buffer) () {
	if .refs.Add(1) <= 1 {
		panic("Cannot ref freed buffer")
	}
}

func ( *buffer) () {
	 := .refs.Add(-1)
	if  < 0 {
		panic("Cannot free freed buffer")
	}
	if  > 0 {
		return
	}

	.data = nil
	if .rootBuf ==  {
		// This buffer is the owner of the data slice and its ref count reached
		// 0, free the slice.
		if .pool != nil {
			.pool.Put(.origData)
			.pool = nil
		}
		.origData = nil
	} else {
		// This buffer doesn't own the data slice, decrement a ref on the root
		// buffer.
		.rootBuf.()
	}

	.rootBuf = nil
	bufferObjectPool.Put()
}

func ( *buffer) () int {
	return len(.ReadOnlyData())
}

func ( *buffer) ( int) (Buffer, Buffer) {
	if .rootBuf == nil || .rootBuf.refs.Add(1) <= 1 {
		panic("Cannot split freed buffer")
	}

	 := newBuffer()
	.data = .data[:]
	.rootBuf = .rootBuf
	.refs.Store(1)

	.data = .data[:]

	return , 
}

func ( *buffer) ( []byte) (int, Buffer) {
	if .rootBuf == nil {
		panic("Cannot read freed buffer")
	}

	 := copy(, .data)
	if  == len(.data) {
		.Free()
		return , nil
	}

	.data = .data[:]
	return , 
}

func ( *buffer) () string {
	return fmt.Sprintf("mem.Buffer(%p, data: %p, length: %d)", , .ReadOnlyData(), len(.ReadOnlyData()))
}

// ReadUnsafe reads bytes from the given Buffer into the provided slice.
// It does not perform safety checks.
func ( []byte,  Buffer) (int, Buffer) {
	return .read()
}

// SplitUnsafe modifies the receiver to point to the first n bytes while it
// returns a new reference to the remaining bytes. The returned Buffer
// functions just like a normal reference acquired using Ref().
func ( Buffer,  int) (,  Buffer) {
	return .split()
}

type emptyBuffer struct{}

func ( emptyBuffer) () []byte {
	return nil
}

func ( emptyBuffer) ()  {}
func ( emptyBuffer) () {}

func ( emptyBuffer) () int {
	return 0
}

func ( emptyBuffer) (int) (,  Buffer) {
	return , 
}

func ( emptyBuffer) ([]byte) (int, Buffer) {
	return 0, 
}

// SliceBuffer is a Buffer implementation that wraps a byte slice. It provides
// methods for reading, splitting, and managing the byte slice.
type SliceBuffer []byte

// ReadOnlyData returns the byte slice.
func ( SliceBuffer) () []byte { return  }

// Ref is a noop implementation of Ref.
func ( SliceBuffer) () {}

// Free is a noop implementation of Free.
func ( SliceBuffer) () {}

// Len is a noop implementation of Len.
func ( SliceBuffer) () int { return len() }

func ( SliceBuffer) ( int) (,  Buffer) {
	return [:], [:]
}

func ( SliceBuffer) ( []byte) (int, Buffer) {
	 := copy(, )
	if  == len() {
		return , nil
	}
	return , [:]
}