Skip to main content

submilli_engine/
arena.rs

1//! Checked storage operations shared by the parsed and typed AST arenas.
2//!
3//! Errors carry storage metadata only. A compiler phase adds its stage without
4//! reading a potentially invalid node or using that node's source span.
5//!
6//! IDs remain arena-local indices, not ownership tokens. These operations check
7//! storage bounds and capacity; they do not validate a node's child IDs or spans.
8//! The maximum node count is `u32::MAX`, matching the legacy arena allocation
9//! invariant and leaving `u32::MAX` itself outside the allocated ID range.
10
11use 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// A scoped, thread-local cap exercises real phase allocation failures without
198// allocating billions of nodes or leaking fault injection into parallel tests.
199#[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}