Worksheet improvements

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

@@ -107,7 +107,7 @@ class DottyLanguageServer extends LanguageServer
     if (Memory.isCritical()) CompletableFutures.computeAsync { _ => restart() }
 
   /** The driver instance responsible for compiling `uri` */
-  def driverFor(uri: URI): InteractiveDriver = {
+  def driverFor(uri: URI): InteractiveDriver = thisServer.synchronized {
     val matchingConfig =
       drivers.keys.find(config => config.sourceDirectories.exists(sourceDir =>
         new File(uri.getPath).getCanonicalPath.startsWith(sourceDir.getCanonicalPath)))
@@ -133,10 +133,10 @@ class DottyLanguageServer extends LanguageServer
     CompletableFuture.completedFuture(new Object)
   }
 
-  def computeAsync[R](fun: CancelChecker => R): CompletableFuture[R] =
+  def computeAsync[R](fun: CancelChecker => R, synchronize: Boolean = true): CompletableFuture[R] =
     CompletableFutures.computeAsync { cancelToken =>
       // We do not support any concurrent use of the compiler currently.
-      thisServer.synchronized {
+      def computation(): R = {
         cancelToken.checkCanceled()
         checkMemory()
         try {
@@ -147,6 +147,10 @@ class DottyLanguageServer extends LanguageServer
             throw ex
         }
       }
+      if (synchronize)
+        thisServer.synchronized { computation() }
+      else
+        computation()
     }
 
   override def initialize(params: InitializeParams) = computeAsync { cancelToken =>
@@ -202,10 +206,6 @@ class DottyLanguageServer extends LanguageServer
     val uri = new URI(document.getUri)
     val worksheetMode = isWorksheet(uri)
 
-    if (worksheetMode) {
-      Option(worksheets.get(uri)).foreach(_.cancel(true))
-    }
-
     thisServer.synchronized  {
       checkMemory()
 

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

@@ -73,7 +73,7 @@ object Main {
         .setInput(in)
         .setOutput(out)
         // For debugging JSON messages:
-        // .traceMessages(new java.io.PrintWriter(System.err, true))
+        //.traceMessages(new java.io.PrintWriter(System.err, true))
         .create();
 
     val client = launcher.getRemoteProxy()

language-server/src/dotty/tools/languageserver/worksheet/Evaluator.scala

@@ -28,7 +28,7 @@ private object Evaluator {
    * @param cancelChecker The token that indicates whether evaluation has been cancelled.
    * @return A JVM running the REPL.
    */
-  def get(cancelChecker: CancelChecker)(implicit ctx: Context): Option[Evaluator] = {
+  def get(cancelChecker: CancelChecker)(implicit ctx: Context): Option[Evaluator] = synchronized {
     val classpath = ctx.settings.classpath.value
     previousEvaluator match {
       case Some(cp, evaluator) if evaluator.isAlive() && cp == classpath =>

language-server/src/dotty/tools/languageserver/worksheet/Worksheet.scala

@@ -10,60 +10,65 @@ import dotty.tools.dotc.core.Flags.Synthetic
 
 import org.eclipse.lsp4j.jsonrpc.CancelChecker
 
-import java.util.concurrent.CancellationException
-
 object Worksheet {
 
   /**
    * Run `tree` as a worksheet using the REPL.
    *
    * @param tree          The top level object wrapping the worksheet.
+   * @param treeLock      Object on which to lock when doing operations on trees.
    * @param sendMessage   A mean of communicating the results of evaluation back.
    * @param cancelChecker A token to check whether execution should be cancelled.
    */
   def run(tree: SourceTree,
+          treeLock: Object,
           sendMessage: (Int, String) => Unit,
           cancelChecker: CancelChecker)(
-      implicit ctx: Context): Unit = synchronized {
-
-    Evaluator.get(cancelChecker) match {
-      case None =>
-        sendMessage(1, "Couldn't start JVM.")
-      case Some(evaluator) =>
-        tree.tree match {
-          case td @ TypeDef(_, template: Template) =>
-            val executed = collection.mutable.Set.empty[(Int, Int)]
-
-            template.body.foreach {
-              case statement: DefTree if statement.symbol.is(Synthetic) =>
-                ()
+    implicit ctx: Context): Unit = {
+    // For now, don't try to run multiple evaluators in parallel, this would require
+    // changes to the logic of Evaluator.get among other things.
+    Evaluator.synchronized {
+      Evaluator.get(cancelChecker) match {
+        case None =>
+          sendMessage(1, "Couldn't start the JVM.")
+        case Some(evaluator) =>
+          val queries = treeLock.synchronized {
+            tree.tree match {
+              case td @ TypeDef(_, template: Template) =>
+                val seen = collection.mutable.Set.empty[(Int, Int)]
 
-              case statement if evaluator.isAlive() && executed.add(bounds(statement.pos)) =>
-                try {
-                  cancelChecker.checkCanceled()
-                  val (line, result) = execute(evaluator, statement, tree.source)
-                  if (result.nonEmpty) sendMessage(line, result)
-                } catch { case _: CancellationException => () }
-
-              case _ =>
-                ()
+                template.body.flatMap {
+                  case statement: DefTree if statement.symbol.is(Synthetic) =>
+                    None
+                  case statement if seen.add(bounds(statement.pos)) =>
+                    Some(query(statement, tree.source))
+                  case _ =>
+                    None
+                }
             }
-        }
+          }
+          queries.foreach { (line, code) =>
+            cancelChecker.checkCanceled()
+            val res = evaluator.eval(code).getOrElse("")
+            cancelChecker.checkCanceled()
+            if (res.nonEmpty)
+              sendMessage(line, res)
+          }
+      }
     }
   }
 
   /**
-   * Extract `tree` from the source and evaluate it in the REPL.
+   * Extract the line number and source code corresponding to this tree
    *
    * @param evaluator  The JVM that runs the REPL.
    * @param tree       The compiled tree to evaluate.
    * @param sourcefile The sourcefile of the worksheet.
-   * @return The line in the sourcefile that corresponds to `tree`, and the result.
    */
-  private def execute(evaluator: Evaluator, tree: Tree, sourcefile: SourceFile): (Int, String) = {
-    val source = sourcefile.content.slice(tree.pos.start, tree.pos.end).mkString
+  private def query(tree: Tree, sourcefile: SourceFile): (Int, String) = {
     val line = sourcefile.offsetToLine(tree.pos.end)
-    (line, evaluator.eval(source).getOrElse(""))
+    val source = sourcefile.content.slice(tree.pos.start, tree.pos.end).mkString
+    (line, source)
   }
 
   private def bounds(pos: Position): (Int, Int) = (pos.start, pos.end)

language-server/src/dotty/tools/languageserver/worksheet/WorksheetService.scala

@@ -13,28 +13,23 @@ import java.util.concurrent.{CompletableFuture, ConcurrentHashMap}
 @JsonSegment("worksheet")
 trait WorksheetService { thisServer: DottyLanguageServer =>
 
-  val worksheets: ConcurrentHashMap[URI, CompletableFuture[_]] = new ConcurrentHashMap()
-
   @JsonRequest
-  def run(params: WorksheetRunParams): CompletableFuture[WorksheetRunResult] = thisServer.synchronized {
-    val uri = new URI(params.textDocument.getUri)
-    val future =
-      computeAsync { cancelChecker =>
-        try {
-          val driver = driverFor(uri)
-          val sendMessage = (line: Int, msg: String) => client.publishOutput(WorksheetRunOutput(params.textDocument, line, msg))
-          runWorksheet(driver, uri, sendMessage, cancelChecker)(driver.currentCtx)
-          WorksheetRunResult(success = true)
-        } catch {
-          case _: Throwable =>
-            WorksheetRunResult(success = false)
-        } finally {
-          worksheets.remove(uri)
-        }
+  def run(params: WorksheetRunParams): CompletableFuture[WorksheetRunResult] =
+    computeAsync(synchronize = false, fun = { cancelChecker =>
+      val uri = new URI(params.textDocument.getUri)
+      try {
+        val driver = driverFor(uri)
+        val sendMessage =
+          (line: Int, msg: String) => client.publishOutput(WorksheetRunOutput(params.textDocument, line, msg))
+
+        runWorksheet(driver, uri, sendMessage, cancelChecker)(driver.currentCtx)
+        cancelChecker.checkCanceled()
+        WorksheetRunResult(success = true)
+      } catch {
+        case _: Throwable =>
+          WorksheetRunResult(success = false)
       }
-    worksheets.put(uri, future)
-    future
-  }
+    })
 
   /**
    * Run the worksheet at `uri`.
@@ -45,13 +40,13 @@ trait WorksheetService { thisServer: DottyLanguageServer =>
    * @param cancelChecker Token to check whether evaluation was cancelled
    */
   private def runWorksheet(driver: InteractiveDriver,
-                                uri: URI,
-                                sendMessage: (Int, String) => Unit,
-                                cancelChecker: CancelChecker)(
+                           uri: URI,
+                           sendMessage: (Int, String) => Unit,
+                           cancelChecker: CancelChecker)(
       implicit ctx: Context): Unit = {
-    val trees = driver.openedTrees(uri)
-    trees.headOption.foreach { tree =>
-      Worksheet.run(tree, sendMessage, cancelChecker)
+    val treeOpt = thisServer.synchronized {
+      driver.openedTrees(uri).headOption
     }
+    treeOpt.foreach(tree => Worksheet.run(tree, thisServer, sendMessage, cancelChecker))
   }
 }

vscode-dotty/package.json

@@ -59,7 +59,7 @@
       },
       {
         "command": "dotty.worksheet.cancel",
-        "title": "Cancel worksheet evaluation",
+        "title": "Cancel running worksheet",
         "category": "Scala"
       }
     ],