[FIR] Merge loop and lambda assignment look-ahead into single class

There are two classes that essentially do the same thing: look-ahead at
future assignments to determine smartcast stability. One class is used
primarily for loops (PreliminaryLoopVisitor), and the other is used
primarily for lambdas (FirLocalVariableAssignmentAnalyzer). Merge their
behavior to avoid smartcast inconsistencies between loops and lambdas.

^KT-57678 Fixed
^KT-59688 Fixed

Author: bnorm

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

@@ -44,8 +44,6 @@ import org.jetbrains.kotlin.util.OperatorNameConventions
 
 class DataFlowAnalyzerContext(session: FirSession) {
     val graphBuilder = ControlFlowGraphBuilder()
-    val preliminaryLoopVisitor = PreliminaryLoopVisitor()
-    val variablesClearedBeforeLoop = stackOf<List<RealVariable>>()
     internal val variableAssignmentAnalyzer = FirLocalVariableAssignmentAnalyzer()
 
     var variableStorage = VariableStorageImpl(session)
@@ -59,8 +57,6 @@ class DataFlowAnalyzerContext(session: FirSession) {
 
     fun reset() {
         graphBuilder.reset()
-        preliminaryLoopVisitor.resetState()
-        variablesClearedBeforeLoop.reset()
         variableAssignmentAnalyzer.reset()
         variableStorage = variableStorage.clear()
     }
@@ -181,6 +177,8 @@ abstract class FirDataFlowAnalyzer(
     fun enterFunction(function: FirFunction) {
         if (function is FirDefaultPropertyAccessor) return
 
+        val assignedInside = context.variableAssignmentAnalyzer.enterFunction(function)
+
         val (localFunctionNode, functionEnterNode) = if (function is FirAnonymousFunction) {
             null to graphBuilder.enterAnonymousFunction(function)
         } else {
@@ -192,22 +190,17 @@ abstract class FirDataFlowAnalyzer(
              * Anonymous functions which can be revisited, either in-place or not in-place, are treated as repeatable statements. This
              * causes any assignments to local variables within the anonymous function body to clear type statements for those local
              * variables.
-             * TODO(KT-57678): is it possible for FirLocalVariableAssignmentAnalyzer to handle this for both lambdas and loops?
              */
             if (function is FirAnonymousFunction && function.invocationKind?.canBeRevisited() != false) {
-                enterRepeatableStatement(flow, function)
+                enterRepeatableStatement(flow, assignedInside)
             }
         }
-        context.variableAssignmentAnalyzer.enterFunction(function)
     }
 
     fun exitFunction(function: FirFunction): FirControlFlowGraphReference? {
         if (function is FirDefaultPropertyAccessor) return null
 
         context.variableAssignmentAnalyzer.exitFunction()
-        if (function is FirAnonymousFunction && function.invocationKind?.canBeRevisited() != false) {
-            exitRepeatableStatement(function)
-        }
 
         if (function is FirAnonymousFunction) {
             val (functionExitNode, postponedLambdaExitNode, graph) = graphBuilder.exitAnonymousFunction(function)
@@ -699,9 +692,10 @@ abstract class FirDataFlowAnalyzer(
     // ----------------------------------- While Loop -----------------------------------
 
     fun enterWhileLoop(loop: FirLoop) {
+        val assignedInside = context.variableAssignmentAnalyzer.enterLoop(loop)
         val (loopEnterNode, loopConditionEnterNode) = graphBuilder.enterWhileLoop(loop)
         loopEnterNode.mergeIncomingFlow()
-        loopConditionEnterNode.mergeIncomingFlow { _, flow -> enterRepeatableStatement(flow, loop) }
+        loopConditionEnterNode.mergeIncomingFlow { _, flow -> enterRepeatableStatement(flow, assignedInside) }
     }
 
     fun exitWhileLoopCondition(loop: FirLoop) {
@@ -716,17 +710,25 @@ abstract class FirDataFlowAnalyzer(
     }
 
     fun exitWhileLoop(loop: FirLoop) {
+        val assignedInside = context.variableAssignmentAnalyzer.exitLoop()
         val (conditionEnterNode, blockExitNode, exitNode) = graphBuilder.exitWhileLoop(loop)
         blockExitNode.mergeIncomingFlow()
         exitNode.mergeIncomingFlow { path, flow ->
-            processWhileLoopExit(path, flow, exitNode, conditionEnterNode)
+            processWhileLoopExit(path, flow, exitNode, conditionEnterNode, assignedInside)
             processLoopExit(flow, exitNode, exitNode.firstPreviousNode as LoopConditionExitNode)
         }
     }
 
-    private fun processWhileLoopExit(path: FlowPath, flow: MutableFlow, node: LoopExitNode, conditionEnterNode: LoopConditionEnterNode) {
-        val possiblyChangedVariables = exitRepeatableStatement(node.fir)
-        if (possiblyChangedVariables.isNullOrEmpty()) return
+    private fun processWhileLoopExit(
+        path: FlowPath,
+        flow: MutableFlow,
+        node: LoopExitNode,
+        conditionEnterNode: LoopConditionEnterNode,
+        reassigned: Set<FirPropertySymbol>,
+    ) {
+        if (reassigned.isEmpty()) return
+        val possiblyChangedVariables = variableStorage.realVariables.values.filter { it.identifier.symbol in reassigned }
+        if (possiblyChangedVariables.isEmpty()) return
         // While analyzing the loop we might have added some backwards jumps to `conditionEnterNode` which weren't
         // there at the time its flow was computed - which is why we erased all information about `possiblyChangedVariables`
         // from it. Now that we have those edges, we can restore type information for the code after the loop.
@@ -758,33 +760,20 @@ abstract class FirDataFlowAnalyzer(
         }
     }
 
-    private fun enterRepeatableStatement(flow: MutableFlow, statement: FirStatement) {
-        val reassignedNames = context.preliminaryLoopVisitor.enterCapturingStatement(statement)
-        if (reassignedNames.isEmpty()) return
-        // TODO: only choose the innermost variable for each name, KT-59688
-        val possiblyChangedVariables = variableStorage.realVariables.values.filter {
-            val identifier = it.identifier
-            val symbol = identifier.symbol
-            // Non-local vars can never produce stable smart casts anyway.
-            identifier.dispatchReceiver == null && identifier.extensionReceiver == null &&
-                    symbol is FirPropertySymbol && symbol.isVar && symbol.name in reassignedNames
-        }
+    private fun enterRepeatableStatement(flow: MutableFlow, reassigned: Set<FirPropertySymbol>) {
+        if (reassigned.isEmpty()) return
+        val possiblyChangedVariables = variableStorage.realVariables.values.filter { it.identifier.symbol in reassigned }
         for (variable in possiblyChangedVariables) {
             logicSystem.recordNewAssignment(flow, variable, context.newAssignmentIndex())
         }
-        context.variablesClearedBeforeLoop.push(possiblyChangedVariables)
-    }
-
-    private fun exitRepeatableStatement(statement: FirStatement): List<RealVariable>? {
-        if (context.preliminaryLoopVisitor.exitCapturingStatement(statement).isEmpty()) return null
-        return context.variablesClearedBeforeLoop.pop()
     }
 
     // ----------------------------------- Do while Loop -----------------------------------
 
     fun enterDoWhileLoop(loop: FirLoop) {
+        val assignedInside = context.variableAssignmentAnalyzer.enterLoop(loop)
         val (loopEnterNode, loopBlockEnterNode) = graphBuilder.enterDoWhileLoop(loop)
-        loopEnterNode.mergeIncomingFlow { _, flow -> enterRepeatableStatement(flow, loop) }
+        loopEnterNode.mergeIncomingFlow { _, flow -> enterRepeatableStatement(flow, assignedInside) }
         loopBlockEnterNode.mergeIncomingFlow()
     }
 
@@ -795,12 +784,12 @@ abstract class FirDataFlowAnalyzer(
     }
 
     fun exitDoWhileLoop(loop: FirLoop) {
+        context.variableAssignmentAnalyzer.exitLoop()
         val (loopConditionExitNode, loopExitNode) = graphBuilder.exitDoWhileLoop(loop)
         loopConditionExitNode.mergeIncomingFlow()
         loopExitNode.mergeIncomingFlow { _, flow ->
             processLoopExit(flow, loopExitNode, loopConditionExitNode)
         }
-        exitRepeatableStatement(loop)
     }
 
     // ----------------------------------- Try-catch-finally -----------------------------------

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

@@ -5,17 +5,18 @@
 
 package org.jetbrains.kotlin.fir.resolve.dfa
 
+import org.jetbrains.kotlin.KtFakeSourceElementKind
 import org.jetbrains.kotlin.contracts.description.isInPlace
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.expressions.explicitReceiver
 import org.jetbrains.kotlin.fir.references.FirNamedReference
 import org.jetbrains.kotlin.fir.references.FirReference
 import org.jetbrains.kotlin.fir.references.toResolvedPropertySymbol
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.typeContext
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
@@ -31,7 +32,7 @@ import org.jetbrains.kotlin.types.AbstractTypeChecker
  **/
 internal class FirLocalVariableAssignmentAnalyzer {
     private var rootFunction: FirFunctionSymbol<*>? = null
-    private var assignedLocalVariablesByDeclaration: Map<FirBasedSymbol<*>, Fork>? = null
+    private var assignedLocalVariablesByDeclaration: Map<Any /* FirBasedSymbol<*> | FirLoop */, Fork>? = null
     private var variableAssignments: Map<FirProperty, List<Assignment>>? = null
 
     private val scopes: Stack<Pair<Fork?, VariableAssignments>> = stackOf()
@@ -82,14 +83,14 @@ internal class FirLocalVariableAssignmentAnalyzer {
         return assignments.all { AbstractTypeChecker.isSubtypeOf(session.typeContext, it.type!!, targetType) }
     }
 
-    private fun getInfoForDeclaration(symbol: FirBasedSymbol<*>): Fork? {
+    private fun getInfoForDeclaration(symbol: Any): Fork? {
         val root = rootFunction ?: return null
         if (root == symbol) return null
         val cachedMap = buildInfoForRoot(root)
         return cachedMap[symbol]
     }
 
-    private fun buildInfoForRoot(root: FirFunctionSymbol<*>): Map<FirBasedSymbol<*>, Fork> {
+    private fun buildInfoForRoot(root: FirFunctionSymbol<*>): Map<Any, Fork> {
         assignedLocalVariablesByDeclaration?.let { return it }
 
         val data = MiniCfgBuilder.MiniCfgData()
@@ -102,7 +103,7 @@ internal class FirLocalVariableAssignmentAnalyzer {
     }
 
     private fun enterScope(
-        symbol: FirBasedSymbol<*>,
+        symbol: Any,
         evaluatedInPlace: Boolean,
     ): Pair<Fork?, VariableAssignments> {
         val currentInfo = getInfoForDeclaration(symbol)
@@ -134,11 +135,15 @@ internal class FirLocalVariableAssignmentAnalyzer {
         return scopes.top()
     }
 
-    fun enterFunction(function: FirFunction) {
+    /**
+     * Enters an [FirFunction] and returns all [FirPropertySymbol]s which are defined before the function that will be modified within the
+     * function.
+     */
+    fun enterFunction(function: FirFunction): Set<FirPropertySymbol> {
         if (rootFunction == null) {
             rootFunction = function.symbol
             scopes.push(null to VariableAssignments())
-            return
+            return emptySet()
         }
         val (info, prohibitSmartCasts) =
             enterScope(function.symbol, function is FirAnonymousFunction && function.invocationKind.isInPlace)
@@ -149,6 +154,7 @@ internal class FirLocalVariableAssignmentAnalyzer {
                 }
             }
         }
+        return scopes.top().first?.assignedInside?.getAssignedProperties().orEmpty()
     }
 
     fun exitFunction() {
@@ -202,6 +208,24 @@ internal class FirLocalVariableAssignmentAnalyzer {
         }
     }
 
+    /**
+     * Enters an [FirLoop] and returns all [FirPropertySymbol]s which are defined before the loop that will be modified within the loop.
+     */
+    fun enterLoop(loop: FirLoop): Set<FirPropertySymbol> {
+        if (rootFunction == null) return emptySet()
+        val (info, _) = enterScope(loop, evaluatedInPlace = true)
+        return info?.assignedInside?.getAssignedProperties().orEmpty()
+    }
+
+    /**
+     * Exits an [FirLoop] and returns all [FirPropertySymbol]s which were defined before the loop that were modified within the loop.
+     */
+    fun exitLoop(): Set<FirPropertySymbol> {
+        if (rootFunction == null) return emptySet()
+        val (info, _) = scopes.pop()
+        return info?.assignedInside?.getAssignedProperties().orEmpty()
+    }
+
     fun visitAssignment(property: FirProperty, type: ConeKotlinType) {
         buildInfoForRoot(rootFunction ?: return)
         val assignments = variableAssignments?.get(property) ?: return
@@ -293,6 +317,7 @@ internal class FirLocalVariableAssignmentAnalyzer {
         )
 
         private class Assignment(
+            val operatorAssignment: Boolean,
             var type: ConeKotlinType? = null,
         )
 
@@ -329,6 +354,13 @@ internal class FirLocalVariableAssignmentAnalyzer {
             }
 
             fun isEmpty(): Boolean = assignments.isEmpty()
+
+            fun getAssignedProperties(): Set<FirPropertySymbol> {
+                return assignments.entries
+                    // TODO(KT-57563): Operator assignments should be treated just like any other assignment.
+                    .filter { (_, v) -> v.any { !it.operatorAssignment } }
+                    .mapTo(mutableSetOf()) { (k, _) -> k.symbol }
+            }
         }
 
         private class MiniFlow(val parents: Set<MiniFlow>) {
@@ -412,8 +444,8 @@ internal class FirLocalVariableAssignmentAnalyzer {
                 // All forks in the loop should have the same set of variables assigned later, equal to the set
                 // at the start of the loop.
                 data.flow.recordAssignments(assignedInside)
-                // The loop flows are detached from the entry flow, so we need to re-join them.
-                data.flow = setOf(entry, data.flow).join()
+                data.flow = entry.fork()
+                data.forks[loop] = Fork(data.flow.assignedLater, assignedInside)
             }
 
             override fun visitWhileLoop(whileLoop: FirWhileLoop, data: MiniCfgData) =
@@ -457,21 +489,24 @@ internal class FirLocalVariableAssignmentAnalyzer {
             override fun visitVariableAssignment(variableAssignment: FirVariableAssignment, data: MiniCfgData) {
                 visitElement(variableAssignment, data)
                 if (variableAssignment.explicitReceiver != null) return
-                variableAssignment.calleeReference?.let { data.recordAssignment(it) }
+                variableAssignment.calleeReference?.let {
+                    val operatorAssignment = variableAssignment.source?.kind is KtFakeSourceElementKind.DesugaredIncrementOrDecrement
+                    data.recordAssignment(it, operatorAssignment)
+                }
             }
 
             override fun visitAssignmentOperatorStatement(assignmentOperatorStatement: FirAssignmentOperatorStatement, data: MiniCfgData) {
                 visitElement(assignmentOperatorStatement, data)
                 val lhs = assignmentOperatorStatement.leftArgument as? FirQualifiedAccessExpression ?: return
                 if (lhs.explicitReceiver != null) return
-                data.recordAssignment(lhs.calleeReference)
+                data.recordAssignment(lhs.calleeReference, operatorAssignment = true)
             }
 
-            private fun MiniCfgData.recordAssignment(reference: FirReference) {
+            private fun MiniCfgData.recordAssignment(reference: FirReference, operatorAssignment: Boolean) {
                 val name = (reference as? FirNamedReference)?.name ?: return
                 val property = variableDeclarations.lastOrNull { name in it }?.get(name) ?: return
 
-                val assignment = Assignment()
+                val assignment = Assignment(operatorAssignment)
                 assignments.getOrPut(property) { mutableListOf() }.add(assignment)
                 flow.recordAssignment(property, assignment)
             }
@@ -492,7 +527,7 @@ internal class FirLocalVariableAssignmentAnalyzer {
                 var flow: MiniFlow = MiniFlow.start()
                 val variableDeclarations: ArrayDeque<MutableMap<Name, FirProperty>> = ArrayDeque(listOf(mutableMapOf()))
                 val assignments: MutableMap<FirProperty, MutableList<Assignment>> = mutableMapOf()
-                val forks: MutableMap<FirBasedSymbol<*>, Fork> = mutableMapOf()
+                val forks: MutableMap<Any /* FirBasedSymbol<*> | FirLoop */, Fork> = mutableMapOf()
             }
         }
     }