1use std::collections::BTreeMap;
2use std::error::Error;
3use std::fmt;
4
5#[derive(Clone, Copy, Debug, Eq, PartialEq)]
7pub enum TraceContractVersion {
8 V1,
10}
11
12#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd, Hash)]
14pub struct ManagedId(u64);
15
16impl ManagedId {
17 pub const fn allocation_ordinal(self) -> u64 {
19 self.0
20 }
21}
22
23#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd, Hash)]
25pub struct RootId(u64);
26
27#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
29pub struct ManagedHandle {
30 id: ManagedId,
31}
32
33impl ManagedHandle {
34 pub const fn id(self) -> ManagedId {
36 self.id
37 }
38
39 pub const fn downgrade(self) -> WeakHandle {
41 WeakHandle { id: self.id }
42 }
43}
44
45#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
48pub struct RootedHandle {
49 root: RootId,
50 handle: ManagedHandle,
51}
52
53impl RootedHandle {
54 pub const fn root_id(self) -> RootId {
56 self.root
57 }
58
59 pub const fn handle(self) -> ManagedHandle {
61 self.handle
62 }
63}
64
65#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
67pub struct WeakHandle {
68 id: ManagedId,
69}
70
71impl WeakHandle {
72 pub const fn id(self) -> ManagedId {
74 self.id
75 }
76}
77
78#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd, Hash)]
80pub struct EdgeId(pub u32);
81
82pub trait EdgeVisitor {
84 fn strong(&mut self, edge: EdgeId, target: ManagedId);
86
87 fn weak(&mut self, edge: EdgeId, target: ManagedId);
89
90 fn ephemeron(&mut self, edge: EdgeId, key: ManagedId, value: ManagedId);
92}
93
94pub trait ManagedObject {
96 fn trace_edges(&self, visitor: &mut dyn EdgeVisitor);
98
99 fn clear_weak_edge(&mut self, edge: EdgeId, expected: ManagedId) -> bool;
105
106 fn clear_ephemeron_edge(
109 &mut self,
110 _edge: EdgeId,
111 _expected_key: ManagedId,
112 _expected_value: ManagedId,
113 ) -> bool {
114 false
115 }
116}
117
118#[derive(Clone, Copy, Debug, Eq, PartialEq)]
121pub struct HardCappedRetainPolicy {
122 max_objects: usize,
123}
124
125impl HardCappedRetainPolicy {
126 pub fn new(max_objects: usize) -> Result<Self, ArenaError> {
128 if max_objects == 0 {
129 return Err(ArenaError::InvalidCap);
130 }
131 Ok(Self { max_objects })
132 }
133
134 pub const fn max_objects(self) -> usize {
136 self.max_objects
137 }
138}
139
140#[derive(Clone, Copy, Debug, Eq, PartialEq)]
142pub enum ArenaError {
143 InvalidCap,
145 CapacityExceeded {
147 cap: usize,
149 },
150 IdentityExhausted,
152 StaleHandle(ManagedId),
154 StaleRoot(RootId),
156 ObjectRooted(ManagedId),
158 MutationEpochChanged {
160 expected: u64,
162 actual: u64,
164 },
165}
166
167impl fmt::Display for ArenaError {
168 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
169 match self {
170 Self::InvalidCap => f.write_str("managed arena cap must be non-zero"),
171 Self::CapacityExceeded { cap } => write!(f, "managed arena hard cap {cap} reached"),
172 Self::IdentityExhausted => f.write_str("managed arena identity space exhausted"),
173 Self::StaleHandle(id) => write!(f, "stale managed handle {}", id.0),
174 Self::StaleRoot(id) => write!(f, "stale managed root {}", id.0),
175 Self::ObjectRooted(id) => write!(f, "managed object {} is rooted", id.0),
176 Self::MutationEpochChanged { expected, actual } => write!(
177 f,
178 "managed arena mutation epoch changed from {expected} to {actual}"
179 ),
180 }
181 }
182}
183
184impl Error for ArenaError {}
185
186pub struct TraceSnapshot<'a, T> {
188 roots: Vec<ManagedId>,
189 kept_alive: Vec<ManagedId>,
190 objects: &'a BTreeMap<ManagedId, T>,
191 mutation_epoch: u64,
192}
193
194impl<T: ManagedObject> TraceSnapshot<'_, T> {
195 pub const fn mutation_epoch(&self) -> u64 {
197 self.mutation_epoch
198 }
199 pub fn roots(&self) -> impl ExactSizeIterator<Item = ManagedId> + '_ {
201 self.roots.iter().copied()
202 }
203
204 pub fn kept_alive(&self) -> impl ExactSizeIterator<Item = ManagedId> + '_ {
206 self.kept_alive.iter().copied()
207 }
208
209 pub fn objects(&self) -> impl ExactSizeIterator<Item = ManagedId> + '_ {
211 self.objects.keys().copied()
212 }
213
214 pub fn visit_edges(
216 &self,
217 owner: ManagedId,
218 visitor: &mut dyn EdgeVisitor,
219 ) -> Result<(), ArenaError> {
220 self.objects
221 .get(&owner)
222 .ok_or(ArenaError::StaleHandle(owner))?
223 .trace_edges(visitor);
224 Ok(())
225 }
226}
227
228#[derive(Clone, Debug, Eq, PartialEq)]
230pub struct SafepointReceipt {
231 pub sequence: u64,
233 pub roots: Vec<ManagedId>,
235 pub objects: Vec<ManagedId>,
237}
238
239#[derive(Clone, Debug, Eq, PartialEq)]
241pub struct TeardownReceipt {
242 pub objects: Vec<ManagedId>,
244 pub roots: Vec<RootId>,
246}
247
248pub struct CollectionMutationReceipt {
250 pub cleared_weak: Vec<(ManagedId, EdgeId)>,
252 pub cleared_ephemerons: Vec<(ManagedId, EdgeId)>,
254 pub swept: Vec<ManagedId>,
256}
257
258pub struct ManagedArena<T> {
260 policy: HardCappedRetainPolicy,
261 next_id: u64,
262 next_root: u64,
263 next_safepoint: u64,
264 mutation_epoch: u64,
265 objects: BTreeMap<ManagedId, T>,
266 roots: BTreeMap<RootId, ManagedId>,
267 kept_alive: BTreeMap<ManagedId, u64>,
268}
269
270impl<T> ManagedArena<T> {
271 pub fn new(policy: HardCappedRetainPolicy) -> Self {
273 Self {
274 policy,
275 next_id: 0,
276 next_root: 0,
277 next_safepoint: 0,
278 mutation_epoch: 0,
279 objects: BTreeMap::new(),
280 roots: BTreeMap::new(),
281 kept_alive: BTreeMap::new(),
282 }
283 }
284
285 pub const fn trace_contract_version(&self) -> TraceContractVersion {
287 TraceContractVersion::V1
288 }
289
290 pub fn len(&self) -> usize {
292 self.objects.len()
293 }
294
295 pub fn is_empty(&self) -> bool {
297 self.objects.is_empty()
298 }
299
300 pub const fn mutation_epoch(&self) -> u64 {
302 self.mutation_epoch
303 }
304
305 fn advance_mutation_epoch(&mut self) -> Result<(), ArenaError> {
306 self.mutation_epoch = self
307 .mutation_epoch
308 .checked_add(1)
309 .ok_or(ArenaError::IdentityExhausted)?;
310 Ok(())
311 }
312
313 pub fn allocate(&mut self, object: T) -> Result<ManagedHandle, ArenaError> {
315 if self.objects.len() >= self.policy.max_objects {
316 return Err(ArenaError::CapacityExceeded {
317 cap: self.policy.max_objects,
318 });
319 }
320 let next = self
321 .next_id
322 .checked_add(1)
323 .ok_or(ArenaError::IdentityExhausted)?;
324 let id = ManagedId(self.next_id);
325 self.advance_mutation_epoch()?;
326 self.objects.insert(id, object);
327 self.next_id = next;
328 Ok(ManagedHandle { id })
329 }
330
331 pub fn get(&self, handle: ManagedHandle) -> Result<&T, ArenaError> {
333 self.objects
334 .get(&handle.id)
335 .ok_or(ArenaError::StaleHandle(handle.id))
336 }
337
338 pub fn get_mut(&mut self, handle: ManagedHandle) -> Result<&mut T, ArenaError> {
340 if !self.objects.contains_key(&handle.id) {
341 return Err(ArenaError::StaleHandle(handle.id));
342 }
343 self.advance_mutation_epoch()?;
344 Ok(self
345 .objects
346 .get_mut(&handle.id)
347 .expect("validated managed id"))
348 }
349
350 pub fn upgrade(&mut self, weak: WeakHandle) -> Result<ManagedHandle, ArenaError> {
352 if !self.objects.contains_key(&weak.id) {
353 return Err(ArenaError::StaleHandle(weak.id));
354 }
355 self.kept_alive.insert(weak.id, self.mutation_epoch);
356 Ok(ManagedHandle { id: weak.id })
357 }
358
359 pub fn handle(&self, id: ManagedId) -> Result<ManagedHandle, ArenaError> {
361 self.objects
362 .contains_key(&id)
363 .then_some(ManagedHandle { id })
364 .ok_or(ArenaError::StaleHandle(id))
365 }
366
367 pub fn root(&mut self, handle: ManagedHandle) -> Result<RootedHandle, ArenaError> {
369 self.get(handle)?;
370 let next = self
371 .next_root
372 .checked_add(1)
373 .ok_or(ArenaError::IdentityExhausted)?;
374 let root = RootId(self.next_root);
375 self.advance_mutation_epoch()?;
376 self.roots.insert(root, handle.id);
377 self.next_root = next;
378 Ok(RootedHandle { root, handle })
379 }
380
381 pub fn release_root(&mut self, rooted: RootedHandle) -> Result<ManagedHandle, ArenaError> {
383 match self.roots.get(&rooted.root) {
384 Some(id) if *id == rooted.handle.id => {
385 self.advance_mutation_epoch()?;
386 self.roots.remove(&rooted.root);
387 Ok(rooted.handle)
388 }
389 _ => Err(ArenaError::StaleRoot(rooted.root)),
390 }
391 }
392
393 pub fn remove(&mut self, handle: ManagedHandle) -> Result<T, ArenaError> {
395 if self.roots.values().any(|id| *id == handle.id) {
396 return Err(ArenaError::ObjectRooted(handle.id));
397 }
398 if !self.objects.contains_key(&handle.id) {
399 return Err(ArenaError::StaleHandle(handle.id));
400 }
401 self.advance_mutation_epoch()?;
402 let removed = self
403 .objects
404 .remove(&handle.id)
405 .expect("validated managed id");
406 Ok(removed)
407 }
408
409 pub fn clear_weak_edge(
411 &mut self,
412 owner: ManagedHandle,
413 edge: EdgeId,
414 expected: WeakHandle,
415 ) -> Result<bool, ArenaError>
416 where
417 T: ManagedObject,
418 {
419 if !self.objects.contains_key(&owner.id) {
420 return Err(ArenaError::StaleHandle(owner.id));
421 }
422 self.advance_mutation_epoch()?;
423 let cleared = self
424 .objects
425 .get_mut(&owner.id)
426 .expect("validated managed id")
427 .clear_weak_edge(edge, expected.id);
428 Ok(cleared)
429 }
430
431 pub fn sweep_at_epoch(
435 &mut self,
436 expected_epoch: u64,
437 objects: &[ManagedId],
438 ) -> Result<Vec<ManagedId>, ArenaError> {
439 if self.mutation_epoch != expected_epoch {
440 return Err(ArenaError::MutationEpochChanged {
441 expected: expected_epoch,
442 actual: self.mutation_epoch,
443 });
444 }
445 for id in objects {
446 if !self.objects.contains_key(id) {
447 return Err(ArenaError::StaleHandle(*id));
448 }
449 if self.roots.values().any(|rooted| rooted == id) {
450 return Err(ArenaError::ObjectRooted(*id));
451 }
452 }
453 if !objects.is_empty() {
454 self.advance_mutation_epoch()?;
455 }
456 for id in objects {
457 self.objects.remove(id);
458 }
459 Ok(objects.to_vec())
460 }
461
462 pub fn apply_collection_at_epoch(
468 &mut self,
469 expected_epoch: u64,
470 weak: &[(ManagedId, EdgeId, ManagedId)],
471 ephemerons: &[(ManagedId, EdgeId, ManagedId, ManagedId)],
472 swept: &[ManagedId],
473 ) -> Result<CollectionMutationReceipt, ArenaError>
474 where
475 T: ManagedObject,
476 {
477 if self.mutation_epoch != expected_epoch {
478 return Err(ArenaError::MutationEpochChanged {
479 expected: expected_epoch,
480 actual: self.mutation_epoch,
481 });
482 }
483 let kept = self
484 .kept_alive
485 .iter()
486 .filter_map(|(id, epoch)| (*epoch == expected_epoch).then_some(*id))
487 .collect::<std::collections::BTreeSet<_>>();
488 let actual_swept = swept
489 .iter()
490 .copied()
491 .filter(|id| !kept.contains(id))
492 .collect::<Vec<_>>();
493 for id in &actual_swept {
494 if !self.objects.contains_key(id) {
495 return Err(ArenaError::StaleHandle(*id));
496 }
497 if self.roots.values().any(|rooted| rooted == id) {
498 return Err(ArenaError::ObjectRooted(*id));
499 }
500 }
501 if !weak.is_empty() || !ephemerons.is_empty() || !actual_swept.is_empty() {
502 self.advance_mutation_epoch()?;
503 }
504 let mut cleared_weak = Vec::new();
505 for &(owner, edge, target) in weak {
506 if let Some(object) = self.objects.get_mut(&owner)
507 && object.clear_weak_edge(edge, target)
508 {
509 cleared_weak.push((owner, edge));
510 }
511 }
512 let mut cleared_ephemerons = Vec::new();
513 for &(owner, edge, key, value) in ephemerons {
514 if let Some(object) = self.objects.get_mut(&owner)
515 && object.clear_ephemeron_edge(edge, key, value)
516 {
517 cleared_ephemerons.push((owner, edge));
518 }
519 }
520 for id in &actual_swept {
521 self.objects.remove(id);
522 }
523 self.kept_alive
524 .retain(|id, epoch| self.objects.contains_key(id) && *epoch != expected_epoch);
525 Ok(CollectionMutationReceipt {
526 cleared_weak,
527 cleared_ephemerons,
528 swept: actual_swept,
529 })
530 }
531
532 pub fn safepoint<R>(
534 &mut self,
535 trace: impl FnOnce(&TraceSnapshot<'_, T>) -> R,
536 ) -> Result<(R, SafepointReceipt), ArenaError>
537 where
538 T: ManagedObject,
539 {
540 let next = self
541 .next_safepoint
542 .checked_add(1)
543 .ok_or(ArenaError::IdentityExhausted)?;
544 let roots = self.roots.values().copied().collect::<Vec<_>>();
545 let snapshot = TraceSnapshot {
546 roots: roots.clone(),
547 kept_alive: self
548 .kept_alive
549 .iter()
550 .filter_map(|(id, epoch)| (*epoch == self.mutation_epoch).then_some(*id))
551 .collect(),
552 objects: &self.objects,
553 mutation_epoch: self.mutation_epoch,
554 };
555 let result = trace(&snapshot);
556 let receipt = SafepointReceipt {
557 sequence: self.next_safepoint,
558 roots,
559 objects: self.objects.keys().copied().collect(),
560 };
561 self.next_safepoint = next;
562 Ok((result, receipt))
563 }
564
565 pub fn teardown(&mut self) -> TeardownReceipt {
567 let receipt = TeardownReceipt {
568 objects: self.objects.keys().copied().collect(),
569 roots: self.roots.keys().copied().collect(),
570 };
571 if !self.objects.is_empty() || !self.roots.is_empty() {
572 self.mutation_epoch = self.mutation_epoch.saturating_add(1);
573 }
574 self.objects.clear();
575 self.roots.clear();
576 self.kept_alive.clear();
577 receipt
578 }
579}