FIR: Simplify JvmBinaryAnnotationDeserializer

Currently, it's anyway created one per class

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

@@ -11,9 +11,7 @@ import org.jetbrains.kotlin.fir.deserialization.AbstractAnnotationDeserializer
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.impl.FirCompositeSymbolProvider
-import org.jetbrains.kotlin.load.kotlin.JvmPackagePartSource
 import org.jetbrains.kotlin.load.kotlin.KotlinJvmBinaryClass
-import org.jetbrains.kotlin.load.kotlin.KotlinJvmBinarySourceElement
 import org.jetbrains.kotlin.load.kotlin.MemberSignature
 import org.jetbrains.kotlin.metadata.ProtoBuf
 import org.jetbrains.kotlin.metadata.deserialization.NameResolver
@@ -28,12 +26,12 @@ import org.jetbrains.kotlin.serialization.deserialization.descriptors.Deserializ
 
 class JvmBinaryAnnotationDeserializer(
     val session: FirSession,
-    private var byteContent: ByteArray?
+    kotlinBinaryClass: KotlinJvmBinaryClass,
+    byteContent: ByteArray?
 ) : AbstractAnnotationDeserializer(session) {
-    private val storage: MutableMap<KotlinJvmBinaryClass, MemberAnnotations> = mutableMapOf()
-
-    // TODO: Rename this once property constants are recorded as well
-    private data class MemberAnnotations(val memberAnnotations: Map<MemberSignature, MutableList<FirAnnotationCall>>)
+    private val annotationInfo by lazy(LazyThreadSafetyMode.PUBLICATION) {
+        session.loadMemberAnnotations(kotlinBinaryClass, byteContent)
+    }
 
     private enum class CallableKind {
         PROPERTY_GETTER,
@@ -133,25 +131,18 @@ class JvmBinaryAnnotationDeserializer(
         containerSource: DeserializedContainerSource?,
         memberSignature: MemberSignature
     ): List<FirAnnotationCall> {
-        val kotlinClass = containerSource?.toKotlinJvmBinaryClass() ?: return emptyList()
-        return loadMemberAnnotations(kotlinClass).memberAnnotations[memberSignature] ?: emptyList()
+        return annotationInfo.memberAnnotations[memberSignature] ?: emptyList()
+    }
 }
 
-    private fun DeserializedContainerSource.toKotlinJvmBinaryClass(): KotlinJvmBinaryClass? =
-        when (this) {
-            is JvmPackagePartSource -> this.knownJvmBinaryClass
-            is KotlinJvmBinarySourceElement -> this.binaryClass
-            else -> null
-        }
+// TODO: Rename this once property constants are recorded as well
+private data class MemberAnnotations(val memberAnnotations: Map<MemberSignature, MutableList<FirAnnotationCall>>)
 
 // TODO: better to be in KotlinDeserializedJvmSymbolsProvider?
-    private fun loadMemberAnnotations(kotlinClass: KotlinJvmBinaryClass): MemberAnnotations {
-        if (storage.containsKey(kotlinClass)) {
-            return storage[kotlinClass] ?: error("$kotlinClass should have been visited and cached.")
-        }
+private fun FirSession.loadMemberAnnotations(kotlinBinaryClass: KotlinJvmBinaryClass, byteContent: ByteArray?): MemberAnnotations {
     val memberAnnotations = hashMapOf<MemberSignature, MutableList<FirAnnotationCall>>()
 
-        kotlinClass.visitMembers(object : KotlinJvmBinaryClass.MemberVisitor {
+    kotlinBinaryClass.visitMembers(object : KotlinJvmBinaryClass.MemberVisitor {
         override fun visitMethod(name: Name, desc: String): KotlinJvmBinaryClass.MethodAnnotationVisitor? {
             return AnnotationVisitorForMethod(MemberSignature.fromMethodNameAndDesc(name.asString(), desc))
         }
@@ -197,21 +188,16 @@ class JvmBinaryAnnotationDeserializer(
         }
     }, byteContent)
 
-        byteContent = null
-
-        val result = MemberAnnotations(memberAnnotations)
-        storage[kotlinClass] = result
-        return result
+    return MemberAnnotations(memberAnnotations)
 }
 
 // TODO: Or, better to migrate annotation deserialization in KotlinDeserializedJvmSymbolsProvider to here?
-    private fun loadAnnotationIfNotSpecial(
+private fun FirSession.loadAnnotationIfNotSpecial(
     annotationClassId: ClassId,
     result: MutableList<FirAnnotationCall>
 ): KotlinJvmBinaryClass.AnnotationArgumentVisitor? =
-        (session.firSymbolProvider as? FirCompositeSymbolProvider)
+    (firSymbolProvider as? FirCompositeSymbolProvider)
         ?.providers
         ?.filterIsInstance<KotlinDeserializedJvmSymbolsProvider>()
         ?.singleOrNull()
         ?.loadAnnotationIfNotSpecial(annotationClassId, result)
-}

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

@@ -126,7 +126,7 @@ class KotlinDeserializedJvmSymbolsProvider(
                 packageProto,
                 FirDeserializationContext.createForPackage(
                     packageFqName, packageProto, nameResolver, session,
-                    JvmBinaryAnnotationDeserializer(session, byteContent),
+                    JvmBinaryAnnotationDeserializer(session, kotlinJvmBinaryClass, byteContent),
                     source
                 ),
                 source,
@@ -338,7 +338,7 @@ class KotlinDeserializedJvmSymbolsProvider(
             val symbol = FirRegularClassSymbol(classId)
             deserializeClassToSymbol(
                 classId, classProto, symbol, nameResolver, session,
-                JvmBinaryAnnotationDeserializer(session, byteContent),
+                JvmBinaryAnnotationDeserializer(session, kotlinJvmBinaryClass, byteContent),
                 kotlinScopeProvider,
                 parentContext, KotlinJvmBinarySourceElement(kotlinJvmBinaryClass),
                 this::findAndDeserializeClass