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