Retain data flow info after array accesses

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

@@ -1289,9 +1289,8 @@ public class BasicExpressionTypingVisitor extends ExpressionTypingVisitor {
 
     @Override
     public JetTypeInfo visitArrayAccessExpression(JetArrayAccessExpression expression, ExpressionTypingContext context) {
-        JetType type = resolveArrayAccessGetMethod(expression, context.replaceExpectedType(NO_EXPECTED_TYPE));
+        JetTypeInfo typeInfo = resolveArrayAccessGetMethod(expression, context.replaceExpectedType(NO_EXPECTED_TYPE));
-        DataFlowUtils.checkType(type, expression, context);
+        return DataFlowUtils.checkType(typeInfo.getType(), expression, context, typeInfo.getDataFlowInfo());
-        return JetTypeInfo.create(type, context.dataFlowInfo);
     }
 
     @NotNull
@@ -1406,38 +1405,51 @@ public class BasicExpressionTypingVisitor extends ExpressionTypingVisitor {
         return JetTypeInfo.create(null, context.dataFlowInfo);
     }
 
-    @Nullable
+    @NotNull
-        /*package*/ JetType resolveArrayAccessSetMethod(@NotNull JetArrayAccessExpression arrayAccessExpression, @NotNull JetExpression rightHandSide, @NotNull ExpressionTypingContext context, @NotNull BindingTrace traceForResolveResult) {
+        /*package*/ JetTypeInfo resolveArrayAccessSetMethod(@NotNull JetArrayAccessExpression arrayAccessExpression, @NotNull JetExpression rightHandSide, @NotNull ExpressionTypingContext context, @NotNull BindingTrace traceForResolveResult) {
         return resolveArrayAccessSpecialMethod(arrayAccessExpression, rightHandSide, context, traceForResolveResult, false);
     }
 
-    @Nullable
+    @NotNull
-        /*package*/ JetType resolveArrayAccessGetMethod(@NotNull JetArrayAccessExpression arrayAccessExpression, @NotNull ExpressionTypingContext context) {
+        /*package*/ JetTypeInfo resolveArrayAccessGetMethod(@NotNull JetArrayAccessExpression arrayAccessExpression, @NotNull ExpressionTypingContext context) {
         return resolveArrayAccessSpecialMethod(arrayAccessExpression, null, context, context.trace, true);
     }
 
-    @Nullable
+    @NotNull
-    private JetType resolveArrayAccessSpecialMethod(@NotNull JetArrayAccessExpression arrayAccessExpression,
+    private JetTypeInfo resolveArrayAccessSpecialMethod(@NotNull JetArrayAccessExpression arrayAccessExpression,
                                                     @Nullable JetExpression rightHandSide, //only for 'set' method
-                                                    @NotNull ExpressionTypingContext context,
+                                                    @NotNull ExpressionTypingContext oldContext,
                                                     @NotNull BindingTrace traceForResolveResult,
                                                     boolean isGet) {
-        JetType arrayType = facade.getTypeInfo(arrayAccessExpression.getArrayExpression(), context).getType();
+        JetTypeInfo arrayTypeInfo = facade.getTypeInfo(arrayAccessExpression.getArrayExpression(), oldContext);
-        if (arrayType == null) return null;
+        JetType arrayType = arrayTypeInfo.getType();
+        if (arrayType == null) return arrayTypeInfo;
 
+        DataFlowInfo dataFlowInfo = arrayTypeInfo.getDataFlowInfo();
+        ExpressionTypingContext context = oldContext.replaceDataFlowInfo(dataFlowInfo);
         ExpressionReceiver receiver = new ExpressionReceiver(arrayAccessExpression.getArrayExpression(), arrayType);
         if (!isGet) assert rightHandSide != null;
+
         OverloadResolutionResults<FunctionDescriptor> functionResults = context.resolveCallWithGivenName(
                 isGet
                 ? CallMaker.makeArrayGetCall(receiver, arrayAccessExpression, Call.CallType.ARRAY_GET_METHOD)
                 : CallMaker.makeArraySetCall(receiver, arrayAccessExpression, rightHandSide, Call.CallType.ARRAY_SET_METHOD),
                 arrayAccessExpression,
                 Name.identifier(isGet ? "get" : "set"));
+
+        List<JetExpression> indices = arrayAccessExpression.getIndexExpressions();
+        // The accumulated data flow info of all index expressions is saved on the last index
+        dataFlowInfo = facade.getTypeInfo(indices.get(indices.size() - 1), context).getDataFlowInfo();
+
+        if (!isGet) {
+            dataFlowInfo = facade.getTypeInfo(rightHandSide, context.replaceDataFlowInfo(dataFlowInfo)).getDataFlowInfo();
+        }
+
         if (!functionResults.isSuccess()) {
             traceForResolveResult.report(isGet ? NO_GET_METHOD.on(arrayAccessExpression) : NO_SET_METHOD.on(arrayAccessExpression));
-            return null;
+            return JetTypeInfo.create(null, dataFlowInfo);
         }
         traceForResolveResult.record(isGet ? INDEXED_LVALUE_GET : INDEXED_LVALUE_SET, arrayAccessExpression, functionResults.getResultingCall());
-        return functionResults.getResultingDescriptor().getReturnType();
+        return JetTypeInfo.create(functionResults.getResultingDescriptor().getReturnType(), dataFlowInfo);
     }
 }

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

@@ -278,9 +278,10 @@ public class ExpressionTypingVisitorForStatements extends ExpressionTypingVisito
         if (left instanceof JetArrayAccessExpression) {
             JetArrayAccessExpression arrayAccessExpression = (JetArrayAccessExpression) left;
             if (right == null) return JetTypeInfo.create(null, context.dataFlowInfo);
-            JetType assignmentType = basic.resolveArrayAccessSetMethod(arrayAccessExpression, right, context, context.trace);
+            JetTypeInfo typeInfo = basic.resolveArrayAccessSetMethod(arrayAccessExpression, right, context, context.trace);
             basic.checkLValue(context.trace, arrayAccessExpression);
-            return JetTypeInfo.create(checkAssignmentType(assignmentType, expression, contextWithExpectedType), context.dataFlowInfo);
+            return JetTypeInfo.create(checkAssignmentType(typeInfo.getType(), expression, contextWithExpectedType),
+                                      typeInfo.getDataFlowInfo());
         }
         JetTypeInfo leftInfo = facade.getTypeInfo(expression.getLeft(), context);
         JetType leftType = leftInfo.getType();

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

@@ -0,0 +1,5 @@
+fun foo(arr: Array<out Number>): Int {
+    val result = (arr as Array<Int>)[0]
+    arr : Array<Int>
+    return result
+}

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

@@ -0,0 +1,26 @@
+trait G {
+    fun get(x: Int, y: Int): Int = x + y
+    fun set(x: Int, y: Int, value: Int) {}
+}
+
+fun foo1(a: Int?, b: G) {
+    b[a!!, a<!UNNECESSARY_NOT_NULL_ASSERTION!>!!<!>] = a
+    a : Int
+}
+
+fun foo2(a: Int?, b: G) {
+    b[0, a!!] = a
+    a : Int
+}
+
+fun foo3(a: Int?, b: G) {
+    val r = b[a!!, a]
+    a : Int
+    r : Int
+}
+
+fun foo4(a: Int?, b: G) {
+    val r = b[0, a!!]
+    a : Int
+    r : Int
+}

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

@@ -0,0 +1,9 @@
+fun foo1(a: Int?, b: Array<Array<Int>>) {
+    b[a!!][a<!UNNECESSARY_NOT_NULL_ASSERTION!>!!<!>] = a
+    a : Int
+}
+
+fun foo2(a: Int?, b: Array<Array<Int>>) {
+    b[0][a!!] = a
+    a : Int
+}

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

@@ -0,0 +1,10 @@
+fun arrayAccessRHS(a: Int?, b: Array<Int>) {
+    b[0] = a!!
+    a : Int
+}
+
+fun arrayAccessLHS(a: Int?, b: Array<Int>) {
+    b[a!!] = a
+    a : Int
+}
+

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

@@ -1141,6 +1141,21 @@ public class JetDiagnosticsTestGenerated extends AbstractDiagnosticsTestWithEage
                 doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/ArrayAccess.kt");
             }
             
+            @TestMetadata("ArrayExpression.kt")
+            public void testArrayExpression() throws Exception {
+                doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/ArrayExpression.kt");
+            }
+            
+            @TestMetadata("ArrayGetSetConvention.kt")
+            public void testArrayGetSetConvention() throws Exception {
+                doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/ArrayGetSetConvention.kt");
+            }
+            
+            @TestMetadata("ArrayIndices.kt")
+            public void testArrayIndices() throws Exception {
+                doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/ArrayIndices.kt");
+            }
+            
             @TestMetadata("Assignment.kt")
             public void testAssignment() throws Exception {
                 doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/Assignment.kt");
@@ -1156,6 +1171,11 @@ public class JetDiagnosticsTestGenerated extends AbstractDiagnosticsTestWithEage
                 doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/AssignmentOperation.kt");
             }
             
+            @TestMetadata("AssignmentToArrayElement.kt")
+            public void testAssignmentToArrayElement() throws Exception {
+                doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/AssignmentToArrayElement.kt");
+            }
+            
             @TestMetadata("BinaryExpression.kt")
             public void testBinaryExpression() throws Exception {
                 doTest("compiler/testData/diagnostics/tests/dataFlowInfoTraversal/BinaryExpression.kt");