Skip to main content

aether_core/
arena.rs

1//! Generational arena allocator.
2//!
3//! Provides O(1) alloc/dealloc with use-after-free detection via generation counters.
4//! All memory is pre-allocated at startup — zero heap activity in the RT thread.
5
6use crate::MAX_NODES;
7
8/// A typed, generational index into the arena.
9/// The generation field prevents ABA problems and dangling references.
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
11pub struct NodeId {
12    pub index: u32,
13    pub generation: u32,
14}
15
16impl NodeId {
17    pub const INVALID: Self = Self {
18        index: u32::MAX,
19        generation: u32::MAX,
20    };
21}
22
23struct Entry<T> {
24    /// The stored value, valid only when `occupied` is true.
25    value: Option<T>,
26    /// Current generation. Incremented on each removal.
27    generation: u32,
28    /// Next free slot index, valid only when `value` is None.
29    next_free: Option<u32>,
30}
31
32/// Fixed-capacity generational arena.
33pub struct Arena<T> {
34    entries: Vec<Entry<T>>,
35    free_head: Option<u32>,
36    len: usize,
37}
38
39impl<T> Arena<T> {
40    /// Allocate a new arena with `capacity` pre-reserved slots.
41    pub fn with_capacity(capacity: usize) -> Self {
42        let mut entries = Vec::with_capacity(capacity);
43        for i in 0..capacity {
44            let next = if i + 1 < capacity {
45                Some((i + 1) as u32)
46            } else {
47                None
48            };
49            entries.push(Entry {
50                value: None,
51                generation: 0,
52                next_free: next,
53            });
54        }
55        Self {
56            entries,
57            free_head: if capacity > 0 { Some(0) } else { None },
58            len: 0,
59        }
60    }
61
62    /// Insert a value, returning its generational id.
63    /// Returns `None` if the arena is full.
64    pub fn insert(&mut self, value: T) -> Option<NodeId> {
65        let index = self.free_head?;
66        let entry = &mut self.entries[index as usize];
67        self.free_head = entry.next_free;
68        let generation = entry.generation;
69        entry.value = Some(value);
70        entry.next_free = None;
71        self.len += 1;
72        Some(NodeId { index, generation })
73    }
74
75    /// Remove a value by id. Returns the value if the id was valid.
76    pub fn remove(&mut self, id: NodeId) -> Option<T> {
77        let entry = self.entries.get_mut(id.index as usize)?;
78        if entry.generation != id.generation || entry.value.is_none() {
79            return None;
80        }
81        let value = entry.value.take();
82        // Bump generation to invalidate all existing ids pointing here.
83        entry.generation = entry.generation.wrapping_add(1);
84        entry.next_free = self.free_head;
85        self.free_head = Some(id.index);
86        self.len -= 1;
87        value
88    }
89
90    /// Get a shared reference. Returns `None` for stale ids.
91    #[inline]
92    pub fn get(&self, id: NodeId) -> Option<&T> {
93        let entry = self.entries.get(id.index as usize)?;
94        if entry.generation == id.generation {
95            entry.value.as_ref()
96        } else {
97            None
98        }
99    }
100
101    /// Get a mutable reference. Returns `None` for stale ids.
102    #[inline]
103    pub fn get_mut(&mut self, id: NodeId) -> Option<&mut T> {
104        let entry = self.entries.get_mut(id.index as usize)?;
105        if entry.generation == id.generation {
106            entry.value.as_mut()
107        } else {
108            None
109        }
110    }
111
112    pub fn len(&self) -> usize {
113        self.len
114    }
115    pub fn is_empty(&self) -> bool {
116        self.len == 0
117    }
118}
119
120/// Default arena sized for the project's node limit.
121pub fn default_node_arena<T>() -> Arena<T> {
122    Arena::with_capacity(MAX_NODES)
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128
129    #[test]
130    fn insert_get_remove() {
131        let mut arena: Arena<i32> = Arena::with_capacity(4);
132        let id = arena.insert(42).unwrap();
133        assert_eq!(*arena.get(id).unwrap(), 42);
134        let val = arena.remove(id).unwrap();
135        assert_eq!(val, 42);
136        assert!(arena.get(id).is_none());
137    }
138
139    #[test]
140    fn generation_prevents_aba() {
141        let mut arena: Arena<i32> = Arena::with_capacity(4);
142        let id1 = arena.insert(1).unwrap();
143        arena.remove(id1).unwrap();
144        let id2 = arena.insert(2).unwrap();
145        // Same slot index, bumped generation.
146        assert_eq!(id1.index, id2.index);
147        assert_ne!(id1.generation, id2.generation);
148        assert!(arena.get(id1).is_none());
149        assert_eq!(*arena.get(id2).unwrap(), 2);
150    }
151
152    #[test]
153    fn capacity_exhaustion() {
154        let mut arena: Arena<i32> = Arena::with_capacity(2);
155        let a = arena.insert(1).unwrap();
156        let _b = arena.insert(2).unwrap();
157        assert!(arena.insert(3).is_none()); // full
158        arena.remove(a).unwrap();
159        assert!(arena.insert(3).is_some()); // slot recycled
160    }
161}