jstd/
stable_arena.rs

1//! Dense payload storage with stable, never-reused logical identifiers.
2//!
3//! [`StableArena`] separates identity from physical position. Removing an entry
4//! eagerly drops its payload with `swap_remove`, while a location table keeps
5//! every surviving identifier valid. Unlike [`Registry`](crate::registry::Registry),
6//! payload addresses are not stable across mutation.
7
8use std::{
9    fmt::{Debug, Formatter},
10    iter::Zip,
11    ops::{Index, IndexMut},
12    slice, vec,
13};
14
15use serde::{Deserialize, Deserializer, Serialize, Serializer};
16
17use crate::registry::{Identified, Identifier};
18
19const DEAD: usize = usize::MAX;
20
21/// Stable logical IDs over compact, dense physical payload storage.
22#[derive(Clone, PartialEq, Eq)]
23pub struct StableArena<Id: Identifier, T> {
24    values: Vec<T>,
25    slot_ids: Vec<Id>,
26    locations: Vec<usize>,
27}
28
29impl<Id: Identifier, T> StableArena<Id, T> {
30    /// Inserts `value` under a fresh, monotonic identifier.
31    pub fn push(&mut self, value: T) -> Id {
32        let raw = self.locations.len();
33        assert!(raw < DEAD, "StableArena identifier space exhausted");
34        let id = Id::from(raw);
35        assert_eq!(
36            Into::<usize>::into(id),
37            raw,
38            "StableArena identifier does not round-trip through its backing type"
39        );
40
41        let slot = self.values.len();
42        assert!(slot < DEAD, "StableArena physical slot space exhausted");
43        self.values.push(value);
44        self.slot_ids.push(id);
45        self.locations.push(slot);
46        id
47    }
48
49    /// Returns the number of live payloads.
50    pub fn len(&self) -> usize {
51        self.values.len()
52    }
53
54    /// Returns the number of logical IDs issued, including removed IDs.
55    pub fn issued_len(&self) -> usize {
56        self.locations.len()
57    }
58
59    /// Returns the payload capacity reserved by the dense value vector.
60    pub fn capacity(&self) -> usize {
61        self.values.capacity()
62    }
63
64    /// Returns bytes reserved by the arena's three structural vectors.
65    ///
66    /// Allocations owned by `T` itself are deliberately excluded.
67    pub fn structural_bytes(&self) -> usize {
68        self.values
69            .capacity()
70            .saturating_mul(std::mem::size_of::<T>())
71            .saturating_add(
72                self.slot_ids
73                    .capacity()
74                    .saturating_mul(std::mem::size_of::<Id>()),
75            )
76            .saturating_add(
77                self.locations
78                    .capacity()
79                    .saturating_mul(std::mem::size_of::<usize>()),
80            )
81    }
82
83    /// Returns `true` when no live payload remains.
84    pub fn is_empty(&self) -> bool {
85        self.values.is_empty()
86    }
87
88    /// Releases structural capacity retained beyond the current live/issued
89    /// lengths.
90    ///
91    /// Identifiers, liveness, and dense physical order are unchanged; only
92    /// allocator capacity is returned. Intended for explicit end-of-mutation
93    /// boundaries — shrinking after every removal would reallocate on the next
94    /// growth.
95    pub fn shrink_to_fit(&mut self) {
96        self.values.shrink_to_fit();
97        self.slot_ids.shrink_to_fit();
98        self.locations.shrink_to_fit();
99    }
100
101    /// Returns whether `id` currently names a live payload.
102    pub fn contains(&self, id: Id) -> bool {
103        self.live_slot(id).is_some()
104    }
105
106    /// Returns the live payload identified by `id`, if any.
107    pub fn get(&self, id: Id) -> Option<Identified<Id, &T>> {
108        let slot = self.live_slot(id)?;
109        Some(Identified::new(id, &self.values[slot]))
110    }
111
112    /// Returns the live payload identified by `id` mutably, if any.
113    pub fn get_mut(&mut self, id: Id) -> Option<Identified<Id, &mut T>> {
114        let slot = self.live_slot(id)?;
115        Some(Identified::new(id, &mut self.values[slot]))
116    }
117
118    /// Removes and returns `id`'s payload while preserving every other ID.
119    ///
120    /// # Panics
121    ///
122    /// Panics when `id` is unknown or was already removed.
123    pub fn remove(&mut self, id: Id) -> T {
124        let raw = id.into();
125        let slot = self.live_slot(id).unwrap_or_else(|| {
126            panic!(
127                "StableArena::remove: dead or unknown id {id:?} (issued {}, live {})",
128                self.issued_len(),
129                self.len()
130            )
131        });
132
133        self.locations[raw] = DEAD;
134        let removed = self.values.swap_remove(slot);
135        let removed_id = self.slot_ids.swap_remove(slot);
136        debug_assert_eq!(removed_id, id);
137
138        if slot < self.values.len() {
139            let moved_id = self.slot_ids[slot];
140            self.locations[Into::<usize>::into(moved_id)] = slot;
141        }
142        removed
143    }
144
145    /// Iterates over live entries in dense physical order.
146    pub fn iter(&self) -> Iter<'_, Id, T> {
147        Iter {
148            inner: self.slot_ids.iter().zip(self.values.iter()),
149        }
150    }
151
152    /// Iterates mutably over live entries in dense physical order.
153    pub fn iter_mut(&mut self) -> IterMut<'_, Id, T> {
154        IterMut {
155            inner: self.slot_ids.iter().zip(self.values.iter_mut()),
156        }
157    }
158
159    #[inline]
160    fn live_slot(&self, id: Id) -> Option<usize> {
161        let slot = *self.locations.get(Into::<usize>::into(id))?;
162        (slot != DEAD).then_some(slot)
163    }
164}
165
166impl<Id: Identifier, T> Default for StableArena<Id, T> {
167    fn default() -> Self {
168        Self {
169            values: Vec::new(),
170            slot_ids: Vec::new(),
171            locations: Vec::new(),
172        }
173    }
174}
175
176impl<Id: Identifier, T: Debug> Debug for StableArena<Id, T> {
177    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
178        f.debug_struct("StableArena")
179            .field("issued", &self.issued_len())
180            .field(
181                "entries",
182                &self
183                    .iter()
184                    .map(|entry| (entry.id, entry.inner))
185                    .collect::<Vec<_>>(),
186            )
187            .finish()
188    }
189}
190
191impl<Id: Identifier, T> Index<Id> for StableArena<Id, T> {
192    type Output = T;
193
194    fn index(&self, id: Id) -> &Self::Output {
195        self.get(id).map(|entry| entry.inner).unwrap_or_else(|| {
196            panic!(
197                "StableArena index: dead or unknown id {id:?} (issued {}, live {})",
198                self.issued_len(),
199                self.len()
200            )
201        })
202    }
203}
204
205impl<Id: Identifier, T> IndexMut<Id> for StableArena<Id, T> {
206    fn index_mut(&mut self, id: Id) -> &mut Self::Output {
207        let issued = self.issued_len();
208        let live = self.len();
209        self.get_mut(id)
210            .map(|entry| entry.inner)
211            .unwrap_or_else(|| {
212                panic!(
213                    "StableArena mutable index: dead or unknown id {id:?} (issued {issued}, live {live})"
214                )
215            })
216    }
217}
218
219/// Immutable dense-order iterator.
220pub struct Iter<'a, Id: Identifier, T> {
221    inner: Zip<slice::Iter<'a, Id>, slice::Iter<'a, T>>,
222}
223
224impl<'a, Id: Identifier, T> Iterator for Iter<'a, Id, T> {
225    type Item = Identified<Id, &'a T>;
226
227    fn next(&mut self) -> Option<Self::Item> {
228        let (id, value) = self.inner.next()?;
229        Some(Identified::new(*id, value))
230    }
231
232    fn size_hint(&self) -> (usize, Option<usize>) {
233        self.inner.size_hint()
234    }
235}
236
237impl<Id: Identifier, T> ExactSizeIterator for Iter<'_, Id, T> {}
238
239/// Mutable dense-order iterator.
240pub struct IterMut<'a, Id: Identifier, T> {
241    inner: Zip<slice::Iter<'a, Id>, slice::IterMut<'a, T>>,
242}
243
244impl<'a, Id: Identifier, T> Iterator for IterMut<'a, Id, T> {
245    type Item = Identified<Id, &'a mut T>;
246
247    fn next(&mut self) -> Option<Self::Item> {
248        let (id, value) = self.inner.next()?;
249        Some(Identified::new(*id, value))
250    }
251
252    fn size_hint(&self) -> (usize, Option<usize>) {
253        self.inner.size_hint()
254    }
255}
256
257impl<Id: Identifier, T> ExactSizeIterator for IterMut<'_, Id, T> {}
258
259/// Consuming dense-order iterator.
260pub struct IntoIter<Id: Identifier, T> {
261    inner: Zip<vec::IntoIter<Id>, vec::IntoIter<T>>,
262}
263
264impl<Id: Identifier, T> Iterator for IntoIter<Id, T> {
265    type Item = Identified<Id, T>;
266
267    fn next(&mut self) -> Option<Self::Item> {
268        let (id, value) = self.inner.next()?;
269        Some(Identified::new(id, value))
270    }
271
272    fn size_hint(&self) -> (usize, Option<usize>) {
273        self.inner.size_hint()
274    }
275}
276
277impl<Id: Identifier, T> ExactSizeIterator for IntoIter<Id, T> {}
278
279impl<'a, Id: Identifier, T> IntoIterator for &'a StableArena<Id, T> {
280    type Item = Identified<Id, &'a T>;
281    type IntoIter = Iter<'a, Id, T>;
282
283    fn into_iter(self) -> Self::IntoIter {
284        self.iter()
285    }
286}
287
288impl<'a, Id: Identifier, T> IntoIterator for &'a mut StableArena<Id, T> {
289    type Item = Identified<Id, &'a mut T>;
290    type IntoIter = IterMut<'a, Id, T>;
291
292    fn into_iter(self) -> Self::IntoIter {
293        self.iter_mut()
294    }
295}
296
297impl<Id: Identifier, T> IntoIterator for StableArena<Id, T> {
298    type Item = Identified<Id, T>;
299    type IntoIter = IntoIter<Id, T>;
300
301    fn into_iter(self) -> Self::IntoIter {
302        IntoIter {
303            inner: self.slot_ids.into_iter().zip(self.values),
304        }
305    }
306}
307
308impl<Id: Identifier, T> FromIterator<T> for StableArena<Id, T> {
309    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
310        let mut arena = Self::default();
311        for value in iter {
312            arena.push(value);
313        }
314        arena
315    }
316}
317
318#[derive(Serialize)]
319struct StableArenaWireRef<'a, T> {
320    next_id: usize,
321    entries: Vec<(usize, &'a T)>,
322}
323
324#[derive(Serialize, Deserialize)]
325struct StableArenaWire<T> {
326    next_id: usize,
327    entries: Vec<(usize, T)>,
328}
329
330impl<Id: Identifier, T: Serialize> Serialize for StableArena<Id, T> {
331    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
332    where
333        S: Serializer,
334    {
335        StableArenaWireRef {
336            next_id: self.issued_len(),
337            entries: self
338                .slot_ids
339                .iter()
340                .copied()
341                .map(Into::<usize>::into)
342                .zip(self.values.iter())
343                .collect(),
344        }
345        .serialize(serializer)
346    }
347}
348
349impl<'de, Id: Identifier, T: Deserialize<'de>> Deserialize<'de> for StableArena<Id, T> {
350    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
351    where
352        D: Deserializer<'de>,
353    {
354        let StableArenaWire { next_id, entries } = StableArenaWire::deserialize(deserializer)?;
355        if next_id == DEAD {
356            return Err(serde::de::Error::custom(
357                "StableArena next ID exhausts the location sentinel",
358            ));
359        }
360        if let Some(last_issued) = next_id.checked_sub(1)
361            && Into::<usize>::into(Id::from(last_issued)) != last_issued
362        {
363            return Err(serde::de::Error::custom(format!(
364                "StableArena next ID {next_id} exceeds its backing type"
365            )));
366        }
367
368        let mut arena = Self {
369            values: Vec::with_capacity(entries.len()),
370            slot_ids: Vec::with_capacity(entries.len()),
371            locations: vec![DEAD; next_id],
372        };
373        for (raw, value) in entries {
374            if raw >= next_id {
375                return Err(serde::de::Error::custom(format!(
376                    "StableArena live ID {raw} is outside next ID {next_id}"
377                )));
378            }
379            let id = Id::from(raw);
380            if Into::<usize>::into(id) != raw {
381                return Err(serde::de::Error::custom(format!(
382                    "StableArena ID {raw} does not fit its backing type"
383                )));
384            }
385            if arena.locations[raw] != DEAD {
386                return Err(serde::de::Error::custom(format!(
387                    "StableArena contains duplicate live ID {raw}"
388                )));
389            }
390            let slot = arena.values.len();
391            arena.locations[raw] = slot;
392            arena.slot_ids.push(id);
393            arena.values.push(value);
394        }
395        Ok(arena)
396    }
397}
398
399#[cfg(test)]
400mod tests {
401    use super::*;
402    use jstd_derive::Identifier;
403
404    #[derive(Identifier)]
405    struct Id(u32);
406
407    #[test]
408    fn removal_repairs_moved_slot_and_never_reuses_ids() {
409        let mut arena = StableArena::<Id, &str>::default();
410        let a = arena.push("a");
411        let b = arena.push("b");
412        let c = arena.push("c");
413
414        assert_eq!(arena.remove(b), "b");
415        assert!(!arena.contains(b));
416        assert_eq!(arena.get(b).map(|entry| **entry), None);
417        assert_eq!(arena[c], "c");
418        assert_eq!(arena[a], "a");
419        assert_eq!(
420            arena.iter().map(|entry| entry.id).collect::<Vec<_>>(),
421            [a, c]
422        );
423
424        let d = arena.push("d");
425        assert_eq!(usize::from(d), 3);
426        assert_eq!(arena.issued_len(), 4);
427        assert_eq!(arena.len(), 3);
428    }
429
430    #[test]
431    fn removes_first_last_and_only_entries() {
432        let mut arena = StableArena::<Id, i32>::default();
433        let a = arena.push(1);
434        let b = arena.push(2);
435        let c = arena.push(3);
436        assert_eq!(arena.remove(a), 1);
437        assert_eq!(arena.remove(c), 3);
438        assert_eq!(arena.remove(b), 2);
439        assert!(arena.is_empty());
440        assert_eq!(arena.issued_len(), 3);
441    }
442
443    #[test]
444    #[should_panic(expected = "dead or unknown id")]
445    fn indexing_removed_id_panics() {
446        let mut arena = StableArena::<Id, i32>::default();
447        let id = arena.push(1);
448        arena.remove(id);
449        let _ = arena[id];
450    }
451
452    #[test]
453    #[should_panic(expected = "dead or unknown id")]
454    fn double_remove_panics() {
455        let mut arena = StableArena::<Id, i32>::default();
456        let id = arena.push(1);
457        arena.remove(id);
458        arena.remove(id);
459    }
460
461    #[test]
462    fn mutable_and_consuming_iteration_keep_stable_ids() {
463        let mut arena = StableArena::<Id, i32>::default();
464        let a = arena.push(1);
465        let b = arena.push(2);
466        let c = arena.push(3);
467        arena.remove(b);
468        for mut entry in &mut arena {
469            **entry += 10;
470        }
471        let entries = arena
472            .into_iter()
473            .map(|entry| (entry.id, entry.inner))
474            .collect::<Vec<_>>();
475        assert_eq!(entries, [(a, 11), (c, 13)]);
476    }
477
478    #[test]
479    fn shrink_to_fit_preserves_ids_order_and_liveness() {
480        let mut arena = StableArena::<Id, u64>::default();
481        let ids = (0..1_000).map(|step| arena.push(step)).collect::<Vec<_>>();
482        for id in &ids[..900] {
483            arena.remove(*id);
484        }
485        let before = arena.clone();
486        arena.shrink_to_fit();
487        assert_eq!(arena, before);
488        assert!(arena.capacity() < 1_000);
489        assert_eq!(arena.len(), 100);
490        assert_eq!(arena.issued_len(), 1_000);
491        assert_eq!(arena[ids[950]], 950);
492        let next = arena.push(1_000);
493        assert_eq!(usize::from(next), 1_000);
494    }
495
496    #[test]
497    fn clone_and_equality_preserve_holes_and_physical_order() {
498        let mut arena = StableArena::<Id, i32>::default();
499        arena.push(1);
500        let removed = arena.push(2);
501        arena.push(3);
502        arena.remove(removed);
503        let clone = arena.clone();
504        assert_eq!(clone, arena);
505        assert_eq!(
506            clone.iter().map(|entry| entry.id).collect::<Vec<_>>(),
507            arena.iter().map(|entry| entry.id).collect::<Vec<_>>()
508        );
509    }
510
511    #[test]
512    fn serde_round_trip_preserves_holes_order_and_cursor() {
513        let mut arena = StableArena::<Id, String>::default();
514        let a = arena.push("a".into());
515        let b = arena.push("b".into());
516        let c = arena.push("c".into());
517        arena.remove(a);
518        arena.remove(b);
519        let bytes = bincode::serde::encode_to_vec(&arena, bincode::config::standard()).unwrap();
520        let (mut decoded, used): (StableArena<Id, String>, _) =
521            bincode::serde::decode_from_slice(&bytes, bincode::config::standard()).unwrap();
522        assert_eq!(used, bytes.len());
523        assert_eq!(decoded, arena);
524        assert_eq!(decoded[c], "c");
525        let d = decoded.push("d".into());
526        assert_eq!(usize::from(d), 3);
527
528        decoded.remove(c);
529        decoded.remove(d);
530        let bytes = bincode::serde::encode_to_vec(&decoded, bincode::config::standard()).unwrap();
531        let (empty, _): (StableArena<Id, String>, _) =
532            bincode::serde::decode_from_slice(&bytes, bincode::config::standard()).unwrap();
533        assert!(empty.is_empty());
534        assert_eq!(empty.issued_len(), 4);
535    }
536
537    #[test]
538    fn malformed_wire_is_rejected() {
539        let duplicate = StableArenaWire {
540            next_id: 2,
541            entries: vec![(0, 1_i32), (0, 2)],
542        };
543        let bytes = bincode::serde::encode_to_vec(duplicate, bincode::config::standard()).unwrap();
544        assert!(
545            bincode::serde::decode_from_slice::<StableArena<Id, i32>, _>(
546                &bytes,
547                bincode::config::standard()
548            )
549            .is_err()
550        );
551
552        let outside = StableArenaWire {
553            next_id: 1,
554            entries: vec![(1, 1_i32)],
555        };
556        let bytes = bincode::serde::encode_to_vec(outside, bincode::config::standard()).unwrap();
557        assert!(
558            bincode::serde::decode_from_slice::<StableArena<Id, i32>, _>(
559                &bytes,
560                bincode::config::standard()
561            )
562            .is_err()
563        );
564
565        let too_wide = StableArenaWire {
566            next_id: u32::MAX as usize + 2,
567            entries: Vec::<(usize, i32)>::new(),
568        };
569        let bytes = bincode::serde::encode_to_vec(too_wide, bincode::config::standard()).unwrap();
570        assert!(
571            bincode::serde::decode_from_slice::<StableArena<Id, i32>, _>(
572                &bytes,
573                bincode::config::standard()
574            )
575            .is_err()
576        );
577    }
578
579    #[test]
580    fn deterministic_mixed_operations_match_option_vec_model() {
581        let mut arena = StableArena::<Id, u64>::default();
582        let mut model: Vec<Option<u64>> = Vec::new();
583        let mut state = 0x1234_5678_9abc_def0_u64;
584
585        for step in 0..2_000_u64 {
586            state = state
587                .wrapping_mul(6_364_136_223_846_793_005)
588                .wrapping_add(1);
589            let live = model
590                .iter()
591                .enumerate()
592                .filter_map(|(id, value)| value.as_ref().map(|_| id))
593                .collect::<Vec<_>>();
594            if live.is_empty() || state & 3 != 0 {
595                let id = arena.push(step);
596                assert_eq!(usize::from(id), model.len());
597                model.push(Some(step));
598            } else {
599                let raw = live[state as usize % live.len()];
600                let expected = model[raw].take().unwrap();
601                assert_eq!(arena.remove(Id::from(raw)), expected);
602            }
603
604            for (raw, expected) in model.iter().enumerate() {
605                assert_eq!(
606                    arena.get(Id::from(raw)).map(|entry| **entry),
607                    *expected,
608                    "model mismatch at step {step}, id {raw}"
609                );
610            }
611        }
612    }
613
614    #[test]
615    fn public_views_iteration_and_capacity_accounting_work() {
616        let mut arena = StableArena::<Id, i32>::default();
617        assert!(arena.is_empty());
618        assert_eq!(arena.structural_bytes(), 0);
619
620        let first = arena.push(1);
621        let second = arena.push(2);
622        let capacity_before = arena.capacity();
623        assert!(arena.structural_bytes() >= capacity_before * std::mem::size_of::<i32>());
624        assert_eq!(arena.get(first).unwrap().id, first);
625        **arena.get_mut(second).unwrap() = 20;
626        arena[first] = 10;
627        assert_eq!(arena[first], 10);
628        assert_eq!(arena[second], 20);
629
630        let mut entries = arena.iter();
631        assert_eq!(entries.size_hint(), (2, Some(2)));
632        assert_eq!(entries.next().unwrap().id, first);
633        assert_eq!(entries.next().unwrap().id, second);
634        assert!(entries.next().is_none());
635        assert!(format!("{arena:?}").contains("issued"));
636
637        let rebuilt: StableArena<Id, i32> = [3, 4].into_iter().collect();
638        assert_eq!(
639            rebuilt.iter().map(|entry| *entry.inner).collect::<Vec<_>>(),
640            [3, 4]
641        );
642    }
643}
jstd/stable_arena.rs

jstd/
stable_arena.rs
