method chain completion: suggest chains for collection of elements

java/compiler/impl/src/com/intellij/compiler/backwardRefs/MethodIncompleteSignature.kt

@@ -15,9 +15,11 @@
  */
 package com.intellij.compiler.backwardRefs
 
+import com.intellij.compiler.chainsSearch.context.TargetType
 import com.intellij.openapi.project.Project
 import com.intellij.psi.*
 import com.intellij.psi.search.GlobalSearchScope
+import com.intellij.psi.util.PsiUtil
 import org.jetbrains.jps.backwardRefs.LightRef
 import org.jetbrains.jps.backwardRefs.SignatureData
 
@@ -67,7 +69,13 @@ class MethodIncompleteSignature(val ref: LightRef.JavaLightMethodRef,
               else -> false
             }
           }
-          SignatureData.ITERATOR_ONE_DIM -> false
+          SignatureData.ITERATOR_ONE_DIM -> {
+            val iteratorKind = TargetType.getIteratorKind(PsiUtil.resolveClassInClassTypeOnly(returnType))
+            when {
+              iteratorKind != null -> PsiUtil.resolveClassInClassTypeOnly(PsiUtil.substituteTypeParameter(returnType, iteratorKind, 0, false))?.qualifiedName == rawReturnType
+              else -> false
+            }
+          }
           SignatureData.ZERO_DIM -> returnType is PsiClassType && returnType.resolve()?.qualifiedName == rawReturnType
           else -> throw IllegalStateException("kind is unsupported ${signatureData.iteratorKind}")
         }

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

@@ -77,27 +77,23 @@ public class TargetType {
     PsiClass resolvedClass = PsiUtil.resolveClassInClassTypeOnly(classType);
     byte iteratorKind = SignatureData.ZERO_DIM;
     if (resolvedClass == null) return null;
-    String iteratorClass = isIterator(resolvedClass);
+    String iteratorClass = getIteratorKind(resolvedClass);
     if (iteratorClass != null) {
-      PsiClassType streamType = (PsiClassType)PsiUtil.substituteTypeParameter(classType, iteratorClass, 0, false);
-      if (streamType == null) return null;
-      PsiType[] parameters = streamType.getParameters();
-      if (parameters.length != 1 || !(parameters[0] instanceof PsiClassType)) return null;
-      resolvedClass = PsiUtil.resolveClassInClassTypeOnly(parameters[0]);
+      resolvedClass = PsiUtil.resolveClassInClassTypeOnly(PsiUtil.substituteTypeParameter(classType, iteratorClass, 0, false));
       if (resolvedClass == null) return null;
       iteratorKind = SignatureData.ITERATOR_ONE_DIM;
     }
+    if (resolvedClass.hasTypeParameters()) {
+      return null;
+    }
     String classQName = resolvedClass.getQualifiedName();
     if (classQName == null) {
       return null;
     }
-    if (resolvedClass.hasTypeParameters()) {
-      return null;
-    }
     return new TargetType(classQName, iteratorKind, classType);
   }
 
-  private static String isIterator(PsiClass resolvedClass) {
+  public static String getIteratorKind(PsiClass resolvedClass) {
     if (InheritanceUtil.isInheritor(resolvedClass, CommonClassNames.JAVA_LANG_ITERABLE)) {
       return CommonClassNames.JAVA_LANG_ITERABLE;
     }

java/java-tests/testData/codeInsight/completion/methodChains/testCollection/TestCompletion.java

@@ -0,0 +1,18 @@
+import java.jang.String;
+import java.util.*;
+
+class PsiMethod {
+}
+
+interface PsiClass {
+  List<PsiMethod> getMethods();
+}
+
+public class TestCompletion {
+
+  PsiClass c;
+
+  public void method() {
+    Collection<PsiMethod> m = <caret>
+  }
+}

java/java-tests/testData/codeInsight/completion/methodChains/testCollection/TestIndex.java

@@ -0,0 +1,23 @@
+import java.util.*;
+
+
+public class TestIndex {
+
+  public void statMethod(PsiClass c) {
+    c.getMethods();
+    c.getMethods();
+    c.getMethods();
+    c.getMethods();
+    c.getMethods();
+    c.getMethods();
+    c.getMethods();
+    c.getMethods();
+  }
+}
+
+class PsiMethod {
+}
+
+interface PsiClass {
+  Collection<PsiMethod> getMethods();
+}

java/java-tests/testSrc/com/intellij/codeInsight/completion/MethodChainsCompletionTest.java

@@ -215,6 +215,11 @@ public class MethodChainsCompletionTest extends AbstractCompilerAwareTest {
     assertEquals("c.getMethods", element.getLookupString());
   }
 
+  public void testCollection() {
+    WeightableChainLookupElement element = assertOneElement(doCompletion());
+    assertEquals("c.getMethods", element.getLookupString());
+  }
+
   public void assertAdvisorLookupElementEquals(String lookupText,
                                                int unreachableParametersCount,
                                                int lastMethodWeight,

jps/jps-builders/src/org/jetbrains/jps/backwardRefs/SignatureData.java

@@ -21,7 +21,7 @@ public class SignatureData {
   public static final byte ZERO_DIM = 0;
   public static final byte ARRAY_ONE_DIM = 1;
   // represents java's Iterator, Iterable and BaseStream
-  public static final byte ITERATOR_ONE_DIM = 2;
+  public static final byte ITERATOR_ONE_DIM = -1;
   @MagicConstant(intValues = {ZERO_DIM, ARRAY_ONE_DIM, ITERATOR_ONE_DIM})
   public @interface IteratorKind {}