From 1870e8987b99860a49209f9a66ec51406a1ae880 Mon Sep 17 00:00:00 2001
From: Jakub Ciesluk <323892@uwr.edu.pl>
Date: Mon, 20 Nov 2023 13:51:55 +0100
Subject: [PATCH] bugfix: Adjust auto import position for shebang

---
 .../pc/ScriptFirstImportPosition.scala        | 23 +++++++++-----
 .../tests/scalacli/ScalaCliActionsSuite.scala | 30 +++++++++++++++++++
 2 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/mtags/src/main/scala/scala/meta/internal/pc/ScriptFirstImportPosition.scala b/mtags/src/main/scala/scala/meta/internal/pc/ScriptFirstImportPosition.scala
index 3eff79e84ea..10cef33ed76 100644
--- a/mtags/src/main/scala/scala/meta/internal/pc/ScriptFirstImportPosition.scala
+++ b/mtags/src/main/scala/scala/meta/internal/pc/ScriptFirstImportPosition.scala
@@ -27,10 +27,14 @@ import scala.meta._
 object ScriptFirstImportPosition {
 
   val usingDirectives: List[String] = List("// using", "//> using")
+  val shebang = "#!"
   val ammHeaders: List[String] = List("// scala", "// ammonite")
 
+  private def adjustShebang(text: String): String =
+    text.replaceFirst(shebang, s"//$shebang")
+
   def ammoniteScStartOffset(text: String): Option[Int] = {
-    val it = tokenize(text).iterator
+    val it = tokenize(adjustShebang(text)).iterator
     startMarkerOffset(it, "/*<start>*/").map { startOffset =>
       val offset =
         skipPrefixesOffset(ammHeaders, it, None)
@@ -41,7 +45,7 @@ object ScriptFirstImportPosition {
   }
 
   def scalaCliScStartOffset(text: String): Option[Int] = {
-    val iterator = tokenize(text).iterator
+    val iterator = tokenize(adjustShebang(text)).iterator
     startMarkerOffset(iterator, "/*<script>*/").map { startOffset =>
       val offset =
         skipPrefixesOffset(usingDirectives, iterator, None)
@@ -52,7 +56,7 @@ object ScriptFirstImportPosition {
   }
 
   def skipUsingDirectivesOffset(text: String): Int =
-    skipPrefixesOffset(usingDirectives, text)
+    skipPrefixesOffset(usingDirectives, adjustShebang(text))
 
   def skipPrefixesOffset(prefixes: List[String], text: String): Int = {
     val it = tokenize(text).iterator
@@ -86,17 +90,20 @@ object ScriptFirstImportPosition {
   private def skipPrefixesOffset(
       prefixes: List[String],
       it: Iterator[Token],
-      lastOffset: Option[Int]
+      lastOffset: Option[Int],
+      foundShebang: Boolean = false
   ): Option[Int] = {
     if (it.hasNext) {
       it.next match {
         case t: Token.Comment
             if prefixes.exists(prefix => t.text.startsWith(prefix)) =>
-          skipPrefixesOffset(prefixes, it, Some(t.pos.end))
+          skipPrefixesOffset(prefixes, it, Some(t.pos.end), foundShebang)
+        case t: Token.Comment if t.value.startsWith(shebang) =>
+          skipPrefixesOffset(prefixes, it, Some(t.pos.end), foundShebang = true)
         case t if isWhitespace(t) =>
-          skipPrefixesOffset(prefixes, it, lastOffset)
-        case _ =>
-          lastOffset
+          skipPrefixesOffset(prefixes, it, lastOffset, foundShebang)
+        case _ if foundShebang => lastOffset.map(_ - 2)
+        case _ => lastOffset
       }
     } else lastOffset
   }
diff --git a/tests/slow/src/test/scala/tests/scalacli/ScalaCliActionsSuite.scala b/tests/slow/src/test/scala/tests/scalacli/ScalaCliActionsSuite.scala
index 6663593eb76..0c0091b1e31 100644
--- a/tests/slow/src/test/scala/tests/scalacli/ScalaCliActionsSuite.scala
+++ b/tests/slow/src/test/scala/tests/scalacli/ScalaCliActionsSuite.scala
@@ -124,4 +124,34 @@ class ScalaCliActionsSuite
     fileName = "A.sc",
   )
 
+  checkScalaCLI(
+    "auto-import-shebang",
+    s"""|#!/usr/bin/env -S scala-cli shebang
+        |
+        |//> using scala "${BuildInfo.scala213}"
+        |//> using lib "org.typelevel::cats-core:2.9.0"
+        |
+        |object A {
+        |  <<Future>>.successful(2)
+        |}
+        |""".stripMargin,
+    s"""|${ImportMissingSymbol.title("Future", "scala.concurrent")}
+        |${ImportMissingSymbol.title("Future", "java.util.concurrent")}
+        |${CreateNewSymbol.title("Future")}
+        |""".stripMargin,
+    s"""|#!/usr/bin/env -S scala-cli shebang
+        |
+        |//> using scala "${BuildInfo.scala213}"
+        |//> using lib "org.typelevel::cats-core:2.9.0"
+        |import scala.concurrent.Future
+        |
+        |object A {
+        |  Future.successful(2)
+        |}
+        |""".stripMargin,
+    scalaCliOptions = List("--actions", "-S", scalaVersion),
+    expectNoDiagnostics = false,
+    fileName = "A.sc",
+  )
+
 }