1use derive_generic_visitor::*;
3use std::collections::{HashMap, HashSet};
4
5use crate::ast::*;
6use crate::transform::TransformCtx;
7use std::fmt::Debug;
8
9use super::ctx::TransformPass;
10
11enum OptionHint<T, H> {
12 Some(T),
13 None,
14 Hint(H),
15}
16
17impl<T, H> OptionHint<T, H> {
18 fn is_some(&self) -> bool {
19 match self {
20 OptionHint::Some(_) => true,
21 OptionHint::None => false,
22 OptionHint::Hint(_) => false,
23 }
24 }
25
26 fn hint_or<'a>(&'a self, hint: &'a H) -> &'a H {
27 match self {
28 OptionHint::Some(_) => hint,
29 OptionHint::None => hint,
30 OptionHint::Hint(h) => h,
31 }
32 }
33}
34
35struct PassData {
36 items: HashMap<(AnyTransId, GenericArgs), OptionHint<AnyTransId, (AnyTransId, BoxedArgs)>>,
39 worklist: Vec<AnyTransId>,
40 visited: HashSet<AnyTransId>,
41}
42
43impl PassData {
44 fn new() -> Self {
45 PassData {
46 items: HashMap::new(),
47 worklist: Vec::new(),
48 visited: HashSet::new(),
49 }
50 }
51}
52
53impl TranslatedCrate {
54 fn find_trait_impl_and_gargs(
56 self: &Self,
57 kind: &TraitRefKind,
58 ) -> Option<(&TraitImpl, GenericArgs)> {
59 match kind {
60 TraitRefKind::TraitImpl(impl_ref) => {
61 let trait_impl = self.trait_impls.get(impl_ref.id)?;
62 Some((trait_impl, impl_ref.generics.as_ref().clone()))
63 }
64 TraitRefKind::ParentClause(p, _, clause) => {
65 let (trait_impl, _) = self.find_trait_impl_and_gargs(p)?;
66 let t_ref = trait_impl.parent_trait_refs.get(*clause)?;
67 self.find_trait_impl_and_gargs(&t_ref.kind)
68 }
69 _ => None,
70 }
71 }
72}
73
74#[derive(Visitor)]
75struct UsageVisitor<'a> {
76 data: &'a mut PassData,
77 krate: &'a TranslatedCrate,
78}
79impl UsageVisitor<'_> {
80 fn found_use(
81 &mut self,
82 id: &AnyTransId,
83 gargs: &GenericArgs,
84 default: OptionHint<AnyTransId, (AnyTransId, BoxedArgs)>,
85 ) {
86 trace!("Mono: Found use: {:?} / {:?}", id, gargs);
87 self.data
88 .items
89 .entry((*id, gargs.clone()))
90 .or_insert(default);
91 }
92 fn found_use_ty(&mut self, tref: &TypeDeclRef) {
93 match tref.id {
94 TypeId::Adt(id) => {
95 self.found_use(&AnyTransId::Type(id), &tref.generics, OptionHint::None)
96 }
97 _ => {}
98 }
99 }
100 fn found_use_fn(&mut self, id: &FunDeclId, gargs: &GenericArgs) {
101 self.found_use(&AnyTransId::Fun(*id), gargs, OptionHint::None);
102 }
103 fn found_use_global_decl_ref(&mut self, id: &GlobalDeclId, gargs: &GenericArgs) {
104 self.found_use(&AnyTransId::Global(*id), gargs, OptionHint::None);
105 }
106 fn found_use_fn_hinted(
107 &mut self,
108 id: &FunDeclId,
109 gargs: &GenericArgs,
110 (h_id, h_args): (FunDeclId, BoxedArgs),
111 ) {
112 self.found_use(
113 &AnyTransId::Fun(*id),
114 gargs,
115 OptionHint::Hint((AnyTransId::Fun(h_id), h_args)),
116 );
117 }
118}
119impl VisitAst for UsageVisitor<'_> {
120 fn visit_item_meta(&mut self, _: &ItemMeta) -> ControlFlow<Infallible> {
122 Continue(())
123 }
124
125 fn enter_aggregate_kind(&mut self, kind: &AggregateKind) {
126 match kind {
127 AggregateKind::Adt(tref, _, _) => self.found_use_ty(tref),
128 _ => {}
129 }
130 }
131
132 fn enter_ty_kind(&mut self, kind: &TyKind) {
133 match kind {
134 TyKind::Adt(tref) => {
135 self.found_use_ty(tref);
136 }
137 _ => {}
138 }
139 }
140
141 fn enter_fn_ptr(&mut self, fn_ptr: &FnPtr) {
142 match fn_ptr.func.as_ref() {
143 FunIdOrTraitMethodRef::Fun(FunId::Regular(id)) => {
144 self.found_use_fn(&id, &fn_ptr.generics)
145 }
146 FunIdOrTraitMethodRef::Trait(t_ref, name, id) => {
147 let Some((trait_impl, impl_gargs)) =
148 self.krate.find_trait_impl_and_gargs(&t_ref.kind)
149 else {
150 return;
151 };
152 let (_, bound_fn) = trait_impl.methods().find(|(n, _)| n == name).unwrap();
153 let fn_ref: Binder<Binder<FunDeclRef>> = Binder::new(
154 BinderKind::Other,
155 trait_impl.generics.clone(),
156 bound_fn.clone(),
157 );
158 let fn_ref = fn_ref.apply(&impl_gargs).apply(&fn_ptr.generics);
161 let gargs_key = fn_ptr
162 .generics
163 .clone()
164 .concat(&t_ref.trait_decl_ref.skip_binder.generics);
165 self.found_use_fn_hinted(&id, &gargs_key, (fn_ref.id, fn_ref.generics))
166 }
167 FunIdOrTraitMethodRef::Fun(FunId::Builtin(..)) => {}
169 }
170 }
171
172 fn enter_global_decl_ref(&mut self, glob: &GlobalDeclRef) {
173 self.found_use_global_decl_ref(&glob.id, &glob.generics);
174 }
175}
176
177#[derive(Visitor)]
181struct SubstVisitor<'a> {
182 data: &'a PassData,
183}
184impl SubstVisitor<'_> {
185 fn subst_use<T, F>(&mut self, id: &mut T, gargs: &mut GenericArgs, of: F)
186 where
187 T: Into<AnyTransId> + Debug + Copy,
188 F: Fn(&AnyTransId) -> Option<&T>,
189 {
190 trace!("Mono: Subst use: {:?} / {:?}", id, gargs);
191 gargs.regions.iter_mut().for_each(|r| *r = Region::Erased);
193 let key = ((*id).into(), gargs.clone());
194 let subst = self.data.items.get(&key);
195 if let Some(OptionHint::Some(any_id)) = subst
196 && let Some(subst_id) = of(any_id)
197 {
198 *id = *subst_id;
199 *gargs = GenericArgs::empty();
200 } else {
201 warn!("Substitution missing for {:?} / {:?}", id, gargs);
202 }
203 }
204 fn subst_use_ty(&mut self, tref: &mut TypeDeclRef) {
205 match &mut tref.id {
206 TypeId::Adt(id) => {
207 self.subst_use(id, &mut tref.generics, AnyTransId::as_type);
208 }
209 _ => {}
210 }
211 }
212 fn subst_use_fun(&mut self, id: &mut FunDeclId, gargs: &mut GenericArgs) {
213 self.subst_use(id, gargs, AnyTransId::as_fun);
214 }
215 fn subst_use_glob(&mut self, id: &mut GlobalDeclId, gargs: &mut GenericArgs) {
216 self.subst_use(id, gargs, AnyTransId::as_global);
217 }
218}
219
220impl VisitAstMut for SubstVisitor<'_> {
221 fn exit_rvalue(&mut self, rval: &mut Rvalue) {
222 if let Rvalue::Discriminant(place, id) = rval
223 && let Some(tref) = place.ty.as_adt()
224 && let TypeId::Adt(new_enum_id) = tref.id
225 {
226 *id = new_enum_id;
230 }
231 }
232
233 fn enter_aggregate_kind(&mut self, kind: &mut AggregateKind) {
234 match kind {
235 AggregateKind::Adt(tref, _, _) => self.subst_use_ty(tref),
236 _ => {}
237 }
238 }
239
240 fn enter_ty_kind(&mut self, kind: &mut TyKind) {
241 match kind {
242 TyKind::Adt(tref) => self.subst_use_ty(tref),
243 _ => {}
244 }
245 }
246
247 fn enter_fn_ptr(&mut self, fn_ptr: &mut FnPtr) {
248 match fn_ptr.func.as_mut() {
249 FunIdOrTraitMethodRef::Fun(FunId::Regular(fun_id)) => {
250 self.subst_use_fun(fun_id, &mut fn_ptr.generics)
251 }
252 FunIdOrTraitMethodRef::Trait(t_ref, _, fun_id) => {
253 let mut gargs_key = fn_ptr
254 .generics
255 .clone()
256 .concat(&t_ref.trait_decl_ref.skip_binder.generics);
257 self.subst_use_fun(fun_id, &mut gargs_key);
258 fn_ptr.generics = Box::new(gargs_key);
259 }
260 FunIdOrTraitMethodRef::Fun(FunId::Builtin(..)) => {}
262 }
263 }
264
265 fn exit_place(&mut self, place: &mut Place) {
266 match &mut place.kind {
267 PlaceKind::Projection(inner, ProjectionElem::Field(FieldProjKind::Adt(id, _), _)) => {
269 let tref = inner.ty.as_adt().unwrap();
272 *id = *tref.id.as_adt().unwrap()
273 }
274 _ => {}
275 }
276 }
277
278 fn enter_global_decl_ref(&mut self, glob: &mut GlobalDeclRef) {
279 self.subst_use_glob(&mut glob.id, &mut glob.generics);
280 }
281}
282
283#[derive(Visitor)]
284#[allow(dead_code)]
285struct MissingIndexChecker<'a> {
286 krate: &'a TranslatedCrate,
287 current_item: Option<AnyTransItem<'a>>,
288}
289impl VisitAst for MissingIndexChecker<'_> {
290 fn enter_fun_decl_id(&mut self, id: &FunDeclId) {
291 if self.krate.fun_decls.get(*id).is_none() {
292 panic!(
293 "Missing function declaration for id: {:?}, in {:?}",
294 id, self.current_item
295 );
296 }
297 }
298
299 fn enter_trait_impl_id(&mut self, id: &TraitImplId) {
300 if self.krate.trait_impls.get(*id).is_none() {
301 panic!(
302 "Missing trait implementation for id: {:?}, in {:?}",
303 id, self.current_item
304 );
305 }
306 }
307
308 fn enter_trait_decl_id(&mut self, id: &TraitDeclId) {
309 if self.krate.trait_decls.get(*id).is_none() {
310 panic!(
311 "Missing trait declaration for id: {:?}, in {:?}",
312 id, self.current_item
313 );
314 }
315 }
316
317 fn enter_type_decl_id(&mut self, id: &TypeDeclId) {
318 if self.krate.type_decls.get(*id).is_none() {
319 panic!(
320 "Missing type declaration for id: {:?}, in {:?}",
321 id, self.current_item
322 );
323 }
324 }
325}
326
327fn find_uses(data: &mut PassData, krate: &TranslatedCrate, item: &AnyTransItem) {
328 let mut visitor = UsageVisitor { data, krate };
329 let _ = item.drive(&mut visitor);
330}
331
332fn subst_uses<T: AstVisitable + Debug>(data: &PassData, item: &mut T) {
333 let mut visitor = SubstVisitor { data };
334 let _ = item.drive_mut(&mut visitor);
335}
336
337pub struct Transform;
353impl TransformPass for Transform {
354 fn transform_ctx(&self, ctx: &mut TransformCtx) {
355 if !ctx.options.monomorphize {
357 return;
358 }
359
360 let mut data = PassData::new();
390
391 let empty_gargs = GenericArgs::empty();
392
393 for (id, item) in ctx.translated.all_items_with_ids() {
395 match item {
396 AnyTransItem::Fun(f) if f.signature.generics.is_empty() => {
397 data.items
398 .insert((id, empty_gargs.clone()), OptionHint::Some(id));
399 data.worklist.push(id);
400 }
401 _ => {}
402 }
403 }
404
405 while let Some(id) = data.worklist.pop() {
407 if data.visited.contains(&id) {
408 continue;
409 }
410 data.visited.insert(id);
411
412 let Some(item) = ctx.translated.get_item(id) else {
414 panic!("Couldn't find item {:} in translated items.", id)
415 };
416 find_uses(&mut data, &ctx.translated, &item);
417
418 for ((id, gargs), mono) in data.items.iter_mut() {
420 if mono.is_some() {
421 continue;
422 }
423
424 let new_mono = if gargs.is_empty() {
426 *id
427 } else {
428 match id {
429 AnyTransId::Fun(_) => {
430 let key_pair = (id.clone(), Box::new(gargs.clone()));
431 let (AnyTransId::Fun(fun_id), gargs) = mono.hint_or(&key_pair) else {
432 panic!("Unexpected ID type in hint_or");
433 };
434 let fun = ctx.translated.fun_decls.get(*fun_id).unwrap();
435 let mut fun_sub = fun.clone().substitute(gargs);
436 fun_sub.signature.generics = GenericParams::empty();
437 fun_sub
438 .item_meta
439 .name
440 .name
441 .push(PathElem::Monomorphized(gargs.clone()));
442
443 let fun_id_sub = ctx.translated.fun_decls.push_with(|id| {
444 fun_sub.def_id = id;
445 fun_sub
446 });
447
448 AnyTransId::Fun(fun_id_sub)
449 }
450 AnyTransId::Type(typ_id) => {
451 let typ = ctx.translated.type_decls.get(*typ_id).unwrap();
452 let mut typ_sub = typ.clone().substitute(gargs);
453 typ_sub.generics = GenericParams::empty();
454 typ_sub
455 .item_meta
456 .name
457 .name
458 .push(PathElem::Monomorphized(gargs.clone().into()));
459
460 let typ_id_sub = ctx.translated.type_decls.push_with(|id| {
461 typ_sub.def_id = id;
462 typ_sub
463 });
464
465 AnyTransId::Type(typ_id_sub)
466 }
467 AnyTransId::Global(g_id) => {
468 let Some(glob) = ctx.translated.global_decls.get(*g_id) else {
469 *mono = OptionHint::Some(*id);
471 warn!("Found a global that has no associated declaration");
472 continue;
473 };
474 let mut glob_sub = glob.clone().substitute(gargs);
475 glob_sub.generics = GenericParams::empty();
476 glob_sub
477 .item_meta
478 .name
479 .name
480 .push(PathElem::Monomorphized(gargs.clone().into()));
481
482 let init = ctx.translated.fun_decls.get(glob.init).unwrap();
483 let mut init_sub = init.clone().substitute(gargs);
484 init_sub.signature.generics = GenericParams::empty();
485 init_sub
486 .item_meta
487 .name
488 .name
489 .push(PathElem::Monomorphized(gargs.clone().into()));
490
491 let init_id_sub = ctx.translated.fun_decls.push_with(|id| {
492 init_sub.def_id = id;
493 glob_sub.init = id;
494 init_sub
495 });
496
497 let g_id_sub = ctx.translated.global_decls.push_with(|id| {
498 glob_sub.def_id = id;
499 glob_sub
500 });
501
502 data.worklist.push(AnyTransId::Fun(init_id_sub));
503
504 AnyTransId::Global(g_id_sub)
505 }
506 _ => todo!("Unhandled monomorphization target ID {:?}", id),
507 }
508 };
509 trace!(
510 "Mono: Monomorphized {:?} with {:?} to {:?}",
511 id,
512 gargs,
513 new_mono
514 );
515 if id != &new_mono {
516 trace!(" - From {:?}", ctx.translated.get_item(id.clone()));
517 trace!(" - To {:?}", ctx.translated.get_item(new_mono.clone()));
518 }
519 *mono = OptionHint::Some(new_mono);
520 data.worklist.push(new_mono);
521
522 let item = ctx.translated.get_item(new_mono).unwrap();
523 ctx.translated
524 .item_names
525 .insert(new_mono, item.item_meta().name.clone());
526 }
527
528 let Some(item) = ctx.translated.get_item_mut(id) else {
530 panic!("Couldn't find item {:} in translated items.", id)
531 };
532 match item {
533 AnyTransItemMut::Fun(f) => subst_uses(&data, f),
534 AnyTransItemMut::Type(t) => subst_uses(&data, t),
535 AnyTransItemMut::TraitImpl(t) => subst_uses(&data, t),
536 AnyTransItemMut::Global(g) => subst_uses(&data, g),
537 AnyTransItemMut::TraitDecl(t) => subst_uses(&data, t),
538 };
539 }
540
541 ctx.translated
544 .fun_decls
545 .retain(|f| data.visited.contains(&AnyTransId::Fun(f.def_id)));
546 ctx.translated
547 .type_decls
548 .retain(|t| data.visited.contains(&AnyTransId::Type(t.def_id)));
549 ctx.translated
550 .global_decls
551 .retain(|g| data.visited.contains(&AnyTransId::Global(g.def_id)));
552 }
559}