FIR CFA: refactor handling of classes a bit

As usual, I'm leaving TODOs wherever I go...

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

@@ -13,7 +13,6 @@ import org.jetbrains.kotlin.fir.contracts.description.ConeConditionalEffectDecla
 import org.jetbrains.kotlin.fir.contracts.description.ConeReturnsEffectDeclaration
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertyAccessor
-import org.jetbrains.kotlin.fir.declarations.utils.isLocal
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirNoReceiverExpression
 import org.jetbrains.kotlin.fir.references.FirControlFlowGraphReference
@@ -204,34 +203,23 @@ abstract class FirDataFlowAnalyzer(
 
     // ----------------------------------- Classes -----------------------------------
 
-    fun enterClass() {
+    fun enterClass(klass: FirClass, buildGraph: Boolean) {
-        graphBuilder.enterClass()
+        graphBuilder.enterClass(klass, buildGraph)?.mergeIncomingFlow()
     }
 
-    fun exitClass() {
+    fun exitClass(): ControlFlowGraph? {
-        graphBuilder.exitClass()
+        // TODO: support capturing of mutable properties
-    }
+        //   var x: String?
-
+        //   x = ""
-    fun exitRegularClass(klass: FirRegularClass): ControlFlowGraph {
+        //   class C { init { x = maybeNull() } }
-        if (klass.isLocal && components.container !is FirClass) return exitLocalClass(klass)
+        //   // C will be initialized at first use - x is no longer safe to smart cast
-        return graphBuilder.exitClass(klass)
+        val (node, graph) = graphBuilder.exitClass() ?: return null
-    }
+        if (node != null) {
-
+            node.mergeIncomingFlow()
-    private fun exitLocalClass(klass: FirRegularClass): ControlFlowGraph {
+        } else {
-        // TODO: support capturing of mutable properties, KT-44877
+            resetReceivers() // to state before class initialization
-        val (node, controlFlowGraph) = graphBuilder.exitLocalClass(klass)
+        }
-        node.mergeIncomingFlow()
+        return graph
-        return controlFlowGraph
-    }
-
-    fun enterAnonymousObject(anonymousObject: FirAnonymousObject) {
-        graphBuilder.enterAnonymousObject(anonymousObject).mergeIncomingFlow()
-    }
-
-    fun exitAnonymousObject(anonymousObject: FirAnonymousObject): ControlFlowGraph {
-        val (node, controlFlowGraph) = graphBuilder.exitAnonymousObject(anonymousObject)
-        node.mergeIncomingFlow()
-        return controlFlowGraph
     }
 
     fun exitAnonymousObjectExpression(anonymousObjectExpression: FirAnonymousObjectExpression) {

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

@@ -66,8 +66,6 @@ class ControlFlowGraphBuilder {
     private val finallyEnterNodes: Stack<FinallyBlockEnterNode> = stackOf()
     private val finallyBlocksInProgress: Stack<FinallyBlockEnterNode> = stackOf()
 
-    private val initBlockExitNodes: Stack<InitBlockExitNode> = stackOf()
-
     private val exitSafeCallNodes: Stack<ExitSafeCallNode> = stackOf()
     private val exitElvisExpressionNodes: Stack<ElvisExitNode> = stackOf()
     private val elvisRhsEnterNodes: Stack<ElvisRhsEnterNode> = stackOf()
@@ -148,9 +146,7 @@ class ControlFlowGraphBuilder {
         if (localFunctionNode != null) {
             addEdge(localFunctionNode, enterNode)
         } else {
-            enterToLocalClassesMembers[function.symbol]?.let {
+            enterToLocalClassesMembers.remove(function.symbol)?.let { addEdge(it, enterNode, preferredKind = EdgeKind.DfgForward) }
-                addEdge(it, enterNode, preferredKind = EdgeKind.DfgForward)
-            }
         }
 
         createFunctionExitNode(function).also {
@@ -164,14 +160,12 @@ class ControlFlowGraphBuilder {
     fun exitFunction(function: FirFunction): Pair<FunctionExitNode, ControlFlowGraph> {
         require(function !is FirAnonymousFunction)
         val exitNode = exitTargetsForReturn.pop()
+        exitTargetsForTry.pop().also { assert(it == exitNode) }
         popAndAddEdge(exitNode)
         val graph = popGraph()
         assert(exitNode == graph.exitNode)
-        exitTargetsForTry.pop().also {
+        exitNode.updateDeadStatus()
-            assert(it == graph.exitNode)
+        return exitNode to graph
-        }
-        graph.exitNode.updateDeadStatus()
-        return graph.exitNode as FunctionExitNode to graph
     }
 
     // ----------------------------------- Anonymous function -----------------------------------
@@ -367,22 +361,70 @@ class ControlFlowGraphBuilder {
 
     // ----------------------------------- Classes -----------------------------------
 
-    fun enterClass() {
+    fun enterClass(klass: FirClass, buildGraph: Boolean): CFGNode<*>? {
-        pushGraph(ControlFlowGraph(null, "STUB_CLASS_GRAPH", ControlFlowGraph.Kind.Stub))
+        if (!buildGraph) {
-    }
+            pushGraph(ControlFlowGraph(null, "STUB_CLASS_GRAPH", ControlFlowGraph.Kind.Stub))
+            return null
+        }
 
-    fun exitClass() {
+        val enterNode = when {
-        popGraph()
+            klass is FirAnonymousObject -> createAnonymousObjectEnterNode(klass)
-    }
+            // Local classes are only initialized on first use, so they look pretty much like named functions:
+            // control flow enters here and never leaves, and assignments invalidate smart casts.
+            klass is FirRegularClass && klass.isLocal && currentGraph.kind.withBody -> createLocalClassExitNode(klass)
+            else -> null
+        }
 
-    fun exitClass(klass: FirClass): ControlFlowGraph {
-        exitClass()
         val name = when (klass) {
             is FirAnonymousObject -> "<anonymous object>"
             is FirRegularClass -> klass.name.asString()
             else -> throw IllegalArgumentException("Unknown class kind: ${klass::class}")
         }
+        pushGraph(ControlFlowGraph(klass, name, ControlFlowGraph.Kind.ClassInitializer))
 
+        val graphEnterNode = createClassEnterNode(klass)
+        if (enterNode != null) {
+            // TODO: anonymous objects are used to represent enum entries - check what happens there
+            //  (likely `exitClass` leaves a node on the stack that will never be consumed)
+            lastNodes.popOrNull()?.let { addEdge(it, enterNode) }
+            lastNodes.push(enterNode)
+            addEdge(enterNode, graphEnterNode, preferredKind = EdgeKind.CfgForward)
+        }
+        return enterNode
+    }
+
+    fun exitClass(): Pair<AnonymousObjectExitNode?, ControlFlowGraph>? {
+        if (currentGraph.kind == ControlFlowGraph.Kind.Stub) {
+            popGraph()
+            return null
+        }
+
+        val graph = popClassGraph()
+        if (graph.declaration !is FirAnonymousObject) {
+            return null to graph
+        }
+
+        val lastNode = lastNodes.pop() as AnonymousObjectEnterNode
+        val exitNode = createAnonymousObjectExitNode(lastNode.fir).also { lastNodes.push(it) }
+        // TODO: should merge data flow from members into the exit node.
+        // TODO: the reason for this liveness trickery is that if control flow is dead, the CFG-only edge from
+        //  `graph.exitNode` gets magically transformed into a CFG+DFG dead edge, and `graph.exitNode` has no data
+        //  flow information attached to it. This should be fixed as soon as data flow is made to go through the object.
+        if (!graph.exitNode.isDead) {
+            // This implies that `lastNode` is not dead either.
+            addEdge(graph.exitNode, exitNode, preferredKind = EdgeKind.CfgForward)
+            addEdge(lastNode, exitNode, preferredKind = EdgeKind.DfgForward)
+        } else {
+            addEdge(lastNode, exitNode, isDead = true)
+        }
+        return exitNode to graph
+    }
+
+    // Members of a class can be visited in any order, so data flow between them is unordered
+    // and we have to recreate the control flow after the fact.
+    private fun popClassGraph(): ControlFlowGraph {
+        assert(currentGraph.kind == ControlFlowGraph.Kind.ClassInitializer)
+        val klass = currentGraph.declaration as FirClass
         val calledInPlace = mutableListOf<ControlFlowGraph>()
         val calledLater = mutableListOf<ControlFlowGraph>()
         for (declaration in klass.declarations) {
@@ -397,24 +439,26 @@ class ControlFlowGraphBuilder {
             }
         }
 
-        pushGraph(ControlFlowGraph(klass, name, ControlFlowGraph.Kind.ClassInitializer))
+        var node: CFGNode<*> = currentGraph.enterNode as ClassEnterNode
-        val exitNode = createClassExitNode(klass)
-        var node: CFGNode<*> = createClassEnterNode(klass)
         var prevInitPartNode: CFGNode<*>? = null
         for (graph in calledInPlace) {
-            createPartOfClassInitializationNode(graph.declaration as FirControlFlowGraphOwner).also {
+            val partNode = createPartOfClassInitializationNode(graph.declaration as FirControlFlowGraphOwner)
-                addEdge(node, it, preferredKind = EdgeKind.CfgForward)
+            // TODO: if one initializer part does not terminate, deadness becomes funky here
-                addEdge(it, graph.enterNode, preferredKind = EdgeKind.CfgForward)
+            addEdge(node, partNode, preferredKind = EdgeKind.CfgForward)
-                node = graph.exitNode
+            addEdge(partNode, graph.enterNode, preferredKind = EdgeKind.CfgForward)
-
+            if (prevInitPartNode != null) {
-                if (prevInitPartNode != null) addEdge(prevInitPartNode!!, it, preferredKind = EdgeKind.DeadForward)
+                // Fake edge to make the nodes in this graph orderable without looking at subgraphs.
-                it.updateDeadStatus()
+                addEdge(prevInitPartNode, partNode, preferredKind = EdgeKind.DeadForward, propagateDeadness = false)
-                prevInitPartNode = it
             }
+            node = graph.exitNode
+            prevInitPartNode = partNode
         }
+
+        val exitNode = createClassExitNode(klass)
         addEdge(node, exitNode, preferredKind = EdgeKind.CfgForward)
-        if (prevInitPartNode != null) addEdge(prevInitPartNode!!, exitNode, preferredKind = EdgeKind.DeadForward)
+        if (prevInitPartNode != null) {
-        exitNode.updateDeadStatus()
+            addEdge(prevInitPartNode, exitNode, preferredKind = EdgeKind.DeadForward, propagateDeadness = false)
+        }
 
         // TODO: Here we're assuming that the methods are called after the object is constructed, which is really not true
         //   (init blocks can call them). But FE1.0 did so too, hence the following code compiles and prints 0:
@@ -433,6 +477,9 @@ class ControlFlowGraphBuilder {
     }
 
     fun prepareForLocalClassMembers(members: Collection<FirDeclaration>) {
+        // TODO: this is called before `enterClass` so the data flow source for objects and local classes
+        //  is not the enter node, but whichever node happens to be before it. This technically works,
+        //  but is ugly.
         members.forEachMember {
             enterToLocalClassesMembers[it.symbol] = lastNodes.topOrNull()
         }
@@ -444,42 +491,8 @@ class ControlFlowGraphBuilder {
         }
     }
 
-    fun exitLocalClass(klass: FirRegularClass): Pair<LocalClassExitNode, ControlFlowGraph> {
-        val graph = exitClass(klass)
-        val node = createLocalClassExitNode(klass).also {
-            addNewSimpleNodeIfPossible(it)
-        }
-        addEdge(node, graph.enterNode, preferredKind = EdgeKind.CfgForward)
-        return node to graph
-    }
-
-    fun enterAnonymousObject(anonymousObject: FirAnonymousObject): AnonymousObjectEnterNode {
-        val enterNode = createAnonymousObjectEnterNode(anonymousObject)
-        // TODO: looks like there was some problem with enum initializers that causes `lastNodes` to be empty
-        lastNodes.popOrNull()?.let { addEdge(it, enterNode, preferredKind = EdgeKind.Forward) }
-        lastNodes.push(enterNode)
-        enterClass()
-        return enterNode
-    }
-
-    fun exitAnonymousObject(anonymousObject: FirAnonymousObject): Pair<AnonymousObjectExitNode, ControlFlowGraph> {
-        val graph = exitClass(anonymousObject)
-        val enterNode = lastNodes.popOrNull()
-        if (enterNode !is AnonymousObjectEnterNode) {
-            throw AssertionError("anonymous object exit should be preceded by anonymous object enter, but got $enterNode")
-        }
-        val exitNode = createAnonymousObjectExitNode(anonymousObject)
-        // TODO: Intentionally not using anonymous object init blocks for data flow? Might've been a FE1.0 bug.
-        addEdge(enterNode, graph.enterNode, preferredKind = EdgeKind.CfgForward)
-        if (!graph.exitNode.isDead) {
-            addEdge(graph.exitNode, exitNode, preferredKind = EdgeKind.CfgForward)
-        }
-        addEdge(enterNode, exitNode, preferredKind = EdgeKind.DfgForward)
-        lastNodes.push(exitNode)
-        return exitNode to graph
-    }
-
     fun exitAnonymousObjectExpression(anonymousObjectExpression: FirAnonymousObjectExpression): AnonymousObjectExpressionExitNode {
+        // TODO: what's AnonymousObjectExitNode for then?
         return createAnonymousObjectExpressionExitNode(anonymousObjectExpression).also {
             addNewSimpleNodeIfPossible(it)
         }
@@ -554,11 +567,7 @@ class ControlFlowGraphBuilder {
         val enterNode = createPropertyInitializerEnterNode(property)
         val exitNode = createPropertyInitializerExitNode(property)
         exitTargetsForTry.push(exitNode)
-
+        enterToLocalClassesMembers.remove(property.symbol)?.let { addEdge(it, enterNode, preferredKind = EdgeKind.DfgForward) }
-        enterToLocalClassesMembers[property.symbol]?.let {
-            addEdge(it, enterNode, preferredKind = EdgeKind.DfgForward)
-        }
-
         lastNodes.push(enterNode)
         return enterNode
     }
@@ -582,11 +591,7 @@ class ControlFlowGraphBuilder {
         val enterNode = createFieldInitializerEnterNode(field)
         val exitNode = createFieldInitializerExitNode(field)
         exitTargetsForTry.push(exitNode)
-
+        enterToLocalClassesMembers.remove(field.symbol)?.let { addEdge(it, enterNode, preferredKind = EdgeKind.DfgForward) }
-        enterToLocalClassesMembers[field.symbol]?.let {
-            addEdge(it, enterNode, preferredKind = EdgeKind.DfgForward)
-        }
-
         lastNodes.push(enterNode)
         return enterNode
     }
@@ -1210,25 +1215,18 @@ class ControlFlowGraphBuilder {
 
     fun enterInitBlock(initBlock: FirAnonymousInitializer): InitBlockEnterNode {
         // TODO: questionable moment that we should pass data flow from init to init
-
         pushGraph(ControlFlowGraph(initBlock, "init block", ControlFlowGraph.Kind.Function))
-        val enterNode = createInitBlockEnterNode(initBlock).also {
-            lastNodes.push(it)
-        }
-        val lastNode = runIf(lastNode is InitBlockExitNode) { lastNodes.pop() } ?: enterToLocalClassesMembers[initBlock.symbol]
-        lastNode?.let { addEdge(it, enterNode, preferredKind = EdgeKind.DfgForward) }
-
-        createInitBlockExitNode(initBlock).also {
-            initBlockExitNodes.push(it)
-            exitTargetsForTry.push(it)
-        }
 
+        val enterNode = createInitBlockEnterNode(initBlock)
+        val exitNode = createInitBlockExitNode(initBlock)
+        exitTargetsForTry.push(exitNode)
+        enterToLocalClassesMembers.remove(initBlock.symbol)?.let { addEdge(it, enterNode, preferredKind = EdgeKind.DfgForward) }
+        lastNodes.push(enterNode)
         return enterNode
     }
 
     fun exitInitBlock(initBlock: FirAnonymousInitializer): Pair<InitBlockExitNode, ControlFlowGraph> {
-        val exitNode = initBlockExitNodes.pop()
+        val exitNode = exitTargetsForTry.pop() as InitBlockExitNode
-        require(exitNode == exitTargetsForTry.pop())
         popAndAddEdge(exitNode)
         val graph = popGraph()
         assert(graph.declaration == initBlock)

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

@@ -540,19 +540,14 @@ open class FirDeclarationsResolveTransformer(transformer: FirAbstractBodyResolve
         regularClass: FirRegularClass,
         data: ResolutionMode
     ): FirRegularClass {
-        dataFlowAnalyzer.enterClass()
+        dataFlowAnalyzer.enterClass(regularClass, !implicitTypeOnly)
-
         val result = context.withRegularClass(regularClass, components) {
             transformDeclarationContent(regularClass, data) as FirRegularClass
         }
-
+        val controlFlowGraph = dataFlowAnalyzer.exitClass()
-        if (!implicitTypeOnly) {
+        if (controlFlowGraph != null) {
-            val controlFlowGraph = dataFlowAnalyzer.exitRegularClass(result)
             result.replaceControlFlowGraphReference(FirControlFlowGraphReferenceImpl(controlFlowGraph))
-        } else {
-            dataFlowAnalyzer.exitClass()
         }
-
         return result
     }
 
@@ -569,19 +564,15 @@ open class FirDeclarationsResolveTransformer(transformer: FirAbstractBodyResolve
                 transformer.firProviderInterceptor
             )
         }
-        if (!implicitTypeOnly && anonymousObject.controlFlowGraphReference == null) {
+        // TODO: why would there be a graph already?
-            dataFlowAnalyzer.enterAnonymousObject(anonymousObject)
+        val buildGraph = !implicitTypeOnly && anonymousObject.controlFlowGraphReference == null
-        } else {
+        dataFlowAnalyzer.enterClass(anonymousObject, buildGraph)
-            dataFlowAnalyzer.enterClass()
-        }
         val result = context.withAnonymousObject(anonymousObject, components) {
             transformDeclarationContent(anonymousObject, data) as FirAnonymousObject
         }
-        if (!implicitTypeOnly && result.controlFlowGraphReference == null) {
+        val graph = dataFlowAnalyzer.exitClass()
-            val graph = dataFlowAnalyzer.exitAnonymousObject(result)
+        if (graph != null) {
             result.replaceControlFlowGraphReference(FirControlFlowGraphReferenceImpl(graph))
-        } else {
-            dataFlowAnalyzer.exitClass()
         }
         return result
     }