jetcrab\memory/
heap.rs

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}