Retain data flow info after conditions in while-statements

compiler/frontend/src/org/jetbrains/jet/lang/types/expressions/ControlStructureTypingVisitor.java

@@ -162,19 +162,17 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
 
         ExpressionTypingContext context = contextWithExpectedType.replaceExpectedType(TypeUtils.NO_EXPECTED_TYPE);
         JetExpression condition = expression.getCondition();
-        checkCondition(context.scope, condition, context);
+        DataFlowInfo dataFlowInfo = checkCondition(context.scope, condition, context);
+
         JetExpression body = expression.getBody();
         if (body != null) {
             WritableScopeImpl scopeToExtend = newWritableScopeImpl(context, "Scope extended in while's condition");
-            DataFlowInfo conditionInfo = condition == null ? context.dataFlowInfo : DataFlowUtils.extractDataFlowInfoFromCondition(condition, true, context);
+            DataFlowInfo conditionInfo = DataFlowUtils.extractDataFlowInfoFromCondition(condition, true, context).and(dataFlowInfo);
             context.expressionTypingServices.getBlockReturnedTypeWithWritableScope(scopeToExtend, Collections.singletonList(body), CoercionStrategy.NO_COERCION, context.replaceDataFlowInfo(conditionInfo), context.trace);
         }
-        DataFlowInfo dataFlowInfo;
+
         if (!containsBreak(expression, context)) {
-            dataFlowInfo = DataFlowUtils.extractDataFlowInfoFromCondition(condition, false, context);
-        }
-        else {
-            dataFlowInfo = context.dataFlowInfo;
+            dataFlowInfo = DataFlowUtils.extractDataFlowInfoFromCondition(condition, false, context).and(dataFlowInfo);
         }
         return DataFlowUtils.checkType(KotlinBuiltIns.getInstance().getUnitType(), expression, contextWithExpectedType, dataFlowInfo);
     }

compiler/testData/diagnostics/tests/dataFlowInfoTraversal/WhileCondition.kt

@@ -0,0 +1,22 @@
+fun simpleWhile(x: Int?, var y: Int) {
+    while (x!! == y) {
+        x : Int
+        y++
+    }
+    x : Int
+}
+
+fun whileWithBreak(x: Int?, var y: Int) {
+    while (x!! == y) {
+        x : Int
+        break
+    }
+    x : Int
+}
+
+fun whileWithNoCondition(x: Int?) {
+    while (<!SYNTAX!><!>) {
+        x!!
+    }
+    <!TYPE_MISMATCH!>x<!> : Int
+}

compiler/tests/org/jetbrains/jet/checkers/JetDiagnosticsTestGenerated.java

@@ -1301,6 +1301,11 @@ public class JetDiagnosticsTestGenerated extends AbstractDiagnosticsTestWithEage
                 doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/While.kt");
             }
             
+            @TestMetadata("WhileCondition.kt")
+            public void testWhileCondition() throws Exception {
+                doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/WhileCondition.kt");
+            }
+            
         }
         
         @TestMetadata("compiler/testData/diagnostics/tests/declarationChecks")