Kapt: Attach generated Kotlin sources in 'compilation' mode (KT-32535)

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/KotlinCoreEnvironment.kt

@@ -293,6 +293,11 @@ class KotlinCoreEnvironment private constructor(
         }
     }
 
+    fun addKotlinSourceRoots(rootDirs: List<File>) {
+        val roots = rootDirs.map { KotlinSourceRoot(it.absolutePath, isCommon = false) }
+        sourceFiles += createSourceFilesFromSourceRoots(configuration, project, roots)
+    }
+
     fun createPackagePartProvider(scope: GlobalSearchScope): JvmPackagePartProvider {
         return JvmPackagePartProvider(configuration.languageVersionSettings, scope).apply {
             addRoots(initialRoots, configuration.getNotNull(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY))

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/KotlinToJVMBytecodeCompiler.kt

@@ -459,7 +459,7 @@ object KotlinToJVMBytecodeCompiler {
         environment: KotlinCoreEnvironment,
         targetDescription: String?
     ): AnalysisResult? {
-        if (result is AnalysisResult.RetryWithAdditionalJavaRoots) {
+        if (result is AnalysisResult.RetryWithAdditionalRoots) {
             val configuration = environment.configuration
 
             val oldReadOnlyValue = configuration.isReadOnly
@@ -471,6 +471,10 @@ object KotlinToJVMBytecodeCompiler {
                 environment.updateClasspath(result.additionalJavaRoots.map { JavaSourceRoot(it, null) })
             }
 
+            if (result.additionalKotlinRoots.isNotEmpty()) {
+                environment.addKotlinSourceRoots(result.additionalKotlinRoots)
+            }
+
             KotlinJavaPsiFacade.getInstance(environment.project).clearPackageCaches()
 
             // Clear all diagnostic messages
@@ -562,7 +566,7 @@ object KotlinToJVMBytecodeCompiler {
 
         val analysisResult = analyzerWithCompilerReport.analysisResult
 
-        return if (!analyzerWithCompilerReport.hasErrors() || analysisResult is AnalysisResult.RetryWithAdditionalJavaRoots)
+        return if (!analyzerWithCompilerReport.hasErrors() || analysisResult is AnalysisResult.RetryWithAdditionalRoots)
             analysisResult
         else
             null

compiler/frontend/src/org/jetbrains/kotlin/analyzer/AnalysisResult.kt

@@ -67,10 +67,11 @@ open class AnalysisResult protected constructor(
         val exception: Throwable
     ) : AnalysisResult(bindingContext, ErrorUtils.getErrorModule())
 
-    class RetryWithAdditionalJavaRoots(
+    class RetryWithAdditionalRoots(
         bindingContext: BindingContext,
         moduleDescriptor: ModuleDescriptor,
         val additionalJavaRoots: List<File>,
+        val additionalKotlinRoots: List<File>,
         val addToEnvironment: Boolean = true
     ) : AnalysisResult(bindingContext, moduleDescriptor)
 

plugins/kapt3/kapt3-cli/test/org/jetbrains/kotlin/kapt/cli/test/KaptToolIntegrationTestGenerated.java

@@ -44,6 +44,16 @@ public class KaptToolIntegrationTestGenerated extends AbstractKaptToolIntegratio
         runTest("plugins/kapt3/kapt3-cli/testData/integration/correctErrorTypesOn/");
     }
 
+    @TestMetadata("kotlinFileGeneration")
+    public void testKotlinFileGeneration() throws Exception {
+        runTest("plugins/kapt3/kapt3-cli/testData/integration/kotlinFileGeneration/");
+    }
+
+    @TestMetadata("kotlinFileGenerationDefaultOutput")
+    public void testKotlinFileGenerationDefaultOutput() throws Exception {
+        runTest("plugins/kapt3/kapt3-cli/testData/integration/kotlinFileGenerationDefaultOutput/");
+    }
+
     @TestMetadata("separateStubAptCompilation")
     public void testSeparateStubAptCompilation() throws Exception {
         runTest("plugins/kapt3/kapt3-cli/testData/integration/separateStubAptCompilation/");

plugins/kapt3/kapt3-cli/testData/integration/kotlinFileGeneration/Test.kt

@@ -0,0 +1,19 @@
+package test
+
+import apt.Anno
+import generated.Test as TestGenerated
+
+@Anno
+class Test {
+    @field:Anno
+    val property: String = ""
+
+    @Anno
+    fun function() {
+
+    }
+}
+
+fun main() {
+    println("Generated class: " + TestGenerated::class.java.name)
+}

plugins/kapt3/kapt3-cli/testData/integration/kotlinFileGeneration/ap/META-INF/services/javax.annotation.processing.Processor

@@ -0,0 +1 @@
+apt.SampleApt

plugins/kapt3/kapt3-cli/testData/integration/kotlinFileGeneration/ap/Processor.kt

@@ -0,0 +1,37 @@
+package apt
+
+import java.io.File
+import javax.annotation.processing.*
+import javax.lang.model.SourceVersion
+import javax.lang.model.element.TypeElement
+import javax.tools.Diagnostic.Kind.*
+
+annotation class Anno
+
+class SampleApt : AbstractProcessor() {
+    private companion object {
+        const val KAPT_KOTLIN_GENERATED_OPTION_NAME = "kapt.kotlin.generated"
+    }
+
+    override fun process(annotations: Set<TypeElement>, roundEnv: RoundEnvironment): Boolean {
+        val kaptKotlinGeneratedDir = processingEnv.options[KAPT_KOTLIN_GENERATED_OPTION_NAME] ?: run {
+            processingEnv.messager.printMessage(ERROR, "Can't find the target directory for generated Kotlin files.")
+            return false
+        }
+
+        val baseDir = File(kaptKotlinGeneratedDir, "generated")
+        baseDir.mkdirs()
+
+        for (element in roundEnv.getElementsAnnotatedWith(Anno::class.java)) {
+            val generatedSimpleName = element.simpleName.toString().capitalize()
+            val file = File(baseDir, "$generatedSimpleName.kt")
+            file.writeText("package generated\nclass $generatedSimpleName")
+        }
+
+        return true
+    }
+
+    override fun getSupportedOptions() = emptySet<String>()
+    override fun getSupportedSourceVersion() = SourceVersion.RELEASE_8
+    override fun getSupportedAnnotationTypes() = setOf("apt.Anno")
+}

plugins/kapt3/kapt3-cli/testData/integration/kotlinFileGeneration/build.txt

@@ -0,0 +1,32 @@
+# mkdir
+output/ap
+output/stubs
+output/classes
+output/sources
+output/kotlin-sources
+
+# kotlinc
+-cp %KOTLIN_STDLIB%
+-d output/ap
+ap/Processor.kt
+
+# copy
+ap/META-INF/services/javax.annotation.processing.Processor
+output/ap/META-INF/services/javax.annotation.processing.Processor
+
+# kapt
+-Kapt-stubs=output/stubs
+-Kapt-classes=output/classes
+-Kapt-sources=output/sources
+-Kapt-classpath=output/ap
+-Kapt-option:kapt.kotlin.generated=output/kotlin-sources
+-d output/classes
+-cp output/ap:%KOTLIN_STDLIB%
+Test.kt
+
+# java
+-cp output/classes:output/ap:%KOTLIN_STDLIB%
+test.TestKt
+
+# after
+Generated class: generated.Test

plugins/kapt3/kapt3-cli/testData/integration/kotlinFileGenerationDefaultOutput/build.txt

@@ -0,0 +1,36 @@
+# copy
+../simple/ap
+ap
+
+# copy
+../simple/Test.kt
+Test.kt
+
+# mkdir
+output/ap
+output/stubs
+output/classes
+output/sources
+
+# kotlinc
+-cp %KOTLIN_STDLIB%
+-d output/ap
+ap/Processor.kt
+
+# kapt
+-Kapt-stubs=output/stubs
+-Kapt-classes=output/classes
+-Kapt-sources=output/sources
+-Kapt-classpath=output/ap
+-Kapt-processors=apt.SampleApt
+-Kapt-option:kapt.test.writeKotlinFiles=true
+-d output/classes
+-cp output/ap:%KOTLIN_STDLIB%
+Test.kt
+
+# java
+-cp output/classes:output/ap:%KOTLIN_STDLIB%
+test.TestKt
+
+# after
+Generated class: generated.Test

plugins/kapt3/kapt3-cli/testData/integration/simple/ap/Processor.kt

@@ -4,14 +4,23 @@ import javax.annotation.processing.*
 import javax.lang.model.SourceVersion
 import javax.lang.model.element.TypeElement
 import javax.tools.Diagnostic.Kind.*
+import javax.tools.StandardLocation
 
 annotation class Anno
 
 class SampleApt : AbstractProcessor() {
     override fun process(annotations: Set<TypeElement>, roundEnv: RoundEnvironment): Boolean {
+        val writeKotlinFiles = processingEnv.options["kapt.test.writeKotlinFiles"] == "true"
+
         for (element in roundEnv.getElementsAnnotatedWith(Anno::class.java)) {
             val generatedSimpleName = element.simpleName.toString().capitalize()
-            processingEnv.filer.createSourceFile("generated.$generatedSimpleName").openWriter().use {
+
+            val file = when (writeKotlinFiles) {
+                true -> processingEnv.filer.createResource(StandardLocation.SOURCE_OUTPUT, "generated", "$generatedSimpleName.kt")
+                false -> processingEnv.filer.createSourceFile("generated.$generatedSimpleName")
+            }
+
+            file.openWriter().use {
                 it.write("package generated;\npublic class $generatedSimpleName {}")
             }
         }
@@ -19,7 +28,7 @@ class SampleApt : AbstractProcessor() {
         return true
     }
 
-    override fun getSupportedOptions() = emptySet<String>()
+    override fun getSupportedOptions() = setOf("kapt.test.writeKotlinFiles")
     override fun getSupportedSourceVersion() = SourceVersion.RELEASE_8
     override fun getSupportedAnnotationTypes() = setOf("apt.Anno")
 }

plugins/kapt3/kapt3-compiler/src/org/jetbrains/kotlin/kapt3/Kapt3Extension.kt

@@ -65,7 +65,8 @@ import java.io.StringWriter
 import java.io.Writer
 import java.net.URLClassLoader
 import javax.annotation.processing.Processor
-import com.sun.tools.javac.util.List as JavacList
+
+private const val KAPT_KOTLIN_GENERATED_OPTION_NAME = "kapt.kotlin.generated"
 
 class ClasspathBasedKapt3Extension(
     options: KaptOptions,
@@ -204,15 +205,21 @@ abstract class AbstractKapt3Extension(
         return if (options.mode != WITH_COMPILATION) {
             doNotGenerateCode()
         } else {
-            AnalysisResult.RetryWithAdditionalJavaRoots(
+            AnalysisResult.RetryWithAdditionalRoots(
                 bindingTrace.bindingContext,
                 module,
                 listOf(options.sourcesOutputDir),
+                listOfNotNull(options.sourcesOutputDir, getKotlinGeneratedSourcesDirectory()),
                 addToEnvironment = true
             )
         }
     }
 
+    private fun getKotlinGeneratedSourcesDirectory(): File? {
+        val value = options.processingOptions[KAPT_KOTLIN_GENERATED_OPTION_NAME] ?: return null
+        return File(value).takeIf { it.exists() }
+    }
+
     private fun runAnnotationProcessing(kaptContext: KaptContext, processors: LoadedProcessors) {
         if (!options.mode.runAnnotationProcessing) return