[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
2024-03-05 10:58:24 +01:00
parent 85a1d67d19
commit 6d8a4a28af
4 changed files with 92 additions and 12 deletions
@@ -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
@@ -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)

@@ -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
 */
@@ -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")
    }
 }