rustc_type_ir/
walk.rs

1//! An iterator over the type substructure.
2//! WARNING: this does not keep track of the region depth.
3
4use smallvec::{SmallVec, smallvec};
5use tracing::debug;
6
7use crate::data_structures::SsoHashSet;
8use crate::inherent::*;
9use crate::{self as ty, Interner};
10
11// The TypeWalker's stack is hot enough that it's worth going to some effort to
12// avoid heap allocations.
13type TypeWalkerStack<I> = SmallVec<[<I as Interner>::GenericArg; 8]>;
14
15pub struct TypeWalker<I: Interner> {
16    stack: TypeWalkerStack<I>,
17    last_subtree: usize,
18    pub visited: SsoHashSet<I::GenericArg>,
19}
20
21/// An iterator for walking the type tree.
22///
23/// It's very easy to produce a deeply
24/// nested type tree with a lot of
25/// identical subtrees. In order to work efficiently
26/// in this situation walker only visits each type once.
27/// It maintains a set of visited types and
28/// skips any types that are already there.
29impl<I: Interner> TypeWalker<I> {
30    pub fn new(root: I::GenericArg) -> Self {
31        Self { stack: smallvec![root], last_subtree: 1, visited: SsoHashSet::new() }
32    }
33
34    /// Skips the subtree corresponding to the last type
35    /// returned by `next()`.
36    ///
37    /// Example: Imagine you are walking `Foo<Bar<i32>, usize>`.
38    ///
39    /// ```ignore (illustrative)
40    /// let mut iter: TypeWalker = ...;
41    /// iter.next(); // yields Foo
42    /// iter.next(); // yields Bar<i32>
43    /// iter.skip_current_subtree(); // skips i32
44    /// iter.next(); // yields usize
45    /// ```
46    pub fn skip_current_subtree(&mut self) {
47        self.stack.truncate(self.last_subtree);
48    }
49}
50
51impl<I: Interner> Iterator for TypeWalker<I> {
52    type Item = I::GenericArg;
53
54    fn next(&mut self) -> Option<I::GenericArg> {
55        debug!("next(): stack={:?}", self.stack);
56        loop {
57            let next = self.stack.pop()?;
58            self.last_subtree = self.stack.len();
59            if self.visited.insert(next) {
60                push_inner::<I>(&mut self.stack, next);
61                debug!("next: stack={:?}", self.stack);
62                return Some(next);
63            }
64        }
65    }
66}
67
68/// We push `GenericArg`s on the stack in reverse order so as to
69/// maintain a pre-order traversal. As of the time of this
70/// writing, the fact that the traversal is pre-order is not
71/// known to be significant to any code, but it seems like the
72/// natural order one would expect (basically, the order of the
73/// types as they are written).
74fn push_inner<I: Interner>(stack: &mut TypeWalkerStack<I>, parent: I::GenericArg) {
75    match parent.kind() {
76        ty::GenericArgKind::Type(parent_ty) => match parent_ty.kind() {
77            ty::Bool
78            | ty::Char
79            | ty::Int(_)
80            | ty::Uint(_)
81            | ty::Float(_)
82            | ty::Str
83            | ty::Infer(_)
84            | ty::Param(_)
85            | ty::Never
86            | ty::Error(_)
87            | ty::Placeholder(..)
88            | ty::Bound(..)
89            | ty::Foreign(..) => {}
90
91            ty::Pat(ty, pat) => {
92                push_ty_pat::<I>(stack, pat);
93                stack.push(ty.into());
94            }
95            ty::Array(ty, len) => {
96                stack.push(len.into());
97                stack.push(ty.into());
98            }
99            ty::Slice(ty) => {
100                stack.push(ty.into());
101            }
102            ty::RawPtr(ty, _) => {
103                stack.push(ty.into());
104            }
105            ty::Ref(lt, ty, _) => {
106                stack.push(ty.into());
107                stack.push(lt.into());
108            }
109            ty::Alias(_, data) => {
110                stack.extend(data.args.iter().rev());
111            }
112            ty::Dynamic(obj, lt, _) => {
113                stack.push(lt.into());
114                stack.extend(
115                    obj.iter()
116                        .rev()
117                        .filter_map(|predicate| {
118                            let (args, opt_ty) = match predicate.skip_binder() {
119                                ty::ExistentialPredicate::Trait(tr) => (tr.args, None),
120                                ty::ExistentialPredicate::Projection(p) => (p.args, Some(p.term)),
121                                ty::ExistentialPredicate::AutoTrait(_) => {
122                                    return None;
123                                }
124                            };
125
126                            Some(args.iter().rev().chain(opt_ty.map(|term| match term.kind() {
127                                ty::TermKind::Ty(ty) => ty.into(),
128                                ty::TermKind::Const(ct) => ct.into(),
129                            })))
130                        })
131                        .flatten(),
132                );
133            }
134            ty::Adt(_, args)
135            | ty::Closure(_, args)
136            | ty::CoroutineClosure(_, args)
137            | ty::Coroutine(_, args)
138            | ty::CoroutineWitness(_, args)
139            | ty::FnDef(_, args) => {
140                stack.extend(args.iter().rev());
141            }
142            ty::Tuple(ts) => stack.extend(ts.iter().rev().map(|ty| ty.into())),
143            ty::FnPtr(sig_tys, _hdr) => {
144                stack.extend(
145                    sig_tys.skip_binder().inputs_and_output.iter().rev().map(|ty| ty.into()),
146                );
147            }
148            ty::UnsafeBinder(bound_ty) => {
149                stack.push(bound_ty.skip_binder().into());
150            }
151        },
152        ty::GenericArgKind::Lifetime(_) => {}
153        ty::GenericArgKind::Const(parent_ct) => match parent_ct.kind() {
154            ty::ConstKind::Infer(_)
155            | ty::ConstKind::Param(_)
156            | ty::ConstKind::Placeholder(_)
157            | ty::ConstKind::Bound(..)
158            | ty::ConstKind::Error(_) => {}
159
160            ty::ConstKind::Value(cv) => stack.push(cv.ty().into()),
161
162            ty::ConstKind::Expr(expr) => stack.extend(expr.args().iter().rev()),
163            ty::ConstKind::Unevaluated(ct) => {
164                stack.extend(ct.args.iter().rev());
165            }
166        },
167    }
168}
169
170fn push_ty_pat<I: Interner>(stack: &mut TypeWalkerStack<I>, pat: I::Pat) {
171    match pat.kind() {
172        ty::PatternKind::Range { start, end } => {
173            stack.push(end.into());
174            stack.push(start.into());
175        }
176        ty::PatternKind::Or(pats) => {
177            for pat in pats.iter() {
178                push_ty_pat::<I>(stack, pat)
179            }
180        }
181    }
182}
rustc_type_ir/walk.rs

rustc_type_ir/
walk.rs
