compiler/rustc_middle/src/ty/walk.rs

   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> {
  28     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     /// ```
  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
  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     /// ```
  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.
  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     /// ```
 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     /// ```
 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).
 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::FnDef(_, substs) => {
 194                 stack.extend(substs.iter().rev());
 195             }
 196             ty::Tuple(ts) => stack.extend(ts.as_substs().iter().rev()),
 197             ty::GeneratorWitness(ts) => {
 198                 stack.extend(ts.skip_binder().iter().rev().map(|ty| ty.into()));
 199             }
 200             ty::FnPtr(sig) => {
 201                 stack.push(sig.skip_binder().output().into());
 202                 stack.extend(sig.skip_binder().inputs().iter().copied().rev().map(|ty| ty.into()));
 203             }
 204         },
 205         GenericArgKind::Lifetime(_) => {}
 206         GenericArgKind::Const(parent_ct) => {
 207             stack.push(parent_ct.ty().into());
 208             match parent_ct.kind() {
 209                 ty::ConstKind::Infer(_)
 210                 | ty::ConstKind::Param(_)
 211                 | ty::ConstKind::Placeholder(_)
 212                 | ty::ConstKind::Bound(..)
 213                 | ty::ConstKind::Value(_)
 214                 | ty::ConstKind::Error(_) => {}
 215
 216                 ty::ConstKind::Expr(expr) => match expr {
 217                     ty::Expr::UnOp(_, v) => push_inner(stack, v.into()),
 218                     ty::Expr::Binop(_, l, r) => {
 219                         push_inner(stack, r.into());
 220                         push_inner(stack, l.into())
 221                     }
 222                     ty::Expr::FunctionCall(func, args) => {
 223                         for a in args.iter().rev() {
 224                             push_inner(stack, a.into());
 225                         }
 226                         push_inner(stack, func.into());
 227                     }
 228                     ty::Expr::Cast(_, c, t) => {
 229                         push_inner(stack, t.into());
 230                         push_inner(stack, c.into());
 231                     }
 232                 },
 233
 234                 ty::ConstKind::Unevaluated(ct) => {
 235                     stack.extend(ct.substs.iter().rev());
 236                 }
 237             }
 238         }
 239     }
 240 }