Initial annotation processing support in Gradle

libraries/tools/kotlin-gradle-plugin-core/src/main/kotlin/org/jetbrains/kotlin/gradle/tasks/Tasks.kt

@@ -31,9 +31,21 @@ import org.jetbrains.kotlin.idea.JetFileType
 import org.jetbrains.kotlin.utils.LibraryUtils
 import com.intellij.openapi.util.io.FileUtil
 import org.apache.commons.io.FilenameUtils
+import org.gradle.api.file.FileCollection
+import org.gradle.api.tasks.compile.JavaCompile
+import org.jetbrains.org.objectweb.asm.ClassWriter
+import java.io.IOException
+import java.lang.ref.WeakReference
+import java.util.zip.ZipFile
+import org.jetbrains.org.objectweb.asm.*
+import org.jetbrains.org.objectweb.asm.Opcodes.*
 
 val DEFAULT_ANNOTATIONS = "org.jebrains.kotlin.gradle.defaultAnnotations"
 
+val ANNOTATIONS_PLUGIN_NAME = "org.jetbrains.kotlin.kapt"
+
+val JAVA_FQNAME_PATTERN = "^([\\p{L}_$][\\p{L}\\p{N}_$]*\\.)*[\\p{L}_$][\\p{L}\\p{N}_$]*$".toRegex()
+
 abstract class AbstractKotlinCompile<T : CommonCompilerArguments>() : AbstractCompile() {
     abstract protected val compiler: CLICompiler<T>
     abstract protected fun createBlankArgs(): T
@@ -98,12 +110,11 @@ public open class KotlinCompile() : AbstractKotlinCompile<K2JVMCompilerArguments
     override val compiler = K2JVMCompiler()
     override fun createBlankArgs(): K2JVMCompilerArguments = K2JVMCompilerArguments()
 
+    public var aptFiles: Set<File> = emptySet()
+
     val srcDirsSources = HashSet<SourceDirectorySet>()
 
     override fun populateTargetSpecificArgs(args: K2JVMCompilerArguments) {
-        val extraProperties = getExtensions().getExtraProperties()
-        args.pluginClasspaths = extraProperties.get("compilerPluginClasspaths") as? Array<String>
-        args.pluginOptions = extraProperties.get("compilerPluginArguments") as? Array<String>
         // show kotlin compiler where to look for java source files
         args.freeArgs = (args.freeArgs + getJavaSourceRoots().map { it.getAbsolutePath() }).toSet().toList()
 
@@ -117,6 +128,20 @@ public open class KotlinCompile() : AbstractKotlinCompile<K2JVMCompilerArguments
             kotlinOptions.destination
         }
 
+        val extraProperties = getExtensions().getExtraProperties()
+        args.pluginClasspaths = extraProperties.get("compilerPluginClasspaths") as? Array<String>
+        val basePluginOptions = (extraProperties.get("compilerPluginArguments") as? Array<String>) ?: array()
+
+        val pluginOptions = arrayListOf(*basePluginOptions)
+
+        if (aptFiles.isNotEmpty()) {
+            val annotationDeclarationsFile = File(args.destination, "annotations.txt")
+            if (annotationDeclarationsFile.exists()) annotationDeclarationsFile.delete()
+            pluginOptions.add("plugin:$ANNOTATIONS_PLUGIN_NAME:output=" + annotationDeclarationsFile)
+        }
+
+        args.pluginOptions = pluginOptions.toTypedArray()
+
         val embeddedAnnotations = getAnnotations(getProject(), getLogger())
         val userAnnotations = kotlinOptions.annotations?.split(File.pathSeparatorChar)?.toList() ?: emptyList()
         val allAnnotations = if (kotlinOptions.noJdkAnnotations) userAnnotations else userAnnotations.plus(embeddedAnnotations.map { it.getPath() })
@@ -149,6 +174,127 @@ public open class KotlinCompile() : AbstractKotlinCompile<K2JVMCompilerArguments
         if (outputDirFile.exists()) {
             FileUtils.copyDirectory(outputDirFile, getDestinationDir())
         }
+
+        [suppress("UNCHECKED_CAST")]
+        val javaTask = (getExtensions().getExtraProperties().get("javaTask") as? WeakReference<JavaCompile>)?.get()
+        if (javaTask != null && aptFiles.isNotEmpty()) {
+            val aptStubsDir = File(outputDirFile, "0apt")
+
+            val annotationDeclarationsFile = File(aptStubsDir, "annotations.txt")
+
+            val javaAptSourceDir = File(aptStubsDir, "java_src")
+            val stubOutputDir = File(outputDirFile, "wrappers")
+            javaAptSourceDir.mkdirs()
+            stubOutputDir.mkdirs()
+
+            javaTask.source(javaAptSourceDir)
+            javaTask.setClasspath(javaTask.getClasspath() + getProject().files(stubOutputDir))
+
+            // Generate a simple Java source file with annotation to launch Java APT even if there's no Java files
+            javaAptSourceDir.mkdirs()
+            val javaHackPackageDir = File(javaAptSourceDir, "__gen/annotation")
+            val javaHackClFile = File(javaHackPackageDir, "Cl.java")
+            javaHackClFile.writeText("package __gen.annotation;" + "class Cl { @javax.inject.Inject boolean v; }")
+
+            val annotationProcessors = lookupAnnotationProcessors(aptFiles)
+            for (processor in annotationProcessors) {
+                generateAnnotationProcessorWrapper(processor, stubOutputDir)
+            }
+
+            val annotationProcessorWrappers = annotationProcessors
+                    .map { "__gen.AnnotationProcessorWrapper_${it.replace('.', '_')}" }
+                    .joinToString(",")
+
+            val compilerArgs = javaTask.getOptions().getCompilerArgs()
+            val processorArgIndex = compilerArgs.indexOfFirst { "-processor" == it }
+
+            // Already has a "-processor" argument (and it is not the last one)
+            if (processorArgIndex >= 0 && compilerArgs.size() > (processorArgIndex + 1)) {
+                compilerArgs[processorArgIndex + 1] =
+                        compilerArgs[processorArgIndex + 1] + "," + annotationProcessorWrappers
+            }
+            else {
+                compilerArgs.add("-processor")
+                compilerArgs.add(annotationProcessorWrappers)
+            }
+
+            javaTask.getOptions().setCompilerArgs(compilerArgs)
+
+            javaTask.doLast {
+                annotationDeclarationsFile.delete()
+                javaAptSourceDir.deleteRecursively()
+            }
+        }
+    }
+
+    private fun lookupAnnotationProcessors(files: Set<File>): Set<String> {
+        fun withZipFile(file: File, job: (ZipFile) -> Unit) {
+            var zipFile: ZipFile? = null
+            try {
+                zipFile = ZipFile(file)
+                job(zipFile)
+            }
+            catch (e: IOException) {
+                // Do nothing (do not continue to search for annotation processors on error)
+            }
+            catch (e: IllegalStateException) {
+                // ZipFile was already closed for some reason
+            }
+            finally {
+                try {
+                    zipFile?.close()
+                }
+                catch (e: IOException) {}
+            }
+        }
+
+        val annotationProcessors = hashSetOf<String>()
+
+        fun processLines(lines: Sequence<String>) {
+            for (line in lines) {
+                if (line.isBlank() || !JAVA_FQNAME_PATTERN.matcher(line).matches()) continue
+                annotationProcessors.add(line)
+            }
+        }
+
+        for (file in files) {
+            withZipFile(file) { zipFile ->
+                val entry = zipFile.getEntry("META-INF/services/javax.annotation.processing.Processor")
+                if (entry != null) {
+                    zipFile.getInputStream(entry).reader().useLines { lines ->
+                        processLines(lines)
+                    }
+                }
+            }
+        }
+
+        return annotationProcessors
+    }
+
+    private fun generateAnnotationProcessorWrapper(processorFqName: String, outputDirectory: File) {
+        val className = "AnnotationProcessorWrapper_${processorFqName.replace('.', '_')}"
+
+        val bytes = with (ClassWriter(0)) {
+            val superClass = "org/kotlin/annotations/AnnotationProcessorWrapper"
+
+            visit(49, ACC_PUBLIC + ACC_SUPER, className, null,
+                    superClass, null)
+
+            visitSource("$className.java", null)
+
+            with (visitMethod(ACC_PUBLIC, "<init>", "()V", null, null)) {
+                visitVarInsn(ALOAD, 0)
+                visitLdcInsn(processorFqName)
+                visitMethodInsn(INVOKESPECIAL, superClass, "<init>", "(Ljava.lang.String;)V", false)
+                visitInsn(RETURN)
+                visitMaxs(2, 1)
+                visitEnd()
+            }
+
+            visitEnd()
+            toByteArray()
+        }
+        File(outputDirectory, "$className.class").writeBytes(bytes)
     }
 
     // override setSource to track source directory sets

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/plugin/KotlinPlugin.kt

@@ -38,16 +38,17 @@ import org.gradle.api.file.SourceDirectorySet
 import kotlin.properties.Delegates
 import org.gradle.api.tasks.Delete
 import groovy.lang.Closure
+import org.gradle.api.artifacts.Configuration
 import org.gradle.api.file.FileCollection
 import org.jetbrains.kotlin.gradle.tasks.KotlinTasksProvider
 import java.util.ServiceLoader
 import org.gradle.api.logging.*
 import java.net.URL
 import java.util.jar.Manifest
+import java.lang.ref.WeakReference
 
 val DEFAULT_ANNOTATIONS = "org.jebrains.kotlin.gradle.defaultAnnotations"
 
-
 abstract class KotlinSourceSetProcessor<T : AbstractCompile>(
         val project: ProjectInternal,
         val javaBasePlugin: JavaBasePlugin,
@@ -300,23 +301,39 @@ open class KotlinAndroidPlugin [Inject] (val scriptHandler: ScriptHandler, val t
 
         (ext as ExtensionAware).getExtensions().add("kotlinOptions", tasksProvider.kotlinJVMOptionsClass)
 
+        fun addConfiguration(name: String): Configuration =
+                project.getConfigurations().create(name)
+
+        val kotlinAptConfiguration = addConfiguration("kotlinApt")
+        val kotlinAptAndroidTestConfiguration = addConfiguration("kotlinAptAndroidTest")
+        val kotlinAptTestConfiguration = addConfiguration("kotlinAptTest")
+
+        val aptConfigurations = mapOf(
+                "base" to kotlinAptConfiguration,
+                "androidTest" to kotlinAptAndroidTestConfiguration,
+                "unitTest" to kotlinAptTestConfiguration)
+
         project afterEvaluate { project ->
             if (project != null) {
                 val plugin = (project.getPlugins().findPlugin("android")
                                 ?: project.getPlugins().findPlugin("android-library")) as BasePlugin
 
-                processVariantData(plugin.getVariantManager().getVariantDataList(), project, ext, plugin)
+                val aptConfigurationDependencyFiles = aptConfigurations.mapValues { it.getValue().resolve() }
+
+                processVariantData(plugin.getVariantManager().getVariantDataList(), project, ext, plugin, aptConfigurationDependencyFiles)
             }
         }
 
-        project.getExtensions().add(DEFAULT_ANNOTATIONS, GradleUtils(scriptHandler, project).resolveKotlinPluginDependency("kotlin-android-sdk-annotations"))
+        project.getExtensions().add(DEFAULT_ANNOTATIONS, GradleUtils(scriptHandler, project)
+                .resolveKotlinPluginDependency("kotlin-android-sdk-annotations"))
     }
 
     private fun processVariantData(
             variantDataList: List<BaseVariantData<out BaseVariantOutputData>>,
             project: Project,
             androidExt: BaseExtension,
-            androidPlugin: BasePlugin
+            androidPlugin: BasePlugin,
+            aptConfigurations: Map<String, Set<File>>
     ) {
         val logger = project.getLogger()
         val kotlinOptions = getExtension<Any?>(androidExt, "kotlinOptions")
@@ -375,12 +392,25 @@ open class KotlinAndroidPlugin [Inject] (val scriptHandler: ScriptHandler, val t
 
             subpluginEnvironment.addSubpluginArguments(project, kotlinTask)
 
+            fun filesForConfigurations(configurationIds: Array<String>): Set<File> {
+                return configurationIds.flatMap { aptConfigurations[it]?.toList() ?: listOf() }.distinct()
+            }
+
+            val aptFiles = filesForConfigurations(when (variantData.getType()) {
+                VariantType.UNIT_TEST -> arrayOf("base", "unitTest")
+                VariantType.ANDROID_TEST -> arrayOf("base", "androidTest")
+                else -> arrayOf("base")
+            })
+
+            kotlinTask.setProperty("aptFiles", aptFiles)
+
             kotlinTask doFirst {
                 val androidRT = project.files(AndroidGradleWrapper.getRuntimeJars(androidPlugin, androidExt))
                 val fullClasspath = (javaTask.getClasspath() + androidRT) - project.files(kotlinTask.property("kotlinDestinationDir"))
                 (it as AbstractCompile).setClasspath(fullClasspath)
             }
 
+            kotlinTask.getExtensions().getExtraProperties().set("javaTask", WeakReference(javaTask))
             javaTask.dependsOn(kotlinTaskName)
 
             javaTask doFirst  {