gruel_air/
param_arena.rs

1//! Arena-based storage for function and method parameter data.
2//!
3//! # Overview
4//!
5//! The `ParamArena` provides centralized, contiguous storage for all parameter data
6//! in the compiler. Instead of each `FunctionInfo` and `MethodInfo` owning their own
7//! `Vec`s of parameter data, they store a `ParamRange` index into this arena.
8//!
9//! # Memory Benefits
10//!
11//! Before (per function with N params):
12//! - `param_names: Vec<Spur>` - 24 bytes header + 4*N bytes
13//! - `param_types: Vec<Type>` - 24 bytes header + 4*N bytes
14//! - `param_modes: Vec<RirParamMode>` - 24 bytes header + N bytes
15//! - `param_comptime: Vec<bool>` - 24 bytes header + N bytes
16//! - Total: ~96 bytes overhead + ~10*N bytes data
17//!
18//! After (per function):
19//! - `ParamRange` - 8 bytes (start + len)
20//! - ~88 bytes saved per function
21//!
22//! # Cache Locality
23//!
24//! All parameter data for all functions is stored contiguously, improving cache
25//! behavior when iterating over function signatures during type checking.
26//!
27//! # Usage
28//!
29//! ```ignore
30//! // During declaration gathering, allocate params for a function:
31//! let range = arena.alloc(
32//!     names.into_iter(),
33//!     types.into_iter(),
34//!     modes.into_iter(),
35//!     comptime.into_iter(),
36//! );
37//!
38//! // Store the range in FunctionInfo instead of the Vec data
39//! let func_info = FunctionInfo {
40//!     params: range,  // Just 8 bytes instead of 4 Vecs
41//!     return_type,
42//!     // ...
43//! };
44//!
45//! // Later, access the data through the arena:
46//! let types = arena.types(range);  // &[Type]
47//! let modes = arena.modes(range);  // &[RirParamMode]
48//! ```
49//!
50//! # Design Decisions
51//!
52//! ## Separate Vecs vs Interleaved Storage
53//!
54//! We use separate Vecs for each field (names, types, modes, comptime) rather than
55//! an interleaved `Vec<ParamData>` for better cache locality when accessing only
56//! one field. Type checking often iterates just over types, while mode checking
57//! iterates just over modes.
58//!
59//! ## Index Type
60//!
61//! Uses u32 indices (4GB of params is sufficient for any reasonable program).
62//! This matches the existing pattern in gruel-air (e.g., `Type` uses u32 indices).
63//!
64//! ## Append-Only Design
65//!
66//! The arena is append-only during declaration gathering. Once all functions and
67//! methods are registered, the data is never modified. This simplifies the
68//! implementation and enables safe shared references.
69//!
70//! ## Method-Specific Allocation
71//!
72//! Methods always store parameter names (needed for named argument checking).
73//! Functions only need names for generic functions (for type substitution).
74//! To handle this, both `alloc` and `alloc_method` store names, but functions
75//! can pass empty iterators if names aren't needed.
76
77use crate::types::Type;
78use gruel_rir::RirParamMode;
79use lasso::Spur;
80
81/// A range of parameters in the `ParamArena`.
82///
83/// This is a lightweight handle (8 bytes) that replaces four `Vec`s (96+ bytes)
84/// in `FunctionInfo` and `MethodInfo`.
85#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
86pub struct ParamRange {
87    /// Starting index into the arena's storage.
88    start: u32,
89    /// Number of parameters in this range.
90    len: u32,
91}
92
93impl ParamRange {
94    /// Creates an empty parameter range.
95    pub const EMPTY: ParamRange = ParamRange { start: 0, len: 0 };
96
97    /// Returns the number of parameters in this range.
98    #[inline]
99    pub fn len(&self) -> usize {
100        self.len as usize
101    }
102
103    /// Returns true if this range contains no parameters.
104    #[inline]
105    pub fn is_empty(&self) -> bool {
106        self.len == 0
107    }
108
109    /// Creates a new range with the given start and length.
110    fn new(start: u32, len: u32) -> Self {
111        Self { start, len }
112    }
113
114    /// Returns the start index.
115    fn start(&self) -> usize {
116        self.start as usize
117    }
118
119    /// Returns the end index (exclusive).
120    fn end(&self) -> usize {
121        self.start as usize + self.len as usize
122    }
123}
124
125/// Central storage for all function/method parameter data.
126///
127/// Stores parameter names, types, modes, and comptime flags in separate
128/// contiguous arrays for optimal cache locality during type checking.
129#[derive(Debug, Default)]
130pub struct ParamArena {
131    /// All parameter names, indexed by position in the arena.
132    /// For functions without named args, this may contain placeholder values.
133    names: Vec<Spur>,
134
135    /// All parameter types, indexed by position in the arena.
136    types: Vec<Type>,
137
138    /// All parameter modes (Normal, Inout, Borrow, Comptime).
139    modes: Vec<RirParamMode>,
140
141    /// Whether each parameter is a comptime parameter.
142    comptime: Vec<bool>,
143}
144
145impl ParamArena {
146    /// Creates a new, empty parameter arena.
147    pub fn new() -> Self {
148        Self::default()
149    }
150
151    /// Creates a new arena with pre-allocated capacity.
152    ///
153    /// Use when the approximate number of total parameters is known.
154    pub fn with_capacity(capacity: usize) -> Self {
155        Self {
156            names: Vec::with_capacity(capacity),
157            types: Vec::with_capacity(capacity),
158            modes: Vec::with_capacity(capacity),
159            comptime: Vec::with_capacity(capacity),
160        }
161    }
162
163    /// Returns the total number of parameters stored in the arena.
164    pub fn total_params(&self) -> usize {
165        self.types.len()
166    }
167
168    /// Allocates storage for a function's parameters.
169    ///
170    /// All iterators must yield the same number of elements.
171    ///
172    /// # Panics
173    ///
174    /// Panics if the iterators yield different numbers of elements.
175    pub fn alloc(
176        &mut self,
177        names: impl IntoIterator<Item = Spur>,
178        types: impl IntoIterator<Item = Type>,
179        modes: impl IntoIterator<Item = RirParamMode>,
180        comptime: impl IntoIterator<Item = bool>,
181    ) -> ParamRange {
182        let start = self.types.len() as u32;
183
184        // Extend all vectors from the iterators
185        self.names.extend(names);
186        let names_len = self.names.len() - start as usize;
187
188        self.types.extend(types);
189        let types_len = self.types.len() - start as usize;
190
191        self.modes.extend(modes);
192        let modes_len = self.modes.len() - start as usize;
193
194        self.comptime.extend(comptime);
195        let comptime_len = self.comptime.len() - start as usize;
196
197        // Verify all lengths match
198        assert_eq!(
199            names_len, types_len,
200            "ParamArena::alloc: names ({}) and types ({}) have different lengths",
201            names_len, types_len
202        );
203        assert_eq!(
204            types_len, modes_len,
205            "ParamArena::alloc: types ({}) and modes ({}) have different lengths",
206            types_len, modes_len
207        );
208        assert_eq!(
209            modes_len, comptime_len,
210            "ParamArena::alloc: modes ({}) and comptime ({}) have different lengths",
211            modes_len, comptime_len
212        );
213
214        ParamRange::new(start, types_len as u32)
215    }
216
217    /// Allocates storage for a method's parameters (without comptime flags).
218    ///
219    /// Methods don't have comptime parameters, so this is a convenience method
220    /// that fills the comptime array with `false` values.
221    pub fn alloc_method(
222        &mut self,
223        names: impl IntoIterator<Item = Spur>,
224        types: impl IntoIterator<Item = Type>,
225    ) -> ParamRange {
226        let start = self.types.len() as u32;
227
228        // Collect names and types
229        self.names.extend(names);
230        let names_len = self.names.len() - start as usize;
231
232        self.types.extend(types);
233        let types_len = self.types.len() - start as usize;
234
235        // Verify lengths match
236        assert_eq!(
237            names_len, types_len,
238            "ParamArena::alloc_method: names ({}) and types ({}) have different lengths",
239            names_len, types_len
240        );
241
242        // Methods use Normal mode and are not comptime by default
243        self.modes
244            .extend(std::iter::repeat_n(RirParamMode::Normal, types_len));
245        self.comptime.extend(std::iter::repeat_n(false, types_len));
246
247        ParamRange::new(start, types_len as u32)
248    }
249
250    /// Returns the parameter names for a range.
251    #[inline]
252    pub fn names(&self, range: ParamRange) -> &[Spur] {
253        &self.names[range.start()..range.end()]
254    }
255
256    /// Returns the parameter types for a range.
257    #[inline]
258    pub fn types(&self, range: ParamRange) -> &[Type] {
259        &self.types[range.start()..range.end()]
260    }
261
262    /// Returns the parameter modes for a range.
263    #[inline]
264    pub fn modes(&self, range: ParamRange) -> &[RirParamMode] {
265        &self.modes[range.start()..range.end()]
266    }
267
268    /// Returns the comptime flags for a range.
269    #[inline]
270    pub fn comptime(&self, range: ParamRange) -> &[bool] {
271        &self.comptime[range.start()..range.end()]
272    }
273
274    /// Returns an iterator over all parameter data for a range.
275    ///
276    /// Useful for cases where you need to iterate all fields together.
277    pub fn iter(
278        &self,
279        range: ParamRange,
280    ) -> impl Iterator<Item = (&Spur, &Type, &RirParamMode, &bool)> {
281        self.names(range)
282            .iter()
283            .zip(self.types(range))
284            .zip(self.modes(range))
285            .zip(self.comptime(range))
286            .map(|(((name, ty), mode), comptime)| (name, ty, mode, comptime))
287    }
288
289    /// Returns an iterator over (name, type) pairs for a range.
290    ///
291    /// Useful for method parameter type checking.
292    pub fn name_type_pairs(&self, range: ParamRange) -> impl Iterator<Item = (&Spur, &Type)> {
293        self.names(range).iter().zip(self.types(range))
294    }
295
296    /// Returns an iterator over (type, mode) pairs for a range.
297    ///
298    /// Useful for function call argument checking.
299    pub fn type_mode_pairs(
300        &self,
301        range: ParamRange,
302    ) -> impl Iterator<Item = (&Type, &RirParamMode)> {
303        self.types(range).iter().zip(self.modes(range))
304    }
305}
306
307#[cfg(test)]
308mod tests {
309    use super::*;
310    use lasso::Rodeo;
311
312    fn make_spur(rodeo: &mut Rodeo, s: &str) -> Spur {
313        rodeo.get_or_intern(s)
314    }
315
316    #[test]
317    fn test_empty_arena() {
318        let arena = ParamArena::new();
319        assert_eq!(arena.total_params(), 0);
320    }
321
322    #[test]
323    fn test_empty_range() {
324        let range = ParamRange::EMPTY;
325        assert!(range.is_empty());
326        assert_eq!(range.len(), 0);
327    }
328
329    #[test]
330    fn test_alloc_single_param() {
331        let mut arena = ParamArena::new();
332        let mut rodeo = Rodeo::default();
333
334        let name = make_spur(&mut rodeo, "x");
335        let range = arena.alloc([name], [Type::I32], [RirParamMode::Normal], [false]);
336
337        assert_eq!(range.len(), 1);
338        assert_eq!(arena.total_params(), 1);
339        assert_eq!(arena.names(range), &[name]);
340        assert_eq!(arena.types(range), &[Type::I32]);
341        assert_eq!(arena.modes(range), &[RirParamMode::Normal]);
342        assert_eq!(arena.comptime(range), &[false]);
343    }
344
345    #[test]
346    fn test_alloc_multiple_params() {
347        let mut arena = ParamArena::new();
348        let mut rodeo = Rodeo::default();
349
350        let x = make_spur(&mut rodeo, "x");
351        let y = make_spur(&mut rodeo, "y");
352        let z = make_spur(&mut rodeo, "z");
353
354        let range = arena.alloc(
355            [x, y, z],
356            [Type::I32, Type::BOOL, Type::I64],
357            [
358                RirParamMode::Normal,
359                RirParamMode::Inout,
360                RirParamMode::Borrow,
361            ],
362            [false, false, true],
363        );
364
365        assert_eq!(range.len(), 3);
366        assert_eq!(arena.types(range), &[Type::I32, Type::BOOL, Type::I64]);
367        assert_eq!(
368            arena.modes(range),
369            &[
370                RirParamMode::Normal,
371                RirParamMode::Inout,
372                RirParamMode::Borrow
373            ]
374        );
375        assert_eq!(arena.comptime(range), &[false, false, true]);
376    }
377
378    #[test]
379    fn test_multiple_functions() {
380        let mut arena = ParamArena::new();
381        let mut rodeo = Rodeo::default();
382
383        // First function with 2 params
384        let a = make_spur(&mut rodeo, "a");
385        let b = make_spur(&mut rodeo, "b");
386        let range1 = arena.alloc(
387            [a, b],
388            [Type::I32, Type::I32],
389            [RirParamMode::Normal, RirParamMode::Normal],
390            [false, false],
391        );
392
393        // Second function with 1 param
394        let c = make_spur(&mut rodeo, "c");
395        let range2 = arena.alloc([c], [Type::BOOL], [RirParamMode::Inout], [false]);
396
397        // Verify they don't overlap
398        assert_eq!(range1.len(), 2);
399        assert_eq!(range2.len(), 1);
400        assert_eq!(arena.total_params(), 3);
401
402        // Verify data is correct for each range
403        assert_eq!(arena.types(range1), &[Type::I32, Type::I32]);
404        assert_eq!(arena.types(range2), &[Type::BOOL]);
405    }
406
407    #[test]
408    fn test_alloc_method() {
409        let mut arena = ParamArena::new();
410        let mut rodeo = Rodeo::default();
411
412        let x = make_spur(&mut rodeo, "x");
413        let y = make_spur(&mut rodeo, "y");
414
415        let range = arena.alloc_method([x, y], [Type::I32, Type::BOOL]);
416
417        assert_eq!(range.len(), 2);
418        assert_eq!(arena.names(range), &[x, y]);
419        assert_eq!(arena.types(range), &[Type::I32, Type::BOOL]);
420        // Methods default to Normal mode and non-comptime
421        assert_eq!(
422            arena.modes(range),
423            &[RirParamMode::Normal, RirParamMode::Normal]
424        );
425        assert_eq!(arena.comptime(range), &[false, false]);
426    }
427
428    #[test]
429    fn test_iter() {
430        let mut arena = ParamArena::new();
431        let mut rodeo = Rodeo::default();
432
433        let x = make_spur(&mut rodeo, "x");
434        let y = make_spur(&mut rodeo, "y");
435
436        let range = arena.alloc(
437            [x, y],
438            [Type::I32, Type::BOOL],
439            [RirParamMode::Normal, RirParamMode::Inout],
440            [false, true],
441        );
442
443        let items: Vec<_> = arena.iter(range).collect();
444        assert_eq!(items.len(), 2);
445        assert_eq!(items[0], (&x, &Type::I32, &RirParamMode::Normal, &false));
446        assert_eq!(items[1], (&y, &Type::BOOL, &RirParamMode::Inout, &true));
447    }
448
449    #[test]
450    fn test_name_type_pairs() {
451        let mut arena = ParamArena::new();
452        let mut rodeo = Rodeo::default();
453
454        let x = make_spur(&mut rodeo, "x");
455        let y = make_spur(&mut rodeo, "y");
456
457        let range = arena.alloc(
458            [x, y],
459            [Type::I32, Type::BOOL],
460            [RirParamMode::Normal, RirParamMode::Normal],
461            [false, false],
462        );
463
464        let pairs: Vec<_> = arena.name_type_pairs(range).collect();
465        assert_eq!(pairs, vec![(&x, &Type::I32), (&y, &Type::BOOL)]);
466    }
467
468    #[test]
469    fn test_type_mode_pairs() {
470        let mut arena = ParamArena::new();
471        let mut rodeo = Rodeo::default();
472
473        let x = make_spur(&mut rodeo, "x");
474        let y = make_spur(&mut rodeo, "y");
475
476        let range = arena.alloc(
477            [x, y],
478            [Type::I32, Type::BOOL],
479            [RirParamMode::Normal, RirParamMode::Inout],
480            [false, false],
481        );
482
483        let pairs: Vec<_> = arena.type_mode_pairs(range).collect();
484        assert_eq!(
485            pairs,
486            vec![
487                (&Type::I32, &RirParamMode::Normal),
488                (&Type::BOOL, &RirParamMode::Inout)
489            ]
490        );
491    }
492
493    #[test]
494    #[should_panic(expected = "names (2) and types (1) have different lengths")]
495    fn test_alloc_mismatched_lengths_panics() {
496        let mut arena = ParamArena::new();
497        let mut rodeo = Rodeo::default();
498
499        let x = make_spur(&mut rodeo, "x");
500        let y = make_spur(&mut rodeo, "y");
501
502        // This should panic - 2 names but only 1 type
503        let _ = arena.alloc([x, y], [Type::I32], [RirParamMode::Normal], [false]);
504    }
505
506    #[test]
507    fn test_with_capacity() {
508        let arena = ParamArena::with_capacity(100);
509        assert_eq!(arena.total_params(), 0);
510        // Can't easily test capacity, but this ensures the method works
511    }
512}