javac ast indices: store data about method with array and "collection" return types

jps/jps-builders-6/src/org/jetbrains/jps/javac/ast/JavacReferenceCollectorListener.java

@@ -269,6 +269,10 @@ final class JavacReferenceCollectorListener implements TaskListener {
       return myTypeUtility;
     }
 
+    JavacNameTable getNameTable() {
+      return myNameTableCache;
+    }
+
     private int decrementRemainDeclarationsAndGet(Tree declarationToProcess) {
       return declarationToProcess == null ? myRemainDeclarations : --myRemainDeclarations;
     }

jps/jps-builders-6/src/org/jetbrains/jps/javac/ast/JavacTreeRefScanner.java

@@ -19,9 +19,11 @@ import com.intellij.util.containers.Stack;
 import com.sun.source.tree.*;
 import com.sun.source.util.TreeScanner;
 import org.jetbrains.jps.javac.ast.api.JavacDef;
+import org.jetbrains.jps.javac.ast.api.JavacNameTable;
 import org.jetbrains.jps.javac.ast.api.JavacRef;
 
 import javax.lang.model.element.*;
+import javax.lang.model.type.ArrayType;
 import javax.lang.model.type.DeclaredType;
 import javax.lang.model.type.TypeKind;
 import javax.lang.model.type.TypeMirror;
@@ -75,11 +77,7 @@ class JavacTreeRefScanner extends TreeScanner<Tree, JavacReferenceCollectorListe
     if (element != null && element.getKind() == ElementKind.FIELD) {
       final JavacRef.JavacElementRefBase ref = refCollector.asJavacRef(element);
       if (ref != null) {
-        refCollector.sinkReference(ref);
-        final JavacRef.JavacElementRefBase returnType = refCollector.asJavacRef(element.asType());
-        if (returnType != null) {
-           refCollector.sinkDeclaration(new JavacDef.JavacMemberDef(ref, returnType, isStatic(element)));
-        }
+        processMemberDefinition(refCollector, ref, element, element.asType());
       }
     }
     return super.visitVariable(node, refCollector);
@@ -106,16 +104,35 @@ class JavacTreeRefScanner extends TreeScanner<Tree, JavacReferenceCollectorListe
     if (element != null) {
       final JavacRef.JavacElementRefBase ref = refCollector.asJavacRef(element);
       if (ref != null) {
-        refCollector.sinkReference(ref);
-        final JavacRef.JavacElementRefBase returnType = refCollector.asJavacRef(((ExecutableElement)element).getReturnType());
-        if (returnType != null) {
-          refCollector.sinkDeclaration(new JavacDef.JavacMemberDef(ref, returnType, isStatic(element)));
-        }
+        processMemberDefinition(refCollector, ref, element, ((ExecutableElement)element).getReturnType());
       }
     }
     return super.visitMethod(node, refCollector);
   }
 
+  private void processMemberDefinition(JavacReferenceCollectorListener.ReferenceCollector refCollector,
+                                       JavacRef.JavacElementRefBase ref,
+                                       Element element,
+                                       TypeMirror retType) {
+    refCollector.sinkReference(ref);
+    byte dimension = 0;
+    if (retType.getKind() == TypeKind.ARRAY) {
+      retType = ((ArrayType)retType).getComponentType();
+      dimension = 1;
+    }
+    else if (retType.getKind() == TypeKind.DECLARED) {
+      List<? extends TypeMirror> typeArguments = ((DeclaredType)retType).getTypeArguments();
+      if (typeArguments.size() == 1 && isIterator((TypeElement)((DeclaredType)retType).asElement(), refCollector)) {
+        dimension = -1;
+        retType = typeArguments.get(0);
+      }
+    }
+    final JavacRef.JavacElementRefBase returnType = refCollector.asJavacRef(retType);
+    if (returnType != null) {
+      refCollector.sinkDeclaration(new JavacDef.JavacMemberDef(ref, returnType, dimension, isStatic(element)));
+    }
+  }
+
   @Override
   public Tree visitMethodInvocation(MethodInvocationTree node, JavacReferenceCollectorListener.ReferenceCollector collector) {
     if (node.getMethodSelect() instanceof IdentifierTree) {
@@ -208,6 +225,23 @@ class JavacTreeRefScanner extends TreeScanner<Tree, JavacReferenceCollectorListe
     return null;
   }
 
+  private static boolean isIterator(TypeElement aClass, JavacReferenceCollectorListener.ReferenceCollector collector) {
+    JavacNameTable table = collector.getNameTable();
+    TypeElement iterable = table.getIterableElement();
+    if (iterable != null && isInheritorOrSelf(aClass, iterable)) {
+      return true;
+    }
+    TypeElement stream = table.getStreamElement();
+    if (stream != null && isInheritorOrSelf(aClass, stream)) {
+      return true;
+    }
+    TypeElement iterator = table.getIteratorElement();
+    if (iterator != null && isInheritorOrSelf(aClass, iterator)) {
+      return true;
+    }
+    return false;
+  }
+
   private static Element getClassOrPackageParent(Element element) {
     element = element.getEnclosingElement();
     while (element != null) {

jps/jps-builders-6/src/org/jetbrains/jps/javac/ast/api/JavacDef.java

@@ -51,11 +51,13 @@ public abstract class JavacDef {
 
   public static class JavacMemberDef extends JavacDef {
     private final JavacRef myRawReturnType;
+    private final byte myArrayDimension;
     private final boolean myStatic;
 
-    public JavacMemberDef(JavacRef element, JavacRef rawReturnType, boolean isStatic) {
+    public JavacMemberDef(JavacRef element, JavacRef rawReturnType, byte dimension, boolean isStatic) {
       super(element);
       myRawReturnType = rawReturnType;
+      myArrayDimension = dimension;
       myStatic = isStatic;
     }
 
@@ -63,6 +65,10 @@ public abstract class JavacDef {
       return myRawReturnType;
     }
 
+    public byte getIteratorKind() {
+      return myArrayDimension;
+    }
+
     public boolean isStatic() {
       return myStatic;
     }

jps/jps-builders-6/src/org/jetbrains/jps/javac/ast/api/JavacFileData.java

@@ -85,6 +85,7 @@ public class JavacFileData {
 
   @NotNull
   public static JavacFileData fromBytes(byte[] bytes) {
+    @SuppressWarnings("IOResourceOpenedButNotSafelyClosed")
     final DataInputStream in = new DataInputStream(new ByteArrayInputStream(bytes));
     try {
       return new JavacFileData(in.readUTF(),
@@ -160,8 +161,9 @@ public class JavacFileData {
         return new JavacDef.JavacFunExprDef(readJavacRef(in));
       case METHOD_MARKER:
         JavacRef retType = readJavacRef(in);
+        byte dimension = in.readByte();
         boolean isStatic = in.readBoolean();
-        return new JavacDef.JavacMemberDef(readJavacRef(in), retType, isStatic);
+        return new JavacDef.JavacMemberDef(readJavacRef(in), retType, dimension, isStatic);
       default:
         throw new IllegalStateException("unknown marker " + marker);
     }
@@ -182,6 +184,7 @@ public class JavacFileData {
     else if (def instanceof JavacDef.JavacMemberDef) {
       out.writeByte(METHOD_MARKER);
       writeJavacRef(out, ((JavacDef.JavacMemberDef)def).getReturnType());
+      out.writeByte(((JavacDef.JavacMemberDef)def).getIteratorKind());
       out.writeBoolean(((JavacDef.JavacMemberDef)def).isStatic());
     }
     else {

jps/jps-builders-6/src/org/jetbrains/jps/javac/ast/api/JavacNameTable.java

@@ -17,6 +17,7 @@ package org.jetbrains.jps.javac.ast.api;
 
 import com.intellij.util.containers.SLRUCache;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 import javax.lang.model.element.Element;
 import javax.lang.model.element.Name;
@@ -27,6 +28,9 @@ public class JavacNameTable {
   private final SLRUCache<Name, String> myParsedNameCache;
   private final Elements myElements;
   private Name myAsterisk;
+  private TypeElement myStreamElement;
+  private TypeElement myIteratorElement;
+  private TypeElement myIterableElement;
 
   public JavacNameTable(Elements elements) {
     myParsedNameCache = new SLRUCache<Name, String>(1000, 1000) {
@@ -56,4 +60,28 @@ public class JavacNameTable {
     }
     return myAsterisk;
   }
+
+  @Nullable("if the type is not loaded to javac name table")
+  public TypeElement getStreamElement() {
+    if (myStreamElement == null) {
+      myStreamElement = myElements.getTypeElement("java.util.stream.Stream");
+    }
+    return myStreamElement;
+  }
+
+  @Nullable("if the type is not loaded to javac name table")
+  public TypeElement getIteratorElement() {
+    if (myIteratorElement == null) {
+      myIteratorElement = myElements.getTypeElement("java.util.Iterator");
+    }
+    return myIteratorElement;
+  }
+
+  @Nullable("if the type is not loaded to javac name table")
+  public TypeElement getIterableElement() {
+    if (myIterableElement == null) {
+      myIterableElement = myElements.getTypeElement("java.lang.Iterable");
+    }
+    return myIterableElement;
+  }
 }

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

@@ -76,7 +76,7 @@ public class BackwardReferenceIndexUtil {
           final LightRef ref = writer.enumerateNames(def.getDefinedElement(), name -> anonymousClassEnumerator.getLightRefIfAnonymous(name));
           final LightRef.JavaLightClassRef returnType = writer.asClassUsage(((JavacDef.JavacMemberDef)def).getReturnType());
           if (ref != null && returnType != null) {
-            final SignatureData data = new SignatureData(returnType.getName(), ((JavacDef.JavacMemberDef)def).isStatic());
+            final SignatureData data = new SignatureData(returnType.getName(), ((JavacDef.JavacMemberDef)def).getIteratorKind(), ((JavacDef.JavacMemberDef)def).isStatic());
             signatureData.computeIfAbsent(data, element -> new SmartList<>()).add(ref);
           }
         }

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

@@ -16,11 +16,18 @@
 package org.jetbrains.jps.backwardRefs;
 
 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;
+
   private final int myRawReturnType;
+  private final byte myArrayDimension;
   private final boolean myStatic;
 
-  public SignatureData(int rawReturnType, boolean isStatic) {
+  public SignatureData(int rawReturnType, byte arrayDimension, boolean isStatic) {
     myRawReturnType = rawReturnType;
+    myArrayDimension = arrayDimension;
     myStatic = isStatic;
   }
 
@@ -28,6 +35,10 @@ public class SignatureData {
     return myRawReturnType;
   }
 
+  public byte getIteratorKind() {
+    return myArrayDimension;
+  }
+
   public boolean isStatic() {
     return myStatic;
   }
@@ -38,12 +49,18 @@ public class SignatureData {
     if (o == null || getClass() != o.getClass()) return false;
 
     SignatureData data = (SignatureData)o;
-    return myRawReturnType == data.myRawReturnType && myStatic == data.myStatic;
+
+    if (myRawReturnType != data.myRawReturnType) return false;
+    if (myArrayDimension != data.myArrayDimension) return false;
+    if (myStatic != data.myStatic) return false;
+
+    return true;
   }
 
   @Override
   public int hashCode() {
     int result = myRawReturnType;
+    result = 31 * result + myArrayDimension;
     result = 31 * result + (myStatic ? 1 : 0);
     return result;
   }

jps/jps-builders/src/org/jetbrains/jps/backwardRefs/index/CompilerIndices.java

@@ -37,7 +37,7 @@ import java.util.List;
 
 public class CompilerIndices {
   //TODO manage version separately
-  public final static int VERSION = 4;
+  public final static int VERSION = 5;
 
   public final static ID<LightRef, Integer> BACK_USAGES = ID.create("back.refs");
   public final static ID<LightRef, Collection<LightRef>> BACK_HIERARCHY = ID.create("back.hierarchy");
@@ -217,12 +217,13 @@ public class CompilerIndices {
       @Override
       public void save(@NotNull DataOutput out, SignatureData value) throws IOException {
         DataInputOutputUtil.writeINT(out, value.getRawReturnType());
+        out.writeByte(value.getIteratorKind());
         out.writeBoolean(value.isStatic());
       }
 
       @Override
       public SignatureData read(@NotNull DataInput in) throws IOException {
-        return new SignatureData(DataInputOutputUtil.readINT(in), in.readBoolean());
+        return new SignatureData(DataInputOutputUtil.readINT(in), in.readByte(), in.readBoolean());
       }
     };
   }

jps/jps-builders/testData/referencesIndex/signatureDataIndex/Foo.java

@@ -14,4 +14,15 @@ public class Foo {
     return null;
   }
 
+  public String[] m4() {
+    return null;
+  }
+
+  public String f1;
+
+  public static List<String> f2;
+
+  public List f3;
+
+  public String[] f4;
 }

jps/jps-builders/testData/referencesIndex/signatureDataIndex/initialIndex.txt

@@ -4,17 +4,23 @@ java.lang.Object -> Foo
 Backward References:
 Foo in Foo occurrences = 1
 Foo.<init>(0) in Foo occurrences = 1
+Foo.f1 in Foo occurrences = 1
+Foo.f2 in Foo occurrences = 1
+Foo.f3 in Foo occurrences = 1
+Foo.f4 in Foo occurrences = 1
 Foo.m1(0) in Foo occurrences = 1
 Foo.m2(0) in Foo occurrences = 1
 Foo.m3(0) in Foo occurrences = 1
+Foo.m4(0) in Foo occurrences = 1
 java.lang.Object.<init>(0) in Foo occurrences = 1
-java.lang.String in Foo occurrences = 2
-java.util.List in Foo occurrences = 3
+java.lang.String in Foo occurrences = 6
+java.util.List in Foo occurrences = 5
 
 Class Definitions:
 Foo in Foo
 
 Members Signatures:
-java.lang.String <- Foo.m1(0)
-java.util.List <- Foo.m3(0)
-static java.util.List <- Foo.m2(0)
+java.lang.String <- Foo.f1 Foo.m1(0)
+java.lang.String[] <- Foo.f4 Foo.m4(0)
+java.util.List <- Foo.f3 Foo.m3(0)
+static java.lang.String iterator <- Foo.f2 Foo.m2(0)

jps/jps-builders/testSrc/org/jetbrains/references/ReferenceIndexTestBase.kt

@@ -185,7 +185,16 @@ abstract class ReferenceIndexTestBase : JpsBuildTestCase() {
       }
 
   private fun SignatureData.asText(nameEnumerator: NameEnumerator): String {
-    return (if (this.isStatic) "static " else "") + this.rawReturnType.asName(nameEnumerator)
+    return (if (this.isStatic) "static " else "") + this.rawReturnType.asName(nameEnumerator) + decodeVectorKind(this.iteratorKind)
+  }
+
+  private fun decodeVectorKind(kind: Byte): String {
+    when (kind) {
+      0.toByte() -> return ""
+      1.toByte() -> return "[]"
+      (-1).toByte() -> return " iterator"
+    }
+    throw IllegalArgumentException()
   }
 
   private fun Int.asFileName(fileNameEnumerator: PersistentStringEnumerator) = FileUtil.getNameWithoutExtension(File(fileNameEnumerator.valueOf(this)).canonicalFile)