FIR JS: Convert KLib symbol provider for handling multiple libs

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

@@ -65,18 +65,14 @@ object FirJsSessionFactory : FirAbstractSessionFactory() {
         registerExtraComponents = { it.registerJsSpecificResolveComponents() },
         createKotlinScopeProvider = { FirKotlinScopeProvider { _, declaredMemberScope, _, _ -> declaredMemberScope } },
         createProviders = { session, builtinsModuleData, kotlinScopeProvider ->
-            val klibProviders = resolvedLibraries.map {
+            listOf(
-                KlibBasedSymbolProvider(session, moduleDataProvider, kotlinScopeProvider, it)
+                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
-            klibProviders +
+                // TODO: consider using "poisoning" provider for builtins to ensure that proper ones are taken from dependencies
-                listOf(
+                // NOTE: it requires precise filtering for true nuiltins, like Function*
-                    FirCloneableSymbolProvider(session, builtinsModuleData, kotlinScopeProvider),
+                FirBuiltinSymbolProvider(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
-                    // NOTE: it requires precise filtering for true nuiltins, like Function*
-                    FirBuiltinSymbolProvider(session, builtinsModuleData, kotlinScopeProvider),
-                )
         }
     )
 

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,31 +69,31 @@ 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 libraryPath = Paths.get(resolvedLibrary.library.libraryFile.path)
-            val moduleData = moduleDataProvider.getModuleData(libraryPath) ?: return@mapNotNull null
+                val moduleData = moduleDataProvider.getModuleData(libraryPath) ?: return@mapNotNull null
-            val packageProto = fragment.`package`
+                val packageProto = fragment.`package`
 
-            val nameResolver = NameResolverImpl(
+                val nameResolver = NameResolverImpl(
-                fragment.strings,
+                    fragment.strings,
-                fragment.qualifiedNames,
+                    fragment.qualifiedNames,
-            )
+                )
 
-            PackagePartsCacheData(
+                PackagePartsCacheData(
-                packageProto,
+                    packageProto,
-                FirDeserializationContext.createForPackage(
+                    FirDeserializationContext.createForPackage(
-                    packageFqName, packageProto, nameResolver, moduleData,
+                        packageFqName, packageProto, nameResolver, moduleData,
-                    annotationDeserializer,
+                        annotationDeserializer,
-                    constDeserializer,
+                        constDeserializer,
-                    null,
+                        null,
-                ),
+                    ),
-            )
+                )
+            }
         }
     }
 
@@ -89,51 +101,52 @@ class KlibBasedSymbolProvider(
     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) {
-            val libraryPath = Paths.get(resolvedLibrary.library.libraryFile.path)
+            for (packageMetadataPart in resolvedLibrary.library.packageMetadataParts(packageStringName)) {
-            val fragment = getPackageFragment(packageStringName, it)
+                val libraryPath = Paths.get(resolvedLibrary.library.libraryFile.path)
+                val fragment = getPackageFragment(resolvedLibrary, packageStringName, packageMetadataPart)
 
-            val nameResolver = NameResolverImpl(
+                val nameResolver = NameResolverImpl(
-                fragment.strings,
+                    fragment.strings,
-                fragment.qualifiedNames,
+                    fragment.qualifiedNames,
-            )
-
-            val finder = KlibMetadataClassDataFinder(fragment, nameResolver)
-            val classProto = finder.findClassData(classId)?.classProto ?: return@forEach
-
-            val moduleData = moduleDataProvider.getModuleData(libraryPath) ?: return null
-
-            return ClassMetadataFindResult.NoMetadata { symbol ->
-                val source = object : DeserializedContainerSource {
-                    override val incompatibility: IncompatibleVersionErrorData<*>? = null
-                    override val isPreReleaseInvisible =
-                        deserializationConfiguration.reportErrorsOnPreReleaseDependencies && (moduleHeader.flags and 1) != 0
-                    override val abiStability = DeserializedContainerAbiStability.STABLE
-                    override val presentableString = "Package '${classId.packageFqName}'"
-
-                    override fun getContainingFile() = SourceFile.NO_SOURCE_FILE
-                }
-
-                deserializeClassToSymbol(
-                    classId,
-                    classProto,
-                    symbol,
-                    nameResolver,
-                    session,
-                    moduleData,
-                    annotationDeserializer,
-                    kotlinScopeProvider,
-                    KlibMetadataSerializerProtocol,
-                    parentContext,
-                    source,
-                    origin = defaultDeserializationOrigin,
-                    deserializeNestedClass = this::getClass,
                 )
-                symbol.fir.isNewPlaceForBodyGeneration = isNewPlaceForBodyGeneration(classProto)
+
+                val finder = KlibMetadataClassDataFinder(fragment, nameResolver)
+                val classProto = finder.findClassData(classId)?.classProto ?: continue
+
+                val moduleData = moduleDataProvider.getModuleData(libraryPath) ?: return null
+
+                return ClassMetadataFindResult.NoMetadata { symbol ->
+                    val source = object : DeserializedContainerSource {
+                        override val incompatibility: IncompatibleVersionErrorData<*>? = null
+                        override val isPreReleaseInvisible =
+                            deserializationConfiguration.reportErrorsOnPreReleaseDependencies &&
+                                    (moduleHeaders[resolvedLibrary]!!.flags and 1) != 0
+                        override val abiStability = DeserializedContainerAbiStability.STABLE
+                        override val presentableString = "Package '${classId.packageFqName}'"
+
+                        override fun getContainingFile() = SourceFile.NO_SOURCE_FILE
+                    }
+
+                    deserializeClassToSymbol(
+                        classId,
+                        classProto,
+                        symbol,
+                        nameResolver,
+                        session,
+                        moduleData,
+                        annotationDeserializer,
+                        kotlinScopeProvider,
+                        KlibMetadataSerializerProtocol,
+                        parentContext,
+                        source,
+                        origin = defaultDeserializationOrigin,
+                        deserializeNestedClass = this::getClass,
+                    )
+                    symbol.fir.isNewPlaceForBodyGeneration = isNewPlaceForBodyGeneration(classProto)
+                }
             }
         }
 
@@ -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