Move main removal to its own pass in --test mode

This handles the case where the #[main] function is buried deeper in
the ast than we search for #[test] functions.  I'm not sure why one
would want to do that, but since it works in standard compilation it
should also work for tests.

1 files changed, 56 insertions, 39 deletions

diff --git a/src/libsyntax/test.rs b/src/libsyntax/test.rs
index 3fbcbd728aa..9975e25f493 100644
--- a/src/libsyntax/test.rs
+++ b/src/libsyntax/test.rs
@@ -175,45 +175,6 @@ impl<'a> fold::Folder for TestHarnessGenerator<'a> {
         let tests = mem::replace(&mut self.tests, tests);
         let tested_submods = mem::replace(&mut self.tested_submods, tested_submods);
 
-        // Remove any #[main] from the AST so it doesn't clash with
-        // the one we're going to add. Only if compiling an executable.
-
-        mod_folded.items = mem::replace(&mut mod_folded.items, vec![]).move_map(|item| {
-            match entry::entry_point_type(&item, self.cx.path.len() + 1) {
-                EntryPointType::MainNamed |
-                EntryPointType::MainAttr |
-                EntryPointType::Start =>
-                    item.map(|ast::Item {id, ident, attrs, node, vis, span}| {
-                        let allow_str = InternedString::new("allow");
-                        let dead_code_str = InternedString::new("dead_code");
-                        let allow_dead_code_item =
-                            attr::mk_list_item(allow_str,
-                                               vec![attr::mk_word_item(dead_code_str)]);
-                        let allow_dead_code = attr::mk_attr_outer(attr::mk_attr_id(),
-                                                                  allow_dead_code_item);
-
-                        ast::Item {
-                            id: id,
-                            ident: ident,
-                            attrs: attrs.into_iter().filter_map(|attr| {
-                                if !attr.check_name("main") {
-                                    Some(attr)
-                                } else {
-                                    None
-                                }
-                            })
-                                .chain(iter::once(allow_dead_code))
-                                .collect(),
-                            node: node,
-                            vis: vis,
-                            span: span
-                        }
-                    }),
-                EntryPointType::None |
-                EntryPointType::OtherMain => item,
-            }
-        });
-
         if !tests.is_empty() || !tested_submods.is_empty() {
             let (it, sym) = mk_reexport_mod(&mut self.cx, tests, tested_submods);
             mod_folded.items.push(it);
@@ -230,6 +191,58 @@ impl<'a> fold::Folder for TestHarnessGenerator<'a> {
     }
 }
 
+struct EntryPointCleaner {
+    // Current depth in the ast
+    depth: usize,
+}
+
+impl fold::Folder for EntryPointCleaner {
+    fn fold_item(&mut self, i: P<ast::Item>) -> SmallVector<P<ast::Item>> {
+        self.depth += 1;
+        let folded = fold::noop_fold_item(i, self).expect_one("noop did something");
+        self.depth -= 1;
+
+        // Remove any #[main] from the AST so it doesn't clash with
+        // the one we're going to add, but mark it as
+        // #[allow(dead_code)] to avoid printing warnings.
+        let folded = match entry::entry_point_type(&*folded, self.depth) {
+            EntryPointType::MainNamed |
+            EntryPointType::MainAttr |
+            EntryPointType::Start =>
+                folded.map(|ast::Item {id, ident, attrs, node, vis, span}| {
+                    let allow_str = InternedString::new("allow");
+                    let dead_code_str = InternedString::new("dead_code");
+                    let allow_dead_code_item =
+                        attr::mk_list_item(allow_str,
+                                           vec![attr::mk_word_item(dead_code_str)]);
+                    let allow_dead_code = attr::mk_attr_outer(attr::mk_attr_id(),
+                                                              allow_dead_code_item);
+
+                    ast::Item {
+                        id: id,
+                        ident: ident,
+                        attrs: attrs.into_iter().filter_map(|attr| {
+                            if !attr.check_name("main") {
+                                Some(attr)
+                            } else {
+                                None
+                            }
+                        })
+                            .chain(iter::once(allow_dead_code))
+                            .collect(),
+                        node: node,
+                        vis: vis,
+                        span: span
+                    }
+                }),
+            EntryPointType::None |
+            EntryPointType::OtherMain => folded,
+        };
+
+        SmallVector::one(folded)
+    }
+}
+
 fn mk_reexport_mod(cx: &mut TestCtxt, tests: Vec<ast::Ident>,
                    tested_submods: Vec<(ast::Ident, ast::Ident)>) -> (P<ast::Item>, ast::Ident) {
     let super_ = token::str_to_ident("super");
@@ -265,6 +278,10 @@ fn generate_test_harness(sess: &ParseSess,
                          krate: ast::Crate,
                          cfg: &ast::CrateConfig,
                          sd: &diagnostic::SpanHandler) -> ast::Crate {
+    // Remove the entry points
+    let mut cleaner = EntryPointCleaner { depth: 0 };
+    let krate = cleaner.fold_crate(krate);
+
     let mut feature_gated_cfgs = vec![];
     let mut cx: TestCtxt = TestCtxt {
         sess: sess,