sim_lib_mutation/managed/
arena.rs1pub struct ManagedArena<T> {
3 policy: HardCappedRetainPolicy,
4 next_id: u64,
5 next_root: u64,
6 next_safepoint: u64,
7 mutation_epoch: u64,
8 objects: BTreeMap<ManagedId, T>,
9 roots: BTreeMap<RootId, ManagedId>,
10 kept_alive: BTreeMap<ManagedId, u64>,
11}
12
13impl<T> ManagedArena<T> {
14 pub fn new(policy: HardCappedRetainPolicy) -> Self {
16 Self {
17 policy,
18 next_id: 0,
19 next_root: 0,
20 next_safepoint: 0,
21 mutation_epoch: 0,
22 objects: BTreeMap::new(),
23 roots: BTreeMap::new(),
24 kept_alive: BTreeMap::new(),
25 }
26 }
27
28 pub const fn trace_contract_version(&self) -> TraceContractVersion {
30 TraceContractVersion::V1
31 }
32
33 pub fn len(&self) -> usize {
35 self.objects.len()
36 }
37
38 pub fn is_empty(&self) -> bool {
40 self.objects.is_empty()
41 }
42
43 pub const fn mutation_epoch(&self) -> u64 {
45 self.mutation_epoch
46 }
47
48 fn advance_mutation_epoch(&mut self) -> Result<(), ArenaError> {
49 self.mutation_epoch = self
50 .mutation_epoch
51 .checked_add(1)
52 .ok_or(ArenaError::IdentityExhausted)?;
53 Ok(())
54 }
55
56 pub fn allocate(&mut self, object: T) -> Result<ManagedHandle, ArenaError> {
58 if self.objects.len() >= self.policy.max_objects {
59 return Err(ArenaError::CapacityExceeded {
60 cap: self.policy.max_objects,
61 });
62 }
63 let next = self
64 .next_id
65 .checked_add(1)
66 .ok_or(ArenaError::IdentityExhausted)?;
67 let id = ManagedId(self.next_id);
68 self.advance_mutation_epoch()?;
69 self.objects.insert(id, object);
70 self.next_id = next;
71 Ok(ManagedHandle { id })
72 }
73
74 pub fn get(&self, handle: ManagedHandle) -> Result<&T, ArenaError> {
76 self.objects
77 .get(&handle.id)
78 .ok_or(ArenaError::StaleHandle(handle.id))
79 }
80
81 pub fn get_mut(&mut self, handle: ManagedHandle) -> Result<&mut T, ArenaError> {
83 if !self.objects.contains_key(&handle.id) {
84 return Err(ArenaError::StaleHandle(handle.id));
85 }
86 self.advance_mutation_epoch()?;
87 Ok(self
88 .objects
89 .get_mut(&handle.id)
90 .expect("validated managed id"))
91 }
92
93 pub fn upgrade(&mut self, weak: WeakHandle) -> Result<ManagedHandle, ArenaError> {
95 if !self.objects.contains_key(&weak.id) {
96 return Err(ArenaError::StaleHandle(weak.id));
97 }
98 self.kept_alive.insert(weak.id, self.mutation_epoch);
99 Ok(ManagedHandle { id: weak.id })
100 }
101
102 pub fn handle(&self, id: ManagedId) -> Result<ManagedHandle, ArenaError> {
104 self.objects
105 .contains_key(&id)
106 .then_some(ManagedHandle { id })
107 .ok_or(ArenaError::StaleHandle(id))
108 }
109
110 pub fn root(&mut self, handle: ManagedHandle) -> Result<RootedHandle, ArenaError> {
112 self.get(handle)?;
113 let next = self
114 .next_root
115 .checked_add(1)
116 .ok_or(ArenaError::IdentityExhausted)?;
117 let root = RootId(self.next_root);
118 self.advance_mutation_epoch()?;
119 self.roots.insert(root, handle.id);
120 self.next_root = next;
121 Ok(RootedHandle { root, handle })
122 }
123
124 pub fn release_root(&mut self, rooted: RootedHandle) -> Result<ManagedHandle, ArenaError> {
126 match self.roots.get(&rooted.root) {
127 Some(id) if *id == rooted.handle.id => {
128 self.advance_mutation_epoch()?;
129 self.roots.remove(&rooted.root);
130 Ok(rooted.handle)
131 }
132 _ => Err(ArenaError::StaleRoot(rooted.root)),
133 }
134 }
135
136 pub fn remove(&mut self, handle: ManagedHandle) -> Result<T, ArenaError> {
138 if self.roots.values().any(|id| *id == handle.id) {
139 return Err(ArenaError::ObjectRooted(handle.id));
140 }
141 if !self.objects.contains_key(&handle.id) {
142 return Err(ArenaError::StaleHandle(handle.id));
143 }
144 self.advance_mutation_epoch()?;
145 let removed = self
146 .objects
147 .remove(&handle.id)
148 .expect("validated managed id");
149 Ok(removed)
150 }
151
152 pub fn clear_weak_edge(
154 &mut self,
155 owner: ManagedHandle,
156 edge: EdgeId,
157 expected: WeakHandle,
158 ) -> Result<bool, ArenaError>
159 where
160 T: ManagedObject,
161 {
162 if !self.objects.contains_key(&owner.id) {
163 return Err(ArenaError::StaleHandle(owner.id));
164 }
165 self.advance_mutation_epoch()?;
166 let cleared = self
167 .objects
168 .get_mut(&owner.id)
169 .expect("validated managed id")
170 .clear_weak_edge(edge, expected.id);
171 Ok(cleared)
172 }
173
174 pub fn sweep_at_epoch(
178 &mut self,
179 expected_epoch: u64,
180 objects: &[ManagedId],
181 ) -> Result<Vec<ManagedId>, ArenaError> {
182 if self.mutation_epoch != expected_epoch {
183 return Err(ArenaError::MutationEpochChanged {
184 expected: expected_epoch,
185 actual: self.mutation_epoch,
186 });
187 }
188 for id in objects {
189 if !self.objects.contains_key(id) {
190 return Err(ArenaError::StaleHandle(*id));
191 }
192 if self.roots.values().any(|rooted| rooted == id) {
193 return Err(ArenaError::ObjectRooted(*id));
194 }
195 }
196 if !objects.is_empty() {
197 self.advance_mutation_epoch()?;
198 }
199 for id in objects {
200 self.objects.remove(id);
201 }
202 Ok(objects.to_vec())
203 }
204
205 pub fn apply_collection_at_epoch(
211 &mut self,
212 expected_epoch: u64,
213 weak: &[(ManagedId, EdgeId, ManagedId)],
214 ephemerons: &[(ManagedId, EdgeId, ManagedId, ManagedId)],
215 swept: &[ManagedId],
216 ) -> Result<CollectionMutationReceipt, ArenaError>
217 where
218 T: ManagedObject,
219 {
220 if self.mutation_epoch != expected_epoch {
221 return Err(ArenaError::MutationEpochChanged {
222 expected: expected_epoch,
223 actual: self.mutation_epoch,
224 });
225 }
226 let kept = self
227 .kept_alive
228 .iter()
229 .filter_map(|(id, epoch)| (*epoch == expected_epoch).then_some(*id))
230 .collect::<std::collections::BTreeSet<_>>();
231 let actual_swept = swept
232 .iter()
233 .copied()
234 .filter(|id| !kept.contains(id))
235 .collect::<Vec<_>>();
236 for id in &actual_swept {
237 if !self.objects.contains_key(id) {
238 return Err(ArenaError::StaleHandle(*id));
239 }
240 if self.roots.values().any(|rooted| rooted == id) {
241 return Err(ArenaError::ObjectRooted(*id));
242 }
243 }
244 if !weak.is_empty() || !ephemerons.is_empty() || !actual_swept.is_empty() {
245 self.advance_mutation_epoch()?;
246 }
247 let mut cleared_weak = Vec::new();
248 for &(owner, edge, target) in weak {
249 if let Some(object) = self.objects.get_mut(&owner)
250 && object.clear_weak_edge(edge, target)
251 {
252 cleared_weak.push((owner, edge));
253 }
254 }
255 let mut cleared_ephemerons = Vec::new();
256 for &(owner, edge, key, value) in ephemerons {
257 if let Some(object) = self.objects.get_mut(&owner)
258 && object.clear_ephemeron_edge(edge, key, value)
259 {
260 cleared_ephemerons.push((owner, edge));
261 }
262 }
263 for id in &actual_swept {
264 self.objects.remove(id);
265 }
266 self.kept_alive
267 .retain(|id, epoch| self.objects.contains_key(id) && *epoch != expected_epoch);
268 Ok(CollectionMutationReceipt {
269 cleared_weak,
270 cleared_ephemerons,
271 swept: actual_swept,
272 })
273 }
274
275 pub fn safepoint<R>(
277 &mut self,
278 trace: impl FnOnce(&TraceSnapshot<'_, T>) -> R,
279 ) -> Result<(R, SafepointReceipt), ArenaError>
280 where
281 T: ManagedObject,
282 {
283 let next = self
284 .next_safepoint
285 .checked_add(1)
286 .ok_or(ArenaError::IdentityExhausted)?;
287 let roots = self.roots.values().copied().collect::<Vec<_>>();
288 let snapshot = TraceSnapshot {
289 roots: roots.clone(),
290 kept_alive: self
291 .kept_alive
292 .iter()
293 .filter_map(|(id, epoch)| (*epoch == self.mutation_epoch).then_some(*id))
294 .collect(),
295 objects: &self.objects,
296 mutation_epoch: self.mutation_epoch,
297 };
298 let result = trace(&snapshot);
299 let receipt = SafepointReceipt {
300 sequence: self.next_safepoint,
301 roots,
302 objects: self.objects.keys().copied().collect(),
303 };
304 self.next_safepoint = next;
305 Ok((result, receipt))
306 }
307
308 pub fn teardown(&mut self) -> TeardownReceipt {
310 let receipt = TeardownReceipt {
311 objects: self.objects.keys().copied().collect(),
312 roots: self.roots.keys().copied().collect(),
313 };
314 if !self.objects.is_empty() || !self.roots.is_empty() {
315 self.mutation_epoch = self.mutation_epoch.saturating_add(1);
316 }
317 self.objects.clear();
318 self.roots.clear();
319 self.kept_alive.clear();
320 receipt
321 }
322}
323
324impl<T: RoleBearingManagedObject> ManagedArena<T> {
325 pub fn replace_role(
327 &mut self,
328 handle: ManagedHandle,
329 role: T::Role,
330 ) -> Result<T::Role, ArenaError> {
331 self.objects
332 .get_mut(&handle.id)
333 .map(|object| object.replace_managed_role(role))
334 .ok_or(ArenaError::StaleHandle(handle.id))
335 }
336
337 pub fn project_roles(
343 &mut self,
344 limit: usize,
345 ) -> Result<RoleProjectionReceipt<T::Role>, RoleProjectionError>
346 where
347 T::Role: Clone,
348 {
349 let (roles, safepoint) = self.safepoint(|snapshot| {
350 let required = snapshot.objects.len();
351 if required > limit {
352 return Err(RoleProjectionError::Limit { limit, required });
353 }
354 let roles = snapshot
355 .objects
356 .iter()
357 .map(|(&id, object)| (id, object.managed_role().clone()))
358 .collect();
359 Ok((snapshot.mutation_epoch(), roles))
360 })?;
361 let (mutation_epoch, roles) = roles?;
362 Ok(RoleProjectionReceipt {
363 safepoint,
364 mutation_epoch,
365 roles,
366 })
367 }
368}