Kapt: support multiple annotation processing steps (KT-13651)

(cherry picked from commit 880e183)
2016-08-30 20:50:06 +03:00
parent 6ceaac63dc
commit 824b778a7b
9 changed files with 298 additions and 136 deletions
@@ -19,18 +19,16 @@ package org.jetbrains.kotlin.annotation
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.StandardFileSystems
 import com.intellij.psi.*
+import com.intellij.psi.impl.PsiModificationTrackerImpl
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.analyzer.AnalysisResult
+import org.jetbrains.kotlin.annotation.processing.RoundAnnotations
 import org.jetbrains.kotlin.annotation.processing.diagnostic.ErrorsAnnotationProcessing
 import org.jetbrains.kotlin.annotation.processing.impl.*
-import org.jetbrains.kotlin.asJava.findFacadeClass
-import org.jetbrains.kotlin.asJava.toLightClass
+import org.jetbrains.kotlin.config.APPEND_JAVA_SOURCE_ROOTS_HANDLER_KEY
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.java.model.elements.JeTypeElement
-import org.jetbrains.kotlin.java.model.internal.getAnnotationsWithInherited
-import org.jetbrains.kotlin.psi.KtClassOrObject
 import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.BindingTrace
 import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisCompletedHandlerExtension
 import java.io.File
@@ -65,6 +63,8 @@ abstract class AbstractAnnotationProcessingExtension(
        if (verboseOutput) messager.printMessage(Diagnostic.Kind.OTHER, "Kapt: " + message())
    }
    
+    private fun Int.count(noun: String) = if (this == 1) "$this $noun" else "$this ${noun}s"
+
    override fun analysisCompleted(
            project: Project,
            module: ModuleDescriptor,
@@ -80,53 +80,146 @@ abstract class AbstractAnnotationProcessingExtension(
            log { "No annotation processors detected, exiting" }
            return null
        }
-        
-        log { "Analysing Kotlin files: " + files.map { it.virtualFile.path } }
-        
-        val analysisContext = AnalysisContext(hashMapOf())
-        analysisContext.analyzeFiles(files)
-        
-        log { "Analysing Java source roots: $javaSourceRoots" }
-        
+
+        val appendJavaSourceRootsHandler = project.getUserData(APPEND_JAVA_SOURCE_ROOTS_HANDLER_KEY)
+        if (appendJavaSourceRootsHandler == null) {
+            log { "Java source root handler is not set, exiting" }
+            return null
+        }
+
        val psiManager = PsiManager.getInstance(project)
-        for (javaSourceRoot in javaSourceRoots) {
-            javaSourceRoot.walk().filter { it.isFile && it.extension == "java" }.forEach {
-                val vFile = StandardFileSystems.local().findFileByPath(it.absolutePath)
-                if (vFile != null) {
-                    val javaFile = psiManager.findFile(vFile) as? PsiJavaFile
-                    if (javaFile != null) {
-                        analysisContext.analyzeFile(javaFile)
+        val javaPsiFacade = JavaPsiFacade.getInstance(project)
+        val projectScope = GlobalSearchScope.projectScope(project)
+
+        val options = emptyMap<String, String>()
+        log { "Options: $options" }
+
+        val filer = KotlinFiler(generatedSourcesOutputDir, classesOutputDir)
+        val types = KotlinTypes(javaPsiFacade, PsiManager.getInstance(project), projectScope)
+        val elements = KotlinElements(javaPsiFacade, projectScope)
+
+        val processingEnvironment = KotlinProcessingEnvironment(
+                elements, types, messager, options, filer, processors,
+                project, psiManager, javaPsiFacade, projectScope, appendJavaSourceRootsHandler)
+
+        val processingResult = processingEnvironment.doAnnotationProcessing(files)
+
+        annotationProcessingComplete = true
+        log {
+            "Annotation processing complete, " + 
+                    processingResult.errorCount.count("error") + ", " + 
+                    processingResult.warningCount.count("warning")
+        }
+
+        if (processingResult.errorCount != 0) {
+            val reportFile = files.firstOrNull()
+            if (reportFile != null) {
+                bindingTrace.report(ErrorsAnnotationProcessing.ANNOTATION_PROCESSING_ERROR.on(reportFile))
+            }
+
+            // Do not restart Kotlin file analysis
+            return null
+        }
+        
+        if (!processingResult.wasAnythingGenerated) {
+            // Nothing was generated, do not need to restart analysis
+            return null
+        }
+
+        return AnalysisResult.RetryWithAdditionalJavaRoots(
+                bindingTrace.bindingContext,
+                module,
+                listOf(generatedSourcesOutputDir),
+                addToEnvironment = false)
+    }
+
+    private fun KotlinProcessingEnvironment.doAnnotationProcessing(files: Collection<KtFile>): ProcessingResult {
+        run initializeProcessors@ {
+            processors.forEach { it.init(this) }
+            log { "Initialized processors: " + processors.joinToString { it.javaClass.name } }
+        }
+
+        val firstRoundAnnotations = RoundAnnotations()
+        
+        run analyzeFilesForFirstRound@ {
+            log { "Analysing Kotlin files: " + files.map { it.virtualFile.path } }
+            firstRoundAnnotations.analyzeFiles(files)
+
+            log { "Analysing Java source roots: $javaSourceRoots" }
+            for (javaSourceRoot in javaSourceRoots) {
+                javaSourceRoot.walk().filter { it.isFile && it.extension == "java" }.forEach {
+                    val vFile = StandardFileSystems.local().findFileByPath(it.absolutePath)
+                    if (vFile != null) {
+                        val javaFile = psiManager.findFile(vFile) as? PsiJavaFile
+                        if (javaFile != null) {
+                            firstRoundAnnotations.analyzeFile(javaFile)
+                        }
                    }
                }
            }
        }
        
-        val options = emptyMap<String, String>()
-        log { "Options: $options" }
+        val finalRoundNumber = run annotationProcessing@ {
+            val firstRoundEnvironment = KotlinRoundEnvironment(firstRoundAnnotations, false, 1)
+            process(firstRoundEnvironment)
+        } + 1
        
-        val javaPsiFacade = JavaPsiFacade.getInstance(project)
-        val projectScope = GlobalSearchScope.projectScope(project)
+        log { "Starting round $finalRoundNumber (final)" }
+        val finalRoundEnvironment = KotlinRoundEnvironment(RoundAnnotations(), true, finalRoundNumber)
+        for (processor in processors) {
+            processor.process(emptySet(), finalRoundEnvironment)
+        }
+
+        return ProcessingResult(messager.errorCount, messager.warningCount, filer.wasAnythingGenerated)
+    }
+    
+    private tailrec fun KotlinProcessingEnvironment.process(roundEnvironment: KotlinRoundEnvironment): Int {
+        val newFiles = doRound(roundEnvironment)
+        val newJavaFiles = newFiles.filter { it.extension.toLowerCase() == "java" }
+        if (newJavaFiles.isEmpty() || messager.errorCount != 0) {
+            return roundEnvironment.roundNumber
+        }
        
-        val filer = KotlinFiler(generatedSourcesOutputDir, classesOutputDir)
-        val types = KotlinTypes(javaPsiFacade, PsiManager.getInstance(project), projectScope)
-        val elements = KotlinElements(javaPsiFacade, projectScope)
+        // Add new Java source roots after the first round
+        if (roundEnvironment.roundNumber == 1) {
+            appendJavaSourceRootsHandler(listOf(generatedSourcesOutputDir))
+        }
        
-        val processingEnvironment = KotlinProcessingEnvironment(elements, types, messager, options, filer)
-        processors.forEach { it.init(processingEnvironment) }
+        // Update the platform caches
+        (PsiManager.getInstance(project).modificationTracker as? PsiModificationTrackerImpl)?.incCounter()
        
-        log { "Initialized processors: " + processors.joinToString { it.javaClass.name } }
+        // Find generated files
+        val localFileSystem = StandardFileSystems.local()
+        val psiFiles = newJavaFiles
+                .map { localFileSystem.findFileByPath(it.absolutePath)?.let { psiManager.findFile(it) } }
+                .filterIsInstance<PsiJavaFile>()
+
+        if (psiFiles.isEmpty()) {
+            log { "Something is strange, files were generated but not found by PsiManager" }
+            return roundEnvironment.roundNumber
+        }
        
-        log { "Starting round 1" }
-        val round1Environment = KotlinRoundEnvironment(analysisContext)
+        // Start the next round
+        val nextRoundAnnotations = RoundAnnotations().apply { analyzeFiles(psiFiles) }
+        val nextRoundEnvironment = KotlinRoundEnvironment(nextRoundAnnotations, false, roundEnvironment.roundNumber + 1)
+        return process(nextRoundEnvironment)
+    }
+    
+    private fun KotlinProcessingEnvironment.doRound(roundEnvironment: KotlinRoundEnvironment): List<File> {
+        log { "Starting round ${roundEnvironment.roundNumber}" }
+        
+        val newFiles = mutableListOf<File>()
+        filer.onFileCreatedHandler = { newFiles += it }
+
        for (processor in processors) {
            val supportedAnnotationNames = processor.supportedAnnotationTypes
-            val supportsAnyAnnotation = supportedAnnotationNames.contains("*")
-            
-            val applicableAnnotationNames = when (supportsAnyAnnotation) {
-                true -> analysisContext.annotationsMap.keys
-                false -> processor.supportedAnnotationTypes.filter { it in analysisContext.annotationsMap } 
+            val acceptsAnyAnnotation = supportedAnnotationNames.contains("*")
+
+            val applicableAnnotationNames = when (acceptsAnyAnnotation) {
+                true -> roundEnvironment.annotationsMap.keys
+                false -> processor.supportedAnnotationTypes.filter { it in roundEnvironment.annotationsMap }
            }
-            
+
            if (applicableAnnotationNames.isEmpty()) {
                log { "Skipping processor " + processor.javaClass.name + ": no relevant annotations" }
                continue
@@ -135,89 +228,20 @@ abstract class AbstractAnnotationProcessingExtension(
            val applicableAnnotations = applicableAnnotationNames
                    .map { javaPsiFacade.findClass(it, projectScope)?.let { JeTypeElement(it) } }
                    .filterNotNullTo(hashSetOf())
-            
+
            log {
                val annotationNames = applicableAnnotations.joinToString { it.qualifiedName.toString() }
                "Processing with " + processor.javaClass.name + " (annotations: " + annotationNames + ")"
            }

-            processor.process(applicableAnnotations, round1Environment)
-        }
-
-        log { "Starting round 2 (final)" }
-        val round2Environment = KotlinRoundEnvironment(AnalysisContext(hashMapOf()), isProcessingOver = true)
-        for (processor in processors) {
-            processor.process(emptySet(), round2Environment)
+            processor.process(applicableAnnotations, roundEnvironment)
        }
        
-        fun Int.count(noun: String) = if (this == 1) "$this $noun" else "$this ${noun}s"
-        log { "Annotation processing complete, ${messager.errorCount.count("error")}, ${messager.warningCount.count("warning")}" }
-        
-        if (messager.errorCount != 0) {
-            val reportFile = files.firstOrNull()
-            if (reportFile != null) {
-                bindingTrace.report(ErrorsAnnotationProcessing.ANNOTATION_PROCESSING_ERROR.on(reportFile))
-            }
-            
-            // Do not restart analysis
-            return null
-        }
-
-        annotationProcessingComplete = true
-        return AnalysisResult.RetryWithAdditionalJavaRoots(bindingTrace.bindingContext, module, listOf(generatedSourcesOutputDir))
+        log { "Round ${roundEnvironment.roundNumber} finished, ${newFiles.size.count("file")} generated (${newFiles.joinToString()})" }
+        return newFiles
    }

    protected abstract fun loadAnnotationProcessors(): List<Processor>
 }

-internal class AnalysisContext(annotationsMap: MutableMap<String, MutableList<PsiModifierListOwner>>) {
-    private val mutableAnnotationsMap = annotationsMap
-    
-    val annotationsMap: Map<String, List<PsiModifierListOwner>>
-        get() = mutableAnnotationsMap
-    
-    fun analyzeFiles(files: Collection<KtFile>) {
-        for (file in files) {
-            analyzeFile(file)
-        }
-    }
-
-    fun analyzeFile(file: KtFile) {
-        val lightClass = file.findFacadeClass()
-
-        if (lightClass != null) {
-            analyzeDeclaration(lightClass)
-        }
-
-        for (declaration in file.declarations) {
-            if (declaration !is KtClassOrObject) continue
-            val clazz = declaration.toLightClass() ?: continue
-            analyzeDeclaration(clazz)
-        }
-    }
-
-    fun analyzeFile(file: PsiJavaFile) {
-        file.classes.forEach { analyzeDeclaration(it) }
-    }
-
-    fun analyzeDeclaration(declaration: PsiElement) {
-        if (declaration !is PsiModifierListOwner) return
-
-        for (annotation in declaration.getAnnotationsWithInherited()) {
-            val fqName = annotation.qualifiedName ?: return
-            mutableAnnotationsMap.getOrPut(fqName, { mutableListOf() }).add(declaration)
-        }
-
-        if (declaration is PsiClass) {
-            declaration.methods.forEach { analyzeDeclaration(it) }
-            declaration.fields.forEach { analyzeDeclaration(it) }
-            declaration.innerClasses.forEach { analyzeDeclaration(it) }
-        }
-
-        if (declaration is PsiMethod) {
-            for (parameter in declaration.parameterList.parameters) {
-                analyzeDeclaration(parameter)
-            }
-        }
-    }
-}
+private class ProcessingResult(val errorCount: Int, val warningCount: Int, val wasAnythingGenerated: Boolean)
@@ -0,0 +1,75 @@
+/*
+ * Copyright 2010-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.annotation.processing
+
+import com.intellij.psi.*
+import org.jetbrains.kotlin.asJava.findFacadeClass
+import org.jetbrains.kotlin.asJava.toLightClass
+import org.jetbrains.kotlin.java.model.internal.getAnnotationsWithInherited
+import org.jetbrains.kotlin.psi.KtClassOrObject
+import org.jetbrains.kotlin.psi.KtFile
+
+internal class RoundAnnotations() {
+    private val mutableAnnotationsMap = mutableMapOf<String, MutableList<PsiModifierListOwner>>()
+
+    val annotationsMap: Map<String, List<PsiModifierListOwner>>
+        get() = mutableAnnotationsMap
+
+    fun analyzeFiles(files: Collection<KtFile>) = files.forEach { analyzeFile(it) }
+    
+    @JvmName("analyzePsiJavaFiles")
+    fun analyzeFiles(files: Collection<PsiJavaFile>) = files.forEach { analyzeFile(it) }
+
+    fun analyzeFile(file: KtFile) {
+        val lightClass = file.findFacadeClass()
+
+        if (lightClass != null) {
+            analyzeDeclaration(lightClass)
+        }
+
+        for (declaration in file.declarations) {
+            if (declaration !is KtClassOrObject) continue
+            val clazz = declaration.toLightClass() ?: continue
+            analyzeDeclaration(clazz)
+        }
+    }
+
+    fun analyzeFile(file: PsiJavaFile) {
+        file.classes.forEach { analyzeDeclaration(it) }
+    }
+
+    fun analyzeDeclaration(declaration: PsiElement) {
+        if (declaration !is PsiModifierListOwner) return
+
+        for (annotation in declaration.getAnnotationsWithInherited()) {
+            val fqName = annotation.qualifiedName ?: return
+            mutableAnnotationsMap.getOrPut(fqName, { mutableListOf() }).add(declaration)
+        }
+
+        if (declaration is PsiClass) {
+            declaration.methods.forEach { analyzeDeclaration(it) }
+            declaration.fields.forEach { analyzeDeclaration(it) }
+            declaration.innerClasses.forEach { analyzeDeclaration(it) }
+        }
+
+        if (declaration is PsiMethod) {
+            for (parameter in declaration.parameterList.parameters) {
+                analyzeDeclaration(parameter)
+            }
+        }
+    }
+}
@@ -24,11 +24,18 @@ import javax.tools.JavaFileManager
 import javax.tools.JavaFileObject
 import javax.tools.StandardLocation

-class KotlinFiler(val generatedSourceDir: File, val classesOutputDir: File) : Filer {
+class KotlinFiler(
+        val generatedSourceDir: File, 
+        val classesOutputDir: File,
+        internal var onFileCreatedHandler: (File) -> Unit = {}
+) : Filer {
    private companion object {
        val PACKAGE_INFO_SUFFIX = ".packageInfo"
    }
    
+    internal var wasAnythingGenerated: Boolean = false
+        private set
+    
    private fun getGeneratedFile(nameCharSequence: CharSequence, extension: String): File {
        val name = nameCharSequence.toString()
        val isPackageInfo = name.endsWith(PACKAGE_INFO_SUFFIX)
@@ -47,8 +54,13 @@ class KotlinFiler(val generatedSourceDir: File, val classesOutputDir: File) : Fi
        return File(packageDir, fileName)
    }
    
+    private fun File.notifyCreated() = apply {
+        wasAnythingGenerated = true
+        onFileCreatedHandler(this)
+    }
+    
    override fun createSourceFile(name: CharSequence, vararg originatingElements: Element?): JavaFileObject {
-        return KotlinJavaFileObject(getGeneratedFile(name, ".java"))
+        return KotlinJavaFileObject(getGeneratedFile(name, ".java").notifyCreated())
    }

    override fun getResource(location: JavaFileManager.Location, pkg: CharSequence, relativeName: CharSequence): FileObject? {
@@ -63,7 +75,7 @@ class KotlinFiler(val generatedSourceDir: File, val classesOutputDir: File) : Fi
    ): FileObject? {
        val resourceFile = getResourceFile(location, pkg, relativeName)
        resourceFile.parentFile.mkdirs()
-        return KotlinFileObject(resourceFile)
+        return KotlinFileObject(resourceFile.notifyCreated())
    }
    
    private fun getResourceFile(location: JavaFileManager.Location, pkg: CharSequence, relativeName: CharSequence): File {
@@ -80,6 +92,6 @@ class KotlinFiler(val generatedSourceDir: File, val classesOutputDir: File) : Fi
    }

    override fun createClassFile(name: CharSequence, vararg originatingElements: Element?): JavaFileObject {
-        return KotlinJavaFileObject(getGeneratedFile(name, ".class"))
+        return KotlinJavaFileObject(getGeneratedFile(name, ".class").notifyCreated())
    }
 }
@@ -16,10 +16,14 @@

 package org.jetbrains.kotlin.annotation.processing.impl

+import com.intellij.openapi.project.Project
+import com.intellij.psi.JavaPsiFacade
+import com.intellij.psi.PsiManager
+import com.intellij.psi.search.GlobalSearchScope
+import java.io.File
 import java.util.*
-import javax.annotation.processing.Filer
-import javax.annotation.processing.Messager
 import javax.annotation.processing.ProcessingEnvironment
+import javax.annotation.processing.Processor
 import javax.lang.model.SourceVersion
 import javax.lang.model.util.Elements
 import javax.lang.model.util.Types
@@ -27,9 +31,17 @@ import javax.lang.model.util.Types
 class KotlinProcessingEnvironment(
        private val elements: Elements,
        private val types: Types,
-        private val messager: Messager,
+        private val messager: KotlinMessager,
        options: Map<String, String>,
-        private val filer: Filer
+        private val filer: KotlinFiler,
+        
+        internal val processors: List<Processor>,
+        
+        internal val project: Project,
+        internal val psiManager: PsiManager,
+        internal val javaPsiFacade: JavaPsiFacade,
+        internal val projectScope: GlobalSearchScope,
+        internal val appendJavaSourceRootsHandler: (List<File>) -> Unit
 ) : ProcessingEnvironment {
    private val options = Collections.unmodifiableMap(options) 
    
@@ -16,23 +16,29 @@

 package org.jetbrains.kotlin.annotation.processing.impl

-import org.jetbrains.kotlin.annotation.AnalysisContext
+import com.intellij.psi.PsiModifierListOwner
+import org.jetbrains.kotlin.annotation.processing.RoundAnnotations
 import org.jetbrains.kotlin.java.model.toJeElement
 import javax.annotation.processing.RoundEnvironment
 import javax.lang.model.element.Element
 import javax.lang.model.element.TypeElement

 internal class KotlinRoundEnvironment(
-        private val context: AnalysisContext,
-        private val isProcessingOver: Boolean = false) : RoundEnvironment {
+        private val roundAnnotations: RoundAnnotations,
+        private val isProcessingOver: Boolean,
+        internal val roundNumber: Int
+) : RoundEnvironment {
    private var isError = false
    
+    internal val annotationsMap: Map<String, List<PsiModifierListOwner>>
+        get() = roundAnnotations.annotationsMap
+    
    override fun getRootElements() = emptySet<Element>()
    
    override fun processingOver() = isProcessingOver

    private fun getElementsAnnotatedWith(fqName: String): Set<Element> {
-        val declarations = context.annotationsMap[fqName] ?: return emptySet()
+        val declarations = roundAnnotations.annotationsMap[fqName] ?: return emptySet()
        return hashSetOf<Element>().apply {
            for (declaration in declarations) {
                declaration.toJeElement()?.let { add(it) }