bugfix: Auto imports in worksheets in Scala 3

Author: kasiaMarek

presentation-compiler/src/main/dotty/tools/pc/AutoImports.scala

@@ -302,11 +302,24 @@ object AutoImports:
           }.headOption
         case _ => None
 
-    def skipUsingDirectivesOffset =
+    def firstMemberDefinitionStart(tree: Tree)(using Context): Option[Int] =
+      tree match
+        case PackageDef(_, stats) =>
+          stats.flatMap {
+            case s: PackageDef => firstMemberDefinitionStart(s)
+            case stat if stat.span.exists => Some(stat.span.start)
+            case _ => None
+          }.headOption
+        case _ => None
+
+
+    def skipUsingDirectivesOffset(
+      firstObjectPos: Int = firstMemberDefinitionStart(tree).getOrElse(0)
+    ): Int =
+      val firstObjectLine = pos.source.offsetToLine(firstObjectPos)
       comments
         .takeWhile(comment =>
-          !comment.isDocComment && comment.span.end < firstObjectBody(tree)
-            .fold(0)(_.span.start)
+          !comment.isDocComment && pos.source.offsetToLine(comment.span.end) + 1 < firstObjectLine
         )
         .lastOption
         .fold(0)(_.span.end + 1)
@@ -318,7 +331,7 @@ object AutoImports:
         val (lineNumber, padTop) = lastImportStatement match
           case Some(stm) => (stm.endPos.line + 1, false)
           case None if pkg.pid.symbol.isEmptyPackage =>
-            (pos.source.offsetToLine(skipUsingDirectivesOffset), false)
+            (pos.source.offsetToLine(skipUsingDirectivesOffset()), false)
           case None =>
             val pos = pkg.pid.endPos
             val line =
@@ -330,7 +343,7 @@ object AutoImports:
         new AutoImportPosition(offset, text, padTop)
       }
 
-    def forScript(isAmmonite: Boolean): Option[AutoImportPosition] =
+    def forScript(path: String): Option[AutoImportPosition] =
       firstObjectBody(tree).map { tmpl =>
         val lastImportStatement =
           tmpl.body.takeWhile(_.isInstanceOf[Import]).lastOption
@@ -340,10 +353,11 @@ object AutoImports:
             offset
           case None =>
             val scriptOffset =
-              if isAmmonite then
-                ScriptFirstImportPosition.ammoniteScStartOffset(text, comments)
-              else
-                ScriptFirstImportPosition.scalaCliScStartOffset(text, comments)
+              if path.isAmmoniteGeneratedFile
+              then ScriptFirstImportPosition.ammoniteScStartOffset(text, comments)
+              else if path.isScalaCLIGeneratedFile
+              then ScriptFirstImportPosition.scalaCliScStartOffset(text, comments)
+              else Some(skipUsingDirectivesOffset(tmpl.span.start))
 
             scriptOffset.getOrElse {
               val tmplPoint = tmpl.self.srcPos.span.point
@@ -359,14 +373,16 @@ object AutoImports:
 
     def fileStart =
       AutoImportPosition(
-        skipUsingDirectivesOffset,
+        skipUsingDirectivesOffset(),
         0,
         padTop = false
       )
 
     val scriptPos =
-      if path.isAmmoniteGeneratedFile then forScript(isAmmonite = true)
-      else if path.isScalaCLIGeneratedFile then forScript(isAmmonite = false)
+      if path.isAmmoniteGeneratedFile ||
+         path.isScalaCLIGeneratedFile ||
+         path.isWorksheet
+      then forScript(path)
       else None
 
     scriptPos

presentation-compiler/test/dotty/tools/pc/base/BaseAutoImportsSuite.scala

@@ -41,6 +41,18 @@ trait BaseAutoImportsSuite extends BaseCodeActionSuite:
       selection
     )
 
+  def checkWorksheetEdit(
+      original: String,
+      expected: String,
+      selection: Int = 0
+  ): Unit =
+    checkEditSelection(
+      "example.worksheet.sc",
+      original,
+      expected,
+      selection
+    )
+
   def checkEditSelection(
       filename: String,
       original: String,

presentation-compiler/test/dotty/tools/pc/tests/edit/AutoImportsSuite.scala

@@ -342,6 +342,69 @@ class AutoImportsSuite extends BaseAutoImportsSuite:
       )
     )
 
+  @Test def `worksheet-import` =
+    checkWorksheetEdit(
+      worksheetPcWrapper(
+        """|//> using scala 3.3.0
+           |
+           |// Some comment
+           |
+           |// Object comment
+           |object A {
+           |  val p: <<Path>> = ???
+           |}
+           |""".stripMargin
+      ),
+      worksheetPcWrapper(
+        """|//> using scala 3.3.0
+           |
+           |// Some comment
+           |import java.nio.file.Path
+           |
+           |// Object comment
+           |object A {
+           |  val p: Path = ???
+           |}
+           |""".stripMargin
+      )
+    )
+
+  @Test def `object-import` =
+    checkEdit(
+      """|object A {
+         |  //some comment
+         |  val p: <<Path>> = ???
+         |}
+         |""".stripMargin,
+      """|import java.nio.file.Path
+         |object A {
+         |  //some comment
+         |  val p: Path = ???
+         |}
+         |""".stripMargin,
+    )
+
+  @Test def `toplevels-import` =
+    checkEdit(
+      """|//some comment
+         |
+         |val p: <<Path>> = ???
+         |
+         |//some other comment
+         |
+         |val v = 1
+         |""".stripMargin,
+      """|//some comment
+         |import java.nio.file.Path
+         |
+         |val p: Path = ???
+         |
+         |//some other comment
+         |
+         |val v = 1
+         |""".stripMargin,
+    )
+
   private def ammoniteWrapper(code: String): String =
     // Vaguely looks like a scala file that Ammonite generates
     // from a sc file.
@@ -359,6 +422,11 @@ class AutoImportsSuite extends BaseAutoImportsSuite:
         |}
         |""".stripMargin
 
+  private def worksheetPcWrapper(code: String): String =
+    s"""|object worksheet{
+        |$code
+        |}""".stripMargin
+
   // https://dotty.epfl.ch/docs/internals/syntax.html#soft-keywords
   @Test
   def `soft-keyword-check-test` =