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
94impl<T> Index<ArenaId<T>> for Arena<T> {
95    type Output = T;
96
97    fn index(&self, index: ArenaId<T>) -> &Self::Output {
98        &self.items[index.id as usize]
99    }
100}
101
102impl<T> Arena<T> {
103    /// Returns a reference to the element at the given `ArenaId`, or `None` if out of bounds.
104    pub fn get(&self, id: ArenaId<T>) -> Option<&T> {
105        self.items.get(id.id as usize)
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112    use rstest::rstest;
113
114    #[rstest]
115    #[case(vec![1, 2, 3], 1, true)]
116    #[case(vec![1, 2, 3], 4, false)]
117    #[case(Vec::new(), 1, false)]
118    fn test_contains(#[case] values: Vec<i32>, #[case] value: i32, #[case] expected: bool) {
119        let mut arena = Arena::new(values.len());
120        for v in values {
121            arena.alloc(v);
122        }
123        assert_eq!(arena.contains(value), expected);
124    }
125
126    #[rstest]
127    #[case(vec![1, 2, 3], 1, 2)]
128    #[case(vec![1, 2, 3], 0, 1)]
129    #[case(vec![1, 2, 3], 2, 3)]
130    fn test_get(#[case] values: Vec<i32>, #[case] index: u32, #[case] expected: i32) {
131        let mut arena = Arena::new(values.len());
132        for v in values {
133            arena.alloc(v);
134        }
135        let id = ArenaId::new(index);
136        assert_eq!(arena[id], expected);
137    }
138
139    #[rstest]
140    #[case(vec![1, 2, 3], 3)]
141    #[case(Vec::new(), 0)]
142    fn test_len(#[case] values: Vec<i32>, #[case] expected: usize) {
143        let mut arena = Arena::new(values.len());
144        for v in values {
145            arena.alloc(v);
146        }
147        assert_eq!(arena.len(), expected);
148    }
149
150    #[rstest]
151    #[case(vec![1, 2, 3], false)]
152    #[case(Vec::new(), true)]
153    fn test_is_empty(#[case] values: Vec<i32>, #[case] expected: bool) {
154        let mut arena = Arena::new(values.len());
155        for v in values {
156            arena.alloc(v);
157        }
158        assert_eq!(arena.is_empty(), expected);
159    }
160
161    #[test]
162    fn test_from() {
163        let id_u32: ArenaId<i32> = 5u32.into();
164        assert_eq!(id_u32.id, 5);
165
166        let id_usize: ArenaId<i32> = 10usize.into();
167        assert_eq!(id_usize.id, 10);
168
169        let id_i32: ArenaId<i32> = 15i32.into();
170        assert_eq!(id_i32.id, 15);
171    }
172}
mq_lang/arena.rs

mq_lang/
arena.rs
