Skip to main content

nxpu_ir/
arena.rs

1//! Arena-based storage with typed handles.
2
3use std::cmp::Ordering;
4use std::collections::HashMap;
5use std::fmt;
6use std::hash::{Hash, Hasher};
7use std::marker::PhantomData;
8use std::ops::{Index, IndexMut};
9
10/// A typed handle into an [`Arena`] or [`UniqueArena`].
11///
12/// Handles are lightweight identifiers (u32 index) that provide
13/// type-safe access to arena-allocated values.
14pub struct Handle<T> {
15    index: u32,
16    _phantom: PhantomData<T>,
17}
18
19impl<T> Clone for Handle<T> {
20    fn clone(&self) -> Self {
21        *self
22    }
23}
24
25impl<T> Copy for Handle<T> {}
26
27impl<T> PartialEq for Handle<T> {
28    fn eq(&self, other: &Self) -> bool {
29        self.index == other.index
30    }
31}
32
33impl<T> Eq for Handle<T> {}
34
35impl<T> PartialOrd for Handle<T> {
36    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
37        Some(self.cmp(other))
38    }
39}
40
41impl<T> Ord for Handle<T> {
42    fn cmp(&self, other: &Self) -> Ordering {
43        self.index.cmp(&other.index)
44    }
45}
46
47impl<T> Hash for Handle<T> {
48    fn hash<H: Hasher>(&self, state: &mut H) {
49        self.index.hash(state);
50    }
51}
52
53impl<T> fmt::Debug for Handle<T> {
54    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
55        write!(f, "[{}]", self.index)
56    }
57}
58
59impl<T> Handle<T> {
60    /// Creates a new handle from a zero-based index.
61    pub(crate) fn new(index: u32) -> Self {
62        Self {
63            index,
64            _phantom: PhantomData,
65        }
66    }
67
68    /// Returns the zero-based index of this handle.
69    pub fn index(self) -> usize {
70        self.index as usize
71    }
72}
73
74/// A half-open range of [`Handle`]s: `[first, last)`.
75pub struct Range<T> {
76    first: u32,
77    last: u32,
78    _phantom: PhantomData<T>,
79}
80
81impl<T> Clone for Range<T> {
82    fn clone(&self) -> Self {
83        *self
84    }
85}
86
87impl<T> Copy for Range<T> {}
88
89impl<T> PartialEq for Range<T> {
90    fn eq(&self, other: &Self) -> bool {
91        self.first == other.first && self.last == other.last
92    }
93}
94
95impl<T> Eq for Range<T> {}
96
97impl<T> Hash for Range<T> {
98    fn hash<H: Hasher>(&self, state: &mut H) {
99        self.first.hash(state);
100        self.last.hash(state);
101    }
102}
103
104impl<T> fmt::Debug for Range<T> {
105    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
106        write!(f, "[{}..{})", self.first, self.last)
107    }
108}
109
110impl<T> Range<T> {
111    /// Creates a new range from two handles.
112    pub fn new(first: Handle<T>, last: Handle<T>) -> Self {
113        Self {
114            first: first.index,
115            last: last.index,
116            _phantom: PhantomData,
117        }
118    }
119
120    /// Creates a range from raw u32 indices.
121    pub fn from_index_range(range: std::ops::Range<u32>) -> Self {
122        Self {
123            first: range.start,
124            last: range.end,
125            _phantom: PhantomData,
126        }
127    }
128
129    /// Returns the first handle in the range.
130    pub fn first(&self) -> Handle<T> {
131        Handle::new(self.first)
132    }
133
134    /// Returns the handle one past the last element.
135    pub fn end(&self) -> Handle<T> {
136        Handle::new(self.last)
137    }
138
139    /// Returns this range as a `std::ops::Range<u32>`.
140    pub fn index_range(&self) -> std::ops::Range<u32> {
141        self.first..self.last
142    }
143
144    /// Returns `true` if the range contains no elements.
145    pub fn is_empty(&self) -> bool {
146        self.first >= self.last
147    }
148}
149
150/// An append-only arena with typed [`Handle`]-based access.
151#[derive(Clone, Debug)]
152pub struct Arena<T> {
153    data: Vec<T>,
154}
155
156impl<T> Default for Arena<T> {
157    fn default() -> Self {
158        Self::new()
159    }
160}
161
162impl<T> Arena<T> {
163    /// Creates an empty arena.
164    pub fn new() -> Self {
165        Self { data: Vec::new() }
166    }
167
168    /// Returns the number of elements in the arena.
169    pub fn len(&self) -> usize {
170        self.data.len()
171    }
172
173    /// Returns `true` if the arena contains no elements.
174    pub fn is_empty(&self) -> bool {
175        self.data.is_empty()
176    }
177
178    /// Returns the handle that will be assigned to the next appended value.
179    pub fn next_handle(&self) -> Handle<T> {
180        let index = u32::try_from(self.data.len()).unwrap_or_else(|_| {
181            panic!("arena overflow: {} items exceeds u32::MAX", self.data.len())
182        });
183        Handle::new(index)
184    }
185
186    /// Appends a value and returns its handle.
187    pub fn append(&mut self, value: T) -> Handle<T> {
188        let index = u32::try_from(self.data.len()).unwrap_or_else(|_| {
189            panic!("arena overflow: {} items exceeds u32::MAX", self.data.len())
190        });
191        self.data.push(value);
192        Handle::new(index)
193    }
194
195    /// Returns a reference to the value if the handle is valid.
196    pub fn try_get(&self, handle: Handle<T>) -> Option<&T> {
197        self.data.get(handle.index())
198    }
199
200    /// Iterates over `(handle, &value)` pairs.
201    pub fn iter(&self) -> impl Iterator<Item = (Handle<T>, &T)> {
202        // Safety: arena size bounded by u32::MAX (enforced in append)
203        self.data
204            .iter()
205            .enumerate()
206            .map(|(i, v)| (Handle::new(i as u32), v))
207    }
208
209    /// Iterates over `(handle, &mut value)` pairs.
210    pub fn iter_mut(&mut self) -> impl Iterator<Item = (Handle<T>, &mut T)> {
211        // Safety: arena size bounded by u32::MAX (enforced in append)
212        self.data
213            .iter_mut()
214            .enumerate()
215            .map(|(i, v)| (Handle::new(i as u32), v))
216    }
217}
218
219impl<T> Index<Handle<T>> for Arena<T> {
220    type Output = T;
221
222    fn index(&self, handle: Handle<T>) -> &T {
223        &self.data[handle.index()]
224    }
225}
226
227impl<T> IndexMut<Handle<T>> for Arena<T> {
228    fn index_mut(&mut self, handle: Handle<T>) -> &mut T {
229        &mut self.data[handle.index()]
230    }
231}
232
233/// A deduplicating arena that returns the same [`Handle`] for equal values.
234#[derive(Clone, Debug)]
235pub struct UniqueArena<T> {
236    data: Vec<T>,
237    map: HashMap<T, u32>,
238}
239
240impl<T: Hash + Eq> Default for UniqueArena<T> {
241    fn default() -> Self {
242        Self::new()
243    }
244}
245
246impl<T: Hash + Eq> UniqueArena<T> {
247    /// Creates an empty deduplicating arena.
248    pub fn new() -> Self {
249        Self {
250            data: Vec::new(),
251            map: HashMap::new(),
252        }
253    }
254
255    /// Returns the number of unique elements in the arena.
256    pub fn len(&self) -> usize {
257        self.data.len()
258    }
259
260    /// Returns `true` if the arena contains no elements.
261    pub fn is_empty(&self) -> bool {
262        self.data.is_empty()
263    }
264
265    /// Inserts a value, returning an existing handle if the value is already present.
266    pub fn insert(&mut self, value: T) -> Handle<T>
267    where
268        T: Clone,
269    {
270        if let Some(&index) = self.map.get(&value) {
271            return Handle::new(index);
272        }
273        let index = u32::try_from(self.data.len()).unwrap_or_else(|_| {
274            panic!("arena overflow: {} items exceeds u32::MAX", self.data.len())
275        });
276        self.map.insert(value.clone(), index);
277        self.data.push(value);
278        Handle::new(index)
279    }
280
281    /// Returns a reference to the value if the handle is valid.
282    pub fn try_get(&self, handle: Handle<T>) -> Option<&T> {
283        self.data.get(handle.index())
284    }
285
286    /// Iterates over `(handle, &value)` pairs.
287    pub fn iter(&self) -> impl Iterator<Item = (Handle<T>, &T)> {
288        // Safety: arena size bounded by u32::MAX (enforced in insert)
289        self.data
290            .iter()
291            .enumerate()
292            .map(|(i, v)| (Handle::new(i as u32), v))
293    }
294}
295
296impl<T> Index<Handle<T>> for UniqueArena<T> {
297    type Output = T;
298
299    fn index(&self, handle: Handle<T>) -> &T {
300        &self.data[handle.index()]
301    }
302}
303
304#[cfg(test)]
305mod tests {
306    use super::*;
307
308    #[test]
309    fn arena_append_and_access() {
310        let mut arena = Arena::new();
311        let h0 = arena.append("hello");
312        let h1 = arena.append("world");
313        assert_eq!(arena[h0], "hello");
314        assert_eq!(arena[h1], "world");
315        assert_eq!(arena.len(), 2);
316    }
317
318    #[test]
319    fn arena_iter() {
320        let mut arena = Arena::new();
321        arena.append(10);
322        arena.append(20);
323        arena.append(30);
324        let items: Vec<_> = arena.iter().map(|(h, &v)| (h.index(), v)).collect();
325        assert_eq!(items, vec![(0, 10), (1, 20), (2, 30)]);
326    }
327
328    #[test]
329    fn arena_next_handle() {
330        let mut arena = Arena::<i32>::new();
331        let h0 = arena.next_handle();
332        assert_eq!(h0.index(), 0);
333        arena.append(42);
334        let h1 = arena.next_handle();
335        assert_eq!(h1.index(), 1);
336    }
337
338    #[test]
339    fn unique_arena_dedup() {
340        let mut arena = UniqueArena::new();
341        let h0 = arena.insert(42);
342        let h1 = arena.insert(99);
343        let h2 = arena.insert(42); // duplicate
344        assert_eq!(h0, h2);
345        assert_ne!(h0, h1);
346        assert_eq!(arena.len(), 2);
347    }
348
349    #[test]
350    fn handle_ordering() {
351        let h0: Handle<u32> = Handle::new(0);
352        let h1: Handle<u32> = Handle::new(1);
353        assert!(h0 < h1);
354        assert_eq!(h0, h0);
355    }
356
357    #[test]
358    fn range_operations() {
359        let range = Range::<u32>::from_index_range(2..5);
360        assert!(!range.is_empty());
361        assert_eq!(range.first().index(), 2);
362        assert_eq!(range.end().index(), 5);
363        assert_eq!(range.index_range(), 2..5);
364
365        let empty = Range::<u32>::from_index_range(3..3);
366        assert!(empty.is_empty());
367    }
368
369    #[test]
370    fn arena_try_get() {
371        let mut arena = Arena::new();
372        let h0 = arena.append(42);
373        assert_eq!(arena.try_get(h0), Some(&42));
374        assert_eq!(arena.try_get(Handle::new(99)), None);
375    }
376}