FIR: fix some errors in local variable assignment analyzer

* wrong method was called from FirDataFlowAnalyzer.exitFunctionCall;
 * map from function to affected properties should be keyed by symbol,
   not FirFunction, as the latter may change;
 * arguments of `return` and assignment statements should be visited,
   as they may contain lambdas.

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

@@ -1056,7 +1056,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     fun exitFunctionCall(functionCall: FirFunctionCall, callCompleted: Boolean) {
         val lambdaArgs = functionCall.arguments.mapNotNull { (it as? FirAnonymousFunctionExpression)?.anonymousFunction }
         if (lambdaArgs.size > 1) {
-            getOrCreateLocalVariableAssignmentAnalyzer(lambdaArgs.first())?.enterFunctionCallWithMultipleLambdaArgs(lambdaArgs)
+            getOrCreateLocalVariableAssignmentAnalyzer(lambdaArgs.first())?.exitFunctionCallWithMultipleLambdaArgs()
         }
         if (ignoreFunctionCalls) {
             graphBuilder.exitIgnoredCall(functionCall)

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

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.fir.declarations.utils.referredPropertySymbol
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.references.FirNamedReference
 import org.jetbrains.kotlin.fir.resolve.dfa.FirLocalVariableAssignmentAnalyzer.Companion.MiniFlow.Companion.join
+import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 import org.jetbrains.kotlin.name.Name
 
@@ -23,7 +24,7 @@ import org.jetbrains.kotlin.name.Name
  *  queries after the traversal is done.
  **/
 internal class FirLocalVariableAssignmentAnalyzer(
-    private val assignedLocalVariablesByFunction: Map<FirFunction, AssignedLocalVariables>
+    private val assignedLocalVariablesByFunction: Map<FirFunctionSymbol<*>, AssignedLocalVariables>
 ) {
     /**
      * Stack storing concurrent lambda arguments for the current visited anonymous function. For example
@@ -70,7 +71,7 @@ internal class FirLocalVariableAssignmentAnalyzer(
         // always null and hence there is no need to check it. In addition, since multiple lambda can be passed, we accumulate the
         // effects by appending to `ephemeralConcurrentlyAssignedLocalVariables`.  After the function call is resolved,
         // `exitAnonymousFunction` will be invoked at some point to properly set up the `persistentConcurrentlyAssignedLocalVariables`.
-        assignedLocalVariablesByFunction[anonymousFunction]?.insideLocalFunction?.let {
+        assignedLocalVariablesByFunction[anonymousFunction.symbol]?.insideLocalFunction?.let {
             ephemeralConcurrentlyAssignedLocalVariables.addAll(it)
         }
     }
@@ -81,30 +82,32 @@ internal class FirLocalVariableAssignmentAnalyzer(
     }
 
     fun enterLocalFunction(function: FirFunction) {
-        val eventOccurrencesRange: EventOccurrencesRange? = (function as? FirAnonymousFunction)?.invocationKind
-        // carry on concurrently modified variables from the current scope
         val concurrentlyAssignedLocalVariables = concurrentlyAssignedLocalVariablesStack.last().toMutableSet()
         concurrentlyAssignedLocalVariablesStack.add(concurrentlyAssignedLocalVariables)
+
         val concurrentLambdasInCurrentCall = concurrentLambdaArgsStack.lastOrNull()
         if (concurrentLambdasInCurrentCall != null && function in concurrentLambdasInCurrentCall) {
-            concurrentLambdasInCurrentCall.filter { it != function }.forEach { otherLambda ->
+            for (otherLambda in concurrentLambdasInCurrentCall) {
-                assignedLocalVariablesByFunction[otherLambda]?.insideLocalFunction?.let {
+                // As mentioned in the comment above, we don't know whether other lambda arguments passed to the same call will be
-                    concurrentlyAssignedLocalVariables += it
+                // called before or after this lambda, so their assignments might have executed. Unless they're not called at all.
+                if (otherLambda != function && otherLambda.invocationKind != EventOccurrencesRange.ZERO) {
+                    assignedLocalVariablesByFunction[otherLambda.symbol]?.insideLocalFunction?.let {
+                        concurrentlyAssignedLocalVariables += it
+                    }
                 }
             }
         }
-        when (eventOccurrencesRange) {
+
+        when ((function as? FirAnonymousFunction)?.invocationKind) {
             EventOccurrencesRange.AT_LEAST_ONCE,
-            EventOccurrencesRange.MORE_THAN_ONCE -> assignedLocalVariablesByFunction[function]?.insideLocalFunction?.let {
+            EventOccurrencesRange.MORE_THAN_ONCE ->
-                concurrentlyAssignedLocalVariables += it
+                // The function may be called repeatedly so the assignments may have already executed before we enter it again.
-            }
+                assignedLocalVariablesByFunction[function.symbol]?.insideLocalFunction?.let { concurrentlyAssignedLocalVariables += it }
-            // Add both inside and outside since this local function may be invoked multiple times concurrently.
+            EventOccurrencesRange.UNKNOWN, null ->
-            EventOccurrencesRange.UNKNOWN, null -> assignedLocalVariablesByFunction[function]?.all?.let {
+                // The function may not only be called repeatedly, but also stored and called later, so assignments done outside
-                concurrentlyAssignedLocalVariables += it
+                // its scope after the definition might also have executed.
-            }
+                assignedLocalVariablesByFunction[function.symbol]?.all?.let { concurrentlyAssignedLocalVariables += it }
-            else -> {
+            else -> {} // The function is called at most once so its assignments have not executed yet.
-                // no additional stuff to do for other cases
-            }
         }
     }
 
@@ -112,12 +115,13 @@ internal class FirLocalVariableAssignmentAnalyzer(
         val eventOccurrencesRange: EventOccurrencesRange? = (function as? FirAnonymousFunction)?.invocationKind
         concurrentlyAssignedLocalVariablesStack.removeLast()
         when (eventOccurrencesRange) {
-            EventOccurrencesRange.UNKNOWN, null -> assignedLocalVariablesByFunction[function]?.insideLocalFunction?.let {
+            EventOccurrencesRange.UNKNOWN, null ->
-                concurrentlyAssignedLocalVariablesStack.last() += it
+                // The function may be stored and then called later, so any access to the variables it touches
-            }
+                // is no longer smartcastable ever.
-            else -> {
+                assignedLocalVariablesByFunction[function.symbol]?.insideLocalFunction?.let {
-                // no additional stuff to do for other cases
+                    concurrentlyAssignedLocalVariablesStack.last() += it
-            }
+                }
+            else -> {} // The function is only called inline; this is handled by CFG construction by visiting the function body.
         }
     }
 
@@ -217,7 +221,7 @@ internal class FirLocalVariableAssignmentAnalyzer(
         /**
          * Computes a mini CFG and returns the map tracking assigned local variables at each potentially concurrent local/lambda function.
          */
-        private fun computeAssignedLocalVariables(firFunction: FirFunction): Map<FirFunction, AssignedLocalVariables> {
+        private fun computeAssignedLocalVariables(firFunction: FirFunction): Map<FirFunctionSymbol<*>, AssignedLocalVariables> {
             val startFlow = MiniFlow.start()
             val data = MiniCfgBuilder.MiniCfgData(startFlow)
             MiniCfgBuilder().visitElement(firFunction, data)
@@ -265,7 +269,7 @@ internal class FirLocalVariableAssignmentAnalyzer(
                 functionFork.assignedLocalVariables.retainAll(data.variableDeclarations.flatMap { it.values })
                 // Create another fork for the normal execution
                 val normalExecution = currentFlow.fork()
-                data.localFunctionToAssignedLocalVariables[function] =
+                data.localFunctionToAssignedLocalVariables[function.symbol] =
                     AssignedLocalVariables(normalExecution.assignedLocalVariables, functionFork.assignedLocalVariables)
                 data.flow = normalExecution
             }
@@ -292,11 +296,11 @@ internal class FirLocalVariableAssignmentAnalyzer(
             }
 
             override fun visitReturnExpression(returnExpression: FirReturnExpression, data: MiniCfgData) {
+                super.visitReturnExpression(returnExpression, data)
                 // TODO: consider to also handle `throw`, which would require keeping track of all `try`, `catch` and `finally` constructs.
                 data.flow = null
             }
 
-            @OptIn(ExperimentalStdlibApi::class)
             override fun visitFunctionCall(functionCall: FirFunctionCall, data: MiniCfgData) {
                 val visitor = this
                 with(functionCall) {
@@ -316,18 +320,21 @@ internal class FirLocalVariableAssignmentAnalyzer(
             }
 
             override fun visitProperty(property: FirProperty, data: MiniCfgData) {
+                super.visitProperty(property, data)
                 if (property.isLocal) {
                     data.variableDeclarations.last()[property.name] = property
                 }
             }
 
             override fun visitVariableAssignment(variableAssignment: FirVariableAssignment, data: MiniCfgData) {
+                super.visitVariableAssignment(variableAssignment, data)
                 val flow = data.flow ?: return
                 val name = (variableAssignment.lValue as? FirNamedReference)?.name ?: return
                 flow.recordAssignment(name, data)
             }
 
             override fun visitAssignmentOperatorStatement(assignmentOperatorStatement: FirAssignmentOperatorStatement, data: MiniCfgData) {
+                super.visitAssignmentOperatorStatement(assignmentOperatorStatement, data)
                 val flow = data.flow ?: return
                 val lhs = assignmentOperatorStatement.leftArgument as? FirQualifiedAccessExpression ?: return
                 if (lhs.explicitReceiver != null) return
@@ -350,7 +357,7 @@ internal class FirLocalVariableAssignmentAnalyzer(
 
             class MiniCfgData(var flow: MiniFlow?) {
                 val variableDeclarations: ArrayDeque<MutableMap<Name, FirProperty>> = ArrayDeque(listOf(mutableMapOf()))
-                val localFunctionToAssignedLocalVariables: MutableMap<FirFunction, AssignedLocalVariables> = mutableMapOf()
+                val localFunctionToAssignedLocalVariables: MutableMap<FirFunctionSymbol<*>, AssignedLocalVariables> = mutableMapOf()
                 fun resolveLocalVariable(name: Name): FirProperty? {
                     return variableDeclarations.asReversed().firstNotNullOfOrNull { it[name] }
                 }

compiler/testData/diagnostics/tests/smartCasts/kt30826.fir.kt

@@ -1,18 +0,0 @@
-// Issue: KT-30826
-
-interface I1
-interface I2 {
-    fun foo() {}
-}
-
-class A : I1, I2
-
-fun foo(x: I1?) {
-    var y = x
-    y as I2
-    val bar = {
-        y.foo() // NPE in NI, smartcast impossible in OI
-    }
-    y = null
-    bar()
-}

compiler/testData/diagnostics/tests/smartCasts/kt30826.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // Issue: KT-30826
 
 interface I1

compiler/tests-spec/testData/diagnostics/notLinked/dfa/neg/18.fir.kt

@@ -1,69 +0,0 @@
-// !LANGUAGE: +NewInference
-// !DIAGNOSTICS: -UNUSED_EXPRESSION
-// SKIP_TXT
-
-/*
- * TESTCASE NUMBER: 1
- * ISSUES: KT-30826
- */
-fun case_1(x: Interface1?) {
-    var y = x
-    y as Interface2
-    val foo = {
-        <!DEBUG_INFO_EXPRESSION_TYPE("Interface1? & Interface1 & Interface2")!>y<!>.itest2()
-    }
-    y = null
-    foo()
-}
-
-/*
- * TESTCASE NUMBER: 2
- * ISSUES: KT-30826
- */
-fun case_2(x: Interface1?) {
-    var y = x
-    y as Interface2
-    val foo = fun () {
-        <!DEBUG_INFO_EXPRESSION_TYPE("Interface1? & Interface1 & Interface2")!>y<!>.itest2()
-    }
-    y = null
-    foo()
-}
-
-/*
- * TESTCASE NUMBER: 3
- * ISSUES: KT-30826
- */
-fun case_3(x: Interface1?) {
-    var y = x
-    y as Interface2
-    fun foo() {
-        <!DEBUG_INFO_EXPRESSION_TYPE("Interface1?"), SMARTCAST_IMPOSSIBLE!>y<!>.itest2()
-    }
-    y = null
-    foo()
-}
-
-// TESTCASE NUMBER: 4
-fun case_4(x: Interface1?) {
-    var y = x
-    y as Interface2
-    y = null
-    fun foo() {
-        <!DEBUG_INFO_EXPRESSION_TYPE("Interface1?")!>y<!>.<!UNRESOLVED_REFERENCE!>itest2<!>()
-    }
-    y = x
-    foo()
-}
-
-// TESTCASE NUMBER: 5
-fun case_5(x: Interface1?) {
-    var y = x
-    y as Interface2
-    y = null
-    fun foo() {
-        <!DEBUG_INFO_EXPRESSION_TYPE("Interface1?")!>y<!>.<!UNRESOLVED_REFERENCE!>itest2<!>()
-    }
-    y = x
-    foo()
-}

compiler/tests-spec/testData/diagnostics/notLinked/dfa/neg/18.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !LANGUAGE: +NewInference
 // !DIAGNOSTICS: -UNUSED_EXPRESSION
 // SKIP_TXT

compiler/tests-spec/testData/diagnostics/notLinked/dfa/pos/53.fir.kt

@@ -20,8 +20,8 @@ fun case_2() {
     var x: Int? = 10
     var y = { x = null }
     if (x != null) {
-        val z = case_2(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int? & kotlin.Int")!>x<!>)
+        val z = case_2(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int?")!>x<!>)
-        <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.String")!>z<!>
+        <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int")!>z<!>
     }
 }
 
@@ -47,7 +47,7 @@ fun case_4(x: Int) = ""
 fun case_4(x: Int?) = 10
 fun case_4(y: Case4) {
     if (y.x != null) {
-        val z = case_4(y.x)
+        val z = case_4(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int?")!>y.x<!>)
         <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int")!>z<!>
     }
 }
@@ -61,7 +61,7 @@ fun case_5(x: Int) = ""
 fun case_5(x: Int?) = 10
 fun case_5(y: Case4) {
     if (y.x != null) {
-        val z = case_5(y.x)
+        val z = case_5(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int?")!>y.x<!>)
         <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int")!>z<!>
     }
 }
@@ -75,7 +75,7 @@ fun case_6(x: Int) = ""
 fun case_6(x: Int?) = 10
 fun case_6(y: Case4) {
     if (y.x != null) {
-        val z = case_6(y.x)
+        val z = case_6(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int?")!>y.x<!>)
         <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int")!>z<!>
     }
 }