Kapt: Annotations on enum constants are not kept on the generated stub (KT-21433)

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

@@ -18,7 +18,11 @@ package org.jetbrains.kotlin.kapt3
 
 import com.intellij.ide.ClassUtilCore
 import com.intellij.openapi.project.Project
+import com.sun.tools.javac.code.Flags
 import com.sun.tools.javac.tree.JCTree
+import com.sun.tools.javac.tree.Pretty
+import com.sun.tools.javac.tree.TreeMaker
+import com.sun.tools.javac.util.Context
 import org.jetbrains.kotlin.analyzer.AnalysisResult
 import org.jetbrains.kotlin.backend.common.output.OutputFile
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity.OUTPUT
@@ -44,6 +48,8 @@ import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.BindingTrace
 import org.jetbrains.kotlin.resolve.jvm.extensions.PartialAnalysisHandlerExtension
 import java.io.File
+import java.io.StringWriter
+import java.io.Writer
 import java.net.URLClassLoader
 import java.util.*
 import javax.annotation.processing.Processor
@@ -264,7 +270,7 @@ abstract class AbstractKapt3Extension(
         logger.info { "Java stub generation took $stubGenerationTime ms" }
         logger.info { "Stubs for Kotlin classes: " + kotlinSourceStubs.joinToString { it.sourcefile.name } }
 
-        saveStubs(kotlinSourceStubs)
+        saveStubs(kaptContext, kotlinSourceStubs)
         saveIncrementalData(kaptContext, logger.messageCollector, converter)
     }
 
@@ -277,14 +283,14 @@ abstract class AbstractKapt3Extension(
         return javaFilesFromJavaSourceRoots
     }
 
-    protected open fun saveStubs(stubs: JavacList<JCTree.JCCompilationUnit>) {
+    protected open fun saveStubs(kaptContext: KaptContext<*>, stubs: JavacList<JCTree.JCCompilationUnit>) {
         for (stub in stubs) {
             val className = (stub.defs.first { it is JCTree.JCClassDecl } as JCTree.JCClassDecl).simpleName.toString()
 
             val packageName = stub.getPackageNameJava9Aware()?.toString() ?: ""
             val packageDir = if (packageName.isEmpty()) stubsOutputDir else File(stubsOutputDir, packageName.replace('.', '/'))
             packageDir.mkdirs()
-            File(packageDir, className + ".java").writeText(stub.toString())
+            File(packageDir, className + ".java").writeText(stub.prettyPrint(kaptContext.context))
         }
     }
 
@@ -314,6 +320,25 @@ abstract class AbstractKapt3Extension(
     protected abstract fun loadProcessors(): List<Processor>
 }
 
+internal fun JCTree.prettyPrint(context: Context): String {
+    return StringWriter().apply { PrettyWithWorkarounds(context, this, false).printStat(this@prettyPrint) }.toString()
+}
+
+private class PrettyWithWorkarounds(private val context: Context, out: Writer?, sourceOutput: Boolean) : Pretty(out, sourceOutput) {
+    companion object {
+        private val ENUM = Flags.ENUM.toLong()
+    }
+
+    override fun visitVarDef(tree: JCTree.JCVariableDecl) {
+        if ((tree.mods.flags and ENUM) != 0L) {
+            // Pretty does not print annotations for enum values for some reason
+            printExpr(TreeMaker.instance(context).Modifiers(0, tree.mods.annotations))
+        }
+
+        super.visitVarDef(tree)
+    }
+}
+
 private inline fun <T> measureTimeMillis(block: () -> T) : Pair<Long, T> {
     val start = System.currentTimeMillis()
     val result = block()

plugins/kapt3/kapt3-compiler/test/org/jetbrains/kotlin/kapt3/test/AbstractKotlinKapt3IntegrationTest.kt

@@ -26,6 +26,7 @@ import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.kapt3.AbstractKapt3Extension
 import org.jetbrains.kotlin.kapt3.AptMode.STUBS_AND_APT
 import org.jetbrains.kotlin.kapt3.Kapt3BuilderFactory
+import org.jetbrains.kotlin.kapt3.prettyPrint
 import org.jetbrains.kotlin.kapt3.KaptContext
 import org.jetbrains.kotlin.kapt3.javac.KaptJavaFileObject
 import org.jetbrains.kotlin.kapt3.stubs.ClassFileToSourceStubConverter
@@ -170,17 +171,17 @@ abstract class AbstractKotlinKapt3IntegrationTest : CodegenTestCase() {
 
         override fun loadProcessors() = processors
 
-        override fun saveStubs(stubs: JavacList<JCTree.JCCompilationUnit>) {
+        override fun saveStubs(kaptContext: KaptContext<*>, stubs: JavacList<JCTree.JCCompilationUnit>) {
             if (this.savedStubs != null) {
                 error("Stubs are already saved")
             }
 
             this.savedStubs = stubs
-                    .map { it.toString() }
+                    .map { it.prettyPrint(kaptContext.context) }
                     .sorted()
                     .joinToString(AbstractKotlinKapt3Test.FILE_SEPARATOR)
 
-            super.saveStubs(stubs)
+            super.saveStubs(kaptContext, stubs)
         }
 
         override fun saveIncrementalData(

plugins/kapt3/kapt3-compiler/test/org/jetbrains/kotlin/kapt3/test/AbstractKotlinKapt3Test.kt

@@ -35,6 +35,7 @@ import org.jetbrains.kotlin.kapt3.doAnnotationProcessing
 import org.jetbrains.kotlin.kapt3.javac.KaptJavaFileObject
 import org.jetbrains.kotlin.kapt3.javac.KaptJavaLog
 import org.jetbrains.kotlin.kapt3.parseJavaFiles
+import org.jetbrains.kotlin.kapt3.prettyPrint
 import org.jetbrains.kotlin.kapt3.stubs.ClassFileToSourceStubConverter
 import org.jetbrains.kotlin.kapt3.util.KaptLogger
 import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisHandlerExtension
@@ -116,7 +117,7 @@ abstract class AbstractKotlinKapt3Test : CodegenTestCase() {
         val converter = ClassFileToSourceStubConverter(kaptContext, generateNonExistentClass, correctErrorTypes)
 
         val convertedTrees = converter.convert()
-        val convertedFiles = convertedTrees.map { tree -> createTempFile("stub", ".java", tree.toString()) }
+        val convertedFiles = convertedTrees.map { tree -> createTempFile("stub", ".java", tree.prettyPrint(kaptContext.context)) }
 
         val allJavaFiles = javaFiles + convertedFiles
 
@@ -194,7 +195,10 @@ open class AbstractClassFileToSourceStubConverterTest : AbstractKotlinKapt3Test(
         kaptContext.javaLog.interceptorData.files = convertedFiles.map { it.sourceFile to it }.toMap()
         if (validate) kaptContext.compiler.enterTrees(convertedFiles)
 
-        val actualRaw = convertedFiles.sortedBy { it.sourceFile.name }.joinToString(FILE_SEPARATOR)
+        val actualRaw = convertedFiles
+                .sortedBy { it.sourceFile.name }
+                .joinToString(FILE_SEPARATOR) { it.prettyPrint(kaptContext.context) }
+
         val actual = StringUtil.convertLineSeparators(actualRaw.trim({ it <= ' ' }))
                 .trimTrailingWhitespacesAndAddNewlineAtEOF()
                 .let { removeMetadataAnnotationContents(it) }

plugins/kapt3/kapt3-compiler/test/org/jetbrains/kotlin/kapt3/test/Jclass.java

@@ -0,0 +1,23 @@
+/*
+ * Copyright 2010-2017 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.kapt3.test;
+
+public enum Jclass {
+    BLACK, @Anno WHITE
+}
+
+@interface Anno {}

plugins/kapt3/kapt3-compiler/testData/converter/annotations.kt

@@ -29,3 +29,7 @@ class TestAnno2 {
     @get:Anno3("getter") @set:Anno3("setter") @property:Anno3("property") @field:Anno3("field") @setparam:Anno3("setparam")
     var b: String = "property initializer"
 }
+
+enum class Enum1 {
+    BLACK, @Anno1 WHITE
+}

plugins/kapt3/kapt3-compiler/testData/converter/annotations.txt

@@ -69,6 +69,21 @@ public enum Colors {
 ////////////////////
 
 
+@kotlin.Metadata()
+@kapt.internal.KaptMetadata()
+public enum Enum1 {
+    /*public static final*/ BLACK /* = new Enum1() */,
+    @Anno1()
+    /*public static final*/ WHITE /* = new Enum1() */;
+
+    @kapt.internal.KaptSignature(value = "<init>(Ljava/lang/String;I)V")
+    Enum1() {
+    }
+}
+
+////////////////////
+
+
 @kotlin.Metadata()
 @Anno3(value = "value")
 @Anno2(a = @Anno1(), clazz = TestAnno.class, classes = {TestAnno.class, Anno1.class})

plugins/kapt3/kapt3-compiler/testData/converter/annotations2.txt

@@ -68,7 +68,9 @@ package test;
 @Anno(value = "enum")
 @kapt.internal.KaptMetadata()
 public enum Enum {
+    @Anno(value = "white")
     /*public static final*/ WHITE /* = new Enum() */,
+    @Anno(value = "black")
     /*public static final*/ BLACK /* = new Enum() */;
     private final int x = 0;