Support go to def, rename, etc. in worksheets

Author: Duhemm

language-server/src/dotty/tools/languageserver/DottyLanguageServer.scala

@@ -183,7 +183,7 @@ class DottyLanguageServer extends LanguageServer
     val worksheetMode = isWorksheet(uri)
 
     val (text, positionMapper) =
-      if (worksheetMode) (wrapWorksheet(document.getText), Some(worksheetPositionMapper _))
+      if (worksheetMode) (wrapWorksheet(document.getText), Some(toUnwrappedPosition _))
       else (document.getText, None)
 
     val diags = driver.run(uri, text)
@@ -204,7 +204,7 @@ class DottyLanguageServer extends LanguageServer
     assert(change.getRange == null, "TextDocumentSyncKind.Incremental support is not implemented")
 
     val (text, positionMapper) =
-      if (worksheetMode) (wrapWorksheet(change.getText), Some(worksheetPositionMapper _))
+      if (worksheetMode) (wrapWorksheet(change.getText), Some(toUnwrappedPosition _))
       else (change.getText, None)
 
     val diags = driver.run(uri, text)
@@ -317,7 +317,7 @@ class DottyLanguageServer extends LanguageServer
             (Nil, Include.overriding)
         }
       val defs = Interactive.namedTrees(trees, include, sym)
-      defs.flatMap(d => location(d.namePos)).asJava
+      defs.flatMap(d => location(d.namePos, positionMapperFor(d.source))).asJava
     }
   }
 
@@ -340,7 +340,7 @@ class DottyLanguageServer extends LanguageServer
         Include.references | Include.overriding | (if (includeDeclaration) Include.definitions else 0)
       val refs = Interactive.findTreesMatching(trees, includes, sym)
 
-      refs.flatMap(ref => location(ref.namePos)).asJava
+      refs.flatMap(ref => location(ref.namePos, positionMapperFor(ref.source))).asJava
     }
   }
 
@@ -363,7 +363,7 @@ class DottyLanguageServer extends LanguageServer
       val changes = refs.groupBy(ref => toUri(ref.source).toString)
         .mapValues(refs =>
           refs.flatMap(ref =>
-            range(ref.namePos).map(nameRange => new TextEdit(nameRange, newName))).asJava)
+            range(ref.namePos, positionMapperFor(ref.source)).map(nameRange => new TextEdit(nameRange, newName))).asJava)
 
       new WorkspaceEdit(changes.asJava)
     }
@@ -383,7 +383,7 @@ class DottyLanguageServer extends LanguageServer
       val refs = Interactive.namedTrees(uriTrees, Include.references | Include.overriding, sym)
       (for {
         ref <- refs if !ref.tree.symbol.isConstructor
-        nameRange <- range(ref.namePos)
+        nameRange <- range(ref.namePos, positionMapperFor(ref.source))
       } yield new DocumentHighlight(nameRange, DocumentHighlightKind.Read)).asJava
     }
   }
@@ -416,8 +416,8 @@ class DottyLanguageServer extends LanguageServer
 
     val defs = Interactive.namedTrees(uriTrees, includeReferences = false, _ => true)
     (for {
-      d <- defs
-      info <- symbolInfo(d.tree.symbol, d.namePos)
+      d <- defs if !isWorksheetWrapper(d)
+      info <- symbolInfo(d.tree.symbol, d.namePos, positionMapperFor(d.source))
     } yield JEither.forLeft(info)).asJava
   }
 
@@ -429,7 +429,7 @@ class DottyLanguageServer extends LanguageServer
 
       val trees = driver.allTrees
       val defs = Interactive.namedTrees(trees, includeReferences = false, nameSubstring = query)
-      defs.flatMap(d => symbolInfo(d.tree.symbol, d.namePos))
+      defs.flatMap(d => symbolInfo(d.tree.symbol, d.namePos, positionMapperFor(d.source)))
     }.asJava
   }
 
@@ -473,9 +473,12 @@ object DottyLanguageServer {
 
   /** Convert an lsp4j.Position to a SourcePosition */
   def sourcePosition(driver: InteractiveDriver, uri: URI, pos: lsp4j.Position): SourcePosition = {
+    val actualPosition =
+      if (isWorksheet(uri)) toWrappedPosition(pos)
+      else pos
     val source = driver.openedFiles(uri)
     if (source.exists) {
-      val p = Positions.Position(source.lineToOffset(pos.getLine) + pos.getCharacter)
+      val p = Positions.Position(source.lineToOffset(actualPosition.getLine) + actualPosition.getCharacter)
       new SourcePosition(source, p)
     }
     else NoSourcePosition
@@ -528,6 +531,10 @@ object DottyLanguageServer {
   private def isWorksheet(uri: URI): Boolean =
     uri.toString.endsWith(".sc")
 
+  /** Does this sourcefile represent a worksheet? */
+  private def isWorksheet(sourcefile: SourceFile): Boolean =
+    sourcefile.file.extension == "sc"
+
   /** Wrap the source of a worksheet inside an `object`. */
   private def wrapWorksheet(source: String): String =
     s"""object Worksheet {
@@ -544,13 +551,48 @@ object DottyLanguageServer {
    * @param position The position as seen by the compiler (after wrapping)
    * @return The position in the actual source file (before wrapping).
    */
-  private def worksheetPositionMapper(position: SourcePosition): SourcePosition = {
+  private def toUnwrappedPosition(position: SourcePosition): SourcePosition = {
     new SourcePosition(position.source, position.pos, position.outer) {
       override def startLine: Int = position.startLine - 1
       override def endLine: Int = position.endLine - 1
     }
   }
 
+  /**
+   * Map `position` in an unwrapped worksheet to the same position in the wrapped source.
+   *
+   * Because worksheet are wrapped in an `object`, the positions in the source are one line
+   * above from what the compiler sees.
+   *
+   * @see wrapWorksheet
+   * @param position The position as seen by VSCode (before wrapping)
+   * @return The position as seen by the compiler (after wrapping)
+   */
+  private def toWrappedPosition(position: lsp4j.Position): lsp4j.Position = {
+    new lsp4j.Position(position.getLine + 1, position.getCharacter)
+  }
+
+  /**
+   * Returns the position mapper necessary to unwrap positions for `sourcefile`. If `sourcefile` is
+   * not a worksheet, no mapper is necessary. Otherwise, return `toUnwrappedPosition`.
+   */
+  private def positionMapperFor(sourcefile: SourceFile): Option[SourcePosition => SourcePosition] = {
+    if (isWorksheet(sourcefile)) Some(toUnwrappedPosition _)
+    else None
+  }
+
+  /**
+   * Is `sourceTree` the wrapper object that we put around worksheet sources?
+   *
+   * @see wrapWorksheet
+   */
+  def isWorksheetWrapper(sourceTree: SourceTree)(implicit ctx: Context): Boolean = {
+    val symbol = sourceTree.tree.symbol
+    isWorksheet(sourceTree.source) &&
+      symbol.name.toString == "Worksheet$" &&
+      symbol.owner == ctx.definitions.EmptyPackageClass
+  }
+
   /** Create an lsp4j.CompletionItem from a Symbol */
   def completionItem(sym: Symbol)(implicit ctx: Context): lsp4j.CompletionItem = {
     def completionItemKind(sym: Symbol)(implicit ctx: Context): lsp4j.CompletionItemKind = {
@@ -596,7 +638,7 @@ object DottyLanguageServer {
   }
 
   /** Create an lsp4j.SymbolInfo from a Symbol and a SourcePosition */
-  def symbolInfo(sym: Symbol, pos: SourcePosition)(implicit ctx: Context): Option[lsp4j.SymbolInformation] = {
+  def symbolInfo(sym: Symbol, pos: SourcePosition, positionMapper: Option[SourcePosition => SourcePosition])(implicit ctx: Context): Option[lsp4j.SymbolInformation] = {
     def symbolKind(sym: Symbol)(implicit ctx: Context): lsp4j.SymbolKind = {
       import lsp4j.{SymbolKind => SK}
 
@@ -621,6 +663,6 @@ object DottyLanguageServer {
       else
         null
 
-    location(pos).map(l => new lsp4j.SymbolInformation(name, symbolKind(sym), l, containerName))
+    location(pos, positionMapper).map(l => new lsp4j.SymbolInformation(name, symbolKind(sym), l, containerName))
   }
 }

language-server/test/dotty/tools/languageserver/DefinitionTest.scala

@@ -55,12 +55,6 @@ class DefinitionTest {
   }
 
   @Test def goToDefNamedArgOverload: Unit = {
-    val m9 = new CodeMarker("m9")
-    val m10 = new CodeMarker("m10")
-    val m11 = new CodeMarker("m11")
-    val m12 = new CodeMarker("m12")
-    val m13 = new CodeMarker("m13")
-    val m14 = new CodeMarker("m14")
 
     code"""object Foo {
              def foo(${m1}x${m2}: String): String = ${m3}x${m4}
@@ -88,10 +82,6 @@ class DefinitionTest {
   }
 
   @Test def goToConstructorNamedArgOverload: Unit = {
-    val m9 = new CodeMarker("m9")
-    val m10 = new CodeMarker("m10")
-    val m11 = new CodeMarker("m11")
-    val m12 = new CodeMarker("m12")
 
     withSources(
       code"""class Foo(${m1}x${m2}: String) {
@@ -110,8 +100,6 @@ class DefinitionTest {
   }
 
   @Test def goToParamCopyMethod: Unit = {
-    val m9 = new CodeMarker("m9")
-    val m10 = new CodeMarker("m10")
 
     withSources(
       code"""case class Foo(${m1}x${m2}: Int, ${m3}y${m4}: String)""",

language-server/test/dotty/tools/languageserver/WorksheetTest.scala

@@ -1,8 +1,10 @@
 package dotty.tools.languageserver
 
 import org.junit.Test
+import org.eclipse.lsp4j.{CompletionItemKind, DocumentHighlightKind, SymbolKind}
 
 import dotty.tools.languageserver.util.Code._
+import dotty.tools.languageserver.util.embedded.CodeMarker
 
 class WorksheetTest {
 
@@ -82,4 +84,104 @@ class WorksheetTest {
       .evaluate(m1, "1:val foo: Int = 1\nval bar: Int = 2")
   }
 
+  @Test def worksheetCompletion(): Unit = {
+    ws"""class Foo { def bar = 123 }
+         val x = new Foo
+         x.b${m1}""".withSource
+      .completion(m1, Set(("bar", CompletionItemKind.Method, "=> Int")))
+  }
+
+  @Test def worksheetGoToDefinition(): Unit = {
+
+    withSources(
+      code"""class ${m11}Baz${m12}""",
+      ws"""class ${m1}Foo${m2} { def ${m3}bar${m4} = new ${m5}Baz${m6}  }
+           val x = new ${m7}Foo${m8}
+           x.${m9}bar${m10}"""
+     ).definition(m1 to m2, List(m1 to m2))
+      .definition(m3 to m4, List(m3 to m4))
+      .definition(m5 to m6, List(m11 to m12))
+      .definition(m7 to m8, List(m1 to m2))
+      .definition(m9 to m10, List(m3 to m4))
+      .definition(m11 to m12, List(m11 to m12))
+  }
+
+  @Test def worksheetReferences(): Unit = {
+
+    withSources(
+      code"""class ${m11}Baz${m12}""",
+      ws"""class ${m1}Foo${m2} { def ${m3}bar${m4} = new ${m9}Baz${m10}  }
+           val x = new ${m5}Foo${m6}
+           x.${m7}bar${m8}"""
+    ).references(m1 to m2, List(m5 to m6))
+     .references(m3 to m4, List(m7 to m8))
+     .references(m11 to m12, List(m9 to m10))
+  }
+
+  @Test def worksheetRename(): Unit = {
+
+    def sources =
+      withSources(
+        code"""class ${m9}Baz${m10}""",
+        ws"""class ${m1}Foo${m2}(baz: ${m3}Baz${m4})
+             val x = new ${m5}Foo${m6}(new ${m7}Baz${m8})"""
+      )
+
+    def testRenameFooFrom(m: CodeMarker) =
+      sources.rename(m, "Bar", Set(m1 to m2, m5 to m6))
+
+    def testRenameBazFrom(m: CodeMarker) =
+      sources.rename(m, "Bar", Set(m3 to m4, m7 to m8, m9 to m10))
+
+    testRenameFooFrom(m1)
+    testRenameBazFrom(m3)
+    testRenameFooFrom(m5)
+    testRenameBazFrom(m7)
+    testRenameBazFrom(m9)
+  }
+
+  @Test def worksheetHighlight(): Unit = {
+    ws"""class ${m1}Foo${m2} { def ${m3}bar${m4} = 123  }
+         val x = new ${m5}Foo${m6}
+         x.${m7}bar${m8}""".withSource
+      .highlight(m1 to m2, (m1 to m2, DocumentHighlightKind.Read), (m5 to m6, DocumentHighlightKind.Read))
+      .highlight(m3 to m4, (m3 to m4, DocumentHighlightKind.Read), (m7 to m8, DocumentHighlightKind.Read))
+      .highlight(m5 to m6, (m1 to m2, DocumentHighlightKind.Read), (m5 to m6, DocumentHighlightKind.Read))
+      .highlight(m7 to m8, (m3 to m4, DocumentHighlightKind.Read), (m7 to m8, DocumentHighlightKind.Read))
+  }
+
+  def hoverContent(typeInfo: String, comment: String): Option[String] =
+    Some(s"""```scala
+            |$typeInfo
+            |$comment
+            |```""".stripMargin)
+  @Test def worksheetHover(): Unit = {
+    ws"""/** A class */ class ${m1}Foo${m2} { /** A method */ def ${m3}bar${m4} = 123  }
+         val x = new ${m5}Foo${m6}
+         x.${m7}bar${m8}""".withSource
+      .hover(m1 to m2, hoverContent("Worksheet.Foo", "/** A class */"))
+      .hover(m3 to m4, hoverContent("Int", "/** A method */"))
+      .hover(m5 to m6, hoverContent("Worksheet.Foo", "/** A class */"))
+      .hover(m7 to m8, hoverContent("Int", "/** A method */"))
+  }
+
+  @Test def worksheetDocumentSymbol(): Unit = {
+    ws"""class ${m1}Foo${m2} {
+           def ${m3}bar${m4} = 123
+         }""".withSource
+      .documentSymbol(m1, (m1 to m2).symInfo("Foo", SymbolKind.Class, "Worksheet$"),
+                          (m3 to m4).symInfo("bar", SymbolKind.Method, "Foo"))
+  }
+
+  @Test def worksheetSymbol(): Unit = {
+    withSources(
+      ws"""class ${m1}Foo${m2} {
+             def ${m3}bar${m4} = 123
+           }""",
+      code"""class ${m5}Baz${m6}"""
+    ).symbol("Foo", (m1 to m2).symInfo("Foo", SymbolKind.Class, "Worksheet$"))
+     .symbol("bar", (m3 to m4).symInfo("bar", SymbolKind.Method, "Foo"))
+     .symbol("Baz", (m5 to m6).symInfo("Baz", SymbolKind.Class))
+  }
+
 }

language-server/test/dotty/tools/languageserver/util/Code.scala

@@ -20,6 +20,12 @@ object Code {
   val m6 = new CodeMarker("m6")
   val m7 = new CodeMarker("m7")
   val m8 = new CodeMarker("m8")
+  val m9 = new CodeMarker("m9")
+  val m10 = new CodeMarker("m10")
+  val m11 = new CodeMarker("m11")
+  val m12 = new CodeMarker("m12")
+  val m13 = new CodeMarker("m13")
+  val m14 = new CodeMarker("m14")
 
   implicit class CodeHelper(val sc: StringContext) extends AnyVal {