Improve --test interaction with entry points

src/librustc/middle/entry.rs

@@ -11,20 +11,19 @@
 
 use ast_map;
 use session::{config, Session};
-use syntax::ast::{Name, NodeId, Item, ItemFn};
+use syntax;
+use syntax::ast::{NodeId, Item};
 use syntax::attr;
 use syntax::codemap::Span;
-use syntax::parse::token;
+use syntax::entry::EntryPointType;
 use syntax::visit;
 use syntax::visit::Visitor;
 
-struct EntryContext<'a, 'ast: 'a> {
+struct EntryContext<'a> {
     session: &'a Session,
 
-    ast_map: &'a ast_map::Map<'ast>,
-
-    // The interned Name for "main".
-    main_name: Name,
+    // The current depth in the ast
+    depth: usize,
 
     // The top-level function called 'main'
     main_fn: Option<(NodeId, Span)>,
@@ -40,9 +39,11 @@ struct EntryContext<'a, 'ast: 'a> {
     non_main_fns: Vec<(NodeId, Span)> ,
 }
 
-impl<'a, 'ast, 'v> Visitor<'v> for EntryContext<'a, 'ast> {
+impl<'a, 'v> Visitor<'v> for EntryContext<'a> {
     fn visit_item(&mut self, item: &Item) {
+        self.depth += 1;
         find_item(item, self);
+        self.depth -= 1;
     }
 }
 
@@ -63,8 +64,7 @@ pub fn find_entry_point(session: &Session, ast_map: &ast_map::Map) {
 
     let mut ctxt = EntryContext {
         session: session,
-        main_name: token::intern("main"),
-        ast_map: ast_map,
+        depth: 0,
         main_fn: None,
         attr_main_fn: None,
         start_fn: None,
@@ -77,44 +77,35 @@ pub fn find_entry_point(session: &Session, ast_map: &ast_map::Map) {
 }
 
 fn find_item(item: &Item, ctxt: &mut EntryContext) {
-    match item.node {
-        ItemFn(..) => {
-            if item.ident.name == ctxt.main_name {
-                 ctxt.ast_map.with_path(item.id, |path| {
-                        if path.count() == 1 {
-                            // This is a top-level function so can be 'main'
-                            if ctxt.main_fn.is_none() {
-                                ctxt.main_fn = Some((item.id, item.span));
-                            } else {
-                                span_err!(ctxt.session, item.span, E0136,
-                                          "multiple 'main' functions");
-                            }
-                        } else {
-                            // This isn't main
-                            ctxt.non_main_fns.push((item.id, item.span));
-                        }
-                });
+    match syntax::entry::entry_point_type(item, ctxt.depth) {
+        EntryPointType::MainNamed => {
+            if ctxt.main_fn.is_none() {
+                ctxt.main_fn = Some((item.id, item.span));
+            } else {
+                span_err!(ctxt.session, item.span, E0136,
+                          "multiple 'main' functions");
             }
-
-            if attr::contains_name(&item.attrs, "main") {
-                if ctxt.attr_main_fn.is_none() {
-                    ctxt.attr_main_fn = Some((item.id, item.span));
-                } else {
-                    span_err!(ctxt.session, item.span, E0137,
-                              "multiple functions with a #[main] attribute");
-                }
+        },
+        EntryPointType::OtherMain => {
+            ctxt.non_main_fns.push((item.id, item.span));
+        },
+        EntryPointType::MainAttr => {
+            if ctxt.attr_main_fn.is_none() {
+                ctxt.attr_main_fn = Some((item.id, item.span));
+            } else {
+                span_err!(ctxt.session, item.span, E0137,
+                          "multiple functions with a #[main] attribute");
             }
-
-            if attr::contains_name(&item.attrs, "start") {
-                if ctxt.start_fn.is_none() {
-                    ctxt.start_fn = Some((item.id, item.span));
-                } else {
-                    span_err!(ctxt.session, item.span, E0138,
-                              "multiple 'start' functions");
-                }
+        },
+        EntryPointType::Start => {
+            if ctxt.start_fn.is_none() {
+                ctxt.start_fn = Some((item.id, item.span));
+            } else {
+                span_err!(ctxt.session, item.span, E0138,
+                          "multiple 'start' functions");
             }
-        }
-        _ => ()
+        },
+        EntryPointType::None => ()
     }
 
     visit::walk_item(ctxt, item);

src/libsyntax/entry.rs

@@ -0,0 +1,42 @@
+// Copyright 2012-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.
+
+use attr;
+use ast::{Item, ItemFn};
+
+pub enum EntryPointType {
+    None,
+    MainNamed,
+    MainAttr,
+    Start,
+    OtherMain, // Not an entry point, but some other function named main
+}
+
+pub fn entry_point_type(item: &Item, depth: usize) -> EntryPointType {
+    match item.node {
+        ItemFn(..) => {
+            if attr::contains_name(&item.attrs, "start") {
+                EntryPointType::Start
+            } else if attr::contains_name(&item.attrs, "main") {
+                EntryPointType::MainAttr
+            } else if item.ident.name == "main" {
+                if depth == 1 {
+                    // This is a top-level function so can be 'main'
+                    EntryPointType::MainNamed
+                } else {
+                    EntryPointType::OtherMain
+                }
+            } else {
+                EntryPointType::None
+            }
+        }
+        _ => EntryPointType::None,
+    }
+}

src/libsyntax/lib.rs

@@ -90,6 +90,7 @@ pub mod attr;
 pub mod codemap;
 pub mod config;
 pub mod diagnostic;
+pub mod entry;
 pub mod feature_gate;
 pub mod fold;
 pub mod owned_slice;

src/libsyntax/test.rs

@@ -14,6 +14,7 @@
 #![allow(unused_imports)]
 use self::HasTestSignature::*;
 
+use std::iter;
 use std::slice;
 use std::mem;
 use std::vec;
@@ -24,6 +25,7 @@ use codemap::{DUMMY_SP, Span, ExpnInfo, NameAndSpan, MacroAttribute};
 use codemap;
 use diagnostic;
 use config;
+use entry::{self, EntryPointType};
 use ext::base::ExtCtxt;
 use ext::build::AstBuilder;
 use ext::expand::ExpansionConfig;
@@ -173,28 +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| {
-            item.map(|ast::Item {id, ident, attrs, node, vis, span}| {
-                ast::Item {
-                    id: id,
-                    ident: ident,
-                    attrs: attrs.into_iter().filter_map(|attr| {
-                        if !attr.check_name("main") {
-                            Some(attr)
-                        } else {
-                            None
-                        }
-                    }).collect(),
-                    node: node,
-                    vis: vis,
-                    span: span
-                }
-            })
-        });
-
         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);
@@ -211,6 +191,55 @@ 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] or #[start] 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(|attr| {
+                                !attr.check_name("main") && !attr.check_name("start")
+                            })
+                            .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");
@@ -246,6 +275,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/compile-fail/test-warns-dead-code.rs

@@ -0,0 +1,17 @@
+// 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
+
+#![deny(dead_code)]
+
+fn dead() {} //~ error: function is never used: `dead`
+
+fn main() {}

src/test/run-pass/test-main-not-dead-attr.rs

@@ -0,0 +1,18 @@
+// 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)]
+
+#![deny(dead_code)]
+
+#[main]
+fn foo() { panic!(); }

src/test/run-pass/test-main-not-dead.rs

@@ -0,0 +1,15 @@
+// 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
+
+#![deny(dead_code)]
+
+fn main() { panic!(); }

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!(); }
+        };
+    }
+}

src/test/run-pass/test-runner-hides-start.rs

@@ -0,0 +1,16 @@
+// 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(start)]
+
+#[start]
+fn start(_: isize, _: *const *const u8) -> isize { panic!(); }