github.com/imgk/memory-go@v0.0.0-20220328012817-37cdd311f1a3/memory.go (about)

     1  //go:build !malloc_cgo
     2  
     3  package memory
     4  
     5  import (
     6  	"errors"
     7  	"sync"
     8  	"unsafe"
     9  )
    10  
    11  // DefaultAllocator is ...
    12  var DefaultAllocator = NewAllocator()
    13  
    14  // pointer is ...
    15  type pointer struct {
    16  	Pointer *[]byte
    17  }
    18  
    19  // Alloc is ..
    20  func Alloc[T any](n int) (pointer, []T) {
    21  	p := DefaultAllocator.Get(n * int(unsafe.Sizeof(*(new(T)))))
    22  	return p, unsafe.Slice((*T)(unsafe.Pointer(*(*uintptr)(unsafe.Pointer(p.Pointer)))), n)
    23  }
    24  
    25  // Free is ...
    26  func Free(p pointer) {
    27  	DefaultAllocator.Put(p)
    28  }
    29  
    30  // mem is ...
    31  type mem int
    32  
    33  func (n mem) Alloc() any {
    34  	// Return a *[]byte instead of []byte ensures that
    35  	// the []byte is not copied, which would cause a heap
    36  	// allocation on every call to sync.Pool.Put
    37  	buf := make([]byte, int(n))
    38  	return &buf
    39  }
    40  
    41  // MemoryPool is ...
    42  type MemoryPool[T any] struct {
    43  	sync.Pool
    44  }
    45  
    46  // Get is ...
    47  func (p *MemoryPool[T]) Get() *T {
    48  	return p.Pool.Get().(*T)
    49  }
    50  
    51  // Put is ...
    52  func (p *MemoryPool[T]) Put(t *T) {
    53  	p.Pool.Put(t)
    54  }
    55  
    56  // Allocator for incoming frames, optimized to prevent overwriting after zeroing
    57  type Allocator struct {
    58  	buffers []MemoryPool[[]byte]
    59  }
    60  
    61  // NewAllocator initiates a []byte allocator for frames less than 65536 bytes,
    62  // the waste(memory fragmentation) of space allocation is guaranteed to be
    63  // no more than 50%.
    64  func NewAllocator() *Allocator {
    65  	alloc := &Allocator{}
    66  	alloc.buffers = make([]MemoryPool[[]byte], 17) // 1B -> 64K
    67  	for k := range alloc.buffers {
    68  		i := k
    69  		alloc.buffers[k].Pool.New = mem(1 << uint32(i)).Alloc
    70  	}
    71  	return alloc
    72  }
    73  
    74  // Get a []byte from pool with most appropriate cap
    75  func (alloc *Allocator) Get(n int) pointer {
    76  	if n <= 0 {
    77  		return pointer{}
    78  	}
    79  	if n > 65536 {
    80  		b := make([]byte, n)
    81  		return pointer{Pointer: &b}
    82  	}
    83  
    84  	bits := msb(n)
    85  	if n == 1<<bits {
    86  		p := alloc.buffers[bits].Get()
    87  		return pointer{Pointer: p}
    88  	} else {
    89  		p := alloc.buffers[bits+1].Get()
    90  		return pointer{Pointer: p}
    91  	}
    92  }
    93  
    94  // Put returns a []byte to pool for future use,
    95  // which the cap must be exactly 2^n
    96  func (alloc *Allocator) Put(p pointer) error {
    97  	buf := *p.Pointer
    98  	bits := msb(cap(buf))
    99  	if cap(buf) == 0 || cap(buf) > 65536 || cap(buf) != 1<<bits {
   100  		return errors.New("allocator Put() incorrect buffer size")
   101  	}
   102  	alloc.buffers[bits].Put(p.Pointer)
   103  	return nil
   104  }
   105  
   106  // msb return the pos of most significiant bit
   107  // http://supertech.csail.mit.edu/papers/debruijn.pdf
   108  func msb(size int) byte {
   109  	// var debruijinPos = [...]byte{0, 9, 1, 10, 13, 21, 2, 29, 11, 14, 16, 18, 22, 25, 3, 30, 8, 12, 20, 28, 15, 17, 24, 7, 19, 27, 23, 6, 26, 5, 4, 31}
   110  	const debruijinPos = "\x00\x09\x01\x0a\x0d\x15\x02\x1d\x0b\x0e\x10\x12\x16\x19\x03\x1e\x08\x0c\x14\x1c\x0f\x11\x18\x07\x13\x1b\x17\x06\x1a\x05\x04\x1f"
   111  	v := uint32(size)
   112  	v |= v >> 1
   113  	v |= v >> 2
   114  	v |= v >> 4
   115  	v |= v >> 8
   116  	v |= v >> 16
   117  	return debruijinPos[(v*0x07C4ACDD)>>27]
   118  }