Keep code lenses after error until next successful compilation

metals/src/main/scala/scala/meta/internal/metals/BuildTargetClasses.scala

@@ -1,11 +1,9 @@
 package scala.meta.internal.metals
 
-import java.util.concurrent.CancellationException
 import ch.epfl.scala.{bsp4j => b}
 import scala.collection.concurrent.TrieMap
 import scala.concurrent.ExecutionContext
 import scala.concurrent.Future
-import scala.concurrent.Promise
 import scala.meta.internal.metals.BuildTargetClasses.Classes
 import scala.meta.internal.metals.MetalsEnrichments._
 import scala.meta.internal.semanticdb.Scala.Descriptor
@@ -17,18 +15,17 @@ import scala.meta.internal.semanticdb.Scala.Symbols
 final class BuildTargetClasses(
     buildServer: () => Option[BuildServerConnection]
 )(implicit val ec: ExecutionContext) {
-  private val index = TrieMap.empty[b.BuildTargetIdentifier, Promise[Classes]]
+  private val index = TrieMap.empty[b.BuildTargetIdentifier, Classes]
 
   val rebuildIndex: BatchedFunction[b.BuildTargetIdentifier, Unit] =
     BatchedFunction.fromFuture(fetchClasses)
 
-  def classesOf(target: b.BuildTargetIdentifier): Future[Classes] = {
-    index.getOrElseUpdate(target, Promise()).future
+  def classesOf(target: b.BuildTargetIdentifier): Classes = {
+    index.getOrElse(target, new Classes)
   }
 
   def invalidate(target: b.BuildTargetIdentifier): Unit = {
-    val previous = index.put(target, Promise())
-    previous.foreach(_.tryFailure(new CancellationException()))
+    index.put(target, new Classes)
   }
 
   private def fetchClasses(
@@ -55,10 +52,7 @@ final class BuildTargetClasses(
           _ <- updateTestClasses
         } yield {
           classes.foreach {
-            case (id, classes) =>
-              index
-                .getOrElseUpdate(id, Promise())
-                .success(classes)
+            case (id, classes) => index.put(id, classes)
           }
         }
 

metals/src/main/scala/scala/meta/internal/metals/CodeLensProvider.scala

@@ -5,7 +5,6 @@ import ch.epfl.scala.{bsp4j => b}
 import com.google.gson.JsonElement
 import org.eclipse.{lsp4j => l}
 import scala.concurrent.ExecutionContext
-import scala.concurrent.Future
 import scala.meta.internal.metals.ClientCommands.StartDebugSession
 import scala.meta.internal.metals.CodeLensProvider._
 import scala.meta.internal.metals.MetalsEnrichments._
@@ -20,17 +19,15 @@ final class CodeLensProvider(
     semanticdbs: Semanticdbs
 )(implicit ec: ExecutionContext) {
   // code lenses will be refreshed after compilation or when workspace gets indexed
-  def findLenses(path: AbsolutePath): Future[Seq[l.CodeLens]] = {
+  def findLenses(path: AbsolutePath): Seq[l.CodeLens] = {
     val lenses = buildTargets
       .inverseSources(path)
-      .filterNot(compilations.isCurrentlyCompiling)
       .map { buildTarget =>
-        for {
-          classes <- buildTargetClasses.classesOf(buildTarget)
-        } yield codeLenses(path, buildTarget, classes)
+        val classes = buildTargetClasses.classesOf(buildTarget)
+        codeLenses(path, buildTarget, classes)
       }
 
-    lenses.getOrElse(Future.successful(Nil))
+    lenses.getOrElse(Nil)
   }
 
   private def codeLenses(

metals/src/main/scala/scala/meta/internal/metals/Compilations.scala

@@ -8,12 +8,15 @@ import scala.meta.internal.metals.MetalsEnrichments._
 import scala.meta.io.AbsolutePath
 import scala.util.Failure
 import scala.util.Success
+import scala.util.Try
 
 final class Compilations(
     buildTargets: BuildTargets,
     classes: BuildTargetClasses,
     workspace: () => AbsolutePath,
-    buildServer: () => Option[BuildServerConnection]
+    buildServer: () => Option[BuildServerConnection],
+    languageClient: MetalsLanguageClient,
+    isCurrentlyFocused: b.BuildTargetIdentifier => Boolean
 )(implicit ec: ExecutionContext) {
 
   // we are maintaining a separate queue for cascade compilation since those must happen ASAP
@@ -92,14 +95,28 @@ final class Compilations(
     targets.foreach(target => isCompiling(target) = true)
     val compilation = connection.compile(params)
 
-    val result = compilation.asScala.andThen {
+    val result = compilation.asScala
+      .andThen {
+        case result =>
+          updateCompiledTargetState(result)
+
+          // See https://github.com/scalacenter/bloop/issues/1067
+          classes.rebuildIndex(targets).foreach { _ =>
+            if (targets.exists(isCurrentlyFocused)) {
+              languageClient.refreshModel()
+            }
+          }
+      }
+
+    CancelableFuture(result, Cancelable(() => compilation.cancel(false)))
+  }
+
+  private def updateCompiledTargetState(result: Try[b.CompileResult]): Unit =
+    result match {
       case Failure(_) =>
         isCompiling.clear()
       case Success(_) =>
         lastCompile = isCompiling.keySet
         isCompiling.clear()
     }
-
-    CancelableFuture(result, Cancelable(() => compilation.cancel(false)))
-  }
 }

metals/src/main/scala/scala/meta/internal/metals/ForwardingMetalsBuildClient.scala

@@ -27,8 +27,7 @@ final class ForwardingMetalsBuildClient(
     statusBar: StatusBar,
     time: Time,
     didCompile: CompileReport => Unit,
-    treeViewProvider: () => TreeViewProvider,
-    isCurrentlyFocused: b.BuildTargetIdentifier => Boolean
+    treeViewProvider: () => TreeViewProvider
 )(implicit ec: ExecutionContext)
     extends MetalsBuildClient
     with Cancelable {
@@ -141,10 +140,6 @@ final class ForwardingMetalsBuildClient(
             scribe.info(s"time: compiled $name in ${compilation.timer}")
           }
           if (isSuccess) {
-            buildTargetClasses.rebuildIndex(target).foreach { _ =>
-              if (isCurrentlyFocused(target)) languageClient.refreshModel()
-            }
-
             if (hasReportedError.contains(target)) {
               // Only report success compilation if it fixes a previous compile error.
               statusBar.addMessage(message)

metals/src/main/scala/scala/meta/internal/metals/MetalsLanguageServer.scala

@@ -117,7 +117,9 @@ class MetalsLanguageServer(
     buildTargets,
     buildTargetClasses,
     () => workspace,
-    () => buildServer
+    () => buildServer,
+    languageClient,
+    buildTarget => focusedDocumentBuildTarget.get() == buildTarget
   )
   private val fileWatcher = register(
     new FileWatcher(
@@ -248,8 +250,7 @@ class MetalsLanguageServer(
       statusBar,
       time,
       report => compilers.didCompile(report),
-      () => treeView,
-      buildTarget => focusedDocumentBuildTarget.get() == buildTarget
+      () => treeView
     )
     trees = new Trees(buffers, diagnostics)
     documentSymbolProvider = new DocumentSymbolProvider(trees)
@@ -1013,10 +1014,10 @@ class MetalsLanguageServer(
   def codeLens(
       params: CodeLensParams
   ): CompletableFuture[util.List[CodeLens]] =
-    CancelTokens.future { _ =>
-      codeLensProvider
-        .findLenses(params.getTextDocument.getUri.toAbsolutePath)
-        .map(_.asJava)
+    CancelTokens { _ =>
+      val path = params.getTextDocument.getUri.toAbsolutePath
+      val lenses = codeLensProvider.findLenses(path)
+      lenses.asJava
     }
 
   @JsonRequest("textDocument/foldingRange")
@@ -1316,6 +1317,7 @@ class MetalsLanguageServer(
   private def connectToNewBuildServer(
       build: BuildServerConnection
   ): Future[BuildChange] = {
+    scribe.info(s"Connected to Build server v${build.version}")
     cancelables.add(build)
     compilers.cancel()
     buildServer = Some(build)

tests/unit/src/main/scala/tests/BaseLspSuite.scala

@@ -85,6 +85,8 @@ abstract class BaseLspSuite(suiteName: String) extends BaseSuite {
       .resolve("e2e")
       .resolve(suiteName)
       .resolve(name.replace(' ', '-'))
+
+    RecursivelyDelete(path)
     Files.createDirectories(path.toNIO)
     path
   }

tests/unit/src/main/scala/tests/TestingClient.scala

@@ -35,6 +35,7 @@ import scala.meta.internal.builds.GradleBuildTool
 import scala.meta.internal.builds.SbtBuildTool
 import scala.meta.internal.builds.MavenBuildTool
 import scala.meta.internal.builds.MillBuildTool
+import scala.meta.internal.metals.ClientCommands
 import scala.meta.internal.metals.NoopLanguageClient
 import scala.meta.internal.tvp.TreeViewDidChangeParams
 
@@ -64,10 +65,22 @@ final class TestingClient(workspace: AbsolutePath, buffers: Buffers)
       None
   }
 
+  private var refreshedOnIndex = false
+  var refreshModelHandler: () => Unit = () => {}
+
   override def metalsExecuteClientCommand(
       params: ExecuteCommandParams
   ): Unit = {
     clientCommands.addLast(params)
+    params.getCommand match {
+      case ClientCommands.RefreshModel.id =>
+        if (refreshedOnIndex) {
+          refreshModelHandler()
+        } else {
+          refreshedOnIndex = true
+        }
+      case _ =>
+    }
   }
 
   override def applyEdit(

tests/unit/src/main/scala/tests/TestingServer.scala

@@ -78,6 +78,7 @@ import com.google.gson.JsonObject
 import com.google.gson.JsonPrimitive
 import scala.meta.internal.tvp.TreeViewProvider
 import org.eclipse.lsp4j.DocumentRangeFormattingParams
+import scala.concurrent.Promise
 import scala.meta.internal.metals.ServerCommands
 import scala.meta.internal.metals.debug.TestDebugger
 import scala.meta.internal.metals.DebugSession
@@ -459,13 +460,44 @@ final class TestingServer(
     }
   }
 
-  def codeLenses(filename: String): Future[String] = {
+  def codeLenses(filename: String)(maxRetries: Int): Future[String] = {
     val path = toPath(filename)
     val uri = path.toURI.toString
     val params = new CodeLensParams(new TextDocumentIdentifier(uri))
+
+    // see https://github.com/scalacenter/bloop/issues/1067
+    // because bloop does not notify us when we can access the main/test classes,
+    // we have to try until we finally get them.
+    // Following handler runs on the refresh-model notification from the server
+    // (basically once the compilation finishes and classes are fetched)
+    // it retries the compilation until we finally can get desired lenses
+    // or fails if it could nat be achieved withing [[maxRetries]] number of tries
+    var retries = maxRetries
+    val codeLenses = Promise[List[l.CodeLens]]()
+    val handler = { () =>
+      for {
+        lenses <- server.codeLens(params).asScala.map(_.asScala)
+      } {
+        if (lenses.nonEmpty) codeLenses.trySuccess(lenses.toList)
+        else if (retries > 0) {
+          retries -= 1
+          server.compilations.compileFiles(List(path))
+        } else {
+          val error = s"Could not fetch any code lenses in $maxRetries tries"
+          codeLenses.tryFailure(new NoSuchElementException(error))
+        }
+      }
+    }
+
     for {
-      lenses <- server.codeLens(params).asScala
-      textEdits = CodeLensesTextEdits(lenses.asScala)
+      _ <- server
+        .didFocus(uri)
+        .asScala // model is refreshed only for focused document
+      _ = client.refreshModelHandler = handler
+      // first compilation, to trigger the handler
+      _ <- server.compilations.compileFiles(List(path))
+      lenses <- codeLenses.future
+      textEdits = CodeLensesTextEdits(lenses)
     } yield TextEdits.applyEdits(textContents(filename), textEdits)
   }
 

tests/unit/src/test/scala/tests/CodeLensesLspSuite.scala

@@ -115,12 +115,7 @@ object CodeLensesLspSuite extends BaseLspSuite("codeLenses") {
            |object Main
            |""".stripMargin
       )
-      _ <- server.didOpen("a/src/main/scala/Main.scala") // compile `a` to populate its cache
-      _ <- assertCodeLenses(
-        "b/src/main/scala/Main.scala",
-        """|object Main
-           |""".stripMargin
-      )
+      _ <- assertNoCodeLenses("b/src/main/scala/Main.scala")
     } yield ()
   }
 
@@ -148,19 +143,46 @@ object CodeLensesLspSuite extends BaseLspSuite("codeLenses") {
       _ <- server.didSave("a/src/main/scala/Main.scala")(
         text => text.replace("object Main", "class Main")
       )
+      _ <- assertNoCodeLenses("a/src/main/scala/Main.scala")
+
+    } yield ()
+  }
+
+  testAsync("keep-after-error") {
+    for {
+      _ <- server.initialize(
+        """|/metals.json
+           |{ "a": { } }
+           |
+           |/a/src/main/scala/Main.scala
+           |object Main {
+           |  def main(args: Array[String]): Unit = ???
+           |}""".stripMargin
+      )
       _ <- assertCodeLenses(
         "a/src/main/scala/Main.scala",
-        """class Main {
-          |  def main(args: Array[String]): Unit = {}
+        """<<run>>
+          |object Main {
+          |  def main(args: Array[String]): Unit = ???
           |}
           |""".stripMargin
       )
-
+      _ <- server.didSave("a/src/main/scala/Main.scala")(
+        text => text.replace("}", "")
+      )
+      _ <- assertCodeLenses(
+        "a/src/main/scala/Main.scala",
+        """<<run>>
+          |object Main {
+          |  def main(args: Array[String]): Unit = ???
+          |
+          |""".stripMargin
+      )
     } yield ()
   }
 
   def check(name: String, library: String = "")(expected: String): Unit = {
-    ignore(name) {
+    testAsync(name) {
       val original = expected.replaceAll("<<.*>>\\W", "")
 
       val sourceFile = {
@@ -196,13 +218,24 @@ object CodeLensesLspSuite extends BaseLspSuite("codeLenses") {
 
   private def assertCodeLenses(
       relativeFile: String,
-      expected: String
+      expected: String,
+      maxRetries: Int = 4
   ): Future[Unit] = {
-    val path = server.toPath(relativeFile)
-    for {
-      _ <- server.server.compilations.compileFiles(List(path))
-      _ <- server.server.compilations.compileFiles(List(path))
-      obtained <- server.codeLenses(relativeFile)
-    } yield assertNoDiff(obtained, expected)
+    val obtained = server.codeLenses(relativeFile)(maxRetries).recover {
+      case _: NoSuchElementException =>
+        server.textContents(relativeFile)
+    }
+
+    obtained.map(assertNoDiff(_, expected))
+  }
+
+  private def assertNoCodeLenses(
+      relativeFile: String,
+      maxRetries: Int = 4
+  ): Future[Unit] = {
+    server.codeLenses(relativeFile)(maxRetries).failed.flatMap {
+      case _: NoSuchElementException => Future.unit
+      case e => Future.failed(e)
+    }
   }
 }