github.com/15mga/kiwi@v0.0.2-0.20240324021231-b95d5c3ac751/util/pool.go (about)

     1  package util
     2  
     3  type (
     4  	poolOption[T any] struct {
     5  		minCap int
     6  		maxCap int
     7  		spawn  func() T
     8  	}
     9  	PoolOption[T any] func(o *poolOption[T])
    10  )
    11  
    12  func PoolMinCap[T any](c int) PoolOption[T] {
    13  	return func(o *poolOption[T]) {
    14  		o.minCap = c
    15  	}
    16  }
    17  
    18  func PoolMaxCap[T any](c int) PoolOption[T] {
    19  	return func(o *poolOption[T]) {
    20  		o.maxCap = c
    21  	}
    22  }
    23  
    24  func PoolSpawn[T any](spawn func() T) PoolOption[T] {
    25  	return func(o *poolOption[T]) {
    26  		o.spawn = spawn
    27  	}
    28  }
    29  
    30  func NewPool[T any](opts ...PoolOption[T]) *Pool[T] {
    31  	opt := &poolOption[T]{
    32  		minCap: 16,
    33  		maxCap: 256,
    34  	}
    35  	for _, o := range opts {
    36  		o(opt)
    37  	}
    38  	return &Pool[T]{
    39  		opt:    opt,
    40  		cap:    opt.minCap,
    41  		values: make([]T, opt.minCap),
    42  	}
    43  }
    44  
    45  type Pool[T any] struct {
    46  	opt    *poolOption[T]
    47  	values []T
    48  	idx    int
    49  	cap    int
    50  }
    51  
    52  func (p *Pool[T]) Shift() (v T) {
    53  	if p.idx == 0 {
    54  		v = p.opt.spawn()
    55  		return
    56  	}
    57  	p.idx--
    58  	v = p.values[p.idx]
    59  	return
    60  }
    61  
    62  func (p *Pool[T]) Push(value T) {
    63  	if p.idx == p.cap {
    64  		if p.cap == p.opt.maxCap {
    65  			return
    66  		}
    67  		p.resetCap(p.cap << 1)
    68  	}
    69  	p.values[p.idx] = value
    70  	p.idx++
    71  }
    72  
    73  func (p *Pool[T]) resetCap(c int) {
    74  	p.cap = c
    75  	values := make([]T, p.cap)
    76  	copy(values, p.values[:p.idx])
    77  	p.values = values
    78  }