Use SmallVec for TypeWalker's stack.

The change also adds the missing `SmallVec::truncate` method.

3 files changed, 27 insertions, 9 deletions

diff --git a/src/librustc/ty/mod.rs b/src/librustc/ty/mod.rs
index e94e93158c4..ebf8c96def7 100644
--- a/src/librustc/ty/mod.rs
+++ b/src/librustc/ty/mod.rs
@@ -48,6 +48,7 @@ use syntax::symbol::{Symbol, InternedString};
 use syntax_pos::{DUMMY_SP, Span};
 
 use rustc_const_math::ConstInt;
+use rustc_data_structures::accumulate_vec::IntoIter as AccIntoIter;
 
 use hir;
 use hir::itemlikevisit::ItemLikeVisitor;
@@ -1887,7 +1888,7 @@ impl<'tcx> TyS<'tcx> {
     /// Iterator that walks the immediate children of `self`.  Hence
     /// `Foo<Bar<i32>, u32>` yields the sequence `[Bar<i32>, u32]`
     /// (but not `i32`, like `walk`).
-    pub fn walk_shallow(&'tcx self) -> IntoIter<Ty<'tcx>> {
+    pub fn walk_shallow(&'tcx self) -> AccIntoIter<walk::TypeWalkerArray<'tcx>> {
         walk::walk_shallow(self)
     }
 
diff --git a/src/librustc/ty/walk.rs b/src/librustc/ty/walk.rs
index a6ecfd2fb70..2f9468dbe58 100644
--- a/src/librustc/ty/walk.rs
+++ b/src/librustc/ty/walk.rs
@@ -12,17 +12,22 @@
 //! WARNING: this does not keep track of the region depth.
 
 use ty::{self, Ty};
-use std::iter::Iterator;
-use std::vec::IntoIter;
+use rustc_data_structures::small_vec::SmallVec;
+use rustc_data_structures::accumulate_vec::IntoIter as AccIntoIter;
+
+// The TypeWalker's stack is hot enough that it's worth going to some effort to
+// avoid heap allocations.
+pub type TypeWalkerArray<'tcx> = [Ty<'tcx>; 8];
+pub type TypeWalkerStack<'tcx> = SmallVec<TypeWalkerArray<'tcx>>;
 
 pub struct TypeWalker<'tcx> {
-    stack: Vec<Ty<'tcx>>,
+    stack: TypeWalkerStack<'tcx>,
     last_subtree: usize,
 }
 
 impl<'tcx> TypeWalker<'tcx> {
     pub fn new(ty: Ty<'tcx>) -> TypeWalker<'tcx> {
-        TypeWalker { stack: vec![ty], last_subtree: 1, }
+        TypeWalker { stack: SmallVec::one(ty), last_subtree: 1, }
     }
 
     /// Skips the subtree of types corresponding to the last type
@@ -61,8 +66,8 @@ impl<'tcx> Iterator for TypeWalker<'tcx> {
     }
 }
 
-pub fn walk_shallow<'tcx>(ty: Ty<'tcx>) -> IntoIter<Ty<'tcx>> {
-    let mut stack = vec![];
+pub fn walk_shallow<'tcx>(ty: Ty<'tcx>) -> AccIntoIter<TypeWalkerArray<'tcx>> {
+    let mut stack = SmallVec::new();
     push_subtypes(&mut stack, ty);
     stack.into_iter()
 }
@@ -73,7 +78,7 @@ pub fn walk_shallow<'tcx>(ty: Ty<'tcx>) -> IntoIter<Ty<'tcx>> {
 // known to be significant to any code, but it seems like the
 // natural order one would expect (basically, the order of the
 // types as they are written).
-fn push_subtypes<'tcx>(stack: &mut Vec<Ty<'tcx>>, parent_ty: Ty<'tcx>) {
+fn push_subtypes<'tcx>(stack: &mut TypeWalkerStack<'tcx>, parent_ty: Ty<'tcx>) {
     match parent_ty.sty {
         ty::TyBool | ty::TyChar | ty::TyInt(_) | ty::TyUint(_) | ty::TyFloat(_) |
         ty::TyStr | ty::TyInfer(_) | ty::TyParam(_) | ty::TyNever | ty::TyError => {
@@ -112,7 +117,7 @@ fn push_subtypes<'tcx>(stack: &mut Vec<Ty<'tcx>>, parent_ty: Ty<'tcx>) {
     }
 }
 
-fn push_sig_subtypes<'tcx>(stack: &mut Vec<Ty<'tcx>>, sig: &ty::PolyFnSig<'tcx>) {
+fn push_sig_subtypes<'tcx>(stack: &mut TypeWalkerStack<'tcx>, sig: &ty::PolyFnSig<'tcx>) {
     stack.push(sig.0.output);
     stack.extend(sig.0.inputs.iter().cloned().rev());
 }
diff --git a/src/librustc_data_structures/small_vec.rs b/src/librustc_data_structures/small_vec.rs
index 565a3c443a3..4e2b3786021 100644
--- a/src/librustc_data_structures/small_vec.rs
+++ b/src/librustc_data_structures/small_vec.rs
@@ -130,6 +130,18 @@ impl<A: Array> SmallVec<A> {
             self.set_len(len + 1);
         }
     }
+
+    pub fn truncate(&mut self, len: usize) {
+        unsafe {
+            while len < self.len() {
+                // Decrement len before the drop_in_place(), so a panic on Drop
+                // doesn't re-drop the just-failed value.
+                let newlen = self.len() - 1;
+                self.set_len(newlen);
+                ::std::ptr::drop_in_place(self.get_unchecked_mut(newlen));
+            }
+        }
+    }
 }
 
 impl<A: Array> Deref for SmallVec<A> {