[FIR] Move out all JVM specific deserialization logic from AbstractFirDeserializedSymbolsProvider

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

@@ -171,7 +171,7 @@ object FirSessionFactory {
         }
 
         return KotlinDeserializedJvmSymbolsProvider(
-            this, packagePartProvider, kotlinClassFinder, kotlinScopeProvider, javaSymbolProvider, javaClassFinder
+            this, kotlinScopeProvider, packagePartProvider, kotlinClassFinder, javaSymbolProvider, javaClassFinder
         )
     }
 

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

@@ -7,18 +7,12 @@ package org.jetbrains.kotlin.fir.deserialization
 
 import com.intellij.openapi.progress.ProcessCanceledException
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.caches.createCache
+import org.jetbrains.kotlin.fir.caches.*
-import org.jetbrains.kotlin.fir.caches.firCachesFactory
-import org.jetbrains.kotlin.fir.caches.getValue
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProviderInternals
 import org.jetbrains.kotlin.fir.scopes.FirKotlinScopeProvider
 import org.jetbrains.kotlin.fir.symbols.impl.*
-import org.jetbrains.kotlin.load.kotlin.KotlinClassFinder
-import org.jetbrains.kotlin.load.kotlin.PackagePartProvider
-import org.jetbrains.kotlin.load.kotlin.header.KotlinClassHeader
 import org.jetbrains.kotlin.metadata.ProtoBuf
-import org.jetbrains.kotlin.metadata.deserialization.Flags
 import org.jetbrains.kotlin.metadata.deserialization.NameResolver
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
@@ -27,44 +21,69 @@ import org.jetbrains.kotlin.name.isOneSegmentFQN
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedContainerSource
 import org.jetbrains.kotlin.serialization.deserialization.getName
 
+class PackagePartsCacheData(
+    val proto: ProtoBuf.Package,
+    val context: FirDeserializationContext,
+) {
+    val topLevelFunctionNameIndex by lazy {
+        proto.functionList.withIndex()
+            .groupBy({ context.nameResolver.getName(it.value.name) }) { (index) -> index }
+    }
+    val topLevelPropertyNameIndex by lazy {
+        proto.propertyList.withIndex()
+            .groupBy({ context.nameResolver.getName(it.value.name) }) { (index) -> index }
+    }
+    val typeAliasNameIndex by lazy {
+        proto.typeAliasList.withIndex()
+            .groupBy({ context.nameResolver.getName(it.value.name) }) { (index) -> index }
+    }
+}
+
+typealias DeserializedClassPostProcessor = (FirRegularClassSymbol) -> Unit
+
 abstract class AbstractFirDeserializedSymbolsProvider(
     session: FirSession,
-    val packagePartProvider: PackagePartProvider,
+    val kotlinScopeProvider: FirKotlinScopeProvider
-    val kotlinClassFinder: KotlinClassFinder,
-    val kotlinScopeProvider: FirKotlinScopeProvider,
 ) : FirSymbolProvider(session) {
 
     // ------------------------ Caches ------------------------
 
     private val packagePartsCache = session.firCachesFactory.createCache(::tryComputePackagePartInfos)
     private val typeAliasCache = session.firCachesFactory.createCache(::findAndDeserializeTypeAlias)
-    private val classCache =
+    private val classCache: FirCache<ClassId, FirRegularClassSymbol?, FirDeserializationContext?> =
-        session.firCachesFactory.createCacheWithPostCompute<ClassId, FirRegularClassSymbol?, FirDeserializationContext?, KotlinClassFinder.Result.KotlinClass?>(
+        session.firCachesFactory.createCacheWithPostCompute(
             createValue = { classId, context -> findAndDeserializeClass(classId, context) },
-            postCompute = { _, symbol, result ->
+            postCompute = { _, symbol, postProcessor ->
-                if (result != null && symbol != null) {
+                if (postProcessor != null && symbol != null) {
-                    postProcessDeserializedClass(result, symbol)
+                    postProcessor.invoke(symbol)
                 }
             }
         )
 
     // ------------------------ Abstract members ------------------------
 
-    protected abstract val knownNameInPackageCache: KnownNameInPackageCache
-
-    protected abstract fun readClassFromClassFile(
-        classId: ClassId,
-        classFile: KotlinClassFinder.Result.ClassFileContent
-    ): FirRegularClassSymbol?
-
-    protected abstract fun KotlinClassFinder.Result.KotlinClass.extractMetadata(): Pair<NameResolver, ProtoBuf.Class>?
     protected abstract fun computePackagePartsInfos(packageFqName: FqName): List<PackagePartsCacheData>
-    protected abstract fun createAnnotationDeserializer(kotlinClass: KotlinClassFinder.Result.KotlinClass): AbstractAnnotationDeserializer
+    protected abstract fun extractClassMetadata(
-    protected abstract fun createSourceElement(kotlinClass: KotlinClassFinder.Result.KotlinClass): DeserializedContainerSource
+        classId: ClassId,
-    protected open fun postProcessDeserializedClass(kotlinClass: KotlinClassFinder.Result.KotlinClass, symbol: FirRegularClassSymbol) {}
+        parentContext: FirDeserializationContext? = null
+    ): ClassMetadataFindResult?
 
     // ------------------------ Deserialization methods ------------------------
 
+    sealed class ClassMetadataFindResult {
+        data class Metadata(
+            val nameResolver: NameResolver,
+            val classProto: ProtoBuf.Class,
+            val annotationDeserializer: AbstractAnnotationDeserializer,
+            val sourceElement: DeserializedContainerSource,
+            val classPostProcessor: DeserializedClassPostProcessor
+        ) : ClassMetadataFindResult()
+
+        class ClassWithoutMetadata(val symbol: FirRegularClassSymbol?) : ClassMetadataFindResult()
+
+        object ShouldDeserializeViaParent : ClassMetadataFindResult()
+    }
+
     private fun tryComputePackagePartInfos(packageFqName: FqName): List<PackagePartsCacheData> {
         return try {
             computePackagePartsInfos(packageFqName)
@@ -85,36 +104,25 @@ abstract class AbstractFirDeserializedSymbolsProvider(
     private fun findAndDeserializeClass(
         classId: ClassId,
         parentContext: FirDeserializationContext? = null
-    ): Pair<FirRegularClassSymbol?, KotlinClassFinder.Result.KotlinClass?> {
+    ): Pair<FirRegularClassSymbol?, DeserializedClassPostProcessor?> {
-        if (knownNameInPackageCache.hasNoTopLevelClassOf(classId)) return null to null
+        return when (val result = extractClassMetadata(classId, parentContext)) {
-        val result = try {
+            is ClassMetadataFindResult.Metadata -> {
-            kotlinClassFinder.findKotlinClassOrContent(classId)
+                val (nameResolver, classProto, annotationDeserializer, sourceElement, postProcessor) = result
-        } catch (e: ProcessCanceledException) {
+                val symbol = FirRegularClassSymbol(classId)
-            return null to null
+                deserializeClassToSymbol(
+                    classId, classProto, symbol, nameResolver, session,
+                    annotationDeserializer,
+                    kotlinScopeProvider,
+                    parentContext,
+                    sourceElement,
+                    deserializeNestedClass = this::getClass,
+                )
+                symbol to postProcessor
+            }
+            is ClassMetadataFindResult.ClassWithoutMetadata -> result.symbol to null
+            ClassMetadataFindResult.ShouldDeserializeViaParent -> findAndDeserializeClassViaParent(classId) to null
+            null -> null to null
         }
-        val kotlinClass = when (result) {
-            is KotlinClassFinder.Result.KotlinClass -> result
-            is KotlinClassFinder.Result.ClassFileContent -> return readClassFromClassFile(classId, result) to null
-            null -> return findAndDeserializeClassViaParent(classId) to null
-        }
-        if (kotlinClass.kotlinJvmBinaryClass.classHeader.kind != KotlinClassHeader.Kind.CLASS) return null to null
-        val (nameResolver, classProto) = kotlinClass.extractMetadata() ?: return null to null
-
-        if (parentContext == null && Flags.CLASS_KIND.get(classProto.flags) == ProtoBuf.Class.Kind.COMPANION_OBJECT) {
-            return findAndDeserializeClassViaParent(classId) to null
-        }
-
-        val symbol = FirRegularClassSymbol(classId)
-        deserializeClassToSymbol(
-            classId, classProto, symbol, nameResolver, session,
-            createAnnotationDeserializer(kotlinClass),
-            kotlinScopeProvider,
-            parentContext,
-            createSourceElement(kotlinClass),
-            deserializeNestedClass = this::getClass,
-        )
-
-        return symbol to kotlinClass
     }
 
     private fun findAndDeserializeClassViaParent(classId: ClassId): FirRegularClassSymbol? {
@@ -183,32 +191,4 @@ abstract class AbstractFirDeserializedSymbolsProvider(
     }
 
     override fun getPackage(fqName: FqName): FqName? = null
-
-    // ------------------------ Additional classes ------------------------
-
-    protected class PackagePartsCacheData(
-        val proto: ProtoBuf.Package,
-        val context: FirDeserializationContext,
-    ) {
-        val topLevelFunctionNameIndex by lazy {
-            proto.functionList.withIndex()
-                .groupBy({ context.nameResolver.getName(it.value.name) }) { (index) -> index }
-        }
-        val topLevelPropertyNameIndex by lazy {
-            proto.propertyList.withIndex()
-                .groupBy({ context.nameResolver.getName(it.value.name) }) { (index) -> index }
-        }
-        val typeAliasNameIndex by lazy {
-            proto.typeAliasList.withIndex()
-                .groupBy({ context.nameResolver.getName(it.value.name) }) { (index) -> index }
-        }
-    }
-
-    protected abstract class KnownNameInPackageCache {
-        /**
-         * This function returns true if we are sure that no top-level class with this id is available
-         * If it returns false, it means we can say nothing about this id
-         */
-        abstract fun hasNoTopLevelClassOf(classId: ClassId): Boolean
-    }
 }

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/deserialization/KotlinDeserializedJvmSymbolsProvider.kt

@@ -5,15 +5,13 @@
 
 package org.jetbrains.kotlin.fir.java.deserialization
 
+import com.intellij.openapi.progress.ProcessCanceledException
 import org.jetbrains.kotlin.descriptors.SourceElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.ThreadSafeMutableState
 import org.jetbrains.kotlin.fir.caches.*
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.deserialization.AbstractAnnotationDeserializer
+import org.jetbrains.kotlin.fir.deserialization.*
-import org.jetbrains.kotlin.fir.deserialization.AbstractFirDeserializedSymbolsProvider
-import org.jetbrains.kotlin.fir.deserialization.FirConstDeserializer
-import org.jetbrains.kotlin.fir.deserialization.FirDeserializationContext
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.*
 import org.jetbrains.kotlin.fir.java.JavaSymbolProvider
@@ -23,7 +21,9 @@ import org.jetbrains.kotlin.fir.scopes.FirKotlinScopeProvider
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.load.java.JavaClassFinder
 import org.jetbrains.kotlin.load.kotlin.*
+import org.jetbrains.kotlin.load.kotlin.header.KotlinClassHeader
 import org.jetbrains.kotlin.metadata.ProtoBuf
+import org.jetbrains.kotlin.metadata.deserialization.Flags
 import org.jetbrains.kotlin.metadata.deserialization.NameResolver
 import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmMetadataVersion
 import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmProtoBufUtil
@@ -31,30 +31,19 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.resolve.jvm.JvmClassName
 import org.jetbrains.kotlin.serialization.deserialization.IncompatibleVersionErrorData
-import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedContainerSource
 
 @ThreadSafeMutableState
 class KotlinDeserializedJvmSymbolsProvider(
     session: FirSession,
-    packagePartProvider: PackagePartProvider,
-    kotlinClassFinder: KotlinClassFinder,
     kotlinScopeProvider: FirKotlinScopeProvider,
+    private val packagePartProvider: PackagePartProvider,
+    private val kotlinClassFinder: KotlinClassFinder,
     private val javaSymbolProvider: JavaSymbolProvider,
     javaClassFinder: JavaClassFinder,
-) : AbstractFirDeserializedSymbolsProvider(session, packagePartProvider, kotlinClassFinder, kotlinScopeProvider) {
+) : AbstractFirDeserializedSymbolsProvider(session, kotlinScopeProvider) {
-    override val knownNameInPackageCache: KnownNameInPackageCache = JvmKnownNameInPackageCache(session, javaClassFinder)
+    private val knownNameInPackageCache = KnownNameInPackageCache(session, javaClassFinder)
     private val annotationsLoader = AnnotationsLoader(session)
 
-    override fun readClassFromClassFile(classId: ClassId, classFile: KotlinClassFinder.Result.ClassFileContent): FirRegularClassSymbol? {
-        return javaSymbolProvider.getFirJavaClass(classId, classFile)
-    }
-
-    override fun KotlinClassFinder.Result.KotlinClass.extractMetadata(): Pair<NameResolver, ProtoBuf.Class>? {
-        val data = kotlinJvmBinaryClass.classHeader.data ?: return null
-        val strings = kotlinJvmBinaryClass.classHeader.strings ?: return null
-        return JvmProtoBufUtil.readClassDataFrom(data, strings)
-    }
-
     override fun computePackagePartsInfos(packageFqName: FqName): List<PackagePartsCacheData> {
         return packagePartProvider.findPackageParts(packageFqName.asString()).mapNotNull { partName ->
             val classId = ClassId.topLevel(JvmClassName.byInternalName(partName).fqNameForTopLevelClassMaybeWithDollars)
@@ -98,7 +87,37 @@ class KotlinDeserializedJvmSymbolsProvider(
     private val KotlinJvmBinaryClass.isPreReleaseInvisible: Boolean
         get() = classHeader.isPreRelease
 
-    override fun postProcessDeserializedClass(
+    override fun extractClassMetadata(classId: ClassId, parentContext: FirDeserializationContext?): ClassMetadataFindResult? {
+        if (knownNameInPackageCache.hasNoTopLevelClassOf(classId)) return null
+        val result = try {
+            kotlinClassFinder.findKotlinClassOrContent(classId)
+        } catch (e: ProcessCanceledException) {
+            return null
+        }
+        val kotlinClass = when (result) {
+            is KotlinClassFinder.Result.KotlinClass -> result
+            is KotlinClassFinder.Result.ClassFileContent -> {
+                return ClassMetadataFindResult.ClassWithoutMetadata(readClassFromClassFile(classId, result))
+            }
+            null -> return ClassMetadataFindResult.ShouldDeserializeViaParent
+        }
+        if (kotlinClass.kotlinJvmBinaryClass.classHeader.kind != KotlinClassHeader.Kind.CLASS) return null
+        val (nameResolver, classProto) = kotlinClass.extractMetadata() ?: return null
+
+        if (parentContext == null && Flags.CLASS_KIND.get(classProto.flags) == ProtoBuf.Class.Kind.COMPANION_OBJECT) {
+            return ClassMetadataFindResult.ShouldDeserializeViaParent
+        }
+
+        return ClassMetadataFindResult.Metadata(
+            nameResolver,
+            classProto,
+            JvmBinaryAnnotationDeserializer(session, kotlinClass.kotlinJvmBinaryClass, kotlinClassFinder, kotlinClass.byteContent),
+            KotlinJvmBinarySourceElement(kotlinClass.kotlinJvmBinaryClass),
+            classPostProcessor = { loadAnnotationsFromClassFile(kotlinClass, it) }
+        )
+    }
+
+    private fun loadAnnotationsFromClassFile(
         kotlinClass: KotlinClassFinder.Result.KotlinClass,
         symbol: FirRegularClassSymbol
     ) {
@@ -117,29 +136,29 @@ class KotlinDeserializedJvmSymbolsProvider(
         (symbol.fir.annotations as MutableList<FirAnnotationCall>) += annotations
     }
 
-    override fun createAnnotationDeserializer(kotlinClass: KotlinClassFinder.Result.KotlinClass): AbstractAnnotationDeserializer {
+    private fun readClassFromClassFile(classId: ClassId, classFile: KotlinClassFinder.Result.ClassFileContent): FirRegularClassSymbol? {
-        return JvmBinaryAnnotationDeserializer(session, kotlinClass.kotlinJvmBinaryClass, kotlinClassFinder, kotlinClass.byteContent)
+        return javaSymbolProvider.getFirJavaClass(classId, classFile)
     }
 
-    override fun createSourceElement(kotlinClass: KotlinClassFinder.Result.KotlinClass): DeserializedContainerSource {
+    private fun KotlinClassFinder.Result.KotlinClass.extractMetadata(): Pair<NameResolver, ProtoBuf.Class>? {
-        return KotlinJvmBinarySourceElement(kotlinClass.kotlinJvmBinaryClass)
+        val data = kotlinJvmBinaryClass.classHeader.data ?: return null
+        val strings = kotlinJvmBinaryClass.classHeader.strings ?: return null
+        return JvmProtoBufUtil.readClassDataFrom(data, strings)
     }
 
-    private class JvmKnownNameInPackageCache(
+    private class KnownNameInPackageCache(
         session: FirSession,
         private val javaClassFinder: JavaClassFinder
-    ) : KnownNameInPackageCache() {
+    ) {
         private val knownClassNamesInPackage = session.firCachesFactory.createCache(javaClassFinder::knownClassNamesInPackage)
 
         /**
          * This function returns true if we are sure that no top-level class with this id is available
          * If it returns false, it means we can say nothing about this id
          */
-        override fun hasNoTopLevelClassOf(classId: ClassId): Boolean {
+        fun hasNoTopLevelClassOf(classId: ClassId): Boolean {
             val knownNames = knownClassNamesInPackage.getValue(classId.packageFqName) ?: return false
             return classId.relativeClassName.topLevelName() !in knownNames
         }
     }
 }
-
-

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/sessions/FirIdeSessionFactory.kt

@@ -187,9 +187,9 @@ internal object FirIdeSessionFactory {
                         FirThreadSafeSymbolProviderWrapper(
                             KotlinDeserializedJvmSymbolsProvider(
                                 this@apply,
+                                kotlinScopeProvider,
                                 packagePartProvider,
                                 kotlinClassFinder,
-                                kotlinScopeProvider,
                                 javaSymbolProvider,
                                 javaClassFinder
                             )