FIR CFG: remove liveness analysis from lambda reassignment analyzer

Handling try/catch and break/continue is too hard here, and failing to
do that properly can create incorrect smartcast. Better be conservative
and assume all code is live.

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

@@ -253,10 +253,6 @@ internal class FirLocalVariableAssignmentAnalyzer(
          *   referenced outside the lambda. It's possible to track the scope at each node and remove the unneeded `a` in "entry" and "then"
          *   nodes. But doing that seems to be more expensive than simply letting it propagate.
          *
-         * There is some liveness analysis present in this construction which errs on the conservative side: some code that is dead
-         * is considered to be live. This restricts some smart casts, which is better than the opposite (permitting unsafe smart casts
-         * if code turns out to be only mostly dead).
-         *
          * Because names are not resolved at this point, we manually track local variable declarations and resolve them along the way
          * so that shadowed names are handled correctly. This works because local variables at any scope have higher priority
          * than members on implicit receivers, even if the implicit receiver is introduced by a later scope.
@@ -277,15 +273,8 @@ internal class FirLocalVariableAssignmentAnalyzer(
         )
 
         private class MiniFlow(val parents: Set<MiniFlow>) {
-            val children: MutableSet<MiniFlow> = mutableSetOf()
             val assignedLater: MutableSet<FirProperty> = mutableSetOf()
 
-            init {
-                for (parent in parents) {
-                    parent.children.add(this)
-                }
-            }
-
             fun fork(): MiniFlow = MiniFlow(setOf(this))
 
             companion object {
@@ -305,10 +294,9 @@ internal class FirLocalVariableAssignmentAnalyzer(
                 visitFunction(simpleFunction, data)
 
             override fun visitFunction(function: FirFunction, data: MiniCfgData) {
-                val currentFlow = data.flow ?: return
                 val freeVariables = data.variableDeclarations.flatMapTo(mutableSetOf()) { it.values }
-                val flowInto = currentFlow.fork()
+                val flowInto = data.flow.fork()
-                val flowAfter = currentFlow.fork()
+                val flowAfter = data.flow.fork()
                 data.flow = flowInto
                 function.acceptChildren(this, data)
                 data.flow = flowAfter
@@ -318,95 +306,50 @@ internal class FirLocalVariableAssignmentAnalyzer(
 
             override fun visitWhenExpression(whenExpression: FirWhenExpression, data: MiniCfgData) {
                 (whenExpression.subjectVariable ?: whenExpression.subject)?.accept(this, data)
-                val flow = data.flow ?: return
+                val flow = data.flow
                 // Also collect `flow` here for the case when none of the branches execute.
-                val branches = whenExpression.branches.mapNotNullTo(mutableSetOf(flow)) {
+                data.flow = whenExpression.branches.mapTo(mutableSetOf(flow)) {
                     // No need to create a fork - it'll not be observed anywhere anyway.
                     data.flow = flow
                     it.accept(this, data)
                     data.flow
-                }
+                }.join()
-                data.flow = branches.join()
             }
 
             override fun visitTryExpression(tryExpression: FirTryExpression, data: MiniCfgData) {
-                if (data.flow == null) return
                 tryExpression.tryBlock.accept(this, data)
-                val catchFlow = data.lastLiveFlow // descendant of flow at the start
+                val flow = data.flow
-                val returnFlows = mutableSetOf<MiniFlow>()
+                data.flow = tryExpression.catches.mapTo(mutableSetOf(flow)) {
-                data.flow?.let(returnFlows::add)
+                    data.flow = flow
-                val finallyFlows = tryExpression.catches.mapTo(mutableSetOf(catchFlow)) {
-                    data.flow = catchFlow
                     it.accept(this, data)
-                    data.flow?.let(returnFlows::add)
+                    data.flow
-                    // Throwing/returning inside a catch clause goes through finally as well.
+                }.join()
-                    data.lastLiveFlow
+                tryExpression.finallyBlock?.accept(this, data)
-                }
-                val finally = tryExpression.finallyBlock
-                if (finally != null) {
-                    data.flow = finallyFlows.join()
-                    finally.accept(this, data)
-                    if (returnFlows.isEmpty()) {
-                        data.flow = null
-                    }
-                } else {
-                    data.flow = returnFlows.join()
-                }
             }
 
-            private fun Set<MiniFlow>.join(): MiniFlow? =
+            private fun Set<MiniFlow>.join(): MiniFlow =
-                when (size) {
+                singleOrNull() ?: MiniFlow(this)
-                    0 -> null
-                    1 -> single()
-                    else -> MiniFlow(this)
-                }
 
             override fun visitWhileLoop(whileLoop: FirWhileLoop, data: MiniCfgData) {
                 // Loop entry is a merge point, so need a new node.
-                val start = data.flow?.fork() ?: return
+                val start = data.flow.fork()
                 data.flow = start
                 whileLoop.condition.accept(this, data)
-                if (data.flow == null) return
                 whileLoop.block.accept(this, data)
-                // There may have been a conditional break/continue before the return, or the condition
+                data.flow.addBackEdgeTo(start)
-                // may have always been false.
-                data.flow = data.lastLiveFlow // descendant of flow after condition
-                data.flow?.addBackEdgeTo(start)
             }
 
             override fun visitDoWhileLoop(doWhileLoop: FirDoWhileLoop, data: MiniCfgData) {
-                val start = data.flow?.fork() ?: return
+                val start = data.flow.fork()
-                data.flow = start
                 doWhileLoop.block.accept(this, data)
-                // Like above, there might have been a break/continue, so the fact that the block does not
-                // terminate doesn't actually mean much.
-                data.flow = data.lastLiveFlow // descendant of flow before the block
                 doWhileLoop.condition.accept(this, data)
-                data.flow?.addBackEdgeTo(start)
+                data.flow.addBackEdgeTo(start)
             }
 
-            override fun visitBreakExpression(breakExpression: FirBreakExpression, data: MiniCfgData) {
+            // TODO: liveness analysis - return/throw/break/continue terminate the flow.
-                visitElement(breakExpression, data)
+            //   This is somewhat problematic though because try-catch and loops can restore it.
-                // Can treat this as an unconditional return if looping constructs reset the flow anyway.
+            //   It is not possible to implement liveness analysis partially - otherwise combinations of
-                // TODO: check which loop this is targeting for more precise liveness analysis?
+            //   control flow structures can cause incorrect smartcasts.
-                data.flow = null
-            }
-
-            override fun visitContinueExpression(continueExpression: FirContinueExpression, data: MiniCfgData) {
-                visitElement(continueExpression, data)
-                // Same comment as for `break`.
-                data.flow = null
-            }
-
-            override fun visitReturnExpression(returnExpression: FirReturnExpression, data: MiniCfgData) {
-                visitElement(returnExpression, data)
-                data.flow = null
-            }
-
-            override fun visitThrowExpression(throwExpression: FirThrowExpression, data: MiniCfgData) {
-                visitElement(throwExpression, data)
-                data.flow = null
-            }
 
             override fun visitFunctionCall(functionCall: FirFunctionCall, data: MiniCfgData) {
                 val visitor = this
@@ -449,40 +392,23 @@ internal class FirLocalVariableAssignmentAnalyzer(
             private fun MiniCfgData.recordAssignment(reference: FirReference) {
                 val name = (reference as? FirNamedReference)?.name ?: return
                 val property = variableDeclarations.lastOrNull { name in it }?.get(name) ?: return
-                flow?.recordAssignment(property, mutableSetOf())
+                flow.recordAssignments(setOf(property), mutableSetOf())
             }
 
-            private fun MiniFlow.recordAssignment(property: FirProperty, visited: MutableSet<MiniFlow>) {
+            private fun MiniFlow.recordAssignments(properties: Set<FirProperty>, visited: MutableSet<MiniFlow>) {
                 if (!visited.add(this)) return
-                assignedLater += property
+                assignedLater += properties
-                parents.forEach { it.recordAssignment(property, visited) }
+                parents.forEach { it.recordAssignments(properties, visited) }
             }
 
             private fun MiniFlow.addBackEdgeTo(loopStart: MiniFlow) {
-                children.add(loopStart)
                 // All forks in the loop should have the same set of variables assigned later, equal to the set
                 // at the start of the loop.
-                propagateForward(loopStart.assignedLater, mutableSetOf())
+                recordAssignments(loopStart.assignedLater, mutableSetOf())
-            }
-
-            private fun MiniFlow.propagateForward(properties: Set<FirProperty>, visited: MutableSet<MiniFlow>) {
-                if (!visited.add(this)) return
-                assignedLater.addAll(properties)
-                children.forEach { it.propagateForward(properties, visited) }
             }
 
             class MiniCfgData {
-                var lastLiveFlow: MiniFlow = MiniFlow.start()
+                var flow: MiniFlow = MiniFlow.start()
-                    private set
-
-                var flow: MiniFlow? = lastLiveFlow
-                    set(value) {
-                        if (value != null) {
-                            lastLiveFlow = value
-                        }
-                        field = value
-                    }
-
                 val variableDeclarations: ArrayDeque<MutableMap<Name, FirProperty>> = ArrayDeque(listOf(mutableMapOf()))
                 val functionForks: MutableMap<FirFunctionSymbol<*>, FunctionFork> = mutableMapOf()
             }

compiler/testData/diagnostics/tests/smartCasts/variables/capturedWithControlJumps.fir.kt

@@ -50,6 +50,6 @@ fun test4() {
             return
         }
         x = null
-        lambda?.invoke()
+        lambda<!UNNECESSARY_SAFE_CALL!>?.<!>invoke()
     }
 }