bugfix: properly resolve build targets for references search (scalame…

…ta#6103)

* bugfix: properly resolve build targets for references search

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

@@ -240,35 +240,65 @@ final class BuildTargets private (
     }
   }
 
+  /**
+   * Returns all build targets containing this source file.
+   */
+  def inverseSourcesAll(
+      source: AbsolutePath
+  ): List[BuildTargetIdentifier] = {
+    val buildTargets = sourceBuildTargets(source)
+    val orSbtBuildTarget =
+      buildTargets.getOrElse(sbtBuildScalaTarget(source).toIterable)
+    if (orSbtBuildTarget.isEmpty) {
+      inferBuildTargets(source)
+    } else orSbtBuildTarget.toList
+  }
+
   def inverseSourcesBsp(
       source: AbsolutePath
   )(implicit ec: ExecutionContext): Future[Option[BuildTargetIdentifier]] = {
     inverseSources(source) match {
       case None =>
-        val identifier = new TextDocumentIdentifier(
-          source.toTextDocumentIdentifier.getUri()
-        )
-        val params = new InverseSourcesParams(identifier)
-        val connections =
-          data.fromIterators(_.targetToConnection.values.toIterator).distinct
-        val queries = connections.map { connection =>
-          connection
-            .buildTargetInverseSources(params)
-            .map(_.getTargets.asScala.toList)
-        }
-        Future.sequence(queries).map { results =>
-          val target = results.flatten.maxByOption(buildTargetsOrder)
-          for {
-            tgt <- target
-            data <- targetData(tgt)
-          } data.addSourceItem(source, tgt)
-          target
-        }
+        bspInverseSources(source).map(_.maxByOption(buildTargetsOrder))
+      case some =>
+        Future.successful(some)
+    }
+  }
+
+  def inverseSourcesBspAll(
+      source: AbsolutePath
+  )(implicit ec: ExecutionContext): Future[List[BuildTargetIdentifier]] = {
+    inverseSourcesAll(source) match {
+      case Nil => bspInverseSources(source).map(_.toList)
       case some =>
         Future.successful(some)
     }
   }
 
+  private def bspInverseSources(
+      source: AbsolutePath
+  )(implicit ec: ExecutionContext) = {
+    val identifier = new TextDocumentIdentifier(
+      source.toTextDocumentIdentifier.getUri()
+    )
+    val params = new InverseSourcesParams(identifier)
+    val connections =
+      data.fromIterators(_.targetToConnection.values.toIterator).distinct
+    val queries = connections.map { connection =>
+      connection
+        .buildTargetInverseSources(params)
+        .map(_.getTargets.asScala.toList)
+    }
+    Future.sequence(queries).map { results =>
+      val targets = results.flatten
+      for {
+        tgt <- targets
+        data <- targetData(tgt)
+      } data.addSourceItem(source, tgt)
+      targets
+    }
+  }
+
   def scalaVersion(source: AbsolutePath): Option[String] = {
     for {
       id <- inverseSources(source)
@@ -306,14 +336,14 @@ final class BuildTargets private (
    *
    * This approach is not glamorous but it seems to work reasonably well.
    */
-  def inferBuildTarget(
+  def inferBuildTargets(
       source: AbsolutePath
-  ): Option[BuildTargetIdentifier] = {
+  ): List[BuildTargetIdentifier] = {
     if (source.isJarFileSystem) {
       for {
-        jarName <- source.jarPath.map(_.filename)
-        sourceJarFile <- sourceJarFile(jarName)
-        buildTargetId <- inverseDependencySource(sourceJarFile).headOption
+        jarName <- source.jarPath.map(_.filename).toList
+        sourceJarFile <- sourceJarFile(jarName).toList
+        buildTargetId <- inverseDependencySource(sourceJarFile)
       } yield buildTargetId
     } else {
       val readonly = workspace.resolve(Directories.readonly)
@@ -323,16 +353,18 @@ final class BuildTargets private (
           names match {
             case Directories.dependenciesName :: jarName :: _ =>
               // match build target by source jar name
-              sourceJarFile(jarName)
-                .flatMap(inverseDependencySource(_).headOption)
-            case _ => None
+              sourceJarFile(jarName).toList
+                .flatMap(inverseDependencySource(_))
+            case _ => Nil
           }
         case None =>
           // else it can be a source file inside a jar
-          val fromJar = jarPath(source)
+          val fromJar = jarPath(source).toList
             .flatMap { jar =>
-              allBuildTargetIdsInternal.find { case (_, id) =>
-                targetJarClasspath(id).exists(_.contains(jar))
+              allBuildTargetIdsInternal.collect {
+                case pair @ (_, id)
+                    if targetJarClasspath(id).exists(_.contains(jar)) =>
+                  pair
               }
             }
           fromJar.map { case (data0, id) =>
@@ -343,6 +375,11 @@ final class BuildTargets private (
     }
   }
 
+  def inferBuildTarget(
+      source: AbsolutePath
+  ): Option[BuildTargetIdentifier] =
+    inferBuildTargets(source).maxByOption(buildTargetsOrder)
+
   def findByDisplayName(name: String): Option[BuildTarget] = {
     data
       .fromIterators(_.buildTargetInfo.valuesIterator)

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

@@ -742,6 +742,7 @@ class PackageProvider(
         path.value,
         decl.symbols,
         isIncludeDeclaration = false,
+        sourceContainsDefinition = true,
       )
       .map { loc =>
         Reference(decl, loc.getRange(), loc.getUri())

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

@@ -262,22 +262,26 @@ final class ReferenceProvider(
    * Return all paths to files which contain at least one symbol from isSymbol set.
    */
   private def pathsFor(
-      buildTarget: BuildTargetIdentifier,
+      buildTargetSet: Set[BuildTargetIdentifier],
       isSymbol: Set[String],
   ): Iterator[AbsolutePath] = {
-    val allowedBuildTargets = buildTargets.allInverseDependencies(buildTarget)
-    val visited = scala.collection.mutable.Set.empty[AbsolutePath]
-    val result = for {
-      (path, entry) <- index.iterator
-      if allowedBuildTargets.contains(entry.id) &&
-        isSymbol.exists(entry.bloom.mightContain)
-      sourcePath = AbsolutePath(path)
-      if !visited(sourcePath)
-      _ = visited.add(sourcePath)
-      if sourcePath.exists
-    } yield sourcePath
-
-    result
+    if (buildTargetSet.isEmpty) Iterator.empty
+    else {
+      val allowedBuildTargets =
+        buildTargetSet.flatMap(buildTargets.allInverseDependencies)
+      val visited = scala.collection.mutable.Set.empty[AbsolutePath]
+      val result = for {
+        (path, entry) <- index.iterator
+        if allowedBuildTargets(entry.id) &&
+          isSymbol.exists(entry.bloom.mightContain)
+        sourcePath = AbsolutePath(path)
+        if !visited(sourcePath)
+        _ = visited.add(sourcePath)
+        if sourcePath.exists
+      } yield sourcePath
+
+      result
+    }
   }
 
   def workspaceReferences(
@@ -286,44 +290,61 @@ final class ReferenceProvider(
       isIncludeDeclaration: Boolean,
       findRealRange: AdjustRange = noAdjustRange,
       includeSynthetics: Synthetic => Boolean = _ => true,
+      sourceContainsDefinition: Boolean = false,
   ): Seq[Location] = {
-    buildTargets.inverseSources(source) match {
-      case None => Seq.empty
-      case Some(id) =>
-        val result = for {
-          sourcePath <- pathsFor(id, isSymbol)
-          semanticdb <-
-            semanticdbs
-              .textDocument(sourcePath)
-              .documentIncludingStale
-              .iterator
-          semanticdbDistance = buffers.tokenEditDistance(
-            sourcePath,
-            semanticdb.text,
-            trees,
-          )
-          uri = sourcePath.toURI.toString
-          reference <-
-            try {
-              referenceLocations(
-                semanticdb,
-                isSymbol,
-                semanticdbDistance,
-                uri,
-                isIncludeDeclaration,
-                findRealRange,
-                includeSynthetics,
-                sourcePath.isJava,
-              )
-            } catch {
-              case NonFatal(e) =>
-                // Can happen for example if the SemanticDB text is empty for some reason.
-                scribe.error(s"reference: $sourcePath", e)
-                Nil
+    val definitionPaths =
+      if (sourceContainsDefinition) Set(source)
+      else {
+        val foundDefinitionLocations =
+          isSymbol
+            .flatMap { sym =>
+              definition.destinationProvider
+                .definition(sym, Some(source))
+                .map(_.path)
             }
-        } yield reference
-        result.toSeq
-    }
+
+        if (foundDefinitionLocations.isEmpty) Set(source)
+        else foundDefinitionLocations
+      }
+
+    val definitionBuildTargets =
+      definitionPaths.flatMap { path =>
+        buildTargets.inverseSourcesAll(path).toSet
+      }
+
+    val result = for {
+      sourcePath <- pathsFor(definitionBuildTargets, isSymbol)
+      semanticdb <-
+        semanticdbs
+          .textDocument(sourcePath)
+          .documentIncludingStale
+          .iterator
+      semanticdbDistance = buffers.tokenEditDistance(
+        sourcePath,
+        semanticdb.text,
+        trees,
+      )
+      uri = sourcePath.toURI.toString
+      reference <-
+        try {
+          referenceLocations(
+            semanticdb,
+            isSymbol,
+            semanticdbDistance,
+            uri,
+            isIncludeDeclaration,
+            findRealRange,
+            includeSynthetics,
+            sourcePath.isJava,
+          )
+        } catch {
+          case NonFatal(e) =>
+            // Can happen for example if the SemanticDB text is empty for some reason.
+            scribe.error(s"reference: $sourcePath", e)
+            Nil
+        }
+    } yield reference
+    result.toSeq
   }
 
   /**
@@ -334,12 +355,8 @@ final class ReferenceProvider(
       isSymbol: Set[String],
   )(implicit ec: ExecutionContext): Future[Set[AbsolutePath]] = {
     buildTargets
-      .inverseSourcesBsp(source)
-      .map {
-        case None => Set.empty
-        case Some(id) =>
-          pathsFor(id, isSymbol).toSet
-      }
+      .inverseSourcesBspAll(source)
+      .map(buildTargets => pathsFor(buildTargets.toSet, isSymbol).toSet)
   }
 
   private def references(
@@ -380,15 +397,22 @@ final class ReferenceProvider(
       else Seq.empty
 
     val workspaceRefs =
-      if (!isLocal)
+      if (!isLocal) {
+        val sourceContainsDefinition =
+          occ.role.isDefinition || snapshot.symbols.exists(
+            _.symbol == occ.symbol
+          )
+
         workspaceReferences(
           source,
           isSymbol,
           isIncludeDeclaration,
           findRealRange,
           includeSynthetics,
+          sourceContainsDefinition,
         )
-      else
+
+      } else
         Seq.empty
 
     workspaceRefs ++ local

tests/unit/src/test/scala/tests/ReferenceLspSuite.scala

@@ -398,6 +398,57 @@ class ReferenceLspSuite extends BaseRangesSuite("reference") {
        |""".stripMargin,
   )
 
+  test("i6101") {
+    for {
+      _ <- initialize(
+        """
+          |/metals.json
+          |{
+          |  "a": { "libraryDependencies": ["com.lihaoyi::sourcecode:0.1.7"] },
+          |  "b": { "libraryDependencies": ["com.lihaoyi::sourcecode:0.1.7"] }
+          |}
+          |/a/src/main/scala/a/A.scala
+          |package a
+          |
+          |import sourcecode.Line
+          |
+          |object A {
+          |  def line = Line
+          |}
+          |/b/src/main/scala/b/B.scala
+          |package b
+          |
+          |//additional line for unambiguous sorting
+          |import sourcecode.Line
+          |
+          |object B {
+          |  def line = Line
+          |}
+          |""".stripMargin
+      )
+      _ <- server.didOpen("a/src/main/scala/a/A.scala")
+      _ <- server.didOpen("b/src/main/scala/b/B.scala")
+      _ = assertNoDiagnostics()
+      references <- server.references("a/src/main/scala/a/A.scala", "Line")
+      _ = assertNoDiff(
+        references,
+        """|a/src/main/scala/a/A.scala:3:19: info: reference
+           |import sourcecode.Line
+           |                  ^^^^
+           |b/src/main/scala/b/B.scala:4:19: info: reference
+           |import sourcecode.Line
+           |                  ^^^^
+           |a/src/main/scala/a/A.scala:6:14: info: reference
+           |  def line = Line
+           |             ^^^^
+           |b/src/main/scala/b/B.scala:7:14: info: reference
+           |  def line = Line
+           |             ^^^^
+           |""".stripMargin,
+      )
+    } yield ()
+  }
+
   override def assertCheck(
       filename: String,
       edit: String,