improvement: Use dependencyModule for resolving sources

Previously, we would assume that the source jar is in the same directory as the normal jar. Now we use dependencyModules for that and only fallback to the previous logic if we can't find the jar in dependency modules.

This extracts part of the lazy classpath PR

Fixes #6182

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

@@ -3,6 +3,7 @@ package bench
 import java.util.concurrent.TimeUnit
 
 import scala.meta.dialects
+import scala.meta.internal.metals.BuildTargets
 import scala.meta.internal.metals.EmptyReportContext
 import scala.meta.internal.metals.MetalsEnrichments._
 import scala.meta.internal.tvp.IndexedSymbols
@@ -41,6 +42,7 @@ class ClasspathSymbolsBench {
       isStatisticsEnabled = false,
       trees,
       buffers,
+      BuildTargets.empty,
     )
     classpath.foreach { jar =>
       jars.jarSymbols(jar, "cats/", dialects.Scala213)

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

@@ -132,6 +132,10 @@ class BuildServerConnection private (
   def isDependencySourcesSupported: Boolean =
     capabilities.getDependencySourcesProvider()
 
+  // Scala CLI breaks when we try to use the `buildTarget/dependencyModules` request
+  def isDependencyModulesSupported: Boolean =
+    capabilities.getDependencyModulesProvider() && !isScalaCLI
+
   /* Currently only Bloop and sbt support running single test cases
    * and ScalaCLI uses Bloop underneath.
    */
@@ -389,6 +393,17 @@ class BuildServerConnection private (
       )
   }
 
+  def buildTargetDependencyModules(
+      params: DependencyModulesParams
+  ): Future[DependencyModulesResult] = {
+    if (isDependencyModulesSupported)
+      register(server => server.buildTargetDependencyModules(params)).asScala
+    else
+      Future.successful(
+        new DependencyModulesResult(List.empty[DependencyModulesItem].asJava)
+      )
+  }
+
   private val cancelled = new AtomicBoolean(false)
 
   override def cancel(): Unit = {

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

@@ -6,6 +6,7 @@ import scala.collection.mutable.ListBuffer
 
 import scala.meta.internal.metals.MetalsEnrichments._
 import scala.meta.internal.mtags.URIEncoderDecoder
+import scala.meta.io.AbsolutePath
 
 import ch.epfl.scala.bsp4j.BuildTarget
 import ch.epfl.scala.bsp4j.BuildTargetIdentifier
@@ -133,10 +134,16 @@ class BuildTargetInfo(buildTargets: BuildTargets) {
     val scalaClassPath = scalaInfo.map(_.fullClasspath).getOrElse(Nil)
     val javaClassPath = javaInfo.map(_.fullClasspath).getOrElse(Nil)
     if (scalaClassPath == javaClassPath)
-      output ++= getSection("Classpath", getClassPath(scalaClassPath))
+      output ++= getSection("Classpath", getClassPath(scalaClassPath, targetId))
     else {
-      output ++= getSection("Java Classpath", getClassPath(javaClassPath))
-      output ++= getSection("Scala Classpath", getClassPath(scalaClassPath))
+      output ++= getSection(
+        "Java Classpath",
+        getClassPath(javaClassPath, targetId),
+      )
+      output ++= getSection(
+        "Scala Classpath",
+        getClassPath(scalaClassPath, targetId),
+      )
     }
     output += ""
     output.mkString(System.lineSeparator())
@@ -158,23 +165,21 @@ class BuildTargetInfo(buildTargets: BuildTargets) {
   ): String =
     if (hasCapability) capability else s"$capability <- NOT SUPPORTED"
 
-  private def jarHasSource(jarName: String): Boolean = {
-    val sourceJarName = jarName.replace(".jar", "-sources.jar")
-    buildTargets
-      .sourceJarFile(sourceJarName)
-      .exists(_.toFile.exists())
-  }
-
   private def getSingleClassPathInfo(
       path: Path,
       shortPath: Path,
       maxFilenameSize: Int,
+      buildTargetId: BuildTargetIdentifier,
   ): String = {
     val filename = shortPath.toString()
     val padding = " " * (maxFilenameSize - filename.size)
     val status = if (path.toFile.exists) {
       val blankWarning = " " * 9
-      if (path.toFile().isDirectory() || jarHasSource(filename))
+      if (
+        path.toFile().isDirectory() || buildTargets
+          .sourceJarFor(buildTargetId, AbsolutePath(path))
+          .nonEmpty
+      )
         blankWarning
       else
         "NO SOURCE"
@@ -184,7 +189,8 @@ class BuildTargetInfo(buildTargets: BuildTargets) {
   }
 
   private def getClassPath(
-      classPath: List[Path]
+      classPath: List[Path],
+      buildTargetId: BuildTargetIdentifier,
   ): List[String] = {
     def shortenPath(path: Path): Path = {
       if (path.toFile.isFile)
@@ -200,6 +206,7 @@ class BuildTargetInfo(buildTargets: BuildTargets) {
           path,
           shortenPath(path),
           maxFilenameSize,
+          buildTargetId,
         )
       )
     } else

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

@@ -348,8 +348,7 @@ final class BuildTargets private (
   ): List[BuildTargetIdentifier] = {
     if (source.isJarFileSystem) {
       for {
-        jarName <- source.jarPath.map(_.filename).toList
-        sourceJarFile <- sourceJarFile(jarName).toList
+        sourceJarFile <- source.jarPath.toList
         buildTargetId <- inverseDependencySource(sourceJarFile)
       } yield buildTargetId
     } else {
@@ -402,6 +401,7 @@ final class BuildTargets private (
       .find(_.getDisplayName() == name)
   }
 
+  @deprecated("Jar and source jar might not always be in the same directory")
   private def jarPath(source: AbsolutePath): Option[AbsolutePath] = {
     source.jarPath.map { sourceJarPath =>
       sourceJarPath.parent.resolve(
@@ -410,6 +410,24 @@ final class BuildTargets private (
     }
   }
 
+  /**
+   * Try to resolve source jar for a jar, this should not be use
+   * in other capacity than as a fallback, since both source jar
+   * and normal jar might not be in the same directory.
+   *
+   * @param sourceJarPath path to the nromaljar
+   * @return path to the source jar for that jar
+   */
+  private def sourceJarPathFallback(
+      sourceJarPath: AbsolutePath
+  ): Option[AbsolutePath] = {
+    val fallback = sourceJarPath.parent.resolve(
+      sourceJarPath.filename.replace(".jar", "-sources.jar")
+    )
+    if (fallback.exists) Some(fallback)
+    else None
+  }
+
   /**
    * Returns meta build target for `*.sbt` or `*.scala`  files.
    * It selects build target by directory of its connection
@@ -439,6 +457,7 @@ final class BuildTargets private (
       jar: AbsolutePath,
       symbol: String,
       id: BuildTargetIdentifier,
+      sourceJar: Option[AbsolutePath],
   )
   def inferBuildTarget(
       toplevels: Iterable[Symbol]
@@ -469,7 +488,12 @@ final class BuildTargets private (
           val path = resource.toAbsolutePath
           classpaths.collectFirst {
             case (id, classpath) if classpath.contains(path) =>
-              InferredBuildTarget(path, toplevel.value, id)
+              InferredBuildTarget(
+                path,
+                toplevel.value,
+                id,
+                data.fromOptions(_.findSourceJarOf(path, Some(id))),
+              )
           }
       }
     } catch {
@@ -522,9 +546,29 @@ final class BuildTargets private (
     )
   }
 
+  @deprecated(
+    "This might return false positives since names of jars could repeat."
+  )
   def sourceJarFile(sourceJarName: String): Option[AbsolutePath] =
     data.fromOptions(_.sourceJarNameToJarFile.get(sourceJarName))
 
+  def sourceJarFor(
+      id: BuildTargetIdentifier,
+      jar: AbsolutePath,
+  ): Option[AbsolutePath] = {
+    data
+      .fromOptions(_.findSourceJarOf(jar, Some(id)))
+      .orElse(sourceJarPathFallback(jar))
+  }
+
+  def sourceJarFor(
+      jar: AbsolutePath
+  ): Option[AbsolutePath] = {
+    data
+      .fromOptions(_.findSourceJarOf(jar, targetId = None))
+      .orElse(sourceJarPathFallback(jar))
+  }
+
   def inverseDependencySource(
       sourceJar: AbsolutePath
   ): collection.Set[BuildTargetIdentifier] = {

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

@@ -21,6 +21,7 @@ case class ImportedBuild(
     sources: SourcesResult,
     dependencySources: DependencySourcesResult,
     wrappedSources: WrappedSourcesResult,
+    dependencyModules: DependencyModulesResult,
 ) {
   def ++(other: ImportedBuild): ImportedBuild = {
     val updatedBuildTargets = new WorkspaceBuildTargetsResult(
@@ -41,13 +42,17 @@ case class ImportedBuild(
     val updatedWrappedSources = new WrappedSourcesResult(
       (wrappedSources.getItems.asScala ++ other.wrappedSources.getItems.asScala).asJava
     )
+    val updatedDependencyModules = new DependencyModulesResult(
+      (dependencyModules.getItems.asScala ++ other.dependencyModules.getItems.asScala).asJava
+    )
     ImportedBuild(
       updatedBuildTargets,
       updatedScalacOptions,
       updatedJavacOptions,
       updatedSources,
       updatedDependencySources,
       updatedWrappedSources,
+      updatedDependencyModules,
     )
   }
 
@@ -66,6 +71,7 @@ object ImportedBuild {
       new SourcesResult(ju.Collections.emptyList()),
       new DependencySourcesResult(ju.Collections.emptyList()),
       new WrappedSourcesResult(ju.Collections.emptyList()),
+      new DependencyModulesResult(ju.Collections.emptyList()),
     )
 
   def fromConnection(
@@ -93,6 +99,9 @@ object ImportedBuild {
       wrappedSources <- conn.buildTargetWrappedSources(
         new WrappedSourcesParams(ids)
       )
+      dependencyModules <- conn.buildTargetDependencyModules(
+        new DependencyModulesParams(ids)
+      )
     } yield {
       ImportedBuild(
         workspaceBuildTargets,
@@ -101,6 +110,7 @@ object ImportedBuild {
         sources,
         dependencySources,
         wrappedSources,
+        dependencyModules,
       )
     }
 

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

@@ -226,6 +226,10 @@ final case class Indexer(
           bspSession().map(_.mainConnection),
         )
 
+        data.addDependencyModules(
+          importedBuild.dependencyModules
+        )
+
         // For "wrapped sources", we create dedicated TargetData.MappedSource instances,
         // able to convert back and forth positions from the user-facing file to
         // the compiler-facing actual underlying source.

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

@@ -131,6 +131,15 @@ object MetalsEnrichments
     }
   }
 
+  implicit class XtensionDependencyModule(module: b.DependencyModule) {
+    def asMavenDependencyModule: Option[b.MavenDependencyModule] = {
+      if (module.getDataKind() == b.DependencyModuleDataKind.MAVEN)
+        decodeJson(module.getData, classOf[b.MavenDependencyModule])
+      else
+        None
+    }
+  }
+
   implicit class XtensionTaskStart(task: b.TaskStartParams) {
     def asCompileTask: Option[b.CompileTask] = {
       decodeJson(task.getData, classOf[b.CompileTask])

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

@@ -236,7 +236,7 @@ object StandaloneSymbolSearch {
 
     download(scalaVersion).toSeq
       .map(path => AbsolutePath(path))
-      .partition(_.toString.endsWith("-sources.jar"))
+      .partition(_.isSourcesJar)
   }
 
 }

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

@@ -9,14 +9,17 @@ import scala.collection.concurrent.TrieMap
 import scala.collection.mutable
 import scala.collection.mutable.ListBuffer
 import scala.collection.mutable.{Map => MMap}
+import scala.util.Properties
 
 import scala.meta.inputs.Input
 import scala.meta.internal.metals.MetalsEnrichments._
 import scala.meta.io.AbsolutePath
 
 import ch.epfl.scala.bsp4j.BuildTarget
 import ch.epfl.scala.bsp4j.BuildTargetIdentifier
+import ch.epfl.scala.bsp4j.DependencyModulesResult
 import ch.epfl.scala.bsp4j.JavacOptionsResult
+import ch.epfl.scala.bsp4j.MavenDependencyModule
 import ch.epfl.scala.bsp4j.ScalacOptionsResult
 import ch.epfl.scala.bsp4j.SourceItem
 import ch.epfl.scala.bsp4j.SourceItemKind.DIRECTORY
@@ -40,6 +43,9 @@ final class TargetData {
     TrieMap.empty[BuildTargetIdentifier, ListBuffer[BuildTargetIdentifier]]
   val buildTargetSources: MMap[BuildTargetIdentifier, util.Set[AbsolutePath]] =
     TrieMap.empty[BuildTargetIdentifier, util.Set[AbsolutePath]]
+  val buildTargetDependencyModules
+      : MMap[BuildTargetIdentifier, List[MavenDependencyModule]] =
+    TrieMap.empty[BuildTargetIdentifier, List[MavenDependencyModule]]
   val inverseDependencySources: MMap[AbsolutePath, Set[BuildTargetIdentifier]] =
     TrieMap.empty[AbsolutePath, Set[BuildTargetIdentifier]]
   val buildTargetGeneratedDirs: MMap[AbsolutePath, Unit] =
@@ -140,6 +146,35 @@ final class TargetData {
       .orElse(javacData)
   }
 
+  def findSourceJarOf(
+      jar: AbsolutePath,
+      targetId: Option[BuildTargetIdentifier],
+  ): Option[AbsolutePath] = {
+    val jarUri = jar.toURI.toString()
+    def depModules: Iterator[MavenDependencyModule] = targetId match {
+      case None => buildTargetDependencyModules.values.flatten.iterator
+      case Some(id) => buildTargetDependencyModules.get(id).iterator.flatten
+    }
+
+    /**
+     * For windows file:///C:/Users/runneradmin/AppData/Local/Coursier/Cache and
+     * file:///C:/Users/runneradmin/AppData/Local/Coursier/cache is equivalent
+     */
+    def isUriEqual(uri: String, otherUri: String) = {
+      Properties.isWin && uri.toLowerCase() == otherUri
+        .toLowerCase() || uri == otherUri
+    }
+    val allFound = for {
+      module <- depModules
+      artifacts = module.getArtifacts().asScala
+      if artifacts.exists(artifact => isUriEqual(artifact.getUri(), jarUri))
+      sourceJar <- artifacts.find(_.getClassifier() == "sources")
+      sourceJarPath = sourceJar.getUri().toAbsolutePath
+      if sourceJarPath.exists
+    } yield sourceJarPath
+    allFound.headOption
+  }
+
   def targetClassDirectories(id: BuildTargetIdentifier): List[String] = {
     val scalacData =
       scalaTarget(id).map(_.scalac.getClassDirectory).filter(_.nonEmpty).toList
@@ -291,6 +326,18 @@ final class TargetData {
   ): Unit =
     actualSources(path) = mapped
 
+  def addDependencyModules(
+      dependencyModules: DependencyModulesResult
+  ): Unit = {
+    dependencyModules.getItems().asScala.groupBy(_.getTarget()).foreach {
+      case (id, items) =>
+        val modules = items
+          .flatMap(_.getModules().asScala)
+          .flatMap(_.asMavenDependencyModule)
+        buildTargetDependencyModules.put(id, modules.toList)
+    }
+  }
+
   def resetConnections(
       idToConn: List[(BuildTargetIdentifier, BuildServerConnection)]
   ): Unit = {

metals/src/main/scala/scala/meta/internal/metals/debug/SourcePathAdapter.scala

@@ -26,7 +26,7 @@ private[debug] final class SourcePathAdapter(
       )
     ) {
       // if sourcePath is a dependency source file
-      // we retrieve the original source jar and we build the uri innside the source jar filesystem
+      // we retrieve the original source jar and we build the uri inside the source jar filesystem
       for {
         dependencySource <- sourcePath.toRelativeInside(dependencies)
         dependencyFolder <- dependencySource.toNIO.iterator.asScala.headOption