From b756a80224e1a2df0b604bf081addc578f5ac35a Mon Sep 17 00:00:00 2001
From: Nicholas Nethercote <n.nethercote@gmail.com>
Date: Fri, 29 Sep 2023 16:54:48 +1000
Subject: [PATCH] Factor out the two `entry_point_type` functions.

They are very similar, and each one has a comment about the importance
of being kept in sync with the other. This commit removes the
duplication.
---
 compiler/rustc_ast/src/entry.rs               | 26 ++++++++++++++++
 .../rustc_builtin_macros/src/test_harness.rs  | 18 ++---------
 compiler/rustc_passes/src/entry.rs            | 31 +++----------------
 3 files changed, 33 insertions(+), 42 deletions(-)

diff --git a/compiler/rustc_ast/src/entry.rs b/compiler/rustc_ast/src/entry.rs
index 3370146193a52..426d748b52d13 100644
--- a/compiler/rustc_ast/src/entry.rs
+++ b/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,25 @@ pub enum EntryPointType {
     Start,
     OtherMain, // Not an entry point, but some other function named main
 }
+
+pub fn entry_point_type<F>(attrs: &[Attribute], at_root: bool, get_name: F) -> EntryPointType
+where
+    F: FnOnce() -> Option<Symbol>,
+{
+    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) = get_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
+        }
+    }
+}
diff --git a/compiler/rustc_builtin_macros/src/test_harness.rs b/compiler/rustc_builtin_macros/src/test_harness.rs
index 026b4a2a6f641..bc7326dd3d944 100644
--- a/compiler/rustc_builtin_macros/src/test_harness.rs
+++ b/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, 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 at_root {
-                    // 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> {
diff --git a/compiler/rustc_passes/src/entry.rs b/compiler/rustc_passes/src/entry.rs
index 4f71704b8853d..25604ec037b1e 100644
--- a/compiler/rustc_passes/src/entry.rs
+++ b/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,11 @@ 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 });