1use std::ops::DerefMut;
11use std::panic;
12
13use rustc_data_structures::flat_map_in_place::FlatMapInPlace;
14use rustc_span::source_map::Spanned;
15use rustc_span::{Ident, Span};
16use smallvec::{SmallVec, smallvec};
17use thin_vec::ThinVec;
18
19use crate::ast::*;
20use crate::ptr::P;
21use crate::tokenstream::*;
22use crate::visit::{AssocCtxt, BoundKind, FnCtxt, VisitorResult, try_visit, visit_opt, walk_list};
23
24mod sealed {
25 use rustc_ast_ir::visit::VisitorResult;
26
27 pub trait MutVisitorResult {
29 type Result: VisitorResult;
30 }
31
32 impl<T> MutVisitorResult for T {
33 type Result = ();
34 }
35}
36
37use sealed::MutVisitorResult;
38
39super::common_visitor_and_walkers!((mut) MutVisitor);
40
41macro_rules! generate_flat_map_visitor_fns {
42 ($($name:ident, $Ty:ty, $flat_map_fn:ident$(, $param:ident: $ParamTy:ty)*;)+) => {
43 $(
44 fn $name<V: MutVisitor>(
45 vis: &mut V,
46 values: &mut ThinVec<$Ty>,
47 $(
48 $param: $ParamTy,
49 )*
50 ) {
51 values.flat_map_in_place(|value| vis.$flat_map_fn(value$(,$param)*));
52 }
53 )+
54 }
55}
56
57generate_flat_map_visitor_fns! {
58 visit_items, P<Item>, flat_map_item;
59 visit_foreign_items, P<ForeignItem>, flat_map_foreign_item;
60 visit_generic_params, GenericParam, flat_map_generic_param;
61 visit_stmts, Stmt, flat_map_stmt;
62 visit_exprs, P<Expr>, filter_map_expr;
63 visit_expr_fields, ExprField, flat_map_expr_field;
64 visit_pat_fields, PatField, flat_map_pat_field;
65 visit_variants, Variant, flat_map_variant;
66 visit_assoc_items, P<AssocItem>, flat_map_assoc_item, ctxt: AssocCtxt;
67 visit_where_predicates, WherePredicate, flat_map_where_predicate;
68 visit_params, Param, flat_map_param;
69 visit_field_defs, FieldDef, flat_map_field_def;
70 visit_arms, Arm, flat_map_arm;
71}
72
73pub fn walk_flat_map_pat_field<T: MutVisitor>(
74 vis: &mut T,
75 mut fp: PatField,
76) -> SmallVec<[PatField; 1]> {
77 vis.visit_pat_field(&mut fp);
78 smallvec![fp]
79}
80
81fn visit_nested_use_tree<V: MutVisitor>(
82 vis: &mut V,
83 nested_tree: &mut UseTree,
84 nested_id: &mut NodeId,
85) {
86 vis.visit_id(nested_id);
87 vis.visit_use_tree(nested_tree);
88}
89
90macro_rules! generate_walk_flat_map_fns {
91 ($($fn_name:ident($Ty:ty$(,$extra_name:ident: $ExtraTy:ty)*) => $visit_fn_name:ident;)+) => {$(
92 pub fn $fn_name<V: MutVisitor>(vis: &mut V, mut value: $Ty$(,$extra_name: $ExtraTy)*) -> SmallVec<[$Ty; 1]> {
93 vis.$visit_fn_name(&mut value$(,$extra_name)*);
94 smallvec![value]
95 }
96 )+};
97}
98
99generate_walk_flat_map_fns! {
100 walk_flat_map_arm(Arm) => visit_arm;
101 walk_flat_map_variant(Variant) => visit_variant;
102 walk_flat_map_param(Param) => visit_param;
103 walk_flat_map_generic_param(GenericParam) => visit_generic_param;
104 walk_flat_map_where_predicate(WherePredicate) => visit_where_predicate;
105 walk_flat_map_field_def(FieldDef) => visit_field_def;
106 walk_flat_map_expr_field(ExprField) => visit_expr_field;
107 walk_flat_map_item(P<Item>) => visit_item;
108 walk_flat_map_foreign_item(P<ForeignItem>) => visit_foreign_item;
109 walk_flat_map_assoc_item(P<AssocItem>, ctxt: AssocCtxt) => visit_assoc_item;
110}
111
112fn walk_ty_alias_where_clauses<T: MutVisitor>(vis: &mut T, tawcs: &mut TyAliasWhereClauses) {
113 let TyAliasWhereClauses { before, after, split: _ } = tawcs;
114 let TyAliasWhereClause { has_where_token: _, span: span_before } = before;
115 let TyAliasWhereClause { has_where_token: _, span: span_after } = after;
116 vis.visit_span(span_before);
117 vis.visit_span(span_after);
118}
119
120pub fn walk_filter_map_expr<T: MutVisitor>(vis: &mut T, mut e: P<Expr>) -> Option<P<Expr>> {
121 vis.visit_expr(&mut e);
122 Some(e)
123}
124
125pub fn walk_flat_map_stmt<T: MutVisitor>(
126 vis: &mut T,
127 Stmt { kind, span, mut id }: Stmt,
128) -> SmallVec<[Stmt; 1]> {
129 vis.visit_id(&mut id);
130 let mut stmts: SmallVec<[Stmt; 1]> = walk_flat_map_stmt_kind(vis, kind)
131 .into_iter()
132 .map(|kind| Stmt { id, kind, span })
133 .collect();
134 match &mut stmts[..] {
135 [] => {}
136 [stmt] => vis.visit_span(&mut stmt.span),
137 _ => panic!(
138 "cloning statement `NodeId`s is prohibited by default, \
139 the visitor should implement custom statement visiting"
140 ),
141 }
142 stmts
143}
144
145fn walk_flat_map_stmt_kind<T: MutVisitor>(vis: &mut T, kind: StmtKind) -> SmallVec<[StmtKind; 1]> {
146 match kind {
147 StmtKind::Let(mut local) => smallvec![StmtKind::Let({
148 vis.visit_local(&mut local);
149 local
150 })],
151 StmtKind::Item(item) => vis.flat_map_item(item).into_iter().map(StmtKind::Item).collect(),
152 StmtKind::Expr(expr) => vis.filter_map_expr(expr).into_iter().map(StmtKind::Expr).collect(),
153 StmtKind::Semi(expr) => vis.filter_map_expr(expr).into_iter().map(StmtKind::Semi).collect(),
154 StmtKind::Empty => smallvec![StmtKind::Empty],
155 StmtKind::MacCall(mut mac) => {
156 let MacCallStmt { mac: mac_, style: _, attrs, tokens: _ } = mac.deref_mut();
157 for attr in attrs {
158 vis.visit_attribute(attr);
159 }
160 vis.visit_mac_call(mac_);
161 smallvec![StmtKind::MacCall(mac)]
162 }
163 }
164}