From ae77bf0a4ed94b56a16f3a9b25307eb0faee6d7f Mon Sep 17 00:00:00 2001
From: Alexander Udalov <Alexander.Udalov@jetbrains.com>
Date: Tue, 13 Nov 2012 16:52:26 +0400
Subject: [PATCH] Retain data flow info after conditions in if-statements

---
 .../ControlStructureTypingVisitor.java        | 15 ++++++----
 .../dataFlowInfoTraversal/IfStatement.kt      | 29 +++++++++++++++++++
 .../checkers/JetDiagnosticsTestGenerated.java |  5 ++++
 3 files changed, 44 insertions(+), 5 deletions(-)
 create mode 100644 compiler/testData/diagnostics/tests/dataFlowInfoTraversal/IfStatement.kt

diff --git a/compiler/frontend/src/org/jetbrains/jet/lang/types/expressions/ControlStructureTypingVisitor.java b/compiler/frontend/src/org/jetbrains/jet/lang/types/expressions/ControlStructureTypingVisitor.java
index 510e0e0a90f..71121aac5ec 100644
--- a/compiler/frontend/src/org/jetbrains/jet/lang/types/expressions/ControlStructureTypingVisitor.java
+++ b/compiler/frontend/src/org/jetbrains/jet/lang/types/expressions/ControlStructureTypingVisitor.java
@@ -63,14 +63,19 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
         super(facade);
     }
 
-    private void checkCondition(@NotNull JetScope scope, @Nullable JetExpression condition, ExpressionTypingContext context) {
+    @NotNull
+    private DataFlowInfo checkCondition(@NotNull JetScope scope, @Nullable JetExpression condition, ExpressionTypingContext context) {
         if (condition != null) {
-            JetType conditionType = facade.getTypeInfo(condition, context.replaceScope(scope)).getType();
+            JetTypeInfo typeInfo = facade.getTypeInfo(condition, context.replaceScope(scope));
+            JetType conditionType = typeInfo.getType();
 
             if (conditionType != null && !isBoolean(conditionType)) {
                 context.trace.report(TYPE_MISMATCH_IN_CONDITION.on(condition, conditionType));
             }
+
+            return typeInfo.getDataFlowInfo();
         }
+        return context.dataFlowInfo;
     }
 
     ////////////////////////////////////////////////////////////////////////////////////////////////////
@@ -84,15 +89,15 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
     public JetTypeInfo visitIfExpression(JetIfExpression expression, ExpressionTypingContext contextWithExpectedType, boolean isStatement) {
         ExpressionTypingContext context = contextWithExpectedType.replaceExpectedType(TypeUtils.NO_EXPECTED_TYPE);
         JetExpression condition = expression.getCondition();
-        checkCondition(context.scope, condition, context);
+        DataFlowInfo conditionDataFlowInfo = checkCondition(context.scope, condition, context);
 
         JetExpression elseBranch = expression.getElse();
         JetExpression thenBranch = expression.getThen();
 
         WritableScopeImpl thenScope = newWritableScopeImpl(context, "Then scope");
         WritableScopeImpl elseScope = newWritableScopeImpl(context, "Else scope");
-        DataFlowInfo thenInfo = DataFlowUtils.extractDataFlowInfoFromCondition(condition, true, context);
-        DataFlowInfo elseInfo = DataFlowUtils.extractDataFlowInfoFromCondition(condition, false, context);
+        DataFlowInfo thenInfo = DataFlowUtils.extractDataFlowInfoFromCondition(condition, true, context).and(conditionDataFlowInfo);
+        DataFlowInfo elseInfo = DataFlowUtils.extractDataFlowInfoFromCondition(condition, false, context).and(conditionDataFlowInfo);
 
         if (elseBranch == null) {
             if (thenBranch != null) {
diff --git a/compiler/testData/diagnostics/tests/dataFlowInfoTraversal/IfStatement.kt b/compiler/testData/diagnostics/tests/dataFlowInfoTraversal/IfStatement.kt
new file mode 100644
index 00000000000..d50c887f76f
--- /dev/null
+++ b/compiler/testData/diagnostics/tests/dataFlowInfoTraversal/IfStatement.kt
@@ -0,0 +1,29 @@
+fun ifThen(x: Int?) {
+    if (x!! == 0) {
+        x : Int
+    }
+    x : Int
+}
+
+fun ifElse(x: Int?) {
+    if (x!! == 0) else {
+        x : Int
+    }
+    x : Int
+}
+
+fun ifThenElse(x: Int?) {
+    if (x!! == 0) {
+        x : Int
+    } else {
+        x : Int
+    }
+    x : Int
+}
+
+fun ifIs(x: Int?, cond: Boolean) {
+    if ((x is Int) == cond) {
+        <!TYPE_MISMATCH!>x<!> : Int
+    }
+    <!TYPE_MISMATCH!>x<!> : Int
+}
diff --git a/compiler/tests/org/jetbrains/jet/checkers/JetDiagnosticsTestGenerated.java b/compiler/tests/org/jetbrains/jet/checkers/JetDiagnosticsTestGenerated.java
index 86ef6b9add7..bdbecdf0b52 100644
--- a/compiler/tests/org/jetbrains/jet/checkers/JetDiagnosticsTestGenerated.java
+++ b/compiler/tests/org/jetbrains/jet/checkers/JetDiagnosticsTestGenerated.java
@@ -1231,6 +1231,11 @@ public class JetDiagnosticsTestGenerated extends AbstractDiagnosticsTestWithEage
                 doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/FunctionLiteral.kt");
             }
             
+            @TestMetadata("IfStatement.kt")
+            public void testIfStatement() throws Exception {
+                doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/IfStatement.kt");
+            }
+            
             @TestMetadata("IfThenElse.kt")
             public void testIfThenElse() throws Exception {
                 doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/IfThenElse.kt");