FIR DFA: restructure handling of += to avoid losing data flow

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt

@@ -133,9 +133,6 @@ abstract class FirDataFlowAnalyzer(
     private val any = components.session.builtinTypes.anyType.type
     private val nullableNothing = components.session.builtinTypes.nullableNothingType.type
 
-    @PrivateForInline
-    var ignoreFunctionCalls: Boolean = false
-
     // ----------------------------------- Requests -----------------------------------
 
     fun isAccessToUnstableLocalVariable(expression: FirExpression): Boolean {
@@ -162,17 +159,6 @@ abstract class FirDataFlowAnalyzer(
         graphBuilder.dropSubgraphFromCall(call)
     }
 
-    @OptIn(PrivateForInline::class)
-    inline fun <T> withIgnoreFunctionCalls(block: () -> T): T {
-        val oldValue = ignoreFunctionCalls
-        ignoreFunctionCalls = true
-        return try {
-            block()
-        } finally {
-            ignoreFunctionCalls = oldValue
-        }
-    }
-
     // ----------------------------------- Named function -----------------------------------
 
     fun enterFunction(function: FirFunction) {
@@ -841,28 +827,52 @@ abstract class FirDataFlowAnalyzer(
     }
 
     private var functionCallLevel = 0
+    private var resolvingAugmentedAssignmentOptions: Boolean = false
+
+    // The expected sequence of calls for augmented assignment:
+    //  1. enterAugmentedAssignmentCall()
+    //  2. resolve arguments
+    //  3. enterSelectAugmentedAssignmentCall()
+    //  4. resolve all options for calls in context-independent mode
+    //  5. exitSelectAugmentedAssignmentCall()
+    //  6. complete the chosen version
+    //  7. exitFunctionCall(top-level call in the chosen option)
+    fun enterAugmentedAssignmentCall() {
+        graphBuilder.enterCall()
+        // Add an extra space in the local variable assignment analyzer. That way all postponed
+        // lambdas from all alternatives ([get+]plusAssign/[get+]plus[+set]) will be collected
+        // into that space and only removed when `exitFunctionCall` finalizes the chosen option.
+        functionCallLevel++
+    }
+
+    fun enterSelectAugmentedAssignmentCall() {
+        assert(!resolvingAugmentedAssignmentOptions) { "resolving augmented assignment while resolving augmented assignment?" }
+        resolvingAugmentedAssignmentOptions = true
+    }
+
+    fun exitSelectAugmentedAssignmentCall() {
+        assert(resolvingAugmentedAssignmentOptions) { "no enterSelectAugmentedAssignmentCall before exitSelectAugmentedAssignmentCall" }
+        resolvingAugmentedAssignmentOptions = false
+    }
 
-    @OptIn(PrivateForInline::class)
     fun enterFunctionCall(functionCall: FirFunctionCall) {
-        if (!ignoreFunctionCalls) {
-            graphBuilder.enterCall()
-        }
         val lambdaArgs = functionCall.arguments.mapNotNullTo(mutableSetOf()) { it.getAnonymousFunction() }
         val localVariableAssignmentAnalyzer = context.firLocalVariableAssignmentAnalyzer
             ?: if (lambdaArgs.isNotEmpty()) getOrCreateLocalVariableAssignmentAnalyzer(lambdaArgs.first()) else null
         localVariableAssignmentAnalyzer?.enterFunctionCall(lambdaArgs, functionCallLevel)
         functionCallLevel++
+        if (!resolvingAugmentedAssignmentOptions) {
+            graphBuilder.enterCall()
+        }
     }
 
-    @OptIn(PrivateForInline::class)
     fun exitFunctionCall(functionCall: FirFunctionCall, callCompleted: Boolean) {
         functionCallLevel--
         context.firLocalVariableAssignmentAnalyzer?.exitFunctionCall(callCompleted)
-        if (ignoreFunctionCalls) {
+        if (!resolvingAugmentedAssignmentOptions) {
-            return
+            graphBuilder.exitFunctionCall(functionCall, callCompleted).mergeIncomingFlow {
-        }
+                processConditionalContract(it, functionCall)
-        graphBuilder.exitFunctionCall(functionCall, callCompleted).mergeIncomingFlow {
+            }
-            processConditionalContract(it, functionCall)
         }
     }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt

@@ -247,6 +247,7 @@ class ControlFlowGraphBuilder {
         postponedAnonymousFunctionNodes[symbol] = enterNode to exitNode
         addNewSimpleNode(enterNode)
         addNewSimpleNode(exitNode, preferredKind = EdgeKind.DfgForward)
+        require(dataFlowSourcesForNextCompletedCall.isNotEmpty)
         return enterNode to exitNode
     }
 

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

@@ -533,10 +533,12 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
         require(operation != FirOperation.ASSIGN)
 
         assignmentOperatorStatement.transformAnnotations(transformer, ResolutionMode.ContextIndependent)
+        dataFlowAnalyzer.enterAugmentedAssignmentCall()
         val leftArgument = assignmentOperatorStatement.leftArgument.transformSingle(transformer, ResolutionMode.ContextIndependent)
         val rightArgument = assignmentOperatorStatement.rightArgument.transformSingle(transformer, ResolutionMode.ContextDependent)
 
         val generator = GeneratorOfPlusAssignCalls(assignmentOperatorStatement, operation, leftArgument, rightArgument)
+        dataFlowAnalyzer.enterSelectAugmentedAssignmentCall()
 
         // x.plusAssign(y)
         val assignOperatorCall = generator.createAssignOperatorCall()
@@ -554,6 +556,8 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
         val operatorCallReference = resolvedOperatorCall.calleeReference as? FirNamedReferenceWithCandidate
         val operatorIsSuccessful = operatorCallReference?.isError == false
 
+        dataFlowAnalyzer.exitSelectAugmentedAssignmentCall()
+
         fun operatorReturnTypeMatches(candidate: Candidate): Boolean {
             // After KT-45503, non-assign flavor of operator is checked more strictly: the return type must be assignable to the variable.
             val operatorCallReturnType = resolvedOperatorCall.typeRef.coneType
@@ -574,14 +578,12 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
         val lhsIsVar = lhsVariable?.isVar == true
 
         fun chooseAssign(): FirStatement {
-            dataFlowAnalyzer.enterFunctionCall(resolvedAssignCall)
             callCompleter.completeCall(resolvedAssignCall, noExpectedType)
             dataFlowAnalyzer.exitFunctionCall(resolvedAssignCall, callCompleted = true)
             return resolvedAssignCall
         }
 
         fun chooseOperator(): FirStatement {
-            dataFlowAnalyzer.enterFunctionCall(resolvedAssignCall)
             callCompleter.completeCall(
                 resolvedOperatorCall,
                 lhsVariable?.returnTypeRef ?: noExpectedType,
@@ -669,11 +671,9 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
     }
 
     private inline fun <T> resolveCandidateForAssignmentOperatorCall(block: () -> T): T {
-        return dataFlowAnalyzer.withIgnoreFunctionCalls {
+        return callResolver.withNoArgumentsTransform {
-            callResolver.withNoArgumentsTransform {
+            context.withInferenceSession(InferenceSessionForAssignmentOperatorCall) {
-                context.withInferenceSession(InferenceSessionForAssignmentOperatorCall) {
+                block()
-                    block()
-                }
             }
         }
     }
@@ -1230,11 +1230,14 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
         augmentedArraySetCall.transformAnnotations(transformer, data)
 
         // transformedLhsCall: a.get(index)
+        dataFlowAnalyzer.enterAugmentedAssignmentCall()
         val transformedLhsCall = augmentedArraySetCall.lhsGetCall.transformSingle(transformer, ResolutionMode.ContextIndependent)
         val transformedRhs = augmentedArraySetCall.rhs.transformSingle(transformer, ResolutionMode.ContextDependent)
 
         val generator = GeneratorOfPlusAssignCalls(augmentedArraySetCall, operation, transformedLhsCall, transformedRhs)
 
+        dataFlowAnalyzer.enterSelectAugmentedAssignmentCall()
+
         // a.get(b).plusAssign(c)
         val assignOperatorCall = generator.createAssignOperatorCall()
         val resolvedAssignCall = resolveCandidateForAssignmentOperatorCall {
@@ -1243,24 +1246,23 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
         val assignCallReference = resolvedAssignCall.calleeReference as? FirNamedReferenceWithCandidate
         val assignIsSuccessful = assignCallReference?.isError == false
 
-        fun completeAssignCall() {
+        fun chooseAssign(): FirFunctionCall {
-            dataFlowAnalyzer.enterFunctionCall(resolvedAssignCall)
             callCompleter.completeCall(resolvedAssignCall, noExpectedType)
             dataFlowAnalyzer.exitFunctionCall(resolvedAssignCall, callCompleted = true)
-        }
-
-        fun chooseAssign(): FirStatement {
-            completeAssignCall()
             return resolvedAssignCall
         }
 
         // prefer a "simpler" variant for dynamics
         if (transformedLhsCall.calleeReference.resolvedSymbol?.origin == FirDeclarationOrigin.DynamicScope) {
+            dataFlowAnalyzer.exitSelectAugmentedAssignmentCall()
             return chooseAssign()
         }
 
         // <array>.set(<index_i>, <array>.get(<index_i>).plus(c))
         val info = tryResolveAugmentedArraySetCallAsSetGetBlock(augmentedArraySetCall, transformedLhsCall, transformedRhs)
+
+        dataFlowAnalyzer.exitSelectAugmentedAssignmentCall()
+
         val resolvedOperatorCall = info.operatorCall
         val operatorCallReference = resolvedOperatorCall.calleeReference as? FirNamedReferenceWithCandidate
         val operatorIsSuccessful = operatorCallReference?.isError == false
@@ -1276,27 +1278,23 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
         val setIsSuccessful = setCallReference?.isError == false
 
         fun chooseSetOperator(): FirStatement {
-            dataFlowAnalyzer.enterFunctionCall(resolvedSetCall)
-            dataFlowAnalyzer.enterFunctionCall(resolvedOperatorCall)
             callCompleter.completeCall(
                 resolvedSetCall,
                 noExpectedType,
                 expectedTypeMismatchIsReportedInChecker = true
             )
-            dataFlowAnalyzer.exitFunctionCall(resolvedOperatorCall, callCompleted = true)
             dataFlowAnalyzer.exitFunctionCall(resolvedSetCall, callCompleted = true)
             return info.toBlock()
         }
 
         fun reportError(diagnostic: ConeDiagnostic): FirStatement {
-            completeAssignCall()
+            return chooseAssign().also {
-            val errorReference = buildErrorNamedReference {
+                val errorReference = buildErrorNamedReference {
-                source = augmentedArraySetCall.source
+                    source = augmentedArraySetCall.source
-                this.diagnostic = diagnostic
+                    this.diagnostic = diagnostic
+                }
+                it.replaceCalleeReference(errorReference)
             }
-            resolvedAssignCall.replaceCalleeReference(errorReference)
-
-            return resolvedAssignCall
         }
 
         fun reportAmbiguity(