Factor out duplicated entry_point_type functions

compiler/rustc_ast/src/entry.rs

@@ -1,3 +1,7 @@
+use crate::{attr, Attribute};
+use rustc_span::symbol::sym;
+use rustc_span::Symbol;
+
 #[derive(Debug)]
 pub enum EntryPointType {
     None,
@@ -6,3 +10,26 @@ pub enum EntryPointType {
     Start,
     OtherMain, // Not an entry point, but some other function named main
 }
+
+pub fn entry_point_type(
+    attrs: &[Attribute],
+    at_root: bool,
+    name: Option<Symbol>,
+) -> EntryPointType {
+    if attr::contains_name(attrs, sym::start) {
+        EntryPointType::Start
+    } else if attr::contains_name(attrs, sym::rustc_main) {
+        EntryPointType::RustcMainAttr
+    } else {
+        if let Some(name) = name && name == sym::main {
+            if at_root {
+                // This is a top-level function so it can be `main`.
+                EntryPointType::MainNamed
+            } else {
+                EntryPointType::OtherMain
+            }
+        } else {
+            EntryPointType::None
+        }
+    }
+}

compiler/rustc_builtin_macros/src/test_harness.rs

@@ -169,29 +169,15 @@ impl<'a> Visitor<'a> for InnerItemLinter<'_> {
     }
 }
 
-// Beware, this is duplicated in librustc_passes/entry.rs (with
-// `rustc_hir::Item`), so make sure to keep them in sync.
-fn entry_point_type(item: &ast::Item, depth: usize) -> EntryPointType {
+fn entry_point_type(item: &ast::Item, at_root: bool) -> EntryPointType {
     match item.kind {
         ast::ItemKind::Fn(..) => {
-            if attr::contains_name(&item.attrs, sym::start) {
-                EntryPointType::Start
-            } else if attr::contains_name(&item.attrs, sym::rustc_main) {
-                EntryPointType::RustcMainAttr
-            } else if item.ident.name == sym::main {
-                if depth == 0 {
-                    // This is a top-level function so can be 'main'
-                    EntryPointType::MainNamed
-                } else {
-                    EntryPointType::OtherMain
-                }
-            } else {
-                EntryPointType::None
-            }
+            rustc_ast::entry::entry_point_type(&item.attrs, at_root, Some(item.ident.name))
         }
         _ => EntryPointType::None,
     }
 }
+
 /// A folder used to remove any entry points (like fn main) because the harness
 /// generator will provide its own
 struct EntryPointCleaner<'a> {
@@ -210,7 +196,7 @@ impl<'a> MutVisitor for EntryPointCleaner<'a> {
         // Remove any #[rustc_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 item = match entry_point_type(&item, self.depth) {
+        let item = match entry_point_type(&item, self.depth == 0) {
             EntryPointType::MainNamed | EntryPointType::RustcMainAttr | EntryPointType::Start => {
                 item.map(|ast::Item { id, ident, attrs, kind, vis, span, tokens }| {
                     let allow_dead_code = attr::mk_attr_nested_word(

compiler/rustc_passes/src/entry.rs

@@ -52,31 +52,6 @@ fn entry_fn(tcx: TyCtxt<'_>, (): ()) -> Option<(DefId, EntryFnType)> {
     configure_main(tcx, &ctxt)
 }
 
-// Beware, this is duplicated in `librustc_builtin_macros/test_harness.rs`
-// (with `ast::Item`), so make sure to keep them in sync.
-// A small optimization was added so that hir::Item is fetched only when needed.
-// An equivalent optimization was not applied to the duplicated code in test_harness.rs.
-fn entry_point_type(ctxt: &EntryContext<'_>, id: ItemId, at_root: bool) -> EntryPointType {
-    let attrs = ctxt.tcx.hir().attrs(id.hir_id());
-    if attr::contains_name(attrs, sym::start) {
-        EntryPointType::Start
-    } else if attr::contains_name(attrs, sym::rustc_main) {
-        EntryPointType::RustcMainAttr
-    } else {
-        if let Some(name) = ctxt.tcx.opt_item_name(id.owner_id.to_def_id())
-            && name == sym::main {
-            if at_root {
-                // This is a top-level function so can be `main`.
-                EntryPointType::MainNamed
-            } else {
-                EntryPointType::OtherMain
-            }
-        } else {
-            EntryPointType::None
-        }
-    }
-}
-
 fn attr_span_by_symbol(ctxt: &EntryContext<'_>, id: ItemId, sym: Symbol) -> Option<Span> {
     let attrs = ctxt.tcx.hir().attrs(id.hir_id());
     attr::find_by_name(attrs, sym).map(|attr| attr.span)
@@ -85,7 +60,13 @@ fn attr_span_by_symbol(ctxt: &EntryContext<'_>, id: ItemId, sym: Symbol) -> Opti
 fn find_item(id: ItemId, ctxt: &mut EntryContext<'_>) {
     let at_root = ctxt.tcx.opt_local_parent(id.owner_id.def_id) == Some(CRATE_DEF_ID);
 
-    match entry_point_type(ctxt, id, at_root) {
+    let attrs = ctxt.tcx.hir().attrs(id.hir_id());
+    let entry_point_type = rustc_ast::entry::entry_point_type(
+        attrs,
+        at_root,
+        ctxt.tcx.opt_item_name(id.owner_id.to_def_id()),
+    );
+    match entry_point_type {
         EntryPointType::None => {
             if let Some(span) = attr_span_by_symbol(ctxt, id, sym::unix_sigpipe) {
                 ctxt.tcx.sess.emit_err(AttrOnlyOnMain { span, attr: sym::unix_sigpipe });