1package wazevoapi
2
3const poolPageSize = 128
4
5// Pool is a pool of T that can be allocated and reset.
6// This is useful to avoid unnecessary allocations.
7type 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.
15func 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.
23func (p *Pool[T]) Allocated() int {
24 return p.allocated
25}
26
27// Allocate allocates a new T from the pool.
28func (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 if p.resetFn != nil {
43 p.resetFn(ret)
44 }
45 p.index++
46 p.allocated++
47 return ret
48}
49
50// View returns the pointer to i-th item from the pool.
51func (p *Pool[T]) View(i int) *T {
52 page, index := i/poolPageSize, i%poolPageSize
53 return &p.pages[page][index]
54}
55
56// Reset resets the pool.
57func (p *Pool[T]) Reset() {
58 p.pages = p.pages[:0]
59 p.index = poolPageSize
60 p.allocated = 0
61}
62
63// IDedPool is a pool of T that can be allocated and reset, with a way to get T by an ID.
64type IDedPool[T any] struct {
65 pool Pool[T]
66 idToItems []*T
67 maxIDEncountered int
68}
69
70// NewIDedPool returns a new IDedPool.
71func NewIDedPool[T any](resetFn func(*T)) IDedPool[T] {
72 return IDedPool[T]{pool: NewPool[T](resetFn), maxIDEncountered: -1}
73}
74
75// GetOrAllocate returns the T with the given id.
76func (p *IDedPool[T]) GetOrAllocate(id int) *T {
77 if p.maxIDEncountered < id {
78 p.maxIDEncountered = id
79 }
80 if id >= len(p.idToItems) {
81 p.idToItems = append(p.idToItems, make([]*T, id-len(p.idToItems)+1)...)
82 }
83 if p.idToItems[id] == nil {
84 p.idToItems[id] = p.pool.Allocate()
85 }
86 return p.idToItems[id]
87}
88
89// Get returns the T with the given id, or nil if it's not allocated.
90func (p *IDedPool[T]) Get(id int) *T {
91 if id >= len(p.idToItems) {
92 return nil
93 }
94 return p.idToItems[id]
95}
96
97// Reset resets the pool.
98func (p *IDedPool[T]) Reset() {
99 p.pool.Reset()
100 for i := 0; i <= p.maxIDEncountered; i++ {
101 p.idToItems[i] = nil
102 }
103 p.maxIDEncountered = -1
104}
105
106// MaxIDEncountered returns the maximum id encountered so far.
107func (p *IDedPool[T]) MaxIDEncountered() int {
108 return p.maxIDEncountered
109}
110
111// arraySize is the size of the array used in VarLengthPool's arrayPool.
112// This is chosen to be 8, which is empirically a good number among 8, 12, 16 and 20.
113const arraySize = 8
114
115// VarLengthPool is a pool of VarLength[T] that can be allocated and reset.
116type (
117 VarLengthPool[T any] struct {
118 arrayPool Pool[varLengthPoolArray[T]]
119 slicePool Pool[[]T]
120 }
121 // varLengthPoolArray wraps an array and keeps track of the next index to be used to avoid the heap allocation.
122 varLengthPoolArray[T any] struct {
123 arr [arraySize]T
124 next int
125 }
126)
127
128// VarLength is a variable length array that can be reused via a pool.
129type VarLength[T any] struct {
130 arr *varLengthPoolArray[T]
131 slc *[]T
132}
133
134// NewVarLengthPool returns a new VarLengthPool.
135func NewVarLengthPool[T any]() VarLengthPool[T] {
136 return VarLengthPool[T]{
137 arrayPool: NewPool[varLengthPoolArray[T]](func(v *varLengthPoolArray[T]) {
138 v.next = 0
139 }),
140 slicePool: NewPool[[]T](func(i *[]T) {
141 *i = (*i)[:0]
142 }),
143 }
144}
145
146// NewNilVarLength returns a new VarLength[T] with a nil backing.
147func NewNilVarLength[T any]() VarLength[T] {
148 return VarLength[T]{}
149}
150
151// Allocate allocates a new VarLength[T] from the pool.
152func (p *VarLengthPool[T]) Allocate(knownMin int) VarLength[T] {
153 if knownMin <= arraySize {
154 arr := p.arrayPool.Allocate()
155 return VarLength[T]{arr: arr}
156 }
157 slc := p.slicePool.Allocate()
158 return VarLength[T]{slc: slc}
159}
160
161// Reset resets the pool.
162func (p *VarLengthPool[T]) Reset() {
163 p.arrayPool.Reset()
164 p.slicePool.Reset()
165}
166
167// Append appends items to the backing slice just like the `append` builtin function in Go.
168func (i VarLength[T]) Append(p *VarLengthPool[T], items ...T) VarLength[T] {
169 if i.slc != nil {
170 *i.slc = append(*i.slc, items...)
171 return i
172 }
173
174 if i.arr == nil {
175 i.arr = p.arrayPool.Allocate()
176 }
177
178 arr := i.arr
179 if arr.next+len(items) <= arraySize {
180 for _, item := range items {
181 arr.arr[arr.next] = item
182 arr.next++
183 }
184 } else {
185 slc := p.slicePool.Allocate()
186 // Copy the array to the slice.
187 for ptr := 0; ptr < arr.next; ptr++ {
188 *slc = append(*slc, arr.arr[ptr])
189 }
190 i.slc = slc
191 *i.slc = append(*i.slc, items...)
192 }
193 return i
194}
195
196// View returns the backing slice.
197func (i VarLength[T]) View() []T {
198 if i.slc != nil {
199 return *i.slc
200 } else if i.arr != nil {
201 arr := i.arr
202 return arr.arr[:arr.next]
203 }
204 return nil
205}
206
207// Cut cuts the backing slice to the given length.
208// Precondition: n <= len(i.backing).
209func (i VarLength[T]) Cut(n int) {
210 if i.slc != nil {
211 *i.slc = (*i.slc)[:n]
212 } else if i.arr != nil {
213 i.arr.next = n
214 }
215}