Retain data flow info after conditions in if-statements

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) {

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
+}

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");