1use std::{collections::TryReserveError, ops::Range};
12
13use crate::compiler_error::{CompilerFailure, CompilerStage};
14
15#[derive(Clone, Copy, Debug, PartialEq, Eq)]
16pub enum ArenaKind {
17 Expressions,
18 Statements,
19 TypedExpressions,
20 TypedStatements,
21}
22
23#[derive(Clone, Copy, Debug, PartialEq, Eq)]
24pub enum ArenaOperation {
25 Read,
26 Mutate,
27 Allocate,
28 Iterate,
29}
30
31#[derive(Debug)]
32pub enum ArenaError {
33 InvalidId {
34 arena: ArenaKind,
35 operation: ArenaOperation,
36 id: u32,
37 len: usize,
38 },
39 Capacity {
40 arena: ArenaKind,
41 operation: ArenaOperation,
42 len: usize,
43 max_nodes: u32,
44 },
45 Allocation {
46 arena: ArenaKind,
47 len: usize,
48 source: TryReserveError,
49 },
50}
51
52impl ArenaError {
53 pub fn into_compiler_failure(self, stage: CompilerStage) -> CompilerFailure {
54 let message = self.to_string();
55 match self {
56 Self::InvalidId { .. } => CompilerFailure::Internal {
57 stage,
58 span: None,
59 message,
60 },
61 Self::Capacity { .. } | Self::Allocation { .. } => CompilerFailure::Limit {
62 stage,
63 span: None,
64 message,
65 help: vec!["reduce the size of the program or generated compiler work".into()],
66 },
67 }
68 }
69}
70
71impl std::fmt::Display for ArenaError {
72 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
73 match self {
74 Self::InvalidId {
75 arena,
76 operation,
77 id,
78 len,
79 } => {
80 write!(
81 f,
82 "{arena:?} {operation:?}: invalid node ID {id} (length {len})"
83 )
84 }
85 Self::Capacity {
86 arena,
87 operation,
88 len,
89 max_nodes,
90 } => {
91 write!(
92 f,
93 "{arena:?} {operation:?}: node capacity exceeded (length {len}, maximum {max_nodes})"
94 )
95 }
96 Self::Allocation { arena, len, source } => {
97 write!(
98 f,
99 "{arena:?} Allocate: cannot reserve node storage (length {len}): {source}"
100 )
101 }
102 }
103 }
104}
105
106impl std::error::Error for ArenaError {
107 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
108 match self {
109 Self::Allocation { source, .. } => Some(source),
110 _ => None,
111 }
112 }
113}
114
115pub(crate) fn get<T>(nodes: &[T], id: u32, arena: ArenaKind) -> Result<&T, ArenaError> {
116 usize::try_from(id)
117 .ok()
118 .and_then(|index| nodes.get(index))
119 .ok_or(ArenaError::InvalidId {
120 arena,
121 operation: ArenaOperation::Read,
122 id,
123 len: nodes.len(),
124 })
125}
126
127pub(crate) fn get_mut<T>(nodes: &mut [T], id: u32, arena: ArenaKind) -> Result<&mut T, ArenaError> {
128 let len = nodes.len();
129 usize::try_from(id)
130 .ok()
131 .and_then(|index| nodes.get_mut(index))
132 .ok_or(ArenaError::InvalidId {
133 arena,
134 operation: ArenaOperation::Mutate,
135 id,
136 len,
137 })
138}
139
140pub(crate) fn push<T>(nodes: &mut Vec<T>, node: T, arena: ArenaKind) -> Result<u32, ArenaError> {
141 #[cfg(test)]
142 let limit = TEST_NODE_LIMIT.with(|limit| {
143 limit
144 .get()
145 .filter(|(kind, _)| *kind == arena)
146 .map_or(u32::MAX, |(_, limit)| limit)
147 });
148 #[cfg(not(test))]
149 let limit = u32::MAX;
150 push_with_limit(nodes, node, arena, limit)
151}
152
153pub(crate) fn ids(len: usize, arena: ArenaKind) -> Result<Range<u32>, ArenaError> {
154 let end = u32::try_from(len).map_err(|_| ArenaError::Capacity {
155 arena,
156 operation: ArenaOperation::Iterate,
157 len,
158 max_nodes: u32::MAX,
159 })?;
160 Ok(0..end)
161}
162
163fn push_with_limit<T>(
164 nodes: &mut Vec<T>,
165 node: T,
166 arena: ArenaKind,
167 max_nodes: u32,
168) -> Result<u32, ArenaError> {
169 let id = next_id(nodes.len(), arena, max_nodes)?;
170 reserve(nodes, 1, arena)?;
171 nodes.push(node);
172 Ok(id)
173}
174
175fn next_id(len: usize, arena: ArenaKind, max_nodes: u32) -> Result<u32, ArenaError> {
176 u32::try_from(len)
177 .ok()
178 .filter(|&id| id < max_nodes)
179 .ok_or(ArenaError::Capacity {
180 arena,
181 operation: ArenaOperation::Allocate,
182 len,
183 max_nodes,
184 })
185}
186
187fn reserve<T>(nodes: &mut Vec<T>, additional: usize, arena: ArenaKind) -> Result<(), ArenaError> {
188 nodes
189 .try_reserve(additional)
190 .map_err(|source| ArenaError::Allocation {
191 arena,
192 len: nodes.len(),
193 source,
194 })
195}
196
197#[cfg(test)]
200thread_local! {
201 static TEST_NODE_LIMIT: std::cell::Cell<Option<(ArenaKind, u32)>> = const { std::cell::Cell::new(None) };
202}
203
204#[cfg(test)]
205pub(crate) fn with_node_limit<T>(arena: ArenaKind, limit: u32, run: impl FnOnce() -> T) -> T {
206 struct Restore(Option<(ArenaKind, u32)>);
207 impl Drop for Restore {
208 fn drop(&mut self) {
209 TEST_NODE_LIMIT.with(|limit| limit.set(self.0));
210 }
211 }
212 let _restore = Restore(TEST_NODE_LIMIT.with(|cell| cell.replace(Some((arena, limit)))));
213 run()
214}
215
216#[cfg(test)]
217mod tests {
218 use super::*;
219
220 #[test]
221 fn checked_arena_capacity_failure_preserves_nodes_and_allows_reads() {
222 let arena = ArenaKind::Expressions;
223 let mut nodes = Vec::new();
224 assert_eq!(push_with_limit(&mut nodes, 11, arena, 2).unwrap(), 0);
225 assert_eq!(push_with_limit(&mut nodes, 22, arena, 2).unwrap(), 1);
226 let capacity = nodes.capacity();
227 assert!(matches!(
228 push_with_limit(&mut nodes, 33, arena, 2),
229 Err(ArenaError::Capacity {
230 len: 2,
231 max_nodes: 2,
232 ..
233 })
234 ));
235 assert_eq!(nodes, [11, 22]);
236 assert_eq!(nodes.capacity(), capacity);
237 assert_eq!(*get(&nodes, 1, arena).unwrap(), 22);
238 *get_mut(&mut nodes, 0, arena).unwrap() = 44;
239 assert_eq!(nodes, [44, 22]);
240 assert!(push_with_limit(&mut Vec::new(), 1, arena, 0).is_err());
241 }
242
243 #[test]
244 fn checked_arena_id_boundaries_do_not_need_large_allocations() {
245 let arena = ArenaKind::TypedStatements;
246 assert_eq!(
247 next_id(u32::MAX as usize - 1, arena, u32::MAX).unwrap(),
248 u32::MAX - 1
249 );
250 assert!(next_id(u32::MAX as usize, arena, u32::MAX).is_err());
251 assert_eq!(ids(0, arena).unwrap(), 0..0);
252 assert_eq!(
253 ids(u32::MAX as usize, arena).unwrap().next_back(),
254 Some(u32::MAX - 1)
255 );
256 if let Some(overflow) = (u32::MAX as usize).checked_add(1) {
257 assert!(next_id(overflow, arena, u32::MAX).is_err());
258 assert!(ids(overflow, arena).is_err());
259 }
260 }
261
262 #[test]
263 fn checked_arena_reservation_failure_is_a_limit_without_mutation() {
264 let mut nodes = vec![7_u8];
265 let error = reserve(&mut nodes, usize::MAX, ArenaKind::Statements).unwrap_err();
266 assert!(matches!(&error, ArenaError::Allocation { len: 1, .. }));
267 assert!(std::error::Error::source(&error).is_some());
268 assert!(matches!(
269 error.into_compiler_failure(CompilerStage::Parse),
270 CompilerFailure::Limit {
271 stage: CompilerStage::Parse,
272 span: None,
273 ..
274 }
275 ));
276 assert_eq!(nodes, [7]);
277 assert_eq!(push(&mut nodes, 8, ArenaKind::Statements).unwrap(), 1);
278 }
279
280 #[test]
281 fn checked_arena_failures_map_without_node_or_span_access() {
282 let error = get::<u8>(&[], u32::MAX, ArenaKind::TypedExpressions).unwrap_err();
283 let message = error.to_string();
284 assert!(message.contains("4294967295"));
285 assert!(message.contains("length 0"));
286 assert!(matches!(
287 error.into_compiler_failure(CompilerStage::Codegen),
288 CompilerFailure::Internal {
289 stage: CompilerStage::Codegen,
290 span: None,
291 ..
292 }
293 ));
294 let error = next_id(2, ArenaKind::Expressions, 2).unwrap_err();
295 assert!(matches!(
296 error.into_compiler_failure(CompilerStage::Infer),
297 CompilerFailure::Limit {
298 stage: CompilerStage::Infer,
299 span: None,
300 ..
301 }
302 ));
303 }
304}