spirt/passes/
link.rs

1use crate::transform::{InnerTransform, Transformed, Transformer};
2use crate::visit::{InnerVisit, Visitor};
3use crate::{
4    AttrSet, Const, Context, DataInstForm, DeclDef, ExportKey, Exportee, Func, FxIndexSet,
5    GlobalVar, Import, Module, Type,
6};
7use rustc_hash::{FxHashMap, FxHashSet};
8use std::collections::VecDeque;
9
10// FIXME(eddyb) maybe make an export pruning pass that keeps some exports as
11// roots and then only other exports if they're used by imports.
12
13/// Remove exports which aren't "roots" (`is_root(export_key)` returns `false`),
14/// and which aren't otherwise kept alive by a "root" (through [`Import::LinkName`]
15/// declarations, with `name` matching [`ExportKey::LinkName`]), either directly
16/// or transitively (including through any number of imports).
17///
18/// In essence, other than the "root" exports, `minimize_exports` only keeps the
19/// exports that `resolve_imports` would use, and is recommended to first call
20/// `minimize_exports` before using [`resolve_imports`], to reduce work.
21///
22/// Note that the "dead" definitions are not removed from the module, and any
23/// external references to them could still be used (e.g. from a clone of the
24/// `module.exports` map, before calling `minimize_exports`).
25//
26// FIXME(eddyb) make this operate on multiple modules.
27pub fn minimize_exports(module: &mut Module, is_root: impl Fn(&ExportKey) -> bool) {
28    let mut collector = LiveExportCollector {
29        cx: module.cx_ref(),
30        module,
31
32        live_exports: FxIndexSet::default(),
33
34        seen_types: FxHashSet::default(),
35        seen_consts: FxHashSet::default(),
36        seen_data_inst_forms: FxHashSet::default(),
37        seen_global_vars: FxHashSet::default(),
38        seen_funcs: FxHashSet::default(),
39    };
40    for (export_key, &exportee) in &module.exports {
41        if is_root(export_key) && collector.live_exports.insert(export_key.clone()) {
42            exportee.inner_visit_with(&mut collector);
43        }
44    }
45    module.exports = collector
46        .live_exports
47        .into_iter()
48        .map(|export_key| {
49            let exportee = module.exports[&export_key];
50            (export_key, exportee)
51        })
52        .collect();
53}
54
55struct LiveExportCollector<'a> {
56    cx: &'a Context,
57    module: &'a Module,
58
59    live_exports: FxIndexSet<ExportKey>,
60
61    // FIXME(eddyb) build some automation to avoid ever repeating these.
62    seen_types: FxHashSet<Type>,
63    seen_consts: FxHashSet<Const>,
64    seen_data_inst_forms: FxHashSet<DataInstForm>,
65    seen_global_vars: FxHashSet<GlobalVar>,
66    seen_funcs: FxHashSet<Func>,
67}
68
69impl Visitor<'_> for LiveExportCollector<'_> {
70    // FIXME(eddyb) build some automation to avoid ever repeating these.
71    fn visit_attr_set_use(&mut self, _attrs: AttrSet) {
72        // FIXME(eddyb) if `AttrSet`s are ignored, why not `Type`s too?
73    }
74    fn visit_type_use(&mut self, ty: Type) {
75        if self.seen_types.insert(ty) {
76            self.visit_type_def(&self.cx[ty]);
77        }
78    }
79    fn visit_const_use(&mut self, ct: Const) {
80        if self.seen_consts.insert(ct) {
81            self.visit_const_def(&self.cx[ct]);
82        }
83    }
84    fn visit_data_inst_form_use(&mut self, data_inst_form: DataInstForm) {
85        if self.seen_data_inst_forms.insert(data_inst_form) {
86            self.visit_data_inst_form_def(&self.cx[data_inst_form]);
87        }
88    }
89
90    fn visit_global_var_use(&mut self, gv: GlobalVar) {
91        if self.seen_global_vars.insert(gv) {
92            self.visit_global_var_decl(&self.module.global_vars[gv]);
93        }
94    }
95    fn visit_func_use(&mut self, func: Func) {
96        if self.seen_funcs.insert(func) {
97            self.visit_func_decl(&self.module.funcs[func]);
98        }
99    }
100
101    fn visit_import(&mut self, import: &Import) {
102        match *import {
103            Import::LinkName(name) => {
104                let export_key = ExportKey::LinkName(name);
105                if let Some(&exportee) = self.module.exports.get(&export_key) {
106                    if self.live_exports.insert(export_key) {
107                        exportee.inner_visit_with(self);
108                    }
109                }
110            }
111        }
112    }
113}
114
115/// Remap [`Import::LinkName`] to definitions exported as [`ExportKey::LinkName`].
116///
117/// To reduce the work performed, calling [`minimize_exports`] first is recommended.
118//
119// FIXME(eddyb) make this operate on multiple modules.
120pub fn resolve_imports(module: &mut Module) {
121    let (resolved_global_vars, resolved_funcs) = {
122        let mut collector = ImportResolutionCollector {
123            cx: module.cx_ref(),
124            module,
125
126            resolved_global_vars: FxHashMap::default(),
127            resolved_funcs: FxHashMap::default(),
128
129            seen_types: FxHashSet::default(),
130            seen_consts: FxHashSet::default(),
131            seen_data_inst_forms: FxHashSet::default(),
132            seen_global_vars: FxHashSet::default(),
133            seen_funcs: FxHashSet::default(),
134        };
135        collector.visit_module(module);
136        (collector.resolved_global_vars, collector.resolved_funcs)
137    };
138
139    let mut resolver = ImportResolver {
140        cx: &module.cx(),
141
142        resolved_global_vars: &resolved_global_vars,
143        resolved_funcs: &resolved_funcs,
144
145        transformed_types: FxHashMap::default(),
146        transformed_consts: FxHashMap::default(),
147        transformed_data_inst_forms: FxHashMap::default(),
148        transformed_global_vars: FxHashMap::default(),
149        global_var_queue: VecDeque::new(),
150        transformed_funcs: FxHashMap::default(),
151        func_queue: VecDeque::new(),
152    };
153
154    // Seed the queues starting from the module exports.
155    for exportee in module.exports.values_mut() {
156        exportee.inner_transform_with(&mut resolver).apply_to(exportee);
157    }
158
159    // Process the queues until they're all empty.
160    while !resolver.global_var_queue.is_empty() || !resolver.func_queue.is_empty() {
161        while let Some(gv) = resolver.global_var_queue.pop_front() {
162            resolver.in_place_transform_global_var_decl(&mut module.global_vars[gv]);
163        }
164        while let Some(func) = resolver.func_queue.pop_front() {
165            resolver.in_place_transform_func_decl(&mut module.funcs[func]);
166        }
167    }
168}
169
170// FIXME(eddyb) figure out if this step can be skipped by somehow letting
171// `ImportResolver` access declarations while mutating definitions.
172struct ImportResolutionCollector<'a> {
173    cx: &'a Context,
174    module: &'a Module,
175
176    resolved_global_vars: FxHashMap<GlobalVar, GlobalVar>,
177    resolved_funcs: FxHashMap<Func, Func>,
178
179    // FIXME(eddyb) build some automation to avoid ever repeating these.
180    seen_types: FxHashSet<Type>,
181    seen_consts: FxHashSet<Const>,
182    seen_data_inst_forms: FxHashSet<DataInstForm>,
183    seen_global_vars: FxHashSet<GlobalVar>,
184    seen_funcs: FxHashSet<Func>,
185}
186
187impl Visitor<'_> for ImportResolutionCollector<'_> {
188    // FIXME(eddyb) build some automation to avoid ever repeating these.
189    fn visit_attr_set_use(&mut self, _attrs: AttrSet) {
190        // FIXME(eddyb) if `AttrSet`s are ignored, why not `Type`s too?
191    }
192    fn visit_type_use(&mut self, ty: Type) {
193        if self.seen_types.insert(ty) {
194            self.visit_type_def(&self.cx[ty]);
195        }
196    }
197    fn visit_const_use(&mut self, ct: Const) {
198        if self.seen_consts.insert(ct) {
199            self.visit_const_def(&self.cx[ct]);
200        }
201    }
202    fn visit_data_inst_form_use(&mut self, data_inst_form: DataInstForm) {
203        if self.seen_data_inst_forms.insert(data_inst_form) {
204            self.visit_data_inst_form_def(&self.cx[data_inst_form]);
205        }
206    }
207
208    fn visit_global_var_use(&mut self, gv: GlobalVar) {
209        if self.seen_global_vars.insert(gv) {
210            let gv_decl = &self.module.global_vars[gv];
211            self.visit_global_var_decl(gv_decl);
212
213            // FIXME(eddyb) if the export is missing (or the wrong kind), it will
214            // simply not get remapped - perhaps some kind of diagnostic is in
215            // order? (maybe an entire pass infrastructure that can report errors)
216            if let DeclDef::Imported(Import::LinkName(name)) = gv_decl.def {
217                if let Some(&Exportee::GlobalVar(def_gv)) =
218                    self.module.exports.get(&ExportKey::LinkName(name))
219                {
220                    self.resolved_global_vars.insert(gv, def_gv);
221                }
222            }
223        }
224    }
225    fn visit_func_use(&mut self, func: Func) {
226        if self.seen_funcs.insert(func) {
227            let func_decl = &self.module.funcs[func];
228            self.visit_func_decl(func_decl);
229
230            // FIXME(eddyb) if the export is missing (or the wrong kind), it will
231            // simply not get remapped - perhaps some kind of diagnostic is in
232            // order? (maybe an entire pass infrastructure that can report errors)
233            if let DeclDef::Imported(Import::LinkName(name)) = func_decl.def {
234                if let Some(&Exportee::Func(def_func)) =
235                    self.module.exports.get(&ExportKey::LinkName(name))
236                {
237                    self.resolved_funcs.insert(func, def_func);
238                }
239            }
240        }
241    }
242}
243
244struct ImportResolver<'a> {
245    cx: &'a Context,
246
247    resolved_global_vars: &'a FxHashMap<GlobalVar, GlobalVar>,
248    resolved_funcs: &'a FxHashMap<Func, Func>,
249
250    // FIXME(eddyb) build some automation to avoid ever repeating these.
251    transformed_types: FxHashMap<Type, Transformed<Type>>,
252    transformed_consts: FxHashMap<Const, Transformed<Const>>,
253    transformed_data_inst_forms: FxHashMap<DataInstForm, Transformed<DataInstForm>>,
254    transformed_global_vars: FxHashMap<GlobalVar, Transformed<GlobalVar>>,
255    global_var_queue: VecDeque<GlobalVar>,
256    transformed_funcs: FxHashMap<Func, Transformed<Func>>,
257    func_queue: VecDeque<Func>,
258}
259
260impl Transformer for ImportResolver<'_> {
261    // FIXME(eddyb) build some automation to avoid ever repeating these.
262    fn transform_type_use(&mut self, ty: Type) -> Transformed<Type> {
263        if let Some(&cached) = self.transformed_types.get(&ty) {
264            return cached;
265        }
266        let transformed =
267            self.transform_type_def(&self.cx[ty]).map(|ty_def| self.cx.intern(ty_def));
268        self.transformed_types.insert(ty, transformed);
269        transformed
270    }
271    fn transform_const_use(&mut self, ct: Const) -> Transformed<Const> {
272        if let Some(&cached) = self.transformed_consts.get(&ct) {
273            return cached;
274        }
275        let transformed =
276            self.transform_const_def(&self.cx[ct]).map(|ct_def| self.cx.intern(ct_def));
277        self.transformed_consts.insert(ct, transformed);
278        transformed
279    }
280    fn transform_data_inst_form_use(
281        &mut self,
282        data_inst_form: DataInstForm,
283    ) -> Transformed<DataInstForm> {
284        if let Some(&cached) = self.transformed_data_inst_forms.get(&data_inst_form) {
285            return cached;
286        }
287        let transformed = self
288            .transform_data_inst_form_def(&self.cx[data_inst_form])
289            .map(|data_inst_form_def| self.cx.intern(data_inst_form_def));
290        self.transformed_data_inst_forms.insert(data_inst_form, transformed);
291        transformed
292    }
293
294    fn transform_global_var_use(&mut self, gv: GlobalVar) -> Transformed<GlobalVar> {
295        if let Some(&cached) = self.transformed_global_vars.get(&gv) {
296            return cached;
297        }
298        let transformed = match self.resolved_global_vars.get(&gv).copied() {
299            Some(mut new_gv) => {
300                self.transform_global_var_use(new_gv).apply_to(&mut new_gv);
301                Transformed::Changed(new_gv)
302            }
303            None => {
304                self.global_var_queue.push_back(gv);
305                Transformed::Unchanged
306            }
307        };
308        self.transformed_global_vars.insert(gv, transformed);
309        transformed
310    }
311    fn transform_func_use(&mut self, func: Func) -> Transformed<Func> {
312        if let Some(&cached) = self.transformed_funcs.get(&func) {
313            return cached;
314        }
315        let transformed = match self.resolved_funcs.get(&func).copied() {
316            Some(mut new_func) => {
317                self.transform_func_use(new_func).apply_to(&mut new_func);
318                Transformed::Changed(new_func)
319            }
320            None => {
321                self.func_queue.push_back(func);
322                Transformed::Unchanged
323            }
324        };
325        self.transformed_funcs.insert(func, transformed);
326        transformed
327    }
328}
spirt/passes/link.rs

spirt/passes/
link.rs
