spirt/
visit.rs

1//! Immutable IR traversal.
2
3use crate::func_at::FuncAt;
4use crate::qptr::{self, QPtrAttr, QPtrMemUsage, QPtrMemUsageKind, QPtrOp, QPtrUsage};
5use crate::{
6    AddrSpace, Attr, AttrSet, AttrSetDef, Const, ConstDef, ConstKind, ControlNode, ControlNodeDef,
7    ControlNodeKind, ControlNodeOutputDecl, ControlRegion, ControlRegionDef,
8    ControlRegionInputDecl, DataInstDef, DataInstForm, DataInstFormDef, DataInstKind, DeclDef,
9    DiagMsgPart, EntityListIter, ExportKey, Exportee, Func, FuncDecl, FuncDefBody, FuncParam,
10    GlobalVar, GlobalVarDecl, GlobalVarDefBody, Import, Module, ModuleDebugInfo, ModuleDialect,
11    SelectionKind, Type, TypeDef, TypeKind, TypeOrConst, Value, cfg, spv,
12};
13
14// FIXME(eddyb) `Sized` bound shouldn't be needed but removing it requires
15// writing `impl Visitor<'a> + ?Sized` in `fn inner_visit_with` signatures.
16pub trait Visitor<'a>: Sized {
17    // Context-interned leaves (no default provided).
18    // FIXME(eddyb) treat these separately somehow and allow e.g. automatic deep
19    // visiting (with a set to avoid repeat visits) if a `Rc<Context>` is provided.
20    fn visit_attr_set_use(&mut self, attrs: AttrSet);
21    fn visit_type_use(&mut self, ty: Type);
22    fn visit_const_use(&mut self, ct: Const);
23    fn visit_data_inst_form_use(&mut self, data_inst_form: DataInstForm);
24
25    // Module-stored entity leaves (no default provided).
26    fn visit_global_var_use(&mut self, gv: GlobalVar);
27    fn visit_func_use(&mut self, func: Func);
28
29    // Leaves (noop default behavior).
30    fn visit_spv_dialect(&mut self, _dialect: &spv::Dialect) {}
31    fn visit_spv_module_debug_info(&mut self, _debug_info: &spv::ModuleDebugInfo) {}
32    fn visit_import(&mut self, _import: &Import) {}
33
34    // Non-leaves (defaulting to calling `.inner_visit_with(self)`).
35    fn visit_module(&mut self, module: &'a Module) {
36        module.inner_visit_with(self);
37    }
38    fn visit_module_dialect(&mut self, dialect: &'a ModuleDialect) {
39        dialect.inner_visit_with(self);
40    }
41    fn visit_module_debug_info(&mut self, debug_info: &'a ModuleDebugInfo) {
42        debug_info.inner_visit_with(self);
43    }
44    fn visit_attr_set_def(&mut self, attrs_def: &'a AttrSetDef) {
45        attrs_def.inner_visit_with(self);
46    }
47    fn visit_attr(&mut self, attr: &'a Attr) {
48        attr.inner_visit_with(self);
49    }
50    fn visit_type_def(&mut self, ty_def: &'a TypeDef) {
51        ty_def.inner_visit_with(self);
52    }
53    fn visit_const_def(&mut self, ct_def: &'a ConstDef) {
54        ct_def.inner_visit_with(self);
55    }
56    fn visit_global_var_decl(&mut self, gv_decl: &'a GlobalVarDecl) {
57        gv_decl.inner_visit_with(self);
58    }
59    fn visit_func_decl(&mut self, func_decl: &'a FuncDecl) {
60        func_decl.inner_visit_with(self);
61    }
62    fn visit_control_region_def(&mut self, func_at_control_region: FuncAt<'a, ControlRegion>) {
63        func_at_control_region.inner_visit_with(self);
64    }
65    fn visit_control_node_def(&mut self, func_at_control_node: FuncAt<'a, ControlNode>) {
66        func_at_control_node.inner_visit_with(self);
67    }
68    fn visit_data_inst_def(&mut self, data_inst_def: &'a DataInstDef) {
69        data_inst_def.inner_visit_with(self);
70    }
71    fn visit_data_inst_form_def(&mut self, data_inst_form_def: &'a DataInstFormDef) {
72        data_inst_form_def.inner_visit_with(self);
73    }
74    fn visit_value_use(&mut self, v: &'a Value) {
75        v.inner_visit_with(self);
76    }
77}
78
79/// Trait implemented on "visitable" types (shallowly visitable, at least).
80///
81/// That is, an `impl Visit for X` will call the relevant [`Visitor`] method for
82/// `X`, typically named `Visitor::visit_X` or `Visitor::visit_X_use`.
83//
84// FIXME(eddyb) use this more (e.g. in implementing `InnerVisit`).
85pub trait Visit {
86    fn visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>);
87}
88
89macro_rules! impl_visit {
90    (
91        by_val { $($by_val_method:ident($by_val_ty:ty)),* $(,)? }
92        by_ref { $($by_ref_method:ident($by_ref_ty:ty)),* $(,)? }
93        forward_to_inner_visit { $($forward_to_inner_visit_ty:ty),* $(,)? }
94    ) => {
95        $(impl Visit for $by_val_ty {
96            fn visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
97                visitor.$by_val_method(*self);
98            }
99        })*
100        $(impl Visit for $by_ref_ty {
101            fn visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
102                visitor.$by_ref_method(self);
103            }
104        })*
105        $(impl Visit for $forward_to_inner_visit_ty {
106            fn visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
107                self.inner_visit_with(visitor);
108            }
109        })*
110    };
111}
112
113impl_visit! {
114    by_val {
115        visit_attr_set_use(AttrSet),
116        visit_type_use(Type),
117        visit_const_use(Const),
118        visit_global_var_use(GlobalVar),
119        visit_func_use(Func),
120    }
121    by_ref {
122        visit_spv_dialect(spv::Dialect),
123        visit_spv_module_debug_info(spv::ModuleDebugInfo),
124        visit_import(Import),
125        visit_module(Module),
126        visit_module_dialect(ModuleDialect),
127        visit_module_debug_info(ModuleDebugInfo),
128        visit_attr_set_def(AttrSetDef),
129        visit_attr(Attr),
130        visit_type_def(TypeDef),
131        visit_const_def(ConstDef),
132        visit_global_var_decl(GlobalVarDecl),
133        visit_func_decl(FuncDecl),
134        visit_data_inst_def(DataInstDef),
135        visit_value_use(Value),
136    }
137    forward_to_inner_visit {
138        // NOTE(eddyb) the interpolated parts of `Attr::Diagnostics` aren't visited
139        // by default (as they're "inert data"), this is only for `print`'s usage.
140        Vec<DiagMsgPart>,
141    }
142}
143
144/// Trait implemented on "deeply visitable" types, to further "explore" a type
145/// by visiting its "interior" (i.e. variants and/or fields).
146///
147/// That is, an `impl InnerVisit for X` will call the relevant [`Visitor`] method
148/// for each `X` field, effectively performing a single level of a deep visit.
149/// Also, if `Visitor::visit_X` exists for a given `X`, its default should be to
150/// call `X::inner_visit_with` (i.e. so that visiting is mostly-deep by default).
151pub trait InnerVisit {
152    // FIXME(eddyb) the naming here isn't great, can it be improved?
153    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>);
154}
155
156/// Dynamic dispatch version of [`InnerVisit`].
157///
158/// `dyn DynInnerVisit<'a, V>` is possible, unlike `dyn InnerVisit`, because of
159/// the `trait`-level type parameter `V`, which replaces the method parameter.
160pub trait DynInnerVisit<'a, V> {
161    fn dyn_inner_visit_with(&'a self, visitor: &mut V);
162}
163
164impl<'a, T: InnerVisit, V: Visitor<'a>> DynInnerVisit<'a, V> for T {
165    fn dyn_inner_visit_with(&'a self, visitor: &mut V) {
166        self.inner_visit_with(visitor);
167    }
168}
169
170// FIXME(eddyb) should the impls be here, or next to definitions? (maybe derived?)
171impl InnerVisit for Module {
172    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
173        // FIXME(eddyb) this can't be exhaustive because of the private `cx` field.
174        let Self { dialect, debug_info, global_vars: _, funcs: _, exports, .. } = self;
175
176        visitor.visit_module_dialect(dialect);
177        visitor.visit_module_debug_info(debug_info);
178        for (export_key, exportee) in exports {
179            export_key.inner_visit_with(visitor);
180            exportee.inner_visit_with(visitor);
181        }
182    }
183}
184
185impl InnerVisit for ModuleDialect {
186    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
187        match self {
188            Self::Spv(dialect) => visitor.visit_spv_dialect(dialect),
189        }
190    }
191}
192
193impl InnerVisit for ModuleDebugInfo {
194    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
195        match self {
196            Self::Spv(debug_info) => {
197                visitor.visit_spv_module_debug_info(debug_info);
198            }
199        }
200    }
201}
202
203impl InnerVisit for ExportKey {
204    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
205        match self {
206            Self::LinkName(_) => {}
207
208            Self::SpvEntryPoint { imms: _, interface_global_vars } => {
209                for &gv in interface_global_vars {
210                    visitor.visit_global_var_use(gv);
211                }
212            }
213        }
214    }
215}
216
217impl InnerVisit for Exportee {
218    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
219        match *self {
220            Self::GlobalVar(gv) => visitor.visit_global_var_use(gv),
221            Self::Func(func) => visitor.visit_func_use(func),
222        }
223    }
224}
225
226impl InnerVisit for AttrSetDef {
227    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
228        let Self { attrs } = self;
229
230        for attr in attrs {
231            visitor.visit_attr(attr);
232        }
233    }
234}
235
236impl InnerVisit for Attr {
237    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
238        match self {
239            Attr::Diagnostics(_)
240            | Attr::SpvAnnotation(_)
241            | Attr::SpvDebugLine { .. }
242            | Attr::SpvBitflagsOperand(_) => {}
243
244            Attr::QPtr(attr) => match attr {
245                QPtrAttr::ToSpvPtrInput { input_idx: _, pointee }
246                | QPtrAttr::FromSpvPtrOutput { addr_space: _, pointee } => {
247                    visitor.visit_type_use(pointee.0);
248                }
249
250                QPtrAttr::Usage(usage) => usage.0.inner_visit_with(visitor),
251            },
252        }
253    }
254}
255
256// NOTE(eddyb) the interpolated parts of `Attr::Diagnostics` aren't visited
257// by default (as they're "inert data"), this is only for `print`'s usage.
258impl InnerVisit for Vec<DiagMsgPart> {
259    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
260        for part in self {
261            match part {
262                DiagMsgPart::Plain(_) => {}
263                &DiagMsgPart::Attrs(attrs) => visitor.visit_attr_set_use(attrs),
264                &DiagMsgPart::Type(ty) => visitor.visit_type_use(ty),
265                &DiagMsgPart::Const(ct) => visitor.visit_const_use(ct),
266                DiagMsgPart::QPtrUsage(usage) => usage.inner_visit_with(visitor),
267            }
268        }
269    }
270}
271
272impl InnerVisit for QPtrUsage {
273    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
274        match self {
275            &QPtrUsage::Handles(qptr::shapes::Handle::Opaque(ty)) => {
276                visitor.visit_type_use(ty);
277            }
278            QPtrUsage::Handles(qptr::shapes::Handle::Buffer(_, data_usage)) => {
279                data_usage.inner_visit_with(visitor);
280            }
281            QPtrUsage::Memory(usage) => usage.inner_visit_with(visitor),
282        }
283    }
284}
285
286impl InnerVisit for QPtrMemUsage {
287    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
288        let Self { max_size: _, kind } = self;
289        kind.inner_visit_with(visitor);
290    }
291}
292
293impl InnerVisit for QPtrMemUsageKind {
294    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
295        match self {
296            Self::Unused => {}
297            &Self::StrictlyTyped(ty) | &Self::DirectAccess(ty) => {
298                visitor.visit_type_use(ty);
299            }
300            Self::OffsetBase(entries) => {
301                for sub_usage in entries.values() {
302                    sub_usage.inner_visit_with(visitor);
303                }
304            }
305            Self::DynOffsetBase { element, stride: _ } => {
306                element.inner_visit_with(visitor);
307            }
308        }
309    }
310}
311
312impl InnerVisit for TypeDef {
313    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
314        let Self { attrs, kind } = self;
315
316        visitor.visit_attr_set_use(*attrs);
317        match kind {
318            TypeKind::QPtr | TypeKind::SpvStringLiteralForExtInst => {}
319
320            TypeKind::SpvInst { spv_inst: _, type_and_const_inputs } => {
321                for &ty_or_ct in type_and_const_inputs {
322                    match ty_or_ct {
323                        TypeOrConst::Type(ty) => visitor.visit_type_use(ty),
324                        TypeOrConst::Const(ct) => visitor.visit_const_use(ct),
325                    }
326                }
327            }
328        }
329    }
330}
331
332impl InnerVisit for ConstDef {
333    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
334        let Self { attrs, ty, kind } = self;
335
336        visitor.visit_attr_set_use(*attrs);
337        visitor.visit_type_use(*ty);
338        match kind {
339            &ConstKind::PtrToGlobalVar(gv) => visitor.visit_global_var_use(gv),
340            ConstKind::SpvInst { spv_inst_and_const_inputs } => {
341                let (_spv_inst, const_inputs) = &**spv_inst_and_const_inputs;
342                for &ct in const_inputs {
343                    visitor.visit_const_use(ct);
344                }
345            }
346            ConstKind::SpvStringLiteralForExtInst(_) => {}
347        }
348    }
349}
350
351impl<D: InnerVisit> InnerVisit for DeclDef<D> {
352    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
353        match self {
354            Self::Imported(import) => visitor.visit_import(import),
355            Self::Present(def) => def.inner_visit_with(visitor),
356        }
357    }
358}
359
360impl InnerVisit for GlobalVarDecl {
361    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
362        let Self { attrs, type_of_ptr_to, shape, addr_space, def } = self;
363
364        visitor.visit_attr_set_use(*attrs);
365        visitor.visit_type_use(*type_of_ptr_to);
366        if let Some(shape) = shape {
367            match shape {
368                qptr::shapes::GlobalVarShape::TypedInterface(ty) => visitor.visit_type_use(*ty),
369                qptr::shapes::GlobalVarShape::Handles { .. }
370                | qptr::shapes::GlobalVarShape::UntypedData(_) => {}
371            }
372        }
373        match addr_space {
374            AddrSpace::Handles | AddrSpace::SpvStorageClass(_) => {}
375        }
376        def.inner_visit_with(visitor);
377    }
378}
379
380impl InnerVisit for GlobalVarDefBody {
381    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
382        let Self { initializer } = self;
383
384        if let Some(initializer) = *initializer {
385            visitor.visit_const_use(initializer);
386        }
387    }
388}
389
390impl InnerVisit for FuncDecl {
391    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
392        let Self { attrs, ret_type, params, def } = self;
393
394        visitor.visit_attr_set_use(*attrs);
395        visitor.visit_type_use(*ret_type);
396        for param in params {
397            param.inner_visit_with(visitor);
398        }
399        def.inner_visit_with(visitor);
400    }
401}
402
403impl InnerVisit for FuncParam {
404    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
405        let Self { attrs, ty } = *self;
406
407        visitor.visit_attr_set_use(attrs);
408        visitor.visit_type_use(ty);
409    }
410}
411
412impl InnerVisit for FuncDefBody {
413    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
414        match &self.unstructured_cfg {
415            None => visitor.visit_control_region_def(self.at_body()),
416            Some(cfg) => {
417                for region in cfg.rev_post_order(self) {
418                    visitor.visit_control_region_def(self.at(region));
419
420                    if let Some(control_inst) = cfg.control_inst_on_exit_from.get(region) {
421                        control_inst.inner_visit_with(visitor);
422                    }
423                }
424            }
425        }
426    }
427}
428
429// FIXME(eddyb) this can't implement `InnerVisit` because of the `&'a self`
430// requirement, whereas this has `'a` in `self: FuncAt<'a, ControlRegion>`.
431impl<'a> FuncAt<'a, ControlRegion> {
432    pub fn inner_visit_with(self, visitor: &mut impl Visitor<'a>) {
433        let ControlRegionDef { inputs, children, outputs } = self.def();
434
435        for input in inputs {
436            input.inner_visit_with(visitor);
437        }
438        self.at(*children).into_iter().inner_visit_with(visitor);
439        for v in outputs {
440            visitor.visit_value_use(v);
441        }
442    }
443}
444
445impl InnerVisit for ControlRegionInputDecl {
446    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
447        let Self { attrs, ty } = *self;
448
449        visitor.visit_attr_set_use(attrs);
450        visitor.visit_type_use(ty);
451    }
452}
453
454// FIXME(eddyb) this can't implement `InnerVisit` because of the `&'a self`
455// requirement, whereas this has `'a` in `self: FuncAt<'a, ...>`.
456impl<'a> FuncAt<'a, EntityListIter<ControlNode>> {
457    pub fn inner_visit_with(self, visitor: &mut impl Visitor<'a>) {
458        for func_at_control_node in self {
459            visitor.visit_control_node_def(func_at_control_node);
460        }
461    }
462}
463
464// FIXME(eddyb) this can't implement `InnerVisit` because of the `&'a self`
465// requirement, whereas this has `'a` in `self: FuncAt<'a, ControlNode>`.
466impl<'a> FuncAt<'a, ControlNode> {
467    pub fn inner_visit_with(self, visitor: &mut impl Visitor<'a>) {
468        let ControlNodeDef { kind, outputs } = self.def();
469
470        match kind {
471            ControlNodeKind::Block { insts } => {
472                for func_at_inst in self.at(*insts) {
473                    visitor.visit_data_inst_def(func_at_inst.def());
474                }
475            }
476            ControlNodeKind::Select {
477                kind: SelectionKind::BoolCond | SelectionKind::SpvInst(_),
478                scrutinee,
479                cases,
480            } => {
481                visitor.visit_value_use(scrutinee);
482                for &case in cases {
483                    visitor.visit_control_region_def(self.at(case));
484                }
485            }
486            ControlNodeKind::Loop { initial_inputs, body, repeat_condition } => {
487                for v in initial_inputs {
488                    visitor.visit_value_use(v);
489                }
490                visitor.visit_control_region_def(self.at(*body));
491                visitor.visit_value_use(repeat_condition);
492            }
493            ControlNodeKind::ExitInvocation {
494                kind: cfg::ExitInvocationKind::SpvInst(_),
495                inputs,
496            } => {
497                for v in inputs {
498                    visitor.visit_value_use(v);
499                }
500            }
501        }
502        for output in outputs {
503            output.inner_visit_with(visitor);
504        }
505    }
506}
507
508impl InnerVisit for ControlNodeOutputDecl {
509    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
510        let Self { attrs, ty } = *self;
511
512        visitor.visit_attr_set_use(attrs);
513        visitor.visit_type_use(ty);
514    }
515}
516
517impl InnerVisit for DataInstDef {
518    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
519        let Self { attrs, form, inputs } = self;
520
521        visitor.visit_attr_set_use(*attrs);
522        visitor.visit_data_inst_form_use(*form);
523        for v in inputs {
524            visitor.visit_value_use(v);
525        }
526    }
527}
528
529impl InnerVisit for DataInstFormDef {
530    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
531        let Self { kind, output_type } = self;
532
533        match kind {
534            &DataInstKind::FuncCall(func) => visitor.visit_func_use(func),
535            DataInstKind::QPtr(op) => match *op {
536                QPtrOp::FuncLocalVar(_)
537                | QPtrOp::HandleArrayIndex
538                | QPtrOp::BufferData
539                | QPtrOp::BufferDynLen { .. }
540                | QPtrOp::Offset(_)
541                | QPtrOp::DynOffset { .. }
542                | QPtrOp::Load
543                | QPtrOp::Store => {}
544            },
545            DataInstKind::SpvInst(_) | DataInstKind::SpvExtInst { .. } => {}
546        }
547        if let Some(ty) = *output_type {
548            visitor.visit_type_use(ty);
549        }
550    }
551}
552
553impl InnerVisit for cfg::ControlInst {
554    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
555        let Self { attrs, kind, inputs, targets: _, target_inputs } = self;
556
557        visitor.visit_attr_set_use(*attrs);
558        match kind {
559            cfg::ControlInstKind::Unreachable
560            | cfg::ControlInstKind::Return
561            | cfg::ControlInstKind::ExitInvocation(cfg::ExitInvocationKind::SpvInst(_))
562            | cfg::ControlInstKind::Branch
563            | cfg::ControlInstKind::SelectBranch(
564                SelectionKind::BoolCond | SelectionKind::SpvInst(_),
565            ) => {}
566        }
567        for v in inputs {
568            visitor.visit_value_use(v);
569        }
570        for inputs in target_inputs.values() {
571            for v in inputs {
572                visitor.visit_value_use(v);
573            }
574        }
575    }
576}
577
578impl InnerVisit for Value {
579    fn inner_visit_with<'a>(&'a self, visitor: &mut impl Visitor<'a>) {
580        match *self {
581            Self::Const(ct) => visitor.visit_const_use(ct),
582            Self::ControlRegionInput { region: _, input_idx: _ }
583            | Self::ControlNodeOutput { control_node: _, output_idx: _ }
584            | Self::DataInstOutput(_) => {}
585        }
586    }
587}
spirt/visit.rs

spirt/
visit.rs
