1use crate::types::Type;
78use gruel_rir::RirParamMode;
79use lasso::Spur;
80
81#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
86pub struct ParamRange {
87 start: u32,
89 len: u32,
91}
92
93impl ParamRange {
94 pub const EMPTY: ParamRange = ParamRange { start: 0, len: 0 };
96
97 #[inline]
99 pub fn len(&self) -> usize {
100 self.len as usize
101 }
102
103 #[inline]
105 pub fn is_empty(&self) -> bool {
106 self.len == 0
107 }
108
109 fn new(start: u32, len: u32) -> Self {
111 Self { start, len }
112 }
113
114 fn start(&self) -> usize {
116 self.start as usize
117 }
118
119 fn end(&self) -> usize {
121 self.start as usize + self.len as usize
122 }
123}
124
125#[derive(Debug, Default)]
130pub struct ParamArena {
131 names: Vec<Spur>,
134
135 types: Vec<Type>,
137
138 modes: Vec<RirParamMode>,
140
141 comptime: Vec<bool>,
143}
144
145impl ParamArena {
146 pub fn new() -> Self {
148 Self::default()
149 }
150
151 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 pub fn total_params(&self) -> usize {
165 self.types.len()
166 }
167
168 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 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 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 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 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 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 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 #[inline]
252 pub fn names(&self, range: ParamRange) -> &[Spur] {
253 &self.names[range.start()..range.end()]
254 }
255
256 #[inline]
258 pub fn types(&self, range: ParamRange) -> &[Type] {
259 &self.types[range.start()..range.end()]
260 }
261
262 #[inline]
264 pub fn modes(&self, range: ParamRange) -> &[RirParamMode] {
265 &self.modes[range.start()..range.end()]
266 }
267
268 #[inline]
270 pub fn comptime(&self, range: ParamRange) -> &[bool] {
271 &self.comptime[range.start()..range.end()]
272 }
273
274 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 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 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 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 let c = make_spur(&mut rodeo, "c");
395 let range2 = arena.alloc([c], [Type::BOOL], [RirParamMode::Inout], [false]);
396
397 assert_eq!(range1.len(), 2);
399 assert_eq!(range2.len(), 1);
400 assert_eq!(arena.total_params(), 3);
401
402 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 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 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 }
512}