KT-28456 generate index arguments per expression

In the desugaring for compound assignment to a collection element,
argument expression 'i' is mapped to value parameters 'iG' and 'iS' of
corresponding 'get' and 'set' operators.
In general, these value parameters can have different indices.

This requires extra machinery in argument generation - that is, to be
able to generate a particular expression argument using an arbitrary
callback. In the vast majority of the cases this callback will just use
the corresponding StatementGenerator to generate IR subtree for the
provided expression. In case of 'get' and 'set' operator calls for an
augmented assignment expression this will map corresponding argument
expressions to pregenerated temporary variables.

Thus, in the following context:
```
  class A

  operator fun A.get(vararg xs: Int) = 0
  operator fun A.set(i: Int, j: Int, v: Int) {}
```

statement `a[1, 2] += 3` will be desugared as (in a really pseudo
Kotlin):
```
  {
    val tmp_array = a
    val tmp_index0 = 1
    val tmp_index1 = 2
    tmp_array.set(
      i = tmp_index0,
      j = tmp_index1,
      v = tmp_array.get(xs = [tmp_index0, tmp_index1]).plus(3)
    )
  }
```

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

@@ -26,6 +26,7 @@ import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrExpressionWithCopy
 import org.jetbrains.kotlin.ir.expressions.impl.*
 import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
 import org.jetbrains.kotlin.psi.psiUtil.startOffsetSkippingComments
 import org.jetbrains.kotlin.psi2ir.intermediate.*
@@ -38,7 +39,6 @@ import org.jetbrains.kotlin.resolve.calls.model.*
 import org.jetbrains.kotlin.resolve.scopes.receivers.*
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeSubstitutor
-import java.lang.AssertionError
 
 fun StatementGenerator.generateReceiverOrNull(ktDefaultElement: KtElement, receiver: ReceiverValue?): IntermediateValue? =
     receiver?.let { generateReceiver(ktDefaultElement, receiver) }
@@ -209,9 +209,10 @@ private fun StatementGenerator.generateReceiverForCalleeImportedFromObject(
     }
 }
 
-fun StatementGenerator.generateVarargExpression(
+fun StatementGenerator.generateVarargExpressionUsing(
     varargArgument: VarargValueArgument,
-    valueParameter: ValueParameterDescriptor
+    valueParameter: ValueParameterDescriptor,
+    generateArgumentExpression: (KtExpression) -> IrExpression?
 ): IrExpression? {
     if (varargArgument.arguments.isEmpty()) {
         return null
@@ -232,14 +233,16 @@ fun StatementGenerator.generateVarargExpression(
     for (argument in varargArgument.arguments) {
         val ktArgumentExpression = argument.getArgumentExpression()
             ?: throw AssertionError("No argument expression for vararg element ${argument.asElement().text}")
+        val irArgumentExpression = generateArgumentExpression(ktArgumentExpression)
+            ?: throw AssertionError("'generateArgumentExpression' should return non-null for vararg element ${ktArgumentExpression.text}")
         val irVarargElement =
             if (argument.getSpreadElement() != null)
                 IrSpreadElementImpl(
                     ktArgumentExpression.startOffsetSkippingComments, ktArgumentExpression.endOffset,
-                    generateExpression(ktArgumentExpression)
+                    irArgumentExpression
                 )
             else
-                generateExpression(ktArgumentExpression)
+                irArgumentExpression
 
         irVararg.addElement(irVarargElement)
     }
@@ -247,17 +250,23 @@ fun StatementGenerator.generateVarargExpression(
     return irVararg
 }
 
-fun StatementGenerator.generateValueArgument(
+private fun StatementGenerator.generateValueArgument(
     valueArgument: ResolvedValueArgument,
     valueParameter: ValueParameterDescriptor
+) = generateValueArgumentUsing(valueArgument, valueParameter) { generateExpression(it) }
+
+fun StatementGenerator.generateValueArgumentUsing(
+    valueArgument: ResolvedValueArgument,
+    valueParameter: ValueParameterDescriptor,
+    generateArgumentExpression: (KtExpression) -> IrExpression?
 ): IrExpression? =
     when (valueArgument) {
         is DefaultValueArgument ->
             null
         is ExpressionValueArgument ->
-            generateExpression(valueArgument.valueArgument!!.getArgumentExpression()!!)
+            generateArgumentExpression(valueArgument.valueArgument!!.getArgumentExpression()!!)
         is VarargValueArgument ->
-            generateVarargExpression(valueArgument, valueParameter)
+            generateVarargExpressionUsing(valueArgument, valueParameter, generateArgumentExpression)
         else ->
             TODO("Unexpected valueArgument: ${valueArgument::class.java.simpleName}")
     }
@@ -357,9 +366,19 @@ private fun ResolvedCall<*>.isExtensionInvokeCall(): Boolean {
 }
 
 private fun StatementGenerator.pregenerateValueArguments(call: CallBuilder, resolvedCall: ResolvedCall<*>) {
+    pregenerateValueArgumentsUsing(call, resolvedCall) {
+        generateExpression(it)
+    }
+}
+
+fun StatementGenerator.pregenerateValueArgumentsUsing(
+    call: CallBuilder,
+    resolvedCall: ResolvedCall<*>,
+    generateArgumentExpression: (KtExpression) -> IrExpression?
+) {
     resolvedCall.valueArgumentsByIndex!!.forEachIndexed { index, valueArgument ->
         val valueParameter = call.descriptor.valueParameters[index]
-        call.irValueArgumentsByIndex[index] = generateValueArgument(valueArgument, valueParameter)
+        call.irValueArgumentsByIndex[index] = generateValueArgumentUsing(valueArgument, valueParameter, generateArgumentExpression)
     }
 }
 

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

@@ -268,6 +268,7 @@ class AssignmentGenerator(statementGenerator: StatementGenerator) : StatementGen
 
         return ArrayAccessAssignmentReceiver(
             ktLeft.arrayExpression!!.genExpr(),
+            ktLeft.indexExpressions,
             ktLeft.indexExpressions.map { it.genExpr() },
             indexedGetResolvedCall,
             indexedSetResolvedCall,

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

@@ -22,13 +22,16 @@ import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
 import org.jetbrains.kotlin.ir.expressions.impl.IrBlockImpl
 import org.jetbrains.kotlin.ir.expressions.impl.inlineStatement
 import org.jetbrains.kotlin.ir.expressions.isAssignmentOperatorWithResult
+import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi2ir.generators.CallGenerator
+import org.jetbrains.kotlin.psi2ir.generators.pregenerateValueArgumentsUsing
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.types.KotlinType
 
 class ArrayAccessAssignmentReceiver(
     private val irArray: IrExpression,
-    private val irIndices: List<IrExpression>,
+    private val ktIndexExpressions: List<KtExpression>,
+    private val irIndexExpressions: List<IrExpression>,
     private val indexedGetResolvedCall: ResolvedCall<FunctionDescriptor>?,
     private val indexedSetResolvedCall: ResolvedCall<FunctionDescriptor>?,
     private val indexedGetCall: () -> CallBuilder?,
@@ -57,15 +60,18 @@ class ArrayAccessAssignmentReceiver(
 
         val irArrayValue = callGenerator.scope.createTemporaryVariableInBlock(callGenerator.context, irArray, irBlock, "array")
 
-        val irIndexValues = irIndices.mapIndexed { i, irIndex ->
-            callGenerator.scope.createTemporaryVariableInBlock(callGenerator.context, irIndex, irBlock, "index$i")
+        val ktExpressionToIrIndexValue = HashMap<KtExpression, IntermediateValue>()
+
+        for ((i, irIndex) in irIndexExpressions.withIndex()) {
+            ktExpressionToIrIndexValue[ktIndexExpressions[i]] =
+                    callGenerator.scope.createTemporaryVariableInBlock(callGenerator.context, irIndex, irBlock, "index$i")
         }
 
         val irLValue = LValueWithGetterAndSetterCalls(
             callGenerator,
             descriptor,
-            { indexedGetCall()?.fillArrayAndIndexArguments(irArrayValue, irIndexValues) },
-            { indexedSetCall()?.fillArrayAndIndexArguments(irArrayValue, irIndexValues) },
+            { indexedGetCall()?.fillArrayAndIndexArguments(irArrayValue, indexedGetResolvedCall!!, ktExpressionToIrIndexValue) },
+            { indexedSetCall()?.fillArrayAndIndexArguments(irArrayValue, indexedSetResolvedCall!!, ktExpressionToIrIndexValue) },
             callGenerator.translateType(kotlinType),
             startOffset, endOffset, origin
         )
@@ -76,18 +82,24 @@ class ArrayAccessAssignmentReceiver(
 
     override fun assign(value: IrExpression): IrExpression {
         val call = indexedSetCall() ?: throw AssertionError("Array access without indexed-get call")
+        val ktExpressionToIrIndexExpression = ktIndexExpressions.zip(irIndexExpressions).toMap()
         call.setExplicitReceiverValue(OnceExpressionValue(irArray))
-        irIndices.forEachIndexed { i, irIndex ->
-            call.irValueArgumentsByIndex[i] = irIndex
+        callGenerator.statementGenerator.pregenerateValueArgumentsUsing(call, indexedSetResolvedCall!!) {
+            ktExpressionToIrIndexExpression[it]
         }
         call.lastArgument = value
         return callGenerator.generateCall(startOffset, endOffset, call, IrStatementOrigin.EQ)
     }
 
-    private fun CallBuilder.fillArrayAndIndexArguments(arrayValue: IntermediateValue, indexValues: List<IntermediateValue>) = apply {
+    private fun CallBuilder.fillArrayAndIndexArguments(
+        arrayValue: IntermediateValue,
+        resolvedCall: ResolvedCall<FunctionDescriptor>,
+        ktExpressionToIrIndexValue: Map<KtExpression, IntermediateValue>
+    ) = apply {
         setExplicitReceiverValue(arrayValue)
-        indexValues.forEachIndexed { i, irIndexValue ->
-            irValueArgumentsByIndex[i] = irIndexValue.load()
+
+        callGenerator.statementGenerator.pregenerateValueArgumentsUsing(this, resolvedCall) { ktExpression ->
+            ktExpressionToIrIndexValue[ktExpression]?.load()
         }
     }
 }