jetcrab\memory/
collector.rs1use crate::memory::heap::MemoryHeap;
2use crate::vm::types::{AllocationCount, HeapId, MemorySize, ObjectId};
3use crate::vm::value::Value;
4use std::collections::HashSet;
5
6#[derive(Debug, Clone, PartialEq)]
7pub enum GCStrategy {
8 MarkAndSweep,
9 Generational,
10 Incremental,
11}
12
13#[derive(Debug, Clone, Default)]
14pub struct GCStats {
15 pub total_collections: AllocationCount,
16 pub total_bytes_freed: MemorySize,
17 pub total_collection_time_ms: u64,
18 pub last_collection_time_ms: u64,
19 pub objects_marked: AllocationCount,
20 pub objects_swept: AllocationCount,
21}
22
23pub struct GarbageCollector {
24 heap: MemoryHeap,
25
26 strategy: GCStrategy,
27
28 stats: GCStats,
29
30 marked_objects: HashSet<ObjectId>,
31
32 root_objects: HashSet<ObjectId>,
33
34 is_collecting: bool,
35
36 collection_threshold: MemorySize,
37
38 allocated_bytes: MemorySize,
39}
40
41impl GarbageCollector {
42 pub fn new(heap: MemoryHeap) -> Self {
43 Self {
44 heap,
45 strategy: GCStrategy::MarkAndSweep,
46 stats: GCStats::default(),
47 marked_objects: HashSet::new(),
48 root_objects: HashSet::new(),
49 is_collecting: false,
50 collection_threshold: MemorySize::new(1024 * 1024),
51 allocated_bytes: MemorySize::new(0),
52 }
53 }
54
55 pub fn with_strategy(mut self, strategy: GCStrategy) -> Self {
56 self.strategy = strategy;
57 self
58 }
59
60 pub fn with_threshold(mut self, threshold: MemorySize) -> Self {
61 self.collection_threshold = threshold;
62 self
63 }
64
65 pub fn allocate(&mut self, value: Value) -> ObjectId {
66 let object_id = self.heap.allocate(value);
67 let size = self.heap.total_allocated() - self.allocated_bytes;
68
69 self.allocated_bytes = self.allocated_bytes.add(size);
70
71 if self.allocated_bytes.as_usize() > self.collection_threshold.as_usize() {
72 self.collect();
73 }
74
75 ObjectId::new(object_id.as_usize())
76 }
77
78 pub fn add_root(&mut self, object_id: ObjectId) {
79 self.root_objects.insert(object_id);
80 }
81
82 pub fn remove_root(&mut self, object_id: ObjectId) {
83 self.root_objects.remove(&object_id);
84 }
85
86 pub fn get(&self, object_id: ObjectId) -> Option<&Value> {
87 self.heap.get(HeapId::from(object_id.as_usize()))
88 }
89
90 pub fn get_mut(&mut self, object_id: ObjectId) -> Option<&mut Value> {
91 self.heap.get_mut(HeapId::from(object_id.as_usize()))
92 }
93
94 pub fn collect(&mut self) -> GCStats {
95 if self.is_collecting {
96 return self.stats.clone();
97 }
98
99 self.is_collecting = true;
100 let start_time = std::time::Instant::now();
101
102 match self.strategy {
103 GCStrategy::MarkAndSweep => self.mark_and_sweep(),
104 GCStrategy::Generational => self.generational_collect(),
105 GCStrategy::Incremental => self.incremental_collect(),
106 }
107
108 let collection_time = start_time.elapsed().as_millis() as u64;
109 self.stats.last_collection_time_ms = collection_time;
110 self.stats.total_collection_time_ms += collection_time;
111 self.stats.total_collections.increment();
112
113 self.is_collecting = false;
114 self.stats.clone()
115 }
116
117 fn mark_and_sweep(&mut self) {
118 self.mark_phase();
119
120 let freed_bytes = self.sweep_phase();
121
122 self.stats.total_bytes_freed = self.stats.total_bytes_freed.add(freed_bytes);
123 self.allocated_bytes = self.allocated_bytes.sub(freed_bytes);
124 }
125
126 fn mark_phase(&mut self) {
127 self.marked_objects.clear();
128
129 let root_objects: Vec<ObjectId> = self.root_objects.iter().cloned().collect();
130 for root_id in root_objects {
131 self.mark_object(root_id);
132 }
133
134 self.stats.objects_marked = AllocationCount::new(self.marked_objects.len());
135 }
136
137 fn mark_object(&mut self, object_id: ObjectId) {
138 if self.marked_objects.contains(&object_id) {
139 return;
140 }
141
142 self.marked_objects.insert(object_id);
143
144 if let Some(value) = self.heap.get(HeapId::from(object_id.as_usize())) {
145 match value {
146 Value::Object(object_id) => self.mark_object(ObjectId::new(object_id.as_usize())),
147 Value::Array(array_id) => self.mark_object(ObjectId::new(array_id.as_usize())),
148 Value::Function(function_id) => {
149 self.mark_object(ObjectId::new(function_id.as_usize()))
150 }
151 _ => {}
152 }
153 }
154 }
155
156 fn sweep_phase(&mut self) -> MemorySize {
157 let mut freed_bytes = MemorySize::new(0);
158 let mut to_remove = Vec::new();
159
160 for (object_id, _) in self.heap.iter_objects() {
161 if !self
162 .marked_objects
163 .contains(&ObjectId::new(object_id.as_usize()))
164 {
165 to_remove.push(object_id);
166 }
167 }
168
169 let objects_to_remove = to_remove.len();
170
171 for object_id in to_remove {
172 if let Some(_value) = self.heap.get(object_id) {
173 freed_bytes = freed_bytes.add(MemorySize::new(
174 self.heap.total_allocated().as_usize()
175 / self.heap.object_count().as_usize().max(1),
176 ));
177 }
178 self.heap.remove(object_id);
179 }
180
181 self.stats.objects_swept = AllocationCount::new(objects_to_remove);
182 freed_bytes
183 }
184
185 fn generational_collect(&mut self) {
186 self.mark_and_sweep();
187 }
188
189 fn incremental_collect(&mut self) {
190 self.mark_and_sweep();
191 }
192
193 pub fn stats(&self) -> &GCStats {
194 &self.stats
195 }
196
197 pub fn heap(&self) -> &MemoryHeap {
198 &self.heap
199 }
200
201 pub fn heap_mut(&mut self) -> &mut MemoryHeap {
202 &mut self.heap
203 }
204
205 pub fn is_marked(&self, object_id: ObjectId) -> bool {
206 self.marked_objects.contains(&object_id)
207 }
208
209 pub fn allocated_bytes(&self) -> MemorySize {
210 self.allocated_bytes
211 }
212
213 pub fn set_threshold(&mut self, threshold: MemorySize) {
214 self.collection_threshold = threshold;
215 }
216
217 pub fn force_collect(&mut self) -> GCStats {
218 self.collect()
219 }
220}
221
222impl Default for GarbageCollector {
223 fn default() -> Self {
224 Self::new(MemoryHeap::new())
225 }
226}