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::{self, 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 /// We push `GenericArg`s on the stack in reverse order so as to
 116 /// maintain a pre-order traversal. As of the time of this
 117 /// writing, the fact that the traversal is pre-order is not
 118 /// known to be significant to any code, but it seems like the
 119 /// natural order one would expect (basically, the order of the
 120 /// types as they are written).
 121 fn push_inner<'tcx>(stack: &mut TypeWalkerStack<'tcx>, parent: GenericArg<'tcx>) {
 122     match parent.unpack() {
 123         GenericArgKind::Type(parent_ty) => match *parent_ty.kind() {
 124             ty::Bool
 125             | ty::Char
 126             | ty::Int(_)
 127             | ty::Uint(_)
 128             | ty::Float(_)
 129             | ty::Str
 130             | ty::Infer(_)
 131             | ty::Param(_)
 132             | ty::Never
 133             | ty::Error(_)
 134             | ty::Placeholder(..)
 135             | ty::Bound(..)
 136             | ty::Foreign(..) => {}
 137
 138             ty::Array(ty, len) => {
 139                 stack.push(len.into());
 140                 stack.push(ty.into());
 141             }
 142             ty::Slice(ty) => {
 143                 stack.push(ty.into());
 144             }
 145             ty::RawPtr(mt) => {
 146                 stack.push(mt.ty.into());
 147             }
 148             ty::Ref(lt, ty, _) => {
 149                 stack.push(ty.into());
 150                 stack.push(lt.into());
 151             }
 152             ty::Projection(data) => {
 153                 stack.extend(data.substs.iter().rev());
 154             }
 155             ty::Dynamic(obj, lt) => {
 156                 stack.push(lt.into());
 157                 stack.extend(obj.iter().rev().flat_map(|predicate| {
 158                     let (substs, opt_ty) = match predicate.skip_binder() {
 159                         ty::ExistentialPredicate::Trait(tr) => (tr.substs, None),
 160                         ty::ExistentialPredicate::Projection(p) => (p.substs, Some(p.term)),
 161                         ty::ExistentialPredicate::AutoTrait(_) =>
 162                         // Empty iterator
 163                         {
 164                             (ty::InternalSubsts::empty(), None)
 165                         }
 166                     };
 167
 168                     substs.iter().rev().chain(opt_ty.map(|term| match term {
 169                         ty::Term::Ty(ty) => ty.into(),
 170                         ty::Term::Const(ct) => ct.into(),
 171                     }))
 172                 }));
 173             }
 174             ty::Adt(_, substs)
 175             | ty::Opaque(_, substs)
 176             | ty::Closure(_, substs)
 177             | ty::Generator(_, substs, _)
 178             | ty::FnDef(_, substs) => {
 179                 stack.extend(substs.iter().rev());
 180             }
 181             ty::Tuple(ts) => stack.extend(ts.as_substs().iter().rev()),
 182             ty::GeneratorWitness(ts) => {
 183                 stack.extend(ts.skip_binder().iter().rev().map(|ty| ty.into()));
 184             }
 185             ty::FnPtr(sig) => {
 186                 stack.push(sig.skip_binder().output().into());
 187                 stack.extend(sig.skip_binder().inputs().iter().copied().rev().map(|ty| ty.into()));
 188             }
 189         },
 190         GenericArgKind::Lifetime(_) => {}
 191         GenericArgKind::Const(parent_ct) => {
 192             stack.push(parent_ct.ty().into());
 193             match parent_ct.kind() {
 194                 ty::ConstKind::Infer(_)
 195                 | ty::ConstKind::Param(_)
 196                 | ty::ConstKind::Placeholder(_)
 197                 | ty::ConstKind::Bound(..)
 198                 | ty::ConstKind::Value(_)
 199                 | ty::ConstKind::Error(_) => {}
 200
 201                 ty::ConstKind::Unevaluated(ct) => {
 202                     stack.extend(ct.substs.iter().rev());
 203                 }
 204             }
 205         }
 206     }
 207 }