JVM IR: change generation scheme of property $delegate methods

Generate $delegate method as instance method in
PropertyReferenceDelegationLowering, and remove dispatch receiver later
in MakePropertyDelegateMethodsStatic. The method needs to be static to
be non-overridable (see delegateMethodIsNonOverridable.kt), and public
to be accessible in reflection.

Otherwise we generated incorrect IR where a static function accessed an
instance field of the containing class, which failed in multiple places
including LocalDeclarationsLowering.

 #KT-48350 Fixed

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/LocalDeclarationsLowering.kt

@@ -185,7 +185,8 @@ class LocalDeclarationsLowering(
         override fun irGet(startOffset: Int, endOffset: Int, valueDeclaration: IrValueDeclaration): IrExpression? {
             val field = classContext.capturedValueToField[valueDeclaration] ?: return null
 
-            val receiver = member.dispatchReceiverParameter!!
+            val receiver = member.dispatchReceiverParameter
+                ?: error("No dispatch receiver parameter for ${member.render()}")
             return IrGetFieldImpl(
                 startOffset, endOffset, field.symbol, field.type,
                 receiver = IrGetValueImpl(startOffset, endOffset, receiver.type, receiver.symbol)

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmLower.kt

@@ -416,6 +416,7 @@ private val jvmFilePhases = listOf(
     typeOperatorLowering,
     replaceKFunctionInvokeWithFunctionInvokePhase,
     kotlinNothingValueExceptionPhase,
+    makePropertyDelegateMethodsStaticPhase,
 
     renameFieldsPhase,
     fakeInliningLocalVariablesLowering,

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/JvmPropertiesLowering.kt

@@ -27,7 +27,9 @@ import org.jetbrains.kotlin.ir.expressions.IrCall
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrFieldAccessExpression
 import org.jetbrains.kotlin.ir.expressions.impl.IrBlockBodyImpl
-import org.jetbrains.kotlin.ir.types.*
+import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.types.classifierOrNull
+import org.jetbrains.kotlin.ir.types.makeNotNull
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.load.java.JvmAbi
@@ -185,14 +187,15 @@ class JvmPropertiesLowering(private val backendContext: JvmBackendContext) : IrE
             this.origin = origin
             this.returnType = returnType
         }.apply {
-            var index = 0
-            valueParameters = listOfNotNull(
-                declaration.getter?.dispatchReceiverParameter?.takeIf { !isStatic }?.let {
+            if (!isStatic) {
+                dispatchReceiverParameter = declaration.getter?.dispatchReceiverParameter?.let {
                     // Synthetic methods don't get generic type signatures anyway, so not exactly useful to preserve type parameters.
-                    it.copyTo(this, type = it.type.eraseTypeParameters(), index = index++)
-                },
+                    it.copyTo(this, type = it.type.eraseTypeParameters())
+                }
+            }
+            valueParameters = listOfNotNull(
                 declaration.getter?.extensionReceiverParameter?.let {
-                    it.copyTo(this, type = it.type.eraseTypeParameters(), index = index)
+                    it.copyTo(this, type = it.type.eraseTypeParameters(), index = 0)
                 }
             )
             parent = declaration.parent

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/MakePropertyDelegateMethodsStaticLowering.kt

@@ -0,0 +1,73 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.backend.jvm.lower
+
+import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.ir.copyTo
+import org.jetbrains.kotlin.backend.common.lower.VariableRemapper
+import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
+import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.backend.jvm.localDeclarationsPhase
+import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
+import org.jetbrains.kotlin.ir.expressions.IrCall
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.util.render
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
+import org.jetbrains.kotlin.load.java.JvmAbi
+
+// This phase is needed to support correct generation of synthetic `$delegate` methods for optimized delegated properties, in the case
+// when receiver of the optimized property reference is a field from an outer class.
+//
+// Since PropertyReferenceDelegationLowering runs before LocalDeclarationsLowering, fields for captured this (aka `this$0`) are not
+// generated yet. And there's no other way to obtain the instance of the outer class on an arbitrary value of an inner class.
+// However, we need `$delegate` methods to be static to be non-overridable (and public to be visible in reflection and external tools).
+//
+// So PropertyReferenceDelegationLowering generates `$delegate` methods for optimized property references as instance methods,
+// and this phase, which runs _after_ LocalDeclarationsLowering, transforms them to static methods.
+internal val makePropertyDelegateMethodsStaticPhase = makeIrFilePhase(
+    ::MakePropertyDelegateMethodsStaticLowering,
+    name = "MakePropertyDelegateMethodsStatic",
+    description = "Make `\$delegate` methods for optimized delegated properties static",
+    prerequisite = setOf(propertyReferenceDelegationPhase, localDeclarationsPhase)
+)
+
+private class MakePropertyDelegateMethodsStaticLowering(val context: JvmBackendContext) : IrElementTransformerVoid(), FileLoweringPass {
+    override fun lower(irFile: IrFile) {
+        irFile.transform(this, null)
+    }
+
+    override fun visitSimpleFunction(declaration: IrSimpleFunction): IrStatement {
+        if (!declaration.isSyntheticDelegateMethod()) return super.visitSimpleFunction(declaration)
+
+        val oldParameter = declaration.dispatchReceiverParameter ?: return super.visitSimpleFunction(declaration)
+        val newParameter = oldParameter.copyTo(declaration, index = 0)
+
+        return declaration.apply {
+            valueParameters =
+                listOf(newParameter) + valueParameters.map { it.copyTo(this, index = it.index + 1) }
+            dispatchReceiverParameter = null
+            body = body?.transform(VariableRemapper(mapOf(oldParameter to newParameter)), null)
+        }
+    }
+
+    override fun visitCall(expression: IrCall): IrExpression {
+        // There should be no calls to synthetic `$delegate` methods because they aren't accessible in the source code, and we don't
+        // generate any calls in the IR. Otherwise we would need to remap arguments in those calls.
+        if (expression.symbol.owner.isSyntheticDelegateMethod()) {
+            error(
+                "`\$delegate` method should not be called. Please either remove the call, or support remapping of dispatch receiver " +
+                        "in MakePropertyDelegateMethodsStaticLowering: ${expression.symbol.owner.render()}"
+            )
+        }
+        return super.visitCall(expression)
+    }
+
+    private fun IrSimpleFunction.isSyntheticDelegateMethod(): Boolean =
+        origin == IrDeclarationOrigin.PROPERTY_DELEGATE && name.asString().endsWith(JvmAbi.DELEGATED_PROPERTY_NAME_SUFFIX)
+}

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/PropertyReferenceDelegationLowering.kt

@@ -151,11 +151,11 @@ private class PropertyReferenceDelegationTransformer(val context: JvmBackendCont
         getter?.apply { body = accessorBody(delegate, backingField ?: receiver?.inline(originalThis, dispatchReceiverParameter)) }
         setter?.apply { body = accessorBody(delegate, backingField ?: receiver?.inline(originalThis, dispatchReceiverParameter)) }
 
+        // The `$delegate` method is generated as instance method here, see MakePropertyDelegateMethodsStaticLowering.
         val delegateMethod = context.createSyntheticMethodForPropertyDelegate(this).apply {
             body = context.createJvmIrBuilder(symbol).run {
-                val propertyOwner = if (getter?.dispatchReceiverParameter != null) valueParameters[0] else null
-                val boundReceiver = backingField?.let { irGetField(propertyOwner?.let(::irGet), it) }
-                    ?: receiver?.inline(originalThis, propertyOwner)
+                val boundReceiver = backingField?.let { irGetField(dispatchReceiverParameter?.let(::irGet), it) }
+                    ?: receiver?.inline(originalThis, dispatchReceiverParameter)
                 irExprBody(with(delegate) {
                     val origin = PropertyReferenceLowering.REFLECTED_PROPERTY_REFERENCE
                     IrPropertyReferenceImpl(startOffset, endOffset, type, symbol, typeArgumentsCount, field, getter, setter, origin)