KT-5155 Auto-casts do not work with when

#KT-5155 Fixed

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

@@ -100,21 +100,14 @@ public class PatternMatchingTypingVisitor extends ExpressionTypingVisitor {
         DataFlowInfo elseDataFlowInfo = context.dataFlowInfo;
         for (JetWhenEntry whenEntry : expression.getEntries()) {
             DataFlowInfos infosForCondition = getDataFlowInfosForEntryCondition(
-                    whenEntry, context, subjectExpression, subjectType, subjectDataFlowValue);
-            DataFlowInfo dataFlowInfoForEntryBody;
-            if (infosForCondition == null) {
-                dataFlowInfoForEntryBody = elseDataFlowInfo;
-            }
-            else {
-                dataFlowInfoForEntryBody = infosForCondition.thenInfo.and(elseDataFlowInfo);
-                elseDataFlowInfo = elseDataFlowInfo.and(infosForCondition.elseInfo);
-            }
+                    whenEntry, context.replaceDataFlowInfo(elseDataFlowInfo), subjectExpression, subjectType, subjectDataFlowValue);
+            elseDataFlowInfo = elseDataFlowInfo.and(infosForCondition.elseInfo);
 
             JetExpression bodyExpression = whenEntry.getExpression();
             if (bodyExpression != null) {
                 WritableScope scopeToExtend = newWritableScopeImpl(context, "Scope extended in when entry");
                 ExpressionTypingContext newContext = contextWithExpectedType
-                        .replaceScope(scopeToExtend).replaceDataFlowInfo(dataFlowInfoForEntryBody).replaceContextDependency(INDEPENDENT);
+                        .replaceScope(scopeToExtend).replaceDataFlowInfo(infosForCondition.thenInfo).replaceContextDependency(INDEPENDENT);
                 CoercionStrategy coercionStrategy = isStatement ? CoercionStrategy.COERCION_TO_UNIT : CoercionStrategy.NO_COERCION;
                 JetTypeInfo typeInfo = components.expressionTypingServices.getBlockReturnedTypeWithWritableScope(
                         scopeToExtend, Collections.singletonList(bodyExpression), coercionStrategy, newContext, context.trace);
@@ -141,7 +134,7 @@ public class PatternMatchingTypingVisitor extends ExpressionTypingVisitor {
         return JetTypeInfo.create(null, commonDataFlowInfo);
     }
 
-    @Nullable
+    @NotNull
     private DataFlowInfos getDataFlowInfosForEntryCondition(
             @NotNull JetWhenEntry whenEntry,
             @NotNull ExpressionTypingContext context,
@@ -149,13 +142,12 @@ public class PatternMatchingTypingVisitor extends ExpressionTypingVisitor {
             @NotNull JetType subjectType,
             @NotNull DataFlowValue subjectDataFlowValue
     ) {
-        JetWhenCondition[] conditions = whenEntry.getConditions();
         if (whenEntry.isElse()) {
-            return null;
+            return new DataFlowInfos(context.dataFlowInfo);
         }
 
         DataFlowInfos infos = null;
-        for (JetWhenCondition condition : conditions) {
+        for (JetWhenCondition condition : whenEntry.getConditions()) {
             DataFlowInfos conditionInfos = checkWhenCondition(subjectExpression, subjectExpression == null, subjectType, condition,
                                                               context, subjectDataFlowValue);
             if (infos != null) {
@@ -165,7 +157,7 @@ public class PatternMatchingTypingVisitor extends ExpressionTypingVisitor {
                 infos = conditionInfos;
             }
         }
-        return infos;
+        return infos != null ? infos : new DataFlowInfos(context.dataFlowInfo);
     }
 
     private DataFlowInfos checkWhenCondition(
@@ -238,6 +230,10 @@ public class PatternMatchingTypingVisitor extends ExpressionTypingVisitor {
             this.thenInfo = thenInfo;
             this.elseInfo = elseInfo;
         }
+
+        private DataFlowInfos(DataFlowInfo info) {
+            this(info, info);
+        }
     }
 
     private DataFlowInfos checkTypeForExpressionCondition(

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

@@ -0,0 +1,14 @@
+//KT-5155 Auto-casts do not work with when
+
+fun foo(s: String?) {
+    when {
+        s == null -> 1
+        <!DEBUG_INFO_AUTOCAST!>s<!>.foo() -> 2
+        else -> 3
+    }
+}
+
+fun String.foo() = true
+
+
+

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

@@ -2103,6 +2103,11 @@ public class JetDiagnosticsTestGenerated extends AbstractJetDiagnosticsTest {
                 doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/kt4332WhenBranches.kt");
             }
             
+            @TestMetadata("kt5155WhenBranches.kt")
+            public void testKt5155WhenBranches() throws Exception {
+                doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/kt5155WhenBranches.kt");
+            }
+            
             @TestMetadata("kt5182WhenBranches.kt")
             public void testKt5182WhenBranches() throws Exception {
                 doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/kt5182WhenBranches.kt");