qanglang_core/memory/
closure_arena.rs1use crate::{ClosureObject, debug_log};
2use crate::{
3 arena::{Arena, Index},
4 memory::UpvalueHandle,
5};
6
7pub const INLINE_UPVALUE_COUNT: usize = 4;
8const OVERFLOW_CHUNK_SIZE: usize = 8;
9
10pub type ClosureHandle = Index;
11pub type OverflowHandle = Index;
12
13#[derive(Debug, Clone, PartialEq, Copy)]
14#[repr(u8)]
15pub enum UpvalueSlot {
16 Open(usize),
17 Closed(UpvalueHandle),
18}
19
20impl Default for UpvalueSlot {
21 fn default() -> Self {
22 Self::Open(0)
23 }
24}
25
26#[derive(Debug, Clone)]
27pub struct OverflowChunk {
28 pub entries: [UpvalueSlot; OVERFLOW_CHUNK_SIZE],
29 pub count: usize,
30 pub next: Option<OverflowHandle>,
31 pub is_marked: bool,
32}
33
34impl Default for OverflowChunk {
35 fn default() -> Self {
36 Self {
37 entries: [UpvalueSlot::default(); OVERFLOW_CHUNK_SIZE],
38 count: 0,
39 next: None,
40 is_marked: false,
41 }
42 }
43}
44
45#[derive(Debug, Default, Clone)]
46pub struct ClosureArena {
47 closures: Arena<ClosureObject>,
48 overflow_chunks: Arena<OverflowChunk>,
49 is_debug: bool,
50}
51
52impl ClosureArena {
53 pub fn new() -> Self {
54 Self::with_capacity(32)
55 }
56
57 pub fn with_capacity(capacity: usize) -> Self {
58 Self {
59 closures: Arena::with_capacity(capacity),
60 overflow_chunks: Arena::with_capacity(capacity),
61 is_debug: false,
62 }
63 }
64
65 pub fn set_debug(mut self, is_debug: bool) -> Self {
66 self.is_debug = is_debug;
67 self
68 }
69
70 pub fn allocate_closure(&mut self, closure: ClosureObject) -> ClosureHandle {
71 let handle = self.closures.insert(closure);
72 debug_log!(
73 self.is_debug,
74 "Allocated {} bytes for closure: {:?}",
75 std::mem::size_of::<ClosureObject>(),
76 handle
77 );
78 handle
79 }
80
81 pub fn get_closure(&self, handle: ClosureHandle) -> &ClosureObject {
82 &self.closures[handle]
83 }
84
85 pub fn get_closure_mut(&mut self, handle: ClosureHandle) -> &mut ClosureObject {
86 &mut self.closures[handle]
87 }
88
89 pub fn free_closure(&mut self, handle: ClosureHandle) {
90 self.closures.remove(handle);
91 debug_log!(
92 self.is_debug,
93 "Freed {} bytes for closure: {:?}",
94 std::mem::size_of::<ClosureObject>(),
95 handle
96 );
97 }
98
99 pub fn get_upvalue(&self, closure_handle: ClosureHandle, index: usize) -> Option<UpvalueSlot> {
100 let closure = self.closures.get(closure_handle)?;
101
102 if index >= closure.upvalue_count {
103 return None;
104 }
105
106 if index < INLINE_UPVALUE_COUNT {
107 Some(closure.inline_upvalues[index])
108 } else {
109 let mut overflow_index = index - INLINE_UPVALUE_COUNT;
110 let mut current_handle = closure.overflow_handle?;
111
112 loop {
113 let chunk = self.get_overflow_chunk(current_handle)?;
114 if overflow_index < chunk.count {
115 return Some(chunk.entries[overflow_index]);
116 }
117 overflow_index -= chunk.count;
118 current_handle = chunk.next?;
119 }
120 }
121 }
122
123 pub fn set_upvalue(
124 &mut self,
125 closure_handle: ClosureHandle,
126 index: usize,
127 value: UpvalueSlot,
128 ) -> bool {
129 let upvalue_count = if let Some(closure) = self.closures.get(closure_handle) {
130 closure.upvalue_count
131 } else {
132 return false;
133 };
134
135 if index >= upvalue_count {
136 return false;
137 }
138
139 if index < INLINE_UPVALUE_COUNT {
140 if let Some(closure) = self.closures.get_mut(closure_handle) {
141 closure.inline_upvalues[index] = value;
142 true
143 } else {
144 false
145 }
146 } else {
147 let mut overflow_index = index - INLINE_UPVALUE_COUNT;
148
149 let first_handle = if let Some(closure) = self.closures.get(closure_handle) {
151 if let Some(handle) = closure.overflow_handle {
152 handle
153 } else {
154 let handle = self.allocate_overflow_chunk();
155 if let Some(closure) = self.closures.get_mut(closure_handle) {
156 closure.overflow_handle = Some(handle);
157 }
158 handle
159 }
160 } else {
161 return false;
162 };
163
164 let mut current_handle = first_handle;
165 loop {
166 if overflow_index < OVERFLOW_CHUNK_SIZE {
167 if let Some(chunk) = self.get_overflow_chunk_mut(current_handle) {
168 chunk.entries[overflow_index] = value;
169 chunk.count = chunk.count.max(overflow_index + 1);
170 return true;
171 } else {
172 return false;
173 }
174 }
175
176 overflow_index -= OVERFLOW_CHUNK_SIZE;
177
178 let next_handle = if let Some(chunk) = self.get_overflow_chunk(current_handle) {
179 if let Some(next) = chunk.next {
180 next
181 } else {
182 let new_handle = self.allocate_overflow_chunk();
183 if let Some(chunk) = self.get_overflow_chunk_mut(current_handle) {
184 chunk.next = Some(new_handle);
185 }
186 new_handle
187 }
188 } else {
189 return false;
190 };
191
192 current_handle = next_handle;
193 }
194 }
195 }
196
197 pub fn allocate_overflow_chunk(&mut self) -> OverflowHandle {
198 self.overflow_chunks.insert(OverflowChunk::default())
199 }
200
201 pub fn get_overflow_chunk(&self, handle: OverflowHandle) -> Option<&OverflowChunk> {
202 self.overflow_chunks.get(handle)
203 }
204
205 pub fn get_overflow_chunk_mut(&mut self, handle: OverflowHandle) -> Option<&mut OverflowChunk> {
206 self.overflow_chunks.get_mut(handle)
207 }
208
209 pub fn mark_closure(&mut self, handle: ClosureHandle) {
210 if let Some(closure) = self.closures.get_mut(handle) {
211 closure.is_marked = true;
212 if let Some(overflow_handle) = closure.overflow_handle {
213 self.mark_overflow_chain(overflow_handle);
214 }
215 }
216 }
217
218 pub fn mark_overflow_chain(&mut self, handle: OverflowHandle) {
219 let mut current_handle = Some(handle);
220 while let Some(h) = current_handle {
221 if let Some(chunk) = self.overflow_chunks.get_mut(h) {
222 if chunk.is_marked {
223 break;
224 }
225 chunk.is_marked = true;
226 current_handle = chunk.next;
227 } else {
228 break;
229 }
230 }
231 }
232
233 pub fn trace_closure_references(
234 &mut self,
235 handle: ClosureHandle,
236 upvalue_tracer: &mut dyn FnMut(UpvalueHandle),
237 ) {
238 let (should_mark, overflow_handle, upvalue_count) =
239 if let Some(closure) = self.closures.get(handle) {
240 (
241 !closure.is_marked,
242 closure.overflow_handle,
243 closure.upvalue_count,
244 )
245 } else {
246 return;
247 };
248
249 if should_mark {
250 debug_log!(self.is_debug, "Marking closure: {:?}", handle);
251 if let Some(closure) = self.closures.get_mut(handle) {
252 closure.is_marked = true;
253 }
254
255 if let Some(overflow_handle) = overflow_handle {
256 self.mark_overflow_chain(overflow_handle);
257 }
258
259 debug_log!(self.is_debug, "Blackening closure: {:?}", handle);
260
261 for i in 0..upvalue_count {
262 if let Some(UpvalueSlot::Closed(upvalue_handle)) = self.get_upvalue(handle, i) {
263 upvalue_tracer(upvalue_handle);
264 }
265 }
266 }
267 }
268
269 pub fn collect_garbage(&mut self) {
270 let estimated_deletions = (self.closures.len() / 4).max(8);
271 let mut handles_to_remove = Vec::with_capacity(estimated_deletions);
272
273 for (handle, closure) in self.closures.iter_mut() {
274 if closure.is_marked {
275 closure.is_marked = false;
276 } else {
277 handles_to_remove.push(handle);
278 }
279 }
280
281 for handle in &handles_to_remove {
282 self.closures.remove(*handle);
283 }
284
285 handles_to_remove.clear();
286
287 for (handle, chunk) in self.overflow_chunks.iter_mut() {
288 if chunk.is_marked {
289 chunk.is_marked = false;
290 } else {
291 handles_to_remove.push(handle);
292 }
293 }
294
295 for handle in &handles_to_remove {
296 self.overflow_chunks.remove(*handle);
297 }
298 }
299
300 pub fn get_allocated_bytes(&self) -> usize {
301 let closure_bytes = self.closures.len() * std::mem::size_of::<ClosureObject>();
302 let overflow_bytes = self.overflow_chunks.len() * std::mem::size_of::<OverflowChunk>();
303 closure_bytes + overflow_bytes
304 }
305
306 pub fn len(&self) -> usize {
307 self.closures.len()
308 }
309
310 pub fn is_empty(&self) -> bool {
311 self.closures.is_empty()
312 }
313
314 pub fn iter(&self) -> impl Iterator<Item = (ClosureHandle, &ClosureObject)> {
315 self.closures.iter()
316 }
317
318 pub fn iter_mut(&mut self) -> impl Iterator<Item = (ClosureHandle, &mut ClosureObject)> {
319 self.closures.iter_mut()
320 }
321}