wesl/pass/
visit.rs

1use wesl_macros::{query, query_mut};
2
3use wgsl_parse::syntax::*;
4
5pub trait Visit<T> {
6    /// Visit each child node of type `T` in the subtree of `Self`.
7    ///
8    /// Implementations of Visit do not recurse past `T`, meaning that if you really want
9    /// to visit all children of type T you would have to call `<T as Visit<T>>::visit` on
10    /// each visited `T`. Alternatively, use [`Self::visit_rec`] which solves this
11    /// exact problem.
12    fn visit<'a>(&'a self) -> impl Iterator<Item = &'a T>
13    where
14        T: 'a;
15
16    /// Mutable version of [`Self::visit`].
17    fn visit_mut<'a>(&'a mut self) -> impl Iterator<Item = &'a mut T>
18    where
19        T: 'a;
20
21    /// Visit each child node of type `T` in the subtree of `Self`, recursively in depth-first order.
22    ///
23    /// Due to Rust's aliasing model, you can't iterate mutably on a node and its
24    /// children. This function allows visiting recursively by passing a closure instead.
25    #[allow(unused)]
26    fn visit_rec<'a, F>(&'a self, f: &mut F)
27    where
28        T: Visit<T> + 'a,
29        F: FnMut(&T),
30    {
31        Visit::<T>::visit(self).for_each(|x| {
32            f(x);
33            x.visit_rec(f);
34        });
35    }
36
37    /// Mutable version of [`Self::visit_rec`].
38    #[allow(unused)]
39    fn visit_rec_mut<'a, F>(&'a mut self, f: &mut F)
40    where
41        T: Visit<T> + 'a,
42        F: FnMut(&mut T),
43    {
44        Visit::<T>::visit_mut(self).for_each(|x| {
45            f(x);
46            x.visit_rec_mut(f);
47        });
48    }
49}
50
51macro_rules! impl_visit {
52    ($type:ty => $visited:ty, $expr:tt) => {
53        impl Visit<$visited> for $type {
54            fn visit<'a>(&'a self) -> impl Iterator<Item = &'a $visited>
55            where
56                $visited: 'a,
57            {
58                #[allow(unused)]
59                fn visit<'a, T: Visit<U>, U: 'a>(expr: &'a T) -> impl Iterator<Item = &'a U> {
60                    Visit::<U>::visit(expr)
61                }
62
63                #[allow(unused)]
64                fn recurse(expr: &$type) -> impl Iterator<Item = &$visited> {
65                    Visit::<$visited>::visit(expr)
66                }
67
68                let root: &$type = self;
69                query!(root.$expr)
70            }
71            fn visit_mut<'a>(&'a mut self) -> impl Iterator<Item = &'a mut $visited>
72            where
73                $visited: 'a,
74            {
75                #[allow(unused)]
76                fn visit<'a, T: Visit<U>, U: 'a>(
77                    expr: &'a mut T,
78                ) -> impl Iterator<Item = &'a mut U> {
79                    Visit::<U>::visit_mut(expr)
80                }
81
82                #[allow(unused)]
83                fn recurse(expr: &mut $type) -> impl Iterator<Item = &mut $visited> {
84                    Visit::<$visited>::visit_mut(expr)
85                }
86
87                let root: &mut $type = self;
88                query_mut!(root.$expr)
89            }
90        }
91    };
92}
93
94impl_visit! { Expression => ExpressionNode,
95    {
96        Expression::Parenthesized.expression.(x => recurse(x)),
97        Expression::NamedComponent.base.(x => recurse(x)),
98        Expression::Indexing.{
99            base.(x => recurse(x)),
100            index.(x => recurse(x)),
101        },
102        Expression::Unary.operand.(x => recurse(x)),
103        Expression::Binary.{
104            left.(x => recurse(x)),
105            right.(x => recurse(x)),
106        },
107        Expression::FunctionCall.arguments.[].(x => recurse(x)),
108    }
109}
110
111impl_visit! { Expression => TypeExpression,
112    {
113        Expression::Parenthesized.expression.(x => recurse(x)),
114        Expression::NamedComponent.base.(x => recurse(x)),
115        Expression::Indexing.{ base.(x => recurse(x)), index.(x => recurse(x)) },
116        Expression::Unary.operand.(x => recurse(x)),
117        Expression::Binary.{ left.(x => recurse(x)), right.(x => recurse(x)) },
118        Expression::FunctionCall.{
119            ty,
120            arguments.[].(x => recurse(x))
121        },
122        Expression::TypeOrIdentifier,
123    }
124}
125
126impl_visit! { TypeExpression => TypeExpression,
127    {
128        template_args.[].[].expression.(x => visit::<Expression, TypeExpression>(x))
129    }
130}
131
132impl_visit! { Statement => Attributes,
133    {
134        Statement::Compound.{ attributes, statements.[].(x => recurse(x)) },
135        Statement::If.{
136            attributes,
137            else_if_clauses.[].body.statements.[].(x => recurse(x)),
138            else_clause.[].body.statements.[].(x => recurse(x)),
139        },
140        Statement::Switch.{
141            attributes,
142            clauses.[].{
143                attributes,
144                body.statements.[].(x => recurse(x))
145            },
146        },
147        Statement::Loop.{
148            attributes,
149            body.statements.[].(x => recurse(x)),
150            continuing.[].{
151                attributes,
152                body.statements.[].(x => recurse(x)),
153                break_if.[].attributes
154            },
155        },
156        Statement::For.{
157            attributes,
158            body.statements.[].(x => recurse(x)),
159        },
160        Statement::While.{
161            attributes,
162            body.statements.[].(x => recurse(x)),
163        },
164        Statement::Break.attributes,
165        Statement::Continue.attributes,
166        Statement::Return.attributes,
167        Statement::Discard.attributes,
168        Statement::FunctionCall.attributes,
169        Statement::ConstAssert.attributes,
170        Statement::Declaration.attributes,
171    }
172}
173
174impl_visit! { Statement => TypeExpression,
175    {
176        Statement::Compound.{
177            attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
178            statements.[].(x => recurse(x)),
179        },
180        Statement::Assignment.{
181            attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
182            lhs.(x => visit::<Expression, TypeExpression>(x)),
183            rhs.(x => visit::<Expression, TypeExpression>(x)),
184        },
185        Statement::Increment.{
186            attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
187            expression.(x => visit::<Expression, TypeExpression>(x)),
188        },
189        Statement::Decrement.{
190            attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
191            expression.(x => visit::<Expression, TypeExpression>(x)),
192        },
193        Statement::If.{
194            attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
195            if_clause.{
196                expression.(x => visit::<Expression, TypeExpression>(x)),
197                body.{
198                    attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
199                    statements.[].(x => recurse(x)),
200                }
201            },
202            else_if_clauses.[].{
203                attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
204                expression.(x => visit::<Expression, TypeExpression>(x)),
205                body.{
206                    attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
207                    statements.[].(x => recurse(x)),
208                }
209            },
210            else_clause.[].{
211                attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
212                body.{
213                    attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
214                    statements.[].(x => recurse(x)),
215                }
216            },
217        },
218        Statement::Switch.{
219            attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
220            expression.(x => visit::<Expression, TypeExpression>(x)),
221            body_attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
222            clauses.[].{
223                attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
224                case_selectors.[].CaseSelector::Expression.(x => visit::<Expression, TypeExpression>(x)),
225                body.{
226                    attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
227                    statements.[].(x => recurse(x)),
228                }
229            }
230        },
231        Statement::Loop.{
232            attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
233            body.{
234                attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
235                statements.[].(x => recurse(x)),
236            },
237            continuing.[].{
238                attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
239                body.{
240                    attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
241                    statements.[].(x => recurse(x)),
242                },
243                break_if.[].{
244                    attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
245                    expression.(x => visit::<Expression, TypeExpression>(x)),
246                }
247            }
248        },
249        Statement::For.{
250            attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
251            initializer.[].(x => recurse(x)),
252            condition.[].(x => visit::<Expression, TypeExpression>(x)),
253            update.[].(x => recurse(x)),
254            body.{
255                attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
256                statements.[].(x => recurse(x)),
257            },
258        },
259        Statement::While.{
260            attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
261            condition.(x => visit::<Expression, TypeExpression>(x)),
262            body.{
263                attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
264                statements.[].(x => recurse(x)),
265            },
266        },
267        Statement::Break.{
268            attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
269        },
270        Statement::Continue.{
271            attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
272        },
273        Statement::Return.{
274            attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
275            expression.[].(x => visit::<Expression, TypeExpression>(x)),
276        },
277        Statement::Discard.{
278            attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
279        },
280        Statement::FunctionCall.{
281            attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
282            call.{
283                ty,
284                arguments.[].(x => visit::<Expression, TypeExpression>(x)),
285            }
286        },
287        Statement::ConstAssert.{
288            attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
289            expression.(x => visit::<Expression, TypeExpression>(x)),
290        },
291        Statement::Declaration.{
292            attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
293            ty.[],
294            initializer.[].(x => visit::<Expression, TypeExpression>(x)),
295        },
296    }
297}
298
299impl_visit! { Statement => ExpressionNode,
300    {
301        Statement::Compound.statements.[].(x => recurse(x)),
302        Statement::Assignment.{ lhs, rhs },
303        Statement::Increment.expression,
304        Statement::Decrement.expression,
305        Statement::If.{
306            if_clause.{
307                expression,
308                body.statements.[].(x => recurse(x)),
309            },
310            else_if_clauses.[].{
311                expression,
312                body.statements.[].(x => recurse(x)),
313            },
314            else_clause.[].body.statements.[].(x => recurse(x)),
315        },
316        Statement::Switch.{
317            expression,
318            clauses.[].{
319                case_selectors.[].CaseSelector::Expression,
320                body.statements.[].(x => recurse(x)),
321            }
322        },
323        Statement::Loop.{
324            body.statements.[].(x => recurse(x)),
325            continuing.[].{
326                body.statements.[].(x => recurse(x)),
327                break_if.[].expression,
328            }
329        },
330        Statement::For.{
331            initializer.[].(x => recurse(x)),
332            condition.[],
333            update.[].(x => recurse(x)),
334            body.statements.[].(x => recurse(x)),
335        },
336        Statement::While.{
337            condition,
338            body.statements.[].(x => recurse(x)),
339        },
340        Statement::Return.expression.[],
341        Statement::FunctionCall.call.arguments.[],
342        Statement::ConstAssert.expression,
343        Statement::Declaration.initializer.[],
344    }
345}
346
347impl_visit! { Statement => StatementNode,
348    {
349        Statement::Compound.statements.[],
350        Statement::If.{
351            if_clause.body.statements.[],
352            else_if_clauses.[].body.statements.[],
353            else_clause.[].body.statements.[],
354        },
355        Statement::Switch.clauses.[].body.statements.[],
356        Statement::Loop.{
357            body.statements.[],
358            continuing.[].body.statements.[],
359        },
360        Statement::For.{
361            initializer.[],
362            update.[],
363            body.statements.[],
364        },
365        Statement::While.body.statements.[],
366    }
367}
368
369impl_visit! { Attribute => TypeExpression,
370    {
371        Attribute::Align.(x => visit::<Expression, TypeExpression>(x)),
372        Attribute::Binding.(x => visit::<Expression, TypeExpression>(x)),
373        Attribute::BlendSrc.(x => visit::<Expression, TypeExpression>(x)),
374        Attribute::Group.(x => visit::<Expression, TypeExpression>(x)),
375        Attribute::Id.(x => visit::<Expression, TypeExpression>(x)),
376        Attribute::Location.(x => visit::<Expression, TypeExpression>(x)),
377        Attribute::Size.(x => visit::<Expression, TypeExpression>(x)),
378        Attribute::WorkgroupSize.{
379            x.(x => visit::<Expression, TypeExpression>(x)),
380            y.[].(x => visit::<Expression, TypeExpression>(x)),
381            z.[].(x => visit::<Expression, TypeExpression>(x)),
382        },
383        #[cfg(feature = "naga-ext")]
384        Attribute::Payload.(x => visit::<Expression, TypeExpression>(x)),
385        #[cfg(feature = "naga-ext")]
386        Attribute::Mesh.(x => visit::<Expression, TypeExpression>(x)),
387        #[cfg(feature = "naga-ext")]
388        Attribute::IncomingPayload.(x => visit::<Expression, TypeExpression>(x)),
389        #[cfg(feature = "generics")]
390        Attribute::Type.variants.[],
391        Attribute::Custom.arguments.[].[].(x => visit::<Expression, TypeExpression>(x))
392    }
393}
394
395impl_visit! { TranslationUnit => ExpressionNode,
396    {
397        global_declarations.[].(x => visit::<GlobalDeclaration, ExpressionNode>(x))
398    }
399}
400
401impl_visit! { GlobalDeclaration => ExpressionNode,
402    {
403        GlobalDeclaration::Declaration.{
404            initializer.[],
405        },
406        GlobalDeclaration::Function.{
407            body.statements.[].(x => visit::<Statement, ExpressionNode>(x)),
408        },
409        GlobalDeclaration::Compound.body.[].(x => visit::<GlobalDeclaration, ExpressionNode>(x)),
410    }
411}
412
413impl_visit! { TranslationUnit => StatementNode,
414    {
415        global_declarations.[].(x => visit::<GlobalDeclaration, StatementNode>(x)),
416    }
417}
418
419impl_visit! { GlobalDeclaration => StatementNode,
420    {
421        GlobalDeclaration::Function.body.statements.[]
422    }
423}
424
425impl_visit! { TranslationUnit => Attributes,
426    {
427        imports.[].attributes,
428        global_directives.[].{
429            GlobalDirective::Diagnostic.attributes,
430            GlobalDirective::Enable.attributes,
431            GlobalDirective::Requires.attributes,
432        },
433        global_declarations.[].(x => visit::<GlobalDeclaration, Attributes>(x)),
434    }
435}
436
437impl_visit! { GlobalDeclaration => Attributes,
438    {
439        GlobalDeclaration::Declaration.attributes,
440        GlobalDeclaration::TypeAlias.attributes,
441        GlobalDeclaration::Struct.{
442            attributes,
443            members.[].attributes,
444        },
445        GlobalDeclaration::Function.{
446            attributes,
447            parameters.[].attributes,
448            body.{ attributes, statements.[].(x => visit::<Statement, Attributes>(x)) }
449        },
450        GlobalDeclaration::ConstAssert.attributes,
451        GlobalDeclaration::Compound.{
452            attributes,
453            body.[].(x => visit::<GlobalDeclaration, Attributes>(x))
454        }
455    }
456}
457
458impl_visit! { TranslationUnit => TypeExpression,
459    {
460        global_declarations.[].(x => visit::<GlobalDeclaration, TypeExpression>(x))
461    }
462}
463
464impl_visit! { GlobalDeclaration => TypeExpression,
465    {
466        GlobalDeclaration::Declaration.(x => visit::<Declaration, TypeExpression>(x)),
467        GlobalDeclaration::TypeAlias.(x => visit::<TypeAlias, TypeExpression>(x)),
468        GlobalDeclaration::Struct.(x => visit::<Struct, TypeExpression>(x)),
469        GlobalDeclaration::Function.(x => visit::<Function, TypeExpression>(x)),
470        GlobalDeclaration::ConstAssert.(x => visit::<ConstAssert, TypeExpression>(x)),
471        GlobalDeclaration::Compound.{
472            attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
473            body.[].(x => visit::<GlobalDeclaration, TypeExpression>(x))
474        }
475    }
476}
477
478impl_visit! { Declaration => TypeExpression,
479    {
480        attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
481        ty.[],
482        initializer.[].(x => visit::<Expression, TypeExpression>(x)),
483    }
484}
485impl_visit! { TypeAlias => TypeExpression,
486    {
487        attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
488        ty,
489    }
490}
491impl_visit! { Struct => TypeExpression,
492    {
493        attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
494        members.[].(x => visit::<StructMember, TypeExpression>(x)),
495    }
496}
497impl_visit! { StructMember => TypeExpression,
498    {
499        attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
500        ty,
501    }
502}
503impl_visit! { Function => TypeExpression,
504    {
505        attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
506        parameters.[].{
507            attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
508            ty,
509        },
510        return_attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
511        return_type.[],
512        body.{
513            attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
514            statements.[].(x => visit::<Statement, TypeExpression>(x)),
515        }
516    }
517}
518impl_visit! { ConstAssert => TypeExpression,
519    {
520        attributes.[].(x => visit::<Attribute, TypeExpression>(x)),
521        expression.(x => visit::<Expression, TypeExpression>(x)),
522    }
523}
524
525#[test]
526fn test_visit_if_else() {
527    let mut statement: Statement =
528        "if false { return 1; } else if false { return 2; } else { return 3; }"
529            .parse()
530            .expect("parse failure");
531
532    assert_eq!(
533        Visit::<StatementNode>::visit(&statement)
534            .map(|statement| statement.to_string())
535            .collect::<Vec<_>>(),
536        ["return 1;", "return 2;", "return 3;"]
537    );
538    assert_eq!(
539        Visit::<StatementNode>::visit_mut(&mut statement)
540            .map(|statement| statement.to_string())
541            .collect::<Vec<_>>(),
542        ["return 1;", "return 2;", "return 3;"]
543    );
544    assert_eq!(
545        Visit::<ExpressionNode>::visit(&statement)
546            .map(|expression| expression.to_string())
547            .collect::<Vec<_>>(),
548        ["false", "1", "false", "2", "3"]
549    );
550    assert_eq!(
551        Visit::<ExpressionNode>::visit_mut(&mut statement)
552            .map(|expression| expression.to_string())
553            .collect::<Vec<_>>(),
554        ["false", "1", "false", "2", "3"]
555    );
556}
wesl/pass/visit.rs

wesl/pass/
visit.rs
