Skip to main content

rl_lang/ast/
arena.rs

1//! Generic index-based arena for AST nodes.
2//!
3//! Replaces `Box<T>` with a `Copy` integer handle (`Id<T>`) into a flat
4//! `Vec<T>` owned by an `Arena<T>`. Every `Id<T>` is tagged with the id of
5//! the arena it was allocated from, so using it against the wrong arena
6//! panics immediately at the access site instead of silently reading
7//! unrelated data or panicking later with a confusing out-of-bounds error.
8
9use std::fmt;
10use std::marker::PhantomData;
11use std::sync::atomic::{AtomicU32, Ordering};
12
13/// Global counter handing out a unique id to every `Arena<T>` that's ever
14/// constructed, so `Id<T>`s can be checked against their origin arena.
15static NEXT_ARENA_ID: AtomicU32 = AtomicU32::new(0);
16
17/// A typed, `Copy` handle into a specific [`Arena<T>`].
18///
19/// Two `Id<T>`s are only meaningfully comparable if they came from the same
20/// arena - `arena_id` is what lets `Arena::get`/`get_mut` detect and reject
21/// a handle that was allocated somewhere else.
22pub struct Id<T> {
23    index: u32,
24    arena_id: u32,
25    _marker: PhantomData<T>,
26}
27
28impl<T> Id<T> {
29    #[inline]
30    fn new(index: u32, arena_id: u32) -> Self {
31        Self {
32            index,
33            arena_id,
34            _marker: PhantomData,
35        }
36    }
37
38    #[inline]
39    pub fn index(self) -> usize {
40        self.index as usize
41    }
42
43    #[inline]
44    pub fn arena_id(self) -> u32 {
45        self.arena_id
46    }
47
48    /// Rebases this id onto a different arena: shifts its index by `offset`
49    /// and retags it with `new_arena_id`. Used only by arena-merge logic -
50    /// this is the one place it's legal to construct an `Id` that wasn't
51    /// handed out by `Arena::alloc`.
52    #[inline]
53    pub fn rebase(self, offset: u32, new_arena_id: u32) -> Self {
54        Self {
55            index: self.index + offset,
56            arena_id: new_arena_id,
57            _marker: PhantomData,
58        }
59    }
60}
61
62impl<T> Clone for Id<T> {
63    fn clone(&self) -> Self {
64        *self
65    }
66}
67impl<T> Copy for Id<T> {}
68impl<T> PartialEq for Id<T> {
69    fn eq(&self, other: &Self) -> bool {
70        self.index == other.index && self.arena_id == other.arena_id
71    }
72}
73impl<T> Eq for Id<T> {}
74impl<T> std::hash::Hash for Id<T> {
75    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
76        self.index.hash(state);
77        self.arena_id.hash(state);
78    }
79}
80impl<T> fmt::Debug for Id<T> {
81    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
82        write!(f, "Id({}@arena{})", self.index, self.arena_id)
83    }
84}
85
86/// Flat backing store for `T`, indexed by [`Id<T>`].
87#[derive(Debug)]
88pub struct Arena<T> {
89    id: u32,
90    items: Vec<T>,
91}
92
93impl<T> Arena<T> {
94    pub fn new() -> Self {
95        Self {
96            id: NEXT_ARENA_ID.fetch_add(1, Ordering::Relaxed),
97            items: Vec::new(),
98        }
99    }
100
101    pub fn with_capacity(cap: usize) -> Self {
102        Self {
103            id: NEXT_ARENA_ID.fetch_add(1, Ordering::Relaxed),
104            items: Vec::with_capacity(cap),
105        }
106    }
107
108    /// This arena's unique id, e.g. for logging which arena an `Id` came from.
109    pub fn id(&self) -> u32 {
110        self.id
111    }
112
113    /// Insert a node, returning its id.
114    pub fn alloc(&mut self, item: T) -> Id<T> {
115        let index = self.items.len() as u32;
116        self.items.push(item);
117        Id::new(index, self.id)
118    }
119
120    pub fn get(&self, id: Id<T>) -> &T {
121        self.check_owner(id, "get");
122        let index = id.index();
123        if index >= self.items.len() {
124            panic!(
125                "Arena<{}>: index {} out of bounds (len {})",
126                std::any::type_name::<T>(),
127                index,
128                self.items.len()
129            );
130        }
131        &self.items[index]
132    }
133
134    pub fn get_mut(&mut self, id: Id<T>) -> &mut T {
135        self.check_owner(id, "get_mut");
136        let index = id.index();
137        let len = self.items.len();
138        if index >= len {
139            panic!(
140                "Arena<{}>: index {} out of bounds (len {})",
141                std::any::type_name::<T>(),
142                index,
143                len
144            );
145        }
146        &mut self.items[index]
147    }
148
149    pub fn len(&self) -> usize {
150        self.items.len()
151    }
152
153    pub fn is_empty(&self) -> bool {
154        self.items.is_empty()
155    }
156
157    /// Panics with a clear message if `id` didn't come from this arena.
158    /// This is the check that turns cross-arena bugs into an obvious,
159    /// immediate panic instead of silent corruption or a confusing
160    /// out-of-bounds error somewhere else.
161    #[inline]
162    fn check_owner(&self, id: Id<T>, op: &str) {
163        if id.arena_id != self.id {
164            panic!(
165                "Arena<{}>::{}: Id belongs to arena {} but this arena is {} - cross-arena access",
166                std::any::type_name::<T>(),
167                op,
168                id.arena_id,
169                self.id
170            );
171        }
172    }
173
174    /// Low-level access to this arena's id and backing `Vec`, used only by
175    /// `Ast::merge_statements` to implement arena merging. Bypasses the
176    /// normal `Id`-checked API - not for general use.
177    pub fn raw_parts_mut(&mut self) -> (u32, &mut Vec<T>) {
178        (self.id, &mut self.items)
179    }
180}
181
182impl<T> Default for Arena<T> {
183    fn default() -> Self {
184        Self::new()
185    }
186}
187
188impl<T> std::ops::Index<Id<T>> for Arena<T> {
189    type Output = T;
190    fn index(&self, id: Id<T>) -> &T {
191        self.get(id)
192    }
193}
194
195impl<T> std::ops::IndexMut<Id<T>> for Arena<T> {
196    fn index_mut(&mut self, id: Id<T>) -> &mut T {
197        self.get_mut(id)
198    }
199}