Kapt: Provide SourceRetentionAnnotationHandler for incremental compilation.

Collect annotations with the "SOURCE" retention. (cherry picked from commit 6ef66e7)
2016-09-02 20:15:32 +03:00
parent c3ae66bc9c
commit 975364b2ed
17 changed files with 217 additions and 26 deletions
@@ -17,16 +17,20 @@
 package org.jetbrains.kotlin.incremental

 import org.jetbrains.kotlin.incremental.components.LookupTracker
+import org.jetbrains.kotlin.incremental.components.SourceRetentionAnnotationHandler
 import org.jetbrains.kotlin.load.kotlin.incremental.components.IncrementalCache
 import org.jetbrains.kotlin.load.kotlin.incremental.components.IncrementalCompilationComponents
 import org.jetbrains.kotlin.modules.TargetId

 class IncrementalCompilationComponentsImpl(
        private val caches: Map<TargetId, IncrementalCache>,
-        private val lookupTracker: LookupTracker
+        private val lookupTracker: LookupTracker,
+        private val sourceRetentionAnnotationHandler: SourceRetentionAnnotationHandler?
 ): IncrementalCompilationComponents {
    override fun getIncrementalCache(target: TargetId): IncrementalCache =
            caches[target] ?: throw Exception("Incremental cache for target ${target.name} not found")

    override fun getLookupTracker(): LookupTracker = lookupTracker
+
+    override fun getSourceRetentionAnnotationHandler() = sourceRetentionAnnotationHandler 
 }
@@ -0,0 +1,30 @@
+/*
+ * 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.incremental
+
+import org.jetbrains.kotlin.incremental.components.SourceRetentionAnnotationHandler
+
+class SourceRetentionAnnotationHandlerImpl : SourceRetentionAnnotationHandler {
+    private val mutableSourceRetentionAnnotations = mutableSetOf<String>()
+    
+    val sourceRetentionAnnotations: Set<String>
+        get() = mutableSourceRetentionAnnotations
+    
+    override fun register(internalName: String) {
+        mutableSourceRetentionAnnotations += internalName
+    }
+}
@@ -30,6 +30,7 @@ import org.jetbrains.kotlin.compilerRunner.OutputItemsCollectorImpl
 import org.jetbrains.kotlin.config.IncrementalCompilation
 import org.jetbrains.kotlin.config.Services
 import org.jetbrains.kotlin.incremental.components.LookupTracker
+import org.jetbrains.kotlin.incremental.components.SourceRetentionAnnotationHandler
 import org.jetbrains.kotlin.load.kotlin.incremental.components.IncrementalCache
 import org.jetbrains.kotlin.load.kotlin.incremental.components.IncrementalCompilationComponents
 import org.jetbrains.kotlin.modules.KotlinModuleXmlBuilder
@@ -70,9 +71,17 @@ fun makeCompileServices(
        incrementalCaches: Map<TargetId, IncrementalCache>,
        lookupTracker: LookupTracker,
        compilationCanceledStatus: CompilationCanceledStatus?
+) = makeCompileServices(incrementalCaches, lookupTracker, compilationCanceledStatus, null)
+
+fun makeCompileServices(
+        incrementalCaches: Map<TargetId, IncrementalCache>,
+        lookupTracker: LookupTracker,
+        compilationCanceledStatus: CompilationCanceledStatus?,
+        sourceRetentionAnnotationHandler: SourceRetentionAnnotationHandler?
 ): Services =
    with(Services.Builder()) {
-        register(IncrementalCompilationComponents::class.java, IncrementalCompilationComponentsImpl(incrementalCaches, lookupTracker))
+        register(IncrementalCompilationComponents::class.java, 
+                 IncrementalCompilationComponentsImpl(incrementalCaches, lookupTracker, sourceRetentionAnnotationHandler))
        compilationCanceledStatus?.let {
            register(CompilationCanceledStatus::class.java, it)
        }
@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.modules.TargetId
 import org.jetbrains.kotlin.daemon.common.CompilerCallbackServicesFacade
 import org.jetbrains.kotlin.daemon.common.DummyProfiler
 import org.jetbrains.kotlin.daemon.common.Profiler
+import org.jetbrains.kotlin.incremental.components.SourceRetentionAnnotationHandler


 class RemoteIncrementalCompilationComponentsClient(val facade: CompilerCallbackServicesFacade, eventManger: EventManger, val profiler: Profiler = DummyProfiler()) : IncrementalCompilationComponents {
@@ -31,4 +32,6 @@ class RemoteIncrementalCompilationComponentsClient(val facade: CompilerCallbackS
    override fun getIncrementalCache(target: TargetId): IncrementalCache = RemoteIncrementalCacheClient(facade, target, profiler)

    override fun getLookupTracker(): LookupTracker = remoteLookupTrackerClient
+
+    override fun getSourceRetentionAnnotationHandler() = null
 }
@@ -17,9 +17,11 @@
 package org.jetbrains.kotlin.load.kotlin.incremental.components

 import org.jetbrains.kotlin.incremental.components.LookupTracker
+import org.jetbrains.kotlin.incremental.components.SourceRetentionAnnotationHandler
 import org.jetbrains.kotlin.modules.TargetId

 interface IncrementalCompilationComponents {
    fun getIncrementalCache(target: TargetId): IncrementalCache
    fun getLookupTracker(): LookupTracker
+    fun getSourceRetentionAnnotationHandler(): SourceRetentionAnnotationHandler?
 }
@@ -0,0 +1,21 @@
+/*
+ * 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.incremental.components
+
+interface SourceRetentionAnnotationHandler {
+    fun register(internalName: String)
+}
@@ -57,6 +57,7 @@ import org.jetbrains.kotlin.config.Services
 import org.jetbrains.kotlin.daemon.common.isDaemonEnabled
 import org.jetbrains.kotlin.incremental.*
 import org.jetbrains.kotlin.incremental.components.LookupTracker
+import org.jetbrains.kotlin.incremental.components.SourceRetentionAnnotationHandler
 import org.jetbrains.kotlin.jps.JpsKotlinCompilerSettings
 import org.jetbrains.kotlin.jps.incremental.*
 import org.jetbrains.kotlin.load.kotlin.incremental.components.IncrementalCache
@@ -206,7 +207,8 @@ class KotlinBuilder : ModuleLevelBuilder(BuilderCategory.SOURCE_PROCESSOR) {
        val project = projectDescriptor.project
        val lookupTracker = getLookupTracker(project)
        val incrementalCaches = getIncrementalCaches(chunk, context)
-        val environment = createCompileEnvironment(incrementalCaches, lookupTracker, context)
+        val sourceRetentionAnnotationHandler = SourceRetentionAnnotationHandlerImpl()
+        val environment = createCompileEnvironment(incrementalCaches, lookupTracker, sourceRetentionAnnotationHandler, context)
        if (!environment.success()) {
            environment.reportErrorsTo(messageCollector)
            return ABORT
@@ -409,11 +411,13 @@ class KotlinBuilder : ModuleLevelBuilder(BuilderCategory.SOURCE_PROCESSOR) {
    private fun createCompileEnvironment(
            incrementalCaches: Map<ModuleBuildTarget, IncrementalCache>,
            lookupTracker: LookupTracker,
+            sourceRetentionAnnotationHandler: SourceRetentionAnnotationHandler?,
            context: CompileContext
    ): CompilerEnvironment {
        val compilerServices = Services.Builder()
                .register(IncrementalCompilationComponents::class.java,
-                          IncrementalCompilationComponentsImpl(incrementalCaches.mapKeys { TargetId(it.key) }, lookupTracker))
+                          IncrementalCompilationComponentsImpl(incrementalCaches.mapKeys { TargetId(it.key) },
+                                                               lookupTracker, sourceRetentionAnnotationHandler))
                .register(CompilationCanceledStatus::class.java, object : CompilationCanceledStatus {
                    override fun checkCanceled() {
                        if (context.cancelStatus.isCanceled) throw CompilationCanceledException()
@@ -16,5 +16,6 @@
    <orderEntry type="module" module-name="util" />
    <orderEntry type="library" name="kotlin-runtime" level="project" />
    <orderEntry type="module" module-name="cli" />
+    <orderEntry type="module" module-name="backend" />
  </component>
 </module>
@@ -25,9 +25,15 @@ 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.codegen.ClassBuilderMode
+import org.jetbrains.kotlin.codegen.state.IncompatibleClassTracker
+import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper
 import org.jetbrains.kotlin.config.APPEND_JAVA_SOURCE_ROOTS_HANDLER_KEY
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
+import org.jetbrains.kotlin.fileClasses.NoResolveFileClassesProvider
 import org.jetbrains.kotlin.java.model.elements.JeTypeElement
+import org.jetbrains.kotlin.load.java.JvmAbi
+import org.jetbrains.kotlin.load.kotlin.incremental.components.IncrementalCompilationComponents
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.BindingTrace
 import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisCompletedHandlerExtension
@@ -44,9 +50,11 @@ class ClasspathBasedAnnotationProcessingExtension(
        classesOutputDir: File,
        javaSourceRoots: List<File>,
        verboseOutput: Boolean,
-        incrementalDataFile: File?
+        incrementalDataFile: File?,
+        incrementalCompilationComponents: IncrementalCompilationComponents?
 ) : AbstractAnnotationProcessingExtension(generatedSourcesOutputDir, 
-                                          classesOutputDir, javaSourceRoots, verboseOutput, incrementalDataFile) {
+                                          classesOutputDir, javaSourceRoots, verboseOutput,
+                                          incrementalDataFile, incrementalCompilationComponents) {
    override fun loadAnnotationProcessors(): List<Processor> {
        val classLoader = URLClassLoader(annotationProcessingClasspath.map { it.toURI().toURL() }.toTypedArray())
        return ServiceLoader.load(Processor::class.java, classLoader).toList()
@@ -58,7 +66,8 @@ abstract class AbstractAnnotationProcessingExtension(
        val classesOutputDir: File,
        val javaSourceRoots: List<File>,
        val verboseOutput: Boolean,
-        val incrementalDataFile: File? = null
+        val incrementalDataFile: File? = null,
+        val incrementalCompilationComponents: IncrementalCompilationComponents? = null
 ) : AnalysisCompletedHandlerExtension {
    private companion object {
        val LINE_SEPARATOR = System.getProperty("line.separator") ?: "\n"
@@ -116,7 +125,7 @@ abstract class AbstractAnnotationProcessingExtension(

        val processingEnvironment = KotlinProcessingEnvironment(
                elements, types, messager, options, filer, processors,
-                project, psiManager, javaPsiFacade, projectScope, appendJavaSourceRootsHandler)
+                project, psiManager, javaPsiFacade, projectScope, bindingTrace.bindingContext, appendJavaSourceRootsHandler)

        val processingResult = processingEnvironment.doAnnotationProcessing(files)

@@ -148,6 +157,11 @@ abstract class AbstractAnnotationProcessingExtension(
                listOf(generatedSourcesOutputDir),
                addToEnvironment = false)
    }
+    
+    private fun KotlinProcessingEnvironment.createTypeMapper(): KotlinTypeMapper {
+        return KotlinTypeMapper(bindingContext, ClassBuilderMode.full(false), NoResolveFileClassesProvider,
+                         null, IncompatibleClassTracker.DoNothing, JvmAbi.DEFAULT_MODULE_NAME)
+    }

    private fun KotlinProcessingEnvironment.doAnnotationProcessing(files: Collection<KtFile>): ProcessingResult {
        val allSupportedAnnotationFqNames = run initializeProcessors@ {
@@ -156,7 +170,11 @@ abstract class AbstractAnnotationProcessingExtension(
            processors.flatMapTo(mutableSetOf()) { it.supportedAnnotationTypes }
        }

-        val firstRoundAnnotations = RoundAnnotations(allSupportedAnnotationFqNames)
+        val firstRoundAnnotations = RoundAnnotations(
+                allSupportedAnnotationFqNames, 
+                incrementalCompilationComponents?.getSourceRetentionAnnotationHandler(),
+                bindingContext,
+                createTypeMapper())
        
        run analyzeFilesForFirstRound@ {
            log { "Analysing Kotlin files: " + files.map { it.virtualFile.path } }
@@ -223,7 +241,7 @@ abstract class AbstractAnnotationProcessingExtension(
        } + 1
        
        log { "Starting round $finalRoundNumber (final)" }
-        val finalRoundEnvironment = KotlinRoundEnvironment(RoundAnnotations(allSupportedAnnotationFqNames), true, finalRoundNumber)
+        val finalRoundEnvironment = KotlinRoundEnvironment(firstRoundAnnotations.copy(), true, finalRoundNumber)
        for (processor in processors) {
            processor.process(emptySet(), finalRoundEnvironment)
        }
@@ -258,7 +276,7 @@ abstract class AbstractAnnotationProcessingExtension(
        }
        
        // Start the next round
-        val nextRoundAnnotations = RoundAnnotations(roundEnvironment.supportedAnnotationFqNames).apply { analyzeFiles(psiFiles) }
+        val nextRoundAnnotations = roundEnvironment.roundAnnotations.copy().apply { analyzeFiles(psiFiles) }
        val nextRoundEnvironment = KotlinRoundEnvironment(nextRoundAnnotations, false, roundEnvironment.roundNumber + 1)
        return process(nextRoundEnvironment)
    }
@@ -274,8 +292,8 @@ abstract class AbstractAnnotationProcessingExtension(
            val acceptsAnyAnnotation = supportedAnnotationNames.contains("*")

            val applicableAnnotationNames = when (acceptsAnyAnnotation) {
-                true -> roundEnvironment.annotationsMap.keys
-                false -> processor.supportedAnnotationTypes.filter { it in roundEnvironment.annotationsMap }
+                true -> roundEnvironment.roundAnnotations.annotationsMap.keys
+                false -> processor.supportedAnnotationTypes.filter { it in roundEnvironment.roundAnnotations.annotationsMap }
            }

            if (applicableAnnotationNames.isEmpty()) {
@@ -122,8 +122,11 @@ class AnnotationProcessingComponentRegistrar : ComponentRegistrar {
        // Annotations with the "SOURCE" retention will be written to class files
        project.putUserData(IS_KAPT2_ENABLED_KEY, true)
        
+        val incrementalCompilationComponents = configuration[JVMConfigurationKeys.INCREMENTAL_COMPILATION_COMPONENTS]
+        
        val annotationProcessingExtension = ClasspathBasedAnnotationProcessingExtension(
-                classpath, generatedOutputDirFile, classesOutputDir, javaRoots, verboseOutput, incrementalDataFile)
+                classpath, generatedOutputDirFile, classesOutputDir, javaRoots, verboseOutput, 
+                incrementalDataFile, incrementalCompilationComponents)
        
        AnalysisCompletedHandlerExtension.registerExtension(project, annotationProcessingExtension)
    }
@@ -18,13 +18,22 @@ package org.jetbrains.kotlin.annotation.processing

 import com.intellij.psi.*
 import com.intellij.psi.util.PsiTreeUtil
+import org.jetbrains.kotlin.asJava.elements.KtLightAnnotation
 import org.jetbrains.kotlin.asJava.findFacadeClass
 import org.jetbrains.kotlin.asJava.toLightClass
+import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper
+import org.jetbrains.kotlin.incremental.components.SourceRetentionAnnotationHandler
 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

-internal class RoundAnnotations(val supportedAnnotationFqNames: Set<String>) {
+internal class RoundAnnotations(
+        val supportedAnnotationFqNames: Set<String>,
+        val sourceRetentionAnnotationHandler: SourceRetentionAnnotationHandler?,
+        val bindingContext: BindingContext,
+        val typeMapper: KotlinTypeMapper
+) {
    private val mutableAnnotationsMap = mutableMapOf<String, MutableList<PsiModifierListOwner>>()
    private val mutableAnalyzedClasses = mutableSetOf<String>()
    
@@ -35,6 +44,8 @@ internal class RoundAnnotations(val supportedAnnotationFqNames: Set<String>) {
    
    val analyzedClasses: Set<String>
        get() = mutableAnalyzedClasses
+    
+    fun copy() = RoundAnnotations(supportedAnnotationFqNames, sourceRetentionAnnotationHandler, bindingContext, typeMapper)

    fun analyzeFiles(files: Collection<KtFile>) = files.forEach { analyzeFile(it) }
    
@@ -63,6 +74,18 @@ internal class RoundAnnotations(val supportedAnnotationFqNames: Set<String>) {
        is PsiClass -> containingClass?.let { it.getTopLevelClassParent() } ?: this
        else -> PsiTreeUtil.getParentOfType(this, PsiClass::class.java, true)?.getTopLevelClassParent()
    }
+    
+    private val PsiAnnotation.hasSourceRetention: Boolean
+        get() {
+            val annotationDeclaration = nameReferenceElement?.resolve() as? PsiClass ?: return false
+            val metaAnnotations = annotationDeclaration.modifierList?.annotations ?: return false
+            return metaAnnotations.any { anno ->
+                val declaration = anno.nameReferenceElement?.resolve() as? PsiClass ?: return@any false
+                if (declaration.qualifiedName != "java.lang.annotation.Retention") return@any false
+                val value = (anno.findAttributeValue("value") as? PsiReferenceExpression)?.resolve() ?: return@any false
+                value is PsiEnumConstant && value.name == "SOURCE"
+            }
+        }

    fun analyzeDeclaration(declaration: PsiElement): Boolean {
        if (declaration !is PsiModifierListOwner) return false
@@ -72,6 +95,16 @@ internal class RoundAnnotations(val supportedAnnotationFqNames: Set<String>) {

        for (annotation in declaration.getAnnotationsWithInherited()) {
            val fqName = annotation.qualifiedName ?: continue
+            
+            val ktLightAnnotation = annotation as? KtLightAnnotation
+            if (ktLightAnnotation != null && sourceRetentionAnnotationHandler != null && ktLightAnnotation.hasSourceRetention) { 
+                val type = bindingContext[BindingContext.ANNOTATION, ktLightAnnotation.kotlinOrigin]?.type
+                if (type != null) {
+                    val internalName = typeMapper.mapType(type).internalName
+                    sourceRetentionAnnotationHandler.register(internalName)
+                }
+            }
+            
            if (!acceptsAnyAnnotation && fqName !in supportedAnnotationFqNames) continue
            mutableAnnotationsMap.getOrPut(fqName, { mutableListOf() }).add(declaration)

@@ -20,6 +20,7 @@ import com.intellij.openapi.project.Project
 import com.intellij.psi.JavaPsiFacade
 import com.intellij.psi.PsiManager
 import com.intellij.psi.search.GlobalSearchScope
+import org.jetbrains.kotlin.resolve.BindingContext
 import java.io.File
 import java.util.*
 import javax.annotation.processing.ProcessingEnvironment
@@ -41,6 +42,7 @@ class KotlinProcessingEnvironment(
        internal val psiManager: PsiManager,
        internal val javaPsiFacade: JavaPsiFacade,
        internal val projectScope: GlobalSearchScope,
+        internal val bindingContext: BindingContext,
        internal val appendJavaSourceRootsHandler: (List<File>) -> Unit
 ) : ProcessingEnvironment {
    private val options = Collections.unmodifiableMap(options) 
@@ -16,7 +16,6 @@

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

-import com.intellij.psi.PsiModifierListOwner
 import org.jetbrains.kotlin.annotation.processing.RoundAnnotations
 import org.jetbrains.kotlin.java.model.toJeElement
 import javax.annotation.processing.RoundEnvironment
@@ -24,18 +23,12 @@ import javax.lang.model.element.Element
 import javax.lang.model.element.TypeElement

 internal class KotlinRoundEnvironment(
-        private val roundAnnotations: RoundAnnotations,
+        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
-    
-    internal val supportedAnnotationFqNames: Set<String>
-        get() = roundAnnotations.supportedAnnotationFqNames
-    
    override fun getRootElements() = emptySet<Element>()
    
    override fun processingOver() = isProcessingOver
@@ -0,0 +1,42 @@
+@Retention(AnnotationRetention.SOURCE)
+annotation class Source1
+
+@Retention(AnnotationRetention.SOURCE)
+annotation class Source2
+
+@Retention(AnnotationRetention.SOURCE)
+annotation class Source3
+
+@Retention(AnnotationRetention.SOURCE)
+annotation class Source4
+
+@Retention(AnnotationRetention.BINARY)
+annotation class Binary
+
+@Retention(AnnotationRetention.RUNTIME)
+annotation class Runtime
+
+@Source1
+class Test
+
+class Test2 {
+    @Source2
+    fun t() {}
+}
+
+class Test3 {
+    @field:Source3
+    val p: String = "A"
+}
+
+class Test4 {
+    fun t(@Source4 a: String) {}
+}
+
+class Test5 {
+    @Retention(AnnotationRetention.SOURCE)
+    annotation class Source5
+
+    @Source5
+    fun t() {}
+}
@@ -34,5 +34,6 @@
    <orderEntry type="module" module-name="light-classes" scope="TEST" />
    <orderEntry type="library" scope="TEST" name="junit-plugin" level="project" />
    <orderEntry type="module" module-name="annotation-processing" scope="TEST" />
+    <orderEntry type="module" module-name="build-common" scope="TEST" />
  </component>
 </module>
@@ -21,9 +21,12 @@ import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import org.jetbrains.kotlin.codegen.AbstractBytecodeTextTest
 import org.jetbrains.kotlin.codegen.CodegenTestUtil
+import org.jetbrains.kotlin.incremental.SourceRetentionAnnotationHandlerImpl
 import org.jetbrains.kotlin.java.model.elements.JeAnnotationMirror
 import org.jetbrains.kotlin.java.model.elements.JeMethodExecutableElement
 import org.jetbrains.kotlin.java.model.elements.JeTypeElement
+import org.jetbrains.kotlin.load.kotlin.incremental.components.IncrementalCompilationComponents
+import org.jetbrains.kotlin.modules.TargetId
 import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisCompletedHandlerExtension
 import org.jetbrains.kotlin.test.ConfigurationKind
 import org.jetbrains.kotlin.test.KotlinTestUtils
@@ -39,7 +42,8 @@ import javax.lang.model.element.*

 class AnnotationProcessingExtensionForTests(
        val processors: List<Processor>
-) : AbstractAnnotationProcessingExtension(createTempDir(), createTempDir(), listOf(), true, createIncrementalDataFile()) {
+) : AbstractAnnotationProcessingExtension(createTempDir(), createTempDir(), listOf(), true,
+                                          createIncrementalDataFile(), StubIncrementalCompilationComponents()) {
    override fun loadAnnotationProcessors() = processors
    
    private companion object {
@@ -51,6 +55,14 @@ class AnnotationProcessingExtensionForTests(
            deleteOnExit()
        }
    }
+    
+    private class StubIncrementalCompilationComponents : IncrementalCompilationComponents {
+        private val sourceRetentionAnnotationHandler = SourceRetentionAnnotationHandlerImpl()
+        
+        override fun getIncrementalCache(target: TargetId) = null!!
+        override fun getLookupTracker() = null!!
+        override fun getSourceRetentionAnnotationHandler() = sourceRetentionAnnotationHandler
+    }
 }

 abstract class AbstractProcessorTest : AbstractBytecodeTextTest() {
@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.annotation.processing.test.processor

 import org.intellij.lang.annotations.Language
+import org.jetbrains.kotlin.incremental.SourceRetentionAnnotationHandlerImpl
 import org.jetbrains.kotlin.java.model.elements.*
 import org.jetbrains.kotlin.java.model.types.JeDeclaredType
 import org.jetbrains.kotlin.java.model.types.JeMethodExecutableTypeMirror
@@ -191,13 +192,25 @@ class ProcessorTests : AbstractProcessorTest() {
            "IncrementalDataSimple",
            "i Intf, i Test, i Test2, i Test3, i Test8")
    
+    private fun getKapt2Extension() = AnalysisCompletedHandlerExtension.getInstances(myEnvironment.project)
+            .firstIsInstance<AnnotationProcessingExtensionForTests>()
+    
    private fun incrementalDataTest(fileName: String, @Language("TEXT") expectedText: String) {
        test(fileName, "Anno", "Anno2", "Anno3") { set, roundEnv, env -> }
-        val ext = AnalysisCompletedHandlerExtension.getInstances(myEnvironment.project)
-                .firstIsInstance<AnnotationProcessingExtensionForTests>()
+        val ext = getKapt2Extension()
        val incrementalDataFile = ext.incrementalDataFile
        assertNotNull(incrementalDataFile)
        val text = incrementalDataFile!!.readText().lines().sorted().joinToString(", ")
        assertEquals(expectedText, text)
    }
+    
+    fun testSourceRetention() {
+        test("SourceRetention", "*") { set, roundEnv, env -> }
+        val ext = getKapt2Extension()
+        val incrementalCompilationComponents = ext.incrementalCompilationComponents
+        assertNotNull(incrementalCompilationComponents)
+        val annotationHandler = incrementalCompilationComponents!!.getSourceRetentionAnnotationHandler()
+        val annotations = (annotationHandler as SourceRetentionAnnotationHandlerImpl).sourceRetentionAnnotations.sorted()
+        assertEquals("Source1, Source2, Source3, Source4, Test5\$Source5", annotations.joinToString())
+    }
 }