Implement an iterator for walking types rather than the old callback code.

4 files changed, 218 insertions, 48 deletions

diff --git a/src/librustc/lib.rs b/src/librustc/lib.rs
index 4647c92e3d1..3d99a880f68 100644
--- a/src/librustc/lib.rs
+++ b/src/librustc/lib.rs
@@ -98,6 +98,7 @@ pub mod middle {
     pub mod traits;
     pub mod ty;
     pub mod ty_fold;
+    pub mod ty_walk;
     pub mod weak_lang_items;
 }
 
diff --git a/src/librustc/middle/ty.rs b/src/librustc/middle/ty.rs
index 8e7470f5084..c84217d956f 100644
--- a/src/librustc/middle/ty.rs
+++ b/src/librustc/middle/ty.rs
@@ -59,6 +59,7 @@ use middle::subst::{mod, Subst, Substs, VecPerParamSpace};
 use middle::traits;
 use middle::ty;
 use middle::ty_fold::{mod, TypeFoldable, TypeFolder};
+use middle::ty_walk::TypeWalker;
 use util::ppaux::{note_and_explain_region, bound_region_ptr_to_string};
 use util::ppaux::{trait_store_to_string, ty_to_string};
 use util::ppaux::{Repr, UserString};
@@ -2806,55 +2807,59 @@ pub fn mk_param_from_def<'tcx>(cx: &ctxt<'tcx>, def: &TypeParameterDef) -> Ty<'t
 
 pub fn mk_open<'tcx>(cx: &ctxt<'tcx>, ty: Ty<'tcx>) -> Ty<'tcx> { mk_t(cx, ty_open(ty)) }
 
-pub fn walk_ty<'tcx, F>(ty: Ty<'tcx>, mut f: F) where
-    F: FnMut(Ty<'tcx>),
+impl<'tcx> TyS<'tcx> {
+    /// Iterator that walks `self` and any types reachable from
+    /// `self`, in depth-first order. Note that just walks the types
+    /// that appear in `self`, it does not descend into the fields of
+    /// structs or variants. For example:
+    ///
+    /// ```notrust
+    /// int => { int }
+    /// Foo<Bar<int>> => { Foo<Bar<int>>, Bar<int>, int }
+    /// [int] => { [int], int }
+    /// ```
+    pub fn walk(&'tcx self) -> TypeWalker<'tcx> {
+        TypeWalker::new(self)
+    }
+
+    /// Iterator that walks types reachable from `self`, in
+    /// depth-first order. Note that this is a shallow walk. For
+    /// example:
+    ///
+    /// ```notrust
+    /// int => { }
+    /// Foo<Bar<int>> => { Bar<int>, int }
+    /// [int] => { int }
+    /// ```
+    pub fn walk_children(&'tcx self) -> TypeWalker<'tcx> {
+        // Walks type reachable from `self` but not `self
+        let mut walker = self.walk();
+        let r = walker.next();
+        assert_eq!(r, Some(self));
+        walker
+    }
+}
+
+pub fn walk_ty<'tcx, F>(ty_root: Ty<'tcx>, mut f: F)
+    where F: FnMut(Ty<'tcx>),
 {
-    maybe_walk_ty(ty, |ty| { f(ty); true });
+    for ty in ty_root.walk() {
+        f(ty);
+    }
 }
 
-pub fn maybe_walk_ty<'tcx, F>(ty: Ty<'tcx>, mut f: F) where F: FnMut(Ty<'tcx>) -> bool {
-    // FIXME(#19596) This is a workaround, but there should be a better way to do this
-    fn maybe_walk_ty_<'tcx, F>(ty: Ty<'tcx>, f: &mut F) where F: FnMut(Ty<'tcx>) -> bool {
-        if !(*f)(ty) {
-            return;
-        }
-        match ty.sty {
-            ty_bool | ty_char | ty_int(_) | ty_uint(_) | ty_float(_) |
-            ty_str | ty_infer(_) | ty_param(_) | ty_err => {}
-            ty_uniq(ty) | ty_vec(ty, _) | ty_open(ty) => maybe_walk_ty_(ty, f),
-            ty_ptr(ref tm) | ty_rptr(_, ref tm) => {
-                maybe_walk_ty_(tm.ty, f);
-            }
-            ty_trait(box TyTrait { ref principal, .. }) => {
-                for subty in principal.0.substs.types.iter() {
-                    maybe_walk_ty_(*subty, f);
-                }
-            }
-            ty_projection(ProjectionTy { ref trait_ref, .. }) => {
-                for subty in trait_ref.substs.types.iter() {
-                    maybe_walk_ty_(*subty, f);
-                }
-            }
-            ty_enum(_, ref substs) |
-            ty_struct(_, ref substs) |
-            ty_unboxed_closure(_, _, ref substs) => {
-                for subty in substs.types.iter() {
-                    maybe_walk_ty_(*subty, f);
-                }
-            }
-            ty_tup(ref ts) => { for tt in ts.iter() { maybe_walk_ty_(*tt, f); } }
-            ty_bare_fn(_, ref ft) => {
-                for a in ft.sig.0.inputs.iter() { maybe_walk_ty_(*a, f); }
-                if let ty::FnConverging(output) = ft.sig.0.output {
-                    maybe_walk_ty_(output, f);
-                }
-            }
-            ty_closure(ref ft) => {
-                for a in ft.sig.0.inputs.iter() { maybe_walk_ty_(*a, f); }
-                if let ty::FnConverging(output) = ft.sig.0.output {
-                    maybe_walk_ty_(output, f);
-                }
-            }
+/// Walks `ty` and any types appearing within `ty`, invoking the
+/// callback `f` on each type. If the callback returns false, then the
+/// children of the current type are ignored.
+///
+/// Note: prefer `ty.walk()` where possible.
+pub fn maybe_walk_ty<'tcx,F>(ty_root: Ty<'tcx>, mut f: F)
+    where F : FnMut(Ty<'tcx>) -> bool
+{
+    let mut walker = ty_root.walk();
+    while let Some(ty) = walker.next() {
+        if !f(ty) {
+            walker.skip_current_subtree();
         }
     }
 
diff --git a/src/librustc/middle/ty_walk.rs b/src/librustc/middle/ty_walk.rs
new file mode 100644
index 00000000000..406ebf4bc38
--- /dev/null
+++ b/src/librustc/middle/ty_walk.rs
@@ -0,0 +1,112 @@
+// Copyright 2012-2014 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+//! An iterator over the type substructure.
+
+use middle::ty::{mod, Ty};
+use std::iter::Iterator;
+
+pub struct TypeWalker<'tcx> {
+    stack: Vec<Ty<'tcx>>,
+    last_subtree: uint,
+}
+
+impl<'tcx> TypeWalker<'tcx> {
+    pub fn new(ty: Ty<'tcx>) -> TypeWalker<'tcx> {
+        TypeWalker { stack: vec!(ty), last_subtree: 1, }
+    }
+
+    fn push_subtypes(&mut self, parent_ty: Ty<'tcx>) {
+        match parent_ty.sty {
+            ty::ty_bool | ty::ty_char | ty::ty_int(_) | ty::ty_uint(_) | ty::ty_float(_) |
+            ty::ty_str | ty::ty_infer(_) | ty::ty_param(_) | ty::ty_err => {
+            }
+            ty::ty_uniq(ty) | ty::ty_vec(ty, _) | ty::ty_open(ty) => {
+                self.stack.push(ty);
+            }
+            ty::ty_ptr(ref mt) | ty::ty_rptr(_, ref mt) => {
+                self.stack.push(mt.ty);
+            }
+            ty::ty_projection(ref data) => {
+                self.push_reversed(data.trait_ref.substs.types.as_slice());
+            }
+            ty::ty_trait(box ty::TyTrait { ref principal, .. }) => {
+                self.push_reversed(principal.substs().types.as_slice());
+            }
+            ty::ty_enum(_, ref substs) |
+            ty::ty_struct(_, ref substs) |
+            ty::ty_unboxed_closure(_, _, ref substs) => {
+                self.push_reversed(substs.types.as_slice());
+            }
+            ty::ty_tup(ref ts) => {
+                self.push_reversed(ts.as_slice());
+            }
+            ty::ty_bare_fn(_, ref ft) => {
+                self.push_sig_subtypes(&ft.sig);
+            }
+            ty::ty_closure(ref ft) => {
+                self.push_sig_subtypes(&ft.sig);
+            }
+        }
+    }
+
+    fn push_sig_subtypes(&mut self, sig: &ty::PolyFnSig<'tcx>) {
+        match sig.0.output {
+            ty::FnConverging(output) => { self.stack.push(output); }
+            ty::FnDiverging => { }
+        }
+        self.push_reversed(sig.0.inputs.as_slice());
+    }
+
+    fn push_reversed(&mut self, tys: &[Ty<'tcx>]) {
+        // We push slices on the stack in reverse order so as to
+        // maintain a pre-order traversal. As of the time of this
+        // writing, the fact that the traversal is pre-order is not
+        // 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).
+        for &ty in tys.iter().rev() {
+            self.stack.push(ty);
+        }
+    }
+
+    /// Skips the subtree of types corresponding to the last type
+    /// returned by `next()`.
+    ///
+    /// Example: Imagine you are walking `Foo<Bar<int>, uint>`.
+    ///
+    /// ```rust
+    /// let mut iter: TypeWalker = ...;
+    /// iter.next(); // yields Foo
+    /// iter.next(); // yields Bar<int>
+    /// iter.skip_current_subtree(); // skips int
+    /// iter.next(); // yields uint
+    /// ```
+    pub fn skip_current_subtree(&mut self) {
+        self.stack.truncate(self.last_subtree);
+    }
+}
+
+impl<'tcx> Iterator<Ty<'tcx>> for TypeWalker<'tcx> {
+    fn next(&mut self) -> Option<Ty<'tcx>> {
+        debug!("next(): stack={}", self.stack);
+        match self.stack.pop() {
+            None => {
+                return None;
+            }
+            Some(ty) => {
+                self.last_subtree = self.stack.len();
+                self.push_subtypes(ty);
+                debug!("next: stack={}", self.stack);
+                Some(ty)
+            }
+        }
+    }
+}
diff --git a/src/librustc_driver/test.rs b/src/librustc_driver/test.rs
index 6329acfb578..eddcc750068 100644
--- a/src/librustc_driver/test.rs
+++ b/src/librustc_driver/test.rs
@@ -34,8 +34,6 @@ use syntax::codemap::{Span, CodeMap, DUMMY_SP};
 use syntax::diagnostic::{Level, RenderSpan, Bug, Fatal, Error, Warning, Note, Help};
 use syntax::parse::token;
 
-use arena::TypedArena;
-
 struct Env<'a, 'tcx: 'a> {
     infcx: &'a infer::InferCtxt<'a, 'tcx>,
 }
@@ -831,3 +829,57 @@ fn subst_region_renumber_region() {
         assert_eq!(t_substituted, t_expected);
     })
 }
+
+#[test]
+fn walk_ty() {
+    test_env(EMPTY_SOURCE_STR, errors(&[]), |env| {
+        let tcx = env.infcx.tcx;
+        let int_ty = tcx.types.int;
+        let uint_ty = tcx.types.uint;
+        let tup1_ty = ty::mk_tup(tcx, vec!(int_ty, uint_ty, int_ty, uint_ty));
+        let tup2_ty = ty::mk_tup(tcx, vec!(tup1_ty, tup1_ty, uint_ty));
+        let uniq_ty = ty::mk_uniq(tcx, tup2_ty);
+        let walked: Vec<_> = uniq_ty.walk().collect();
+        assert_eq!(vec!(uniq_ty,
+                        tup2_ty,
+                        tup1_ty, int_ty, uint_ty, int_ty, uint_ty,
+                        tup1_ty, int_ty, uint_ty, int_ty, uint_ty,
+                        uint_ty),
+                   walked);
+    })
+}
+
+#[test]
+fn walk_ty_skip_subtree() {
+    test_env(EMPTY_SOURCE_STR, errors(&[]), |env| {
+        let tcx = env.infcx.tcx;
+        let int_ty = tcx.types.int;
+        let uint_ty = tcx.types.uint;
+        let tup1_ty = ty::mk_tup(tcx, vec!(int_ty, uint_ty, int_ty, uint_ty));
+        let tup2_ty = ty::mk_tup(tcx, vec!(tup1_ty, tup1_ty, uint_ty));
+        let uniq_ty = ty::mk_uniq(tcx, tup2_ty);
+
+        // types we expect to see (in order), plus a boolean saying
+        // whether to skip the subtree.
+        let mut expected = vec!((uniq_ty, false),
+                                (tup2_ty, false),
+                                (tup1_ty, false),
+                                (int_ty, false),
+                                (uint_ty, false),
+                                (int_ty, false),
+                                (uint_ty, false),
+                                (tup1_ty, true), // skip the int/uint/int/uint
+                                (uint_ty, false));
+        expected.reverse();
+
+        let mut walker = uniq_ty.walk();
+        while let Some(t) = walker.next() {
+            debug!("walked to {}", t);
+            let (expected_ty, skip) = expected.pop().unwrap();
+            assert_eq!(t, expected_ty);
+            if skip { walker.skip_current_subtree(); }
+        }
+
+        assert!(expected.is_empty());
+    })
+}