FIR JS: Convert KLib symbol provider for handling multiple libs

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/FirJsSessionFactory.kt

@@ -65,12 +65,8 @@ object FirJsSessionFactory : FirAbstractSessionFactory() {
         registerExtraComponents = { it.registerJsSpecificResolveComponents() },
         createKotlinScopeProvider = { FirKotlinScopeProvider { _, declaredMemberScope, _, _ -> declaredMemberScope } },
         createProviders = { session, builtinsModuleData, kotlinScopeProvider ->
-            val klibProviders = resolvedLibraries.map {
-                KlibBasedSymbolProvider(session, moduleDataProvider, kotlinScopeProvider, it)
-            }
-
-            klibProviders +
             listOf(
+                KlibBasedSymbolProvider(session, moduleDataProvider, kotlinScopeProvider, resolvedLibraries),
                 FirCloneableSymbolProvider(session, builtinsModuleData, kotlinScopeProvider),
                 // (Most) builtins should be taken from the dependencies in JS compilation, therefore builtins provider is the last one
                 // TODO: consider using "poisoning" provider for builtins to ensure that proper ones are taken from dependencies

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/KlibBasedSymbolProvider.kt

@@ -24,32 +24,44 @@ import org.jetbrains.kotlin.resolve.CompilerDeserializationConfiguration
 import org.jetbrains.kotlin.serialization.deserialization.IncompatibleVersionErrorData
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedContainerAbiStability
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedContainerSource
+import org.jetbrains.kotlin.utils.SmartList
 import java.nio.file.Paths
 
 class KlibBasedSymbolProvider(
     session: FirSession,
     moduleDataProvider: ModuleDataProvider,
     kotlinScopeProvider: FirKotlinScopeProvider,
-    private val resolvedLibrary: KotlinResolvedLibrary,
+    private val resolvedLibraries: Collection<KotlinResolvedLibrary>,
     defaultDeserializationOrigin: FirDeclarationOrigin = FirDeclarationOrigin.Library
 ) : AbstractFirDeserializedSymbolProvider(
     session, moduleDataProvider, kotlinScopeProvider, defaultDeserializationOrigin, KlibMetadataSerializerProtocol
 ) {
-    private val moduleHeader by lazy {
+    private val moduleHeaders by lazy {
-        resolvedLibrary.loadModuleHeader(resolvedLibrary.library)
+        resolvedLibraries.associate { it to it.loadModuleHeader(it.library) }
     }
 
-    private val fragmentNameList by lazy {
+    private val fragmentNamesInLibraries: Map<String, List<KotlinResolvedLibrary>> by lazy {
-        moduleHeader.packageFragmentNameList.toSet()
+        buildMap<String, SmartList<KotlinResolvedLibrary>> {
+            for ((library, header) in moduleHeaders) {
+                for (fragmentName in header.packageFragmentNameList) {
+                    getOrPut(fragmentName) { SmartList() }
+                        .add(library)
+                }
+            }
+        }
     }
 
     private val annotationDeserializer = KlibBasedAnnotationDeserializer(session)
     private val constDeserializer = FirConstDeserializer(session, KlibMetadataSerializerProtocol)
     private val deserializationConfiguration = CompilerDeserializationConfiguration(session.languageVersionSettings)
-    private val cachedFragments = mutableMapOf<Pair<String, String>, ProtoBuf.PackageFragment>()
+    private val cachedFragments = mutableMapOf<KotlinResolvedLibrary, MutableMap<Pair<String, String>, ProtoBuf.PackageFragment>>()
 
-    private fun getPackageFragment(packageStringName: String, packageMetadataPart: String): ProtoBuf.PackageFragment {
+    private fun getPackageFragment(
-        return cachedFragments.getOrPut(packageStringName to packageMetadataPart) {
+        resolvedLibrary: KotlinResolvedLibrary, packageStringName: String, packageMetadataPart: String
+    ): ProtoBuf.PackageFragment {
+        return cachedFragments.getOrPut(resolvedLibrary) {
+            mutableMapOf()
+        }.getOrPut(packageStringName to packageMetadataPart) {
             resolvedLibrary.loadPackageFragment(resolvedLibrary.library, packageStringName, packageMetadataPart)
         }
     }
@@ -57,12 +69,11 @@ class KlibBasedSymbolProvider(
     override fun computePackagePartsInfos(packageFqName: FqName): List<PackagePartsCacheData> {
         val packageStringName = if (packageFqName.isRoot) "" else packageFqName.asString()
 
-        if (packageStringName !in fragmentNameList) {
+        val librariesWithFragment = fragmentNamesInLibraries[packageStringName] ?: return emptyList()
-            return emptyList()
-        }
 
-        return resolvedLibrary.library.packageMetadataParts(packageStringName).mapNotNull {
+        return librariesWithFragment.flatMap { resolvedLibrary ->
-            val fragment = getPackageFragment(packageStringName, it)
+            resolvedLibrary.library.packageMetadataParts(packageStringName).mapNotNull {
+                val fragment = getPackageFragment(resolvedLibrary, packageStringName, it)
 
                 val libraryPath = Paths.get(resolvedLibrary.library.libraryFile.path)
                 val moduleData = moduleDataProvider.getModuleData(libraryPath) ?: return@mapNotNull null
@@ -84,18 +95,18 @@ class KlibBasedSymbolProvider(
                 )
             }
         }
+    }
 
     @OptIn(SymbolInternals::class)
     override fun extractClassMetadata(classId: ClassId, parentContext: FirDeserializationContext?): ClassMetadataFindResult? {
         val packageStringName = classId.packageFqName.asString()
 
-        if (packageStringName !in fragmentNameList) {
+        val librariesWithFragment = fragmentNamesInLibraries[packageStringName] ?: return null
-            return null
-        }
 
-        resolvedLibrary.library.packageMetadataParts(packageStringName).forEach {
+        for (resolvedLibrary in librariesWithFragment) {
+            for (packageMetadataPart in resolvedLibrary.library.packageMetadataParts(packageStringName)) {
                 val libraryPath = Paths.get(resolvedLibrary.library.libraryFile.path)
-            val fragment = getPackageFragment(packageStringName, it)
+                val fragment = getPackageFragment(resolvedLibrary, packageStringName, packageMetadataPart)
 
                 val nameResolver = NameResolverImpl(
                     fragment.strings,
@@ -103,7 +114,7 @@ class KlibBasedSymbolProvider(
                 )
 
                 val finder = KlibMetadataClassDataFinder(fragment, nameResolver)
-            val classProto = finder.findClassData(classId)?.classProto ?: return@forEach
+                val classProto = finder.findClassData(classId)?.classProto ?: continue
 
                 val moduleData = moduleDataProvider.getModuleData(libraryPath) ?: return null
 
@@ -111,7 +122,8 @@ class KlibBasedSymbolProvider(
                     val source = object : DeserializedContainerSource {
                         override val incompatibility: IncompatibleVersionErrorData<*>? = null
                         override val isPreReleaseInvisible =
-                        deserializationConfiguration.reportErrorsOnPreReleaseDependencies && (moduleHeader.flags and 1) != 0
+                            deserializationConfiguration.reportErrorsOnPreReleaseDependencies &&
+                                    (moduleHeaders[resolvedLibrary]!!.flags and 1) != 0
                         override val abiStability = DeserializedContainerAbiStability.STABLE
                         override val presentableString = "Package '${classId.packageFqName}'"
 
@@ -136,6 +148,7 @@ class KlibBasedSymbolProvider(
                     symbol.fir.isNewPlaceForBodyGeneration = isNewPlaceForBodyGeneration(classProto)
                 }
             }
+        }
 
         return null
     }
@@ -143,7 +156,7 @@ class KlibBasedSymbolProvider(
     override fun isNewPlaceForBodyGeneration(classProto: ProtoBuf.Class) = false
 
     override fun getPackage(fqName: FqName): FqName? {
-        return if (fqName.toString() in fragmentNameList) {
+        return if (fqName.toString() in fragmentNamesInLibraries) {
             fqName
         } else {
             null