nandi.uk
Fork of daniellemaywood.uk/gleam — Wasm codegen work
26 kB
644 lines
1use std::collections::{HashMap, HashSet, VecDeque};
2
3use ecow::EcoString;
4
5use crate::cranelift::mir::{
6 ast::{CompleteType, Expression, Function, FunctionParameter, IncompleteType, Module, Var},
7 visit::Visit,
8};
9
10pub fn lower(modules: Vec<Module<IncompleteType>>) -> Vec<Module<CompleteType>> {
11 let (generics, functions): (HashMap<_, _>, HashMap<_, _>) = modules
12 .into_iter()
13 .flat_map(|module| {
14 module
15 .functions
16 .into_iter()
17 .map(|function| ((module.name.clone(), function.name.clone()), function))
18 .collect::<Vec<_>>()
19 })
20 .partition(|(_, function)| is_generic(function));
21
22 let mut functions = functions
23 .into_iter()
24 .map(|((module, name), function)| {
25 let function = Function::<CompleteType>::try_from(function)?;
26
27 let instantiation = Instantiation {
28 argument_types: function
29 .parameters
30 .iter()
31 .map(|parameter| parameter.type_.clone())
32 .collect(),
33 return_type: function.return_type.clone(),
34 };
35
36 Ok(((module, name, instantiation), function))
37 })
38 .collect::<Result<
39 HashMap<(EcoString, EcoString, Instantiation), Function<CompleteType>>,
40 NonGenericError,
41 >>()
42 .expect("expected all functions to be non-generic");
43
44 let mut monomorphizer = Monomorphizer::new(&generics);
45 for (_, function) in functions.iter_mut() {
46 monomorphizer.visit_function(function);
47 }
48
49 let mut queue: VecDeque<_> = monomorphizer.instantiations.into_iter().collect();
50
51 while let Some((module, name, instantiation)) = queue.pop_front() {
52 if let Some(function) = generics.get(&(module.clone(), name.clone())).cloned() {
53 let mut function = FunctionMonomorphizer::default()
54 .translate_function(function, instantiation.clone());
55
56 _ = functions.insert((module, name, instantiation), function.clone());
57
58 let mut monomorphizer = Monomorphizer::new(&generics);
59 monomorphizer.visit_function(&mut function);
60
61 for instantiation in monomorphizer.instantiations {
62 if !functions.contains_key(&instantiation) {
63 queue.push_back(instantiation);
64 }
65 }
66 }
67 }
68
69 let mut module_map: HashMap<EcoString, Vec<Function<CompleteType>>> = HashMap::new();
70 for ((module, _, _), function) in functions {
71 module_map
72 .entry(module)
73 .or_insert_with(Vec::new)
74 .push(function);
75 }
76
77 module_map
78 .into_iter()
79 .map(|(name, functions)| Module { name, functions })
80 .collect()
81}
82
83#[derive(Debug, Clone, PartialEq, Eq, Hash)]
84struct Instantiation {
85 argument_types: Vec<CompleteType>,
86 return_type: CompleteType,
87}
88
89#[derive(Default)]
90struct FunctionMonomorphizer {
91 generic_id_map: HashMap<u64, CompleteType>,
92}
93
94impl FunctionMonomorphizer {
95 fn collect_type_ids(&mut self, incomplete: &IncompleteType, complete: &CompleteType) {
96 match incomplete {
97 IncompleteType::Generic { id } => {
98 _ = self.generic_id_map.insert(*id, complete.clone());
99 }
100
101 IncompleteType::Func { arguments, returns } => {
102 let CompleteType::Func {
103 arguments: complete_arguments,
104 returns: complete_returns,
105 } = complete
106 else {
107 unreachable!("types should not be diverged")
108 };
109
110 for (argument, complete) in arguments.iter().zip(complete_arguments.iter()) {
111 self.collect_type_ids(argument, complete);
112 }
113
114 self.collect_type_ids(returns, complete_returns);
115 }
116 IncompleteType::Struct { elements } => {
117 let CompleteType::Struct { elements: complete } = complete else {
118 unreachable!("types should not be diverged")
119 };
120
121 for (element, complete) in elements.iter().zip(complete.iter()) {
122 self.collect_type_ids(element, complete);
123 }
124 }
125 IncompleteType::List(element_type) => {
126 let CompleteType::List(complete) = complete else {
127 unreachable!("types should not be diverged")
128 };
129
130 self.collect_type_ids(element_type, complete);
131 }
132
133 IncompleteType::Int
134 | IncompleteType::Float
135 | IncompleteType::Bool
136 | IncompleteType::String => {}
137 }
138 }
139
140 fn translate_function(
141 mut self,
142 function: Function<IncompleteType>,
143 instantiation: Instantiation,
144 ) -> Function<CompleteType> {
145 for (parameter, argument_type) in function
146 .parameters
147 .iter()
148 .zip(instantiation.argument_types.iter())
149 {
150 self.collect_type_ids(¶meter.type_, argument_type);
151 }
152
153 self.collect_type_ids(&function.return_type, &instantiation.return_type);
154
155 Function {
156 name: function.name,
157 return_type: instantiation.return_type,
158 parameters: function
159 .parameters
160 .into_iter()
161 .zip(instantiation.argument_types.into_iter())
162 .map(|(parameter, argument_type)| FunctionParameter {
163 name: parameter.name,
164 type_: argument_type,
165 })
166 .collect(),
167 body: self.translate_expression(function.body),
168 }
169 }
170
171 fn translate_expression(
172 &mut self,
173 expression: Expression<IncompleteType>,
174 ) -> Expression<CompleteType> {
175 match expression {
176 Expression::Block(expressions) => Expression::Block(
177 expressions
178 .into_iter()
179 .map(|expression| self.translate_expression(expression))
180 .collect(),
181 ),
182 Expression::FunctionRef {
183 module,
184 name,
185 arity,
186 type_,
187 } => Expression::FunctionRef {
188 module,
189 name,
190 arity,
191 type_: self.translate_type(type_),
192 },
193 Expression::Var(var) => Expression::Var(Var {
194 name: var.name,
195 type_: self.translate_type(var.type_),
196 }),
197 Expression::Int { value } => Expression::Int { value },
198 Expression::Float { value } => Expression::Float { value },
199 Expression::Bool { value } => Expression::Bool { value },
200 Expression::String { value } => Expression::String { value },
201 Expression::Equals { lhs, rhs } => Expression::Equals {
202 lhs: Box::new(self.translate_expression(*lhs)),
203 rhs: Box::new(self.translate_expression(*rhs)),
204 },
205 Expression::NotEquals { lhs, rhs } => Expression::NotEquals {
206 lhs: Box::new(self.translate_expression(*lhs)),
207 rhs: Box::new(self.translate_expression(*rhs)),
208 },
209 Expression::IntGt { lhs, rhs } => Expression::IntGt {
210 lhs: Box::new(self.translate_expression(*lhs)),
211 rhs: Box::new(self.translate_expression(*rhs)),
212 },
213 Expression::IntGtEq { lhs, rhs } => Expression::IntGtEq {
214 lhs: Box::new(self.translate_expression(*lhs)),
215 rhs: Box::new(self.translate_expression(*rhs)),
216 },
217 Expression::IntLt { lhs, rhs } => Expression::IntLt {
218 lhs: Box::new(self.translate_expression(*lhs)),
219 rhs: Box::new(self.translate_expression(*rhs)),
220 },
221 Expression::IntLtEq { lhs, rhs } => Expression::IntLtEq {
222 lhs: Box::new(self.translate_expression(*lhs)),
223 rhs: Box::new(self.translate_expression(*rhs)),
224 },
225 Expression::IntAdd { lhs, rhs } => Expression::IntAdd {
226 lhs: Box::new(self.translate_expression(*lhs)),
227 rhs: Box::new(self.translate_expression(*rhs)),
228 },
229 Expression::IntSub { lhs, rhs } => Expression::IntSub {
230 lhs: Box::new(self.translate_expression(*lhs)),
231 rhs: Box::new(self.translate_expression(*rhs)),
232 },
233 Expression::IntMul { lhs, rhs } => Expression::IntMul {
234 lhs: Box::new(self.translate_expression(*lhs)),
235 rhs: Box::new(self.translate_expression(*rhs)),
236 },
237 Expression::IntDiv { lhs, rhs } => Expression::IntDiv {
238 lhs: Box::new(self.translate_expression(*lhs)),
239 rhs: Box::new(self.translate_expression(*rhs)),
240 },
241 Expression::IntRem { lhs, rhs } => Expression::IntRem {
242 lhs: Box::new(self.translate_expression(*lhs)),
243 rhs: Box::new(self.translate_expression(*rhs)),
244 },
245 Expression::FloatGt { lhs, rhs } => Expression::FloatGt {
246 lhs: Box::new(self.translate_expression(*lhs)),
247 rhs: Box::new(self.translate_expression(*rhs)),
248 },
249 Expression::FloatGtEq { lhs, rhs } => Expression::FloatGtEq {
250 lhs: Box::new(self.translate_expression(*lhs)),
251 rhs: Box::new(self.translate_expression(*rhs)),
252 },
253 Expression::FloatLt { lhs, rhs } => Expression::FloatLt {
254 lhs: Box::new(self.translate_expression(*lhs)),
255 rhs: Box::new(self.translate_expression(*rhs)),
256 },
257 Expression::FloatLtEq { lhs, rhs } => Expression::FloatLtEq {
258 lhs: Box::new(self.translate_expression(*lhs)),
259 rhs: Box::new(self.translate_expression(*rhs)),
260 },
261 Expression::FloatAdd { lhs, rhs } => Expression::FloatAdd {
262 lhs: Box::new(self.translate_expression(*lhs)),
263 rhs: Box::new(self.translate_expression(*rhs)),
264 },
265 Expression::FloatSub { lhs, rhs } => Expression::FloatSub {
266 lhs: Box::new(self.translate_expression(*lhs)),
267 rhs: Box::new(self.translate_expression(*rhs)),
268 },
269 Expression::FloatMul { lhs, rhs } => Expression::FloatMul {
270 lhs: Box::new(self.translate_expression(*lhs)),
271 rhs: Box::new(self.translate_expression(*rhs)),
272 },
273 Expression::FloatDiv { lhs, rhs } => Expression::FloatDiv {
274 lhs: Box::new(self.translate_expression(*lhs)),
275 rhs: Box::new(self.translate_expression(*rhs)),
276 },
277 Expression::StringConcat { lhs, rhs } => Expression::StringConcat {
278 lhs: Box::new(self.translate_expression(*lhs)),
279 rhs: Box::new(self.translate_expression(*rhs)),
280 },
281 Expression::List { items, tail, type_ } => Expression::List {
282 items: items
283 .into_iter()
284 .map(|item| self.translate_expression(item))
285 .collect(),
286 tail: tail.map(|t| Box::new(self.translate_expression(*t))),
287 type_: self.translate_type(type_),
288 },
289 Expression::Struct { tag, items, type_ } => Expression::Struct {
290 tag,
291 items: items
292 .into_iter()
293 .map(|item| self.translate_expression(item))
294 .collect(),
295 type_: self.translate_type(type_),
296 },
297 Expression::StructTag { value } => Expression::StructTag {
298 value: Box::new(self.translate_expression(*value)),
299 },
300 Expression::StructAccess {
301 value,
302 index,
303 type_,
304 } => Expression::StructAccess {
305 value: Box::new(self.translate_expression(*value)),
306 index,
307 type_: self.translate_type(type_),
308 },
309 Expression::Set { name, value } => Expression::Set {
310 name: Var {
311 name: name.name,
312 type_: self.translate_type(name.type_),
313 },
314 value: Box::new(self.translate_expression(*value)),
315 },
316 Expression::If { cond, then, else_ } => Expression::If {
317 cond: Box::new(self.translate_expression(*cond)),
318 then: Box::new(self.translate_expression(*then)),
319 else_: Box::new(self.translate_expression(*else_)),
320 },
321 Expression::Call {
322 target,
323 args,
324 type_,
325 } => Expression::Call {
326 target: Box::new(self.translate_expression(*target)),
327 args: args
328 .into_iter()
329 .map(|arg| self.translate_expression(arg))
330 .collect(),
331 type_: self.translate_type(type_),
332 },
333 Expression::Panic { type_ } => Expression::Panic {
334 type_: self.translate_type(type_),
335 },
336 }
337 }
338
339 fn translate_type(&mut self, type_: IncompleteType) -> CompleteType {
340 match type_ {
341 IncompleteType::Int => CompleteType::Int,
342 IncompleteType::Float => CompleteType::Float,
343 IncompleteType::Bool => CompleteType::Bool,
344 IncompleteType::String => CompleteType::String,
345 IncompleteType::Func { arguments, returns } => CompleteType::Func {
346 arguments: arguments
347 .into_iter()
348 .map(|argument| self.translate_type(argument))
349 .collect(),
350 returns: Box::new(self.translate_type(*returns)),
351 },
352 IncompleteType::Struct { elements } => CompleteType::Struct {
353 elements: elements
354 .into_iter()
355 .map(|element| self.translate_type(element))
356 .collect(),
357 },
358 IncompleteType::List(element_type) => {
359 CompleteType::List(Box::new(self.translate_type(*element_type)))
360 }
361 IncompleteType::Generic { id } => self
362 .generic_id_map
363 .get(&id)
364 .cloned()
365 .expect("expected generic id to have a matching type"),
366 }
367 }
368}
369
370struct Monomorphizer<'g> {
371 generics: &'g HashMap<(EcoString, EcoString), Function<IncompleteType>>,
372 instantiations: HashSet<(EcoString, EcoString, Instantiation)>,
373}
374
375impl<'g> Monomorphizer<'g> {
376 fn new(generics: &'g HashMap<(EcoString, EcoString), Function<IncompleteType>>) -> Self {
377 Self {
378 generics,
379 instantiations: Default::default(),
380 }
381 }
382}
383
384impl<'mir, 'g> Visit<'mir, CompleteType> for Monomorphizer<'g> {
385 fn visit_expression_function_ref(
386 &mut self,
387 module: &mut EcoString,
388 name: &mut EcoString,
389 _arity: &mut usize,
390 type_: &mut CompleteType,
391 ) {
392 if !self.generics.contains_key(&(module.clone(), name.clone())) {
393 return;
394 }
395
396 let CompleteType::Func { arguments, returns } = type_ else {
397 unreachable!("a function ref should always have a func type")
398 };
399
400 let instantiation = Instantiation {
401 argument_types: arguments.clone(),
402 return_type: *returns.clone(),
403 };
404
405 _ = self
406 .instantiations
407 .insert((module.clone(), name.clone(), instantiation));
408 }
409}
410
411fn is_generic(function: &Function<IncompleteType>) -> bool {
412 let has_generic_parameter = function
413 .parameters
414 .iter()
415 .any(|parameter| matches!(parameter.type_, IncompleteType::Generic { .. }));
416
417 has_generic_parameter || matches!(function.return_type, IncompleteType::Generic { .. })
418}
419
420#[derive(Debug, Clone, Copy)]
421pub struct NonGenericError;
422
423impl TryFrom<Function<IncompleteType>> for Function<CompleteType> {
424 type Error = NonGenericError;
425
426 fn try_from(value: Function<IncompleteType>) -> Result<Self, Self::Error> {
427 Ok(Function {
428 name: value.name,
429 return_type: value.return_type.try_into()?,
430 parameters: value
431 .parameters
432 .into_iter()
433 .map(|parameter| {
434 Ok(FunctionParameter {
435 name: parameter.name,
436 type_: CompleteType::try_from(parameter.type_)?,
437 })
438 })
439 .collect::<Result<Vec<_>, _>>()?,
440 body: Expression::<CompleteType>::try_from(value.body)?,
441 })
442 }
443}
444
445impl TryFrom<Expression<IncompleteType>> for Expression<CompleteType> {
446 type Error = NonGenericError;
447
448 fn try_from(value: Expression<IncompleteType>) -> Result<Self, Self::Error> {
449 match value {
450 Expression::Block(expressions) => Ok(Expression::Block(
451 expressions
452 .into_iter()
453 .map(Expression::try_from)
454 .collect::<Result<Vec<_>, _>>()?,
455 )),
456 Expression::FunctionRef {
457 module,
458 name,
459 arity,
460 type_,
461 } => Ok(Expression::FunctionRef {
462 module,
463 name,
464 arity,
465 type_: type_.try_into()?,
466 }),
467 Expression::Var(var) => Ok(Expression::Var(Var {
468 name: var.name,
469 type_: var.type_.try_into()?,
470 })),
471 Expression::Int { value } => Ok(Expression::Int { value }),
472 Expression::Float { value } => Ok(Expression::Float { value }),
473 Expression::Bool { value } => Ok(Expression::Bool { value }),
474 Expression::String { value } => Ok(Expression::String { value }),
475 Expression::Equals { lhs, rhs } => Ok(Expression::Equals {
476 lhs: Box::new(Expression::try_from(*lhs)?),
477 rhs: Box::new(Expression::try_from(*rhs)?),
478 }),
479 Expression::NotEquals { lhs, rhs } => Ok(Expression::NotEquals {
480 lhs: Box::new(Expression::try_from(*lhs)?),
481 rhs: Box::new(Expression::try_from(*rhs)?),
482 }),
483 Expression::IntGt { lhs, rhs } => Ok(Expression::IntGt {
484 lhs: Box::new(Expression::try_from(*lhs)?),
485 rhs: Box::new(Expression::try_from(*rhs)?),
486 }),
487 Expression::IntGtEq { lhs, rhs } => Ok(Expression::IntGtEq {
488 lhs: Box::new(Expression::try_from(*lhs)?),
489 rhs: Box::new(Expression::try_from(*rhs)?),
490 }),
491 Expression::IntLt { lhs, rhs } => Ok(Expression::IntLt {
492 lhs: Box::new(Expression::try_from(*lhs)?),
493 rhs: Box::new(Expression::try_from(*rhs)?),
494 }),
495 Expression::IntLtEq { lhs, rhs } => Ok(Expression::IntLtEq {
496 lhs: Box::new(Expression::try_from(*lhs)?),
497 rhs: Box::new(Expression::try_from(*rhs)?),
498 }),
499 Expression::IntAdd { lhs, rhs } => Ok(Expression::IntAdd {
500 lhs: Box::new(Expression::try_from(*lhs)?),
501 rhs: Box::new(Expression::try_from(*rhs)?),
502 }),
503 Expression::IntSub { lhs, rhs } => Ok(Expression::IntSub {
504 lhs: Box::new(Expression::try_from(*lhs)?),
505 rhs: Box::new(Expression::try_from(*rhs)?),
506 }),
507 Expression::IntMul { lhs, rhs } => Ok(Expression::IntMul {
508 lhs: Box::new(Expression::try_from(*lhs)?),
509 rhs: Box::new(Expression::try_from(*rhs)?),
510 }),
511 Expression::IntDiv { lhs, rhs } => Ok(Expression::IntDiv {
512 lhs: Box::new(Expression::try_from(*lhs)?),
513 rhs: Box::new(Expression::try_from(*rhs)?),
514 }),
515 Expression::IntRem { lhs, rhs } => Ok(Expression::IntRem {
516 lhs: Box::new(Expression::try_from(*lhs)?),
517 rhs: Box::new(Expression::try_from(*rhs)?),
518 }),
519 Expression::FloatGt { lhs, rhs } => Ok(Expression::FloatGt {
520 lhs: Box::new(Expression::try_from(*lhs)?),
521 rhs: Box::new(Expression::try_from(*rhs)?),
522 }),
523 Expression::FloatGtEq { lhs, rhs } => Ok(Expression::FloatGtEq {
524 lhs: Box::new(Expression::try_from(*lhs)?),
525 rhs: Box::new(Expression::try_from(*rhs)?),
526 }),
527 Expression::FloatLt { lhs, rhs } => Ok(Expression::FloatLt {
528 lhs: Box::new(Expression::try_from(*lhs)?),
529 rhs: Box::new(Expression::try_from(*rhs)?),
530 }),
531 Expression::FloatLtEq { lhs, rhs } => Ok(Expression::FloatLtEq {
532 lhs: Box::new(Expression::try_from(*lhs)?),
533 rhs: Box::new(Expression::try_from(*rhs)?),
534 }),
535 Expression::FloatAdd { lhs, rhs } => Ok(Expression::FloatAdd {
536 lhs: Box::new(Expression::try_from(*lhs)?),
537 rhs: Box::new(Expression::try_from(*rhs)?),
538 }),
539 Expression::FloatSub { lhs, rhs } => Ok(Expression::FloatSub {
540 lhs: Box::new(Expression::try_from(*lhs)?),
541 rhs: Box::new(Expression::try_from(*rhs)?),
542 }),
543 Expression::FloatMul { lhs, rhs } => Ok(Expression::FloatMul {
544 lhs: Box::new(Expression::try_from(*lhs)?),
545 rhs: Box::new(Expression::try_from(*rhs)?),
546 }),
547 Expression::FloatDiv { lhs, rhs } => Ok(Expression::FloatDiv {
548 lhs: Box::new(Expression::try_from(*lhs)?),
549 rhs: Box::new(Expression::try_from(*rhs)?),
550 }),
551 Expression::StringConcat { lhs, rhs } => Ok(Expression::StringConcat {
552 lhs: Box::new(Expression::try_from(*lhs)?),
553 rhs: Box::new(Expression::try_from(*rhs)?),
554 }),
555 Expression::List { items, tail, type_ } => Ok(Expression::List {
556 items: items
557 .into_iter()
558 .map(Expression::try_from)
559 .collect::<Result<Vec<_>, _>>()?,
560 tail: tail
561 .map(|t| Ok(Box::new(Expression::try_from(*t)?)))
562 .transpose()?,
563 type_: type_.try_into()?,
564 }),
565 Expression::Struct { tag, items, type_ } => Ok(Expression::Struct {
566 tag,
567 items: items
568 .into_iter()
569 .map(Expression::try_from)
570 .collect::<Result<Vec<_>, _>>()?,
571 type_: type_.try_into()?,
572 }),
573 Expression::StructTag { value } => Ok(Expression::StructTag {
574 value: Box::new(Expression::try_from(*value)?),
575 }),
576 Expression::StructAccess {
577 value,
578 index,
579 type_,
580 } => Ok(Expression::StructAccess {
581 value: Box::new(Expression::try_from(*value)?),
582 index,
583 type_: type_.try_into()?,
584 }),
585 Expression::Set { name, value } => Ok(Expression::Set {
586 name: Var {
587 name: name.name,
588 type_: name.type_.try_into()?,
589 },
590 value: Box::new(Expression::try_from(*value)?),
591 }),
592 Expression::If { cond, then, else_ } => Ok(Expression::If {
593 cond: Box::new(Expression::try_from(*cond)?),
594 then: Box::new(Expression::try_from(*then)?),
595 else_: Box::new(Expression::try_from(*else_)?),
596 }),
597 Expression::Call {
598 target,
599 args,
600 type_,
601 } => Ok(Expression::Call {
602 target: Box::new(Expression::try_from(*target)?),
603 args: args
604 .into_iter()
605 .map(Expression::try_from)
606 .collect::<Result<Vec<_>, _>>()?,
607 type_: type_.try_into()?,
608 }),
609 Expression::Panic { type_ } => Ok(Expression::Panic {
610 type_: type_.try_into()?,
611 }),
612 }
613 }
614}
615
616impl TryFrom<IncompleteType> for CompleteType {
617 type Error = NonGenericError;
618
619 fn try_from(value: IncompleteType) -> Result<Self, Self::Error> {
620 match value {
621 IncompleteType::Int => Ok(Self::Int),
622 IncompleteType::Float => Ok(Self::Float),
623 IncompleteType::Bool => Ok(Self::Bool),
624 IncompleteType::String => Ok(Self::String),
625 IncompleteType::Func { arguments, returns } => Ok(Self::Func {
626 arguments: arguments
627 .into_iter()
628 .map(Self::try_from)
629 .collect::<Result<Vec<_>, _>>()?,
630 returns: Box::new(Self::try_from(*returns)?),
631 }),
632 IncompleteType::Struct { elements } => Ok(Self::Struct {
633 elements: elements
634 .into_iter()
635 .map(Self::try_from)
636 .collect::<Result<Vec<_>, _>>()?,
637 }),
638 IncompleteType::List(element_type) => {
639 Ok(Self::List(Box::new(Self::try_from(*element_type)?)))
640 }
641 IncompleteType::Generic { id: _ } => Err(NonGenericError),
642 }
643 }
644}