001/*
002 * Licensed under the Apache License, Version 2.0 (the "License");
003 * you may not use this file except in compliance with the License.
004 * You may obtain a copy of the License at
005 *
006 * http://www.apache.org/licenses/LICENSE-2.0
007 *
008 * Unless required by applicable law or agreed to in writing, software
009 * distributed under the License is distributed on an "AS IS" BASIS,
010 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
011 * See the License for the specific language governing permissions and
012 * limitations under the License.
013 */
014package org.atteo.classindex.processor;
015
016import java.io.BufferedReader;
017import java.io.File;
018import java.io.FileNotFoundException;
019import java.io.FileReader;
020import java.io.IOException;
021import java.io.Reader;
022import java.io.Writer;
023import java.lang.annotation.Annotation;
024import java.lang.annotation.Inherited;
025import java.util.Collections;
026import java.util.HashMap;
027import java.util.HashSet;
028import java.util.Map;
029import java.util.Set;
030import java.util.TreeSet;
031
032import javax.annotation.processing.AbstractProcessor;
033import javax.annotation.processing.Filer;
034import javax.annotation.processing.Messager;
035import javax.annotation.processing.ProcessingEnvironment;
036import javax.annotation.processing.RoundEnvironment;
037import javax.lang.model.SourceVersion;
038import javax.lang.model.element.AnnotationMirror;
039import javax.lang.model.element.Element;
040import javax.lang.model.element.PackageElement;
041import javax.lang.model.element.TypeElement;
042import javax.lang.model.type.DeclaredType;
043import javax.lang.model.type.TypeKind;
044import javax.lang.model.type.TypeMirror;
045import javax.lang.model.util.ElementScanner6;
046import javax.lang.model.util.Elements;
047import javax.lang.model.util.Types;
048import javax.tools.Diagnostic;
049import javax.tools.FileObject;
050import javax.tools.StandardLocation;
051
052import org.atteo.classindex.ClassIndex;
053import org.atteo.classindex.IndexAnnotated;
054import org.atteo.classindex.IndexSubclasses;
055
056/**
057 * Generates index files for {@link ClassIndex}.
058 */
059public class ClassIndexProcessor extends AbstractProcessor {
060    private Map<TypeElement, Set<String>> subclassMap = new HashMap<>();
061    private Map<TypeElement, Set<String>> annotatedMap = new HashMap<>();
062    private Map<PackageElement, Set<String>> packageMap = new HashMap<>();
063
064    private boolean annotationDriven = true;
065    private Set<String> indexedAnnotations = new HashSet<>();
066    private Set<String> indexedSuperclasses = new HashSet<>();
067    private Set<String> indexedPackages = new HashSet<>();
068    private Set<TypeElement> javadocAlreadyStored = new HashSet<>();
069
070    private Types types;
071    private Filer filer;
072    private Elements elementUtils;
073    private Messager messager;
074
075    public ClassIndexProcessor() {
076    }
077
078    /**
079     * Used when creating subclasses of the processor which will index some annotations
080     * which cannot be itself annotated with {@link IndexAnnotated} or {@link IndexSubclasses}.
081     *
082     * @param classes list of classes which the processor will be indexing
083     */
084    protected ClassIndexProcessor(Class<?>... classes) {
085        if (classes.length == 0) {
086            return;
087        }
088        annotationDriven = false;
089        for (Class<?> klass : classes) {
090            indexedAnnotations.add(klass.getCanonicalName());
091        }
092    }
093
094    /**
095     * Adds given annotations for indexing.
096     */
097    protected final void indexAnnotations(Class<?>... classes) {
098        for (Class<?> klass : classes) {
099            indexedAnnotations.add(klass.getCanonicalName());
100        }
101        annotationDriven = false;
102    }
103
104    /**
105     * Adds given classes for subclass indexing.
106     */
107    protected final void indexSubclasses(Class<?>... classes) {
108        for (Class<?> klass : classes) {
109            indexedSuperclasses.add(klass.getCanonicalName());
110        }
111        annotationDriven = false;
112    }
113
114    /**
115     * Adds given package for indexing.
116     */
117    protected final void indexPackages(String... packages) {
118        Collections.addAll(indexedPackages, packages);
119        annotationDriven = false;
120    }
121
122    @Override
123    public SourceVersion getSupportedSourceVersion() {
124        return SourceVersion.latest();
125    }
126
127    @Override
128    public Set<String> getSupportedAnnotationTypes() {
129        return Collections.singleton("*");
130    }
131
132    @Override
133    public synchronized void init(ProcessingEnvironment processingEnv) {
134        super.init(processingEnv);
135        types = processingEnv.getTypeUtils();
136        filer = processingEnv.getFiler();
137        elementUtils = processingEnv.getElementUtils();
138        messager = processingEnv.getMessager();
139    }
140
141    @Override
142    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
143        try {
144            for (Element element : roundEnv.getRootElements()) {
145                if (!(element instanceof TypeElement)) {
146                    continue;
147                }
148                final PackageElement packageElement = getPackage(element);
149                element.accept(new ElementScanner6<Void, Void>() {
150                    @Override
151                    public Void visitType(TypeElement typeElement, Void o) {
152                        try {
153                            for (AnnotationMirror mirror : typeElement.getAnnotationMirrors()) {
154                                final TypeElement annotationElement = (TypeElement) mirror.getAnnotationType().asElement();
155                                storeAnnotation(annotationElement, typeElement);
156                            }
157                            indexSupertypes(typeElement, typeElement);
158                            if (packageElement != null) {
159                                storeClassFromPackage(packageElement, typeElement);
160                            }
161                        } catch (IOException e) {
162                            messager.printMessage(Diagnostic.Kind.ERROR, "[ClassIndexProcessor] " + e.getMessage());
163                        }
164                        return super.visitType(typeElement, o);
165                    }
166                }, null);
167            }
168
169            if (!roundEnv.processingOver()) {
170                return false;
171            }
172
173            writeIndexFiles(ClassIndex.SUBCLASS_INDEX_PREFIX, subclassMap);
174            writeIndexFiles(ClassIndex.ANNOTATED_INDEX_PREFIX, annotatedMap);
175
176            for (Map.Entry<PackageElement, Set<String>> entry : packageMap.entrySet()) {
177                writeSimpleNameIndexFile(entry.getValue(), entry.getKey().getQualifiedName().toString()
178                        .replace(".", "/")
179                        + "/" + ClassIndex.PACKAGE_INDEX_NAME);
180            }
181        } catch (IOException e) {
182            messager.printMessage(Diagnostic.Kind.ERROR, "[ClassIndexProcessor] Can't write index file: " + e.getMessage());
183        } catch (Throwable e) {
184            e.printStackTrace();
185            messager.printMessage(Diagnostic.Kind.ERROR, "[ClassIndexProcessor] Internal error: " + e.getMessage());
186        }
187
188        return false;
189    }
190
191    private void writeIndexFiles(String prefix, Map<TypeElement, Set<String>> indexMap) throws IOException {
192        for (Map.Entry<TypeElement, Set<String>> entry : indexMap.entrySet()) {
193            writeSimpleNameIndexFile(entry.getValue(), prefix + entry.getKey().getQualifiedName().toString());
194        }
195    }
196
197    private void readOldIndexFile(Set<String> entries, String resourceName) throws IOException {
198        Reader reader = null;
199        try {
200            final FileObject resource = filer.getResource(StandardLocation.CLASS_OUTPUT, "", resourceName);
201            reader = resource.openReader(true);
202            readOldIndexFile(entries, reader);
203        } catch (FileNotFoundException e) {
204            /**
205             * Ugly hack for Intellij IDEA incremental compilation.
206             * The problem is that it throws on the files from the existence of FileNotFoundException, if they were not created during the
207             * current session of compilation.
208             */
209            final String realPath = e.getMessage();
210            if (new File(realPath).exists()) {
211                try (Reader fileReader = new FileReader(realPath)) {
212                    readOldIndexFile(entries, fileReader);
213                }
214            }
215        } catch (IOException e) {
216            // Thrown by Eclipse JDT when not found
217        } catch (UnsupportedOperationException e) {
218            // Java6 does not support reading old index files
219        } finally {
220            if (reader != null) {
221                reader.close();
222            }
223        }
224    }
225
226    private static void readOldIndexFile(Set<String> entries, Reader reader) throws IOException {
227        try (BufferedReader bufferedReader = new BufferedReader(reader)) {
228            String line = bufferedReader.readLine();
229            while (line != null) {
230                entries.add(line);
231                line = bufferedReader.readLine();
232            }
233        }
234    }
235
236    private void writeIndexFile(Set<String> entries, String resourceName) throws IOException {
237        FileObject file = filer.createResource(StandardLocation.CLASS_OUTPUT, "", resourceName);
238        try (Writer writer = file.openWriter()) {
239            for (String entry : entries) {
240                writer.write(entry);
241                writer.write("\n");
242            }
243        }
244    }
245
246    private void writeSimpleNameIndexFile(Set<String> elementList, String resourceName)
247            throws IOException {
248        readOldIndexFile(elementList, resourceName);
249        writeIndexFile(elementList, resourceName);
250    }
251
252    private void writeFile(String content, String resourceName) throws IOException {
253        FileObject file = filer.createResource(StandardLocation.CLASS_OUTPUT, "", resourceName);
254        try (Writer writer = file.openWriter()) {
255            writer.write(content);
256        }
257    }
258
259    /**
260     * Index super types for {@link IndexSubclasses} and any {@link IndexAnnotated}
261     * additionally accompanied by {@link Inherited}.
262     */
263    private void indexSupertypes(TypeElement rootElement, TypeElement element) throws IOException {
264
265        for (TypeMirror mirror : types.directSupertypes(element.asType())) {
266            if (mirror.getKind() != TypeKind.DECLARED) {
267                continue;
268            }
269
270            DeclaredType superType = (DeclaredType) mirror;
271            TypeElement superTypeElement = (TypeElement) superType.asElement();
272            storeSubclass(superTypeElement, rootElement);
273
274            for (AnnotationMirror annotationMirror : superTypeElement.getAnnotationMirrors()) {
275                TypeElement annotationElement = (TypeElement) annotationMirror.getAnnotationType()
276                        .asElement();
277
278                if (hasAnnotation(annotationElement, Inherited.class)) {
279                    storeAnnotation(annotationElement, rootElement);
280                }
281            }
282
283            indexSupertypes(rootElement, superTypeElement);
284        }
285    }
286
287    private boolean hasAnnotation(TypeElement element, Class<? extends Annotation> inheritedClass) {
288        try {
289            for (AnnotationMirror annotationMirror : element.getAnnotationMirrors()) {
290                if (annotationMirror.getAnnotationType().toString().equals(inheritedClass.getName())) {
291                    return true;
292                }
293            }
294        } catch (RuntimeException e) {
295            if (!e.getClass().getName().equals("com.sun.tools.javac.code.Symbol$CompletionFailure")) {
296                messager.printMessage(Diagnostic.Kind.ERROR, "[ClassIndexProcessor] Can't check annotation: " + e.getMessage());
297            }
298        }
299        return false;
300    }
301
302    private void storeAnnotation(TypeElement annotationElement, TypeElement rootElement) throws IOException {
303        if (indexedAnnotations.contains(annotationElement.getQualifiedName().toString())) {
304            putElement(annotatedMap, annotationElement, rootElement);
305        } else if (annotationDriven) {
306            IndexAnnotated indexAnnotated = annotationElement.getAnnotation(IndexAnnotated.class);
307            if (indexAnnotated != null) {
308                putElement(annotatedMap, annotationElement, rootElement);
309                if (indexAnnotated.storeJavadoc()) {
310                    storeJavadoc(rootElement);
311                }
312            }
313        }
314    }
315
316    private void storeSubclass(TypeElement superTypeElement, TypeElement rootElement) throws IOException {
317        if (indexedSuperclasses.contains(superTypeElement.getQualifiedName().toString())) {
318            putElement(subclassMap, superTypeElement, rootElement);
319        } else if (annotationDriven) {
320            IndexSubclasses indexSubclasses = superTypeElement.getAnnotation(IndexSubclasses.class);
321            if (indexSubclasses != null) {
322                putElement(subclassMap, superTypeElement, rootElement);
323
324                if (indexSubclasses.storeJavadoc()) {
325                    storeJavadoc(rootElement);
326                }
327            }
328        }
329        if (indexedSuperclasses.contains(superTypeElement.getQualifiedName().toString())
330                || (annotationDriven && superTypeElement.getAnnotation(IndexSubclasses.class) != null)) {
331            putElement(subclassMap, superTypeElement, rootElement);
332        }
333    }
334
335    private void storeClassFromPackage(PackageElement packageElement, TypeElement rootElement) throws IOException {
336        if (indexedPackages.contains(packageElement.getQualifiedName().toString())) {
337            putElement(packageMap, packageElement, rootElement);
338        } else if (annotationDriven) {
339            IndexSubclasses indexSubclasses = packageElement.getAnnotation(IndexSubclasses.class);
340            if (indexSubclasses != null) {
341                String simpleName = getShortName(rootElement);
342                if (simpleName != null) {
343                    putElement(packageMap, packageElement, simpleName);
344                    if (indexSubclasses.storeJavadoc()) {
345                        storeJavadoc(rootElement);
346                    }
347                }
348            }
349        }
350    }
351
352    private <K> void putElement(Map<K, Set<String>> map, K keyElement, TypeElement valueElement) {
353        final String fullName = getFullName(valueElement);
354        if (fullName != null) {
355            putElement(map, keyElement, fullName);
356        }
357    }
358
359    private <K> void putElement(Map<K, Set<String>> map, K keyElement, String valueElement) {
360        Set<String> set = map.get(keyElement);
361        if (set == null) {
362            set = new TreeSet<>();
363            map.put(keyElement, set);
364        }
365        set.add(valueElement);
366    }
367
368    private String getFullName(TypeElement typeElement) {
369        switch (typeElement.getNestingKind()) {
370            case TOP_LEVEL:
371                return typeElement.getQualifiedName().toString();
372            case MEMBER:
373                final Element enclosingElement = typeElement.getEnclosingElement();
374                if (enclosingElement instanceof TypeElement) {
375                    final String enclosingName = getFullName(((TypeElement) enclosingElement));
376                    if (enclosingName != null) {
377                        return enclosingName + '$' + typeElement.getSimpleName().toString();
378                    }
379                }
380                return null;
381            case ANONYMOUS:
382            case LOCAL:
383            default:
384                return null;
385        }
386    }
387
388    private String getShortName(TypeElement typeElement) {
389        switch (typeElement.getNestingKind()) {
390            case TOP_LEVEL:
391                return typeElement.getSimpleName().toString();
392            case MEMBER:
393                final Element enclosingElement = typeElement.getEnclosingElement();
394                if (enclosingElement instanceof TypeElement) {
395                    final String enclosingName = getShortName(((TypeElement) enclosingElement));
396                    if (enclosingName != null) {
397                        return enclosingName + '$' + typeElement.getSimpleName().toString();
398                    }
399                }
400                return null;
401            case ANONYMOUS:
402            case LOCAL:
403            default:
404                return null;
405        }
406    }
407
408    private PackageElement getPackage(Element typeElement) {
409        Element element = typeElement;
410        while (element != null) {
411            if (element instanceof PackageElement) {
412                return (PackageElement) element;
413            }
414            element = element.getEnclosingElement();
415        }
416        return null;
417    }
418
419    private void storeJavadoc(TypeElement element) throws IOException {
420        if (javadocAlreadyStored.contains(element)) {
421            return;
422        }
423        javadocAlreadyStored.add(element);
424
425        String docComment = elementUtils.getDocComment(element);
426        if (docComment == null) {
427            return;
428        }
429        writeFile(docComment, ClassIndex.JAVADOC_PREFIX + element.getQualifiedName().toString());
430    }
431
432
433}