flux_middle/fhir/
visit.rs

1use super::{
2    AliasReft, AssocItemConstraint, AssocItemConstraintKind, BaseTy, BaseTyKind, Ensures, EnumDef,
3    Expr, ExprKind, FieldDef, FieldExpr, FluxItem, FnDecl, FnOutput, FnSig, ForeignItem,
4    ForeignItemKind, FuncSort, GenericArg, GenericBound, Generics, Impl, ImplAssocReft, ImplItem,
5    ImplItemKind, Item, ItemKind, Lifetime, Lit, OpaqueTy, OwnerNode, Path, PathExpr, PathSegment,
6    PolyFuncSort, PolyTraitRef, QPath, Qualifier, RefineParam, Requires, Sort, SortPath, SpecFunc,
7    StructDef, TraitAssocReft, TraitItem, TraitItemKind, Ty, TyAlias, TyKind, VariantDef,
8    VariantRet, WhereBoundPredicate,
9};
10use crate::fhir::StructKind;
11
12#[macro_export]
13macro_rules! walk_list {
14    ($visitor: expr, $method: ident, $list: expr $(, $($extra_args: expr),* )?) => {
15        {
16            #[allow(for_loops_over_fallibles)]
17            for elem in $list {
18                $visitor.$method(elem $(, $($extra_args,)* )?)
19            }
20        }
21    }
22}
23
24pub trait Visitor<'v>: Sized {
25    fn visit_flux_item(&mut self, item: &FluxItem<'v>) {
26        walk_flux_item(self, item);
27    }
28
29    fn visit_qualifier(&mut self, qualifier: &Qualifier<'v>) {
30        walk_qualifier(self, qualifier);
31    }
32
33    fn visit_func(&mut self, func: &SpecFunc<'v>) {
34        walk_func(self, func);
35    }
36
37    fn visit_node(&mut self, node: &OwnerNode<'v>) {
38        walk_node(self, node);
39    }
40
41    fn visit_item(&mut self, item: &Item<'v>) {
42        walk_item(self, item);
43    }
44
45    fn visit_trait_item(&mut self, trait_item: &TraitItem<'v>) {
46        walk_trait_item(self, trait_item);
47    }
48
49    fn visit_impl_item(&mut self, impl_item: &ImplItem<'v>) {
50        walk_impl_item(self, impl_item);
51    }
52
53    fn visit_foreign_item(&mut self, foreign_item: &ForeignItem<'v>) {
54        walk_foreign_item(self, foreign_item);
55    }
56
57    fn visit_generics(&mut self, generics: &Generics<'v>) {
58        walk_generics(self, generics);
59    }
60
61    fn visit_where_predicate(&mut self, predicate: &WhereBoundPredicate<'v>) {
62        walk_where_predicate(self, predicate);
63    }
64
65    fn visit_impl(&mut self, impl_: &Impl<'v>) {
66        walk_impl(self, impl_);
67    }
68
69    fn visit_impl_assoc_reft(&mut self, assoc_reft: &ImplAssocReft<'v>) {
70        walk_impl_assoc_reft(self, assoc_reft);
71    }
72
73    fn visit_trait_assoc_reft(&mut self, assoc_reft: &TraitAssocReft<'v>) {
74        walk_trait_assoc_reft(self, assoc_reft);
75    }
76
77    fn visit_struct_def(&mut self, struct_def: &StructDef<'v>) {
78        walk_struct_def(self, struct_def);
79    }
80
81    fn visit_enum_def(&mut self, enum_def: &EnumDef<'v>) {
82        walk_enum_def(self, enum_def);
83    }
84
85    fn visit_variant(&mut self, variant: &VariantDef<'v>) {
86        walk_variant(self, variant);
87    }
88
89    fn visit_field_def(&mut self, field: &FieldDef<'v>) {
90        walk_field_def(self, field);
91    }
92
93    fn visit_variant_ret(&mut self, ret: &VariantRet<'v>) {
94        walk_variant_ret(self, ret);
95    }
96
97    fn visit_ty_alias(&mut self, ty_alias: &TyAlias<'v>) {
98        walk_ty_alias(self, ty_alias);
99    }
100
101    fn visit_opaque_ty(&mut self, opaque_ty: &OpaqueTy<'v>) {
102        walk_opaque_ty(self, opaque_ty);
103    }
104
105    fn visit_generic_bound(&mut self, bound: &GenericBound<'v>) {
106        walk_generic_bound(self, bound);
107    }
108
109    fn visit_poly_trait_ref(&mut self, trait_ref: &PolyTraitRef<'v>) {
110        walk_poly_trait_ref(self, trait_ref);
111    }
112
113    fn visit_fn_sig(&mut self, sig: &FnSig<'v>) {
114        walk_fn_sig(self, sig);
115    }
116
117    fn visit_fn_decl(&mut self, decl: &FnDecl<'v>) {
118        walk_fn_decl(self, decl);
119    }
120
121    fn visit_refine_param(&mut self, param: &RefineParam<'v>) {
122        walk_refine_param(self, param);
123    }
124
125    fn visit_requires(&mut self, requires: &Requires<'v>) {
126        walk_requires(self, requires);
127    }
128
129    fn visit_ensures(&mut self, ensures: &Ensures<'v>) {
130        walk_ensures(self, ensures);
131    }
132
133    fn visit_fn_output(&mut self, output: &FnOutput<'v>) {
134        walk_fn_output(self, output);
135    }
136
137    fn visit_generic_arg(&mut self, arg: &GenericArg<'v>) {
138        walk_generic_arg(self, arg);
139    }
140
141    fn visit_lifetime(&mut self, _lft: &Lifetime) {}
142
143    fn visit_ty(&mut self, ty: &Ty<'v>) {
144        walk_ty(self, ty);
145    }
146
147    fn visit_bty(&mut self, bty: &BaseTy<'v>) {
148        walk_bty(self, bty);
149    }
150
151    fn visit_qpath(&mut self, qpath: &QPath<'v>) {
152        walk_qpath(self, qpath);
153    }
154
155    fn visit_path(&mut self, path: &Path<'v>) {
156        walk_path(self, path);
157    }
158
159    fn visit_path_segment(&mut self, segment: &PathSegment<'v>) {
160        walk_path_segment(self, segment);
161    }
162
163    fn visit_assoc_item_constraint(&mut self, constraint: &AssocItemConstraint<'v>) {
164        walk_assoc_item_constraint(self, constraint);
165    }
166
167    fn visit_sort(&mut self, sort: &Sort<'v>) {
168        walk_sort(self, sort);
169    }
170
171    fn visit_sort_path(&mut self, path: &SortPath<'v>) {
172        walk_sort_path(self, path);
173    }
174
175    fn visit_poly_func_sort(&mut self, func: &PolyFuncSort<'v>) {
176        walk_poly_func_sort(self, func);
177    }
178
179    fn visit_func_sort(&mut self, func: &FuncSort<'v>) {
180        walk_func_sort(self, func);
181    }
182
183    fn visit_expr(&mut self, expr: &Expr<'v>) {
184        walk_expr(self, expr);
185    }
186
187    fn visit_field_expr(&mut self, expr: &FieldExpr<'v>) {
188        walk_field_expr(self, expr);
189    }
190
191    fn visit_alias_reft(&mut self, alias_reft: &AliasReft<'v>) {
192        walk_alias_reft(self, alias_reft);
193    }
194
195    fn visit_literal(&mut self, _lit: &Lit) {}
196
197    fn visit_path_expr(&mut self, _path: &PathExpr<'v>) {}
198}
199
200fn walk_func<'v, V: Visitor<'v>>(vis: &mut V, func: &SpecFunc<'v>) {
201    walk_list!(vis, visit_refine_param, func.args);
202    vis.visit_sort(&func.sort);
203    if let Some(body) = &func.body {
204        vis.visit_expr(body);
205    }
206}
207
208fn walk_qualifier<'v, V: Visitor<'v>>(vis: &mut V, qualifier: &Qualifier<'v>) {
209    walk_list!(vis, visit_refine_param, qualifier.args);
210    vis.visit_expr(&qualifier.expr);
211}
212
213fn walk_flux_item<'v, V: Visitor<'v>>(vis: &mut V, item: &FluxItem<'v>) {
214    match item {
215        FluxItem::Qualifier(qualifier) => {
216            vis.visit_qualifier(qualifier);
217        }
218        FluxItem::Func(func) => {
219            vis.visit_func(func);
220        }
221    }
222}
223
224pub fn walk_impl<'v, V: Visitor<'v>>(vis: &mut V, impl_: &Impl<'v>) {
225    walk_list!(vis, visit_impl_assoc_reft, impl_.assoc_refinements);
226}
227
228pub fn walk_struct_def<'v, V: Visitor<'v>>(vis: &mut V, struct_def: &StructDef<'v>) {
229    walk_list!(vis, visit_refine_param, struct_def.params);
230    walk_list!(vis, visit_expr, struct_def.invariants);
231    if let StructKind::Transparent { fields } = struct_def.kind {
232        walk_list!(vis, visit_field_def, fields);
233    }
234}
235
236pub fn walk_enum_def<'v, V: Visitor<'v>>(vis: &mut V, enum_def: &EnumDef<'v>) {
237    walk_list!(vis, visit_refine_param, enum_def.params);
238    walk_list!(vis, visit_variant, enum_def.variants);
239    walk_list!(vis, visit_expr, enum_def.invariants);
240}
241
242pub fn walk_variant<'v, V: Visitor<'v>>(vis: &mut V, variant: &VariantDef<'v>) {
243    walk_list!(vis, visit_refine_param, variant.params);
244    walk_list!(vis, visit_field_def, variant.fields);
245    vis.visit_variant_ret(&variant.ret);
246}
247
248pub fn walk_field_def<'v, V: Visitor<'v>>(vis: &mut V, field: &FieldDef<'v>) {
249    let FieldDef { ty, lifted: _ } = field;
250    vis.visit_ty(ty);
251}
252
253pub fn walk_variant_ret<'v, V: Visitor<'v>>(vis: &mut V, ret: &VariantRet<'v>) {
254    let VariantRet { idx, enum_id: _ } = ret;
255    vis.visit_expr(idx);
256}
257
258pub fn walk_ty_alias<'v, V: Visitor<'v>>(vis: &mut V, ty_alias: &TyAlias<'v>) {
259    if let Some(index) = &ty_alias.index {
260        vis.visit_refine_param(index);
261    }
262    vis.visit_ty(&ty_alias.ty);
263}
264
265pub fn walk_opaque_ty<'v, V: Visitor<'v>>(vis: &mut V, opaque_ty: &OpaqueTy<'v>) {
266    walk_list!(vis, visit_generic_bound, opaque_ty.bounds);
267}
268
269pub fn walk_generic_bound<'v, V: Visitor<'v>>(vis: &mut V, bound: &GenericBound<'v>) {
270    match bound {
271        GenericBound::Trait(trait_ref) => vis.visit_poly_trait_ref(trait_ref),
272        GenericBound::Outlives(_) => {}
273    }
274}
275
276pub fn walk_poly_trait_ref<'v, V: Visitor<'v>>(vis: &mut V, trait_ref: &PolyTraitRef<'v>) {
277    walk_list!(vis, visit_refine_param, trait_ref.refine_params);
278    vis.visit_path(&trait_ref.trait_ref);
279}
280
281pub fn walk_node<'v, V: Visitor<'v>>(vis: &mut V, node: &OwnerNode<'v>) {
282    match node {
283        OwnerNode::Item(item) => vis.visit_item(item),
284        OwnerNode::TraitItem(trait_item) => vis.visit_trait_item(trait_item),
285        OwnerNode::ImplItem(impl_item) => vis.visit_impl_item(impl_item),
286        OwnerNode::ForeignItem(foreign_item) => vis.visit_foreign_item(foreign_item),
287    }
288}
289
290pub fn walk_item<'v, V: Visitor<'v>>(vis: &mut V, item: &Item<'v>) {
291    vis.visit_generics(&item.generics);
292    match &item.kind {
293        ItemKind::Enum(enum_def) => vis.visit_enum_def(enum_def),
294        ItemKind::Struct(struct_def) => vis.visit_struct_def(struct_def),
295        ItemKind::TyAlias(ty_alias) => vis.visit_ty_alias(ty_alias),
296        ItemKind::Trait(trait_) => {
297            walk_list!(vis, visit_trait_assoc_reft, trait_.assoc_refinements);
298        }
299        ItemKind::Impl(impl_) => vis.visit_impl(impl_),
300        ItemKind::Fn(fn_sig) => vis.visit_fn_sig(fn_sig),
301        ItemKind::Const(info) => {
302            if let Some(expr) = info {
303                vis.visit_expr(expr);
304            }
305        }
306    }
307}
308
309pub fn walk_trait_item<'v, V: Visitor<'v>>(vis: &mut V, trait_item: &TraitItem<'v>) {
310    vis.visit_generics(&trait_item.generics);
311    match &trait_item.kind {
312        TraitItemKind::Fn(fn_sig) => vis.visit_fn_sig(fn_sig),
313        TraitItemKind::Type => {}
314        TraitItemKind::Const => {}
315    }
316}
317
318pub fn walk_impl_item<'v, V: Visitor<'v>>(vis: &mut V, impl_item: &ImplItem<'v>) {
319    vis.visit_generics(&impl_item.generics);
320    match &impl_item.kind {
321        ImplItemKind::Fn(fn_sig) => vis.visit_fn_sig(fn_sig),
322        ImplItemKind::Const => {}
323        ImplItemKind::Type => {}
324    }
325}
326
327pub fn walk_foreign_item<'v, V: Visitor<'v>>(vis: &mut V, impl_item: &ForeignItem<'v>) {
328    match &impl_item.kind {
329        ForeignItemKind::Fn(fn_sig, generics) => {
330            vis.visit_generics(generics);
331            vis.visit_fn_sig(fn_sig);
332        }
333    }
334}
335
336pub fn walk_trait_assoc_reft<'v, V: Visitor<'v>>(vis: &mut V, assoc_reft: &TraitAssocReft<'v>) {
337    walk_list!(vis, visit_refine_param, assoc_reft.params);
338    vis.visit_sort(&assoc_reft.output);
339    if let Some(expr) = &assoc_reft.body {
340        vis.visit_expr(expr);
341    }
342}
343
344pub fn walk_impl_assoc_reft<'v, V: Visitor<'v>>(vis: &mut V, assoc_reft: &ImplAssocReft<'v>) {
345    walk_list!(vis, visit_refine_param, assoc_reft.params);
346    vis.visit_sort(&assoc_reft.output);
347    vis.visit_expr(&assoc_reft.body);
348}
349
350pub fn walk_generics<'v, V: Visitor<'v>>(vis: &mut V, generics: &Generics<'v>) {
351    walk_list!(vis, visit_refine_param, generics.refinement_params);
352    if let Some(predicates) = generics.predicates {
353        walk_list!(vis, visit_where_predicate, predicates);
354    }
355}
356
357pub fn walk_where_predicate<'v, V: Visitor<'v>>(vis: &mut V, predicate: &WhereBoundPredicate<'v>) {
358    vis.visit_ty(&predicate.bounded_ty);
359    walk_list!(vis, visit_generic_bound, predicate.bounds);
360}
361
362pub fn walk_fn_sig<'v, V: Visitor<'v>>(vis: &mut V, sig: &FnSig<'v>) {
363    vis.visit_fn_decl(sig.decl);
364}
365
366pub fn walk_fn_decl<'v, V: Visitor<'v>>(vis: &mut V, decl: &FnDecl<'v>) {
367    walk_list!(vis, visit_requires, decl.requires);
368    walk_list!(vis, visit_ty, decl.inputs);
369    vis.visit_fn_output(&decl.output);
370}
371
372pub fn walk_refine_param<'v, V: Visitor<'v>>(vis: &mut V, param: &RefineParam<'v>) {
373    vis.visit_sort(&param.sort);
374}
375
376pub fn walk_requires<'v, V: Visitor<'v>>(vis: &mut V, requires: &Requires<'v>) {
377    walk_list!(vis, visit_refine_param, requires.params);
378    vis.visit_expr(&requires.pred);
379}
380
381pub fn walk_ensures<'v, V: Visitor<'v>>(vis: &mut V, constraint: &Ensures<'v>) {
382    match constraint {
383        Ensures::Type(loc, ty) => {
384            vis.visit_path_expr(loc);
385            vis.visit_ty(ty);
386        }
387        Ensures::Pred(pred) => {
388            vis.visit_expr(pred);
389        }
390    }
391}
392
393pub fn walk_fn_output<'v, V: Visitor<'v>>(vis: &mut V, output: &FnOutput<'v>) {
394    walk_list!(vis, visit_refine_param, output.params);
395    vis.visit_ty(&output.ret);
396    walk_list!(vis, visit_ensures, output.ensures);
397}
398
399pub fn walk_generic_arg<'v, V: Visitor<'v>>(vis: &mut V, arg: &GenericArg<'v>) {
400    match arg {
401        GenericArg::Lifetime(lft) => vis.visit_lifetime(lft),
402        GenericArg::Type(ty) => vis.visit_ty(ty),
403        GenericArg::Const(_) | GenericArg::Infer => {}
404    }
405}
406
407pub fn walk_ty<'v, V: Visitor<'v>>(vis: &mut V, ty: &Ty<'v>) {
408    match ty.kind {
409        TyKind::BaseTy(bty) => vis.visit_bty(&bty),
410        TyKind::Indexed(bty, idx) => {
411            vis.visit_bty(&bty);
412            vis.visit_expr(&idx);
413        }
414        TyKind::Exists(params, ty) => {
415            walk_list!(vis, visit_refine_param, params);
416            vis.visit_ty(ty);
417        }
418        TyKind::Constr(pred, ty) => {
419            vis.visit_expr(&pred);
420            vis.visit_ty(ty);
421        }
422        TyKind::StrgRef(lft, loc, ty) => {
423            vis.visit_lifetime(&lft);
424            vis.visit_path_expr(loc);
425            vis.visit_ty(ty);
426        }
427        TyKind::Ref(lft, mty) => {
428            vis.visit_lifetime(&lft);
429            vis.visit_ty(mty.ty);
430        }
431        TyKind::BareFn(bare_fn) => vis.visit_fn_decl(bare_fn.decl),
432        TyKind::Tuple(tys) => {
433            walk_list!(vis, visit_ty, tys);
434        }
435        TyKind::Array(ty, _len) => {
436            vis.visit_ty(ty);
437        }
438        TyKind::RawPtr(ty, _mtblt) => {
439            vis.visit_ty(ty);
440        }
441        TyKind::OpaqueDef(opaque_ty) => {
442            vis.visit_opaque_ty(opaque_ty);
443        }
444        TyKind::TraitObject(poly_traits, lft, _) => {
445            walk_list!(vis, visit_poly_trait_ref, poly_traits);
446            vis.visit_lifetime(&lft);
447        }
448        TyKind::Never | TyKind::Infer | TyKind::Err(_) => {}
449    }
450}
451
452pub fn walk_bty<'v, V: Visitor<'v>>(vis: &mut V, bty: &BaseTy<'v>) {
453    match &bty.kind {
454        BaseTyKind::Path(path) => vis.visit_qpath(path),
455        BaseTyKind::Slice(ty) => vis.visit_ty(ty),
456        BaseTyKind::Err(_) => {}
457    }
458}
459
460pub fn walk_qpath<'v, V: Visitor<'v>>(vis: &mut V, qpath: &QPath<'v>) {
461    match qpath {
462        QPath::Resolved(qself, path) => {
463            if let Some(self_ty) = qself {
464                vis.visit_ty(self_ty);
465            }
466            vis.visit_path(path);
467        }
468        QPath::TypeRelative(qself, assoc) => {
469            vis.visit_ty(qself);
470            vis.visit_path_segment(assoc);
471        }
472    }
473}
474
475pub fn walk_path<'v, V: Visitor<'v>>(vis: &mut V, path: &Path<'v>) {
476    walk_list!(vis, visit_path_segment, path.segments);
477    walk_list!(vis, visit_expr, path.refine);
478}
479
480pub fn walk_path_segment<'v, V: Visitor<'v>>(vis: &mut V, segment: &PathSegment<'v>) {
481    walk_list!(vis, visit_generic_arg, segment.args);
482    walk_list!(vis, visit_assoc_item_constraint, segment.constraints);
483}
484
485pub fn walk_assoc_item_constraint<'v, V: Visitor<'v>>(
486    vis: &mut V,
487    constraint: &AssocItemConstraint<'v>,
488) {
489    match &constraint.kind {
490        AssocItemConstraintKind::Equality { term } => {
491            vis.visit_ty(term);
492        }
493    }
494}
495
496pub fn walk_sort<'v, V: Visitor<'v>>(vis: &mut V, sort: &Sort<'v>) {
497    match sort {
498        Sort::Path(path) => vis.visit_sort_path(path),
499        Sort::Func(func) => vis.visit_poly_func_sort(func),
500        Sort::SortOf(bty) => vis.visit_bty(bty),
501        Sort::Loc | Sort::BitVec(_) | Sort::Infer | Sort::Err(_) => {}
502    }
503}
504
505pub fn walk_sort_path<'v, V: Visitor<'v>>(vis: &mut V, path: &SortPath<'v>) {
506    walk_list!(vis, visit_sort, path.args);
507}
508
509pub fn walk_poly_func_sort<'v, V: Visitor<'v>>(vis: &mut V, func: &PolyFuncSort<'v>) {
510    let PolyFuncSort { params: _, fsort } = func;
511    vis.visit_func_sort(fsort);
512}
513
514pub fn walk_func_sort<'v, V: Visitor<'v>>(vis: &mut V, func: &FuncSort<'v>) {
515    walk_list!(vis, visit_sort, func.inputs_and_output);
516}
517
518pub fn walk_alias_reft<'v, V: Visitor<'v>>(vis: &mut V, alias: &AliasReft<'v>) {
519    vis.visit_ty(alias.qself);
520    vis.visit_path(&alias.path);
521}
522
523pub fn walk_field_expr<'v, V: Visitor<'v>>(vis: &mut V, expr: &FieldExpr<'v>) {
524    vis.visit_expr(&expr.expr);
525}
526
527pub fn walk_expr<'v, V: Visitor<'v>>(vis: &mut V, expr: &Expr<'v>) {
528    match expr.kind {
529        ExprKind::Var(path, _) => vis.visit_path_expr(&path),
530        ExprKind::Dot(base, _fld) => {
531            vis.visit_expr(base);
532        }
533        ExprKind::Literal(lit) => vis.visit_literal(&lit),
534        ExprKind::BinaryOp(_op, e1, e2) => {
535            vis.visit_expr(e1);
536            vis.visit_expr(e2);
537        }
538        ExprKind::UnaryOp(_op, e) => vis.visit_expr(e),
539        ExprKind::App(_func, args) => {
540            walk_list!(vis, visit_expr, args);
541        }
542        ExprKind::Alias(alias_reft, args) => {
543            vis.visit_alias_reft(&alias_reft);
544            walk_list!(vis, visit_expr, args);
545        }
546        ExprKind::IfThenElse(e1, e2, e3) => {
547            vis.visit_expr(e1);
548            vis.visit_expr(e2);
549            vis.visit_expr(e3);
550        }
551        ExprKind::Abs(refine_params, body) => {
552            walk_list!(vis, visit_refine_param, refine_params);
553            vis.visit_expr(body);
554        }
555        ExprKind::Record(fields) => {
556            walk_list!(vis, visit_expr, fields);
557        }
558        ExprKind::Constructor(path, exprs, spread) => {
559            if let Some(path) = path {
560                vis.visit_path_expr(&path);
561            }
562            walk_list!(vis, visit_field_expr, exprs);
563            if let Some(s) = spread {
564                vis.visit_expr(&s.expr);
565            }
566        }
567        ExprKind::BoundedQuant(_, param, _, expr) => {
568            vis.visit_refine_param(&param);
569            vis.visit_expr(expr);
570        }
571        ExprKind::Block(decls, body) => {
572            for decl in decls {
573                vis.visit_expr(&decl.init);
574                vis.visit_refine_param(&decl.param);
575            }
576            vis.visit_expr(body);
577        }
578        ExprKind::Err(_) => {}
579    }
580}