1use std::collections::HashMap;
13use std::sync::Arc as Rc;
14
15use crate::ast::*;
16
17struct TypeInfo {
22 variants: HashMap<(String, String), Vec<String>>,
31 variant_parents: HashMap<String, Vec<String>>,
37 #[allow(dead_code)]
42 records: HashMap<String, Vec<(String, String)>>,
43}
44
45fn build_type_info(items: &[TopLevel]) -> TypeInfo {
46 let mut variants: HashMap<(String, String), Vec<String>> = HashMap::new();
47 let mut variant_parents: HashMap<String, Vec<String>> = HashMap::new();
48 let mut records: HashMap<String, Vec<(String, String)>> = HashMap::new();
49 for item in items {
50 match item {
51 TopLevel::TypeDef(TypeDef::Sum {
52 name: parent,
53 variants: vs,
54 ..
55 }) => {
56 for v in vs {
57 variants.insert((parent.clone(), v.name.clone()), v.fields.clone());
58 variant_parents
59 .entry(v.name.clone())
60 .or_default()
61 .push(parent.clone());
62 }
63 }
64 TopLevel::TypeDef(TypeDef::Product { name, fields, .. }) => {
65 records.insert(name.clone(), fields.clone());
66 }
67 _ => {}
68 }
69 }
70 TypeInfo {
71 variants,
72 variant_parents,
73 records,
74 }
75}
76
77pub fn resolve_program(items: &mut [TopLevel]) {
82 let type_info = build_type_info(items);
83 for item in items.iter_mut() {
84 if let TopLevel::FnDef(fd) = item {
85 resolve_fn(fd, &type_info);
86 }
87 }
88}
89
90fn resolve_fn(fd: &mut FnDef, type_info: &TypeInfo) {
101 let mut state = ResolverState::new(type_info);
102
103 state.scopes.push(HashMap::new());
109 for (param_name, ty_str) in &fd.params {
110 let ty = crate::types::parse_type_str_strict(ty_str).unwrap_or(Type::Invalid);
111 state.declare_param(param_name, ty);
112 }
113
114 let mut body = fd.body.as_ref().clone();
118 state.walk_stmts(body.stmts_mut());
119 fd.body = Rc::new(body);
120
121 let next_slot = state.next_slot;
122 let last_alloc = state.last_alloc;
123 let slot_types = state.slot_types;
124 state.scopes.pop();
125
126 fd.resolution = Some(FnResolution {
127 local_count: next_slot,
128 local_slots: Rc::new(last_alloc),
129 local_slot_types: Rc::new(slot_types),
130 aliased_slots: Rc::new(vec![false; next_slot as usize]),
131 });
132}
133
134struct ResolverState<'a> {
138 type_info: &'a TypeInfo,
139 next_slot: u16,
140 slot_types: Vec<Type>,
141 scopes: Vec<HashMap<String, u16>>,
146 last_alloc: HashMap<String, u16>,
153}
154
155impl<'a> ResolverState<'a> {
156 fn new(type_info: &'a TypeInfo) -> Self {
157 Self {
158 type_info,
159 next_slot: 0,
160 slot_types: Vec::new(),
161 scopes: Vec::new(),
162 last_alloc: HashMap::new(),
163 }
164 }
165
166 fn alloc(&mut self, ty: Type) -> u16 {
168 let idx = self.next_slot;
169 self.next_slot += 1;
170 self.slot_types.push(ty);
171 idx
172 }
173
174 fn declare(&mut self, name: &str, ty: Type) -> u16 {
188 if name == "_" {
189 return u16::MAX;
190 }
191 let slot = self.alloc(ty);
192 if let Some(scope) = self.scopes.last_mut() {
193 scope.insert(name.to_string(), slot);
194 }
195 self.last_alloc.insert(name.to_string(), slot);
196 slot
197 }
198
199 fn declare_param(&mut self, name: &str, ty: Type) -> u16 {
206 let slot = self.alloc(ty);
207 if name != "_" {
208 if let Some(scope) = self.scopes.last_mut() {
209 scope.insert(name.to_string(), slot);
210 }
211 self.last_alloc.insert(name.to_string(), slot);
212 }
213 slot
214 }
215
216 fn lookup(&self, name: &str) -> Option<u16> {
220 for scope in self.scopes.iter().rev() {
221 if let Some(&s) = scope.get(name) {
222 return Some(s);
223 }
224 }
225 None
226 }
227
228 fn walk_stmts(&mut self, stmts: &mut [Stmt]) {
229 for stmt in stmts {
230 match stmt {
231 Stmt::Binding(name, _annot, expr) => {
232 let ty = expr.ty().cloned().unwrap_or(Type::Invalid);
239 self.declare(name, ty);
240 self.walk_expr(expr);
241 }
242 Stmt::Expr(expr) => self.walk_expr(expr),
243 }
244 }
245 }
246
247 fn walk_expr(&mut self, expr: &mut Spanned<Expr>) {
248 match &mut expr.node {
249 Expr::Ident(name) => {
250 if let Some(slot) = self.lookup(name) {
251 expr.node = Expr::Resolved {
252 slot,
253 name: name.clone(),
254 last_use: AnnotBool(false),
255 };
256 }
257 }
258 Expr::Match { subject, arms } => {
259 self.walk_expr(subject);
260 let subject_ty = subject.ty().cloned();
261 for arm in arms.iter_mut() {
262 self.scopes.push(HashMap::new());
263 let slots = self.allocate_pattern(&arm.pattern, subject_ty.as_ref());
264 let _ = arm.binding_slots.set(slots);
265 self.walk_expr(&mut arm.body);
266 self.scopes.pop();
267 }
268 }
269 Expr::FnCall(func, args) => {
270 self.walk_expr(func);
271 for arg in args {
272 self.walk_expr(arg);
273 }
274 }
275 Expr::BinOp(_, l, r) => {
276 self.walk_expr(l);
277 self.walk_expr(r);
278 }
279 Expr::Neg(inner) => self.walk_expr(inner),
280 Expr::Attr(obj, _) => self.walk_expr(obj),
281 Expr::ErrorProp(inner) => self.walk_expr(inner),
282 Expr::Constructor(_, Some(inner)) => self.walk_expr(inner),
283 Expr::Constructor(_, None) => {}
284 Expr::List(items) | Expr::Tuple(items) | Expr::IndependentProduct(items, _) => {
285 for it in items {
286 self.walk_expr(it);
287 }
288 }
289 Expr::MapLiteral(entries) => {
290 for (k, v) in entries {
291 self.walk_expr(k);
292 self.walk_expr(v);
293 }
294 }
295 Expr::InterpolatedStr(parts) => {
296 for part in parts {
297 if let StrPart::Parsed(e) = part {
298 self.walk_expr(e);
299 }
300 }
301 }
302 Expr::RecordCreate { fields, .. } => {
303 for (_, e) in fields {
304 self.walk_expr(e);
305 }
306 }
307 Expr::RecordUpdate { base, updates, .. } => {
308 self.walk_expr(base);
309 for (_, e) in updates {
310 self.walk_expr(e);
311 }
312 }
313 Expr::TailCall(boxed) => {
314 for a in &mut boxed.args {
315 self.walk_expr(a);
316 }
317 }
318 Expr::Literal(_) | Expr::Resolved { .. } => {}
319 }
320 }
321
322 fn allocate_pattern(&mut self, pattern: &Pattern, subject_ty: Option<&Type>) -> Vec<u16> {
328 match pattern {
329 Pattern::Ident(name) => {
330 let ty = subject_ty.cloned().unwrap_or(Type::Invalid);
331 vec![self.declare(name, ty)]
332 }
333 Pattern::Cons(head, tail) => {
334 let elem_ty = match subject_ty {
335 Some(Type::List(inner)) => (**inner).clone(),
336 _ => Type::Invalid,
337 };
338 let list_ty = Type::List(Box::new(elem_ty.clone()));
339 vec![self.declare(head, elem_ty), self.declare(tail, list_ty)]
340 }
341 Pattern::Constructor(name, bindings) => {
342 let bare = name.rsplit('.').next().unwrap_or(name);
343 let parent_hint: Option<String> = match (subject_ty, name.split_once('.')) {
344 (Some(Type::Named { name: parent, .. }), _) => Some(parent.clone()),
345 (_, Some((parent, _))) => Some(parent.to_string()),
346 _ => self
347 .type_info
348 .variant_parents
349 .get(bare)
350 .and_then(|parents| {
351 if parents.len() == 1 {
352 Some(parents[0].clone())
353 } else {
354 None
355 }
356 }),
357 };
358 let field_tys: Vec<Type> = match (bare, subject_ty) {
359 ("Ok", Some(Type::Result(t, _))) => vec![(**t).clone()],
360 ("Err", Some(Type::Result(_, e))) => vec![(**e).clone()],
361 ("Some", Some(Type::Option(inner))) => vec![(**inner).clone()],
362 ("None", _) => Vec::new(),
363 _ => parent_hint
364 .and_then(|p| self.type_info.variants.get(&(p, bare.to_string())))
365 .map(|fields| {
366 fields
367 .iter()
368 .map(|s| {
369 crate::types::parse_type_str_strict(s).unwrap_or(Type::Invalid)
370 })
371 .collect()
372 })
373 .unwrap_or_else(|| vec![Type::Invalid; bindings.len()]),
374 };
375 bindings
376 .iter()
377 .enumerate()
378 .map(|(i, name)| {
379 let ty = field_tys.get(i).cloned().unwrap_or(Type::Invalid);
380 self.declare(name, ty)
381 })
382 .collect()
383 }
384 Pattern::Tuple(items) => {
385 let elem_tys: Vec<Type> = match subject_ty {
386 Some(Type::Tuple(elems)) if elems.len() == items.len() => elems.to_vec(),
387 _ => vec![Type::Invalid; items.len()],
388 };
389 let mut slots = Vec::new();
390 for (item, elem_ty) in items.iter().zip(elem_tys.iter()) {
391 slots.extend(self.allocate_pattern(item, Some(elem_ty)));
392 }
393 slots
394 }
395 Pattern::Wildcard | Pattern::Literal(_) | Pattern::EmptyList => Vec::new(),
396 }
397 }
398}
399
400#[cfg(test)]
401mod tests {
402 use super::*;
403
404 #[test]
405 fn resolves_param_to_slot() {
406 let mut fd = FnDef {
407 name: "add".to_string(),
408 line: 1,
409 params: vec![
410 ("a".to_string(), "Int".to_string()),
411 ("b".to_string(), "Int".to_string()),
412 ],
413 return_type: "Int".to_string(),
414 effects: vec![],
415 desc: None,
416 body: Rc::new(FnBody::from_expr(Spanned::bare(Expr::BinOp(
417 BinOp::Add,
418 Box::new(Spanned::bare(Expr::Ident("a".to_string()))),
419 Box::new(Spanned::bare(Expr::Ident("b".to_string()))),
420 )))),
421 resolution: None,
422 };
423 resolve_fn(
424 &mut fd,
425 &TypeInfo {
426 variants: HashMap::new(),
427 variant_parents: HashMap::new(),
428 records: HashMap::new(),
429 },
430 );
431 let res = fd.resolution.as_ref().unwrap();
432 assert_eq!(res.local_slots["a"], 0);
433 assert_eq!(res.local_slots["b"], 1);
434 assert_eq!(res.local_count, 2);
435
436 match fd.body.tail_expr() {
437 Some(Spanned {
438 node: Expr::BinOp(_, left, right),
439 ..
440 }) => {
441 assert_eq!(
442 left.node,
443 Expr::Resolved {
444 slot: 0,
445 name: "a".to_string(),
446 last_use: AnnotBool(false)
447 }
448 );
449 assert_eq!(
450 right.node,
451 Expr::Resolved {
452 slot: 1,
453 name: "b".to_string(),
454 last_use: AnnotBool(false)
455 }
456 );
457 }
458 other => panic!("unexpected body: {:?}", other),
459 }
460 }
461
462 #[test]
463 fn wildcard_param_still_claims_slot() {
464 let mut fd = FnDef {
471 name: "ignore".to_string(),
472 line: 1,
473 params: vec![("_".to_string(), "Int".to_string())],
474 return_type: "Int".to_string(),
475 effects: vec![],
476 desc: None,
477 body: Rc::new(FnBody::from_expr(Spanned::bare(Expr::Literal(
478 Literal::Int(42),
479 )))),
480 resolution: None,
481 };
482 resolve_fn(
483 &mut fd,
484 &TypeInfo {
485 variants: HashMap::new(),
486 variant_parents: HashMap::new(),
487 records: HashMap::new(),
488 },
489 );
490 let res = fd.resolution.as_ref().unwrap();
491 assert_eq!(res.local_count, 1, "wildcard param must claim a slot");
492 assert!(
493 !res.local_slots.contains_key("_"),
494 "wildcard param must not bind a readable name"
495 );
496 }
497
498 #[test]
499 fn leaves_globals_as_ident() {
500 let mut fd = FnDef {
501 name: "f".to_string(),
502 line: 1,
503 params: vec![("x".to_string(), "Int".to_string())],
504 return_type: "Int".to_string(),
505 effects: vec![],
506 desc: None,
507 body: Rc::new(FnBody::from_expr(Spanned::bare(Expr::FnCall(
508 Box::new(Spanned::bare(Expr::Ident("Console".to_string()))),
509 vec![Spanned::bare(Expr::Ident("x".to_string()))],
510 )))),
511 resolution: None,
512 };
513 resolve_fn(
514 &mut fd,
515 &TypeInfo {
516 variants: HashMap::new(),
517 variant_parents: HashMap::new(),
518 records: HashMap::new(),
519 },
520 );
521 match fd.body.tail_expr() {
522 Some(Spanned {
523 node: Expr::FnCall(func, args),
524 ..
525 }) => {
526 assert_eq!(func.node, Expr::Ident("Console".to_string()));
527 assert_eq!(
528 args[0].node,
529 Expr::Resolved {
530 slot: 0,
531 name: "x".to_string(),
532 last_use: AnnotBool(false)
533 }
534 );
535 }
536 other => panic!("unexpected body: {:?}", other),
537 }
538 }
539
540 #[test]
541 fn resolves_val_in_block_body() {
542 let mut fd = FnDef {
543 name: "f".to_string(),
544 line: 1,
545 params: vec![("x".to_string(), "Int".to_string())],
546 return_type: "Int".to_string(),
547 effects: vec![],
548 desc: None,
549 body: Rc::new(FnBody::Block(vec![
550 Stmt::Binding(
551 "y".to_string(),
552 None,
553 Spanned::bare(Expr::BinOp(
554 BinOp::Add,
555 Box::new(Spanned::bare(Expr::Ident("x".to_string()))),
556 Box::new(Spanned::bare(Expr::Literal(Literal::Int(1)))),
557 )),
558 ),
559 Stmt::Expr(Spanned::bare(Expr::Ident("y".to_string()))),
560 ])),
561 resolution: None,
562 };
563 resolve_fn(
564 &mut fd,
565 &TypeInfo {
566 variants: HashMap::new(),
567 variant_parents: HashMap::new(),
568 records: HashMap::new(),
569 },
570 );
571 let res = fd.resolution.as_ref().unwrap();
572 assert_eq!(res.local_slots["x"], 0);
573 assert_eq!(res.local_slots["y"], 1);
574 assert_eq!(res.local_count, 2);
575
576 let stmts = fd.body.stmts();
577 match &stmts[0] {
579 Stmt::Binding(
580 _,
581 _,
582 Spanned {
583 node: Expr::BinOp(_, left, _),
584 ..
585 },
586 ) => {
587 assert_eq!(
588 left.node,
589 Expr::Resolved {
590 slot: 0,
591 name: "x".to_string(),
592 last_use: AnnotBool(false)
593 }
594 );
595 }
596 other => panic!("unexpected stmt: {:?}", other),
597 }
598 match &stmts[1] {
600 Stmt::Expr(Spanned {
601 node: Expr::Resolved { slot: 1, .. },
602 ..
603 }) => {}
604 other => panic!("unexpected stmt: {:?}", other),
605 }
606 }
607
608 #[test]
609 fn resolves_match_pattern_bindings() {
610 let mut fd = FnDef {
612 name: "f".to_string(),
613 line: 1,
614 params: vec![("x".to_string(), "Int".to_string())],
615 return_type: "Int".to_string(),
616 effects: vec![],
617 desc: None,
618 body: Rc::new(FnBody::from_expr(Spanned::new(
619 Expr::Match {
620 subject: Box::new(Spanned::bare(Expr::Ident("x".to_string()))),
621 arms: vec![
622 MatchArm {
623 pattern: Pattern::Constructor(
624 "Result.Ok".to_string(),
625 vec!["v".to_string()],
626 ),
627 body: Box::new(Spanned::bare(Expr::Ident("v".to_string()))),
628 binding_slots: std::sync::OnceLock::new(),
629 },
630 MatchArm {
631 pattern: Pattern::Wildcard,
632 body: Box::new(Spanned::bare(Expr::Literal(Literal::Int(0)))),
633 binding_slots: std::sync::OnceLock::new(),
634 },
635 ],
636 },
637 1,
638 ))),
639 resolution: None,
640 };
641 resolve_fn(
642 &mut fd,
643 &TypeInfo {
644 variants: HashMap::new(),
645 variant_parents: HashMap::new(),
646 records: HashMap::new(),
647 },
648 );
649 let res = fd.resolution.as_ref().unwrap();
650 assert_eq!(res.local_slots["v"], 1);
652
653 match fd.body.tail_expr() {
654 Some(Spanned {
655 node: Expr::Match { arms, .. },
656 ..
657 }) => {
658 assert_eq!(
659 arms[0].body.node,
660 Expr::Resolved {
661 slot: 1,
662 name: "v".to_string(),
663 last_use: AnnotBool(false)
664 }
665 );
666 }
667 other => panic!("unexpected body: {:?}", other),
668 }
669 }
670
671 #[test]
672 fn resolves_match_pattern_bindings_inside_binding_initializer() {
673 let mut fd = FnDef {
674 name: "f".to_string(),
675 line: 1,
676 params: vec![("x".to_string(), "Int".to_string())],
677 return_type: "Int".to_string(),
678 effects: vec![],
679 desc: None,
680 body: Rc::new(FnBody::Block(vec![
681 Stmt::Binding(
682 "result".to_string(),
683 None,
684 Spanned::bare(Expr::Match {
685 subject: Box::new(Spanned::bare(Expr::Ident("x".to_string()))),
686 arms: vec![
687 MatchArm {
688 pattern: Pattern::Constructor(
689 "Option.Some".to_string(),
690 vec!["v".to_string()],
691 ),
692 body: Box::new(Spanned::bare(Expr::Ident("v".to_string()))),
693 binding_slots: std::sync::OnceLock::new(),
694 },
695 MatchArm {
696 pattern: Pattern::Wildcard,
697 body: Box::new(Spanned::bare(Expr::Literal(Literal::Int(0)))),
698 binding_slots: std::sync::OnceLock::new(),
699 },
700 ],
701 }),
702 ),
703 Stmt::Expr(Spanned::bare(Expr::Ident("result".to_string()))),
704 ])),
705 resolution: None,
706 };
707
708 resolve_fn(
709 &mut fd,
710 &TypeInfo {
711 variants: HashMap::new(),
712 variant_parents: HashMap::new(),
713 records: HashMap::new(),
714 },
715 );
716 let res = fd.resolution.as_ref().unwrap();
717 assert_eq!(res.local_slots["x"], 0);
718 assert_eq!(res.local_slots["result"], 1);
719 assert_eq!(res.local_slots["v"], 2);
720
721 let stmts = fd.body.stmts();
722 match &stmts[0] {
723 Stmt::Binding(
724 _,
725 _,
726 Spanned {
727 node: Expr::Match { arms, .. },
728 ..
729 },
730 ) => {
731 assert_eq!(
732 arms[0].body.node,
733 Expr::Resolved {
734 slot: 2,
735 name: "v".to_string(),
736 last_use: AnnotBool(false)
737 }
738 );
739 }
740 other => panic!("unexpected stmt: {:?}", other),
741 }
742
743 match &stmts[1] {
744 Stmt::Expr(Spanned {
745 node: Expr::Resolved { slot: 1, .. },
746 ..
747 }) => {}
748 other => panic!("unexpected stmt: {:?}", other),
749 }
750 }
751}