Rollup merge of #71372 - ayushmishra2005:shebang_stripping, r=estebank

Fix #! (shebang) stripping account space issue

 #70528

2 files changed, 24 insertions, 1 deletions

diff --git a/src/librustc_lexer/src/lib.rs b/src/librustc_lexer/src/lib.rs
index 5ccfc1b276b..be85a34bd39 100644
--- a/src/librustc_lexer/src/lib.rs
+++ b/src/librustc_lexer/src/lib.rs
@@ -236,12 +236,17 @@ pub enum Base {
 /// (e.g. "#![deny(missing_docs)]").
 pub fn strip_shebang(input: &str) -> Option<usize> {
     debug_assert!(!input.is_empty());
-    if !input.starts_with("#!") || input.starts_with("#![") {
+    let s: &str = &remove_whitespace(input);
+    if !s.starts_with("#!") || s.starts_with("#![") {
         return None;
     }
     Some(input.find('\n').unwrap_or(input.len()))
 }
 
+fn remove_whitespace(s: &str) -> String {
+    s.chars().filter(|c| !c.is_whitespace()).collect()
+}
+
 /// Parses the first token from the provided input string.
 pub fn first_token(input: &str) -> Token {
     debug_assert!(!input.is_empty());
diff --git a/src/librustc_lexer/src/tests.rs b/src/librustc_lexer/src/tests.rs
index 06fc159fe25..065e8f3f646 100644
--- a/src/librustc_lexer/src/tests.rs
+++ b/src/librustc_lexer/src/tests.rs
@@ -145,4 +145,22 @@ mod tests {
             }),
         );
     }
+
+    #[test]
+    fn test_valid_shebang() {
+        // https://github.com/rust-lang/rust/issues/70528
+        let input = "#!/usr/bin/rustrun";
+        let actual = strip_shebang(input);
+        let expected: Option<usize> = Some(18);
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn test_invalid_shebang_valid_rust_syntax() {
+        // https://github.com/rust-lang/rust/issues/70528
+        let input = "#!    [bad_attribute]";
+        let actual = strip_shebang(input);
+        let expected: Option<usize> = None;
+        assert_eq!(expected, actual);
+    }
 }