spirt/
transform.rs

1//! Mutable IR traversal.
2
3use crate::func_at::FuncAtMut;
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, DataInst, DataInstDef, DataInstForm, DataInstFormDef, DataInstKind,
9    DeclDef, EntityListIter, ExportKey, Exportee, Func, FuncDecl, FuncDefBody, FuncParam,
10    GlobalVar, GlobalVarDecl, GlobalVarDefBody, Import, Module, ModuleDebugInfo, ModuleDialect,
11    OrdAssertEq, SelectionKind, Type, TypeDef, TypeKind, TypeOrConst, Value, cfg, spv,
12};
13use std::cmp::Ordering;
14use std::rc::Rc;
15use std::slice;
16
17/// The result of a transformation (which is not in-place).
18#[must_use]
19#[derive(Copy, Clone)]
20pub enum Transformed<T> {
21    /// The original `T` value remains as it was, at no cost.
22    Unchanged,
23
24    /// Some part of the original `T` value was transformed, and a new `T` value
25    /// had to be constructed. This change will propagate in any "outer" value.
26    Changed(T),
27}
28
29impl<T> Transformed<T> {
30    pub fn map<U>(self, f: impl FnOnce(T) -> U) -> Transformed<U> {
31        match self {
32            Transformed::Unchanged => Transformed::Unchanged,
33            Transformed::Changed(new) => Transformed::Changed(f(new)),
34        }
35    }
36
37    pub fn apply_to(self, dest: &mut T) {
38        match self {
39            Transformed::Unchanged => {}
40            Transformed::Changed(new) => *dest = new,
41        }
42    }
43}
44
45// HACK(eddyb) the `Self` type here doesn't matter, it only exists so that we
46// can call e.g. `Transformed::foo(...)` methods.
47impl Transformed<()> {
48    /// Map every element of an iterator through `transform_elem` and return
49    /// [`Transformed::Changed(new_iter)`] if any `transform_elem` call returned
50    /// [`Transformed::Changed`], with `new_iter` containing a combination of the
51    /// changed elements, and clones of the unchanged elements.
52    pub fn map_iter<'a, T: 'a + Clone>(
53        iter: impl Iterator<Item = &'a T> + Clone + 'a,
54        mut transform_elem: impl FnMut(&'a T) -> Transformed<T> + 'a,
55    ) -> Transformed<impl Iterator<Item = T> + 'a> {
56        for (i, elem_i) in iter.clone().enumerate() {
57            if let Transformed::Changed(new_elem_i) = transform_elem(elem_i) {
58                let mut new_elem_i = Some(new_elem_i);
59                return Transformed::Changed(iter.enumerate().map(move |(j, elem_j)| {
60                    match j.cmp(&i) {
61                        // Earlier elements, for which `transform_elem` was called
62                        // already, and had returned `Unchanged`.
63                        Ordering::Less => elem_j.clone(),
64
65                        // The first element for which `transform_elem` returned
66                        // `Changed`, resulting in the `Changed` iterator.
67                        Ordering::Equal => new_elem_i.take().unwrap(),
68
69                        // Later elements, for which only now `transform_elem`
70                        // gets called, and may be `Unchanged` or `Changed`.
71                        Ordering::Greater => match transform_elem(elem_j) {
72                            Transformed::Unchanged => elem_j.clone(),
73                            Transformed::Changed(new_elem_j) => new_elem_j,
74                        },
75                    }
76                }));
77            }
78        }
79        Transformed::Unchanged
80    }
81}
82
83/// Helper type for [`transform!`] - not public as it's easy to misuse.
84enum TransformedWithOriginal<'a, T> {
85    Original(&'a T),
86    Changed(T),
87}
88
89impl<T> Transformed<T> {
90    fn with_original(self, original: &T) -> TransformedWithOriginal<'_, T> {
91        match self {
92            Transformed::Unchanged => TransformedWithOriginal::Original(original),
93            Transformed::Changed(new) => TransformedWithOriginal::Changed(new),
94        }
95    }
96}
97
98impl<T: Clone> TransformedWithOriginal<'_, T> {
99    fn is_changed(&self) -> bool {
100        matches!(self, TransformedWithOriginal::Changed(_))
101    }
102    fn changed_or_original_cloned(self) -> T {
103        match self {
104            TransformedWithOriginal::Original(original) => original.clone(),
105            TransformedWithOriginal::Changed(new) => new,
106        }
107    }
108}
109
110// HACK(eddyb) `transform!` needs auto-ref-like behavior for inputs.
111trait AutoRef {
112    fn auto_ref(&self) -> &Self {
113        self
114    }
115}
116
117impl<T> AutoRef for T {}
118
119/// Helper macro to create a combined [`Transformed`] out of several variables,
120/// each with their own transformation, where any [`Transformed::Changed`] input
121/// will result in a [`Transformed::Changed`] output, using a combination of the
122/// changed inputs, and clones of the unchanged inputs.
123macro_rules! transform {
124    ({ $($input:ident -> $input_transformed:expr),+ $(,)? } => $output:expr) => {{
125        let ($($input,)+) = ($($input_transformed.with_original($input.auto_ref()),)+);
126        if $($input.is_changed())||+ {
127            let ($($input,)*) = ($($input.changed_or_original_cloned(),)+);
128            Transformed::Changed($output)
129        } else {
130            Transformed::Unchanged
131        }
132    }};
133}
134
135// FIXME(eddyb) `Sized` bound shouldn't be needed but removing it requires
136// writing `impl Transformer + ?Sized` in `fn inner_transform_with` signatures.
137pub trait Transformer: Sized {
138    // Context-interned leaves (noop default behavior).
139    fn transform_attr_set_use(&mut self, _attrs: AttrSet) -> Transformed<AttrSet> {
140        Transformed::Unchanged
141    }
142    fn transform_type_use(&mut self, _ty: Type) -> Transformed<Type> {
143        Transformed::Unchanged
144    }
145    fn transform_const_use(&mut self, _ct: Const) -> Transformed<Const> {
146        Transformed::Unchanged
147    }
148    fn transform_data_inst_form_use(
149        &mut self,
150        _data_inst_form: DataInstForm,
151    ) -> Transformed<DataInstForm> {
152        Transformed::Unchanged
153    }
154
155    // Module-stored entity leaves (noop default behavior).
156    fn transform_global_var_use(&mut self, _gv: GlobalVar) -> Transformed<GlobalVar> {
157        Transformed::Unchanged
158    }
159    fn transform_func_use(&mut self, _func: Func) -> Transformed<Func> {
160        Transformed::Unchanged
161    }
162
163    // Leaves transformed in-place (noop default behavior).
164    fn in_place_transform_spv_dialect(&mut self, _dialect: &mut spv::Dialect) {}
165    fn in_place_transform_spv_module_debug_info(&mut self, _debug_info: &mut spv::ModuleDebugInfo) {
166    }
167
168    // Non-leaves (defaulting to calling `.inner_transform_with(self)`).
169    fn transform_attr_set_def(&mut self, attrs_def: &AttrSetDef) -> Transformed<AttrSetDef> {
170        attrs_def.inner_transform_with(self)
171    }
172    fn transform_attr(&mut self, attr: &Attr) -> Transformed<Attr> {
173        attr.inner_transform_with(self)
174    }
175    fn transform_type_def(&mut self, ty_def: &TypeDef) -> Transformed<TypeDef> {
176        ty_def.inner_transform_with(self)
177    }
178    fn transform_const_def(&mut self, ct_def: &ConstDef) -> Transformed<ConstDef> {
179        ct_def.inner_transform_with(self)
180    }
181    fn transform_data_inst_form_def(
182        &mut self,
183        data_inst_form_def: &DataInstFormDef,
184    ) -> Transformed<DataInstFormDef> {
185        data_inst_form_def.inner_transform_with(self)
186    }
187    fn transform_value_use(&mut self, v: &Value) -> Transformed<Value> {
188        v.inner_transform_with(self)
189    }
190
191    // Non-leaves transformed in-place (defaulting to calling
192    // `.inner_in_place_transform_with(self)`).
193    fn in_place_transform_module(&mut self, module: &mut Module) {
194        module.inner_in_place_transform_with(self);
195    }
196    fn in_place_transform_module_dialect(&mut self, dialect: &mut ModuleDialect) {
197        dialect.inner_in_place_transform_with(self);
198    }
199    fn in_place_transform_module_debug_info(&mut self, debug_info: &mut ModuleDebugInfo) {
200        debug_info.inner_in_place_transform_with(self);
201    }
202    fn in_place_transform_global_var_decl(&mut self, gv_decl: &mut GlobalVarDecl) {
203        gv_decl.inner_in_place_transform_with(self);
204    }
205    fn in_place_transform_func_decl(&mut self, func_decl: &mut FuncDecl) {
206        func_decl.inner_in_place_transform_with(self);
207    }
208    fn in_place_transform_control_node_def(
209        &mut self,
210        mut func_at_control_node: FuncAtMut<'_, ControlNode>,
211    ) {
212        func_at_control_node.inner_in_place_transform_with(self);
213    }
214    fn in_place_transform_data_inst_def(&mut self, mut func_at_data_inst: FuncAtMut<'_, DataInst>) {
215        func_at_data_inst.inner_in_place_transform_with(self);
216    }
217}
218
219/// Trait implemented on "transformable" types, to further "elaborate" a type by
220/// transforming its "interior" (i.e. variants and/or fields).
221///
222/// That is, an `impl InnerTransform for X` will call the relevant [`Transformer`]
223/// method for each `X` field, effectively performing a single level of a deep
224/// transform.
225/// Also, if `Transformer::transform_X` exists for a given `X`, its default should
226/// be to call `X::inner_transform_with` (i.e. so that transforming is mostly-deep
227/// by default).
228pub trait InnerTransform: Sized {
229    // FIXME(eddyb) the naming here isn't great, can it be improved?
230    // FIXME(eddyb) should this be `self -> Self` instead of `&mut self -> ()`?
231    fn inner_transform_with(&self, transformer: &mut impl Transformer) -> Transformed<Self>;
232}
233
234/// Like [`InnerTransform`], but only for the `in_place_transform_X` cases.
235pub trait InnerInPlaceTransform {
236    fn inner_in_place_transform_with(&mut self, transformer: &mut impl Transformer);
237}
238
239// FIXME(eddyb) should the impls be here, or next to definitions? (maybe derived?)
240impl InnerInPlaceTransform for Module {
241    fn inner_in_place_transform_with(&mut self, transformer: &mut impl Transformer) {
242        // FIXME(eddyb) this can't be exhaustive because of the private `cx` field.
243        let Self { dialect, debug_info, global_vars: _, funcs: _, exports, .. } = self;
244
245        transformer.in_place_transform_module_dialect(dialect);
246        transformer.in_place_transform_module_debug_info(debug_info);
247
248        // HACK(eddyb) this takes two passes, once for values and once for keys,
249        // to be able to use in-place mutable iteration for the former, and
250        // `Transformed::map_iter` (i.e. immutable iteration) for the latter.
251        for exportee in exports.values_mut() {
252            exportee.inner_transform_with(transformer).apply_to(exportee);
253        }
254        Transformed::map_iter(exports.keys(), |export_key| {
255            export_key.inner_transform_with(transformer)
256        })
257        .map(|new_keys_iter| {
258            // Recombine the new keys with the existing values.
259            new_keys_iter.zip(exports.values().cloned()).collect()
260        })
261        .apply_to(exports);
262    }
263}
264
265impl InnerInPlaceTransform for ModuleDialect {
266    fn inner_in_place_transform_with(&mut self, transformer: &mut impl Transformer) {
267        match self {
268            Self::Spv(dialect) => transformer.in_place_transform_spv_dialect(dialect),
269        }
270    }
271}
272
273impl InnerInPlaceTransform for ModuleDebugInfo {
274    fn inner_in_place_transform_with(&mut self, transformer: &mut impl Transformer) {
275        match self {
276            Self::Spv(debug_info) => {
277                transformer.in_place_transform_spv_module_debug_info(debug_info);
278            }
279        }
280    }
281}
282
283impl InnerTransform for ExportKey {
284    fn inner_transform_with(&self, transformer: &mut impl Transformer) -> Transformed<Self> {
285        match self {
286            Self::LinkName(_) => Transformed::Unchanged,
287
288            Self::SpvEntryPoint { imms, interface_global_vars } => transform!({
289                imms -> Transformed::Unchanged,
290                interface_global_vars -> Transformed::map_iter(
291                    interface_global_vars.iter(),
292                    |&gv| transformer.transform_global_var_use(gv),
293                ).map(|new_iter| new_iter.collect()),
294            } => Self::SpvEntryPoint {
295                imms,
296                interface_global_vars,
297            }),
298        }
299    }
300}
301
302impl InnerTransform for Exportee {
303    fn inner_transform_with(&self, transformer: &mut impl Transformer) -> Transformed<Self> {
304        match *self {
305            Self::GlobalVar(gv) => transform!({
306                gv -> transformer.transform_global_var_use(gv),
307            } => Self::GlobalVar(gv)),
308
309            Self::Func(func) => transform!({
310                func -> transformer.transform_func_use(func),
311            } => Self::Func(func)),
312        }
313    }
314}
315
316impl InnerTransform for AttrSetDef {
317    fn inner_transform_with(&self, transformer: &mut impl Transformer) -> Transformed<Self> {
318        let Self { attrs } = self;
319
320        transform!({
321            attrs -> Transformed::map_iter(
322                attrs.iter(),
323                |attr| transformer.transform_attr(attr),
324            ).map(|new_iter| new_iter.collect()),
325        } => Self {
326            attrs,
327        })
328    }
329}
330
331impl InnerTransform for Attr {
332    fn inner_transform_with(&self, transformer: &mut impl Transformer) -> Transformed<Self> {
333        match self {
334            Attr::Diagnostics(_)
335            | Attr::SpvAnnotation(_)
336            | Attr::SpvDebugLine { .. }
337            | Attr::SpvBitflagsOperand(_) => Transformed::Unchanged,
338
339            Attr::QPtr(attr) => transform!({
340                attr -> match attr {
341                    &QPtrAttr::ToSpvPtrInput { input_idx, pointee } => transform!({
342                        pointee -> transformer.transform_type_use(pointee.0).map(OrdAssertEq),
343                    } => QPtrAttr::ToSpvPtrInput { input_idx, pointee }),
344
345                    &QPtrAttr::FromSpvPtrOutput {
346                        addr_space,
347                        pointee,
348                    } => transform!({
349                        pointee -> transformer.transform_type_use(pointee.0).map(OrdAssertEq),
350                    } => QPtrAttr::FromSpvPtrOutput {
351                        addr_space,
352                        pointee,
353                    }),
354
355                    QPtrAttr::Usage(OrdAssertEq(usage)) => transform!({
356                        usage -> match usage {
357                            &QPtrUsage::Handles(qptr::shapes::Handle::Opaque(ty)) => transform!({
358                                ty -> transformer.transform_type_use(ty),
359                            } => QPtrUsage::Handles(qptr::shapes::Handle::Opaque(ty))),
360                            QPtrUsage::Handles(qptr::shapes::Handle::Buffer(addr_space, data_usage)) => transform!({
361                                data_usage -> data_usage.inner_transform_with(transformer),
362                            } => QPtrUsage::Handles(qptr::shapes::Handle::Buffer(*addr_space, data_usage))),
363                            QPtrUsage::Memory(usage) => transform!({
364                                usage -> usage.inner_transform_with(transformer),
365                            } => QPtrUsage::Memory(usage)),
366                        }
367                    } => QPtrAttr::Usage(OrdAssertEq(usage))),
368                }
369            } => Attr::QPtr(attr)),
370        }
371    }
372}
373
374// FIXME(eddyb) this should maybe be in a more general spot.
375impl<T: InnerTransform> InnerTransform for Rc<T> {
376    fn inner_transform_with(&self, transformer: &mut impl Transformer) -> Transformed<Self> {
377        (**self).inner_transform_with(transformer).map(Rc::new)
378    }
379}
380
381impl InnerTransform for QPtrMemUsage {
382    fn inner_transform_with(&self, transformer: &mut impl Transformer) -> Transformed<Self> {
383        let Self { max_size, kind } = self;
384
385        transform!({
386            kind -> kind.inner_transform_with(transformer)
387        } => Self {
388            max_size: *max_size,
389            kind,
390        })
391    }
392}
393
394impl InnerTransform for QPtrMemUsageKind {
395    fn inner_transform_with(&self, transformer: &mut impl Transformer) -> Transformed<Self> {
396        match self {
397            Self::Unused => Transformed::Unchanged,
398            &Self::StrictlyTyped(ty) => transform!({
399                ty -> transformer.transform_type_use(ty),
400            } => Self::StrictlyTyped(ty)),
401            &Self::DirectAccess(ty) => transform!({
402                ty -> transformer.transform_type_use(ty),
403            } => Self::DirectAccess(ty)),
404            Self::OffsetBase(entries) => transform!({
405                entries -> Transformed::map_iter(entries.values(), |sub_usage| {
406                    sub_usage.inner_transform_with(transformer)
407                }).map(|new_iter| {
408                    // HACK(eddyb) this is a bit inefficient but `Transformed::map_iter`
409                    // limits us here in how it handles the whole `Clone` thing.
410                    entries.keys().copied().zip(new_iter).collect()
411                }).map(Rc::new)
412            } => Self::OffsetBase(entries)),
413            Self::DynOffsetBase { element, stride } => transform!({
414                element -> element.inner_transform_with(transformer),
415            } => Self::DynOffsetBase { element, stride: *stride }),
416        }
417    }
418}
419
420impl InnerTransform for TypeDef {
421    fn inner_transform_with(&self, transformer: &mut impl Transformer) -> Transformed<Self> {
422        let Self { attrs, kind } = self;
423
424        transform!({
425            attrs -> transformer.transform_attr_set_use(*attrs),
426            kind -> match kind {
427                TypeKind::QPtr | TypeKind::SpvStringLiteralForExtInst => Transformed::Unchanged,
428
429                TypeKind::SpvInst { spv_inst, type_and_const_inputs } => Transformed::map_iter(
430                    type_and_const_inputs.iter(),
431                    |ty_or_ct| match *ty_or_ct {
432                        TypeOrConst::Type(ty) => transform!({
433                            ty -> transformer.transform_type_use(ty),
434                        } => TypeOrConst::Type(ty)),
435
436                        TypeOrConst::Const(ct) => transform!({
437                            ct -> transformer.transform_const_use(ct),
438                        } => TypeOrConst::Const(ct)),
439                    },
440                ).map(|new_iter| TypeKind::SpvInst {
441                    spv_inst: spv_inst.clone(),
442                    type_and_const_inputs: new_iter.collect(),
443                }),
444            },
445        } => Self {
446            attrs,
447            kind,
448        })
449    }
450}
451
452impl InnerTransform for ConstDef {
453    fn inner_transform_with(&self, transformer: &mut impl Transformer) -> Transformed<Self> {
454        let Self { attrs, ty, kind } = self;
455
456        transform!({
457            attrs -> transformer.transform_attr_set_use(*attrs),
458            ty -> transformer.transform_type_use(*ty),
459            kind -> match kind {
460                ConstKind::PtrToGlobalVar(gv) => transform!({
461                    gv -> transformer.transform_global_var_use(*gv),
462                } => ConstKind::PtrToGlobalVar(gv)),
463
464                ConstKind::SpvInst { spv_inst_and_const_inputs } => {
465                    let (spv_inst, const_inputs) = &**spv_inst_and_const_inputs;
466                    Transformed::map_iter(
467                        const_inputs.iter(),
468                        |&ct| transformer.transform_const_use(ct),
469                    ).map(|new_iter| ConstKind::SpvInst {
470                        spv_inst_and_const_inputs: Rc::new((spv_inst.clone(), new_iter.collect())),
471                    })
472                }
473                ConstKind::SpvStringLiteralForExtInst(_) => Transformed::Unchanged
474            },
475        } => Self {
476            attrs,
477            ty,
478            kind,
479        })
480    }
481}
482
483impl<D: InnerInPlaceTransform> InnerInPlaceTransform for DeclDef<D> {
484    fn inner_in_place_transform_with(&mut self, transformer: &mut impl Transformer) {
485        match self {
486            Self::Imported(import) => match import {
487                Import::LinkName(_) => {}
488            },
489            Self::Present(def) => def.inner_in_place_transform_with(transformer),
490        }
491    }
492}
493
494impl InnerInPlaceTransform for GlobalVarDecl {
495    fn inner_in_place_transform_with(&mut self, transformer: &mut impl Transformer) {
496        let Self { attrs, type_of_ptr_to, shape, addr_space, def } = self;
497
498        transformer.transform_attr_set_use(*attrs).apply_to(attrs);
499        transformer.transform_type_use(*type_of_ptr_to).apply_to(type_of_ptr_to);
500        if let Some(shape) = shape {
501            match shape {
502                qptr::shapes::GlobalVarShape::TypedInterface(ty) => {
503                    transformer.transform_type_use(*ty).apply_to(ty);
504                }
505                qptr::shapes::GlobalVarShape::Handles { .. }
506                | qptr::shapes::GlobalVarShape::UntypedData(_) => {}
507            }
508        }
509        match addr_space {
510            AddrSpace::Handles | AddrSpace::SpvStorageClass(_) => {}
511        }
512        def.inner_in_place_transform_with(transformer);
513    }
514}
515
516impl InnerInPlaceTransform for GlobalVarDefBody {
517    fn inner_in_place_transform_with(&mut self, transformer: &mut impl Transformer) {
518        let Self { initializer } = self;
519
520        if let Some(initializer) = initializer {
521            transformer.transform_const_use(*initializer).apply_to(initializer);
522        }
523    }
524}
525
526impl InnerInPlaceTransform for FuncDecl {
527    fn inner_in_place_transform_with(&mut self, transformer: &mut impl Transformer) {
528        let Self { attrs, ret_type, params, def } = self;
529
530        transformer.transform_attr_set_use(*attrs).apply_to(attrs);
531        transformer.transform_type_use(*ret_type).apply_to(ret_type);
532        for param in params {
533            param.inner_transform_with(transformer).apply_to(param);
534        }
535        def.inner_in_place_transform_with(transformer);
536    }
537}
538
539impl InnerTransform for FuncParam {
540    fn inner_transform_with(&self, transformer: &mut impl Transformer) -> Transformed<Self> {
541        let Self { attrs, ty } = self;
542
543        transform!({
544            attrs -> transformer.transform_attr_set_use(*attrs),
545            ty -> transformer.transform_type_use(*ty),
546        } => Self {
547            attrs,
548            ty,
549        })
550    }
551}
552
553impl InnerInPlaceTransform for FuncDefBody {
554    fn inner_in_place_transform_with(&mut self, transformer: &mut impl Transformer) {
555        match &self.unstructured_cfg {
556            None => self.at_mut_body().inner_in_place_transform_with(transformer),
557            Some(cfg) => {
558                // HACK(eddyb) have to compute this before borrowing any `self` fields.
559                let rpo = cfg.rev_post_order(self);
560
561                for region in rpo {
562                    self.at_mut(region).inner_in_place_transform_with(transformer);
563
564                    let cfg = self.unstructured_cfg.as_mut().unwrap();
565                    if let Some(control_inst) = cfg.control_inst_on_exit_from.get_mut(region) {
566                        control_inst.inner_in_place_transform_with(transformer);
567                    }
568                }
569            }
570        }
571    }
572}
573
574impl InnerInPlaceTransform for FuncAtMut<'_, ControlRegion> {
575    fn inner_in_place_transform_with(&mut self, transformer: &mut impl Transformer) {
576        // HACK(eddyb) handle the fields of `ControlRegion` separately, to
577        // allow reborrowing `FuncAtMut` (for recursing into `ControlNode`s).
578        let ControlRegionDef { inputs, children: _, outputs: _ } = self.reborrow().def();
579        for input in inputs {
580            input.inner_transform_with(transformer).apply_to(input);
581        }
582
583        self.reborrow().at_children().into_iter().inner_in_place_transform_with(transformer);
584
585        let ControlRegionDef { inputs: _, children: _, outputs } = self.reborrow().def();
586        for v in outputs {
587            transformer.transform_value_use(v).apply_to(v);
588        }
589    }
590}
591
592impl InnerTransform for ControlRegionInputDecl {
593    fn inner_transform_with(&self, transformer: &mut impl Transformer) -> Transformed<Self> {
594        let Self { attrs, ty } = self;
595
596        transform!({
597            attrs -> transformer.transform_attr_set_use(*attrs),
598            ty -> transformer.transform_type_use(*ty),
599        } => Self {
600            attrs,
601            ty,
602        })
603    }
604}
605
606impl InnerInPlaceTransform for FuncAtMut<'_, EntityListIter<ControlNode>> {
607    fn inner_in_place_transform_with(&mut self, transformer: &mut impl Transformer) {
608        let mut iter = self.reborrow();
609        while let Some(func_at_control_node) = iter.next() {
610            transformer.in_place_transform_control_node_def(func_at_control_node);
611        }
612    }
613}
614
615impl FuncAtMut<'_, ControlNode> {
616    fn child_regions(&mut self) -> &mut [ControlRegion] {
617        match &mut self.reborrow().def().kind {
618            ControlNodeKind::Block { .. } | ControlNodeKind::ExitInvocation { .. } => &mut [][..],
619
620            ControlNodeKind::Select { cases, .. } => cases,
621            ControlNodeKind::Loop { body, .. } => slice::from_mut(body),
622        }
623    }
624}
625
626impl InnerInPlaceTransform for FuncAtMut<'_, ControlNode> {
627    fn inner_in_place_transform_with(&mut self, transformer: &mut impl Transformer) {
628        // HACK(eddyb) handle pre-child-regions parts of `kind` separately to
629        // allow reborrowing `FuncAtMut` (for the child region recursion).
630        match &mut self.reborrow().def().kind {
631            &mut ControlNodeKind::Block { insts } => {
632                let mut func_at_inst_iter = self.reborrow().at(insts).into_iter();
633                while let Some(func_at_inst) = func_at_inst_iter.next() {
634                    transformer.in_place_transform_data_inst_def(func_at_inst);
635                }
636            }
637            ControlNodeKind::Select {
638                kind: SelectionKind::BoolCond | SelectionKind::SpvInst(_),
639                scrutinee,
640                cases: _,
641            } => {
642                transformer.transform_value_use(scrutinee).apply_to(scrutinee);
643            }
644            ControlNodeKind::Loop { initial_inputs: inputs, body: _, repeat_condition: _ }
645            | ControlNodeKind::ExitInvocation {
646                kind: cfg::ExitInvocationKind::SpvInst(_),
647                inputs,
648            } => {
649                for v in inputs {
650                    transformer.transform_value_use(v).apply_to(v);
651                }
652            }
653        }
654
655        // FIXME(eddyb) represent the list of child regions without having them
656        // in a `Vec` (or `SmallVec`), which requires workarounds like this.
657        for child_region_idx in 0..self.child_regions().len() {
658            let child_region = self.child_regions()[child_region_idx];
659            self.reborrow().at(child_region).inner_in_place_transform_with(transformer);
660        }
661
662        let ControlNodeDef { kind, outputs } = self.reborrow().def();
663
664        match kind {
665            // Fully handled above, before recursing into any child regions.
666            ControlNodeKind::Block { insts: _ }
667            | ControlNodeKind::Select { kind: _, scrutinee: _, cases: _ }
668            | ControlNodeKind::ExitInvocation {
669                kind: cfg::ExitInvocationKind::SpvInst(_),
670                inputs: _,
671            } => {}
672
673            ControlNodeKind::Loop { initial_inputs: _, body: _, repeat_condition } => {
674                transformer.transform_value_use(repeat_condition).apply_to(repeat_condition);
675            }
676        };
677
678        for output in outputs {
679            output.inner_transform_with(transformer).apply_to(output);
680        }
681    }
682}
683
684impl InnerTransform for ControlNodeOutputDecl {
685    fn inner_transform_with(&self, transformer: &mut impl Transformer) -> Transformed<Self> {
686        let Self { attrs, ty } = self;
687
688        transform!({
689            attrs -> transformer.transform_attr_set_use(*attrs),
690            ty -> transformer.transform_type_use(*ty),
691        } => Self {
692            attrs,
693            ty,
694        })
695    }
696}
697
698impl InnerInPlaceTransform for FuncAtMut<'_, DataInst> {
699    fn inner_in_place_transform_with(&mut self, transformer: &mut impl Transformer) {
700        let DataInstDef { attrs, form, inputs } = self.reborrow().def();
701
702        transformer.transform_attr_set_use(*attrs).apply_to(attrs);
703        transformer.transform_data_inst_form_use(*form).apply_to(form);
704        for v in inputs {
705            transformer.transform_value_use(v).apply_to(v);
706        }
707    }
708}
709
710impl InnerTransform for DataInstFormDef {
711    fn inner_transform_with(&self, transformer: &mut impl Transformer) -> Transformed<Self> {
712        let Self { kind, output_type } = self;
713
714        transform!({
715            kind -> match kind {
716                DataInstKind::FuncCall(func) => transformer.transform_func_use(*func).map(DataInstKind::FuncCall),
717                DataInstKind::QPtr(op) => match op {
718                    QPtrOp::FuncLocalVar(_)
719                    | QPtrOp::HandleArrayIndex
720                    | QPtrOp::BufferData
721                    | QPtrOp::BufferDynLen { .. }
722                    | QPtrOp::Offset(_)
723                    | QPtrOp::DynOffset { .. }
724                    | QPtrOp::Load
725                    | QPtrOp::Store => Transformed::Unchanged,
726                },
727                DataInstKind::SpvInst(_) | DataInstKind::SpvExtInst { .. } => Transformed::Unchanged,
728            },
729            // FIXME(eddyb) this should be replaced with an impl of `InnerTransform`
730            // for `Option<T>` or some other helper, to avoid "manual transpose".
731            output_type -> output_type.map(|ty| transformer.transform_type_use(ty))
732                .map_or(Transformed::Unchanged, |t| t.map(Some)),
733        } => Self {
734            kind,
735            output_type,
736        })
737    }
738}
739
740impl InnerInPlaceTransform for cfg::ControlInst {
741    fn inner_in_place_transform_with(&mut self, transformer: &mut impl Transformer) {
742        let Self { attrs, kind, inputs, targets: _, target_inputs } = self;
743
744        transformer.transform_attr_set_use(*attrs).apply_to(attrs);
745        match kind {
746            cfg::ControlInstKind::Unreachable
747            | cfg::ControlInstKind::Return
748            | cfg::ControlInstKind::ExitInvocation(cfg::ExitInvocationKind::SpvInst(_))
749            | cfg::ControlInstKind::Branch
750            | cfg::ControlInstKind::SelectBranch(
751                SelectionKind::BoolCond | SelectionKind::SpvInst(_),
752            ) => {}
753        }
754        for v in inputs {
755            transformer.transform_value_use(v).apply_to(v);
756        }
757        for inputs in target_inputs.values_mut() {
758            for v in inputs {
759                transformer.transform_value_use(v).apply_to(v);
760            }
761        }
762    }
763}
764
765impl InnerTransform for Value {
766    fn inner_transform_with(&self, transformer: &mut impl Transformer) -> Transformed<Self> {
767        match self {
768            Self::Const(ct) => transform!({
769                ct -> transformer.transform_const_use(*ct),
770            } => Self::Const(ct)),
771
772            Self::ControlRegionInput { region: _, input_idx: _ }
773            | Self::ControlNodeOutput { control_node: _, output_idx: _ }
774            | Self::DataInstOutput(_) => Transformed::Unchanged,
775        }
776    }
777}
spirt/transform.rs

spirt/
transform.rs
