From 1ff968eca2157dc104274535fcfd8f7217587136 Mon Sep 17 00:00:00 2001
From: pyos <pyos@google.com>
Date: Tue, 8 Nov 2022 15:53:23 +0100
Subject: [PATCH] FIR DFA: extract `approveStatementsInsideFlow(variable notEq
 null)`

I'm not sure how this fixes a test. Magic!
---
 .../fir/resolve/dfa/FirDataFlowAnalyzer.kt    | 48 +++++--------------
 .../loops/safeCallBreakInsideDoWhile.fir.kt   |  2 +-
 2 files changed, 13 insertions(+), 37 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 9208c511dc1..7cfc073edd8 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
@@ -447,12 +447,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
                     flow.addTypeStatement(operandVariable typeEq type)
                 }
                 if (!type.canBeNull) {
-                    logicSystem.approveStatementsInsideFlow(
-                        flow,
-                        operandVariable notEq null,
-                        shouldRemoveSynthetics = true,
-                        shouldForkFlow = false
-                    )
+                    flow.assumeNotNull(operandVariable, shouldForkFlow = false, shouldRemoveSynthetics = true)
                 } else {
                     val expressionVariable = variableStorage.createSyntheticVariable(typeOperatorCall)
                     flow.addImplication((expressionVariable notEq null) implies (operandVariable notEq null))
@@ -525,9 +520,6 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         val operandVariable = variableStorage.getOrCreateVariable(node.previousFlow, operand)
         // expression == const -> expression != null
         flow.addImplication((expressionVariable eq isEq) implies (operandVariable notEq null))
-        if (operandVariable.isReal()) {
-            flow.addImplication((expressionVariable eq isEq) implies (operandVariable typeEq any))
-        }
 
         // propagating facts for (... == true) and (... == false)
         when (const.kind) {
@@ -691,17 +683,11 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         unionNode?.let { unionFlowFromArguments(it) }
     }
 
-    fun FirExpression.propagateNotNullInfo(node: CFGNode<*>) {
+    private fun FirExpression.propagateNotNullInfo(node: CFGNode<*>) {
         val symbol = this.symbol
         if (symbol != null) {
             variableStorage.getOrCreateRealVariable(node.previousFlow, symbol, this)?.let { operandVariable ->
-                node.flow.addTypeStatement(operandVariable typeEq any)
-                logicSystem.approveStatementsInsideFlow(
-                    node.flow,
-                    operandVariable notEq null,
-                    shouldRemoveSynthetics = true,
-                    shouldForkFlow = false
-                )
+                node.flow.assumeNotNull(operandVariable, shouldForkFlow = false, shouldRemoveSynthetics = true)
             }
         }
         when (this) {
@@ -714,6 +700,13 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         }
     }
 
+    private fun FLOW.assumeNotNull(variable: DataFlowVariable, shouldForkFlow: Boolean, shouldRemoveSynthetics: Boolean): FLOW =
+        logicSystem.approveStatementsInsideFlow(this, variable notEq null, shouldForkFlow, shouldRemoveSynthetics,).also {
+            if (variable is RealVariable) {
+                it.addTypeStatement(variable typeEq any andTypeNotEq nullableNothing)
+            }
+        }
+
     // ----------------------------------- When -----------------------------------
 
     fun enterWhenExpression(whenExpression: FirWhenExpression) {
@@ -943,12 +936,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
                 }
                 flow.addTypeStatement(variable typeEq type)
             }
-            flow = logicSystem.approveStatementsInsideFlow(
-                flow,
-                variable notEq null,
-                shouldFork,
-                shouldRemoveSynthetics = false
-            )
+            flow = flow.assumeNotNull(variable, shouldFork, shouldRemoveSynthetics = false)
         }
 
         node.flow = flow
@@ -966,9 +954,6 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
             is SyntheticVariable -> variableStorage.getOrCreateVariable(flow, safeCall.receiver)
         }
         flow.addImplication((variable notEq null) implies (receiverVariable notEq null))
-        if (receiverVariable.isReal()) {
-            flow.addImplication((variable notEq null) implies (receiverVariable typeEq any))
-        }
     }
 
     fun exitResolvedQualifierNode(resolvedQualifier: FirResolvedQualifier) {
@@ -1379,16 +1364,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         lhsExitNode.mergeIncomingFlow()
         val flow = lhsExitNode.flow
         val lhsVariable = variableStorage.getOrCreateVariable(flow, elvisExpression.lhs)
-        lhsIsNotNullNode.flow = logicSystem.approveStatementsInsideFlow(
-            flow,
-            lhsVariable notEq null,
-            shouldForkFlow = true,
-            shouldRemoveSynthetics = false
-        ).also {
-            if (lhsVariable.isReal()) {
-                it.addTypeStatement(lhsVariable typeEq any)
-            }
-        }
+        lhsIsNotNullNode.flow = flow.assumeNotNull(lhsVariable, shouldForkFlow = true, shouldRemoveSynthetics = false)
         rhsEnterNode.flow = logicSystem.approveStatementsInsideFlow(
             flow,
             lhsVariable eq null,
diff --git a/compiler/testData/diagnostics/tests/smartCasts/loops/safeCallBreakInsideDoWhile.fir.kt b/compiler/testData/diagnostics/tests/smartCasts/loops/safeCallBreakInsideDoWhile.fir.kt
index c4a3c5e1d0b..2c3d36fffa2 100644
--- a/compiler/testData/diagnostics/tests/smartCasts/loops/safeCallBreakInsideDoWhile.fir.kt
+++ b/compiler/testData/diagnostics/tests/smartCasts/loops/safeCallBreakInsideDoWhile.fir.kt
@@ -8,7 +8,7 @@ fun calc(x: String?, y: String?): Int {
         // x is not null in condition but we do not see it yet
     } while (x.length > 0)
     // y is nullable because of break
-    y.length
+    y<!UNSAFE_CALL!>.<!>length
     // x is not null, at least in theory
     return x.length
 }