github.com/bananabytelabs/wazero@v0.0.0-20240105073314-54b22a776da8/internal/engine/wazevo/wazevoapi/pool.go (about)

     1  package wazevoapi
     2  
     3  const poolPageSize = 128
     4  
     5  // Pool is a pool of T that can be allocated and reset.
     6  // This is useful to avoid unnecessary allocations.
     7  type Pool[T any] struct {
     8  	pages            []*[poolPageSize]T
     9  	resetFn          func(*T)
    10  	allocated, index int
    11  }
    12  
    13  // NewPool returns a new Pool.
    14  // resetFn is called when a new T is allocated in Pool.Allocate.
    15  func NewPool[T any](resetFn func(*T)) Pool[T] {
    16  	var ret Pool[T]
    17  	ret.resetFn = resetFn
    18  	ret.Reset()
    19  	return ret
    20  }
    21  
    22  // Allocated returns the number of allocated T currently in the pool.
    23  func (p *Pool[T]) Allocated() int {
    24  	return p.allocated
    25  }
    26  
    27  // Allocate allocates a new T from the pool.
    28  func (p *Pool[T]) Allocate() *T {
    29  	if p.index == poolPageSize {
    30  		if len(p.pages) == cap(p.pages) {
    31  			p.pages = append(p.pages, new([poolPageSize]T))
    32  		} else {
    33  			i := len(p.pages)
    34  			p.pages = p.pages[:i+1]
    35  			if p.pages[i] == nil {
    36  				p.pages[i] = new([poolPageSize]T)
    37  			}
    38  		}
    39  		p.index = 0
    40  	}
    41  	ret := &p.pages[len(p.pages)-1][p.index]
    42  	p.resetFn(ret)
    43  	p.index++
    44  	p.allocated++
    45  	return ret
    46  }
    47  
    48  // View returns the pointer to i-th item from the pool.
    49  func (p *Pool[T]) View(i int) *T {
    50  	page, index := i/poolPageSize, i%poolPageSize
    51  	return &p.pages[page][index]
    52  }
    53  
    54  // Reset resets the pool.
    55  func (p *Pool[T]) Reset() {
    56  	p.pages = p.pages[:0]
    57  	p.index = poolPageSize
    58  	p.allocated = 0
    59  }