FIR: Optimize deserialized annotations loading

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

@@ -27,7 +27,8 @@ import org.jetbrains.kotlin.protobuf.MessageLite
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedContainerSource
 
 class JvmBinaryAnnotationDeserializer(
-    val session: FirSession
+    val session: FirSession,
+    private var byteContent: ByteArray?
 ) : AbstractAnnotationDeserializer(session) {
     private val storage: MutableMap<KotlinJvmBinaryClass, MemberAnnotations> = mutableMapOf()
 
@@ -190,7 +191,9 @@ class JvmBinaryAnnotationDeserializer(
                     }
                 }
             }
-        }, null) // TODO: cached file content?
+        }, byteContent)
+
+        byteContent = null
 
         val result = MemberAnnotations(memberAnnotations)
         storage[kotlinClass] = result

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

@@ -16,25 +16,29 @@ import org.jetbrains.kotlin.fir.declarations.FirVariable
 import org.jetbrains.kotlin.fir.declarations.impl.FirSimpleFunctionImpl
 import org.jetbrains.kotlin.fir.deserialization.FirDeserializationContext
 import org.jetbrains.kotlin.fir.deserialization.deserializeClassToSymbol
-import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
-import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
+import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
+import org.jetbrains.kotlin.fir.expressions.FirClassReferenceExpression
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.buildUnaryArgumentList
 import org.jetbrains.kotlin.fir.expressions.builder.*
 import org.jetbrains.kotlin.fir.java.JavaSymbolProvider
 import org.jetbrains.kotlin.fir.java.createConstant
 import org.jetbrains.kotlin.fir.java.topLevelName
 import org.jetbrains.kotlin.fir.references.builder.buildErrorNamedReference
 import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference
-import org.jetbrains.kotlin.fir.resolve.*
+import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.AbstractFirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.getClassDeclaredCallableSymbols
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.KotlinScopeProvider
 import org.jetbrains.kotlin.fir.scopes.impl.nestedClassifierScope
+import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.constructClassType
 import org.jetbrains.kotlin.load.java.JavaClassFinder
 import org.jetbrains.kotlin.load.java.structure.JavaClass
 import org.jetbrains.kotlin.load.kotlin.*
@@ -105,7 +109,8 @@ class KotlinDeserializedJvmSymbolsProvider(
         return packagePartProvider.findPackageParts(packageFqName.asString()).mapNotNull { partName ->
             val classId = ClassId.topLevel(JvmClassName.byInternalName(partName).fqNameForTopLevelClassMaybeWithDollars)
             if (hasNoTopLevelClassOf(classId)) return@mapNotNull null
-            val kotlinJvmBinaryClass = kotlinClassFinder.findKotlinClass(classId) ?: return@mapNotNull null
+            val (kotlinJvmBinaryClass, byteContent) =
+                kotlinClassFinder.findKotlinClassOrContent(classId) as? KotlinClassFinder.Result.KotlinClass ?: return@mapNotNull null
 
             val header = kotlinJvmBinaryClass.classHeader
             val data = header.data ?: header.incompatibleData ?: return@mapNotNull null
@@ -121,7 +126,7 @@ class KotlinDeserializedJvmSymbolsProvider(
                 packageProto,
                 FirDeserializationContext.createForPackage(
                     packageFqName, packageProto, nameResolver, session,
-                    JvmBinaryAnnotationDeserializer(session),
+                    JvmBinaryAnnotationDeserializer(session, byteContent),
                     source
                 ),
                 source,
@@ -164,20 +169,15 @@ class KotlinDeserializedJvmSymbolsProvider(
         return JvmProtoBufUtil.readClassDataFrom(data, strings)
     }
 
-    private fun FirClassifierSymbol<*>?.toDefaultResolvedTypeRef(classId: ClassId): FirResolvedTypeRef {
+    private fun ConeClassLikeLookupTag.toDefaultResolvedTypeRef(classId: ClassId): FirResolvedTypeRef =
-        return this?.let {
+        buildResolvedTypeRef {
-            buildResolvedTypeRef {
+            type = constructClassType(emptyArray(), isNullable = false)
-                type = it.constructType(emptyList(), isNullable = false)
+        }
-            }
-        } ?: buildErrorTypeRef { diagnostic = ConeSimpleDiagnostic("Symbol not found for $classId", DiagnosticKind.Java) }
-
-    }
 
     private fun loadAnnotation(
         annotationClassId: ClassId, result: MutableList<FirAnnotationCall>,
     ): KotlinJvmBinaryClass.AnnotationArgumentVisitor? {
         val lookupTag = ConeClassLikeLookupTagImpl(annotationClassId)
-        val symbol = lookupTag.toSymbol(session)
 
         return object : KotlinJvmBinaryClass.AnnotationArgumentVisitor {
             private val argumentMap = mutableMapOf<Name, FirExpression>()
@@ -190,9 +190,8 @@ class KotlinDeserializedJvmSymbolsProvider(
 
             private fun ClassLiteralValue.toFirClassReferenceExpression(): FirClassReferenceExpression {
                 val literalLookupTag = ConeClassLikeLookupTagImpl(classId)
-                val literalSymbol = literalLookupTag.toSymbol(this@KotlinDeserializedJvmSymbolsProvider.session)
                 return buildClassReferenceExpression {
-                    classTypeRef = literalSymbol.toDefaultResolvedTypeRef(classId)
+                    classTypeRef = literalLookupTag.toDefaultResolvedTypeRef(classId)
                 }
             }
 
@@ -275,7 +274,7 @@ class KotlinDeserializedJvmSymbolsProvider(
 
             override fun visitEnd() {
                 result += buildAnnotationCall {
-                    annotationTypeRef = symbol.toDefaultResolvedTypeRef(annotationClassId)
+                    annotationTypeRef = lookupTag.toDefaultResolvedTypeRef(annotationClassId)
                     argumentList = buildArgumentList {
                         for ((name, expression) in argumentMap) {
                             arguments += buildNamedArgumentExpression {
@@ -315,8 +314,8 @@ class KotlinDeserializedJvmSymbolsProvider(
         } catch (e: ProcessCanceledException) {
             return null
         }
-        val kotlinJvmBinaryClass = when (result) {
+        val kotlinClassWithContent = when (result) {
-            is KotlinClassFinder.Result.KotlinClass -> result.kotlinJvmBinaryClass
+            is KotlinClassFinder.Result.KotlinClass -> result
             is KotlinClassFinder.Result.ClassFileContent -> {
                 handledByJava.add(classId)
                 return try {
@@ -327,17 +326,19 @@ class KotlinDeserializedJvmSymbolsProvider(
             }
             null -> null
         }
-        if (kotlinJvmBinaryClass == null) {
+        if (kotlinClassWithContent == null) {
             val outerClassId = classId.outerClassId ?: return null
             findAndDeserializeClass(outerClassId) ?: return null
         } else {
+            val (kotlinJvmBinaryClass, byteContent) = kotlinClassWithContent
+
             if (kotlinJvmBinaryClass.classHeader.kind != KotlinClassHeader.Kind.CLASS) return null
             val (nameResolver, classProto) = kotlinJvmBinaryClass.readClassDataFrom() ?: return null
 
             val symbol = FirRegularClassSymbol(classId)
             deserializeClassToSymbol(
                 classId, classProto, symbol, nameResolver, session,
-                JvmBinaryAnnotationDeserializer(session),
+                JvmBinaryAnnotationDeserializer(session, byteContent),
                 kotlinScopeProvider,
                 parentContext, KotlinJvmBinarySourceElement(kotlinJvmBinaryClass),
                 this::findAndDeserializeClass
@@ -356,7 +357,7 @@ class KotlinDeserializedJvmSymbolsProvider(
 
 
                 },
-                null,
+                byteContent,
             )
             (symbol.fir.annotations as MutableList<FirAnnotationCall>) += annotations
         }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/deserialization/FirMemberDeserializer.kt

@@ -280,7 +280,7 @@ class FirMemberDeserializer(private val c: FirDeserializationContext) {
         }
     }
 
-    fun loadFunction(proto: ProtoBuf.Function): FirSimpleFunction {
+    fun loadFunction(proto: ProtoBuf.Function, byteContent: ByteArray? = null): FirSimpleFunction {
         val flags = if (proto.hasFlags()) proto.flags else loadOldFlags(proto.oldFlags)
 
         val receiverAnnotations =

compiler/frontend.java/src/org/jetbrains/kotlin/load/kotlin/VirtualFileKotlinClass.kt

@@ -68,7 +68,7 @@ class VirtualFileKotlinClass private constructor(
                             VirtualFileKotlinClass(file, name, classVersion, header, innerClasses)
                         }
 
-                        return@time kotlinJvmBinaryClass?.let(::KotlinClass)
+                        return@time kotlinJvmBinaryClass?.let { KotlinClass(it, byteContent) }
                             ?: KotlinClassFinder.Result.ClassFileContent(byteContent)
                     }
                 }

core/descriptors.jvm/src/org/jetbrains/kotlin/load/kotlin/KotlinClassFinder.kt

@@ -28,8 +28,12 @@ interface KotlinClassFinder : KotlinMetadataFinder {
     sealed class Result {
         fun toKotlinJvmBinaryClass(): KotlinJvmBinaryClass? = (this as? KotlinClass)?.kotlinJvmBinaryClass
 
-        data class KotlinClass(val kotlinJvmBinaryClass: KotlinJvmBinaryClass) : Result()
+        class KotlinClass(val kotlinJvmBinaryClass: KotlinJvmBinaryClass, val byteContent: ByteArray? = null) : Result() {
-        data class ClassFileContent(
+            operator fun component1(): KotlinJvmBinaryClass = kotlinJvmBinaryClass
+            operator fun component2(): ByteArray? = byteContent
+        }
+
+        class ClassFileContent(
             @Suppress("ArrayInDataClass")
             val content: ByteArray
         ) : Result()