github.com/MetalBlockchain/metalgo@v1.11.9/utils/bytes.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package utils
     5  
     6  import (
     7  	"crypto/rand"
     8  	"math/bits"
     9  	"sync"
    10  )
    11  
    12  // RandomBytes returns a slice of n random bytes
    13  // Intended for use in testing
    14  func RandomBytes(n int) []byte {
    15  	b := make([]byte, n)
    16  	_, _ = rand.Read(b)
    17  	return b
    18  }
    19  
    20  // Constant taken from the "math" package
    21  const intSize = 32 << (^uint(0) >> 63) // 32 or 64
    22  
    23  // BytesPool tracks buckets of available buffers to be allocated. Each bucket
    24  // allocates buffers of the following length:
    25  //
    26  // 0
    27  // 1
    28  // 3
    29  // 7
    30  // 15
    31  // 31
    32  // 63
    33  // 127
    34  // ...
    35  // MaxInt
    36  //
    37  // In order to allocate a buffer of length 19 (for example), we calculate the
    38  // number of bits required to represent 19 (5). And therefore allocate a slice
    39  // from bucket 5, which has length 31. This is the bucket which produces the
    40  // smallest slices that are at least length 19.
    41  //
    42  // When replacing a buffer of length 19, we calculate the number of bits
    43  // required to represent 20 (5). And therefore place the slice into bucket 4,
    44  // which has length 15. This is the bucket which produces the largest slices
    45  // that a length 19 slice can be used for.
    46  type BytesPool [intSize]sync.Pool
    47  
    48  func NewBytesPool() *BytesPool {
    49  	var p BytesPool
    50  	for i := range p {
    51  		// uint is used here to avoid overflowing int during the shift
    52  		size := uint(1)<<i - 1
    53  		p[i] = sync.Pool{
    54  			New: func() interface{} {
    55  				// Sync pool needs to return pointer-like values to avoid memory
    56  				// allocations.
    57  				b := make([]byte, size)
    58  				return &b
    59  			},
    60  		}
    61  	}
    62  	return &p
    63  }
    64  
    65  // Get returns a non-nil pointer to a slice with the requested length.
    66  //
    67  // It is not guaranteed for the returned bytes to have been zeroed.
    68  func (p *BytesPool) Get(length int) *[]byte {
    69  	index := bits.Len(uint(length)) // Round up
    70  	bytes := p[index].Get().(*[]byte)
    71  	*bytes = (*bytes)[:length] // Set the length to be the expected value
    72  	return bytes
    73  }
    74  
    75  // Put takes ownership of a non-nil pointer to a slice of bytes.
    76  //
    77  // Note: this function takes ownership of the underlying array. So, the length
    78  // of the provided slice is ignored and only its capacity is used.
    79  func (p *BytesPool) Put(bytes *[]byte) {
    80  	size := cap(*bytes)
    81  	index := bits.Len(uint(size)+1) - 1 // Round down
    82  	p[index].Put(bytes)
    83  }