From 7e1a9e2b9a6157b9d4ecb55ef487419fa77467f5 Mon Sep 17 00:00:00 2001
From: Roman Golyshev <roman.golyshev@jetbrains.com>
Date: Wed, 2 Feb 2022 15:34:01 +0300
Subject: [PATCH] KTIJ-20637 Group files by module in
 `getFacadeClassesInPackage` also

Fixup for 35e2561dab11d2d646f3ba3529121c4a8c91bae4, since these changes
were overlooked in the initial commit.

^KTIJ-20637 Fixed
---
 .../symbol/IDEKotlinAsJavaFirSupport.kt       | 23 +++++++++++--------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/IDEKotlinAsJavaFirSupport.kt b/analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/IDEKotlinAsJavaFirSupport.kt
index ae843a2a955..7d4c5c6dbcf 100644
--- a/analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/IDEKotlinAsJavaFirSupport.kt
+++ b/analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/IDEKotlinAsJavaFirSupport.kt
@@ -91,13 +91,9 @@ class IDEKotlinAsJavaFirSupport(private val project: Project) : KotlinAsJavaSupp
         error("Should not be called")
 
     override fun getFacadeClasses(facadeFqName: FqName, scope: GlobalSearchScope): Collection<PsiClass> {
-        val filesByModule = findFilesForFacade(facadeFqName, scope).groupBy { it.getKtModule(project) }
+        val filesForFacade = findFilesForFacade(facadeFqName, scope)
 
-        val lightClassFacadeCache = project.getService(SymbolLightClassFacadeCache::class.java)
-
-        return filesByModule.values.mapNotNull { files ->
-            lightClassFacadeCache.getOrCreateSymbolLightFacade(files, facadeFqName)
-        }
+        return getFacadeClassesForFiles(facadeFqName, filesForFacade)
     }
 
     override fun getScriptClasses(scriptFqName: FqName, scope: GlobalSearchScope): Collection<PsiClass> =
@@ -112,10 +108,7 @@ class IDEKotlinAsJavaFirSupport(private val project: Project) : KotlinAsJavaSupp
             .asSequence()
             .filter { it.isFromSourceOrLibraryBinary(project) }
             .groupBy { it.javaFileFacadeFqName }
-            .mapNotNull {
-                project.getService(SymbolLightClassFacadeCache::class.java)
-                    .getOrCreateSymbolLightFacade(it.value, it.key)
-            }
+            .flatMap { (fqName, files) -> getFacadeClassesForFiles(fqName, files) }
 
     override fun getFacadeNames(packageFqName: FqName, scope: GlobalSearchScope): Collection<String> =
         project.createDeclarationProvider(scope)
@@ -134,6 +127,16 @@ class IDEKotlinAsJavaFirSupport(private val project: Project) : KotlinAsJavaSupp
 
     override fun createFacadeForSyntheticFile(facadeClassFqName: FqName, file: KtFile): PsiClass =
         TODO("Not implemented")
+
+    private fun getFacadeClassesForFiles(facadeFqName: FqName, allFiles: Collection<KtFile>): Collection<PsiClass> {
+        val filesByModule = allFiles.groupBy { it.getKtModule(project) }
+
+        val lightClassFacadeCache = project.getService(SymbolLightClassFacadeCache::class.java)
+
+        return filesByModule.values.mapNotNull { files ->
+            lightClassFacadeCache.getOrCreateSymbolLightFacade(files, facadeFqName)
+        }
+    }
 }