Retain data flow info after == (equals)

#KT-2825 In Progress

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

@@ -1103,10 +1103,15 @@ public class BasicExpressionTypingVisitor extends ExpressionTypingVisitor {
             if (OperatorConventions.EQUALS_OPERATIONS.contains(operationType)) {
                 Name name = Name.identifier("equals");
                 if (right != null) {
-                    ExpressionReceiver receiver = ExpressionTypingUtils.safeGetExpressionReceiver(facade, left, context.replaceScope(context.scope));
+                    ExpressionReceiver receiver = ExpressionTypingUtils.safeGetExpressionReceiver(facade, left, context);
+
+                    dataFlowInfo = facade.getTypeInfo(left, context).getDataFlowInfo();
+                    ExpressionTypingContext contextWithDataFlow = context.replaceDataFlowInfo(dataFlowInfo);
 
                     OverloadResolutionResults<FunctionDescriptor> resolutionResults = resolveFakeCall(
-                            context, receiver, name, KotlinBuiltIns.getInstance().getNullableAnyType());
+                            contextWithDataFlow, receiver, name, KotlinBuiltIns.getInstance().getNullableAnyType());
+
+                    dataFlowInfo = facade.getTypeInfo(right, contextWithDataFlow).getDataFlowInfo();
 
                     if (resolutionResults.isSuccess()) {
                         FunctionDescriptor equals = resolutionResults.getResultingCall().getResultingDescriptor();

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

@@ -0,0 +1,14 @@
+trait A
+trait B : A
+
+fun foo1(a: A, b: B): Boolean {
+    val result = (a as B) == b
+    a : B
+    return result
+}
+
+fun foo2(a: A, b: B): Boolean {
+    val result = b == (a as B)
+    a : B
+    return result
+}

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

@@ -1151,6 +1151,11 @@ public class JetDiagnosticsTestGenerated extends AbstractDiagnosticsTestWithEage
                 doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/BinaryExpressionCompareToConvention.kt");
             }
             
+            @TestMetadata("BinaryExpressionEqualsConvention.kt")
+            public void testBinaryExpressionEqualsConvention() throws Exception {
+                doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/BinaryExpressionEqualsConvention.kt");
+            }
+            
             @TestMetadata("BinaryExpressionIdentifier.kt")
             public void testBinaryExpressionIdentifier() throws Exception {
                 doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/BinaryExpressionIdentifier.kt");