[FIR] Pass implications from subject of safe call in DFA

^KT-59689 Fixed

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

@@ -39195,6 +39195,12 @@ public class DiagnosticCompilerTestFE10TestdataTestGenerated extends AbstractDia
                     runTest("compiler/testData/diagnostics/testsWithStdLib/contracts/smartcasts/reifiedGeneric.kt");
                 }
 
+                @Test
+                @TestMetadata("returnsImpliesAndSafeCalls.kt")
+                public void testReturnsImpliesAndSafeCalls() throws Exception {
+                    runTest("compiler/testData/diagnostics/testsWithStdLib/contracts/smartcasts/returnsImpliesAndSafeCalls.kt");
+                }
+
                 @Test
                 @TestMetadata("safecallAndReturnsNull.kt")
                 public void testSafecallAndReturnsNull() throws Exception {

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

@@ -39195,6 +39195,12 @@ public class LLFirPreresolvedReversedDiagnosticCompilerFE10TestDataTestGenerated
                     runTest("compiler/testData/diagnostics/testsWithStdLib/contracts/smartcasts/reifiedGeneric.kt");
                 }
 
+                @Test
+                @TestMetadata("returnsImpliesAndSafeCalls.kt")
+                public void testReturnsImpliesAndSafeCalls() throws Exception {
+                    runTest("compiler/testData/diagnostics/testsWithStdLib/contracts/smartcasts/returnsImpliesAndSafeCalls.kt");
+                }
+
                 @Test
                 @TestMetadata("safecallAndReturnsNull.kt")
                 public void testSafecallAndReturnsNull() throws Exception {

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

@@ -39195,6 +39195,12 @@ public class FirLightTreeOldFrontendDiagnosticsTestGenerated extends AbstractFir
                     runTest("compiler/testData/diagnostics/testsWithStdLib/contracts/smartcasts/reifiedGeneric.kt");
                 }
 
+                @Test
+                @TestMetadata("returnsImpliesAndSafeCalls.kt")
+                public void testReturnsImpliesAndSafeCalls() throws Exception {
+                    runTest("compiler/testData/diagnostics/testsWithStdLib/contracts/smartcasts/returnsImpliesAndSafeCalls.kt");
+                }
+
                 @Test
                 @TestMetadata("safecallAndReturnsNull.kt")
                 public void testSafecallAndReturnsNull() throws Exception {

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

@@ -39309,6 +39309,12 @@ public class FirPsiOldFrontendDiagnosticsTestGenerated extends AbstractFirPsiDia
                     runTest("compiler/testData/diagnostics/testsWithStdLib/contracts/smartcasts/reifiedGeneric.kt");
                 }
 
+                @Test
+                @TestMetadata("returnsImpliesAndSafeCalls.kt")
+                public void testReturnsImpliesAndSafeCalls() throws Exception {
+                    runTest("compiler/testData/diagnostics/testsWithStdLib/contracts/smartcasts/returnsImpliesAndSafeCalls.kt");
+                }
+
                 @Test
                 @TestMetadata("safecallAndReturnsNull.kt")
                 public void testSafecallAndReturnsNull() throws Exception {

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

@@ -822,10 +822,41 @@ abstract class FirDataFlowAnalyzer(
             // Otherwise if the result is non-null, then `b` executed, which implies `a` is not null
             // and every statement from `b` holds.
             val expressionVariable = variableStorage.getOrCreate(flow, safeCall)
-            // TODO? all new implications in previous node's flow are valid here if receiver != null
-            //  (that requires a second level of implications: receiver != null => condition => effect).
-            //  KT-59689
-            flow.addAllConditionally(expressionVariable notEq null, node.lastPreviousNode.getFlow(path))
+            val previousFlow = node.lastPreviousNode.getFlow(path)
+
+            flow.addAllConditionally(expressionVariable notEq null, previousFlow)
+
+            /*
+             * If we have some implication about rhs of safe call in the previous flow, then we can expand them to the whole
+             *   safe call variable
+             *
+             * a?.foo() // original call
+             * subj.foo() // rhs of safe call
+             *
+             * previousFlow:
+             *  - subj.foo() == True -> X_1
+             *  - subj.foo() == False -> X_2
+             *  - subj.foo() != Null -> X_3
+             *  - subj.foo() == Null -> X_4
+             *
+             * flow:
+             *  - a?.foo() == True -> X_1
+             *  - a?.foo() == False -> X_2
+             *  - a?.foo() != Null -> X_3
+             *
+             * Note that we don't pass implication with 'subj.foo() == Null' in the condition because there are two different ways
+             *   why `a?.foo()` may be `null` -- it's either `a` is `null` or `subj.foo()` is `null`, and we can't differentiate between
+             *   them
+             *
+             * Also, an implementation note: in the following lines we use `expressionVariable` made on safe call expression when looking
+             *   for implications from previous flow in the subject, because VariableStorage doesn't differ between the whole safe call
+             *   and synthetically generated selector, see [variableStorage.getOrCreate] implementation
+             */
+            previousFlow.implications[expressionVariable]?.forEach {
+                if (it.condition.operation != Operation.EqNull) {
+                    flow.addImplication(it)
+                }
+            }
         }
     }
 
@@ -1073,8 +1104,6 @@ abstract class FirDataFlowAnalyzer(
             // If `left && right` is true, then both are evaluated to true. If `left || right` is false, then both are false.
             // Approved type statements for RHS already contain everything implied by the corresponding value of LHS.
             val bothEvaluated = operatorVariable eq isAnd
-            // TODO? `bothEvaluated` also implies all implications from RHS. This requires a second level
-            //  of implications, which the logic system currently doesn't support. See also safe calls. KT-59689
             flow.addAllConditionally(bothEvaluated, flowFromRight)
             if (rightIsBoolean) {
                 flow.addAllConditionally(bothEvaluated, logicSystem.approveOperationStatement(flowFromRight, rightVariable!! eq isAnd))

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/Flow.kt

@@ -19,7 +19,7 @@ abstract class Flow {
 class PersistentFlow internal constructor(
     private val previousFlow: PersistentFlow?,
     private val approvedTypeStatements: PersistentMap<RealVariable, PersistentTypeStatement>,
-    internal val implications: PersistentMap<DataFlowVariable, PersistentList<Implication>>,
+    val implications: PersistentMap<DataFlowVariable, PersistentList<Implication>>,
     // RealVariable describes a storage in memory; a pair of RealVariable with its assignment
     // index at a particular execution point forms an SSA value corresponding to the result of
     // an initializer.

compiler/testData/diagnostics/testsWithStdLib/contracts/smartcasts/returnsImpliesAndSafeCalls.kt

@@ -0,0 +1,71 @@
+// ISSUE: KT-59689
+// FIR_IDENTICAL
+// DIAGNOSTICS: -DEBUG_INFO_SMARTCAST
+// OPT_IN: kotlin.contracts.ExperimentalContracts
+
+import kotlin.contracts.ExperimentalContracts
+import kotlin.contracts.contract
+
+class A {
+    fun returnsTrue(x: Any?): Boolean {
+        contract { returns(true) implies (x is String) }
+        return x is String
+    }
+
+    fun returnsFalse(x: Any?): Boolean {
+        contract { returns(false) implies (x is String) }
+        return x !is String
+    }
+
+    fun returnsNotNull(x: Any?): Any? {
+        contract { returnsNotNull() implies (x is String) }
+        return when (x) {
+            is String -> ""
+            else -> null
+        }
+    }
+
+    fun returnsNull(x: Any?): Any? {
+        contract { returns(null) implies (x is String) }
+        return when (x) {
+            is String -> null
+            else -> ""
+        }
+    }
+}
+
+fun test_returnsTrue(a: A?, s: Any?, b: Boolean) {
+    if (a?.returnsTrue(s) == true) {
+        s.length
+    }
+    if (a?.returnsTrue(s) == true && b) {
+        s.length
+    }
+}
+
+fun test_returnsFalse(a: A?, s: Any?, b: Boolean) {
+    if (a?.returnsFalse(s) == false) {
+        s.length
+    }
+    if (a?.returnsFalse(s) == false && b) {
+        s.length
+    }
+}
+
+fun test_returnsNotNull(a: A?, s: Any?, b: Boolean) {
+    if (a?.returnsNotNull(s) != null) {
+        s.length
+    }
+    if (a?.returnsNotNull(s) != null && b) {
+        s.length
+    }
+}
+
+fun test_returnsNull(a: A?, s: Any?, b: Boolean) {
+    if (a?.returnsNull(s) == null) {
+        s.<!UNRESOLVED_REFERENCE!>length<!> // should be an error
+    }
+    if (a?.returnsNull(s) == null && b) {
+        s.<!UNRESOLVED_REFERENCE!>length<!> // should be an error
+    }
+}

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

@@ -41231,6 +41231,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                     runTest("compiler/testData/diagnostics/testsWithStdLib/contracts/smartcasts/reifiedGeneric.kt");
                 }
 
+                @Test
+                @TestMetadata("returnsImpliesAndSafeCalls.kt")
+                public void testReturnsImpliesAndSafeCalls() throws Exception {
+                    runTest("compiler/testData/diagnostics/testsWithStdLib/contracts/smartcasts/returnsImpliesAndSafeCalls.kt");
+                }
+
                 @Test
                 @TestMetadata("safecallAndReturnsNull.kt")
                 public void testSafecallAndReturnsNull() throws Exception {