FIR CFG: process called-in-place lambdas as loops

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

@@ -30497,6 +30497,12 @@ public class DiagnosisCompilerTestFE10TestdataTestGenerated extends AbstractDiag
                     runTest("compiler/testData/diagnostics/tests/smartCasts/varnotnull/iterations.kt");
                 }
 
+                @Test
+                @TestMetadata("leakingLambdaInCalledInPlace.kt")
+                public void testLeakingLambdaInCalledInPlace() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/smartCasts/varnotnull/leakingLambdaInCalledInPlace.kt");
+                }
+
                 @Test
                 @TestMetadata("nestedDoWhile.kt")
                 public void testNestedDoWhile() throws Exception {

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

@@ -30497,6 +30497,12 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                     runTest("compiler/testData/diagnostics/tests/smartCasts/varnotnull/iterations.kt");
                 }
 
+                @Test
+                @TestMetadata("leakingLambdaInCalledInPlace.kt")
+                public void testLeakingLambdaInCalledInPlace() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/smartCasts/varnotnull/leakingLambdaInCalledInPlace.kt");
+                }
+
                 @Test
                 @TestMetadata("nestedDoWhile.kt")
                 public void testNestedDoWhile() throws Exception {

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

@@ -30497,6 +30497,12 @@ public class FirOldFrontendDiagnosticsWithLightTreeTestGenerated extends Abstrac
                     runTest("compiler/testData/diagnostics/tests/smartCasts/varnotnull/iterations.kt");
                 }
 
+                @Test
+                @TestMetadata("leakingLambdaInCalledInPlace.kt")
+                public void testLeakingLambdaInCalledInPlace() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/smartCasts/varnotnull/leakingLambdaInCalledInPlace.kt");
+                }
+
                 @Test
                 @TestMetadata("nestedDoWhile.kt")
                 public void testNestedDoWhile() throws Exception {

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.resolve.dfa
 
 import org.jetbrains.kotlin.config.LanguageFeature
+import org.jetbrains.kotlin.contracts.description.EventOccurrencesRange
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.contracts.FirResolvedContractDescription
@@ -284,6 +285,13 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         // TODO: questionable
         postponedLambdaEnterNode?.mergeIncomingFlow()
         functionEnterNode.mergeIncomingFlow(shouldForkFlow = true)
+        when (anonymousFunction.invocationKind) {
+            EventOccurrencesRange.AT_LEAST_ONCE,
+            EventOccurrencesRange.MORE_THAN_ONCE,
+            EventOccurrencesRange.UNKNOWN, null ->
+                enterCapturingStatement(functionEnterNode, anonymousFunction)
+            else -> {}
+        }
         logicSystem.updateAllReceivers(functionEnterNode.flow)
     }
 
@@ -292,6 +300,13 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
             anonymousFunction
         )
         val (functionExitNode, postponedLambdaExitNode, graph) = graphBuilder.exitAnonymousFunction(anonymousFunction)
+        when (anonymousFunction.invocationKind) {
+            EventOccurrencesRange.AT_LEAST_ONCE,
+            EventOccurrencesRange.MORE_THAN_ONCE,
+            EventOccurrencesRange.UNKNOWN, null ->
+                exitCapturingStatement(anonymousFunction)
+            else -> {}
+        }
         // TODO: questionable
         postponedLambdaExitNode?.mergeIncomingFlow()
         functionExitNode.mergeIncomingFlow()

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.resolve.dfa
 
 import org.jetbrains.kotlin.contracts.description.EventOccurrencesRange
+import org.jetbrains.kotlin.contracts.description.isInPlace
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.referredPropertySymbol
@@ -61,6 +62,8 @@ internal class FirLocalVariableAssignmentAnalyzer(
      */
     private val ephemeralConcurrentlyAssignedLocalVariables: MutableSet<FirProperty> = mutableSetOf()
 
+    private val functionStack = mutableListOf<AssignedLocalVariables>()
+
     /** Checks whether the given access is an unstable access to a local variable at this moment. */
     fun isAccessToUnstableLocalVariable(qualifiedAccessExpression: FirQualifiedAccessExpression): Boolean {
         val property = qualifiedAccessExpression.referredPropertySymbol?.fir ?: return false
@@ -101,23 +104,27 @@ internal class FirLocalVariableAssignmentAnalyzer(
             }
         }
 
-        when ((function as? FirAnonymousFunction)?.invocationKind) {
-            EventOccurrencesRange.AT_LEAST_ONCE,
-            EventOccurrencesRange.MORE_THAN_ONCE ->
-                // 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 }
-            EventOccurrencesRange.UNKNOWN, null ->
-                // The function may not only be called repeatedly, but also stored and called later, so assignments done outside
-                // its scope after the definition might also have executed.
-                assignedLocalVariablesByFunction[function.symbol]?.all?.let { concurrentlyAssignedLocalVariables += it }
-            else -> {} // The function is called at most once so its assignments have not executed yet.
+        assignedLocalVariablesByFunction[function.symbol]?.let {
+            functionStack.add(it)
+            if (function !is FirAnonymousFunction || !function.invocationKind.isInPlace) {
+                // The function may be called twice concurrently in an SMT environment, which means any assignment it executes
+                // might in theory happen in between any check it does and a subsequent use of the variable. So if this function
+                // does any assignments, it cannot smartcast the target variables.
+                concurrentlyAssignedLocalVariables += it.insideLocalFunction
+                // The function may also stored and called later, so assignments done outside its scope after the definition
+                // might also have executed.
+                for (outerScope in functionStack) {
+                    concurrentlyAssignedLocalVariables += outerScope.outsideLocalFunction
+                }
+            }
         }
     }
 
     fun exitLocalFunction(function: FirFunction) {
         concurrentlyAssignedLocalVariablesStack.removeLast()
-        when ((function as? FirAnonymousFunction)?.invocationKind) {
-            EventOccurrencesRange.UNKNOWN, null ->
+        assignedLocalVariablesByFunction[function.symbol]?.let {
+            functionStack.popLast()
+            if (function !is FirAnonymousFunction || !function.invocationKind.isInPlace) {
                 // The function may be stored and then called later, so any access to the variables it touches
                 // is no longer smartcastable ever.
                 //
@@ -132,12 +139,10 @@ internal class FirLocalVariableAssignmentAnalyzer(
                 //        p.memberOfSomething // Bad
                 //    }
                 //   FE1.0 has the same behavior.
-                assignedLocalVariablesByFunction[function.symbol]?.insideLocalFunction?.let {
-                    for (outerScope in concurrentlyAssignedLocalVariablesStack) {
-                        outerScope += it
-                    }
+                for (outerScope in concurrentlyAssignedLocalVariablesStack) {
+                    outerScope += it.insideLocalFunction
                 }
-            else -> {} // The function is only called inline; this is handled by CFG construction by visiting the function body.
+            }
         }
     }
 
@@ -266,9 +271,7 @@ internal class FirLocalVariableAssignmentAnalyzer(
             return data.localFunctionToAssignedLocalVariables
         }
 
-        class AssignedLocalVariables(val outsideLocalFunction: Set<FirProperty>, val insideLocalFunction: Set<FirProperty>) {
-            val all get() = outsideLocalFunction + insideLocalFunction
-        }
+        class AssignedLocalVariables(val outsideLocalFunction: Set<FirProperty>, val insideLocalFunction: Set<FirProperty>)
 
         private class MiniFlow(val parents: Set<MiniFlow>) {
             val assignedLocalVariables: MutableSet<FirProperty> = mutableSetOf()
@@ -304,7 +307,7 @@ internal class FirLocalVariableAssignmentAnalyzer(
                 // Only retain local variables declared above the current scope. This way, any local variables declared inside the
                 // function will effectively be treated as distinct variables and, hence, stable (Of course, for nested lambda, things would
                 // just work because inside the lambda assigned local variables are tracked by different nodes).
-                functionFork.assignedLocalVariables.retainAll(data.variableDeclarations.flatMap { it.values })
+                functionFork.assignedLocalVariables.retainAll(data.variableDeclarations.flatMapTo(mutableSetOf()) { it.values })
                 // Create another fork for the normal execution
                 val normalExecution = currentFlow.fork()
                 data.localFunctionToAssignedLocalVariables[function.symbol] =

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

@@ -28,7 +28,7 @@ fun test() {
     var s: String? = null
     s = ""
     atLeastOnce {
-        <!SMARTCAST_IMPOSSIBLE!>s<!>.length // unstable since lambda can be called twice
+        s<!UNSAFE_CALL!>.<!>length // unstable since lambda can be called twice
         s = null
         var s2: String? = null
         s2 = ""

compiler/testData/diagnostics/tests/smartCasts/varnotnull/capturedInClosureModifiedBefore.fir.kt

@@ -30,7 +30,7 @@ fun baz(s: String?) {
             <!SMARTCAST_IMPOSSIBLE!>x<!>.hashCode()
         }
         run {
-            <!SMARTCAST_IMPOSSIBLE!>x<!>.hashCode()
+            x<!UNSAFE_CALL!>.<!>hashCode()
             x = null
         }
     }
@@ -40,7 +40,7 @@ fun gaz(s: String?) {
     var x = s
     if (x != null) {
         run {
-            <!SMARTCAST_IMPOSSIBLE!>x<!>.hashCode()
+            x<!UNSAFE_CALL!>.<!>hashCode()
             x = null
         }
         run {

compiler/testData/diagnostics/tests/smartCasts/varnotnull/leakingLambdaInCalledInPlace.kt

@@ -0,0 +1,11 @@
+// FIR_IDENTICAL
+fun main() {
+    var p: String?
+    var block: () -> Int = { 1 }
+    p = "2"
+    run {
+        block = { <!SMARTCAST_IMPOSSIBLE!>p<!>.length }
+    }
+    p = null
+    block()
+}

compiler/testData/diagnostics/tests/smartCasts/varnotnull/leakingLambdaInCalledInPlace.txt

@@ -0,0 +1,3 @@
+package
+
+public fun main(): kotlin.Unit

compiler/testData/diagnostics/tests/smartCasts/varnotnull/varCapturedInClosure.kt

@@ -8,7 +8,8 @@ public fun foo() {
         } else if (s == null) {
             return -2
         } else {
-            return <!SMARTCAST_IMPOSSIBLE!>s<!>.length // Here smartcast is possible, at least in principle
+            // Smart cast might be unsafe if function is invoked twice concurrently
+            return <!SMARTCAST_IMPOSSIBLE!>s<!>.length
         }
     }
     if (s != null) {

compiler/testData/diagnostics/tests/smartCasts/varnotnull/varCapturedInInlineClosure.fir.kt

@@ -0,0 +1,16 @@
+// See also KT-7186 and forEachSafe.kt
+// Custom `forEach` has no contract but the lambda is inline (not crossinline) so smart cast is safe
+
+inline fun IntArray.forEachIndexed( op: (i: Int, value: Int) -> Unit) {
+    for (i in 0..this.size)
+        op(i, this[i])
+}
+
+fun max(a: IntArray): Int? {
+    var maxI: Int? = null
+    a.forEachIndexed { i, value ->
+        if (maxI == null || value >= a[maxI])
+            maxI = i
+    }
+    return maxI
+}

compiler/testData/diagnostics/tests/smartCasts/varnotnull/varCapturedInInlineClosure.kt

@@ -1,7 +1,7 @@
-// FIR_IDENTICAL
-// See also KT-7186
+// See also KT-7186 and forEachSafe.kt
+// Custom `forEach` has no contract but the lambda is inline (not crossinline) so smart cast is safe
 
-fun IntArray.forEachIndexed( op: (i: Int, value: Int) -> Unit) {
+inline fun IntArray.forEachIndexed( op: (i: Int, value: Int) -> Unit) {
     for (i in 0..this.size)
         op(i, this[i])
 }

compiler/testData/diagnostics/tests/smartCasts/varnotnull/varCapturedInInlineClosure.txt

@@ -1,4 +1,4 @@
 package
 
 public fun max(/*0*/ a: kotlin.IntArray): kotlin.Int?
-public fun kotlin.IntArray.forEachIndexed(/*0*/ op: (i: kotlin.Int, value: kotlin.Int) -> kotlin.Unit): kotlin.Unit
+public inline fun kotlin.IntArray.forEachIndexed(/*0*/ op: (i: kotlin.Int, value: kotlin.Int) -> kotlin.Unit): kotlin.Unit

compiler/testData/diagnostics/testsWithStdLib/smartcasts/forEachSafe.fir.kt

@@ -0,0 +1,12 @@
+// See also KT-7186 and varCapturedInInlineClosure.kt
+// Standard library `forEach` calls lambda in-place by contract so smart cast is safe
+
+fun indexOfMax(a: IntArray): Int? {
+    var maxI: Int? = null
+    a.forEachIndexed { i, value ->
+        if (maxI == null || value >= a[maxI]) {
+            maxI = i
+        }
+    }
+    return maxI
+}

compiler/testData/diagnostics/testsWithStdLib/smartcasts/forEachSafe.kt

@@ -1,5 +1,5 @@
-// FIR_IDENTICAL
-// KT-7186: False "Type mismatch" error
+// See also KT-7186 and varCapturedInInlineClosure.kt
+// Standard library `forEach` calls lambda in-place by contract so smart cast is safe
 
 fun indexOfMax(a: IntArray): Int? {
     var maxI: Int? = null

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

@@ -30587,6 +30587,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                     runTest("compiler/testData/diagnostics/tests/smartCasts/varnotnull/iterations.kt");
                 }
 
+                @Test
+                @TestMetadata("leakingLambdaInCalledInPlace.kt")
+                public void testLeakingLambdaInCalledInPlace() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/smartCasts/varnotnull/leakingLambdaInCalledInPlace.kt");
+                }
+
                 @Test
                 @TestMetadata("nestedDoWhile.kt")
                 public void testNestedDoWhile() throws Exception {