Skip to main content

sim_lib_mutation/
managed.rs

1use std::collections::BTreeMap;
2use std::error::Error;
3use std::fmt;
4
5/// The tracing ABI understood by this arena.
6#[derive(Clone, Copy, Debug, Eq, PartialEq)]
7pub enum TraceContractVersion {
8    /// Roots, strong and weak edges, ephemerons, safepoints, clearing, and teardown.
9    V1,
10}
11
12/// A stable managed-object identity assigned from allocation order.
13#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd, Hash)]
14pub struct ManagedId(u64);
15
16impl ManagedId {
17    /// Returns the zero-based allocation ordinal.
18    pub const fn allocation_ordinal(self) -> u64 {
19        self.0
20    }
21}
22
23/// A stable identity for one root registration.
24#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd, Hash)]
25pub struct RootId(u64);
26
27/// An object handle. It does not itself keep the object rooted.
28#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
29pub struct ManagedHandle {
30    id: ManagedId,
31}
32
33impl ManagedHandle {
34    /// Returns the managed identity.
35    pub const fn id(self) -> ManagedId {
36        self.id
37    }
38
39    /// Produces a non-rooting weak handle.
40    pub const fn downgrade(self) -> WeakHandle {
41        WeakHandle { id: self.id }
42    }
43}
44
45/// A registered root handle. Dropping this value does not mutate the arena;
46/// callers explicitly release it so root changes remain transactional.
47#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
48pub struct RootedHandle {
49    root: RootId,
50    handle: ManagedHandle,
51}
52
53impl RootedHandle {
54    /// Returns the root registration identity.
55    pub const fn root_id(self) -> RootId {
56        self.root
57    }
58
59    /// Returns the underlying object handle.
60    pub const fn handle(self) -> ManagedHandle {
61        self.handle
62    }
63}
64
65/// A non-rooting object handle which may become stale.
66#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
67pub struct WeakHandle {
68    id: ManagedId,
69}
70
71impl WeakHandle {
72    /// Returns the identity without claiming the object is still live.
73    pub const fn id(self) -> ManagedId {
74        self.id
75    }
76}
77
78/// Stable identity of an edge within its owning object.
79#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd, Hash)]
80pub struct EdgeId(pub u32);
81
82/// Receives every outgoing managed edge of an object.
83pub trait EdgeVisitor {
84    /// Visits a retaining edge.
85    fn strong(&mut self, edge: EdgeId, target: ManagedId);
86
87    /// Visits a non-retaining edge that may be cleared after tracing.
88    fn weak(&mut self, edge: EdgeId, target: ManagedId);
89
90    /// Visits a value retained only when `key` is reachable.
91    fn ephemeron(&mut self, edge: EdgeId, key: ManagedId, value: ManagedId);
92}
93
94/// An object stored by [`ManagedArena`].
95pub trait ManagedObject {
96    /// Enumerates all strong, weak, and ephemeron edges exactly once.
97    fn trace_edges(&self, visitor: &mut dyn EdgeVisitor);
98
99    /// Clears one weak edge if it still points at `expected`.
100    ///
101    /// Returning `true` means this invocation performed the clear. Repeating
102    /// the same request must return `false`, giving collectors at-most-once
103    /// weak-clear semantics.
104    fn clear_weak_edge(&mut self, edge: EdgeId, expected: ManagedId) -> bool;
105
106    /// Clears one ephemeron entry if it still has the expected key and value.
107    /// Repeating a successful request must return `false`.
108    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/// The only built-in policy: retain objects until explicit teardown, while
119/// refusing allocations beyond a fixed hard cap.
120#[derive(Clone, Copy, Debug, Eq, PartialEq)]
121pub struct HardCappedRetainPolicy {
122    max_objects: usize,
123}
124
125impl HardCappedRetainPolicy {
126    /// Creates a retain policy with a non-zero object cap.
127    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    /// Returns the allocation cap.
135    pub const fn max_objects(self) -> usize {
136        self.max_objects
137    }
138}
139
140/// Fail-closed arena operation errors.
141#[derive(Clone, Copy, Debug, Eq, PartialEq)]
142pub enum ArenaError {
143    /// A zero-sized arena was requested.
144    InvalidCap,
145    /// Allocation would exceed the hard object cap.
146    CapacityExceeded {
147        /// Configured maximum number of live objects.
148        cap: usize,
149    },
150    /// The allocation or root identity space is exhausted.
151    IdentityExhausted,
152    /// A handle no longer names a live object.
153    StaleHandle(ManagedId),
154    /// A root registration is unknown or does not match the handle.
155    StaleRoot(RootId),
156    /// A rooted object cannot be removed.
157    ObjectRooted(ManagedId),
158    /// Collection was planned against a different graph state.
159    MutationEpochChanged {
160        /// Epoch used to prepare the operation.
161        expected: u64,
162        /// Current arena epoch.
163        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
186/// An immutable, complete tracing view taken at a safepoint.
187pub 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    /// Returns the arena mutation epoch captured by this snapshot.
196    pub const fn mutation_epoch(&self) -> u64 {
197        self.mutation_epoch
198    }
199    /// Enumerates roots in root-registration order.
200    pub fn roots(&self) -> impl ExactSizeIterator<Item = ManagedId> + '_ {
201        self.roots.iter().copied()
202    }
203
204    /// Enumerates successful weak dereferences kept alive for this epoch.
205    pub fn kept_alive(&self) -> impl ExactSizeIterator<Item = ManagedId> + '_ {
206        self.kept_alive.iter().copied()
207    }
208
209    /// Enumerates live objects in allocation order.
210    pub fn objects(&self) -> impl ExactSizeIterator<Item = ManagedId> + '_ {
211        self.objects.keys().copied()
212    }
213
214    /// Visits all edges for a live object.
215    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/// Deterministic evidence for one tracing safepoint.
229#[derive(Clone, Debug, Eq, PartialEq)]
230pub struct SafepointReceipt {
231    /// Monotonic zero-based safepoint sequence.
232    pub sequence: u64,
233    /// Roots in root-registration order.
234    pub roots: Vec<ManagedId>,
235    /// Live objects in allocation order.
236    pub objects: Vec<ManagedId>,
237}
238
239/// Deterministic evidence returned by explicit arena teardown.
240#[derive(Clone, Debug, Eq, PartialEq)]
241pub struct TeardownReceipt {
242    /// Objects removed in allocation order.
243    pub objects: Vec<ManagedId>,
244    /// Root registrations removed in registration order.
245    pub roots: Vec<RootId>,
246}
247
248/// Atomic collector mutation evidence.
249pub struct CollectionMutationReceipt {
250    /// Weak entries cleared as owner and edge identities.
251    pub cleared_weak: Vec<(ManagedId, EdgeId)>,
252    /// Ephemeron entries cleared as owner and edge identities.
253    pub cleared_ephemerons: Vec<(ManagedId, EdgeId)>,
254    /// Objects removed in allocation order.
255    pub swept: Vec<ManagedId>,
256}
257
258/// Bounded storage for managed objects, independent of language and collector policy.
259pub 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    /// Creates an empty arena using the hard-capped retain policy.
272    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    /// Returns the tracing contract version.
286    pub const fn trace_contract_version(&self) -> TraceContractVersion {
287        TraceContractVersion::V1
288    }
289
290    /// Returns the number of live objects.
291    pub fn len(&self) -> usize {
292        self.objects.len()
293    }
294
295    /// Reports whether the arena contains no objects.
296    pub fn is_empty(&self) -> bool {
297        self.objects.is_empty()
298    }
299
300    /// Returns the epoch advanced by every graph-affecting arena mutation.
301    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    /// Allocates atomically after checking the cap and identity space.
314    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    /// Returns a shared object reference, refusing stale handles.
332    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    /// Returns a mutable object reference, refusing stale handles.
339    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    /// Upgrades a weak handle only while its object remains live.
351    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    /// Resolves a tracing identity to a live handle for collector operations.
360    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    /// Registers a root after validating the handle.
368    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    /// Releases exactly one matching root registration.
382    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    /// Removes an unrooted object, making all handles to it stale.
394    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    /// Clears a weak edge through the owning object's at-most-once operation.
410    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    /// Atomically removes an allocation-ordered set selected from `expected_epoch`.
432    ///
433    /// Every identity and root condition is checked before the first slot changes.
434    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    /// Applies a collector plan atomically at `expected_epoch`.
463    ///
464    /// Kept-alive objects from that epoch are retained. Weak and ephemeron
465    /// entries are cleared before unreachable objects are removed, and every
466    /// conditional clear is intrinsically at most once.
467    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    /// Runs a read-only tracing callback at a deterministic safepoint.
533    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    /// Tears down all storage and roots, returning allocation-ordered evidence.
566    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}