get file content on demand

metals-bench/src/main/scala/bench/Inflated.scala

@@ -7,9 +7,12 @@ import scala.meta.internal.io.FileIO
 import scala.meta.io.AbsolutePath
 import scala.meta.io.Classpath
 
-case class Inflated(inputs: List[(Input.VirtualFile, AbsolutePath)], linesOfCode: Long) {
+case class Inflated(
+    inputs: List[(Input.VirtualFile, AbsolutePath)],
+    linesOfCode: Long,
+) {
   def filter(f: Input.VirtualFile => Boolean): Inflated = {
-    val newInputs = inputs.filter{case (input, _) => f(input)}
+    val newInputs = inputs.filter { case (input, _) => f(input) }
     val newLinesOfCode = newInputs.foldLeft(0) { case (accum, input) =>
       accum + input._1.text.linesIterator.length
     }
@@ -19,7 +22,7 @@ case class Inflated(inputs: List[(Input.VirtualFile, AbsolutePath)], linesOfCode
     Inflated(other.inputs ++ inputs, other.linesOfCode + linesOfCode)
 
   def foreach(f: Input.VirtualFile => Unit): Unit =
-    inputs.foreach{ case (file, _) => f(file)}
+    inputs.foreach { case (file, _) => f(file) }
 }
 
 object Inflated {

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

@@ -1,5 +1,6 @@
 package scala.meta.internal.metals
 
+import java.net.URI
 import java.nio.file.Path
 import java.nio.file.Paths
 import java.util.Collections
@@ -29,14 +30,14 @@ import scala.meta.internal.pc.ScalaPresentationCompiler
 import scala.meta.internal.worksheets.WorksheetPcData
 import scala.meta.internal.worksheets.WorksheetProvider
 import scala.meta.io.AbsolutePath
+import scala.meta.pc
 import scala.meta.pc.AutoImportsResult
 import scala.meta.pc.CancelToken
 import scala.meta.pc.HoverSignature
 import scala.meta.pc.OffsetParams
 import scala.meta.pc.PresentationCompiler
 import scala.meta.pc.SymbolSearch
 import scala.meta.pc.SyntheticDecoration
-import scala.meta.pc.VirtualFileParams
 
 import ch.epfl.scala.bsp4j.BuildTargetIdentifier
 import ch.epfl.scala.bsp4j.CompileReport
@@ -47,6 +48,7 @@ import org.eclipse.lsp4j.CompletionParams
 import org.eclipse.lsp4j.Diagnostic
 import org.eclipse.lsp4j.DocumentHighlight
 import org.eclipse.lsp4j.InitializeParams
+import org.eclipse.lsp4j.Location
 import org.eclipse.lsp4j.ReferenceParams
 import org.eclipse.lsp4j.RenameParams
 import org.eclipse.lsp4j.SelectionRange
@@ -88,6 +90,18 @@ class Compilers(
     extends Cancelable {
   val plugins = new CompilerPlugins()
 
+  val pcBuffers: pc.Buffers = new pc.Buffers {
+    override def getFile(
+        uri: URI,
+        scalaVersion: String,
+    ): ju.Optional[pc.PcAdjustFileParams] =
+      Try(sourceAdjustments(uri.toString(), scalaVersion)).toOption
+        .map[pc.PcAdjustFileParams] { case (vFile, _, adjust) =>
+          PcAdjustFileParams(CompilerVirtualFileParams(uri, vFile.text), adjust)
+        }
+        .asJava
+  }
+
   // Not a TrieMap because we want to avoid loading duplicate compilers for the same build target.
   // Not a `j.u.c.ConcurrentHashMap` because it can deadlock in `computeIfAbsent` when the absent
   // function is expensive, which is the case here.
@@ -687,36 +701,20 @@ class Compilers(
 
   def references(
       params: ReferenceParams,
-      targetFiles: Iterator[AbsolutePath],
+      targetFiles: List[AbsolutePath],
       token: CancelToken,
   ): Future[List[ReferencesResult]] = {
-    withPCAndAdjustLsp(params) { (pc, pos, adjust) =>
-      val targets = targetFiles.map { target =>
-        target.toURI.toString -> {
-          val (vFile, _, adjustLsp) =
-            sourceAdjustments(
-              target.toURI.toString(),
-              pc.scalaVersion(),
-            )
-          val params =
-            CompilerVirtualFileParams(target.toURI, vFile.text, token)
-          (params, adjustLsp)
-        }
-      }.toMap
-      val targetFilesParams: List[VirtualFileParams] =
-        targets.values.map(_._1).toList
-      pc.references(
+    withPCAndAdjustLsp(params) { (pc, pos, _) =>
+      val request = PcReferencesRequest(
         CompilerOffsetParamsUtils.fromPos(pos, token),
-        targetFilesParams.asJava,
         params.getContext().isIncludeDeclaration(),
-      ).asScala
+        targetFiles.map(_.toURI).asJava,
+      )
+      pc.references(request)
+        .asScala
         .map(
           _.asScala.toList.map { defRes =>
-            val locations = defRes
-              .locations()
-              .asScala
-              .toList
-              .map(loc => targets(loc.getUri())._2.adjustLocation(loc))
+            val locations = defRes.locations().asScala.toList
             ReferencesResult(defRes.symbol(), locations)
           }
         )
@@ -1182,6 +1180,7 @@ class Compilers(
       .withWorkspace(workspace.toNIO)
       .withScheduledExecutorService(sh)
       .withReportsLoggerLevel(MetalsServerConfig.default.loglevel)
+      .withBuffers(pcBuffers)
       .withConfiguration {
         val options =
           InitializationOptions.from(initializeParams).compilerOptions
@@ -1332,3 +1331,11 @@ object Compilers {
     case object Default extends PresentationCompilerKey
   }
 }
+
+case class PcAdjustFileParams(
+    params: pc.VirtualFileParams,
+    adjustLsp: AdjustLspData,
+) extends pc.PcAdjustFileParams {
+  override def adjustLocation(location: Location): Location =
+    adjustLsp.adjustLocation(location)
+}

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

@@ -37,7 +37,7 @@ class IdentifierIndex {
 
   def collectIdentifiers(
       text: String,
-      dialect: Dialect
+      dialect: Dialect,
   ): Iterable[String] = {
     val identifiers = Set.newBuilder[String]
 

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

@@ -1,7 +1,10 @@
 package scala.meta.internal.metals
 
+import java.net.URI
 import java.nio.charset.StandardCharsets
 import java.nio.file.Path
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.TimeoutException
 
 import scala.collection.concurrent.TrieMap
 import scala.concurrent.ExecutionContext
@@ -11,6 +14,7 @@ import scala.util.matching.Regex
 
 import scala.meta.Importee
 import scala.meta.inputs.Input
+import scala.meta.internal.async.CompletableCancelToken
 import scala.meta.internal.metals.MetalsEnrichments._
 import scala.meta.internal.metals.ResolvedSymbolOccurrence
 import scala.meta.internal.mtags.DefinitionAlternatives.GlobalSymbol
@@ -26,6 +30,8 @@ import scala.meta.internal.semanticdb.TextDocument
 import scala.meta.internal.semanticdb.TextDocuments
 import scala.meta.internal.{semanticdb => s}
 import scala.meta.io.AbsolutePath
+import scala.meta.pc.OffsetParams
+import scala.meta.pc.ReferencesRequest
 import scala.meta.tokens.Token.Ident
 
 import ch.epfl.scala.bsp4j.BuildTargetIdentifier
@@ -50,7 +56,9 @@ final class ReferenceProvider(
   val identifierIndex: IdentifierIndex = new IdentifierIndex
 
   def addIdentifiers(file: AbsolutePath, set: Iterable[String]): Unit =
-    buildTargets.inverseSources(file).map(id => identifierIndex.addIdentifiers(file, id, set))
+    buildTargets
+      .inverseSources(file)
+      .map(id => identifierIndex.addIdentifiers(file, id, set))
 
   def indexIdentifiers(
       path: AbsolutePath,
@@ -304,8 +312,15 @@ final class ReferenceProvider(
         pathsForName(buildTarget, localPath, name).filter(filterTargetFiles)
       }
       if targetFiles.nonEmpty
+      cancelToken = new CompletableCancelToken
     } yield compilers
-      .references(params, targetFiles, EmptyCancelToken)
+      .references(params, targetFiles.toList, cancelToken)
+      .withTimeout(30, TimeUnit.SECONDS)
+      .recover { case _: TimeoutException =>
+        cancelToken.cancel()
+        scribe.warn("pc references search timed out after 30 seconds")
+        Nil
+      }
     result.getOrElse(Future.successful(Nil))
   }
 
@@ -445,7 +460,7 @@ final class ReferenceProvider(
     val isLocal = occ.symbol.isLocal
     if (isLocal)
       compilers
-        .references(params, Iterator(source), EmptyCancelToken)
+        .references(params, List(source), EmptyCancelToken)
         .map(_.flatMap(_.locations))
     else {
       /* search local in the following cases:
@@ -663,3 +678,9 @@ object SyntheticPackageObject {
   def unapply(str: String): Option[String] =
     Option.when(regex.matches(str))(str)
 }
+
+case class PcReferencesRequest(
+    params: OffsetParams,
+    includeDefinition: Boolean,
+    targetUris: java.util.List[URI],
+) extends ReferencesRequest

mtags-interfaces/src/main/java/scala/meta/pc/Buffers.java

@@ -0,0 +1,9 @@
+package scala.meta.pc;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.Optional;
+import java.net.URI;
+
+public interface Buffers {
+  Optional<PcAdjustFileParams> getFile(URI uri, String scalaVersion);
+}

mtags-interfaces/src/main/java/scala/meta/pc/PcAdjustFileParams.java

@@ -0,0 +1,8 @@
+package scala.meta.pc;
+
+import org.eclipse.lsp4j.Location;
+
+public interface PcAdjustFileParams {
+  VirtualFileParams params();
+  Location adjustLocation(Location location);
+}

mtags-interfaces/src/main/java/scala/meta/pc/PresentationCompiler.java

@@ -111,10 +111,14 @@ public CompletableFuture<List<Node>> semanticTokens(VirtualFileParams params) {
 	 */
 	public abstract CompletableFuture<java.util.List<DocumentHighlight>> documentHighlight(OffsetParams params);
 
+	public PresentationCompiler withBuffers(Buffers buffers) {
+		return this;
+	}
+
 	/**
 	 * Returns the references of the symbol under the current position in the target files.
 	 */
-	public CompletableFuture<java.util.List<DefinitionResult>> references(OffsetParams params, java.util.List<VirtualFileParams> targetFiles, boolean includeDefinition) {
+	public CompletableFuture<java.util.List<ReferencesResult>> references(ReferencesRequest params) {
 		return CompletableFuture.completedFuture(Collections.emptyList());
 	}
 

mtags-interfaces/src/main/java/scala/meta/pc/ReferencesRequest.java

@@ -0,0 +1,10 @@
+package scala.meta.pc;
+
+import java.util.List;
+import java.net.URI;
+
+public interface ReferencesRequest {
+  OffsetParams params();
+  boolean includeDefinition();
+  List<URI> targetUris();
+}

mtags-interfaces/src/main/java/scala/meta/pc/ReferencesResult.java

@@ -0,0 +1,9 @@
+package scala.meta.pc;
+
+import java.util.List;
+import org.eclipse.lsp4j.Location;
+
+public interface ReferencesResult {
+    String symbol();
+    List<Location> locations();
+}

mtags-java/src/main/scala/scala/meta/internal/pc/JavaPresentationCompiler.scala

@@ -20,6 +20,8 @@ import scala.meta.pc.OffsetParams
 import scala.meta.pc.PresentationCompiler
 import scala.meta.pc.PresentationCompilerConfig
 import scala.meta.pc.RangeParams
+import scala.meta.pc.ReferencesRequest
+import scala.meta.pc.ReferencesResult
 import scala.meta.pc.SymbolSearch
 import scala.meta.pc.SyntheticDecoration
 import scala.meta.pc.SyntheticDecorationsParams
@@ -99,10 +101,8 @@ case class JavaPresentationCompiler(
     CompletableFuture.completedFuture(Nil.asJava)
 
   override def references(
-      params: OffsetParams,
-      targetFiles: util.List[VirtualFileParams],
-      includeDefinition: Boolean
-  ): CompletableFuture[util.List[DefinitionResult]] =
+      params: ReferencesRequest
+  ): CompletableFuture[util.List[ReferencesResult]] =
     CompletableFuture.completedFuture(Nil.asJava)
 
   override def getTasty(

mtags-shared/src/main/scala/scala/meta/internal/pc/NoopBuffers.scala

@@ -0,0 +1,15 @@
+package scala.meta.internal.pc
+
+import java.net.URI
+import java.util.Optional
+
+import scala.meta.pc.Buffers
+import scala.meta.pc.PcAdjustFileParams
+
+object NoopBuffers extends Buffers {
+  override def getFile(
+      uri: URI,
+      scalaVersion: String
+  ): Optional[PcAdjustFileParams] = Optional.empty()
+
+}

mtags-shared/src/main/scala/scala/meta/internal/pc/ReferencesResultImpl.scala

@@ -0,0 +1,17 @@
+package scala.meta.internal.pc
+
+import java.{util => ju}
+
+import scala.meta.pc.ReferencesResult
+
+import org.eclipse.lsp4j.Location
+
+case class ReferencesResultImpl(
+    symbol: String,
+    locations: ju.List[Location]
+) extends ReferencesResult
+
+object ReferencesResultImpl {
+  def empty: ReferencesResult =
+    ReferencesResultImpl("", ju.Collections.emptyList())
+}