FIR: Rework processing AugmentedArraySetCall

Previously (few commits earlier), it contained two versions
of receiver (lhs) generated separately for each desugaring version
that looked a bit redundant.

Now, at FIR building stage we just don't create desugaring sub-trees,
instead they are being built during bodies transformation and that seems
to be much convenient there, since we don't need to reverse-engineer
get-set-operator version to check if containing calls are successful
(as we just built those calls and retain them)

Semantically, this changes may only change how data flow works
for such statements (see changed compatibilityResolveWithVarargAndOperatorCall.kt)

^KT-50861 Relates

compiler/fir/analysis-tests/testData/resolve/arrays/arraySetWithOperation.fir.txt

@@ -56,8 +56,8 @@ FILE: arraySetWithOperation.kt
 
     }
     public final fun test_3(a: R|A<D>|): R|kotlin/Unit| {
-        ArraySet:[R|<local>/a|.R|SubstitutionOverride</A.get: R|D|>|(Int(0)).R|/D.plusAssign|(R|/D.D|())]
+        ArraySet:[R|<local>/a|.R|SubstitutionOverride</A.get: R|D|>|(Int(0)) += R|/D.D|()]
     }
     public final fun test_4(b: R|B|): R|kotlin/Unit| {
-        ArraySet:[R|<local>/b|.<Unresolved name: get>#(Int(0)).<Unresolved name: plusAssign>#(R|/B.B|())]
+        ArraySet:[R|<local>/b|.<Unresolved name: get>#(Int(0)) += R|/B.B|()]
     }

compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/fir/DestructuringDeclaration.kt

@@ -7,9 +7,9 @@ package org.jetbrains.kotlin.fir.lightTree.fir
 
 import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.fir.FirModuleData
-import org.jetbrains.kotlin.fir.builder.generateTemporaryVariable
 import org.jetbrains.kotlin.fir.declarations.FirVariable
 import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.generateTemporaryVariable
 import org.jetbrains.kotlin.fir.lightTree.converter.generateDestructuringBlock
 import org.jetbrains.kotlin.fir.lightTree.fir.modifier.Modifier
 

compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/safeCallsWithAugmentedAssignment.txt

@@ -3,11 +3,11 @@ FILE: safeCallsWithAugmentedAssignment.kt
         a#?.{ +=($subj$.b#, IntegerLiteral(1)) }
         a#?.{ $subj$.b# }?.{ +=($subj$.c#, IntegerLiteral(1)) }
         +=(a#?.{ $subj$.b# }.c#, IntegerLiteral(1))
-        a#?.{ ArraySet:[$subj$.b#.get#(IntegerLiteral(0)).plusAssign#(IntegerLiteral(1))] }
+        a#?.{ ArraySet:[$subj$.b#.get#(IntegerLiteral(0)) += IntegerLiteral(1)] }
-        a#?.{ $subj$.b# }?.{ ArraySet:[$subj$.c#.get#(IntegerLiteral(0)).plusAssign#(IntegerLiteral(1))] }
+        a#?.{ $subj$.b# }?.{ ArraySet:[$subj$.c#.get#(IntegerLiteral(0)) += IntegerLiteral(1)] }
-        ArraySet:[a#?.{ $subj$.b# }.c#.get#(IntegerLiteral(0)).plusAssign#(IntegerLiteral(1))]
+        ArraySet:[a#?.{ $subj$.b# }.c#.get#(IntegerLiteral(0)) += IntegerLiteral(1)]
-        a#?.{ ArraySet:[$subj$.b#.get#(IntegerLiteral(0)).get#(IntegerLiteral(0)).plusAssign#(IntegerLiteral(1))] }
+        a#?.{ ArraySet:[$subj$.b#.get#(IntegerLiteral(0)).get#(IntegerLiteral(0)) += IntegerLiteral(1)] }
-        a#?.{ $subj$.b# }?.{ ArraySet:[$subj$.c#.get#(IntegerLiteral(0)).get#(IntegerLiteral(0)).plusAssign#(IntegerLiteral(1))] }
+        a#?.{ $subj$.b# }?.{ ArraySet:[$subj$.c#.get#(IntegerLiteral(0)).get#(IntegerLiteral(0)) += IntegerLiteral(1)] }
-        ArraySet:[a#?.{ $subj$.b# }.c#.get#(IntegerLiteral(0)).get#(IntegerLiteral(0)).plusAssign#(IntegerLiteral(1))]
+        ArraySet:[a#?.{ $subj$.b# }.c#.get#(IntegerLiteral(0)).get#(IntegerLiteral(0)) += IntegerLiteral(1)]
         +=(a#?.{ $subj$.b# }.d#(), IntegerLiteral(1))
     }

compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/BaseFirBuilder.kt

@@ -1024,123 +1024,28 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
         rhs: T?,
         convert: T.() -> FirExpression
     ): FirStatement {
+        require(receiver is FirFunctionCall) {
+            "Array access should be desugared to a function call, but $receiver is found"
+        }
+
         return buildAugmentedArraySetCall {
             source = baseSource
             this.operation = operation
-            assignCall = generateAugmentedCallForAugmentedArraySetCall(receiver, baseSource, operation, rhs, convert)
+            this.lhsGetCall = receiver
-            setGetBlock =
+            this.rhs = rhs?.convert() ?: buildErrorExpression(
-                generateSetGetBlockForAugmentedArraySetCall(receiver, baseSource, arrayAccessSource, operation, rhs, convert)
+                null,
-            this.annotations += annotations
+                ConeSimpleDiagnostic("No value for array set", DiagnosticKind.Syntax)
-        }
-    }
-
-    private fun generateAugmentedCallForAugmentedArraySetCall(
-        receiver: FirExpression, // a.get(x,y)
-        baseSource: KtSourceElement?,
-        operation: FirOperation,
-        rhs: T?,
-        convert: T.() -> FirExpression
-    ): FirFunctionCall {
-        /*
-         * Desugarings of a[x, y] += z to
-         * a.get(x, y).plusAssign(z)
-         */
-        return buildFunctionCall {
-            source = baseSource?.fakeElement(KtFakeSourceElementKind.DesugaredCompoundAssignment)
-            calleeReference = buildSimpleNamedReference {
-                name = FirOperationNameConventions.ASSIGNMENTS.getValue(operation)
-            }
-            explicitReceiver = receiver
-            argumentList = buildArgumentList {
-                arguments += rhs?.convert() ?: buildErrorExpression(
-                    null,
-                    ConeSimpleDiagnostic("No value for array set", DiagnosticKind.Syntax)
-                )
-            }
-            origin = FirFunctionCallOrigin.Operator
-        }
-    }
-
-
-    private fun generateSetGetBlockForAugmentedArraySetCall(
-        receiver: FirExpression,
-        baseSource: KtSourceElement?,
-        arrayAccessSource: KtSourceElement?,
-        operation: FirOperation,
-        rhs: T?,
-        convert: T.() -> FirExpression
-    ): FirBlock {
-        /*
-         * Desugarings of a[x, y] += z to
-         * {
-         *     val tmp_a = a
-         *     val tmp_x = x
-         *     val tmp_y = y
-         *     tmp_a.set(tmp_x, tmp_a.get(tmp_x, tmp_y).plus(z))
-         * }
-         */
-        return buildBlock {
-            val baseCall = receiver as FirFunctionCall
-
-            val arrayVariable = generateTemporaryVariable(
-                baseModuleData,
-                source = null,
-                specialName = "<array>",
-                initializer = baseCall.explicitReceiver ?: buildErrorExpression {
-                    source = baseSource?.fakeElement(KtFakeSourceElementKind.DesugaredCompoundAssignment)
-                    diagnostic = ConeSimpleDiagnostic("No receiver for array access", DiagnosticKind.Syntax)
-                }
             )
-            statements += arrayVariable
+            // Second copy of rhs is used because we analyze it twice in different contexts
-            val indexVariables = baseCall.arguments.mapIndexed { i, index ->
+            // and now they should be different expressions instances to make everything work properly.
-                generateTemporaryVariable(baseModuleData, source = null, specialName = "<index_$i>", initializer = index)
+            // But this lead to exponential time already at FIR building stage,
-            }
+            // so we hope this hack will be removed with KT-50861
-            statements += indexVariables
+            this.rhs2 = rhs?.convert() ?: buildErrorExpression(
-            statements += buildFunctionCall {
+                null,
-                source = baseSource?.fakeElement(KtFakeSourceElementKind.DesugaredCompoundAssignment)
+                ConeSimpleDiagnostic("No value for array set", DiagnosticKind.Syntax)
-                explicitReceiver = arrayVariable.toQualifiedAccess()
+            )
-                calleeReference = buildSimpleNamedReference {
+            this.arrayAccessSource = arrayAccessSource
-                    name = OperatorNameConventions.SET
+            this.annotations += annotations
-                }
-                origin = FirFunctionCallOrigin.Operator
-                argumentList = buildArgumentList {
-                    for (indexVariable in indexVariables) {
-                        arguments += indexVariable.toQualifiedAccess()
-                    }
-
-                    val getCall = buildFunctionCall {
-                        source = arrayAccessSource?.fakeElement(KtFakeSourceElementKind.DesugaredCompoundAssignment)
-                        explicitReceiver = arrayVariable.toQualifiedAccess()
-                        calleeReference = buildSimpleNamedReference {
-                            name = OperatorNameConventions.GET
-                        }
-                        argumentList = buildArgumentList {
-                            for (indexVariable in indexVariables) {
-                                arguments += indexVariable.toQualifiedAccess()
-                            }
-                        }
-                        origin = FirFunctionCallOrigin.Operator
-                    }
-
-                    val operatorCall = buildFunctionCall {
-                        calleeReference = buildSimpleNamedReference {
-                            name = FirOperationNameConventions.ASSIGNMENTS_TO_SIMPLE_OPERATOR.getValue(operation)
-                        }
-                        explicitReceiver = getCall
-                        argumentList = buildArgumentList {
-                            arguments += rhs?.convert() ?: buildErrorExpression(
-                                null,
-                                ConeSimpleDiagnostic(
-                                    "No value for array set",
-                                    DiagnosticKind.Syntax
-                                )
-                            )
-                        }
-                        origin = FirFunctionCallOrigin.Operator
-                    }
-                    arguments += operatorCall
-                }
-            }
         }
     }
 
@@ -1262,14 +1167,6 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
         initContainingClassAttr(context)
     }
 
-    private fun FirVariable.toQualifiedAccess(): FirQualifiedAccessExpression = buildPropertyAccessExpression {
-        calleeReference = buildResolvedNamedReference {
-            source = this@toQualifiedAccess.source?.fakeElement(KtFakeSourceElementKind.ReferenceInAtomicQualifiedAccess)
-            name = this@toQualifiedAccess.name
-            resolvedSymbol = this@toQualifiedAccess.symbol
-        }
-    }
-
     protected inline fun <R> withDefaultSourceElementKind(newDefault: KtSourceElementKind, action: () -> R): R {
         val currentForced = context.forcedElementSourceKind
         context.forcedElementSourceKind = newDefault

compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/ConversionUtils.kt

@@ -285,50 +285,6 @@ fun generateResolvedAccessExpression(source: KtSourceElement?, variable: FirVari
         }
     }
 
-fun generateTemporaryVariable(
-    moduleData: FirModuleData,
-    source: KtSourceElement?,
-    name: Name,
-    initializer: FirExpression,
-    typeRef: FirTypeRef? = null,
-    extractedAnnotations: Collection<FirAnnotation>? = null,
-): FirVariable =
-    buildProperty {
-        this.source = source
-        this.moduleData = moduleData
-        origin = FirDeclarationOrigin.Source
-        returnTypeRef = typeRef ?: buildImplicitTypeRef {
-            this.source = source
-        }
-        this.name = name
-        this.initializer = initializer
-        symbol = FirPropertySymbol(name)
-        isVar = false
-        isLocal = true
-        status = FirDeclarationStatusImpl(Visibilities.Local, Modality.FINAL)
-        if (extractedAnnotations != null) {
-            // LT extracts annotations ahead.
-            // PSI extracts annotations on demand. Use a similar util in [PsiConversionUtils]
-            annotations.addAll(extractedAnnotations)
-        }
-    }
-
-fun generateTemporaryVariable(
-    moduleData: FirModuleData,
-    source: KtSourceElement?,
-    specialName: String,
-    initializer: FirExpression,
-    extractedAnnotations: Collection<FirAnnotation>? = null,
-): FirVariable =
-    generateTemporaryVariable(
-        moduleData,
-        source,
-        Name.special("<$specialName>"),
-        initializer,
-        null,
-        extractedAnnotations,
-    )
-
 val FirClassBuilder.ownerRegularOrAnonymousObjectSymbol
     get() = when (this) {
         is FirAnonymousObjectBuilder -> symbol

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt

@@ -13,9 +13,7 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.isLocal
 import org.jetbrains.kotlin.fir.diagnostics.*
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.expressions.builder.buildErrorExpression
+import org.jetbrains.kotlin.fir.expressions.builder.*
-import org.jetbrains.kotlin.fir.expressions.builder.buildFunctionCall
-import org.jetbrains.kotlin.fir.expressions.builder.buildVariableAssignment
 import org.jetbrains.kotlin.fir.expressions.impl.FirResolvedArgumentList
 import org.jetbrains.kotlin.fir.expressions.impl.toAnnotationArgumentMapping
 import org.jetbrains.kotlin.fir.references.*
@@ -46,6 +44,7 @@ import org.jetbrains.kotlin.resolve.calls.tower.CandidateApplicability
 import org.jetbrains.kotlin.types.AbstractTypeChecker
 import org.jetbrains.kotlin.types.ConstantValueKind
 import org.jetbrains.kotlin.types.TypeApproximatorConfiguration
+import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) : FirPartialBodyResolveTransformer(transformer) {
@@ -1043,32 +1042,16 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
         augmentedArraySetCall.transformAnnotations(transformer, data)
         val operatorName = FirOperationNameConventions.ASSIGNMENTS.getValue(augmentedArraySetCall.operation)
 
-        val firstCalls = with(augmentedArraySetCall.setGetBlock.statements.last() as FirFunctionCall) setCall@{
+        val transformedLhsCall =
-            (annotations as MutableList<FirAnnotation>) += augmentedArraySetCall.annotations
+            augmentedArraySetCall.lhsGetCall.transformSingle(transformer, ResolutionMode.ContextIndependent).takeIf { it.isSuccessful() }
-            buildList {
-                add(this@setCall)
-                with(arguments.last() as FirFunctionCall) plusCall@{
-                    add(this@plusCall)
-                    add(explicitReceiver as FirFunctionCall)
-                }
-            }
-        }
-        val secondCalls = listOf(
-            augmentedArraySetCall.assignCall,
-            augmentedArraySetCall.assignCall.explicitReceiver as FirFunctionCall
-        )
 
-        val firstResult = augmentedArraySetCall.setGetBlock.transformSingle(transformer, ResolutionMode.ContextIndependent)
+        val blockForGetSetVersion: FirBlock? =
-        val secondResult = augmentedArraySetCall.assignCall.transformSingle(transformer, ResolutionMode.ContextIndependent)
+            transformedLhsCall?.let { augmentedArraySetCall.tryResolveAugmentedArraySetCallAsSetGetBlock(it) }
-
+        val assignResolvedCall: FirFunctionCall? =
-        fun isSuccessful(functionCall: FirFunctionCall): Boolean =
+            transformedLhsCall?.let { augmentedArraySetCall.tryResolveWithOperatorAssignConvention(it) }
-            functionCall.typeRef !is FirErrorTypeRef && functionCall.calleeReference is FirResolvedNamedReference
-
-        val firstSucceed = firstCalls.all(::isSuccessful)
-        val secondSucceed = secondCalls.all(::isSuccessful)
 
         val result: FirStatement = when {
-            firstSucceed && secondSucceed -> {
+            assignResolvedCall != null && blockForGetSetVersion != null -> {
                 augmentedArraySetCall.also {
                     it.replaceCalleeReference(
                         buildErrorNamedReference {
@@ -1079,16 +1062,19 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
                     )
                 }
             }
-            firstSucceed -> {
+            blockForGetSetVersion != null -> {
                 //checking secondResult leave erroneous nodes in dfa graph,
                 //we add another block so final type of expression will be correct
                 //todo replace this hack with proper graph cleaning
-                transformer.components.dataFlowAnalyzer.enterBlock(augmentedArraySetCall.setGetBlock)
+                transformer.components.dataFlowAnalyzer.enterBlock(blockForGetSetVersion)
-                transformer.components.dataFlowAnalyzer.exitBlock(augmentedArraySetCall.setGetBlock)
+                transformer.components.dataFlowAnalyzer.exitBlock(blockForGetSetVersion)
-                firstResult
+                blockForGetSetVersion
             }
-            secondSucceed -> secondResult
+            assignResolvedCall != null -> assignResolvedCall
             else -> {
+                augmentedArraySetCall.rhs.transformSingle(transformer, ResolutionMode.ContextIndependent)
+                augmentedArraySetCall.rhs2.transformSingle(transformer, ResolutionMode.ContextIndependent)
+
                 augmentedArraySetCall.also {
                     it.replaceCalleeReference(
                         buildErrorNamedReference {
@@ -1102,6 +1088,149 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
         return result
     }
 
+    private fun FirFunctionCall.isSuccessful(): Boolean =
+        typeRef !is FirErrorTypeRef && calleeReference is FirResolvedNamedReference
+
+
+    /**
+     * Desugarings of a[x, y] += z to
+     * a.get(x, y).plusAssign(z)
+     *
+     * @return null if `plusAssign` is unresolved
+     * @return block defined as described above, otherwise
+     */
+    private fun FirAugmentedArraySetCall.tryResolveWithOperatorAssignConvention(lhsGetCall: FirFunctionCall): FirFunctionCall? {
+        val assignCall = buildFunctionCall {
+            source = this@tryResolveWithOperatorAssignConvention.source?.fakeElement(KtFakeSourceElementKind.DesugaredCompoundAssignment)
+            calleeReference = buildSimpleNamedReference {
+                name = FirOperationNameConventions.ASSIGNMENTS.getValue(operation)
+            }
+            explicitReceiver = lhsGetCall
+            argumentList = buildArgumentList {
+                arguments += rhs2
+            }
+            origin = FirFunctionCallOrigin.Operator
+            annotations += this@tryResolveWithOperatorAssignConvention.annotations
+        }
+
+        val transformedCall = assignCall.transformSingle(transformer, ResolutionMode.ContextIndependent)
+        return transformedCall.takeIf { it.isSuccessful() }
+    }
+
+    /**
+     * Desugarings of a[x, y] += z to
+     * {
+     *     val tmp_a = a
+     *     val tmp_x = x
+     *     val tmp_y = y
+     *     tmp_a.set(tmp_x, tmp_y, tmp_a.get(tmp_x, tmp_y).plus(z))
+     * }
+     *
+     * @return null if `set` or `plus` calls are unresolved
+     * @return block defined as described above, otherwise
+     */
+    private fun FirAugmentedArraySetCall.tryResolveAugmentedArraySetCallAsSetGetBlock(lhsGetCall: FirFunctionCall): FirBlock? {
+        val arrayVariable = generateTemporaryVariable(
+            session.moduleData,
+            source = null,
+            specialName = "<array>",
+            initializer = lhsGetCall.explicitReceiver ?: buildErrorExpression {
+                source =
+                    this@tryResolveAugmentedArraySetCallAsSetGetBlock.source
+                        ?.fakeElement(KtFakeSourceElementKind.DesugaredCompoundAssignment)
+                diagnostic = ConeSimpleDiagnostic("No receiver for array access", DiagnosticKind.Syntax)
+            }
+        )
+
+        val indexVariables = lhsGetCall.arguments.flatMap {
+            if (it is FirVarargArgumentsExpression)
+                it.arguments
+            else
+                listOf(it)
+        }.mapIndexed { i, index ->
+            generateTemporaryVariable(session.moduleData, source = null, specialName = "<index_$i>", initializer = index)
+        }
+
+        val getCall = buildFunctionCall {
+            source = arrayAccessSource?.fakeElement(KtFakeSourceElementKind.DesugaredCompoundAssignment)
+            explicitReceiver = arrayVariable.toQualifiedAccess()
+            calleeReference = buildSimpleNamedReference {
+                name = OperatorNameConventions.GET
+            }
+            argumentList = buildArgumentList {
+                for (indexVariable in indexVariables) {
+                    arguments += indexVariable.toQualifiedAccess()
+                }
+            }
+            origin = FirFunctionCallOrigin.Operator
+        }
+
+        val operatorCall = buildFunctionCall {
+            calleeReference = buildSimpleNamedReference {
+                name = FirOperationNameConventions.ASSIGNMENTS_TO_SIMPLE_OPERATOR.getValue(operation)
+            }
+            explicitReceiver = getCall
+            argumentList = buildArgumentList {
+                arguments += rhs
+            }
+            origin = FirFunctionCallOrigin.Operator
+        }
+
+        val setCall = buildFunctionCall {
+            source =
+                this@tryResolveAugmentedArraySetCallAsSetGetBlock.source
+                    ?.fakeElement(KtFakeSourceElementKind.DesugaredCompoundAssignment)
+            explicitReceiver = arrayVariable.toQualifiedAccess()
+            calleeReference = buildSimpleNamedReference {
+                name = OperatorNameConventions.SET
+            }
+            origin = FirFunctionCallOrigin.Operator
+            argumentList = buildArgumentList {
+                for (indexVariable in indexVariables) {
+                    arguments += indexVariable.toQualifiedAccess()
+                }
+
+                arguments += operatorCall
+            }
+
+            annotations += this@tryResolveAugmentedArraySetCallAsSetGetBlock.annotations
+        }
+
+        arrayVariable.transformSingle(transformer, ResolutionMode.ContextIndependent)
+        indexVariables.forEach { it.transformSingle(transformer, ResolutionMode.ContextIndependent) }
+        val transformedSet = setCall.transformSingle(transformer, ResolutionMode.ContextIndependent)
+
+        if (!transformedSet.isSuccessful()) return null
+
+        val transformedOperator =
+            transformedSet.argumentList.arguments.last().let {
+                if (it is FirNamedArgumentExpression)
+                    it.expression
+                else
+                    it
+            }
+
+        require(transformedOperator is FirFunctionCall) {
+            "Last argument of set call should be an operator but $transformedOperator found"
+        }
+
+        if (!transformedOperator.isSuccessful()) return null
+
+        return buildBlock {
+            statements += arrayVariable
+            statements += indexVariables
+            statements += setCall
+        }.also {
+            it.replaceTypeRef(
+                buildResolvedTypeRef {
+                    source = this@tryResolveAugmentedArraySetCallAsSetGetBlock.source
+                    type = session.builtinTypes.unitType.type
+                }
+            )
+        }
+    }
+
+
     override fun transformArrayOfCall(arrayOfCall: FirArrayOfCall, data: ResolutionMode): FirStatement {
         if (data is ResolutionMode.ContextDependent) {
             arrayOfCall.transformChildren(transformer, data)

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/FirAugmentedArraySetCall.kt

@@ -19,10 +19,12 @@ import org.jetbrains.kotlin.fir.visitors.*
 abstract class FirAugmentedArraySetCall : FirPureAbstractElement(), FirStatement {
     abstract override val source: KtSourceElement?
     abstract override val annotations: List<FirAnnotation>
-    abstract val assignCall: FirFunctionCall
+    abstract val lhsGetCall: FirFunctionCall
-    abstract val setGetBlock: FirBlock
+    abstract val rhs: FirExpression
+    abstract val rhs2: FirExpression
     abstract val operation: FirOperation
     abstract val calleeReference: FirReference
+    abstract val arrayAccessSource: KtSourceElement?
 
     override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R = visitor.visitAugmentedArraySetCall(this, data)
 

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/builder/FirAugmentedArraySetCallBuilder.kt

@@ -11,7 +11,7 @@ import org.jetbrains.kotlin.fir.builder.FirAnnotationContainerBuilder
 import org.jetbrains.kotlin.fir.builder.FirBuilderDsl
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirAugmentedArraySetCall
-import org.jetbrains.kotlin.fir.expressions.FirBlock
+import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.expressions.FirOperation
 import org.jetbrains.kotlin.fir.expressions.impl.FirAugmentedArraySetCallImpl
@@ -28,19 +28,23 @@ import org.jetbrains.kotlin.fir.visitors.*
 class FirAugmentedArraySetCallBuilder : FirAnnotationContainerBuilder {
     override var source: KtSourceElement? = null
     override val annotations: MutableList<FirAnnotation> = mutableListOf()
-    lateinit var assignCall: FirFunctionCall
+    lateinit var lhsGetCall: FirFunctionCall
-    lateinit var setGetBlock: FirBlock
+    lateinit var rhs: FirExpression
+    lateinit var rhs2: FirExpression
     lateinit var operation: FirOperation
     var calleeReference: FirReference = FirStubReference
+    var arrayAccessSource: KtSourceElement? = null
 
     override fun build(): FirAugmentedArraySetCall {
         return FirAugmentedArraySetCallImpl(
             source,
             annotations,
-            assignCall,
+            lhsGetCall,
-            setGetBlock,
+            rhs,
+            rhs2,
             operation,
             calleeReference,
+            arrayAccessSource,
         )
     }
 

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/impl/FirAugmentedArraySetCallImpl.kt

@@ -8,7 +8,7 @@ package org.jetbrains.kotlin.fir.expressions.impl
 import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirAugmentedArraySetCall
-import org.jetbrains.kotlin.fir.expressions.FirBlock
+import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.expressions.FirOperation
 import org.jetbrains.kotlin.fir.references.FirReference
@@ -22,22 +22,26 @@ import org.jetbrains.kotlin.fir.visitors.*
 internal class FirAugmentedArraySetCallImpl(
     override val source: KtSourceElement?,
     override val annotations: MutableList<FirAnnotation>,
-    override var assignCall: FirFunctionCall,
+    override var lhsGetCall: FirFunctionCall,
-    override var setGetBlock: FirBlock,
+    override var rhs: FirExpression,
+    override var rhs2: FirExpression,
     override val operation: FirOperation,
     override var calleeReference: FirReference,
+    override val arrayAccessSource: KtSourceElement?,
 ) : FirAugmentedArraySetCall() {
     override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
         annotations.forEach { it.accept(visitor, data) }
-        assignCall.accept(visitor, data)
+        lhsGetCall.accept(visitor, data)
-        setGetBlock.accept(visitor, data)
+        rhs.accept(visitor, data)
+        rhs2.accept(visitor, data)
         calleeReference.accept(visitor, data)
     }
 
     override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirAugmentedArraySetCallImpl {
         transformAnnotations(transformer, data)
-        assignCall = assignCall.transform(transformer, data)
+        lhsGetCall = lhsGetCall.transform(transformer, data)
-        setGetBlock = setGetBlock.transform(transformer, data)
+        rhs = rhs.transform(transformer, data)
+        rhs2 = rhs2.transform(transformer, data)
         calleeReference = calleeReference.transform(transformer, data)
         return this
     }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirGeneration.kt

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir
+
+import org.jetbrains.kotlin.KtFakeSourceElementKind
+import org.jetbrains.kotlin.KtSourceElement
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.fakeElement
+import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
+import org.jetbrains.kotlin.fir.declarations.FirVariable
+import org.jetbrains.kotlin.fir.declarations.builder.buildProperty
+import org.jetbrains.kotlin.fir.declarations.impl.FirDeclarationStatusImpl
+import org.jetbrains.kotlin.fir.expressions.FirAnnotation
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
+import org.jetbrains.kotlin.fir.expressions.builder.buildPropertyAccessExpression
+import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.fir.types.builder.buildImplicitTypeRef
+import org.jetbrains.kotlin.name.Name
+
+fun FirVariable.toQualifiedAccess(): FirQualifiedAccessExpression = buildPropertyAccessExpression {
+    calleeReference = buildResolvedNamedReference {
+        source = this@toQualifiedAccess.source?.fakeElement(KtFakeSourceElementKind.ReferenceInAtomicQualifiedAccess)
+        name = this@toQualifiedAccess.name
+        resolvedSymbol = this@toQualifiedAccess.symbol
+    }
+}
+
+fun generateTemporaryVariable(
+    moduleData: FirModuleData,
+    source: KtSourceElement?,
+    name: Name,
+    initializer: FirExpression,
+    typeRef: FirTypeRef? = null,
+    extractedAnnotations: Collection<FirAnnotation>? = null,
+): FirVariable =
+    buildProperty {
+        this.source = source
+        this.moduleData = moduleData
+        origin = FirDeclarationOrigin.Source
+        returnTypeRef = typeRef ?: buildImplicitTypeRef {
+            this.source = source
+        }
+        this.name = name
+        this.initializer = initializer
+        symbol = FirPropertySymbol(name)
+        isVar = false
+        isLocal = true
+        status = FirDeclarationStatusImpl(Visibilities.Local, Modality.FINAL)
+        if (extractedAnnotations != null) {
+            // LT extracts annotations ahead.
+            // PSI extracts annotations on demand. Use a similar util in [PsiConversionUtils]
+            annotations.addAll(extractedAnnotations)
+        }
+    }
+
+fun generateTemporaryVariable(
+    moduleData: FirModuleData,
+    source: KtSourceElement?,
+    specialName: String,
+    initializer: FirExpression,
+    extractedAnnotations: Collection<FirAnnotation>? = null,
+): FirVariable =
+    generateTemporaryVariable(
+        moduleData,
+        source,
+        Name.special("<$specialName>"),
+        initializer,
+        null,
+        extractedAnnotations,
+    )

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirRenderer.kt

@@ -20,7 +20,10 @@ import org.jetbrains.kotlin.fir.expressions.impl.*
 import org.jetbrains.kotlin.fir.references.*
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.*
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.name.Name
@@ -1302,7 +1305,11 @@ open class FirRenderer(builder: StringBuilder, protected val mode: RenderMode =
     override fun visitAugmentedArraySetCall(augmentedArraySetCall: FirAugmentedArraySetCall) {
         augmentedArraySetCall.annotations.renderAnnotations()
         print("ArraySet:[")
-        augmentedArraySetCall.assignCall.accept(this)
+        augmentedArraySetCall.lhsGetCall.accept(this)
+        print(" ")
+        print(augmentedArraySetCall.operation.operator)
+        print(" ")
+        augmentedArraySetCall.rhs.accept(this)
         print("]")
     }
 

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/NodeConfigurator.kt

@@ -461,10 +461,13 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
         }
 
         augmentedArraySetCall.configure {
-            +field("assignCall", functionCall)
+            +field("lhsGetCall", functionCall)
-            +field("setGetBlock", block)
+            +field("rhs", expression)
+            +field("rhs2", expression)
             +field("operation", operationType)
+            // Used for resolution errors reporting in case
             +field("calleeReference", reference, withReplace = true)
+            +field("arrayAccessSource", sourceElementType, nullable = true)
         }
 
         classReferenceExpression.configure {

compiler/testData/diagnostics/tests/callableReference/compatibilityResolveWithVarargAndOperatorCall.fir.kt

@@ -1,18 +0,0 @@
-// !DIAGNOSTICS: -UNUSED_PARAMETER
-
-fun interface IFoo {
-    fun foo(i: Int)
-}
-
-fun interface IFoo2 : IFoo
-
-object A
-
-operator fun A.get(i: IFoo) = 1
-operator fun A.set(i: IFoo, newValue: Int) {}
-
-fun withVararg(vararg xs: Int) = 42
-
-fun test1() {
-    A[::withVararg] <!UNRESOLVED_REFERENCE!>+=<!> 1
-}

compiler/testData/diagnostics/tests/callableReference/compatibilityResolveWithVarargAndOperatorCall.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !DIAGNOSTICS: -UNUSED_PARAMETER
 
 fun interface IFoo {

compiler/testData/diagnostics/tests/controlFlowAnalysis/deadCode/deadCodeInArrayAccess.fir.kt

@@ -34,7 +34,7 @@ fun testArrayAssignment4(n: Nothing) {
 
 fun testArrayPlusAssign(array: Array<Any>) {
     operator fun Any.plusAssign(a: Any) {}
-    array<!UNREACHABLE_CODE!>[<!>1<!UNREACHABLE_CODE!>] +=<!> todo()
+    array[1] <!UNREACHABLE_CODE!>+=<!> todo()
 }
 
 fun todo(): Nothing = throw Exception()

compiler/testData/diagnostics/tests/nullableTypes/safeCallOperators.fir.kt

@@ -64,7 +64,7 @@ fun foo(a: A?) {
         // 1. All kinds of green code with safe+call + invoke we identified fails with CCE if `a != null`, anyway
         // 2. In case of null value, the behavior is intended (no call performed)
         <!SAFE_CALL_WILL_CHANGE_NULLABILITY!>a<!UNNECESSARY_SAFE_CALL!>?.<!>q()<!>
-        <!SAFE_CALL_WILL_CHANGE_NULLABILITY, SAFE_CALL_WILL_CHANGE_NULLABILITY!>a<!UNNECESSARY_SAFE_CALL, UNNECESSARY_SAFE_CALL!>?.<!>w<!>++
+        <!SAFE_CALL_WILL_CHANGE_NULLABILITY!>a<!UNNECESSARY_SAFE_CALL!>?.<!>w<!>++
 
         (<!SAFE_CALL_WILL_CHANGE_NULLABILITY!>a<!UNNECESSARY_SAFE_CALL!>?.<!>l<!>) += 1
         (<!SAFE_CALL_WILL_CHANGE_NULLABILITY!>a<!UNNECESSARY_SAFE_CALL!>?.<!>l<!>)[0]
@@ -75,7 +75,7 @@ fun foo(a: A?) {
         (<!SAFE_CALL_WILL_CHANGE_NULLABILITY!>a<!UNNECESSARY_SAFE_CALL!>?.<!>ll<!>)[0][0]++
         (<!SAFE_CALL_WILL_CHANGE_NULLABILITY!>a<!UNNECESSARY_SAFE_CALL!>?.<!>ll<!>)[0][0] = 1
         (<!SAFE_CALL_WILL_CHANGE_NULLABILITY!>a<!UNNECESSARY_SAFE_CALL!>?.<!>q<!>)()
-        (<!SAFE_CALL_WILL_CHANGE_NULLABILITY, SAFE_CALL_WILL_CHANGE_NULLABILITY!>a<!UNNECESSARY_SAFE_CALL, UNNECESSARY_SAFE_CALL!>?.<!>w<!>)++
+        (<!SAFE_CALL_WILL_CHANGE_NULLABILITY!>a<!UNNECESSARY_SAFE_CALL!>?.<!>w<!>)++
 
         <!SAFE_CALL_WILL_CHANGE_NULLABILITY!>a<!UNNECESSARY_SAFE_CALL!>?.<!>l<!>.plusAssign(1)
         <!SAFE_CALL_WILL_CHANGE_NULLABILITY!>a<!UNNECESSARY_SAFE_CALL!>?.<!>l<!>.get(0)