mq_lang/
arena.rs

1#[cfg(feature = "ast-json")]
2use serde::{Deserialize, Serialize};
3use std::{marker::PhantomData, ops::Index};
4
5/// A type-safe identifier for elements stored in an [`Arena`].
6///
7/// Uses phantom data to ensure type safety - an `ArenaId<A>` cannot be used
8/// to access elements from an `Arena<B>`.
9#[cfg_attr(feature = "ast-json", derive(Serialize, Deserialize))]
10#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
11pub struct ArenaId<T> {
12    id: u32,
13    _phantom_data: PhantomData<T>,
14}
15
16impl<T> Copy for ArenaId<T> {}
17
18impl<T> Clone for ArenaId<T> {
19    #[inline(always)]
20    fn clone(&self) -> ArenaId<T> {
21        *self
22    }
23}
24
25impl<T> From<u32> for ArenaId<T> {
26    fn from(id: u32) -> Self {
27        Self::new(id)
28    }
29}
30
31impl<T> From<usize> for ArenaId<T> {
32    fn from(id: usize) -> Self {
33        Self::new(id as u32)
34    }
35}
36
37impl<T> From<i32> for ArenaId<T> {
38    fn from(id: i32) -> Self {
39        Self::new(id as u32)
40    }
41}
42
43impl<T> ArenaId<T> {
44    /// Creates a new arena identifier from a raw `u32` index.
45    pub const fn new(id: u32) -> ArenaId<T> {
46        Self {
47            id,
48            _phantom_data: PhantomData,
49        }
50    }
51}
52
53/// An arena allocator for efficiently storing and accessing elements.
54///
55/// The arena allocates elements sequentially and returns type-safe [`ArenaId`]s
56/// that can be used to retrieve elements later. This pattern provides fast allocation
57/// and cache-friendly access.
58#[derive(Debug, Clone, Default)]
59pub struct Arena<T> {
60    items: Vec<T>,
61}
62
63impl<T: Clone + PartialEq> Arena<T> {
64    /// Creates a new arena with the specified initial capacity.
65    pub fn new(size: usize) -> Self {
66        Arena {
67            items: Vec::with_capacity(size),
68        }
69    }
70
71    /// Allocates a value in the arena and returns its identifier.
72    pub fn alloc(&mut self, value: T) -> ArenaId<T> {
73        let arena_id = self.items.len() as u32;
74        self.items.push(value);
75        ArenaId::new(arena_id)
76    }
77
78    /// Returns the number of elements in the arena.
79    pub fn len(&self) -> usize {
80        self.items.len()
81    }
82
83    /// Returns `true` if the arena contains no elements.
84    pub fn is_empty(&self) -> bool {
85        self.len() == 0
86    }
87
88    /// Returns `true` if the arena contains the specified value.
89    pub fn contains(&self, value: T) -> bool {
90        self.items.contains(&value)
91    }
92
93    /// Extends the arena by cloning elements from a slice.
94    pub fn extend_from_slice(&mut self, items: &[T]) {
95        self.items.extend_from_slice(items);
96    }
97}
98
99impl<T> Index<ArenaId<T>> for Arena<T> {
100    type Output = T;
101
102    fn index(&self, index: ArenaId<T>) -> &Self::Output {
103        &self.items[index.id as usize]
104    }
105}
106
107impl<T> Arena<T> {
108    /// Returns a reference to the element at the given `ArenaId`, or `None` if out of bounds.
109    pub fn get(&self, id: ArenaId<T>) -> Option<&T> {
110        self.items.get(id.id as usize)
111    }
112
113    /// Returns a slice of all elements in the arena.
114    pub fn as_slice(&self) -> &[T] {
115        &self.items
116    }
117}
118
119#[cfg(test)]
120mod tests {
121    use super::*;
122    use rstest::rstest;
123
124    #[rstest]
125    #[case(vec![1, 2, 3], 1, true)]
126    #[case(vec![1, 2, 3], 4, false)]
127    #[case(Vec::new(), 1, false)]
128    fn test_contains(#[case] values: Vec<i32>, #[case] value: i32, #[case] expected: bool) {
129        let mut arena = Arena::new(values.len());
130        for v in values {
131            arena.alloc(v);
132        }
133        assert_eq!(arena.contains(value), expected);
134    }
135
136    #[rstest]
137    #[case(vec![1, 2, 3], 1, 2)]
138    #[case(vec![1, 2, 3], 0, 1)]
139    #[case(vec![1, 2, 3], 2, 3)]
140    fn test_get(#[case] values: Vec<i32>, #[case] index: u32, #[case] expected: i32) {
141        let mut arena = Arena::new(values.len());
142        for v in values {
143            arena.alloc(v);
144        }
145        let id = ArenaId::new(index);
146        assert_eq!(arena[id], expected);
147    }
148
149    #[rstest]
150    #[case(vec![1, 2, 3], 3)]
151    #[case(Vec::new(), 0)]
152    fn test_len(#[case] values: Vec<i32>, #[case] expected: usize) {
153        let mut arena = Arena::new(values.len());
154        for v in values {
155            arena.alloc(v);
156        }
157        assert_eq!(arena.len(), expected);
158    }
159
160    #[rstest]
161    #[case(vec![1, 2, 3], false)]
162    #[case(Vec::new(), true)]
163    fn test_is_empty(#[case] values: Vec<i32>, #[case] expected: bool) {
164        let mut arena = Arena::new(values.len());
165        for v in values {
166            arena.alloc(v);
167        }
168        assert_eq!(arena.is_empty(), expected);
169    }
170
171    #[test]
172    fn test_from() {
173        let id_u32: ArenaId<i32> = 5u32.into();
174        assert_eq!(id_u32.id, 5);
175
176        let id_usize: ArenaId<i32> = 10usize.into();
177        assert_eq!(id_usize.id, 10);
178
179        let id_i32: ArenaId<i32> = 15i32.into();
180        assert_eq!(id_i32.id, 15);
181    }
182}
mq_lang/arena.rs

mq_lang/
arena.rs
