KT-51290 Account for context receivers when generating delegate function bodies

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java

@@ -17480,6 +17480,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/kt51284.kt");
             }
 
+            @Test
+            @TestMetadata("kt51290.kt")
+            public void testKt51290() throws Exception {
+                runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/kt51290.kt");
+            }
+
             @Test
             @TestMetadata("kt51475.kt")
             public void testKt51475() throws Exception {

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/ClassGenerator.kt

@@ -28,15 +28,11 @@ import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrGetFieldImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrReturnImpl
-import org.jetbrains.kotlin.ir.expressions.mapValueParameters
 import org.jetbrains.kotlin.ir.expressions.putTypeArguments
 import org.jetbrains.kotlin.ir.expressions.typeParametersCount
 import org.jetbrains.kotlin.ir.symbols.impl.IrFieldSymbolImpl
 import org.jetbrains.kotlin.ir.types.IrSimpleType
-import org.jetbrains.kotlin.ir.util.SYNTHETIC_OFFSET
-import org.jetbrains.kotlin.ir.util.createIrClassFromDescriptor
-import org.jetbrains.kotlin.ir.util.declareSimpleFunctionWithOverrides
-import org.jetbrains.kotlin.ir.util.properties
+import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtClassOrObject
 import org.jetbrains.kotlin.psi.KtDelegatedSuperTypeEntry
@@ -366,10 +362,9 @@ class ClassGenerator(
                     IrGetValueImpl(startOffset, endOffset, extensionReceiver.type, extensionReceiver.symbol)
                 }
 
-            mapValueParameters { overriddenValueParameter ->
-                val delegatedValueParameter = delegatedDescriptor.valueParameters[overriddenValueParameter.index]
-                val irDelegatedValueParameter = irDelegatedFunction.getIrValueParameter(delegatedValueParameter)
-                IrGetValueImpl(startOffset, endOffset, irDelegatedValueParameter.type, irDelegatedValueParameter.symbol)
+            (delegatedDescriptor.contextReceiverParameters + delegatedDescriptor.valueParameters).forEachIndexed { index, param ->
+                val delegatedParameter = irDelegatedFunction.getIrValueParameter(param, index)
+                putValueArgument(index, IrGetValueImpl(startOffset, endOffset, delegatedParameter.type, delegatedParameter.symbol))
             }
         }
 

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/IrDeclarations.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.ir.declarations
 
 import org.jetbrains.kotlin.descriptors.InlineClassRepresentation
 import org.jetbrains.kotlin.descriptors.MultiFieldValueClassRepresentation
+import org.jetbrains.kotlin.descriptors.ParameterDescriptor
 import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
 import org.jetbrains.kotlin.ir.ObsoleteDescriptorBasedAPI
 import org.jetbrains.kotlin.ir.expressions.IrExpressionBody
@@ -38,7 +39,11 @@ val IrFile.name: String get() = File(path).name
 
 @ObsoleteDescriptorBasedAPI
 fun IrFunction.getIrValueParameter(parameter: ValueParameterDescriptor): IrValueParameter =
-    valueParameters.getOrElse(parameter.index) {
+    getIrValueParameter(parameter, parameter.index)
+
+@ObsoleteDescriptorBasedAPI
+fun IrFunction.getIrValueParameter(parameter: ParameterDescriptor, index: Int): IrValueParameter =
+    valueParameters.getOrElse(index) {
         throw AssertionError("No IrValueParameter for $parameter")
     }.also { found ->
         assert(found.descriptor == parameter) {

compiler/testData/codegen/box/extensionFunctions/contextReceivers/kt51290.kt

@@ -0,0 +1,40 @@
+// !LANGUAGE: +ContextReceivers
+// TARGET_BACKEND: JVM_IR
+
+data class Language(var name: String)
+
+interface LoggingContext {
+    fun log(level: Int, message: String)
+}
+
+interface SaveRepository<T> {
+    context(LoggingContext)
+    fun save(content: T)
+}
+
+context(LoggingContext, SaveRepository<Language>)
+fun startBusinessOperation() {
+    log(0, "Operation has started")
+    save(Language("Kotlin"))
+}
+
+class CompositeContext(c1: LoggingContext, c2: SaveRepository<Language>): LoggingContext by c1, SaveRepository<Language> by c2
+
+fun box(): String {
+    val loggingCtx = object : LoggingContext {
+        override fun log(level: Int, message: String) {}
+    }
+
+    val saveCtx = object : SaveRepository<Language> {
+        context(LoggingContext)
+        override fun save(content: Language) {
+            log(message = "Saving $content", level = 123)
+        }
+    }
+
+    with(CompositeContext(loggingCtx, saveCtx)) {
+        startBusinessOperation()
+    }
+
+    return "OK"
+}

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java

@@ -17480,6 +17480,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/kt51284.kt");
             }
 
+            @Test
+            @TestMetadata("kt51290.kt")
+            public void testKt51290() throws Exception {
+                runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/kt51290.kt");
+            }
+
             @Test
             @TestMetadata("kt51475.kt")
             public void testKt51475() throws Exception {