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.

src/libsyntax/test.rs

@@ -175,15 +175,41 @@ 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
+        if !tests.is_empty() || !tested_submods.is_empty() {
-        // the one we're going to add. Only if compiling an executable.
+            let (it, sym) = mk_reexport_mod(&mut self.cx, tests, tested_submods);
+            mod_folded.items.push(it);
 
-        mod_folded.items = mem::replace(&mut mod_folded.items, vec![]).move_map(|item| {
+            if !self.cx.path.is_empty() {
-            match entry::entry_point_type(&item, self.cx.path.len() + 1) {
+                self.tested_submods.push((self.cx.path[self.cx.path.len()-1], sym));
+            } else {
+                debug!("pushing nothing, sym: {:?}", sym);
+                self.cx.toplevel_reexport = Some(sym);
+            }
+        }
+
+        mod_folded
+    }
+}
+
+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 =>
-                    item.map(|ast::Item {id, ident, attrs, node, vis, span}| {
+                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 =
@@ -210,23 +236,10 @@ impl<'a> fold::Folder for TestHarnessGenerator<'a> {
                     }
                 }),
             EntryPointType::None |
-                EntryPointType::OtherMain => item,
+            EntryPointType::OtherMain => folded,
-            }
+        };
-        });
 
-        if !tests.is_empty() || !tested_submods.is_empty() {
+        SmallVector::one(folded)
-            let (it, sym) = mk_reexport_mod(&mut self.cx, tests, tested_submods);
-            mod_folded.items.push(it);
-
-            if !self.cx.path.is_empty() {
-                self.tested_submods.push((self.cx.path[self.cx.path.len()-1], sym));
-            } else {
-                debug!("pushing nothing, sym: {:?}", sym);
-                self.cx.toplevel_reexport = Some(sym);
-            }
-        }
-
-        mod_folded
     }
 }
 
@@ -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,

src/test/run-pass/test-runner-hides-buried-main.rs

@@ -0,0 +1,24 @@
+// Copyright 2015 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.
+
+// compile-flags: --test
+
+#![feature(main)]
+
+#![allow(dead_code)]
+
+mod a {
+    fn b() {
+        || {
+            #[main]
+            fn c() { panic!(); }
+        };
+    }
+}