Fix 'Find symbol' and 'Go to declaration' when jar is used in many libs

SourceNavigationHelper takes into account all the module infos
    that given virtual file is a part of

 #KT-15093 Fixed

compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/KtLightElementBase.kt

@@ -32,7 +32,7 @@ abstract class KtLightElementBase(private val parent: PsiElement): LightElement(
     override fun getText() = kotlinOrigin?.text ?: ""
     override fun getTextRange() = kotlinOrigin?.textRange ?: TextRange.EMPTY_RANGE
     override fun isWritable() = kotlinOrigin?.isWritable ?: false
-    override fun getNavigationElement() = kotlinOrigin ?: this
+    override fun getNavigationElement() = kotlinOrigin?.navigationElement ?: this
     override fun getUseScope() = kotlinOrigin?.useScope ?: super.getUseScope()
     override fun getContainingFile() = parent.containingFile
     override fun getPresentation() = (kotlinOrigin ?: this).let { ItemPresentationProviders.getItemPresentation(it) }

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/getModuleInfo.kt

@@ -17,10 +17,7 @@
 package org.jetbrains.kotlin.idea.caches.resolve
 
 import com.intellij.openapi.project.Project
-import com.intellij.openapi.roots.JdkOrderEntry
-import com.intellij.openapi.roots.LibraryOrderEntry
-import com.intellij.openapi.roots.ModuleRootManager
-import com.intellij.openapi.roots.ProjectFileIndex
+import com.intellij.openapi.roots.*
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.asJava.classes.FakeLightClassForFileOfPackage
@@ -108,27 +105,10 @@ private fun getModuleInfoByVirtualFile(project: Project, virtualFile: VirtualFil
         }
     }
 
-    val orderEntries = projectFileIndex.getOrderEntriesForFile(virtualFile)
-
-    entries@ for (orderEntry in orderEntries) {
-        if (!orderEntry.isValid) continue@entries
-
-        when (orderEntry) {
-            is LibraryOrderEntry -> {
-                val library = orderEntry.library ?: continue@entries
-                if (ProjectRootsUtil.isLibraryClassFile(project, virtualFile) && !treatAsLibrarySource) {
-                    return LibraryInfo(project, library)
-                }
-                else if (ProjectRootsUtil.isLibraryFile(project, virtualFile) || treatAsLibrarySource) {
-                    return LibrarySourceInfo(project, library)
-                }
+    projectFileIndex.getOrderEntriesForFile(virtualFile)
+            .process(project, virtualFile, treatAsLibrarySource) { correspondingModuleInfo ->
+                return correspondingModuleInfo
             }
-            is JdkOrderEntry -> {
-                val sdk = orderEntry.jdk ?: continue@entries
-                return SdkInfo(project, sdk)
-            }
-        }
-    }
 
     val scriptDefinition = getScriptDefinition(virtualFile, project)
     if (scriptDefinition != null) {
@@ -168,3 +148,47 @@ private fun KtLightElement<*, *>.getModuleInfoForLightElement(onFailure: (String
     }
     return element.getModuleInfo()
 }
+
+fun getBinaryLibrariesModuleInfos(project: Project, virtualFile: VirtualFile) = collectModuleInfosByType<BinaryModuleInfo>(project, virtualFile)
+fun getLibrarySourcesModuleInfos(project: Project, virtualFile: VirtualFile) = collectModuleInfosByType<LibrarySourceInfo>(project, virtualFile)
+
+private inline fun <reified T : IdeaModuleInfo> collectModuleInfosByType(project: Project, virtualFile: VirtualFile): Collection<T> {
+    val orderEntries = ProjectFileIndex.SERVICE.getInstance(project).getOrderEntriesForFile(virtualFile)
+
+    val result = linkedSetOf<T?>()
+    orderEntries.process(project, virtualFile, treatAsLibrarySource = false) {
+        result.add(it as? T)
+    }
+    // NOTE: non idea model infos can be obtained this way, like script related infos
+    // only one though, luckily it covers existing cases
+    result.add(getModuleInfoByVirtualFile(project, virtualFile) as? T)
+
+    return result.filterNotNull()
+}
+
+private inline fun List<OrderEntry>.process(
+        project: Project,
+        virtualFile: VirtualFile,
+        treatAsLibrarySource: Boolean = false,
+        body: (IdeaModuleInfo) -> Unit
+) {
+    entries@ for (orderEntry in this) {
+        if (!orderEntry.isValid) continue
+
+        when (orderEntry) {
+            is LibraryOrderEntry -> {
+                val library = orderEntry.library ?: continue@entries
+                if (ProjectRootsUtil.isLibraryClassFile(project, virtualFile) && !treatAsLibrarySource) {
+                    body(LibraryInfo(project, library))
+                }
+                else if (ProjectRootsUtil.isLibraryFile(project, virtualFile) || treatAsLibrarySource) {
+                    body(LibrarySourceInfo(project, library))
+                }
+            }
+            is JdkOrderEntry -> {
+                val sdk = orderEntry.jdk ?: continue@entries
+                body(SdkInfo(project, sdk))
+            }
+        }
+    }
+}

idea/idea-analysis/src/org/jetbrains/kotlin/idea/decompiler/navigation/SourceNavigationHelper.kt

@@ -31,10 +31,7 @@ import gnu.trove.THashSet
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.asJava.toLightClass
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
-import org.jetbrains.kotlin.idea.caches.resolve.BinaryModuleInfo
-import org.jetbrains.kotlin.idea.caches.resolve.SourceForBinaryModuleInfo
-import org.jetbrains.kotlin.idea.caches.resolve.getModuleInfoByVirtualFile
-import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor
+import org.jetbrains.kotlin.idea.caches.resolve.*
 import org.jetbrains.kotlin.idea.decompiler.navigation.MemberMatching.*
 import org.jetbrains.kotlin.idea.stubindex.KotlinFullClassNameIndex
 import org.jetbrains.kotlin.idea.stubindex.KotlinTopLevelFunctionFqnNameIndex
@@ -65,14 +62,18 @@ object SourceNavigationHelper {
     private fun targetScope(declaration: KtNamedDeclaration, navigationKind: NavigationKind): GlobalSearchScope? {
         val containingFile = declaration.containingKtFile
         val vFile = containingFile.virtualFile ?: return null
-        val fromModuleInfo = getModuleInfoByVirtualFile(declaration.project, vFile)
 
         return when (navigationKind) {
-            NavigationKind.CLASS_FILES_TO_SOURCES -> (fromModuleInfo as? BinaryModuleInfo)?.sourcesModuleInfo?.sourceScope()
-            NavigationKind.SOURCES_TO_CLASS_FILES -> (fromModuleInfo as? SourceForBinaryModuleInfo)?.binariesModuleInfo?.binariesScope()
+            NavigationKind.CLASS_FILES_TO_SOURCES -> getBinaryLibrariesModuleInfos(declaration.project, vFile)
+                    .mapNotNull { it.sourcesModuleInfo?.sourceScope() }.union()
+
+            NavigationKind.SOURCES_TO_CLASS_FILES -> getLibrarySourcesModuleInfos(declaration.project, vFile)
+                    .mapNotNull { it.binariesModuleInfo.binariesScope() }.union()
         }
     }
 
+    private fun Collection<GlobalSearchScope>.union() = GlobalSearchScope.union(this.toTypedArray())
+
     private fun haveRenamesInImports(files: Collection<KtFile>) = files.any { it.importDirectives.any { it.aliasName != null } }
 
     private fun findSpecialProperty(memberName: Name, containingClass: KtClass): KtNamedDeclaration? {