1use parking_lot::{Mutex, RwLock};
2use runmat_builtins::Value;
5use runmat_time::Instant;
6use std::collections::HashSet;
7use std::sync::atomic::{AtomicBool, Ordering};
8use std::sync::Arc;
9use thiserror::Error;
10
11pub mod allocator;
12pub mod barriers;
13pub mod collector;
14pub mod config;
15pub mod gc_ptr;
16pub use runmat_gc_api::GcPtr;
17pub mod generations;
18pub mod roots;
19pub mod stats;
20
21use std::collections::HashMap;
23
24pub trait GcFinalizer: Send + Sync {
28 fn finalize(&self);
29}
30
31static FINALIZERS: once_cell::sync::Lazy<
32 parking_lot::Mutex<HashMap<usize, std::sync::Arc<dyn GcFinalizer>>>,
33> = once_cell::sync::Lazy::new(|| parking_lot::Mutex::new(HashMap::new()));
34
35pub fn gc_register_finalizer(ptr: GcPtr<Value>, f: std::sync::Arc<dyn GcFinalizer>) -> Result<()> {
37 let addr = unsafe { ptr.as_raw() } as usize;
38 if addr == 0 {
39 return Ok(());
40 }
41 FINALIZERS.lock().insert(addr, f);
42 Ok(())
43}
44
45pub fn gc_unregister_finalizer(ptr: GcPtr<Value>) -> Result<()> {
47 let addr = unsafe { ptr.as_raw() } as usize;
48 if addr == 0 {
49 return Ok(());
50 }
51 FINALIZERS.lock().remove(&addr);
52 Ok(())
53}
54
55pub(crate) fn gc_run_finalizer_for_addr(addr: usize) {
57 if let Some(f) = FINALIZERS.lock().remove(&addr) {
58 f.finalize();
60 }
61}
62
63#[cfg(feature = "pointer-compression")]
64pub mod compression;
65
66pub use allocator::{AllocatorStats, GenerationalAllocator, SizeClass};
67use barriers::WriteBarrierManager;
68pub use barriers::{CardTable, WriteBarrier};
69pub use collector::MarkSweepCollector;
70pub use config::{GcConfig, GcConfigBuilder};
71pub use generations::{Generation, GenerationStats, GenerationalHeap, GenerationalHeapStats};
73pub use roots::{GcRoot, GlobalRoot, RootId, RootScanner, StackRoot, VariableArrayRoot};
74pub use stats::{CollectionEvent, CollectionType, GcStats};
75
76#[cfg(feature = "pointer-compression")]
77pub use compression::*;
78
79#[derive(Error, Debug)]
81pub enum GcError {
82 #[error("Out of memory: {0}")]
83 OutOfMemory(String),
84
85 #[error("Invalid GC pointer: {0}")]
86 InvalidPointer(String),
87
88 #[error("Collection failed: {0}")]
89 CollectionFailed(String),
90
91 #[error("Root registration failed: {0}")]
92 RootRegistrationFailed(String),
93
94 #[error("Configuration error: {0}")]
95 ConfigError(String),
96
97 #[error("Thread synchronization error: {0}")]
98 SyncError(String),
99}
100
101pub type Result<T> = std::result::Result<T, GcError>;
102
103pub struct GarbageCollector {
107 config: Arc<RwLock<GcConfig>>,
109
110 allocator: Mutex<GenerationalAllocator>,
112
113 collector: Mutex<MarkSweepCollector>,
115
116 root_ptrs: Arc<Mutex<HashSet<usize>>>,
118
119 collection_in_progress: AtomicBool,
121
122 stats: Arc<GcStats>,
124}
125
126impl GarbageCollector {
127 pub fn new() -> Result<Self> {
128 Self::with_config(GcConfig::default())
129 }
130
131 pub fn with_config(config: GcConfig) -> Result<Self> {
132 let allocator = GenerationalAllocator::new(&config);
133 let collector = MarkSweepCollector::new(&config);
134
135 Ok(Self {
136 config: Arc::new(RwLock::new(config)),
137 allocator: Mutex::new(allocator),
138 collector: Mutex::new(collector),
139 root_ptrs: Arc::new(Mutex::new(HashSet::new())),
140 collection_in_progress: AtomicBool::new(false),
141 stats: Arc::new(GcStats::new()),
142 })
143 }
144
145 pub fn allocate(&self, value: Value) -> Result<GcPtr<Value>> {
147 if self.should_collect() {
149 let _ = self.collect_minor();
150 }
151
152 let gpu_handle: Option<runmat_accelerate_api::GpuTensorHandle> =
154 if let Value::GpuTensor(h) = &value {
155 Some(h.clone())
156 } else {
157 None
158 };
159
160 let mut allocator = self.allocator.lock();
161 let ptr = allocator.allocate(value, &self.stats)?;
162 let usage = allocator.young_generation_usage();
163 let alloc_count = allocator.young_allocations_count();
164 let cfg = self.config.read().clone();
165 drop(allocator);
166
167 if let Some(handle) = gpu_handle {
169 struct GpuTensorFinalizer {
170 handle: runmat_accelerate_api::GpuTensorHandle,
171 }
172 impl GcFinalizer for GpuTensorFinalizer {
173 fn finalize(&self) {
174 if let Some(p) = runmat_accelerate_api::provider() {
175 let _ = p.free(&self.handle);
176 runmat_accelerate_api::clear_handle_logical(&self.handle);
177 }
178 }
179 }
180 let fin = std::sync::Arc::new(GpuTensorFinalizer { handle });
181 let _ = gc_register_finalizer(ptr.clone(), fin);
182 }
183
184 if usage > cfg.minor_gc_threshold
188 || (cfg.minor_gc_threshold <= 0.35 && alloc_count > 0 && alloc_count.is_multiple_of(32))
189 {
190 let _ = self.collect_minor();
191 }
192 Ok(ptr)
193 }
194
195 fn should_collect(&self) -> bool {
199 if self.collection_in_progress.load(Ordering::Acquire) {
200 return false;
201 }
202
203 let threshold = self.config.read().minor_gc_threshold;
204 let allocator = self.allocator.lock();
205 let utilization = allocator.young_generation_usage();
206 utilization > threshold
207 }
208
209 pub fn collect_minor(&self) -> Result<usize> {
211 if self
212 .collection_in_progress
213 .compare_exchange(false, true, Ordering::AcqRel, Ordering::Relaxed)
214 .is_err()
215 {
216 return Ok(0);
217 }
218
219 let start_time = Instant::now();
220
221 let mut combined_roots: Vec<GcPtr<Value>> = Vec::new();
223 {
224 let roots = self.root_ptrs.lock();
225 combined_roots.extend(
226 roots
227 .iter()
228 .map(|&addr| unsafe { GcPtr::from_raw(addr as *const Value) }),
229 );
230 }
231 if let Ok(mut ext) = ROOT_SCANNER.scan_roots() {
233 combined_roots.append(&mut ext);
234 }
235 combined_roots.extend(gc_barrier_minor_roots());
236
237 let mut allocator = self.allocator.lock();
239 let mut collector = self.collector.lock();
240 let collected_count =
241 collector.collect_young_generation(&mut allocator, &combined_roots, &self.stats)?;
242
243 WRITE_BARRIERS.clear_after_minor_gc();
245
246 let duration = start_time.elapsed();
247 self.stats
248 .record_minor_collection(collected_count, duration);
249
250 self.collection_in_progress.store(false, Ordering::Release);
251
252 log::info!("Minor GC: collected {collected_count} objects in {duration:?}");
253
254 Ok(collected_count)
255 }
256
257 pub fn collect_major(&self) -> Result<usize> {
268 if self
269 .collection_in_progress
270 .compare_exchange(false, true, Ordering::AcqRel, Ordering::Relaxed)
271 .is_err()
272 {
273 return Ok(0);
274 }
275
276 let start_time = Instant::now();
277
278 let mut combined_roots: Vec<GcPtr<Value>> = Vec::new();
280 {
281 let roots = self.root_ptrs.lock();
282 combined_roots.extend(
283 roots
284 .iter()
285 .map(|&addr| unsafe { GcPtr::from_raw(addr as *const Value) }),
286 );
287 }
288 if let Ok(mut ext) = ROOT_SCANNER.scan_roots() {
289 combined_roots.append(&mut ext);
290 }
291 combined_roots.extend(gc_barrier_minor_roots());
292
293 let mut allocator = self.allocator.lock();
294 let mut collector = self.collector.lock();
295 let collected_count =
296 collector.collect_all_generations(&mut allocator, &combined_roots, &self.stats)?;
297
298 WRITE_BARRIERS.clear_after_major_gc();
300 allocator.clear_promotion_state();
301
302 let duration = start_time.elapsed();
303 self.stats
304 .record_major_collection(collected_count, duration);
305
306 self.collection_in_progress.store(false, Ordering::Release);
307
308 log::info!("Major GC: collected {collected_count} objects in {duration:?}");
309
310 Ok(collected_count)
311 }
312
313 pub fn add_root(&self, root: GcPtr<Value>) -> Result<()> {
315 let value_ptr = unsafe { root.as_raw() } as usize;
316 if value_ptr == 0 {
317 return Ok(());
318 }
319 self.root_ptrs.lock().insert(value_ptr);
320 Ok(())
321 }
322
323 pub fn remove_root(&self, root: GcPtr<Value>) -> Result<()> {
325 let value_ptr = unsafe { root.as_raw() } as usize;
326 if value_ptr == 0 {
327 return Ok(());
328 }
329 self.root_ptrs.lock().remove(&value_ptr);
330 Ok(())
331 }
332
333 pub fn stats(&self) -> GcStats {
337 self.stats.as_ref().clone()
338 }
339
340 pub fn configure(&self, config: GcConfig) -> Result<()> {
342 *self.config.write() = config.clone();
343 {
344 let mut allocator = self.allocator.lock();
346 *allocator = GenerationalAllocator::new(&config);
347 }
348 {
349 let mut collector = self.collector.lock();
350 collector.reconfigure(&config)?;
351 }
352 Ok(())
353 }
354
355 pub fn get_config(&self) -> GcConfig {
357 self.config.read().clone()
358 }
359}
360
361unsafe impl Send for GarbageCollector {}
363unsafe impl Sync for GarbageCollector {}
364
365static GC: once_cell::sync::Lazy<Arc<GarbageCollector>> =
367 once_cell::sync::Lazy::new(|| Arc::new(GarbageCollector::new().expect("Failed to create GC")));
368static ROOT_SCANNER: once_cell::sync::Lazy<Arc<RootScanner>> =
369 once_cell::sync::Lazy::new(|| Arc::new(RootScanner::new()));
370
371static WRITE_BARRIERS: once_cell::sync::Lazy<Arc<WriteBarrierManager>> =
373 once_cell::sync::Lazy::new(|| Arc::new(WriteBarrierManager::new(true, false)));
374
375pub fn gc_deref(ptr: GcPtr<Value>) -> Value {
377 (*ptr).clone()
378}
379
380pub fn gc_allocate(value: Value) -> Result<GcPtr<Value>> {
382 GC.allocate(value)
383}
384
385pub fn gc_collect_minor() -> Result<usize> {
386 GC.collect_minor()
388}
389
390pub fn gc_collect_major() -> Result<usize> {
391 GC.collect_major()
392}
393
394pub fn gc_add_root(root: GcPtr<Value>) -> Result<()> {
395 GC.add_root(root)
396}
397
398pub fn gc_remove_root(root: GcPtr<Value>) -> Result<()> {
399 GC.remove_root(root)
400}
401
402pub fn gc_stats() -> GcStats {
403 GC.stats()
404}
405
406pub fn gc_configure(config: GcConfig) -> Result<()> {
407 GC.configure(config)
408}
409
410pub fn gc_get_config() -> GcConfig {
411 GC.get_config()
412}
413
414pub fn gc_register_root(root: Box<dyn GcRoot>) -> Result<RootId> {
416 ROOT_SCANNER.register_root(root)
417}
418
419pub fn gc_unregister_root(root_id: RootId) -> Result<()> {
420 ROOT_SCANNER.unregister_root(root_id)
421}
422
423pub fn gc_record_write(old: &Value, new: &Value) {
425 let old_ptr = old as *const Value as *const u8;
427 let young_ptr = new as *const Value as *const u8;
428 let alloc = GC.allocator.lock();
430 let old_gen = alloc.logical_generation(old_ptr).unwrap_or(0);
431 let new_gen = alloc.logical_generation(young_ptr).unwrap_or(0);
432 drop(alloc);
433 if old_gen > new_gen {
434 WRITE_BARRIERS.record_reference(old_ptr, young_ptr);
435 }
436}
437
438pub fn gc_barrier_minor_roots() -> Vec<GcPtr<Value>> {
440 WRITE_BARRIERS
441 .get_minor_gc_roots()
442 .into_iter()
443 .map(|p| unsafe { GcPtr::from_raw(p as *const Value) })
444 .collect()
445}
446
447#[cfg(any(test, feature = "debug-gc"))]
449pub fn gc_force_collect() -> Result<usize> {
450 gc_collect_major()
451}
452
453pub fn gc_reset_for_test() -> Result<()> {
455 GC.stats.reset();
457
458 {
460 let config = GcConfig::default();
461 *GC.config.write() = config.clone();
462 let mut alloc = GC.allocator.lock();
463 *alloc = GenerationalAllocator::new(&config);
464 let mut coll = GC.collector.lock();
465 *coll = MarkSweepCollector::new(&config);
466 }
467
468 {
469 let mut roots = GC.root_ptrs.lock();
470 roots.clear();
471 }
472
473 GC.collection_in_progress.store(false, Ordering::Relaxed);
475
476 gc_configure(GcConfig::default())?;
478
479 Ok(())
480}
481
482pub fn gc_test_context<F, R>(test_fn: F) -> R
484where
485 F: FnOnce() -> R,
486{
487 static TEST_MUTEX: std::sync::Mutex<()> = std::sync::Mutex::new(());
489
490 let _guard = match TEST_MUTEX.lock() {
492 Ok(guard) => guard,
493 Err(poisoned) => {
494 poisoned.into_inner()
496 }
497 };
498
499 gc_reset_for_test().expect("GC reset should succeed");
501
502 let result = test_fn();
504
505 let _ = gc_reset_for_test();
507
508 result
509}
510
511#[cfg(test)]
512mod tests {
513 use super::*;
514
515 #[test]
516 fn test_basic_allocation() {
517 gc_test_context(|| {
518 let value = Value::Num(42.0);
519 let ptr = gc_allocate(value).expect("allocation failed");
520 assert_eq!(*ptr, Value::Num(42.0));
521 });
522 }
523
524 #[test]
525 fn test_collection() {
526 gc_test_context(|| {
527 let mut _ptrs = Vec::new();
528 for i in 0..100 {
529 let ptr = gc_allocate(Value::Num(i as f64)).expect("allocation failed");
530 _ptrs.push(ptr);
531 }
532
533 let _collected = gc_collect_minor().expect("collection failed");
534 drop(_ptrs);
535 });
536 }
537
538 #[test]
539 fn test_thread_safety() {
540 use std::thread;
541
542 gc_test_context(|| {
543 let handles: Vec<_> = (0..4)
544 .map(|i| {
545 thread::spawn(move || {
546 let mut ptrs = Vec::new();
547 for j in 0..100 {
548 let value = Value::Num((i * 100 + j) as f64);
549 let ptr = gc_allocate(value).expect("allocation failed");
550 ptrs.push(ptr);
551 }
552 ptrs
553 })
554 })
555 .collect();
556
557 for handle in handles {
558 let _ptrs = handle.join().expect("thread failed");
559 }
560
561 let _ = gc_collect_major();
562 });
563 }
564
565 #[test]
566 fn test_root_protection() {
567 gc_test_context(|| {
568 let protected = gc_allocate(Value::Num(42.0)).expect("allocation failed");
569 gc_add_root(protected.clone()).expect("root registration failed");
570
571 for i in 0..60 {
572 let _ = gc_allocate(Value::String(format!("garbage_{i}")));
573 }
574
575 let _ = gc_collect_minor().expect("collection failed");
576 assert_eq!(*protected, Value::Num(42.0));
577
578 gc_remove_root(protected).expect("root removal failed");
579 });
580 }
581
582 #[test]
583 fn test_gc_allocation_and_roots() {
584 gc_test_context(|| {
585 let v = gc_allocate(Value::Num(7.0)).expect("allocation failed");
586 assert_eq!(*v, Value::Num(7.0));
587
588 gc_add_root(v.clone()).expect("root add failed");
589 let _ = gc_collect_minor().expect("collection failed");
590 assert_eq!(*v, Value::Num(7.0));
591
592 gc_remove_root(v).expect("root remove failed");
593 });
594 }
595}