Auto merge of #15709 - Veykril:runnables-custom-main, r=Veykril

fix: Recognize custom main function as binary entrypoint for runnables

3 files changed, 72 insertions, 23 deletions

diff --git a/crates/hir-def/src/attr.rs b/crates/hir-def/src/attr.rs
index c6454eb9ea0..fa3025e0303 100644
--- a/crates/hir-def/src/attr.rs
+++ b/crates/hir-def/src/attr.rs
@@ -215,6 +215,10 @@ impl Attrs {
         self.doc_exprs().flat_map(|doc_expr| doc_expr.aliases().to_vec())
     }
 
+    pub fn export_name(&self) -> Option<&SmolStr> {
+        self.by_key("export_name").string_value()
+    }
+
     pub fn is_proc_macro(&self) -> bool {
         self.by_key("proc_macro").exists()
     }
diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
index a6c6c0dbb8b..8e48afd6af8 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -1971,6 +1971,17 @@ impl Function {
         db.function_data(self.id).attrs.is_test()
     }
 
+    /// is this a `fn main` or a function with an `export_name` of `main`?
+    pub fn is_main(self, db: &dyn HirDatabase) -> bool {
+        if !self.module(db).is_crate_root() {
+            return false;
+        }
+        let data = db.function_data(self.id);
+
+        data.name.to_smol_str() == "main"
+            || data.attrs.export_name().map(core::ops::Deref::deref) == Some("main")
+    }
+
     /// Does this function have the ignore attribute?
     pub fn is_ignore(self, db: &dyn HirDatabase) -> bool {
         db.function_data(self.id).attrs.is_ignore()
diff --git a/crates/ide/src/runnables.rs b/crates/ide/src/runnables.rs
index 2d528c64255..07cdddd15f8 100644
--- a/crates/ide/src/runnables.rs
+++ b/crates/ide/src/runnables.rs
@@ -308,11 +308,7 @@ pub(crate) fn runnable_fn(
     sema: &Semantics<'_, RootDatabase>,
     def: hir::Function,
 ) -> Option<Runnable> {
-    let name = def.name(sema.db).to_smol_str();
-
-    let root = def.module(sema.db).krate().root_module();
-
-    let kind = if name == "main" && def.module(sema.db) == root {
+    let kind = if def.is_main(sema.db) {
         RunnableKind::Bin
     } else {
         let test_id = || {
@@ -320,7 +316,9 @@ pub(crate) fn runnable_fn(
                 let def: hir::ModuleDef = def.into();
                 def.canonical_path(sema.db)
             };
-            canonical_path.map(TestId::Path).unwrap_or(TestId::Name(name))
+            canonical_path
+                .map(TestId::Path)
+                .unwrap_or(TestId::Name(def.name(sema.db).to_smol_str()))
         };
 
         if def.is_test(sema.db) {
@@ -587,6 +585,9 @@ mod tests {
 $0
 fn main() {}
 
+#[export_name = "main"]
+fn __cortex_m_rt_main_trampoline() {}
+
 #[test]
 fn test_foo() {}
 
@@ -604,7 +605,7 @@ mod not_a_root {
     fn main() {}
 }
 "#,
-            &[TestMod, Bin, Test, Test, Test, Bench],
+            &[TestMod, Bin, Bin, Test, Test, Test, Bench],
             expect![[r#"
                 [
                     Runnable {
@@ -613,7 +614,7 @@ mod not_a_root {
                             file_id: FileId(
                                 0,
                             ),
-                            full_range: 0..190,
+                            full_range: 0..253,
                             name: "",
                             kind: Module,
                         },
@@ -642,8 +643,22 @@ mod not_a_root {
                             file_id: FileId(
                                 0,
                             ),
-                            full_range: 15..39,
-                            focus_range: 26..34,
+                            full_range: 15..76,
+                            focus_range: 42..71,
+                            name: "__cortex_m_rt_main_trampoline",
+                            kind: Function,
+                        },
+                        kind: Bin,
+                        cfg: None,
+                    },
+                    Runnable {
+                        use_name_in_title: false,
+                        nav: NavigationTarget {
+                            file_id: FileId(
+                                0,
+                            ),
+                            full_range: 78..102,
+                            focus_range: 89..97,
                             name: "test_foo",
                             kind: Function,
                         },
@@ -663,8 +678,8 @@ mod not_a_root {
                             file_id: FileId(
                                 0,
                             ),
-                            full_range: 41..92,
-                            focus_range: 73..87,
+                            full_range: 104..155,
+                            focus_range: 136..150,
                             name: "test_full_path",
                             kind: Function,
                         },
@@ -684,8 +699,8 @@ mod not_a_root {
                             file_id: FileId(
                                 0,
                             ),
-                            full_range: 94..128,
-                            focus_range: 115..123,
+                            full_range: 157..191,
+                            focus_range: 178..186,
                             name: "test_foo",
                             kind: Function,
                         },
@@ -705,8 +720,8 @@ mod not_a_root {
                             file_id: FileId(
                                 0,
                             ),
-                            full_range: 130..152,
-                            focus_range: 142..147,
+                            full_range: 193..215,
+                            focus_range: 205..210,
                             name: "bench",
                             kind: Function,
                         },
@@ -1655,12 +1670,18 @@ macro_rules! gen2 {
         }
     }
 }
+macro_rules! gen_main {
+    () => {
+        fn main() {}
+    }
+}
 mod tests {
     gen!();
 }
 gen2!();
+gen_main!();
 "#,
-            &[TestMod, TestMod, Test, Test, TestMod],
+            &[TestMod, TestMod, Test, Test, TestMod, Bin],
             expect![[r#"
                 [
                     Runnable {
@@ -1669,7 +1690,7 @@ gen2!();
                             file_id: FileId(
                                 0,
                             ),
-                            full_range: 0..237,
+                            full_range: 0..315,
                             name: "",
                             kind: Module,
                         },
@@ -1684,8 +1705,8 @@ gen2!();
                             file_id: FileId(
                                 0,
                             ),
-                            full_range: 202..227,
-                            focus_range: 206..211,
+                            full_range: 267..292,
+                            focus_range: 271..276,
                             name: "tests",
                             kind: Module,
                             description: "mod tests",
@@ -1701,7 +1722,7 @@ gen2!();
                             file_id: FileId(
                                 0,
                             ),
-                            full_range: 218..225,
+                            full_range: 283..290,
                             name: "foo_test",
                             kind: Function,
                         },
@@ -1721,7 +1742,7 @@ gen2!();
                             file_id: FileId(
                                 0,
                             ),
-                            full_range: 228..236,
+                            full_range: 293..301,
                             name: "foo_test2",
                             kind: Function,
                         },
@@ -1741,7 +1762,7 @@ gen2!();
                             file_id: FileId(
                                 0,
                             ),
-                            full_range: 228..236,
+                            full_range: 293..301,
                             name: "tests2",
                             kind: Module,
                             description: "mod tests2",
@@ -1751,6 +1772,19 @@ gen2!();
                         },
                         cfg: None,
                     },
+                    Runnable {
+                        use_name_in_title: false,
+                        nav: NavigationTarget {
+                            file_id: FileId(
+                                0,
+                            ),
+                            full_range: 302..314,
+                            name: "main",
+                            kind: Function,
+                        },
+                        kind: Bin,
+                        cfg: None,
+                    },
                 ]
             "#]],
         );