Retain data flow info after when-expressions

This means the subject expression and the common data flow info of all entries'
conditions

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

@@ -1144,8 +1144,9 @@ public class BasicExpressionTypingVisitor extends ExpressionTypingVisitor {
                     result = ErrorUtils.createErrorType("No right argument"); // TODO
                     return JetTypeInfo.create(null, dataFlowInfo);
                 }
-                dataFlowInfo = checkInExpression(expression, expression.getOperationReference(), left, right, context).getSecond();
-                result = booleanType;
+                JetTypeInfo typeInfo = checkInExpression(expression, expression.getOperationReference(), left, right, context);
+                dataFlowInfo = typeInfo.getDataFlowInfo();
+                result = typeInfo.getType();
             }
             else if (OperatorConventions.BOOLEAN_OPERATIONS.containsKey(operationType)) {
                 JetTypeInfo leftTypeInfo = facade.getTypeInfo(left, context);
@@ -1197,7 +1198,7 @@ public class BasicExpressionTypingVisitor extends ExpressionTypingVisitor {
     }
 
     @NotNull
-    public Pair<Boolean, DataFlowInfo> checkInExpression(
+    public JetTypeInfo checkInExpression(
             JetElement callElement,
             @NotNull JetSimpleNameExpression operationSign,
             @Nullable JetExpression left,
@@ -1219,10 +1220,10 @@ public class BasicExpressionTypingVisitor extends ExpressionTypingVisitor {
         ensureBooleanResult(operationSign, name, containsType, context);
 
         if (left != null) {
-            dataFlowInfo = facade.getTypeInfo(left, contextWithDataFlow).getDataFlowInfo();
+            dataFlowInfo = facade.getTypeInfo(left, contextWithDataFlow).getDataFlowInfo().and(dataFlowInfo);
         }
 
-        return Pair.create(resolutionResult.isSuccess(), dataFlowInfo);
+        return JetTypeInfo.create(resolutionResult.isSuccess() ? KotlinBuiltIns.getInstance().getBooleanType() : null, dataFlowInfo);
     }
 
     private void ensureNonemptyIntersectionOfOperandTypes(JetBinaryExpression expression, ExpressionTypingContext context) {

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

@@ -33,7 +33,8 @@ import org.jetbrains.jet.lang.types.JetTypeInfo;
     @NotNull
     JetTypeInfo getSelectorReturnTypeInfo(@NotNull ReceiverValue receiver, @Nullable ASTNode callOperationNode, @NotNull JetExpression selectorExpression, @NotNull ExpressionTypingContext context);
 
-    boolean checkInExpression(JetElement callElement, @NotNull JetSimpleNameExpression operationSign, @Nullable JetExpression left, @NotNull JetExpression right, ExpressionTypingContext context);
+    @NotNull
+    JetTypeInfo checkInExpression(JetElement callElement, @NotNull JetSimpleNameExpression operationSign, @Nullable JetExpression left, @NotNull JetExpression right, ExpressionTypingContext context);
 
     void checkStatementType(@NotNull JetExpression expression, ExpressionTypingContext context);
 }

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

@@ -72,9 +72,10 @@ public class ExpressionTypingVisitorDispatcher extends JetVisitor<JetTypeInfo, E
         return basic.getSelectorReturnTypeInfo(receiver, callOperationNode, selectorExpression, context);
     }
 
+    @NotNull
     @Override
-    public boolean checkInExpression(JetElement callElement, @NotNull JetSimpleNameExpression operationSign, @Nullable JetExpression left, @NotNull JetExpression right, ExpressionTypingContext context) {
-        return basic.checkInExpression(callElement, operationSign, left, right, context).getFirst();
+    public JetTypeInfo checkInExpression(JetElement callElement, @NotNull JetSimpleNameExpression operationSign, @Nullable JetExpression left, @NotNull JetExpression right, ExpressionTypingContext context) {
+        return basic.checkInExpression(callElement, operationSign, left, right, context);
     }
 
     @Override

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

@@ -71,9 +71,15 @@ public class PatternMatchingTypingVisitor extends ExpressionTypingVisitor {
         // TODO :change scope according to the bound value in the when header
         final JetExpression subjectExpression = expression.getSubjectExpression();
 
-        final JetType subjectType = subjectExpression != null
-                                    ? context.expressionTypingServices.safeGetType(context.scope, subjectExpression, TypeUtils.NO_EXPECTED_TYPE, context.dataFlowInfo, context.trace)
-                                    : ErrorUtils.createErrorType("Unknown type");
+        final JetType subjectType;
+        if (subjectExpression == null) {
+            subjectType = ErrorUtils.createErrorType("Unknown type");
+        }
+        else {
+            JetTypeInfo typeInfo = facade.safeGetTypeInfo(subjectExpression, context);
+            subjectType = typeInfo.getType();
+            context = context.replaceDataFlowInfo(typeInfo.getDataFlowInfo());
+        }
         final DataFlowValue variableDescriptor = subjectExpression != null ? DataFlowValueFactory.INSTANCE.createDataFlowValue(subjectExpression, subjectType, context.trace.getBindingContext()) : DataFlowValue.NULL;
 
         // TODO : exhaustive patterns
@@ -158,17 +164,21 @@ public class PatternMatchingTypingVisitor extends ExpressionTypingVisitor {
     ) {
         final Ref<DataFlowInfos> newDataFlowInfo = new Ref<DataFlowInfos>(noChange(context));
         condition.accept(new JetVisitorVoid() {
-
             @Override
             public void visitWhenConditionInRange(JetWhenConditionInRange condition) {
                 JetExpression rangeExpression = condition.getRangeExpression();
                 if (rangeExpression == null) return;
                 if (expectedCondition) {
                     context.trace.report(EXPECTED_CONDITION.on(condition));
-                    facade.getTypeInfo(rangeExpression, context);
+                    DataFlowInfo dataFlowInfo = facade.getTypeInfo(rangeExpression, context).getDataFlowInfo();
+                    newDataFlowInfo.set(new DataFlowInfos(dataFlowInfo, dataFlowInfo));
                     return;
                 }
-                if (!facade.checkInExpression(condition, condition.getOperationReference(), subjectExpression, rangeExpression, context)) {
+                JetTypeInfo typeInfo = facade.checkInExpression(condition, condition.getOperationReference(),
+                                                                subjectExpression, rangeExpression, context);
+                DataFlowInfo dataFlowInfo = typeInfo.getDataFlowInfo();
+                newDataFlowInfo.set(new DataFlowInfos(dataFlowInfo, dataFlowInfo));
+                if (!KotlinBuiltIns.getInstance().getBooleanType().equals(typeInfo.getType())) {
                     context.trace.report(TYPE_MISMATCH_IN_RANGE.on(condition));
                 }
             }
@@ -231,6 +241,7 @@ public class PatternMatchingTypingVisitor extends ExpressionTypingVisitor {
         if (type == null) {
             return noChange(context);
         }
+        context = context.replaceDataFlowInfo(typeInfo.getDataFlowInfo());
         if (conditionExpected) {
             JetType booleanType = KotlinBuiltIns.getInstance().getBooleanType();
             if (!JetTypeChecker.INSTANCE.equalTypes(booleanType, type)) {

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

@@ -0,0 +1,24 @@
+fun foo(x: Number, y: Int) {
+    when (x) {
+        x as Int -> x : Int
+        y -> {}
+        else -> {}
+    }
+    <!TYPE_MISMATCH!>x<!> : Int
+}
+
+fun bar(x: Number) {
+    when (x) {
+        x as Int -> x : Int
+        else -> {}
+    }
+    x : Int
+}
+
+fun whenWithoutSubject(x: Number) {
+    when {
+        (x as Int) == 42 -> x : Int
+        else -> {}
+    }
+    x : Int
+}

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

@@ -0,0 +1,13 @@
+fun foo(x: Int, list: List<Int>?) {
+    when (x) {
+        in list!! -> list : List<Int>
+        else -> {}
+    }
+}
+
+fun whenWithoutSubject(x: Int, list: List<Int>?) {
+    when {
+        x in list!! -> list : List<Int>
+        else -> {}
+    }
+}

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

@@ -0,0 +1,6 @@
+fun foo(x: Number) {
+    when (x as Int) {
+        else -> x : Int
+    }
+    x : Int
+}

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

@@ -1316,6 +1316,21 @@ public class JetDiagnosticsTestGenerated extends AbstractDiagnosticsTestWithEage
                 doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/When.kt");
             }
             
+            @TestMetadata("WhenEntry.kt")
+            public void testWhenEntry() throws Exception {
+                doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/WhenEntry.kt");
+            }
+            
+            @TestMetadata("WhenIn.kt")
+            public void testWhenIn() throws Exception {
+                doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/WhenIn.kt");
+            }
+            
+            @TestMetadata("WhenSubject.kt")
+            public void testWhenSubject() throws Exception {
+                doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/WhenSubject.kt");
+            }
+            
             @TestMetadata("While.kt")
             public void testWhile() throws Exception {
                 doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/While.kt");