Fix lowering of receiver access in IR scripting

#KT-53947 fixed
without the fix, in the presence of the implicit receiver of the same
type as the processed receive, the lowering was incorrectly substituting
the correct receiver with the accessor to the implicit one

Author: ligee

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

@@ -598,7 +598,7 @@ private class ScriptToClassTransformer(
             val ctorDispatchReceiverType = expression.symbol.owner.dispatchReceiverParameter?.type
                 ?: if (capturingClassesConstructors.keys.any { it.symbol == expression.symbol }) scriptClassReceiver.type else null
             if (ctorDispatchReceiverType != null) {
-                getDispatchReceiverExpression(data, expression, ctorDispatchReceiverType, expression.origin)?.let {
+                getDispatchReceiverExpression(data, expression, ctorDispatchReceiverType, expression.origin, null)?.let {
                     expression.dispatchReceiver = it
                 }
             }
@@ -607,18 +607,29 @@ private class ScriptToClassTransformer(
     }
 
     private fun getDispatchReceiverExpression(
-        data: ScriptToClassTransformerContext, expression: IrDeclarationReference, receiverType: IrType, origin: IrStatementOrigin?,
+        data: ScriptToClassTransformerContext,
+        expression: IrDeclarationReference,
+        receiverType: IrType,
+        origin: IrStatementOrigin?,
+        originalReceiverParameter: IrValueParameter?,
     ): IrExpression? {
         return if (receiverType == scriptClassReceiver.type) {
-            getAccessCallForScriptInstance(data, expression.startOffset, expression.endOffset, origin)
+            getAccessCallForScriptInstance(data, expression.startOffset, expression.endOffset, origin, originalReceiverParameter)
         } else {
-            getAccessCallForImplicitReceiver(data, expression, receiverType, origin)
+            getAccessCallForImplicitReceiver(data, expression, receiverType, origin, originalReceiverParameter)
         }
     }
 
     private fun getAccessCallForScriptInstance(
-        data: ScriptToClassTransformerContext, startOffset: Int, endOffset: Int, origin: IrStatementOrigin?
-    ) = when {
+        data: ScriptToClassTransformerContext,
+        startOffset: Int,
+        endOffset: Int,
+        origin: IrStatementOrigin?,
+        originalReceiverParameter: IrValueParameter?
+    ): IrExpression? = when {
+        originalReceiverParameter != null && originalReceiverParameter != scriptClassReceiver ->
+            null
+
         data.fieldForScriptThis != null ->
             IrGetFieldImpl(
                 startOffset, endOffset,
@@ -634,30 +645,36 @@ private class ScriptToClassTransformer(
                         origin
                     )
             }
+
         data.valueParameterForScriptThis != null ->
             IrGetValueImpl(
                 startOffset, endOffset,
                 scriptClassReceiver.type,
                 data.valueParameterForScriptThis,
                 origin
             )
+
         else -> error("Unexpected script transformation state: $data")
     }
 
     private fun getAccessCallForImplicitReceiver(
         data: ScriptToClassTransformerContext,
         expression: IrDeclarationReference,
         receiverType: IrType,
-        expressionOrigin: IrStatementOrigin?
+        expressionOrigin: IrStatementOrigin?,
+        originalReceiverParameter: IrValueParameter?
     ): IrExpression? {
         // implicit receivers has priority (as per descriptor outer scopes)
-        implicitReceiversFieldsWithParameters.firstOrNull { it.second.type == receiverType }?.let { (field, param) ->
+        implicitReceiversFieldsWithParameters.firstOrNull {
+            if (originalReceiverParameter != null) it.second == originalReceiverParameter
+            else it.second.type == receiverType
+        }?.let { (field, param) ->
             val builder = context.createIrBuilder(expression.symbol)
             return if (data.isInScriptConstructor) {
                 builder.irGet(param.type, param.symbol)
             } else {
                 val scriptReceiver =
-                    getAccessCallForScriptInstance(data, expression.startOffset, expression.endOffset, expressionOrigin)
+                    getAccessCallForScriptInstance(data, expression.startOffset, expression.endOffset, expressionOrigin, null)
                 builder.irGetField(scriptReceiver, field)
             }
         }
@@ -692,7 +709,7 @@ private class ScriptToClassTransformer(
                     builder.irGet(objArray.defaultType, irScript.earlierScriptsParameter!!.symbol)
                 } else {
                     val scriptReceiver =
-                        getAccessCallForScriptInstance(data, expression.startOffset, expression.endOffset, expressionOrigin)
+                        getAccessCallForScriptInstance(data, expression.startOffset, expression.endOffset, expressionOrigin, null)
                     builder.irGetField(scriptReceiver, earlierScriptsField!!)
                 }
             val getPrevScriptObjectExpression = builder.irCall(objArrayGet).apply {
@@ -711,40 +728,13 @@ private class ScriptToClassTransformer(
         return null
     }
 
-    override fun visitGetField(expression: IrGetField, data: ScriptToClassTransformerContext): IrExpression {
-        if (irScript.needsReceiverProcessing) {
-            val receiver = expression.receiver
-            if (receiver is IrGetValue && receiver.symbol.owner.name == SpecialNames.THIS) {
-                val newReceiver = getDispatchReceiverExpression(data, expression, receiver.type, expression.origin)
-                if (newReceiver != null) {
-                    val newGetField =
-                        IrGetFieldImpl(expression.startOffset, expression.endOffset, expression.symbol, expression.type, newReceiver)
-                    return super.visitGetField(newGetField, data)
-                }
-            }
-        }
-        return super.visitGetField(expression, data)
-    }
-
-    override fun visitCall(expression: IrCall, data: ScriptToClassTransformerContext): IrExpression {
-        if (irScript.needsReceiverProcessing) {
-            val target = expression.symbol.owner
-            val receiver: IrValueParameter? = target.dispatchReceiverParameter
-            if (receiver?.name == SpecialNames.THIS) {
-                val newReceiver = getDispatchReceiverExpression(data, expression, receiver.type, expression.origin)
-                if (newReceiver != null) {
-                    expression.dispatchReceiver = newReceiver
-                }
-            }
-        }
-        return super.visitCall(expression, data) as IrExpression
-    }
-
     override fun visitGetValue(expression: IrGetValue, data: ScriptToClassTransformerContext): IrExpression {
         if (irScript.needsReceiverProcessing) {
             val getValueParameter = expression.symbol.owner as? IrValueParameter
             if (getValueParameter != null && getValueParameter.name == SpecialNames.THIS) {
-                val newExpression = getDispatchReceiverExpression(data, expression, getValueParameter.type, expression.origin)
+                val newExpression = getDispatchReceiverExpression(
+                    data, expression, getValueParameter.type, expression.origin, getValueParameter
+                )
                 if (newExpression != null) {
                     return super.visitExpression(newExpression, data)
                 }

libraries/scripting/jvm-host-test/test/kotlin/script/experimental/jvmhost/test/CapturingTest.kt

@@ -0,0 +1,74 @@
+/*
+ * Copyright 2010-2022 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 kotlin.script.experimental.jvmhost.test
+
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
+import org.junit.Test
+import kotlin.script.experimental.api.*
+import kotlin.test.assertTrue
+
+class CapturingTest {
+
+    @Test
+    fun testScriptWithImplicitReceiverAndSimpleCapturing() {
+        // Reproducing (a bit extended) scenario from KT-53947: without the fix, in the presence of the implicit receiver
+        // of the same type as the receiver in the `apply` function body, the lowering was incorrectly substituting
+        // the correct receiver with the accessor to the implicit one
+        val res = evalString<ScriptWithImplicitReceiver>(
+            """
+                class C {
+                    fun foo() = receiverString + "."
+                }
+                
+                C().foo()
+            """.trimIndent()
+        ) {
+            implicitReceivers(ImplicitReceiverClass("Ok"))
+        }
+
+        assertTrue(
+            res.safeAs<ResultWithDiagnostics.Success<EvaluationResult>>()?.value?.returnValue?.safeAs<ResultValue.Value>()?.value == "Ok.",
+            "test failed:\n  ${res.render()}"
+        )
+    }
+
+    @Test
+    fun testScriptWithImplicitReceiverAndNoCapturing() {
+        // Reproducing (a bit extended) scenario from KT-53947: without the fix, in the presence of the implicit receiver
+        // of the same type as the receiver in the `C2.apply` function body, the lowering was incorrectly substituting
+        // the correct receiver with the accessor to the implicit one
+        // `C1.foo` tests the similar situation with extension receiver.
+        val res = evalString<ScriptWithImplicitReceiver>(
+            """
+                import kotlin.script.experimental.jvmhost.test.ImplicitReceiverClass
+                
+                class C1 {
+                    fun run(receiver: ImplicitReceiverClass): String = receiver.foo()
+                        
+                    fun ImplicitReceiverClass.foo() = "--" + receiverString
+                }
+                
+                class C2 {
+                    fun apply(receiver: ImplicitReceiverClass): String =
+                        "++" + receiver.receiverString
+                }
+                
+                C2().apply(
+                    ImplicitReceiverClass(
+                        C1().run(ImplicitReceiverClass("Ok"))
+                    )
+                )
+            """.trimIndent()
+        ) {
+            implicitReceivers(ImplicitReceiverClass("Not Ok."))
+        }
+
+        assertTrue(
+            res.safeAs<ResultWithDiagnostics.Success<EvaluationResult>>()?.value?.returnValue?.safeAs<ResultValue.Value>()?.value == "++--Ok",
+            "test failed:\n  ${res.render()}"
+        )
+    }
+}

libraries/scripting/jvm-host-test/test/kotlin/script/experimental/jvmhost/test/ConstructorArgumentOrderTest.kt

@@ -6,10 +6,7 @@
 package kotlin.script.experimental.jvmhost.test
 
 import org.junit.Test
-import kotlin.script.experimental.api.EvaluationResult
-import kotlin.script.experimental.api.ResultWithDiagnostics
-import kotlin.script.experimental.api.implicitReceivers
-import kotlin.script.experimental.api.providedProperties
+import kotlin.script.experimental.api.*
 import kotlin.test.assertTrue
 
 class ConstructorArgumentsOrderTest {
@@ -51,6 +48,7 @@ class ConstructorArgumentsOrderTest {
         )
     }
 
-    private fun ResultWithDiagnostics<EvaluationResult>.render() =
-        reports.joinToString("\n  ") { it.message + if (it.exception == null) "" else ": ${it.exception!!.printStackTrace()}" }
 }
+
+internal fun ResultWithDiagnostics<EvaluationResult>.render() =
+    reports.joinToString("\n  ") { it.message + if (it.exception == null) "" else ": ${it.exception!!.printStackTrace()}" }