[PSI2IR] Pass context receiver values in PropertyLValue

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

@@ -16169,6 +16169,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/inferGenericPropertyType.kt");
             }
 
+            @Test
+            @TestMetadata("plusAssign.kt")
+            public void testPlusAssign() throws Exception {
+                runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/plusAssign.kt");
+            }
+
             @Test
             @TestMetadata("plusMatrix.kt")
             public void testPlusMatrix() throws Exception {

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

@@ -40,7 +40,7 @@ abstract class PropertyLValueBase(
     val superQualifier: IrClassSymbol?
 ) : LValue, AssignmentReceiver {
     override fun assign(withLValue: (LValue) -> IrExpression) =
-        callReceiver.call { dispatchReceiverValue, extensionReceiverValue, _ ->
+        callReceiver.call { dispatchReceiverValue, extensionReceiverValue, contextReceiverValues ->
             val dispatchReceiverVariable2 = dispatchReceiverValue?.let {
                 scope.createTemporaryVariable(dispatchReceiverValue.load(), "this")
             }
@@ -51,11 +51,17 @@ abstract class PropertyLValueBase(
             }
             val extensionReceiverValue2 = extensionReceiverVariable2?.let { VariableLValue(context, it) }
 
-            val irResultExpression = withLValue(withReceiver(dispatchReceiverValue2, extensionReceiverValue2))
+            val contextReceiverVariables2 = contextReceiverValues.mapIndexed { i, value ->
+                scope.createTemporaryVariable(value.load(), "additionalReceiver$i")
+            }
+            val contextReceiversValues2 = contextReceiverVariables2.map { VariableLValue(context, it) }
+
+            val irResultExpression = withLValue(withReceiver(dispatchReceiverValue2, extensionReceiverValue2, contextReceiversValues2))
 
             val irBlock = IrBlockImpl(startOffset, endOffset, irResultExpression.type, origin)
             irBlock.addIfNotNull(dispatchReceiverVariable2)
             irBlock.addIfNotNull(extensionReceiverVariable2)
+            contextReceiverVariables2.forEach { irBlock.addIfNotNull(it) }
             irBlock.statements.add(irResultExpression)
             irBlock
         }
@@ -63,7 +69,7 @@ abstract class PropertyLValueBase(
     override fun assign(value: IrExpression): IrExpression =
         store(value)
 
-    protected abstract fun withReceiver(dispatchReceiver: VariableLValue?, extensionReceiver: VariableLValue?): PropertyLValueBase
+    protected abstract fun withReceiver(dispatchReceiver: VariableLValue?, extensionReceiver: VariableLValue?, contextReceivers: List<VariableLValue>): PropertyLValueBase
 }
 
 class FieldPropertyLValue(
@@ -81,9 +87,9 @@ class FieldPropertyLValue(
     PropertyLValueBase(context, scope, startOffset, endOffset, origin, type, callReceiver, superQualifier) {
 
     override fun load(): IrExpression =
-        // TODO: Use context receiver values
+        callReceiver.call { dispatchReceiverValue, extensionReceiverValue, contextReceiverValues ->
-        callReceiver.call { dispatchReceiverValue, extensionReceiverValue, _ ->
             assert(extensionReceiverValue == null) { "Field can't have an extension receiver: ${field.descriptor}" }
+            assert(contextReceiverValues.isEmpty()) { "Field can't have context receivers: ${field.descriptor}" }
             IrGetFieldImpl(
                 startOffset, endOffset,
                 field,
@@ -95,9 +101,9 @@ class FieldPropertyLValue(
         }
 
     override fun store(irExpression: IrExpression) =
-        // TODO: Use context receiver values
+        callReceiver.call { dispatchReceiverValue, extensionReceiverValue, contextReceiverValues ->
-        callReceiver.call { dispatchReceiverValue, extensionReceiverValue, _ ->
             assert(extensionReceiverValue == null) { "Field can't have an extension receiver: ${field.descriptor}" }
+            assert(contextReceiverValues.isEmpty()) { "Field can't have context receivers: ${field.descriptor}" }
             IrSetFieldImpl(
                 startOffset, endOffset,
                 field,
@@ -109,14 +115,14 @@ class FieldPropertyLValue(
             ).also { context.callToSubstitutedDescriptorMap[it] = descriptor }
         }
 
-    override fun withReceiver(dispatchReceiver: VariableLValue?, extensionReceiver: VariableLValue?): PropertyLValueBase =
+    override fun withReceiver(dispatchReceiver: VariableLValue?, extensionReceiver: VariableLValue?, contextReceivers: List<VariableLValue>): PropertyLValueBase =
         FieldPropertyLValue(
             context,
             scope, startOffset, endOffset, origin,
             field,
             descriptor,
             type,
-            SimpleCallReceiver(dispatchReceiver, extensionReceiver, emptyList()),
+            SimpleCallReceiver(dispatchReceiver, extensionReceiver, contextReceivers),
             superQualifier
         )
 }
@@ -186,13 +192,13 @@ class AccessorPropertyLValue(
             }
         }
 
-    override fun withReceiver(dispatchReceiver: VariableLValue?, extensionReceiver: VariableLValue?): PropertyLValueBase =
+    override fun withReceiver(dispatchReceiver: VariableLValue?, extensionReceiver: VariableLValue?, contextReceivers: List<VariableLValue>): PropertyLValueBase =
         AccessorPropertyLValue(
             context, scope,
             startOffset, endOffset, origin,
             type, getter, getterDescriptor, setter, setterDescriptor,
             typeArguments,
-            SimpleCallReceiver(dispatchReceiver, extensionReceiver, emptyList()),
+            SimpleCallReceiver(dispatchReceiver, extensionReceiver, contextReceivers),
             superQualifier
         )
 }

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

@@ -0,0 +1,72 @@
+// TARGET_BACKEND: JVM_IR
+// IGNORE_BACKEND_FIR: JVM_IR
+// WITH_RUNTIME
+
+@file:Suppress("RESERVED_VAR_PROPERTY_OF_VALUE_CLASS")
+
+open class EntityFactory<E>(val size: Int, val factory: (Int) -> E)
+
+class EntityContext {
+    var d = DoubleArray(16)
+        private set
+    var size: Int = 0
+        private set
+    fun <E> create(entity: EntityFactory<E>): E {
+        val i = size
+        size += entity.size
+        if (size > d.size) d = d.copyOf(maxOf(2 * d.size, size))
+        return entity.factory(i)
+    }
+}
+
+@JvmInline value class EDouble(private val i: Int) {
+    companion object Factory : EntityFactory<EDouble>(1, ::EDouble)
+
+    context(EntityContext)
+    var value: Double
+        get() = d[i]
+        set(value) { d[i] = value }
+}
+
+@JvmInline value class EVec3(private val i: Int) {
+    companion object Factory : EntityFactory<EVec3>(3, ::EVec3)
+
+    context(EntityContext)
+    var x: Double
+        get() = d[i]
+        set(value) { d[i] = value }
+
+    context(EntityContext)
+    var y: Double
+        get() = d[i + 1]
+        set(value) { d[i + 1] = value }
+
+    context(EntityContext)
+    var z: Double
+        get() = d[i + 2]
+        set(value) { d[i + 2] = value }
+}
+
+context(EntityContext)
+fun EVec3.str(): String =
+    "[$x, $y, $z]"
+
+context(EntityContext)
+operator fun EVec3.plusAssign(v: EVec3) {
+    x += v.x
+    y += v.y
+    z += v.z
+}
+
+fun box(): String = with(EntityContext()) {
+    val v0 = create(EVec3)
+    v0.x = 1.0
+    v0.y = 2.0
+    v0.z = 3.0
+    val v1 = create(EVec3)
+    v1.x = 2.0
+    v1.y = 0.0
+    v1.z = 4.0
+    v1 += v0
+    if (v1.x == 3.0 && v1.y == 2.0 && v1.z == 7.0) "OK" else "fail"
+}

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

@@ -16169,6 +16169,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/inferGenericPropertyType.kt");
             }
 
+            @Test
+            @TestMetadata("plusAssign.kt")
+            public void testPlusAssign() throws Exception {
+                runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/plusAssign.kt");
+            }
+
             @Test
             @TestMetadata("plusMatrix.kt")
             public void testPlusMatrix() throws Exception {