1 //! An iterator over the type substructure.
2 //! WARNING: this does not keep track of the region depth.
3
4 use crate::ty::subst::{GenericArg, GenericArgKind};
5 use crate::ty::{self, Ty};
6 use rustc_data_structures::sso::SsoHashSet;
7 use smallvec::SmallVec;
8
9 // The TypeWalker's stack is hot enough that it's worth going to some effort to
10 // avoid heap allocations.
11 type TypeWalkerStack<'tcx> = SmallVec<[GenericArg<'tcx>; 8]>;
12
13 pub struct TypeWalker<'tcx> {
14 stack: TypeWalkerStack<'tcx>,
15 last_subtree: usize,
16 pub visited: SsoHashSet<GenericArg<'tcx>>,
17 }
18
19 /// An iterator for walking the type tree.
20 ///
21 /// It's very easy to produce a deeply
22 /// nested type tree with a lot of
23 /// identical subtrees. In order to work efficiently
24 /// in this situation walker only visits each type once.
25 /// It maintains a set of visited types and
26 /// skips any types that are already there.
27 impl<'tcx> TypeWalker<'tcx> {
new(root: GenericArg<'tcx>) -> Self28 pub fn new(root: GenericArg<'tcx>) -> Self {
29 Self { stack: smallvec![root], last_subtree: 1, visited: SsoHashSet::new() }
30 }
31
32 /// Skips the subtree corresponding to the last type
33 /// returned by `next()`.
34 ///
35 /// Example: Imagine you are walking `Foo<Bar<i32>, usize>`.
36 ///
37 /// ```ignore (illustrative)
38 /// let mut iter: TypeWalker = ...;
39 /// iter.next(); // yields Foo
40 /// iter.next(); // yields Bar<i32>
41 /// iter.skip_current_subtree(); // skips i32
42 /// iter.next(); // yields usize
43 /// ```
skip_current_subtree(&mut self)44 pub fn skip_current_subtree(&mut self) {
45 self.stack.truncate(self.last_subtree);
46 }
47 }
48
49 impl<'tcx> Iterator for TypeWalker<'tcx> {
50 type Item = GenericArg<'tcx>;
51
next(&mut self) -> Option<GenericArg<'tcx>>52 fn next(&mut self) -> Option<GenericArg<'tcx>> {
53 debug!("next(): stack={:?}", self.stack);
54 loop {
55 let next = self.stack.pop()?;
56 self.last_subtree = self.stack.len();
57 if self.visited.insert(next) {
58 push_inner(&mut self.stack, next);
59 debug!("next: stack={:?}", self.stack);
60 return Some(next);
61 }
62 }
63 }
64 }
65
66 impl<'tcx> GenericArg<'tcx> {
67 /// Iterator that walks `self` and any types reachable from
68 /// `self`, in depth-first order. Note that just walks the types
69 /// that appear in `self`, it does not descend into the fields of
70 /// structs or variants. For example:
71 ///
72 /// ```text
73 /// isize => { isize }
74 /// Foo<Bar<isize>> => { Foo<Bar<isize>>, Bar<isize>, isize }
75 /// [isize] => { [isize], isize }
76 /// ```
walk(self) -> TypeWalker<'tcx>77 pub fn walk(self) -> TypeWalker<'tcx> {
78 TypeWalker::new(self)
79 }
80
81 /// Iterator that walks the immediate children of `self`. Hence
82 /// `Foo<Bar<i32>, u32>` yields the sequence `[Bar<i32>, u32]`
83 /// (but not `i32`, like `walk`).
84 ///
85 /// Iterator only walks items once.
86 /// It accepts visited set, updates it with all visited types
87 /// and skips any types that are already there.
walk_shallow( self, visited: &mut SsoHashSet<GenericArg<'tcx>>, ) -> impl Iterator<Item = GenericArg<'tcx>>88 pub fn walk_shallow(
89 self,
90 visited: &mut SsoHashSet<GenericArg<'tcx>>,
91 ) -> impl Iterator<Item = GenericArg<'tcx>> {
92 let mut stack = SmallVec::new();
93 push_inner(&mut stack, self);
94 stack.retain(|a| visited.insert(*a));
95 stack.into_iter()
96 }
97 }
98
99 impl<'tcx> Ty<'tcx> {
100 /// Iterator that walks `self` and any types reachable from
101 /// `self`, in depth-first order. Note that just walks the types
102 /// that appear in `self`, it does not descend into the fields of
103 /// structs or variants. For example:
104 ///
105 /// ```text
106 /// isize => { isize }
107 /// Foo<Bar<isize>> => { Foo<Bar<isize>>, Bar<isize>, isize }
108 /// [isize] => { [isize], isize }
109 /// ```
walk(self) -> TypeWalker<'tcx>110 pub fn walk(self) -> TypeWalker<'tcx> {
111 TypeWalker::new(self.into())
112 }
113 }
114
115 impl<'tcx> ty::Const<'tcx> {
116 /// Iterator that walks `self` and any types reachable from
117 /// `self`, in depth-first order. Note that just walks the types
118 /// that appear in `self`, it does not descend into the fields of
119 /// structs or variants. For example:
120 ///
121 /// ```text
122 /// isize => { isize }
123 /// Foo<Bar<isize>> => { Foo<Bar<isize>>, Bar<isize>, isize }
124 /// [isize] => { [isize], isize }
125 /// ```
walk(self) -> TypeWalker<'tcx>126 pub fn walk(self) -> TypeWalker<'tcx> {
127 TypeWalker::new(self.into())
128 }
129 }
130
131 /// We push `GenericArg`s on the stack in reverse order so as to
132 /// maintain a pre-order traversal. As of the time of this
133 /// writing, the fact that the traversal is pre-order is not
134 /// known to be significant to any code, but it seems like the
135 /// natural order one would expect (basically, the order of the
136 /// types as they are written).
push_inner<'tcx>(stack: &mut TypeWalkerStack<'tcx>, parent: GenericArg<'tcx>)137 fn push_inner<'tcx>(stack: &mut TypeWalkerStack<'tcx>, parent: GenericArg<'tcx>) {
138 match parent.unpack() {
139 GenericArgKind::Type(parent_ty) => match *parent_ty.kind() {
140 ty::Bool
141 | ty::Char
142 | ty::Int(_)
143 | ty::Uint(_)
144 | ty::Float(_)
145 | ty::Str
146 | ty::Infer(_)
147 | ty::Param(_)
148 | ty::Never
149 | ty::Error(_)
150 | ty::Placeholder(..)
151 | ty::Bound(..)
152 | ty::Foreign(..) => {}
153
154 ty::Array(ty, len) => {
155 stack.push(len.into());
156 stack.push(ty.into());
157 }
158 ty::Slice(ty) => {
159 stack.push(ty.into());
160 }
161 ty::RawPtr(mt) => {
162 stack.push(mt.ty.into());
163 }
164 ty::Ref(lt, ty, _) => {
165 stack.push(ty.into());
166 stack.push(lt.into());
167 }
168 ty::Alias(_, data) => {
169 stack.extend(data.substs.iter().rev());
170 }
171 ty::Dynamic(obj, lt, _) => {
172 stack.push(lt.into());
173 stack.extend(obj.iter().rev().flat_map(|predicate| {
174 let (substs, opt_ty) = match predicate.skip_binder() {
175 ty::ExistentialPredicate::Trait(tr) => (tr.substs, None),
176 ty::ExistentialPredicate::Projection(p) => (p.substs, Some(p.term)),
177 ty::ExistentialPredicate::AutoTrait(_) =>
178 // Empty iterator
179 {
180 (ty::InternalSubsts::empty(), None)
181 }
182 };
183
184 substs.iter().rev().chain(opt_ty.map(|term| match term.unpack() {
185 ty::TermKind::Ty(ty) => ty.into(),
186 ty::TermKind::Const(ct) => ct.into(),
187 }))
188 }));
189 }
190 ty::Adt(_, substs)
191 | ty::Closure(_, substs)
192 | ty::Generator(_, substs, _)
193 | ty::GeneratorWitnessMIR(_, substs)
194 | ty::FnDef(_, substs) => {
195 stack.extend(substs.iter().rev());
196 }
197 ty::Tuple(ts) => stack.extend(ts.iter().rev().map(GenericArg::from)),
198 ty::GeneratorWitness(ts) => {
199 stack.extend(ts.skip_binder().iter().rev().map(|ty| ty.into()));
200 }
201 ty::FnPtr(sig) => {
202 stack.push(sig.skip_binder().output().into());
203 stack.extend(sig.skip_binder().inputs().iter().copied().rev().map(|ty| ty.into()));
204 }
205 },
206 GenericArgKind::Lifetime(_) => {}
207 GenericArgKind::Const(parent_ct) => {
208 stack.push(parent_ct.ty().into());
209 match parent_ct.kind() {
210 ty::ConstKind::Infer(_)
211 | ty::ConstKind::Param(_)
212 | ty::ConstKind::Placeholder(_)
213 | ty::ConstKind::Bound(..)
214 | ty::ConstKind::Value(_)
215 | ty::ConstKind::Error(_) => {}
216
217 ty::ConstKind::Expr(expr) => match expr {
218 ty::Expr::UnOp(_, v) => push_inner(stack, v.into()),
219 ty::Expr::Binop(_, l, r) => {
220 push_inner(stack, r.into());
221 push_inner(stack, l.into())
222 }
223 ty::Expr::FunctionCall(func, args) => {
224 for a in args.iter().rev() {
225 push_inner(stack, a.into());
226 }
227 push_inner(stack, func.into());
228 }
229 ty::Expr::Cast(_, c, t) => {
230 push_inner(stack, t.into());
231 push_inner(stack, c.into());
232 }
233 },
234
235 ty::ConstKind::Unevaluated(ct) => {
236 stack.extend(ct.substs.iter().rev());
237 }
238 }
239 }
240 }
241 }
242