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
2022-01-31 11:19:26 +03:00
parent 772579143b
commit adb9dfb256
16 changed files with 306 additions and 244 deletions
@@ -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|()]
    }
@@ -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

@@ -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#?.{ $subj$.b# }?.{ ArraySet:[$subj$.c#.get#(IntegerLiteral(0)).plusAssign#(IntegerLiteral(1))] }
-        ArraySet:[a#?.{ $subj$.b# }.c#.get#(IntegerLiteral(0)).plusAssign#(IntegerLiteral(1))]
-        a#?.{ ArraySet:[$subj$.b#.get#(IntegerLiteral(0)).get#(IntegerLiteral(0)).plusAssign#(IntegerLiteral(1))] }
-        a#?.{ $subj$.b# }?.{ ArraySet:[$subj$.c#.get#(IntegerLiteral(0)).get#(IntegerLiteral(0)).plusAssign#(IntegerLiteral(1))] }
-        ArraySet:[a#?.{ $subj$.b# }.c#.get#(IntegerLiteral(0)).get#(IntegerLiteral(0)).plusAssign#(IntegerLiteral(1))]
+        a#?.{ ArraySet:[$subj$.b#.get#(IntegerLiteral(0)) += IntegerLiteral(1)] }
+        a#?.{ $subj$.b# }?.{ ArraySet:[$subj$.c#.get#(IntegerLiteral(0)) += IntegerLiteral(1)] }
+        ArraySet:[a#?.{ $subj$.b# }.c#.get#(IntegerLiteral(0)) += 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)) += IntegerLiteral(1)] }
+        ArraySet:[a#?.{ $subj$.b# }.c#.get#(IntegerLiteral(0)).get#(IntegerLiteral(0)) += IntegerLiteral(1)]
        +=(a#?.{ $subj$.b# }.d#(), IntegerLiteral(1))
    }
@@ -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)
-            setGetBlock =
-                generateSetGetBlockForAugmentedArraySetCall(receiver, baseSource, arrayAccessSource, operation, rhs, convert)
-            this.annotations += annotations
-        }
-    }
-
-    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)
-                }
+            this.lhsGetCall = receiver
+            this.rhs = rhs?.convert() ?: buildErrorExpression(
+                null,
+                ConeSimpleDiagnostic("No value for array set", DiagnosticKind.Syntax)
            )
-            statements += arrayVariable
-            val indexVariables = baseCall.arguments.mapIndexed { i, index ->
-                generateTemporaryVariable(baseModuleData, source = null, specialName = "<index_$i>", initializer = index)
-            }
-            statements += indexVariables
-            statements += buildFunctionCall {
-                source = baseSource?.fakeElement(KtFakeSourceElementKind.DesugaredCompoundAssignment)
-                explicitReceiver = arrayVariable.toQualifiedAccess()
-                calleeReference = buildSimpleNamedReference {
-                    name = OperatorNameConventions.SET
-                }
-                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
-                }
-            }
+            // Second copy of rhs is used because we analyze it twice in different contexts
+            // and now they should be different expressions instances to make everything work properly.
+            // But this lead to exponential time already at FIR building stage,
+            // so we hope this hack will be removed with KT-50861
+            this.rhs2 = rhs?.convert() ?: buildErrorExpression(
+                null,
+                ConeSimpleDiagnostic("No value for array set", DiagnosticKind.Syntax)
+            )
+            this.arrayAccessSource = arrayAccessSource
+            this.annotations += annotations
        }
    }

@@ -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
@@ -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
@@ -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.buildFunctionCall
-import org.jetbrains.kotlin.fir.expressions.builder.buildVariableAssignment
+import org.jetbrains.kotlin.fir.expressions.builder.*
 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@{
-            (annotations as MutableList<FirAnnotation>) += augmentedArraySetCall.annotations
-            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 transformedLhsCall =
+            augmentedArraySetCall.lhsGetCall.transformSingle(transformer, ResolutionMode.ContextIndependent).takeIf { it.isSuccessful() }

-        val firstResult = augmentedArraySetCall.setGetBlock.transformSingle(transformer, ResolutionMode.ContextIndependent)
-        val secondResult = augmentedArraySetCall.assignCall.transformSingle(transformer, ResolutionMode.ContextIndependent)
-
-        fun isSuccessful(functionCall: FirFunctionCall): Boolean =
-            functionCall.typeRef !is FirErrorTypeRef && functionCall.calleeReference is FirResolvedNamedReference
-
-        val firstSucceed = firstCalls.all(::isSuccessful)
-        val secondSucceed = secondCalls.all(::isSuccessful)
+        val blockForGetSetVersion: FirBlock? =
+            transformedLhsCall?.let { augmentedArraySetCall.tryResolveAugmentedArraySetCallAsSetGetBlock(it) }
+        val assignResolvedCall: FirFunctionCall? =
+            transformedLhsCall?.let { augmentedArraySetCall.tryResolveWithOperatorAssignConvention(it) }

        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.exitBlock(augmentedArraySetCall.setGetBlock)
-                firstResult
+                transformer.components.dataFlowAnalyzer.enterBlock(blockForGetSetVersion)
+                transformer.components.dataFlowAnalyzer.exitBlock(blockForGetSetVersion)
+                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)
@@ -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 setGetBlock: FirBlock
+    abstract val lhsGetCall: FirFunctionCall
+    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)

@@ -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 setGetBlock: FirBlock
+    lateinit var lhsGetCall: FirFunctionCall
+    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,
-            setGetBlock,
+            lhsGetCall,
+            rhs,
+            rhs2,
            operation,
            calleeReference,
+            arrayAccessSource,
        )
    }

@@ -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 setGetBlock: FirBlock,
+    override var lhsGetCall: FirFunctionCall,
+    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)
-        setGetBlock.accept(visitor, data)
+        lhsGetCall.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)
-        setGetBlock = setGetBlock.transform(transformer, data)
+        lhsGetCall = lhsGetCall.transform(transformer, data)
+        rhs = rhs.transform(transformer, data)
+        rhs2 = rhs2.transform(transformer, data)
        calleeReference = calleeReference.transform(transformer, data)
        return this
    }
@@ -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,
+    )
@@ -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("]")
    }

@@ -461,10 +461,13 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
        }

        augmentedArraySetCall.configure {
-            +field("assignCall", functionCall)
-            +field("setGetBlock", block)
+            +field("lhsGetCall", functionCall)
+            +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 {
@@ -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
-}
@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !DIAGNOSTICS: -UNUSED_PARAMETER

 fun interface IFoo {
@@ -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()
@@ -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)