JVM_IR: load serialized data for every affected property of IrLazyFunction

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/lazy/IrLazyConstructor.kt

@@ -42,15 +42,19 @@ class IrLazyConstructor(
 
     override var body: IrBody? = null
 
-    override var returnType: IrType by createReturnType()
+    override var returnType: IrType by lazyVar(stubGenerator.lock) { createReturnType() }
 
     override val initialSignatureFunction: IrFunction? by createInitialSignatureFunction()
 
-    override var dispatchReceiverParameter: IrValueParameter? by createReceiverParameter(descriptor.dispatchReceiverParameter)
+    override var dispatchReceiverParameter: IrValueParameter? by lazyVar(stubGenerator.lock) {
+        createReceiverParameter(descriptor.dispatchReceiverParameter)
+    }
 
-    override var extensionReceiverParameter: IrValueParameter? by createReceiverParameter(descriptor.extensionReceiverParameter)
+    override var extensionReceiverParameter: IrValueParameter? by lazyVar(stubGenerator.lock) {
+        createReceiverParameter(descriptor.extensionReceiverParameter)
+    }
 
-    override var valueParameters: List<IrValueParameter> by createValueParameters()
+    override var valueParameters: List<IrValueParameter> by lazyVar(stubGenerator.lock) { createValueParameters() }
 
     override var metadata: MetadataSource?
         get() = null

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/lazy/IrLazyFunction.kt

@@ -49,29 +49,34 @@ class IrLazyFunction(
 
     override var annotations: List<IrConstructorCall> by createLazyAnnotations()
 
-    override var body: IrBody? by lazyVar {
-        if (!isInline) return@lazyVar null
-        val toplevel = getToplevel()
-        val loaded = (toplevel as? DeserializableClass)?.loadIR() ?: false
-        // Underlying field has been set by IR loading
-        if (loaded) body else null
+    override var body: IrBody? by lazyVar(stubGenerator.lock) {
+        if (tryLoadIr()) body else null
     }
 
-    override var returnType: IrType by createReturnType()
+    override var returnType: IrType by lazyVar(stubGenerator.lock) {
+        if (tryLoadIr()) returnType else createReturnType()
+    }
 
     override val initialSignatureFunction: IrFunction? by createInitialSignatureFunction()
 
-    override var dispatchReceiverParameter: IrValueParameter? by createReceiverParameter(descriptor.dispatchReceiverParameter, true)
+    override var dispatchReceiverParameter: IrValueParameter? by lazyVar(stubGenerator.lock) {
+        if (tryLoadIr()) dispatchReceiverParameter else createReceiverParameter(descriptor.dispatchReceiverParameter, true)
+    }
 
-    override var extensionReceiverParameter: IrValueParameter? by createReceiverParameter(descriptor.extensionReceiverParameter)
+    override var extensionReceiverParameter: IrValueParameter? by lazyVar(stubGenerator.lock) {
+        if (tryLoadIr()) extensionReceiverParameter else createReceiverParameter(descriptor.extensionReceiverParameter)
+    }
 
-    override var valueParameters: List<IrValueParameter> by createValueParameters()
+    override var valueParameters: List<IrValueParameter> by lazyVar(stubGenerator.lock) {
+        if (tryLoadIr()) valueParameters else createValueParameters()
+    }
 
     override var metadata: MetadataSource?
         get() = null
         set(_) = error("We should never need to store metadata of external declarations.")
 
     override var typeParameters: List<IrTypeParameter> by lazyVar(stubGenerator.lock) {
+        if (tryLoadIr()) return@lazyVar typeParameters
         typeTranslator.buildWithScope(this) {
             stubGenerator.symbolTable.withScope(this) {
                 val propertyIfAccessor = descriptor.propertyIfAccessor
@@ -103,15 +108,22 @@ class IrLazyFunction(
     override val containerSource: DeserializedContainerSource?
         get() = (descriptor as? DescriptorWithContainerSource)?.containerSource
 
-    init {
-        symbol.bind(this)
+    private fun tryLoadIr(): Boolean {
+        if (!isInline || isFakeOverride) return false
+        val toplevel = getToplevel()
+        return (toplevel as? DeserializableClass)?.loadIr() ?: false
     }
 
     private fun getToplevel(): IrDeclaration {
-        var current: IrDeclarationParent = parent
-        while (current is IrDeclaration && current.parent !is IrPackageFragment) {
-            current = current.parent
+        var current: IrDeclaration = this
+        while (current.parent !is IrPackageFragment) {
+            current = current.parent as IrDeclaration
         }
-        return current as IrDeclaration
+        return current
+    }
+
+    init {
+        symbol.bind(this)
     }
 }
+

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/lazy/IrLazyFunctionBase.kt

@@ -25,30 +25,24 @@ interface IrLazyFunctionBase : IrLazyDeclarationBase, IrTypeParametersContainer
             descriptor.initialSignatureDescriptor?.takeIf { it != descriptor }?.original?.let(stubGenerator::generateFunctionStub)
         }
 
-    fun createValueParameters(): ReadWriteProperty<Any?, List<IrValueParameter>> =
-        lazyVar(stubGenerator.lock) {
-            typeTranslator.buildWithScope(this) {
-                descriptor.valueParameters.mapTo(arrayListOf()) {
-                    stubGenerator.generateValueParameterStub(it).apply { parent = this@IrLazyFunctionBase }
-                }
+    fun createValueParameters(): List<IrValueParameter> =
+        typeTranslator.buildWithScope(this) {
+            descriptor.valueParameters.mapTo(arrayListOf()) {
+                stubGenerator.generateValueParameterStub(it).apply { parent = this@IrLazyFunctionBase }
             }
         }
 
     fun createReceiverParameter(
         parameter: ReceiverParameterDescriptor?,
         functionDispatchReceiver: Boolean = false,
-    ): ReadWriteProperty<Any?, IrValueParameter?> =
-        lazyVar(stubGenerator.lock) {
-            if (functionDispatchReceiver && stubGenerator.extensions.isStaticFunction(descriptor)) null
-            else typeTranslator.buildWithScope(this) {
-                parameter?.generateReceiverParameterStub()?.also { it.parent = this@IrLazyFunctionBase }
-            }
+    ): IrValueParameter? =
+        if (functionDispatchReceiver && stubGenerator.extensions.isStaticFunction(descriptor)) null
+        else typeTranslator.buildWithScope(this) {
+            parameter?.generateReceiverParameterStub()?.also { it.parent = this@IrLazyFunctionBase }
         }
 
-    fun createReturnType(): ReadWriteProperty<Any?, IrType> =
-        lazyVar(stubGenerator.lock) {
-            typeTranslator.buildWithScope(this) {
-                descriptor.returnType!!.toIrType()
-            }
+    fun createReturnType(): IrType =
+        typeTranslator.buildWithScope(this) {
+            descriptor.returnType!!.toIrType()
         }
 }

compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/serialization/IrDeclarationDeserializer.kt

@@ -227,7 +227,7 @@ class IrDeclarationDeserializer(
                 coordinates.startOffset, coordinates.endOffset,
                 deserializeIrDeclarationOrigin(proto.originName), proto.flags
             )
-            result.annotations += deserializeAnnotations(proto.annotationList)
+            result.annotations = deserializeAnnotations(proto.annotationList)
             if (!skipMutableState) {
                 result.parent = currentParent
             }
@@ -327,7 +327,7 @@ class IrDeclarationDeserializer(
                     withExternalValue(isExternal) {
                         proto.declarationList
                             .filterNot { isSkippableFakeOverride(it, this) }
-                            .mapTo(declarations) { deserializeDeclaration(it) }
+                            .mapNotNullTo(declarations) { declProto -> deserializeDeclaration(declProto) }
                     }
 
                     thisReceiver = deserializeIrValueParameter(proto.thisReceiver, -1)
@@ -518,13 +518,15 @@ class IrDeclarationDeserializer(
 
                     withBodyGuard {
                         valueParameters = deserializeValueParameters(proto.valueParameterList)
-                        if (proto.hasDispatchReceiver())
-                            dispatchReceiverParameter = deserializeIrValueParameter(proto.dispatchReceiver, -1)
-                        if (proto.hasExtensionReceiver())
-                            extensionReceiverParameter = deserializeIrValueParameter(proto.extensionReceiver, -1)
-                        if (proto.hasBody()) {
-                            body = deserializeStatementBody(proto.body) as IrBody
-                        }
+                        dispatchReceiverParameter =
+                            if (proto.hasDispatchReceiver()) deserializeIrValueParameter(proto.dispatchReceiver, -1)
+                            else null
+                        extensionReceiverParameter =
+                            if (proto.hasExtensionReceiver()) deserializeIrValueParameter(proto.extensionReceiver, -1)
+                            else null
+                        body =
+                            if (proto.hasBody()) deserializeStatementBody(proto.body) as IrBody
+                            else null
                     }
                 }
             }
@@ -764,4 +766,4 @@ class IrDeclarationDeserializer(
             else -> false
         }
     }
-}
+}