1use crate::vm::types::{HeapId, MemorySize, ObjectCount, ObjectSize};
2use crate::vm::value::Value;
3use std::collections::HashMap;
4
5#[derive(Debug, Clone)]
6pub struct HeapObject {
7 pub value: Value,
8 pub marked: bool,
9}
10
11impl HeapObject {
12 pub fn new(value: Value) -> Self {
13 Self {
14 value,
15 marked: false,
16 }
17 }
18}
19
20pub struct MemoryHeap {
21 objects: HashMap<HeapId, HeapObject>,
22 next_id: HeapId,
23 total_allocated: MemorySize,
24}
25
26impl Default for MemoryHeap {
27 fn default() -> Self {
28 Self::new()
29 }
30}
31
32impl MemoryHeap {
33 pub fn new() -> Self {
34 Self {
35 objects: HashMap::new(),
36 next_id: HeapId::new(0),
37 total_allocated: MemorySize::new(0),
38 }
39 }
40
41 pub fn allocate(&mut self, value: Value) -> HeapId {
42 let id = self.next_id;
43 self.next_id = self.next_id.next();
44
45 let size = MemorySize::new(self.calculate_size(&value).as_usize());
46 let object = HeapObject::new(value);
47 self.objects.insert(id, object);
48
49 self.total_allocated = self.total_allocated.add(size);
50 id
51 }
52
53 pub fn get(&self, id: HeapId) -> Option<&Value> {
54 self.objects.get(&id).map(|obj| &obj.value)
55 }
56
57 pub fn get_mut(&mut self, id: HeapId) -> Option<&mut Value> {
58 self.objects.get_mut(&id).map(|obj| &mut obj.value)
59 }
60
61 pub fn mark(&mut self, id: HeapId) {
62 if let Some(obj) = self.objects.get_mut(&id) {
63 obj.marked = true;
64 }
65 }
66
67 pub fn sweep(&mut self) -> MemorySize {
68 let mut freed = MemorySize::new(0);
69 let mut to_remove = Vec::new();
70
71 for (id, obj) in &self.objects {
72 if !obj.marked {
73 to_remove.push(*id);
74 freed = freed.add(MemorySize::new(self.calculate_size(&obj.value).as_usize()));
75 }
76 }
77
78 for id in to_remove {
79 self.objects.remove(&id);
80 }
81
82 self.total_allocated = self.total_allocated.sub(freed);
83 freed
84 }
85
86 pub fn remove(&mut self, id: HeapId) -> Option<Value> {
87 if let Some(obj) = self.objects.remove(&id) {
88 let size = MemorySize::new(self.calculate_size(&obj.value).as_usize());
89 self.total_allocated = self.total_allocated.sub(size);
90 Some(obj.value)
91 } else {
92 None
93 }
94 }
95
96 pub fn total_allocated(&self) -> MemorySize {
97 self.total_allocated
98 }
99
100 pub fn object_count(&self) -> ObjectCount {
101 ObjectCount::new(self.objects.len())
102 }
103
104 pub fn iter_objects(&self) -> impl Iterator<Item = (HeapId, &HeapObject)> {
105 self.objects.iter().map(|(&id, obj)| (id, obj))
106 }
107
108 pub fn iter_objects_mut(&mut self) -> impl Iterator<Item = (HeapId, &mut HeapObject)> {
109 self.objects.iter_mut().map(|(&id, obj)| (id, obj))
110 }
111
112 pub fn clear(&mut self) {
113 self.objects.clear();
114 self.total_allocated = MemorySize::new(0);
115 }
116
117 pub fn contains(&self, id: HeapId) -> bool {
118 self.objects.contains_key(&id)
119 }
120
121 pub fn is_marked(&self, id: HeapId) -> bool {
122 self.objects.get(&id).map(|obj| obj.marked).unwrap_or(false)
123 }
124
125 pub fn unmark_all(&mut self) {
126 for obj in self.objects.values_mut() {
127 obj.marked = false;
128 }
129 }
130
131 pub fn mark_all(&mut self) {
132 for obj in self.objects.values_mut() {
133 obj.marked = true;
134 }
135 }
136
137 pub fn get_marked_objects(&self) -> Vec<HeapId> {
138 self.objects
139 .iter()
140 .filter(|(_, obj)| obj.marked)
141 .map(|(&id, _)| id)
142 .collect()
143 }
144
145 pub fn get_unmarked_objects(&self) -> Vec<HeapId> {
146 self.objects
147 .iter()
148 .filter(|(_, obj)| !obj.marked)
149 .map(|(&id, _)| id)
150 .collect()
151 }
152
153 fn calculate_size(&self, value: &Value) -> ObjectSize {
154 match value {
155 Value::Number(_) => ObjectSize::new(8),
156 Value::String(s) => ObjectSize::new(s.len()),
157 Value::Boolean(_) => ObjectSize::new(1),
158 Value::Undefined | Value::Null => ObjectSize::new(0),
159 Value::Object(_) | Value::Array(_) | Value::Function(_) => ObjectSize::new(8),
160 }
161 }
162}