// Copyright 2024 ByteDance Inc.
//
// 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 span

import (
	"math/bits"
	"sync/atomic"

	"github.com/bytedance/gopkg/lang/dirtmake"
)

/* Span Cache: A thread-safe linear allocator

Design:
1. [GC Friendly]: Centralize a batch of small bytes slice into a big size bytes slice to avoid malloc too many objects
2. [Thread Safe]: Multi thread may share a same span, but it should fall back to the native allocator if lock conflict
*/

const (
	spanCacheSize = 10
	minSpanObject = 128                                  // 128 B
	maxSpanObject = (minSpanObject << spanCacheSize) - 1 // 128 KB
	minSpanClass  = 8                                    // = spanClass(minSpanObject)
)

type spanCache struct {
	spans [spanCacheSize]*span
}

// NewSpanCache returns a spanCache with the given spanSize,
// each span is used to allocate a binary of a specific size level.
func NewSpanCache(spanSize int) *spanCache {
	c := new(spanCache)
	for i := 0; i < len(c.spans); i++ {
		c.spans[i] = NewSpan(spanSize)
	}
	return c
}

// Make returns a [:n:n] bytes slice from a cached buffer
// NOTE: Make will not clear the underlay bytes for performance concern. So caller MUST set every byte before read.
func (c *spanCache) Make(n int) []byte {
	sclass := spanClass(n) - minSpanClass
	if sclass < 0 || sclass >= len(c.spans) {
		return dirtmake.Bytes(n, n)
	}
	return c.spans[sclass].Make(n)
}

// Copy is an alias function for make-and-copy pattern
func (c *spanCache) Copy(buf []byte) (p []byte) {
	p = c.Make(len(buf))
	copy(p, buf)
	return p
}

// NewSpan returns a span with given size
func NewSpan(size int) *span {
	sp := new(span)
	sp.size = uint32(size)
	sp.buffer = dirtmake.Bytes(0, size)
	return sp
}

type span struct {
	lock   uint32
	read   uint32 // read index of buffer
	size   uint32 // size of buffer
	buffer []byte
}

// Make returns a [:n:n] bytes slice from a cached buffer
// NOTE: Make will not clear the underlay bytes for performance concern. So caller MUST set every byte before read.
func (b *span) Make(_n int) []byte {
	n := uint32(_n)
	if n >= b.size || !atomic.CompareAndSwapUint32(&b.lock, 0, 1) {
		// fallback path: make a new byte slice if current goroutine cannot get the lock or n is out of size
		return dirtmake.Bytes(int(n), int(n))
	}
START:
	b.read += n
	// fast path
	if b.read <= b.size {
		buf := b.buffer[b.read-n : b.read : b.read]
		atomic.StoreUint32(&b.lock, 0)
		return buf
	}
	// slow path: create a new buffer
	b.buffer = dirtmake.Bytes(int(b.size), int(b.size))
	b.read = 0
	goto START
}

// Copy is an alias function for make-and-copy pattern
func (b *span) Copy(buf []byte) (p []byte) {
	p = b.Make(len(buf))
	copy(p, buf)
	return p
}

// spanClass calc the minimum number of bits required to represent x
// [2^sclass,2^(sclass+1)) bytes in a same span class
func spanClass(size int) int {
	if size == 0 {
		return 0
	}
	return bits.Len(uint(size))
}