From 81416d1c464aba3abd664686c23ae9fbb859323d Mon Sep 17 00:00:00 2001
From: "Denis.Zharkov" <denis.zharkov@jetbrains.com>
Date: Fri, 29 Apr 2022 18:15:07 +0300
Subject: [PATCH] Fix codegen for context-receiver contained super class
 constructor calls

^KT-51889 In Progress
---
 .../FirBlackBoxCodegenTestGenerated.java      |  6 ++++++
 .../generators/ArgumentsGenerationUtils.kt    | 13 ++++++-------
 ...uperClassAndSubClassWithContextReceiver.kt | 19 +++++++++++++++++++
 .../IrBlackBoxCodegenTestGenerated.java       |  6 ++++++
 4 files changed, 37 insertions(+), 7 deletions(-)
 create mode 100644 compiler/testData/codegen/box/extensionFunctions/contextReceivers/superClassAndSubClassWithContextReceiver.kt

diff --git a/compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java b/compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java
index c4579a86856..13d30f6670f 100644
--- a/compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java
+++ b/compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java
@@ -16752,6 +16752,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/simpleCall.kt");
             }
 
+            @Test
+            @TestMetadata("superClassAndSubClassWithContextReceiver.kt")
+            public void testSuperClassAndSubClassWithContextReceiver() throws Exception {
+                runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/superClassAndSubClassWithContextReceiver.kt");
+            }
+
             @Test
             @TestMetadata("suspendContextualWithExtension.kt")
             public void testSuspendContextualWithExtension() throws Exception {
diff --git a/compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/ArgumentsGenerationUtils.kt b/compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/ArgumentsGenerationUtils.kt
index bea78b6b28b..1a2d72dd14d 100644
--- a/compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/ArgumentsGenerationUtils.kt
+++ b/compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/ArgumentsGenerationUtils.kt
@@ -34,10 +34,7 @@ import org.jetbrains.kotlin.ir.symbols.impl.IrSimpleFunctionSymbolImpl
 import org.jetbrains.kotlin.ir.symbols.impl.IrValueParameterSymbolImpl
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.psi.KtConstructorDelegationCall
-import org.jetbrains.kotlin.psi.KtElement
-import org.jetbrains.kotlin.psi.KtExpression
-import org.jetbrains.kotlin.psi.ValueArgument
+import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
 import org.jetbrains.kotlin.psi.psiUtil.startOffsetSkippingComments
 import org.jetbrains.kotlin.psi2ir.intermediate.*
@@ -264,10 +261,12 @@ fun StatementGenerator.generateCallReceiver(
         else -> {
             dispatchReceiverValue = generateReceiverOrNull(ktDefaultElement, dispatchReceiver)
             extensionReceiverValue = generateReceiverOrNull(ktDefaultElement, extensionReceiver)
-            contextReceiverValues = if (ktDefaultElement is KtConstructorDelegationCall) contextReceivers.mapNotNull {
-                generateContextReceiverForDelegatingConstructorCall(ktDefaultElement, it as ContextClassReceiver)
+            contextReceiverValues = when (ktDefaultElement) {
+                is KtConstructorDelegationCall, is KtSuperTypeCallEntry -> contextReceivers.mapNotNull {
+                    generateContextReceiverForDelegatingConstructorCall(ktDefaultElement, it as ContextClassReceiver)
+                }
+                else -> contextReceivers.mapNotNull { generateReceiverOrNull(ktDefaultElement, it) }
             }
-            else contextReceivers.mapNotNull { generateReceiverOrNull(ktDefaultElement, it) }
         }
     }
 
diff --git a/compiler/testData/codegen/box/extensionFunctions/contextReceivers/superClassAndSubClassWithContextReceiver.kt b/compiler/testData/codegen/box/extensionFunctions/contextReceivers/superClassAndSubClassWithContextReceiver.kt
new file mode 100644
index 00000000000..a87f92a8236
--- /dev/null
+++ b/compiler/testData/codegen/box/extensionFunctions/contextReceivers/superClassAndSubClassWithContextReceiver.kt
@@ -0,0 +1,19 @@
+// !LANGUAGE: +ContextReceivers
+// TARGET_BACKEND: JVM_IR
+// IGNORE_BACKEND_FIR: JVM_IR
+
+class Components(val x: String)
+
+context(Components)
+abstract class A {
+    fun foo(): String = x
+}
+
+context(Components)
+class B : A()
+
+fun box(): String {
+    return with(Components("OK")) {
+        B().foo()
+    }
+}
diff --git a/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java b/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java
index 477da51916e..156d3a0c9b6 100644
--- a/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java
+++ b/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java
@@ -16752,6 +16752,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/simpleCall.kt");
             }
 
+            @Test
+            @TestMetadata("superClassAndSubClassWithContextReceiver.kt")
+            public void testSuperClassAndSubClassWithContextReceiver() throws Exception {
+                runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/superClassAndSubClassWithContextReceiver.kt");
+            }
+
             @Test
             @TestMetadata("suspendContextualWithExtension.kt")
             public void testSuspendContextualWithExtension() throws Exception {