[PSI2IR] Generate safe calls correctly

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

@@ -16247,6 +16247,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/plusMatrix.kt");
             }
 
+            @Test
+            @TestMetadata("propertyCompoundAssignment.kt")
+            public void testPropertyCompoundAssignment() throws Exception {
+                runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/propertyCompoundAssignment.kt");
+            }
+
             @Test
             @TestMetadata("simpleCall.kt")
             public void testSimpleCall() throws Exception {

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/intermediate/SafeCallReceiver.kt

@@ -31,7 +31,6 @@ class SafeCallReceiver(
     val startOffset: Int,
     val endOffset: Int,
     val extensionReceiver: IntermediateValue?,
-    // TODO: Use context receivers
     val contextReceivers: List<IntermediateValue>,
     val dispatchReceiver: IntermediateValue?,
     val isStatement: Boolean
@@ -51,7 +50,7 @@ class SafeCallReceiver(
             extensionReceiverValue = null
         }
 
-        val irResult = withDispatchAndExtensionAndContextReceivers(dispatchReceiverValue, extensionReceiverValue, emptyList())
+        val irResult = withDispatchAndExtensionAndContextReceivers(dispatchReceiverValue, extensionReceiverValue, contextReceivers)
 
         val resultType = if (isStatement) generator.context.irBuiltIns.unitType else irResult.type.makeNullable()
 

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

@@ -0,0 +1,33 @@
+// !LANGUAGE: +ContextReceivers
+// TARGET_BACKEND: JVM_IR
+// IGNORE_BACKEND_FIR: JVM_IR
+
+class LoggingCounter {
+    var operationCounter = 0
+}
+
+class A {
+    context(LoggingCounter)
+    var p: Int
+        get(): Int {
+            operationCounter++
+            return 1
+        }
+        set(value: Int) {
+            operationCounter++
+        }
+}
+
+fun foo() = A()
+
+fun box(): String {
+    val loggingCounter = LoggingCounter()
+    with(loggingCounter) {
+        foo().p += 1
+        foo().p = 1
+        foo()?.p = 1
+        foo().p
+    }
+    val operationsTotal = loggingCounter.operationCounter
+    return if (operationsTotal == 5) "OK" else "$operationsTotal"
+}

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

@@ -16247,6 +16247,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/plusMatrix.kt");
             }
 
+            @Test
+            @TestMetadata("propertyCompoundAssignment.kt")
+            public void testPropertyCompoundAssignment() throws Exception {
+                runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/propertyCompoundAssignment.kt");
+            }
+
             @Test
             @TestMetadata("simpleCall.kt")
             public void testSimpleCall() throws Exception {

core/descriptors/src/org/jetbrains/kotlin/descriptors/impl/PropertyDescriptorImpl.java

@@ -23,6 +23,7 @@ import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.descriptors.*;
 import org.jetbrains.kotlin.descriptors.annotations.Annotations;
 import org.jetbrains.kotlin.name.Name;
+import org.jetbrains.kotlin.resolve.scopes.receivers.ContextReceiver;
 import org.jetbrains.kotlin.resolve.scopes.receivers.ExtensionReceiver;
 import org.jetbrains.kotlin.types.*;
 import org.jetbrains.kotlin.utils.SmartSet;
@@ -444,7 +445,7 @@ public class PropertyDescriptorImpl extends VariableDescriptorWithInitializerImp
 
         List<ReceiverParameterDescriptor> substitutedContextReceivers = new ArrayList<ReceiverParameterDescriptor>();
         for (ReceiverParameterDescriptor contextReceiverParameter: contextReceiverParameters) {
-            ReceiverParameterDescriptor substitutedContextReceiver = substituteParameterDescriptor(substitutor, substitutedDescriptor,
+            ReceiverParameterDescriptor substitutedContextReceiver = substituteContextParameterDescriptor(substitutor, substitutedDescriptor,
                                                                                                       contextReceiverParameter);
             if (substitutedContextReceiver != null) {
                 substitutedContextReceivers.add(substitutedContextReceiver);
@@ -541,6 +542,20 @@ public class PropertyDescriptorImpl extends VariableDescriptorWithInitializerImp
         );
     }
 
+    private static ReceiverParameterDescriptor substituteContextParameterDescriptor(
+            TypeSubstitutor substitutor,
+            PropertyDescriptor substitutedPropertyDescriptor,
+            ReceiverParameterDescriptor receiverParameterDescriptor
+    ) {
+        KotlinType substitutedType = substitutor.substitute(receiverParameterDescriptor.getType(), Variance.IN_VARIANCE);
+        if (substitutedType == null) return null;
+        return new ReceiverParameterDescriptorImpl(
+                substitutedPropertyDescriptor,
+                new ContextReceiver(substitutedPropertyDescriptor, substitutedType, receiverParameterDescriptor.getValue()),
+                receiverParameterDescriptor.getAnnotations()
+        );
+    }
+
     private static FunctionDescriptor getSubstitutedInitialSignatureDescriptor(
             @NotNull TypeSubstitutor substitutor,
             @NotNull PropertyAccessorDescriptor accessorDescriptor

native/native.tests/tests-gen/org/jetbrains/kotlin/konan/blackboxtest/NativeExtBlackBoxTestGenerated.java

@@ -16434,6 +16434,12 @@ public class NativeExtBlackBoxTestGenerated extends AbstractNativeBlackBoxTest {
                     runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/plusMatrix.kt");
                 }
 
+                @Test
+                @TestMetadata("propertyCompoundAssignment.kt")
+                public void testPropertyCompoundAssignment() throws Exception {
+                    runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/propertyCompoundAssignment.kt");
+                }
+
                 @Test
                 @TestMetadata("simpleCall.kt")
                 public void testSimpleCall() throws Exception {