method chain completion: suggest arrays/collections on 2nd invocation

java/compiler/impl/src/com/intellij/compiler/chainsSearch/ChainsSearcher.java

@@ -57,7 +57,16 @@ public class ChainsSearcher {
   private static SearchInitializer createInitializer(TargetType target,
                                                      CompilerReferenceServiceEx compilerReferenceServiceEx,
                                                      ChainCompletionContext context) {
-    SortedSet<OccurrencesAware<MethodIncompleteSignature>> methods = compilerReferenceServiceEx.findMethodReferenceOccurrences(target.getClassQName(), target.getArrayKind());
+    SortedSet<OccurrencesAware<MethodIncompleteSignature>> methods = null;
+    for (byte kind : target.getArrayKind()) {
+      SortedSet<OccurrencesAware<MethodIncompleteSignature>> currentMethods =
+        compilerReferenceServiceEx.findMethodReferenceOccurrences(target.getClassQName(), kind);
+      if (methods == null) {
+        methods = currentMethods;
+      } else {
+        methods.addAll(currentMethods);
+      }
+    }
     return new SearchInitializer(methods, context);
   }
 

java/compiler/impl/src/com/intellij/compiler/chainsSearch/completion/MethodsChainsCompletionContributor.java

@@ -116,10 +116,10 @@ public class MethodsChainsCompletionContributor extends CompletionContributor {
     LOG.assertTrue(parent != null, "A completion position should match to a pattern");
 
     if (parent instanceof PsiAssignmentExpression) {
-      return extractContextFromAssignment((PsiAssignmentExpression)parent);
+      return extractContextFromAssignment((PsiAssignmentExpression)parent, parameters);
     }
     if (parent instanceof PsiLocalVariable) {
-      return extractContextFromVariable((PsiLocalVariable)parent);
+      return extractContextFromVariable((PsiLocalVariable)parent, parameters);
     }
     PsiMethod method = ((PsiMethodCallExpression)parent).resolveMethod();
     if (method == null) return null;
@@ -130,23 +130,25 @@ public class MethodsChainsCompletionContributor extends CompletionContributor {
     PsiParameter[] methodParameters = method.getParameterList().getParameters();
     if (exprPosition < methodParameters.length) {
       PsiParameter methodParameter = methodParameters[exprPosition];
-      return ChainCompletionContext.createContext(methodParameter.getType(), PsiTreeUtil.getParentOfType(expression, PsiDeclarationStatement.class));
+      return ChainCompletionContext.createContext(methodParameter.getType(), PsiTreeUtil.getParentOfType(expression, PsiDeclarationStatement.class), suggestIterators(parameters));
     }
     return null;
   }
 
   @Nullable
-  private static ChainCompletionContext extractContextFromVariable(PsiLocalVariable localVariable) {
+  private static ChainCompletionContext extractContextFromVariable(PsiLocalVariable localVariable,
+                                                                   CompletionParameters parameters) {
     PsiDeclarationStatement declaration = PsiTreeUtil.getParentOfType(localVariable, PsiDeclarationStatement.class);
-    return ChainCompletionContext.createContext(localVariable.getType(), declaration);
+    return ChainCompletionContext.createContext(localVariable.getType(), declaration, suggestIterators(parameters));
   }
 
   @Nullable
-  private static ChainCompletionContext extractContextFromAssignment(PsiAssignmentExpression assignmentExpression) {
+  private static ChainCompletionContext extractContextFromAssignment(PsiAssignmentExpression assignmentExpression,
+                                                                     CompletionParameters parameters) {
     if (!(assignmentExpression instanceof PsiReferenceExpression)) return null;
     PsiElement resolved = ((PsiReferenceExpression)assignmentExpression).resolve();
     return resolved instanceof PsiVariable
-           ? ChainCompletionContext.createContext(((PsiVariable)resolved).getType(), assignmentExpression)
+           ? ChainCompletionContext.createContext(((PsiVariable)resolved).getType(), assignmentExpression, suggestIterators(parameters))
            : null;
   }
 

java/compiler/impl/src/com/intellij/compiler/chainsSearch/context/ChainCompletionContext.java

@@ -161,10 +161,13 @@ public class ChainCompletionContext {
 
   @Nullable
   public static ChainCompletionContext createContext(@Nullable PsiType targetType,
-                                                     @Nullable PsiElement containingElement) {
+                                                     @Nullable PsiElement containingElement, boolean suggestIterators) {
     if (containingElement == null) return null;
     TargetType target = TargetType.create(targetType);
     if (target == null) return null;
+    if (suggestIterators) {
+      target = target.toIterators();
+    }
 
     ContextProcessor processor = new ContextProcessor(null, containingElement.getProject(), containingElement);
     PsiScopesUtil.treeWalkUp(processor, containingElement, containingElement.getContainingFile());

java/compiler/impl/src/com/intellij/compiler/chainsSearch/context/TargetType.java

@@ -26,14 +26,14 @@ import org.jetbrains.jps.backwardRefs.SignatureData;
  */
 public class TargetType {
   private final String myClassQName;
-  private final byte myArrayKind;
+  private final byte[] myAcceptedArrayKinds;
   private final PsiType myPsiType;
 
   public TargetType(String classQName,
-                    byte arrayKind,
+                    byte[] arrayKinds,
                     PsiType targetType) {
     myClassQName = classQName;
-    myArrayKind = arrayKind;
+    myAcceptedArrayKinds = arrayKinds;
     myPsiType = targetType;
   }
 
@@ -41,15 +41,21 @@ public class TargetType {
     return myClassQName;
   }
 
-  @SignatureData.IteratorKind
-  public byte getArrayKind() {
-    return myArrayKind;
+  //@SignatureData.IteratorKind
+  public byte[] getArrayKind() {
+    return myAcceptedArrayKinds;
   }
 
   public PsiClass getTargetClass() {
     return PsiUtil.resolveClassInType(myPsiType);
   }
 
+  public TargetType toIterators() {
+    return myAcceptedArrayKinds.length == 1 && myAcceptedArrayKinds[0] == SignatureData.ZERO_DIM ?
+           new TargetType(myClassQName, new byte[]{SignatureData.ARRAY_ONE_DIM, SignatureData.ITERATOR_ONE_DIM}, myPsiType) :
+           this;
+  }
+
   @Nullable
   public static TargetType create(PsiType type) {
     if (type instanceof PsiArrayType) {
@@ -69,7 +75,7 @@ public class TargetType {
     if (aClass == null) return null;
     String targetQName = aClass.getQualifiedName();
     if (targetQName == null) return null;
-    return new TargetType(targetQName, SignatureData.ARRAY_ONE_DIM, arrayType);
+    return new TargetType(targetQName, new byte[] {SignatureData.ARRAY_ONE_DIM}, arrayType);
   }
 
   @Nullable
@@ -90,7 +96,7 @@ public class TargetType {
     if (classQName == null) {
       return null;
     }
-    return new TargetType(classQName, iteratorKind, classType);
+    return new TargetType(classQName, new byte[] {iteratorKind}, classType);
   }
 
   public static String getIteratorKind(PsiClass resolvedClass) {