From 1aae58623848f6deecea35eb908e2be8646d5fd0 Mon Sep 17 00:00:00 2001
From: pyos <pyos@google.com>
Date: Mon, 12 Dec 2022 19:40:53 +0100
Subject: [PATCH] FIR DFA: don't erase statements when entering
 non-call-in-place lambda

Instead, rely on the variable assignment analyzer to properly restrict
smart casts. This makes error messages more consistent, but otherwise
should have no effect.
---
 .../fir/resolve/dfa/FirDataFlowAnalyzer.kt    | 22 +++++++++----------
 .../capturedInClosureModifiedBefore.fir.kt    |  4 ++--
 2 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt
index f88d346878a..7664dcf8bdc 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt
@@ -154,9 +154,8 @@ abstract class FirDataFlowAnalyzer(
         }
         localFunctionNode?.mergeIncomingFlow()
         functionEnterNode.mergeIncomingFlow {
-            // TODO: ||?
-            if (function is FirAnonymousFunction && function.invocationKind?.canBeRevisited() != false) {
-                enterCapturingStatement(it, function)
+            if (function is FirAnonymousFunction && function.invocationKind?.canBeRevisited() == true) {
+                enterRepeatableStatement(it, function)
             }
         }
         context.variableAssignmentAnalyzer.enterFunction(function)
@@ -166,9 +165,8 @@ abstract class FirDataFlowAnalyzer(
         if (function is FirDefaultPropertyAccessor) return null
 
         context.variableAssignmentAnalyzer.exitFunction()
-        // TODO: ||?
-        if (function is FirAnonymousFunction && function.invocationKind?.canBeRevisited() != false) {
-            exitCapturingStatement(function)
+        if (function is FirAnonymousFunction && function.invocationKind?.canBeRevisited() == true) {
+            exitRepeatableStatement(function)
         }
 
         if (function is FirAnonymousFunction) {
@@ -596,7 +594,7 @@ abstract class FirDataFlowAnalyzer(
     fun enterWhileLoop(loop: FirLoop) {
         val (loopEnterNode, loopConditionEnterNode) = graphBuilder.enterWhileLoop(loop)
         loopEnterNode.mergeIncomingFlow()
-        loopConditionEnterNode.mergeIncomingFlow { flow -> enterCapturingStatement(flow, loop) }
+        loopConditionEnterNode.mergeIncomingFlow { flow -> enterRepeatableStatement(flow, loop) }
     }
 
     fun exitWhileLoopCondition(loop: FirLoop) {
@@ -620,7 +618,7 @@ abstract class FirDataFlowAnalyzer(
     }
 
     private fun processWhileLoopExit(flow: MutableFlow, node: LoopExitNode, conditionEnterNode: LoopConditionEnterNode) {
-        val possiblyChangedVariables = exitCapturingStatement(node.fir)
+        val possiblyChangedVariables = exitRepeatableStatement(node.fir)
         if (possiblyChangedVariables.isNullOrEmpty()) 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`
@@ -653,7 +651,7 @@ abstract class FirDataFlowAnalyzer(
         }
     }
 
-    private fun enterCapturingStatement(flow: MutableFlow, statement: FirStatement) {
+    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
@@ -670,7 +668,7 @@ abstract class FirDataFlowAnalyzer(
         context.variablesClearedBeforeLoop.push(possiblyChangedVariables)
     }
 
-    private fun exitCapturingStatement(statement: FirStatement): List<RealVariable>? {
+    private fun exitRepeatableStatement(statement: FirStatement): List<RealVariable>? {
         if (context.preliminaryLoopVisitor.exitCapturingStatement(statement).isEmpty()) return null
         return context.variablesClearedBeforeLoop.pop()
     }
@@ -679,7 +677,7 @@ abstract class FirDataFlowAnalyzer(
 
     fun enterDoWhileLoop(loop: FirLoop) {
         val (loopEnterNode, loopBlockEnterNode) = graphBuilder.enterDoWhileLoop(loop)
-        loopEnterNode.mergeIncomingFlow { flow -> enterCapturingStatement(flow, loop) }
+        loopEnterNode.mergeIncomingFlow { flow -> enterRepeatableStatement(flow, loop) }
         loopBlockEnterNode.mergeIncomingFlow()
     }
 
@@ -693,7 +691,7 @@ abstract class FirDataFlowAnalyzer(
         val (loopConditionExitNode, loopExitNode) = graphBuilder.exitDoWhileLoop(loop)
         loopConditionExitNode.mergeIncomingFlow()
         loopExitNode.mergeIncomingFlow { processLoopExit(it, loopExitNode, loopConditionExitNode) }
-        exitCapturingStatement(loop)
+        exitRepeatableStatement(loop)
     }
 
     // ----------------------------------- Try-catch-finally -----------------------------------
diff --git a/compiler/testData/diagnostics/tests/smartCasts/varnotnull/capturedInClosureModifiedBefore.fir.kt b/compiler/testData/diagnostics/tests/smartCasts/varnotnull/capturedInClosureModifiedBefore.fir.kt
index f4041257f4c..4475e956293 100644
--- a/compiler/testData/diagnostics/tests/smartCasts/varnotnull/capturedInClosureModifiedBefore.fir.kt
+++ b/compiler/testData/diagnostics/tests/smartCasts/varnotnull/capturedInClosureModifiedBefore.fir.kt
@@ -30,7 +30,7 @@ fun baz(s: String?) {
             <!SMARTCAST_IMPOSSIBLE!>x<!>.hashCode()
         }
         run {
-            x<!UNSAFE_CALL!>.<!>hashCode()
+            <!SMARTCAST_IMPOSSIBLE!>x<!>.hashCode()
             x = null
         }
     }
@@ -40,7 +40,7 @@ fun gaz(s: String?) {
     var x = s
     if (x != null) {
         run {
-            x<!UNSAFE_CALL!>.<!>hashCode()
+            <!SMARTCAST_IMPOSSIBLE!>x<!>.hashCode()
             x = null
         }
         run {