javac ast indices: when java file contains >1 top declaration classes only one was...
[idea/community.git] / jps / jps-builders / src / org / jetbrains / jps / javac / ast / JavacReferencesCollector.java
1 /*
2  * Copyright 2000-2016 JetBrains s.r.o.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 package org.jetbrains.jps.javac.ast;
17
18 import com.intellij.util.ReflectionUtil;
19 import com.intellij.util.SmartList;
20 import com.sun.source.tree.Tree;
21 import com.sun.source.util.JavacTask;
22 import com.sun.source.util.TaskEvent;
23 import com.sun.source.util.TaskListener;
24 import com.sun.tools.javac.code.Symbol;
25 import com.sun.tools.javac.tree.JCTree;
26 import com.sun.tools.javac.util.ClientCodeException;
27 import com.sun.tools.javac.util.Name;
28 import gnu.trove.THashSet;
29 import org.jetbrains.jps.javac.ast.api.JavacFileReferencesRegistrar;
30 import org.jetbrains.jps.javac.ast.api.JavacRefSymbol;
31 import org.jetbrains.jps.service.JpsServiceManager;
32
33 import javax.lang.model.element.ElementKind;
34 import javax.lang.model.element.TypeElement;
35 import javax.tools.*;
36 import java.lang.reflect.InvocationTargetException;
37 import java.lang.reflect.Method;
38 import java.util.*;
39
40 public class JavacReferencesCollector {
41   public static void installOn(JavacTask task) {
42     List<JavacFileReferencesRegistrar> fullASTListeners = new SmartList<JavacFileReferencesRegistrar>();
43     List<JavacFileReferencesRegistrar> onlyImportsListeners = new SmartList<JavacFileReferencesRegistrar>();
44     for (JavacFileReferencesRegistrar listener : JpsServiceManager.getInstance().getExtensions(JavacFileReferencesRegistrar.class)) {
45       if (!listener.initialize()) {
46         continue;
47       }
48       (listener.onlyImports() ? onlyImportsListeners : fullASTListeners).add(listener);
49     }
50
51     final JavacFileReferencesRegistrar[] fullASTListenerArray = fullASTListeners.toArray(new JavacFileReferencesRegistrar[fullASTListeners.size()]);
52     final JavacFileReferencesRegistrar[] onlyImportsListenerArray = onlyImportsListeners.toArray(new JavacFileReferencesRegistrar[onlyImportsListeners.size()]);
53     if (fullASTListenerArray.length == 0 && onlyImportsListenerArray.length == 0) {
54       return;
55     }
56
57     Method addTaskMethod = ReflectionUtil.getMethod(JavacTask.class, "addTaskListener", TaskListener.class); // jdk >= 8
58     if (addTaskMethod == null) {
59       addTaskMethod = ReflectionUtil.getMethod(JavacTask.class, "setTaskListener", TaskListener.class); // jdk 6-7
60     }
61     assert addTaskMethod != null;
62
63     try {
64       addTaskMethod.invoke(task, new MyTaskListener(fullASTListenerArray, onlyImportsListenerArray));
65     }
66     catch (IllegalAccessException e) {
67       throw new RuntimeException(e);
68     }
69     catch (InvocationTargetException e) {
70       throw new RuntimeException(e);
71     }
72   }
73
74   private static final class MyTaskListener implements TaskListener {
75     private final JavacFileReferencesRegistrar[] myFullASTListeners;
76     private final JavacFileReferencesRegistrar[] myOnlyImportsListeners;
77     private final JavacTreeRefScanner myAstScanner;
78
79     private Name myAsteriks;
80
81     private int myRemainDecls;
82     private JCTree.JCCompilationUnit myCurrentCompilationUnit;
83     private Set<JavacRefSymbol> myCollectedReferences;
84     private Set<JavacRefSymbol> myCollectedDefinitions;
85
86     public MyTaskListener(JavacFileReferencesRegistrar[] fullASTListenerArray, JavacFileReferencesRegistrar[] importsListenerArray) {
87       myFullASTListeners = fullASTListenerArray;
88       myOnlyImportsListeners = importsListenerArray;
89       myAstScanner = JavacTreeRefScanner.createASTScanner();
90     }
91
92     @Override
93     public void started(TaskEvent e) {
94
95     }
96
97     @Override
98     public void finished(TaskEvent e) {
99       try {
100         if (e.getKind() == TaskEvent.Kind.ANALYZE) {
101           // javac creates an event on each processed top level declared class not file
102           if (myCurrentCompilationUnit != e.getCompilationUnit()) {
103             myCurrentCompilationUnit = (JCTree.JCCompilationUnit)e.getCompilationUnit();
104             myCollectedDefinitions = new THashSet<JavacRefSymbol>();
105             myCollectedReferences = new THashSet<JavacRefSymbol>();
106             myRemainDecls = myCurrentCompilationUnit.getTypeDecls().size() - 1;
107             scanImports(myCurrentCompilationUnit, myCollectedReferences);
108             for(JavacFileReferencesRegistrar r: myOnlyImportsListeners) {
109               r.registerFile(e.getSourceFile(), myCollectedReferences, myCollectedDefinitions);
110             }
111           }
112           else {
113             myRemainDecls--;
114           }
115
116           JavacTreeScannerSink sink = new JavacTreeScannerSink() {
117             @Override
118             public void sinkReference(JavacRefSymbol ref) {
119               myCollectedReferences.add(ref);
120             }
121
122             @Override
123             public void sinkDeclaration(JavacRefSymbol def) {
124               myCollectedDefinitions.add(def);
125             }
126           };
127
128           if (myFullASTListeners.length != 0) {
129             TypeElement analyzedElement = e.getTypeElement();
130             for (JCTree tree : myCurrentCompilationUnit.getTypeDecls()) {
131               if (tree.type != null && tree.type.tsym == analyzedElement) {
132                 myAstScanner.scan(tree, sink);
133               }
134             }
135           }
136
137           if (myRemainDecls == 0) {
138             if (myFullASTListeners.length != 0) {
139               for (JCTree.JCAnnotation annotation : myCurrentCompilationUnit.getPackageAnnotations()) {
140                 myAstScanner.scan(annotation, sink);
141               }
142             }
143
144             for(JavacFileReferencesRegistrar r: myFullASTListeners) {
145               r.registerFile(e.getSourceFile(), myCollectedReferences, myCollectedDefinitions);
146             }
147
148             myCurrentCompilationUnit = null;
149             myCollectedDefinitions = null;
150             myCollectedReferences = null;
151           }
152
153         }
154       }
155       catch (Exception ex) {
156         throw new ClientCodeException(ex);
157       }
158     }
159
160     private Name getAsteriskFromCurrentNameTable(Name tableRepresentative) {
161       if (myAsteriks == null) {
162         myAsteriks = tableRepresentative.table.fromChars(new char[]{'*'}, 0, 1);
163       }
164       return myAsteriks;
165     }
166
167     private void scanImports(JCTree.JCCompilationUnit compilationUnit, Set<JavacRefSymbol> symbols) {
168       for (JCTree.JCImport anImport : compilationUnit.getImports()) {
169         final JCTree.JCFieldAccess id = (JCTree.JCFieldAccess)anImport.getQualifiedIdentifier();
170         final Symbol sym = id.sym;
171         if (sym == null) {
172           final JCTree.JCExpression qExpr = id.getExpression();
173           if (qExpr instanceof JCTree.JCFieldAccess) {
174             final JCTree.JCFieldAccess classImport = (JCTree.JCFieldAccess)qExpr;
175             final Symbol ownerSym = classImport.sym;
176             final Name name = id.getIdentifier();
177             if (name != getAsteriskFromCurrentNameTable(name)) {
178               // member import
179               for (Symbol memberSymbol : ownerSym.members().getElements()) {
180                 if (memberSymbol.getSimpleName() == name) {
181                   symbols.add(new JavacRefSymbol(memberSymbol, Tree.Kind.IMPORT));
182                 }
183               }
184             }
185             collectClassImports(ownerSym, symbols);
186           }
187         } else {
188           // class import
189           collectClassImports(sym, symbols);
190         }
191       }
192     }
193   }
194
195   private static void collectClassImports(Symbol baseImport, Set<JavacRefSymbol> collector) {
196     for (Symbol symbol = baseImport;
197          symbol != null && symbol.getKind() != ElementKind.PACKAGE;
198          symbol = symbol.owner) {
199       collector.add(new JavacRefSymbol(symbol, Tree.Kind.IMPORT));
200     }
201   }
202 }