[FIR] Deserialize Klib KlibMetadataProtoBuf. classFile, functionFile & propertyFile

klibs provide information about the container source file
in the KlibMetadataProtoBuf extensions for functions,
properties and classes.

This information is deserialized and attached to the
`klibSourceFile` extension (stored in FirDeclarationDataRegistry)

^KT-66271 Fixed

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

@@ -7,10 +7,10 @@ package org.jetbrains.kotlin.fir.session
 
 import org.jetbrains.kotlin.fir.FirModuleData
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.caches.FirCache
-import org.jetbrains.kotlin.fir.caches.firCachesFactory
-import org.jetbrains.kotlin.fir.caches.getValue
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
+import org.jetbrains.kotlin.fir.declarations.FirFunction
+import org.jetbrains.kotlin.fir.declarations.FirProperty
+import org.jetbrains.kotlin.fir.declarations.utils.klibSourceFile
 import org.jetbrains.kotlin.fir.deserialization.*
 import org.jetbrains.kotlin.fir.isNewPlaceForBodyGeneration
 import org.jetbrains.kotlin.fir.languageVersionSettings
@@ -19,26 +19,28 @@ import org.jetbrains.kotlin.fir.symbols.SymbolInternals
 import org.jetbrains.kotlin.library.KotlinLibrary
 import org.jetbrains.kotlin.library.MetadataLibrary
 import org.jetbrains.kotlin.library.metadata.*
-import org.jetbrains.kotlin.library.metadata.resolver.KotlinResolvedLibrary
 import org.jetbrains.kotlin.metadata.ProtoBuf
 import org.jetbrains.kotlin.metadata.deserialization.NameResolver
 import org.jetbrains.kotlin.metadata.deserialization.NameResolverImpl
+import org.jetbrains.kotlin.metadata.deserialization.getExtensionOrNull
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.protobuf.GeneratedMessageLite
+import org.jetbrains.kotlin.protobuf.GeneratedMessageLite.GeneratedExtension
 import org.jetbrains.kotlin.resolve.CompilerDeserializationConfiguration
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedContainerSource
 import org.jetbrains.kotlin.serialization.deserialization.getClassId
-import org.jetbrains.kotlin.utils.SmartList
-import java.nio.file.Paths
 
 abstract class MetadataLibraryBasedSymbolProvider<L : MetadataLibrary>(
     session: FirSession,
     moduleDataProvider: ModuleDataProvider,
     kotlinScopeProvider: FirKotlinScopeProvider,
-    defaultDeserializationOrigin: FirDeclarationOrigin = FirDeclarationOrigin.Library
+    defaultDeserializationOrigin: FirDeclarationOrigin = FirDeclarationOrigin.Library,
 ) : AbstractFirDeserializedSymbolProvider(
     session, moduleDataProvider, kotlinScopeProvider, defaultDeserializationOrigin, KlibMetadataSerializerProtocol
 ) {
+    private class MetadataLibraryPackagePartCacheDataExtra(val library: MetadataLibrary) : PackagePartsCacheData.Extra
+
     protected abstract fun moduleData(library: L): FirModuleData?
 
     protected abstract val fragmentNamesInLibraries: Map<String, List<L>>
@@ -85,6 +87,7 @@ abstract class MetadataLibraryBasedSymbolProvider<L : MetadataLibrary>(
                         constDeserializer,
                         createDeserializedContainerSource(resolvedLibrary, packageFqName),
                     ),
+                    MetadataLibraryPackagePartCacheDataExtra(resolvedLibrary)
                 )
             }
         }
@@ -130,6 +133,13 @@ abstract class MetadataLibraryBasedSymbolProvider<L : MetadataLibrary>(
                     origin = defaultDeserializationOrigin,
                     deserializeNestedClass = this::getClass,
                 )
+
+                if (resolvedLibrary is KotlinLibrary) {
+                    symbol.fir.klibSourceFile = loadKlibSourceFileExtensionOrNull(
+                        resolvedLibrary, nameResolver, classProto, KlibMetadataProtoBuf.classFile
+                    )
+                }
+
                 symbol.fir.isNewPlaceForBodyGeneration = isNewPlaceForBodyGeneration(classProto)
             }
         }
@@ -160,6 +170,30 @@ abstract class MetadataLibraryBasedSymbolProvider<L : MetadataLibrary>(
         }
     }
 
+    override fun loadFunctionExtensions(packagePart: PackagePartsCacheData, proto: ProtoBuf.Function, fir: FirFunction) {
+        fir.klibSourceFile = loadKlibSourceFileExtensionOrNull(packagePart, proto, KlibMetadataProtoBuf.functionFile) ?: return
+    }
+
+    override fun loadPropertyExtensions(packagePart: PackagePartsCacheData, proto: ProtoBuf.Property, fir: FirProperty) {
+        fir.klibSourceFile = loadKlibSourceFileExtensionOrNull(packagePart, proto, KlibMetadataProtoBuf.propertyFile) ?: return
+    }
+
+    private fun <T : GeneratedMessageLite.ExtendableMessage<T>> loadKlibSourceFileExtensionOrNull(
+        packagePart: PackagePartsCacheData, proto: T, sourceFileExtension: GeneratedExtension<T, Int>,
+    ): DeserializedSourceFile? {
+        val library = (packagePart.extra as? MetadataLibraryPackagePartCacheDataExtra)?.library as? KotlinLibrary ?: return null
+        return loadKlibSourceFileExtensionOrNull(library, packagePart.context.nameResolver, proto, sourceFileExtension)
+    }
+
+    private fun <T : GeneratedMessageLite.ExtendableMessage<T>> loadKlibSourceFileExtensionOrNull(
+        library: KotlinLibrary, nameResolver: NameResolver, proto: T, sourceFileExtension: GeneratedExtension<T, Int>,
+    ): DeserializedSourceFile? {
+        return proto.getExtensionOrNull(sourceFileExtension)
+            ?.let { fileId -> nameResolver.getString(fileId) }
+            ?.let { fileName -> DeserializedSourceFile(fileName, library) }
+    }
+
+
     protected abstract fun createDeserializedContainerSource(
         resolvedLibrary: L,
         packageFqName: FqName

compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/deserialization/AbstractFirDeserializedSymbolProvider.kt

@@ -12,12 +12,14 @@ import org.jetbrains.kotlin.fir.caches.createCache
 import org.jetbrains.kotlin.fir.caches.firCachesFactory
 import org.jetbrains.kotlin.fir.caches.getValue
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
+import org.jetbrains.kotlin.fir.declarations.FirFunction
+import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.declarations.utils.isExpect
 import org.jetbrains.kotlin.fir.isNewPlaceForBodyGeneration
+import org.jetbrains.kotlin.fir.resolve.providers.FirCachedSymbolNamesProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolNamesProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProviderInternals
-import org.jetbrains.kotlin.fir.resolve.providers.FirCachedSymbolNamesProvider
 import org.jetbrains.kotlin.fir.scopes.FirKotlinScopeProvider
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.metadata.ProtoBuf
@@ -32,7 +34,17 @@ import java.nio.file.Path
 class PackagePartsCacheData(
     val proto: ProtoBuf.Package,
     val context: FirDeserializationContext,
+    val extra: Extra? = null,
 ) {
+    /**
+     * Marker interface for 'extra' data that can be attached to a given [PackagePartsCacheData].
+     * Example: This can be used by a [FirSymbolProvider] to attach data (like the library that this package part came from) to
+     * this particular package
+     *
+     * @see PackagePartsCacheData.extra
+     */
+    interface Extra
+
     val topLevelFunctionNameIndex by lazy {
         proto.functionList.withIndex()
             .groupBy({ context.nameResolver.getName(it.value.name) }) { (index) -> index }
@@ -244,10 +256,13 @@ abstract class AbstractFirDeserializedSymbolProvider(
         return getPackageParts(callableId.packageName).flatMap { part ->
             val functionIds = part.topLevelFunctionNameIndex[callableId.callableName] ?: return@flatMap emptyList()
             functionIds.map {
-                part.context.memberDeserializer.loadFunction(
+                val proto = part.proto.getFunction(it)
+                val fir = part.context.memberDeserializer.loadFunction(
                     part.proto.getFunction(it),
                     deserializationOrigin = defaultDeserializationOrigin
-                ).symbol
+                )
+                loadFunctionExtensions(part, proto, fir)
+                fir.symbol
             }
         }
     }
@@ -256,11 +271,24 @@ abstract class AbstractFirDeserializedSymbolProvider(
         return getPackageParts(callableId.packageName).flatMap { part ->
             val propertyIds = part.topLevelPropertyNameIndex[callableId.callableName] ?: return@flatMap emptyList()
             propertyIds.map {
-                part.context.memberDeserializer.loadProperty(part.proto.getProperty(it)).symbol
+                val proto = part.proto.getProperty(it)
+                val fir = part.context.memberDeserializer.loadProperty(proto)
+                loadPropertyExtensions(part, proto, fir)
+                fir.symbol
             }
         }
     }
 
+    open fun loadFunctionExtensions(
+        packagePart: PackagePartsCacheData, proto: ProtoBuf.Function, fir: FirFunction,
+    ) {
+    }
+
+    open fun loadPropertyExtensions(
+        packagePart: PackagePartsCacheData, proto: ProtoBuf.Property, fir: FirProperty,
+    ) {
+    }
+
     private fun getPackageParts(packageFqName: FqName): Collection<PackagePartsCacheData> =
         packagePartsCache.getValue(packageFqName)
 

compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/utils/declarationAttributes.kt

@@ -7,11 +7,13 @@ package org.jetbrains.kotlin.fir.declarations.utils
 
 import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.descriptors.SourceElement
+import org.jetbrains.kotlin.descriptors.SourceFile
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertyBackingField
 import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertyGetter
 import org.jetbrains.kotlin.fir.expressions.FirPropertyAccessExpression
 import org.jetbrains.kotlin.fir.references.impl.FirPropertyFromParameterResolvedNamedReference
+import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
 import org.jetbrains.kotlin.fir.types.coneType
@@ -24,6 +26,7 @@ private object ComponentFunctionSymbolKey : FirDeclarationDataKey()
 private object SourceElementKey : FirDeclarationDataKey()
 private object ModuleNameKey : FirDeclarationDataKey()
 private object DanglingTypeConstraintsKey : FirDeclarationDataKey()
+private object KlibSourceFile : FirDeclarationDataKey()
 
 var FirProperty.isFromVararg: Boolean? by FirDeclarationDataRegistry.data(IsFromVarargKey)
 var FirProperty.isReferredViaField: Boolean? by FirDeclarationDataRegistry.data(IsReferredViaField)
@@ -32,12 +35,27 @@ var FirProperty.componentFunctionSymbol: FirNamedFunctionSymbol? by FirDeclarati
 var FirClassLikeDeclaration.sourceElement: SourceElement? by FirDeclarationDataRegistry.data(SourceElementKey)
 var FirRegularClass.moduleName: String? by FirDeclarationDataRegistry.data(ModuleNameKey)
 
+/**
+ * @see [FirBasedSymbol.klibSourceFile]
+ */
+var FirDeclaration.klibSourceFile: SourceFile? by FirDeclarationDataRegistry.data(KlibSourceFile)
+
 val FirClassLikeSymbol<*>.sourceElement: SourceElement?
     get() = fir.sourceElement
 
 val FirPropertySymbol.fromPrimaryConstructor: Boolean
     get() = fir.fromPrimaryConstructor ?: false
 
+/**
+ * Declarations like classes, functions, and properties can encode their containing Kotlin source file into .klibs using
+ * klib specific metadata extensions.
+ * If present in the klib and deserialized by the corresponding deserializer/symbol provider,
+ * then this source file is available here
+ * @see FirDeclaration.klibSourceFile
+ */
+val FirBasedSymbol<out FirDeclaration>.klibSourceFile: SourceFile?
+    get() = fir.klibSourceFile
+
 /**
  * Constraint without corresponding type argument
  */

compiler/fir/tree/src/org/jetbrains/kotlin/fir/renderer/FirDeclarationRendererWithFilteredAttributes.kt

@@ -11,6 +11,6 @@ class FirDeclarationRendererWithFilteredAttributes : FirDeclarationRendererWithA
     }
 
     private companion object {
-        private val IGNORED_ATTRIBUTES = setOf("FirVersionRequirementsTableKey", "SourceElementKey")
+        private val IGNORED_ATTRIBUTES = setOf("FirVersionRequirementsTableKey", "SourceElementKey", "KlibSourceFile")
     }
 }