flux_syntax/surface/
visit.rs

1//! A visitor for types in [`surface`]
2//!
3//! In general there's no specific order except that [refinement parameters] should be
4//! visited in an order that matches their scope. Name resolution relies on this gurantee.
5//!
6//! [`surface`]: crate::surface
7//! [refinement parameters]: crate::surface::RefineParam
8use rustc_span::symbol::Ident;
9
10use super::{
11    Async, BaseSort, BaseTy, BaseTyKind, ConstArg, ConstantInfo, ConstructorArg, Ensures, EnumDef,
12    Expr, ExprKind, ExprPath, ExprPathSegment, FieldExpr, FnInput, FnOutput, FnRetTy, FnSig,
13    GenericArg, GenericArgKind, GenericParam, Generics, Impl, ImplAssocReft, Indices, ItemKind,
14    Lit, Path, PathSegment, Qualifier, RefineArg, RefineParam, Sort, SortPath, SpecFunc, StructDef,
15    Trait, TraitAssocReft, TraitRef, Ty, TyAlias, TyKind, VariantDef, VariantRet,
16    WhereBoundPredicate,
17};
18use crate::surface::{FluxItem, ImplItemFn, Item, PrimOpProp, TraitItemFn};
19
20#[macro_export]
21macro_rules! walk_list {
22    ($visitor: expr, $method: ident, $list: expr $(, $($extra_args: expr),* )?) => {
23        {
24            #[allow(for_loops_over_fallibles)]
25            for elem in $list {
26                $visitor.$method(elem $(, $($extra_args,)* )?)
27            }
28        }
29    }
30}
31
32pub trait Visitor: Sized {
33    fn visit_flux_item(&mut self, item: &FluxItem) {
34        walk_flux_item(self, item);
35    }
36
37    fn visit_qualifier(&mut self, qualifier: &Qualifier) {
38        walk_qualifier(self, qualifier);
39    }
40
41    fn visit_defn(&mut self, defn: &SpecFunc) {
42        walk_defn(self, defn);
43    }
44
45    fn visit_primop_prop(&mut self, prop: &PrimOpProp) {
46        walk_primop_prop(self, prop);
47    }
48
49    fn visit_refine_param(&mut self, param: &RefineParam) {
50        walk_refine_param(self, param);
51    }
52
53    fn visit_generic_param(&mut self, param: &GenericParam) {
54        walk_generic_param(self, param);
55    }
56
57    fn visit_sort(&mut self, sort: &Sort) {
58        walk_sort(self, sort);
59    }
60
61    fn visit_item(&mut self, item: &Item) {
62        walk_item(self, item);
63    }
64
65    fn visit_trait_item(&mut self, item: &TraitItemFn) {
66        walk_trait_item(self, item);
67    }
68
69    fn visit_impl_item(&mut self, item: &ImplItemFn) {
70        walk_impl_item(self, item);
71    }
72
73    fn visit_trait(&mut self, trait_: &Trait) {
74        walk_trait(self, trait_);
75    }
76
77    fn visit_trait_assoc_reft(&mut self, assoc_reft: &TraitAssocReft) {
78        walk_trait_assoc_reft(self, assoc_reft);
79    }
80
81    fn visit_impl(&mut self, impl_: &Impl) {
82        walk_impl(self, impl_);
83    }
84
85    fn visit_impl_assoc_reft(&mut self, assoc_reft: &ImplAssocReft) {
86        walk_impl_assoc_reft(self, assoc_reft);
87    }
88
89    fn visit_trait_ref(&mut self, trait_ref: &TraitRef) {
90        walk_trait_ref(self, trait_ref);
91    }
92
93    fn visit_base_sort(&mut self, bsort: &BaseSort) {
94        walk_base_sort(self, bsort);
95    }
96
97    fn visit_sort_path(&mut self, path: &SortPath) {
98        walk_sort_path(self, path);
99    }
100
101    fn visit_ty_alias(&mut self, ty_alias: &TyAlias) {
102        walk_ty_alias(self, ty_alias);
103    }
104
105    fn visit_struct_def(&mut self, struct_def: &StructDef) {
106        walk_struct_def(self, struct_def);
107    }
108
109    fn visit_constant(&mut self, _constant: &ConstantInfo) {
110        walk_constant(self, _constant);
111    }
112
113    fn visit_enum_def(&mut self, enum_def: &EnumDef) {
114        walk_enum_def(self, enum_def);
115    }
116
117    fn visit_variant(&mut self, variant: &VariantDef) {
118        walk_variant(self, variant);
119    }
120
121    fn visit_variant_ret(&mut self, ret: &VariantRet) {
122        walk_variant_ret(self, ret);
123    }
124
125    fn visit_fn_sig(&mut self, fn_sig: &FnSig) {
126        walk_fn_sig(self, fn_sig);
127    }
128
129    fn visit_fn_output(&mut self, fn_output: &FnOutput) {
130        walk_fn_output(self, fn_output);
131    }
132
133    fn visit_async(&mut self, _asyncness: &Async) {}
134
135    fn visit_generics(&mut self, generics: &Generics) {
136        walk_generics(self, generics);
137    }
138
139    fn visit_fn_input(&mut self, input: &FnInput) {
140        walk_fn_input(self, input);
141    }
142
143    fn visit_fn_ret_ty(&mut self, fn_ret_ty: &FnRetTy) {
144        walk_fn_ret_ty(self, fn_ret_ty);
145    }
146
147    fn visit_ensures(&mut self, ensures: &Ensures) {
148        walk_ensures(self, ensures);
149    }
150
151    fn visit_where_predicate(&mut self, predicate: &WhereBoundPredicate) {
152        walk_where_predicate(self, predicate);
153    }
154
155    fn visit_generic_arg(&mut self, arg: &GenericArg) {
156        walk_generic_arg(self, arg);
157    }
158
159    fn visit_refine_arg(&mut self, arg: &RefineArg) {
160        walk_refine_arg(self, arg);
161    }
162
163    fn visit_indices(&mut self, indices: &Indices) {
164        walk_indices(self, indices);
165    }
166
167    fn visit_ty(&mut self, ty: &Ty) {
168        walk_ty(self, ty);
169    }
170
171    fn visit_const_arg(&mut self, _const_arg: &ConstArg) {}
172
173    fn visit_bty(&mut self, bty: &BaseTy) {
174        walk_bty(self, bty);
175    }
176
177    fn visit_path(&mut self, path: &Path) {
178        walk_path(self, path);
179    }
180
181    fn visit_path_segment(&mut self, segment: &PathSegment) {
182        walk_path_segment(self, segment);
183    }
184
185    fn visit_expr(&mut self, expr: &Expr) {
186        walk_expr(self, expr);
187    }
188
189    fn visit_constructor_args(&mut self, expr: &ConstructorArg) {
190        match expr {
191            ConstructorArg::FieldExpr(field_expr) => walk_field_expr(self, field_expr),
192            ConstructorArg::Spread(spread) => self.visit_expr(&spread.expr),
193        }
194    }
195
196    fn visit_path_expr(&mut self, qpath: &ExprPath) {
197        walk_path_expr(self, qpath);
198    }
199
200    fn visit_path_expr_segment(&mut self, segment: &ExprPathSegment) {
201        walk_path_expr_segment(self, segment);
202    }
203
204    fn visit_ident(&mut self, _ident: Ident) {}
205
206    fn visit_literal(&mut self, _lit: Lit) {}
207}
208
209pub fn walk_flux_item<V: Visitor>(vis: &mut V, item: &FluxItem) {
210    match item {
211        FluxItem::Qualifier(qualifier) => vis.visit_qualifier(qualifier),
212        FluxItem::FuncDef(spec_func) => vis.visit_defn(spec_func),
213        FluxItem::SortDecl(_sort_decl) => {}
214        FluxItem::PrimOpProp(prim_op_prop) => vis.visit_primop_prop(prim_op_prop),
215    }
216}
217
218pub fn walk_qualifier<V: Visitor>(vis: &mut V, qualifier: &Qualifier) {
219    vis.visit_ident(qualifier.name);
220    walk_list!(vis, visit_refine_param, &qualifier.params);
221    vis.visit_expr(&qualifier.expr);
222}
223
224pub fn walk_defn<V: Visitor>(vis: &mut V, defn: &SpecFunc) {
225    vis.visit_ident(defn.name);
226    walk_list!(vis, visit_ident, defn.sort_vars.iter().copied());
227    walk_list!(vis, visit_refine_param, &defn.params);
228    vis.visit_sort(&defn.output);
229    if let Some(body) = &defn.body {
230        vis.visit_expr(body);
231    }
232}
233
234pub fn walk_primop_prop<V: Visitor>(vis: &mut V, prop: &PrimOpProp) {
235    vis.visit_ident(prop.name);
236    walk_list!(vis, visit_refine_param, &prop.params);
237    vis.visit_expr(&prop.body);
238}
239
240pub fn walk_refine_param<V: Visitor>(vis: &mut V, param: &RefineParam) {
241    vis.visit_ident(param.ident);
242    vis.visit_sort(&param.sort);
243}
244
245pub fn walk_generic_param<V: Visitor>(vis: &mut V, param: &GenericParam) {
246    vis.visit_ident(param.name);
247}
248
249pub fn walk_sort<V: Visitor>(vis: &mut V, sort: &Sort) {
250    match sort {
251        Sort::Base(bsort) => vis.visit_base_sort(bsort),
252        Sort::Func { inputs, output } => {
253            walk_list!(vis, visit_base_sort, inputs);
254            vis.visit_base_sort(output);
255        }
256        Sort::Infer => {}
257    }
258}
259
260pub fn walk_item<V: Visitor>(vis: &mut V, item: &Item) {
261    match &item.kind {
262        ItemKind::Fn(fn_sig) => {
263            if let Some(fn_sig) = fn_sig {
264                vis.visit_fn_sig(fn_sig);
265            }
266        }
267        ItemKind::Struct(struct_def) => vis.visit_struct_def(struct_def),
268        ItemKind::Enum(enum_def) => vis.visit_enum_def(enum_def),
269        ItemKind::Trait(trait_) => vis.visit_trait(trait_),
270        ItemKind::Impl(impl_) => vis.visit_impl(impl_),
271        ItemKind::Const(cst) => vis.visit_constant(cst),
272        ItemKind::TyAlias(ty_alias) => vis.visit_ty_alias(ty_alias),
273        ItemKind::Mod => {}
274    }
275}
276
277pub fn walk_trait_item<V: Visitor>(vis: &mut V, item: &TraitItemFn) {
278    if let Some(fn_sig) = item.sig.as_ref() {
279        vis.visit_fn_sig(fn_sig);
280    }
281}
282
283pub fn walk_impl_item<V: Visitor>(vis: &mut V, item: &ImplItemFn) {
284    if let Some(fn_sig) = item.sig.as_ref() {
285        vis.visit_fn_sig(fn_sig);
286    }
287}
288
289pub fn walk_trait<V: Visitor>(vis: &mut V, trait_: &Trait) {
290    if let Some(generics) = &trait_.generics {
291        vis.visit_generics(generics);
292    }
293    walk_list!(vis, visit_trait_assoc_reft, &trait_.assoc_refinements);
294}
295
296pub fn walk_trait_assoc_reft<V: Visitor>(vis: &mut V, assoc_reft: &TraitAssocReft) {
297    vis.visit_ident(assoc_reft.name);
298    walk_list!(vis, visit_refine_param, &assoc_reft.params);
299    vis.visit_base_sort(&assoc_reft.output);
300    if let Some(expr) = &assoc_reft.body {
301        vis.visit_expr(expr);
302    }
303}
304
305pub fn walk_impl<V: Visitor>(vis: &mut V, impl_: &Impl) {
306    if let Some(generics) = &impl_.generics {
307        vis.visit_generics(generics);
308    }
309    walk_list!(vis, visit_impl_assoc_reft, &impl_.assoc_refinements);
310}
311
312pub fn walk_impl_assoc_reft<V: Visitor>(vis: &mut V, assoc_reft: &ImplAssocReft) {
313    vis.visit_ident(assoc_reft.name);
314    walk_list!(vis, visit_refine_param, &assoc_reft.params);
315    vis.visit_expr(&assoc_reft.body);
316    vis.visit_base_sort(&assoc_reft.output);
317}
318
319pub fn walk_trait_ref<V: Visitor>(vis: &mut V, trait_ref: &TraitRef) {
320    vis.visit_path(&trait_ref.path);
321}
322
323pub fn walk_base_sort<V: Visitor>(vis: &mut V, bsort: &BaseSort) {
324    match bsort {
325        BaseSort::BitVec(_len) => {}
326        BaseSort::Path(path) => vis.visit_sort_path(path),
327        BaseSort::SortOf(qself, path) => {
328            vis.visit_ty(qself);
329            vis.visit_path(path);
330        }
331    }
332}
333
334pub fn walk_sort_path<V: Visitor>(vis: &mut V, path: &SortPath) {
335    walk_list!(vis, visit_ident, path.segments.iter().copied());
336    walk_list!(vis, visit_base_sort, &path.args);
337}
338
339pub fn walk_ty_alias<V: Visitor>(vis: &mut V, ty_alias: &TyAlias) {
340    vis.visit_ident(ty_alias.ident);
341    vis.visit_generics(&ty_alias.generics);
342    walk_list!(vis, visit_refine_param, &ty_alias.params);
343    if let Some(index) = &ty_alias.index {
344        vis.visit_refine_param(index);
345    }
346    vis.visit_ty(&ty_alias.ty);
347}
348pub fn walk_constant<V: Visitor>(vis: &mut V, constant_info: &ConstantInfo) {
349    if let Some(expr) = &constant_info.expr {
350        vis.visit_expr(expr);
351    }
352}
353
354pub fn walk_struct_def<V: Visitor>(vis: &mut V, struct_def: &StructDef) {
355    if let Some(generics) = &struct_def.generics {
356        vis.visit_generics(generics);
357    }
358    if let Some(refined_by) = &struct_def.refined_by {
359        walk_list!(vis, visit_refine_param, refined_by);
360    }
361    walk_list!(vis, visit_expr, &struct_def.invariants);
362    struct_def.fields.iter().flatten().for_each(|field| {
363        vis.visit_ty(field);
364    });
365}
366
367pub fn walk_enum_def<V: Visitor>(vis: &mut V, enum_def: &EnumDef) {
368    if let Some(refined_by) = &enum_def.refined_by {
369        walk_list!(vis, visit_refine_param, refined_by);
370    }
371    walk_list!(vis, visit_expr, &enum_def.invariants);
372    enum_def
373        .variants
374        .iter()
375        .flatten()
376        .for_each(|variant| vis.visit_variant(variant));
377}
378
379pub fn walk_variant<V: Visitor>(vis: &mut V, variant: &VariantDef) {
380    walk_list!(vis, visit_ty, &variant.fields);
381    if let Some(ret) = &variant.ret {
382        vis.visit_variant_ret(ret);
383    }
384}
385
386pub fn walk_fn_trait_ref<V: Visitor>(vis: &mut V, in_arg: &GenericArg, out_arg: &GenericArg) {
387    vis.visit_generic_arg(in_arg);
388    vis.visit_generic_arg(out_arg);
389}
390
391pub fn walk_variant_ret<V: Visitor>(vis: &mut V, ret: &VariantRet) {
392    vis.visit_path(&ret.path);
393    vis.visit_indices(&ret.indices);
394}
395
396pub fn walk_fn_sig<V: Visitor>(vis: &mut V, fn_sig: &FnSig) {
397    vis.visit_async(&fn_sig.asyncness);
398    vis.visit_generics(&fn_sig.generics);
399    walk_list!(vis, visit_refine_param, &fn_sig.params);
400    for requires in &fn_sig.requires {
401        walk_list!(vis, visit_refine_param, &requires.params);
402        vis.visit_expr(&requires.pred);
403    }
404    walk_list!(vis, visit_fn_input, &fn_sig.inputs);
405    vis.visit_fn_output(&fn_sig.output);
406}
407
408pub fn walk_fn_output<V: Visitor>(vis: &mut V, fn_output: &FnOutput) {
409    vis.visit_fn_ret_ty(&fn_output.returns);
410    walk_list!(vis, visit_ensures, &fn_output.ensures);
411}
412
413pub fn walk_generics<V: Visitor>(vis: &mut V, generics: &Generics) {
414    walk_list!(vis, visit_generic_param, &generics.params);
415    if let Some(predicates) = &generics.predicates {
416        walk_list!(vis, visit_where_predicate, predicates);
417    }
418}
419
420pub fn walk_fn_input<V: Visitor>(vis: &mut V, arg: &FnInput) {
421    match arg {
422        FnInput::Constr(bind, path, pred, _node_id) => {
423            vis.visit_ident(*bind);
424            vis.visit_path(path);
425            vis.visit_expr(pred);
426        }
427        FnInput::StrgRef(bind, ty, _node_id) => {
428            vis.visit_ident(*bind);
429            vis.visit_ty(ty);
430        }
431        FnInput::Ty(bind, ty, _node_id) => {
432            if let Some(bind) = bind {
433                vis.visit_ident(*bind);
434            }
435            vis.visit_ty(ty);
436        }
437    }
438}
439
440pub fn walk_fn_ret_ty<V: Visitor>(vis: &mut V, fn_ret_ty: &FnRetTy) {
441    match fn_ret_ty {
442        FnRetTy::Default(_span) => {}
443        FnRetTy::Ty(ty) => vis.visit_ty(ty),
444    }
445}
446
447pub fn walk_ensures<V: Visitor>(vis: &mut V, ensures: &Ensures) {
448    match ensures {
449        Ensures::Type(bind, ty, _node_id) => {
450            vis.visit_ident(*bind);
451            vis.visit_ty(ty);
452        }
453        Ensures::Pred(pred) => {
454            vis.visit_expr(pred);
455        }
456    }
457}
458
459pub fn walk_where_predicate<V: Visitor>(vis: &mut V, predicate: &WhereBoundPredicate) {
460    vis.visit_ty(&predicate.bounded_ty);
461    walk_list!(vis, visit_trait_ref, &predicate.bounds);
462}
463
464pub fn walk_generic_arg<V: Visitor>(vis: &mut V, arg: &GenericArg) {
465    match &arg.kind {
466        GenericArgKind::Type(ty) => {
467            vis.visit_ty(ty);
468        }
469        GenericArgKind::Constraint(ident, ty) => {
470            vis.visit_ident(*ident);
471            vis.visit_ty(ty);
472        }
473    }
474}
475
476pub fn walk_refine_arg<V: Visitor>(vis: &mut V, arg: &RefineArg) {
477    match arg {
478        RefineArg::Bind(ident, _kind, _span, _node_id) => {
479            vis.visit_ident(*ident);
480        }
481        RefineArg::Expr(e) => {
482            vis.visit_expr(e);
483        }
484        RefineArg::Abs(params, e, _node_id, _span) => {
485            walk_list!(vis, visit_refine_param, params);
486            vis.visit_expr(e);
487        }
488    }
489}
490
491pub fn walk_indices<V: Visitor>(vis: &mut V, indices: &Indices) {
492    walk_list!(vis, visit_refine_arg, &indices.indices);
493}
494
495pub fn walk_ty<V: Visitor>(vis: &mut V, ty: &Ty) {
496    match &ty.kind {
497        TyKind::Base(bty) => vis.visit_bty(bty),
498        TyKind::Indexed { bty, indices } => {
499            vis.visit_indices(indices);
500            vis.visit_bty(bty);
501        }
502        TyKind::Exists { bind, bty, pred } => {
503            vis.visit_ident(*bind);
504            vis.visit_bty(bty);
505            vis.visit_expr(pred);
506        }
507        TyKind::GeneralExists { params, ty, pred } => {
508            walk_list!(vis, visit_refine_param, params);
509            vis.visit_ty(ty);
510            if let Some(pred) = pred {
511                vis.visit_expr(pred);
512            }
513        }
514        TyKind::Ref(_mutbl, ty) => {
515            vis.visit_ty(ty);
516        }
517        TyKind::Constr(pred, ty) => {
518            vis.visit_expr(pred);
519            vis.visit_ty(ty);
520        }
521        TyKind::Tuple(tys) => {
522            walk_list!(vis, visit_ty, tys);
523        }
524        TyKind::Array(ty, len) => {
525            vis.visit_const_arg(len);
526            vis.visit_ty(ty);
527        }
528        TyKind::ImplTrait(_node_id, trait_ref) => {
529            walk_list!(vis, visit_trait_ref, trait_ref);
530        }
531        TyKind::Hole => {}
532    }
533}
534
535pub fn walk_bty<V: Visitor>(vis: &mut V, bty: &BaseTy) {
536    match &bty.kind {
537        BaseTyKind::Path(qself, path) => {
538            if let Some(qself) = qself {
539                vis.visit_ty(qself);
540            }
541            vis.visit_path(path);
542        }
543        BaseTyKind::Slice(ty) => vis.visit_ty(ty),
544    }
545}
546
547pub fn walk_path<V: Visitor>(vis: &mut V, path: &Path) {
548    walk_list!(vis, visit_path_segment, &path.segments);
549    walk_list!(vis, visit_refine_arg, &path.refine);
550}
551
552pub fn walk_path_segment<V: Visitor>(vis: &mut V, segment: &PathSegment) {
553    vis.visit_ident(segment.ident);
554    walk_list!(vis, visit_generic_arg, &segment.args);
555}
556
557pub fn walk_field_expr<V: Visitor>(vis: &mut V, expr: &FieldExpr) {
558    vis.visit_ident(expr.ident);
559    vis.visit_expr(&expr.expr);
560}
561
562pub fn walk_expr<V: Visitor>(vis: &mut V, expr: &Expr) {
563    match &expr.kind {
564        ExprKind::Path(qpath) => vis.visit_path_expr(qpath),
565        ExprKind::Dot(base, fld) => {
566            vis.visit_expr(base);
567            vis.visit_ident(*fld);
568        }
569        ExprKind::Literal(lit) => {
570            vis.visit_literal(*lit);
571        }
572        ExprKind::BinaryOp(_bin_op, box exprs) => {
573            walk_list!(vis, visit_expr, exprs);
574        }
575        ExprKind::UnaryOp(_un_op, e) => {
576            vis.visit_expr(e);
577        }
578        ExprKind::PrimUIF(_) => {}
579        ExprKind::Call(callee, args) => {
580            vis.visit_expr(callee);
581            walk_list!(vis, visit_expr, args);
582        }
583        ExprKind::AssocReft(qself, path, name) => {
584            vis.visit_ty(qself);
585            vis.visit_path(path);
586            vis.visit_ident(*name);
587        }
588        ExprKind::IfThenElse(box exprs) => {
589            walk_list!(vis, visit_expr, exprs);
590        }
591        ExprKind::Constructor(path, exprs) => {
592            if let Some(path) = path {
593                vis.visit_path_expr(path);
594            }
595            walk_list!(vis, visit_constructor_args, exprs);
596        }
597        ExprKind::BoundedQuant(_, i, _, e) => {
598            vis.visit_refine_param(i);
599            vis.visit_expr(e);
600        }
601        ExprKind::Block(decls, body) => {
602            for decl in decls {
603                // the order here is important because the parameter is not in scope
604                // in the initializer
605                vis.visit_expr(&decl.init);
606                vis.visit_refine_param(&decl.param);
607            }
608            vis.visit_expr(body);
609        }
610    }
611}
612
613pub fn walk_path_expr<V: Visitor>(vis: &mut V, qpath: &ExprPath) {
614    walk_list!(vis, visit_path_expr_segment, &qpath.segments);
615}
616
617pub fn walk_path_expr_segment<V: Visitor>(vis: &mut V, segment: &ExprPathSegment) {
618    vis.visit_ident(segment.ident);
619}