Skip to main content

sim_lib_mutation/managed/
arena.rs

1/// Bounded storage for managed objects, independent of language and collector policy.
2pub 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    /// Creates an empty arena using the hard-capped retain policy.
15    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    /// Returns the tracing contract version.
29    pub const fn trace_contract_version(&self) -> TraceContractVersion {
30        TraceContractVersion::V1
31    }
32
33    /// Returns the number of live objects.
34    pub fn len(&self) -> usize {
35        self.objects.len()
36    }
37
38    /// Reports whether the arena contains no objects.
39    pub fn is_empty(&self) -> bool {
40        self.objects.is_empty()
41    }
42
43    /// Returns the epoch advanced by every graph-affecting arena mutation.
44    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    /// Allocates atomically after checking the cap and identity space.
57    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    /// Returns a shared object reference, refusing stale handles.
75    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    /// Returns a mutable object reference, refusing stale handles.
82    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    /// Upgrades a weak handle only while its object remains live.
94    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    /// Resolves a tracing identity to a live handle for collector operations.
103    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    /// Registers a root after validating the handle.
111    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    /// Releases exactly one matching root registration.
125    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    /// Removes an unrooted object, making all handles to it stale.
137    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    /// Clears a weak edge through the owning object's at-most-once operation.
153    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    /// Atomically removes an allocation-ordered set selected from `expected_epoch`.
175    ///
176    /// Every identity and root condition is checked before the first slot changes.
177    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    /// Applies a collector plan atomically at `expected_epoch`.
206    ///
207    /// Kept-alive objects from that epoch are retained. Weak and ephemeron
208    /// entries are cleared before unreachable objects are removed, and every
209    /// conditional clear is intrinsically at most once.
210    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    /// Runs a read-only tracing callback at a deterministic safepoint.
276    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    /// Tears down all storage and roots, returning allocation-ordered evidence.
309    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    /// Replaces caller-owned role evidence without advancing the graph epoch.
326    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    /// Projects optional owner-role evidence at a read-only safepoint.
338    ///
339    /// Admission is checked before invoking `role`, and rows follow managed id
340    /// order. The projection observes objects but is not visible to tracing or
341    /// collector policy.
342    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}