From 4e82fabc7d0d6815a8b283a941b2d5fa26f4e200 Mon Sep 17 00:00:00 2001
From: Yan Zhulanow <yan.zhulanow@jetbrains.com>
Date: Tue, 29 Dec 2015 19:26:34 +0300
Subject: [PATCH] kapt: Pass Kotlin annotations to process()

---
 .../annotation/AnnotationProcessorWrapper.kt   | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/libraries/tools/kotlin-annotation-processing/src/main/kotlin/org/jetbrains/kotlin/annotation/AnnotationProcessorWrapper.kt b/libraries/tools/kotlin-annotation-processing/src/main/kotlin/org/jetbrains/kotlin/annotation/AnnotationProcessorWrapper.kt
index f464ef539ae..fe098136446 100644
--- a/libraries/tools/kotlin-annotation-processing/src/main/kotlin/org/jetbrains/kotlin/annotation/AnnotationProcessorWrapper.kt
+++ b/libraries/tools/kotlin-annotation-processing/src/main/kotlin/org/jetbrains/kotlin/annotation/AnnotationProcessorWrapper.kt
@@ -17,7 +17,6 @@
 package org.jetbrains.kotlin.annotation
 
 import java.io.File
-import java.io.IOException
 import javax.annotation.processing.*
 import javax.lang.model.SourceVersion
 import javax.lang.model.element.AnnotationMirror
@@ -25,7 +24,6 @@ import javax.lang.model.element.Element
 import javax.lang.model.element.ExecutableElement
 import javax.lang.model.element.TypeElement
 import javax.tools.Diagnostic
-import javax.tools.StandardLocation
 import kotlin.properties.Delegates
 
 public class AnnotationProcessorStub : AbstractProcessor() {
@@ -117,12 +115,24 @@ public abstract class AnnotationProcessorWrapper(
         return processor.supportedSourceVersion
     }
 
-    override fun process(annotations: MutableSet<out TypeElement>?, roundEnv: RoundEnvironment): Boolean {
+    override fun process(annotations: Set<TypeElement>?, roundEnv: RoundEnvironment): Boolean {
         roundCounter += 1
 
         val roundEnvironmentWrapper = RoundEnvironmentWrapper(
                 processingEnv, roundEnv, roundCounter, kotlinAnnotationsProvider)
-        processor.process(annotations, roundEnvironmentWrapper)
+
+        val wrappedAnnotations = annotations?.toMutableSet() ?: hashSetOf<TypeElement>()
+        val existingFqNames = wrappedAnnotations.mapTo(hashSetOf<String>()) { it.qualifiedName.toString() }
+
+        if (roundCounter == 1) {
+            for (annotationFqName in kotlinAnnotationsProvider.annotatedKotlinElements.keys) {
+                if (annotationFqName in existingFqNames) continue
+                existingFqNames.add(annotationFqName)
+                processingEnv.elementUtils.getTypeElement(annotationFqName)?.let { wrappedAnnotations += it }
+            }
+        }
+
+        processor.process(wrappedAnnotations, roundEnvironmentWrapper)
         return false
     }