bugfix: Auto imports in worksheets in Scala 3

Previously the import would be inserted above using directives
scalameta · Jan 25, 2024 · 1da8ec3 · 1da8ec3
1 parent a6d28d7
commit 1da8ec3
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 13 deletions.
diff --git a/mtags/src/main/scala-3/scala/meta/internal/pc/AutoImports.scala b/mtags/src/main/scala-3/scala/meta/internal/pc/AutoImports.scala
@@ -288,6 +288,7 @@ object AutoImports:
       text: String,
       tree: Tree,
   )(using Context): AutoImportPosition =
+    def isScala3Worksheet = pos.source.path.isWorksheet
 
     @tailrec
     def lastPackageDef(
@@ -321,7 +322,10 @@ object AutoImports:
           case Some(stm) => (stm.endPos.line + 1, false)
           case None if pkg.pid.symbol.isEmptyPackage =>
             val offset =
-              ScriptFirstImportPosition.skipUsingDirectivesOffset(text)
+              ScriptFirstImportPosition.skipUsingDirectivesOffset(
+                text,
+                isScala3Worksheet,
+              )
             (pos.source.offsetToLine(offset), false)
           case None =>
             val pos = pkg.pid.endPos
@@ -362,7 +366,8 @@ object AutoImports:
 
     def fileStart =
       AutoImportPosition(
-        ScriptFirstImportPosition.skipUsingDirectivesOffset(text),
+        ScriptFirstImportPosition
+          .skipUsingDirectivesOffset(text, isScala3Worksheet),
         0,
         padTop = false,
       )

diff --git a/mtags/src/main/scala/scala/meta/internal/pc/ScriptFirstImportPosition.scala b/mtags/src/main/scala/scala/meta/internal/pc/ScriptFirstImportPosition.scala
@@ -55,8 +55,21 @@ object ScriptFirstImportPosition {
     }
   }
 
-  def skipUsingDirectivesOffset(text: String): Int =
-    skipPrefixesOffset(usingDirectives, adjustShebang(text))
+  def skipUsingDirectivesOffset(
+      text: String,
+      isScala3Worksheet: Boolean = false
+  ): Int =
+    if (isScala3Worksheet) {
+      val iterator = tokenize(adjustShebang(text)).iterator
+      skipScala3WorksheetObject(iterator)
+        .map { startOffset =>
+          val offset =
+            skipPrefixesOffset(usingDirectives, iterator, None)
+              .getOrElse(startOffset)
+          offset + 1
+        }
+        .getOrElse(0)
+    } else skipPrefixesOffset(usingDirectives, adjustShebang(text))
 
   def skipPrefixesOffset(prefixes: List[String], text: String): Int = {
     val it = tokenize(text).iterator
@@ -86,6 +99,19 @@ object ScriptFirstImportPosition {
     } else None
   }
 
+  @tailrec
+  private def skipScala3WorksheetObject(
+      it: Iterator[Token]
+  ): Option[Int] = {
+    if (it.hasNext) {
+      it.next() match {
+        // We skip everything until the brace in `object worksheet {`
+        case t: Token.LeftBrace => Some(t.pos.end)
+        case _ => skipScala3WorksheetObject(it)
+      }
+    } else None
+  }
+
   @tailrec
   private def skipPrefixesOffset(
       prefixes: List[String],

diff --git a/tests/unit/src/main/scala/tests/BaseWorksheetLspSuite.scala b/tests/unit/src/main/scala/tests/BaseWorksheetLspSuite.scala
@@ -870,6 +870,8 @@ abstract class BaseWorksheetLspSuite(
              |/metals.json
              |{"a": {"scalaVersion": "${scalaVersion}"}}
              |/$path
+             |//> using scala $scalaVersion
+             |
              |object A {
              |  val f = Future.successful(42)
              |}
@@ -881,10 +883,12 @@ abstract class BaseWorksheetLspSuite(
           server
             .assertCodeAction(
               path,
-              """|object A {
-                 |  val f = <<Future>>.successful(42)
-                 |}
-                 |""".stripMargin,
+              s"""|//> using scala $scalaVersion
+                  |
+                  |object A {
+                  |  val f = <<Future>>.successful(42)
+                  |}
+                  |""".stripMargin,
               expectedActions,
               Nil,
             )
@@ -895,11 +899,13 @@ abstract class BaseWorksheetLspSuite(
         // Assert if indentation is correct. See `AutoImports.renderImport`
         _ = assertNoDiff(
           server.bufferContents(path),
-          """|import scala.concurrent.Future
-             |object A {
-             |  val f = Future.successful(42)
-             |}
-             |""".stripMargin,
+          s"""|//> using scala $scalaVersion
+              |import scala.concurrent.Future
+              |
+              |object A {
+              |  val f = Future.successful(42)
+              |}
+              |""".stripMargin,
         )
       } yield ()
     }