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