From 09ae60019e1730bc8137a064c2fcc6966ba86da5 Mon Sep 17 00:00:00 2001
From: Fredrik Fornwall <fredrik@fornwall.net>
Date: Mon, 10 Apr 2023 23:28:01 +0200
Subject: [PATCH] Detect main function even with spaces after

Fixes #95.
---
 src/manifest.rs                       | 11 ++++-------
 tests/data/script-main-with-spaces.rs |  4 ++++
 tests/tests/script.rs                 |  9 +++++++++
 3 files changed, 17 insertions(+), 7 deletions(-)
 create mode 100644 tests/data/script-main-with-spaces.rs

diff --git a/src/manifest.rs b/src/manifest.rs
index 46f1797..2c8a1aa 100644
--- a/src/manifest.rs
+++ b/src/manifest.rs
@@ -26,12 +26,9 @@ pub fn split_input(
     script_name: &str,
     toolchain: Option<String>,
 ) -> MainResult<(String, String)> {
-    fn contains_main_method(line: &str) -> bool {
-        let line = line.trim_start();
-        line.starts_with("fn main(")
-            || line.starts_with("pub fn main(")
-            || line.starts_with("async fn main(")
-            || line.starts_with("pub async fn main(")
+    fn contains_main_method(source: &str) -> bool {
+        let re_shebang: Regex = Regex::new(r"(?m)^ *(pub )?(async )?fn main *\(").unwrap();
+        re_shebang.is_match(source)
     }
 
     let (part_mani, source, template, sub_prelude) = match input {
@@ -41,7 +38,7 @@ pub fn split_input(
             let (manifest, source) =
                 find_embedded_manifest(content).unwrap_or((Manifest::Toml(""), content));
 
-            let source = if source.lines().any(contains_main_method) {
+            let source = if contains_main_method(source) {
                 source.to_string()
             } else {
                 format!("fn main() -> Result<(), Box<dyn std::error::Error+Sync+Send>> {{\n    {{\n    {}    }}\n    Ok(())\n}}", source)
diff --git a/tests/data/script-main-with-spaces.rs b/tests/data/script-main-with-spaces.rs
new file mode 100644
index 0000000..d18903b
--- /dev/null
+++ b/tests/data/script-main-with-spaces.rs
@@ -0,0 +1,4 @@
+fn main  () {
+    println!("--output--");
+    println!("Hello, World!");
+}
diff --git a/tests/tests/script.rs b/tests/tests/script.rs
index bdfa523..08ad4c2 100644
--- a/tests/tests/script.rs
+++ b/tests/tests/script.rs
@@ -34,6 +34,15 @@ fn test_script_full_line_without_main() {
     .unwrap()
 }
 
+#[test]
+fn test_script_main_with_space() {
+    let out = rust_script!("tests/data/script-main-with-spaces.rs").unwrap();
+    scan!(out.stdout_output();
+        ("Hello, World!") => ()
+    )
+    .unwrap()
+}
+
 #[test]
 fn test_script_invalid_doc_comment() {
     let out = rust_script!("tests/data/script-invalid-doc-comment.rs").unwrap();