[FIR] Complete data-flow analysis for all control-flow graph nodes

There are conditions where the data-flow analysis for a control-flow
graph node is delayed. Make sure that when completing a graph, all nodes
within the graph have completed their data-flow analysis.

^KT-61794 Fixed

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosticCompilerTestFE10TestdataTestGenerated.java

@@ -6984,6 +6984,12 @@ public class DiagnosticCompilerTestFE10TestdataTestGenerated extends AbstractDia
                 runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/reassignmentInCatch.kt");
             }
 
+            @Test
+            @TestMetadata("reassignmentInLambda.kt")
+            public void testReassignmentInLambda() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/reassignmentInLambda.kt");
+            }
+
             @Test
             @TestMetadata("reassignmentInTryCatch.kt")
             public void testReassignmentInTryCatch() throws Exception {

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/LLFirPreresolvedReversedDiagnosticCompilerFE10TestDataTestGenerated.java

@@ -6984,6 +6984,12 @@ public class LLFirPreresolvedReversedDiagnosticCompilerFE10TestDataTestGenerated
                 runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/reassignmentInCatch.kt");
             }
 
+            @Test
+            @TestMetadata("reassignmentInLambda.kt")
+            public void testReassignmentInLambda() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/reassignmentInLambda.kt");
+            }
+
             @Test
             @TestMetadata("reassignmentInTryCatch.kt")
             public void testReassignmentInTryCatch() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirLightTreeOldFrontendDiagnosticsTestGenerated.java

@@ -6984,6 +6984,12 @@ public class FirLightTreeOldFrontendDiagnosticsTestGenerated extends AbstractFir
                 runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/reassignmentInCatch.kt");
             }
 
+            @Test
+            @TestMetadata("reassignmentInLambda.kt")
+            public void testReassignmentInLambda() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/reassignmentInLambda.kt");
+            }
+
             @Test
             @TestMetadata("reassignmentInTryCatch.kt")
             public void testReassignmentInTryCatch() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirPsiOldFrontendDiagnosticsTestGenerated.java

@@ -6990,6 +6990,12 @@ public class FirPsiOldFrontendDiagnosticsTestGenerated extends AbstractFirPsiDia
                 runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/reassignmentInCatch.kt");
             }
 
+            @Test
+            @TestMetadata("reassignmentInLambda.kt")
+            public void testReassignmentInLambda() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/reassignmentInLambda.kt");
+            }
+
             @Test
             @TestMetadata("reassignmentInTryCatch.kt")
             public void testReassignmentInTryCatch() throws Exception {

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

@@ -24,7 +24,6 @@ import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutorByMap
 import org.jetbrains.kotlin.fir.resolve.substitution.chain
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirAbstractBodyResolveTransformer
-import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
 import org.jetbrains.kotlin.fir.resolve.transformers.unwrapAnonymousFunctionExpression
 import org.jetbrains.kotlin.fir.scopes.getFunctions
 import org.jetbrains.kotlin.fir.scopes.impl.toConeType
@@ -182,6 +181,7 @@ abstract class FirDataFlowAnalyzer(
 
         val (node, graph) = graphBuilder.exitFunction(function)
         node.mergeIncomingFlow()
+        graph.completePostponedNodes()
         if (!graphBuilder.isTopLevel) {
             for (valueParameter in function.valueParameters) {
                 variableStorage.removeRealVariable(valueParameter.symbol)
@@ -211,6 +211,7 @@ abstract class FirDataFlowAnalyzer(
         } else {
             resetReceivers()
         }
+        graph?.completePostponedNodes()
         return graph
     }
 
@@ -231,6 +232,7 @@ abstract class FirDataFlowAnalyzer(
         } else {
             resetReceivers() // to state before class initialization
         }
+        graph?.completePostponedNodes()
         return graph
     }
 
@@ -247,6 +249,7 @@ abstract class FirDataFlowAnalyzer(
     fun exitScript(): ControlFlowGraph {
         val (node, graph) = graphBuilder.exitScript()
         node.mergeIncomingFlow()
+        graph.completePostponedNodes()
         return graph
     }
 
@@ -266,6 +269,7 @@ abstract class FirDataFlowAnalyzer(
     fun exitCodeFragment(): ControlFlowGraph {
         val (node, graph) = graphBuilder.exitCodeFragment()
         node.mergeIncomingFlow()
+        graph.completePostponedNodes()
         return graph
     }
     // ----------------------------------- Value parameters (and it's defaults) -----------------------------------
@@ -280,6 +284,7 @@ abstract class FirDataFlowAnalyzer(
         val (innerNode, outerNode, graph) = graphBuilder.exitValueParameter(valueParameter) ?: return null
         innerNode.mergeIncomingFlow()
         outerNode.mergeIncomingFlow()
+        graph.completePostponedNodes()
         return graph
     }
 
@@ -292,6 +297,7 @@ abstract class FirDataFlowAnalyzer(
     fun exitProperty(property: FirProperty): ControlFlowGraph? {
         val (node, graph) = graphBuilder.exitProperty(property) ?: return null
         node.mergeIncomingFlow()
+        graph.completePostponedNodes()
         return graph
     }
 
@@ -304,6 +310,7 @@ abstract class FirDataFlowAnalyzer(
     fun exitField(field: FirField): ControlFlowGraph? {
         val (node, graph) = graphBuilder.exitField(field) ?: return null
         node.mergeIncomingFlow()
+        graph.completePostponedNodes()
         return graph
     }
 
@@ -1101,6 +1108,7 @@ abstract class FirDataFlowAnalyzer(
     fun exitInitBlock(): ControlFlowGraph {
         val (node, controlFlowGraph) = graphBuilder.exitInitBlock()
         node.mergeIncomingFlow()
+        controlFlowGraph.completePostponedNodes()
         return controlFlowGraph
     }
 
@@ -1310,6 +1318,17 @@ abstract class FirDataFlowAnalyzer(
         }
     }
 
+    private fun ControlFlowGraph.completePostponedNodes() {
+        for (subGraph in subGraphs) {
+            subGraph.completePostponedNodes()
+        }
+        for (node in nodes) {
+            if (node !is ClassExitNode && !node.flowInitialized) {
+                node.mergeIncomingFlow()
+            }
+        }
+    }
+
     // In rare cases (like after exiting functions) after adding more nodes `graphBuilder` will revert the current
     // state to a previously created node, so none of the nodes it returned are `lastNode` and `mergeIncomingFlow`
     // will not ensure consistency. In that case an explicit call to `resetReceivers` is needed to roll back the stack

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

@@ -15,11 +15,9 @@ import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.resolve.dfa.FlowPath
 import org.jetbrains.kotlin.fir.resolve.dfa.PersistentFlow
 import org.jetbrains.kotlin.fir.resolve.dfa.controlFlowGraph
-import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.constructClassLikeType
 import org.jetbrains.kotlin.fir.types.isNothing
-import org.jetbrains.kotlin.fir.types.resolvedType
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 import org.jetbrains.kotlin.name.StandardClassIds
@@ -117,6 +115,7 @@ sealed class CFGNode<out E : FirElement>(val owner: ControlFlowGraph, val level:
      * be used for all type resolutions at this node.
      */
     private var _flow: PersistentFlow? = null
+    open val flowInitialized: Boolean get() = _flow != null
     open var flow: PersistentFlow
         get() = _flow ?: throw IllegalStateException("flow for $this not initialized - traversing nodes in wrong order?")
         @CfgInternals
@@ -264,17 +263,6 @@ class PostponedLambdaExitNode(owner: ControlFlowGraph, override val fir: FirAnon
 }
 
 class MergePostponedLambdaExitsNode(owner: ControlFlowGraph, override val fir: FirElement, level: Int) : CFGNode<FirElement>(owner, level) {
-
-    private var _flowInitialized = false
-    val flowInitialized: Boolean get() = _flowInitialized
-    override var flow: PersistentFlow
-        get() = super.flow
-        @CfgInternals
-        set(value) {
-            super.flow = value
-            _flowInitialized = true
-        }
-
     override fun <R, D> accept(visitor: ControlFlowGraphVisitor<R, D>, data: D): R {
         return visitor.visitMergePostponedLambdaExitsNode(this, data)
     }
@@ -780,6 +768,7 @@ class StubNode(owner: ControlFlowGraph, level: Int) : CFGNode<FirStub>(owner, le
 
     override val fir: FirStub get() = FirStub
 
+    override val flowInitialized: Boolean get() = firstPreviousNode.flowInitialized
     override var flow: PersistentFlow
         get() = firstPreviousNode.flow
         @CfgInternals

compiler/testData/diagnostics/tests/controlFlowAnalysis/reassignmentInLambda.kt

@@ -0,0 +1,9 @@
+// FIR_IDENTICAL
+// ISSUE: KT-61794
+
+private fun createStubFunction(expression: String?): String? {
+    val tmp = expression?.let {
+        it
+    } ?: return null
+    return tmp
+}

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

@@ -6990,6 +6990,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                 runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/reassignmentInCatch.kt");
             }
 
+            @Test
+            @TestMetadata("reassignmentInLambda.kt")
+            public void testReassignmentInLambda() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/reassignmentInLambda.kt");
+            }
+
             @Test
             @TestMetadata("reassignmentInTryCatch.kt")
             public void testReassignmentInTryCatch() throws Exception {

compiler/tests-compiler-utils/tests/org/jetbrains/kotlin/fir/FirCfgConsistencyChecker.kt

@@ -7,10 +7,7 @@ package org.jetbrains.kotlin.fir
 
 import org.jetbrains.kotlin.fir.references.FirControlFlowGraphReference
 import org.jetbrains.kotlin.fir.resolve.dfa.FirControlFlowGraphReferenceImpl
-import org.jetbrains.kotlin.fir.resolve.dfa.cfg.CFGNode
-import org.jetbrains.kotlin.fir.resolve.dfa.cfg.ControlFlowGraph
-import org.jetbrains.kotlin.fir.resolve.dfa.cfg.GraphEnterNodeMarker
-import org.jetbrains.kotlin.fir.resolve.dfa.cfg.GraphExitNodeMarker
+import org.jetbrains.kotlin.fir.resolve.dfa.cfg.*
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.test.Assertions
 
@@ -41,6 +38,7 @@ class FirCfgConsistencyChecker(private val assertions: Assertions) : FirVisitorV
                 assertions.assertContainsElements(from.followingNodes, node)
             }
             assertions.assertFalse(node.followingNodes.isEmpty() && node.previousNodes.isEmpty()) { "Unconnected CFG node: $node" }
+            assertions.assertTrue(node is ClassExitNode || node.flowInitialized) { "All nodes must have a flow: $node" }
         }
     }