Kapt3: Support apt modes

Allow to run kapt in "annotation processing only" and "stub generation only" modes in order to support incremental compilation in Gradle.
Unfortunately, annotation processing can't be done incrementally cause it uses all module files, but the stub generation can be done incrementally.

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

@@ -29,7 +29,10 @@ import org.jetbrains.kotlin.codegen.KotlinCodegenFacade
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.config.CommonConfigurationKeys
 import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.container.ComponentProvider
+import org.jetbrains.kotlin.context.ProjectContext
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
+import org.jetbrains.kotlin.kapt3.AptMode.*
 import org.jetbrains.kotlin.kapt3.diagnostic.ErrorsKapt3
 import org.jetbrains.kotlin.kapt3.diagnostic.KaptError
 import org.jetbrains.kotlin.kapt3.stubs.ClassFileToSourceStubConverter
@@ -55,14 +58,14 @@ class ClasspathBasedKapt3Extension(
         options: Map<String, String>,
         javacOptions: Map<String, String>,
         annotationProcessors: String,
-        aptOnly: Boolean,
+        aptMode: AptMode,
         val useLightAnalysis: Boolean,
         correctErrorTypes: Boolean,
         pluginInitializedTime: Long,
         logger: KaptLogger
 ) : AbstractKapt3Extension(compileClasspath, annotationProcessingClasspath, javaSourceRoots, sourcesOutputDir,
                            classFilesOutputDir, stubsOutputDir, incrementalDataOutputDir, options, javacOptions, annotationProcessors,
-                           aptOnly, pluginInitializedTime, logger, correctErrorTypes) {
+                           aptMode, pluginInitializedTime, logger, correctErrorTypes) {
     override val analyzePartially: Boolean
         get() = useLightAnalysis
 
@@ -108,7 +111,7 @@ abstract class AbstractKapt3Extension(
         val options: Map<String, String>,
         val javacOptions: Map<String, String>,
         val annotationProcessors: String,
-        val aptOnly: Boolean,
+        val aptMode: AptMode,
         val pluginInitializedTime: Long,
         val logger: KaptLogger,
         val correctErrorTypes: Boolean
@@ -125,6 +128,21 @@ abstract class AbstractKapt3Extension(
         return false
     }
 
+    override fun doAnalysis(
+            project: Project,
+            module: ModuleDescriptor,
+            projectContext: ProjectContext,
+            files: Collection<KtFile>,
+            bindingTrace: BindingTrace,
+            componentProvider: ComponentProvider
+    ): AnalysisResult? {
+        if (aptMode == APT_ONLY) {
+            return AnalysisResult.EMPTY
+        }
+
+        return super.doAnalysis(project, module, projectContext, files, bindingTrace, componentProvider)
+    }
+
     override fun analysisCompleted(
             project: Project,
             module: ModuleDescriptor,
@@ -136,24 +154,14 @@ abstract class AbstractKapt3Extension(
         fun doNotGenerateCode() = AnalysisResult.Companion.success(BindingContext.EMPTY, module, shouldGenerateCode = false)
 
         logger.info { "Initial analysis took ${System.currentTimeMillis() - pluginInitializedTime} ms" }
-        logger.info { "Kotlin files to compile: " + files.map { it.virtualFile?.name ?: "<in memory ${it.hashCode()}>" } }
 
         val processors = loadProcessors()
-        if (processors.isEmpty()) return if (aptOnly) doNotGenerateCode() else null
+        if (processors.isEmpty()) return if (aptMode != WITH_COMPILATION) doNotGenerateCode() else null
 
-        val (kaptContext, generationState) = compileStubs(project, module, bindingTrace.bindingContext, files.toList())
+        val kaptContext = generateStubs(project, module, bindingTrace.bindingContext, files)
 
         try {
-            generateKotlinSourceStubs(kaptContext, generationState)
+            runAnnotationProcessing(kaptContext, processors)
-            val javaSourceFiles = collectJavaSourceFiles()
-
-            val (annotationProcessingTime) = measureTimeMillis {
-                kaptContext.doAnnotationProcessing(
-                        javaSourceFiles, processors, compileClasspath, annotationProcessingClasspath,
-                        annotationProcessors, sourcesOutputDir, classFilesOutputDir)
-            }
-
-            logger.info { "Annotation processing took $annotationProcessingTime ms" }
         } catch (thr: Throwable) {
             if (thr !is KaptError || thr.kind != KaptError.Kind.ERROR_RAISED) {
                 logger.exception(thr)
@@ -167,11 +175,10 @@ abstract class AbstractKapt3Extension(
             bindingTrace.report(ErrorsKapt3.KAPT3_PROCESSING_ERROR.on(files.first()))
             return null // Compilation will be aborted anyway because of the error above
         } finally {
-            generationState.destroy()
             kaptContext.close()
         }
 
-        return if (aptOnly) {
+        return if (aptMode != WITH_COMPILATION) {
             doNotGenerateCode()
         } else {
             AnalysisResult.RetryWithAdditionalJavaRoots(
@@ -182,12 +189,38 @@ abstract class AbstractKapt3Extension(
         }
     }
 
+    private fun generateStubs(project: Project, module: ModuleDescriptor, context: BindingContext, files: Collection<KtFile>): KaptContext<*> {
+        if (!aptMode.generateStubs) {
+            return KaptContext(logger, BindingContext.EMPTY, emptyList(), emptyMap(), null, options, javacOptions)
+        }
+
+        logger.info { "Kotlin files to compile: " + files.map { it.virtualFile?.name ?: "<in memory ${it.hashCode()}>" } }
+
+        return compileStubs(project, module, context, files.toList()).apply {
+            generateKotlinSourceStubs(this)
+        }
+    }
+
+    private fun runAnnotationProcessing(kaptContext: KaptContext<*>, processors: List<Processor>) {
+        if (!aptMode.runAnnotationProcessing) return
+
+        val javaSourceFiles = collectJavaSourceFiles()
+
+        val (annotationProcessingTime) = measureTimeMillis {
+            kaptContext.doAnnotationProcessing(
+                    javaSourceFiles, processors, compileClasspath, annotationProcessingClasspath,
+                    annotationProcessors, sourcesOutputDir, classFilesOutputDir)
+        }
+
+        logger.info { "Annotation processing took $annotationProcessingTime ms" }
+    }
+
     private fun compileStubs(
             project: Project,
             module: ModuleDescriptor,
             bindingContext: BindingContext,
             files: List<KtFile>
-    ): Pair<KaptContext, GenerationState> {
+    ): KaptContext<GenerationState> {
         val builderFactory = Kapt3BuilderFactory()
 
         val generationState = GenerationState(
@@ -208,12 +241,11 @@ abstract class AbstractKapt3Extension(
         logger.info { "Stubs compilation took $classFilesCompilationTime ms" }
         logger.info { "Compiled classes: " + compiledClasses.joinToString { it.name } }
 
-        return Pair(KaptContext(logger, bindingContext, compiledClasses, origins, options, javacOptions), generationState)
+        return KaptContext(logger, bindingContext, compiledClasses, origins, generationState, options, javacOptions)
     }
 
-    private fun generateKotlinSourceStubs(kaptContext: KaptContext, generationState: GenerationState) {
+    private fun generateKotlinSourceStubs(kaptContext: KaptContext<GenerationState>) {
-        val converter = ClassFileToSourceStubConverter(kaptContext, generationState.typeMapper,
+        val converter = ClassFileToSourceStubConverter(kaptContext, generateNonExistentClass = true, correctErrorTypes = correctErrorTypes)
-                                                       generateNonExistentClass = true, correctErrorTypes = correctErrorTypes)
 
         val (stubGenerationTime, kotlinSourceStubs) = measureTimeMillis {
             converter.convert()
@@ -223,7 +255,7 @@ abstract class AbstractKapt3Extension(
         logger.info { "Stubs for Kotlin classes: " + kotlinSourceStubs.joinToString { it.sourcefile.name } }
 
         saveStubs(kotlinSourceStubs)
-        saveIncrementalData(generationState, logger.messageCollector, converter)
+        saveIncrementalData(kaptContext, logger.messageCollector, converter)
     }
 
     private fun collectJavaSourceFiles(): List<File> {
@@ -247,13 +279,13 @@ abstract class AbstractKapt3Extension(
     }
 
     protected open fun saveIncrementalData(
-            generationState: GenerationState,
+            kaptContext: KaptContext<GenerationState>,
             messageCollector: MessageCollector,
             converter: ClassFileToSourceStubConverter) {
         val incrementalDataOutputDir = this.incrementalDataOutputDir ?: return
 
-        val reportOutputFiles = generationState.configuration.getBoolean(CommonConfigurationKeys.REPORT_OUTPUT_FILES)
+        val reportOutputFiles = kaptContext.generationState.configuration.getBoolean(CommonConfigurationKeys.REPORT_OUTPUT_FILES)
-        generationState.factory.writeAll(
+        kaptContext.generationState.factory.writeAll(
                 incrementalDataOutputDir,
                 if (!reportOutputFiles) null else fun(file: OutputFile, sources: List<File>, output: File) {
                     val stubFileObject = converter.bindings[file.relativePath.substringBeforeLast(".class", missingDelimiterValue = "")]

plugins/kapt3/src/org/jetbrains/kotlin/kapt3/Kapt3Plugin.kt

@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.cli.common.messages.MessageRenderer
 import org.jetbrains.kotlin.cli.common.messages.PrintingMessageCollector
 import org.jetbrains.kotlin.kapt3.Kapt3ConfigurationKeys.ANNOTATION_PROCESSOR_CLASSPATH
 import org.jetbrains.kotlin.kapt3.Kapt3ConfigurationKeys.APT_OPTIONS
+import org.jetbrains.kotlin.kapt3.Kapt3ConfigurationKeys.JAVAC_CLI_OPTIONS
 import org.jetbrains.kotlin.cli.jvm.config.JavaSourceRoot
 import org.jetbrains.kotlin.cli.jvm.config.JvmContentRoot
 import org.jetbrains.kotlin.compiler.plugin.CliOption
@@ -38,7 +39,6 @@ import org.jetbrains.kotlin.container.ComponentProvider
 import org.jetbrains.kotlin.context.ProjectContext
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.diagnostics.rendering.DefaultErrorMessages
-import org.jetbrains.kotlin.kapt3.Kapt3ConfigurationKeys.JAVAC_CLI_OPTIONS
 import org.jetbrains.kotlin.kapt3.diagnostic.DefaultErrorMessagesKapt3
 import org.jetbrains.kotlin.kapt3.util.KaptLogger
 import org.jetbrains.kotlin.psi.KtFile
@@ -77,9 +77,13 @@ object Kapt3ConfigurationKeys {
     val VERBOSE_MODE: CompilerConfigurationKey<String> =
             CompilerConfigurationKey.create<String>("verbose mode")
 
+    @Deprecated("Use APT_MODE instead.")
     val APT_ONLY: CompilerConfigurationKey<String> =
             CompilerConfigurationKey.create<String>("do only annotation processing")
 
+    val APT_MODE: CompilerConfigurationKey<String> =
+            CompilerConfigurationKey.create<String>("annotation processing mode")
+
     val USE_LIGHT_ANALYSIS: CompilerConfigurationKey<String> =
             CompilerConfigurationKey.create<String>("do not analyze declaration bodies if can")
 
@@ -122,9 +126,15 @@ class Kapt3CommandLineProcessor : CommandLineProcessor {
         val VERBOSE_MODE_OPTION: CliOption =
                 CliOption("verbose", "true | false", "Enable verbose output", required = false)
 
+        @Deprecated("Use APT_MODE_OPTION instead.")
         val APT_ONLY_OPTION: CliOption =
                 CliOption("aptOnly", "true | false", "Run only annotation processing, do not compile Kotlin files", required = false)
 
+        val APT_MODE_OPTION: CliOption =
+                CliOption("aptMode", "apt | stubs | stubsAndApt | compile",
+                          "Annotation processing mode: only apt, only stub generation, both, or with the subsequent compilation",
+                          required = false)
+
         val USE_LIGHT_ANALYSIS_OPTION: CliOption =
                 CliOption("useLightAnalysis", "true | false", "Do not analyze declaration bodies if can", required = false)
 
@@ -136,8 +146,8 @@ class Kapt3CommandLineProcessor : CommandLineProcessor {
 
     override val pluginOptions: Collection<CliOption> =
             listOf(SOURCE_OUTPUT_DIR_OPTION, ANNOTATION_PROCESSOR_CLASSPATH_OPTION, APT_OPTIONS_OPTION, JAVAC_CLI_OPTIONS_OPTION,
-                   CLASS_OUTPUT_DIR_OPTION, VERBOSE_MODE_OPTION, STUBS_OUTPUT_DIR_OPTION, APT_ONLY_OPTION,
+                   CLASS_OUTPUT_DIR_OPTION, VERBOSE_MODE_OPTION, STUBS_OUTPUT_DIR_OPTION, APT_ONLY_OPTION, APT_MODE_OPTION,
-                   USE_LIGHT_ANALYSIS_OPTION, CORRECT_ERROR_TYPES_OPTION, ANNOTATION_PROCESSORS_OPTION)
+                   USE_LIGHT_ANALYSIS_OPTION, CORRECT_ERROR_TYPES_OPTION, ANNOTATION_PROCESSORS_OPTION, INCREMENTAL_DATA_OUTPUT_DIR_OPTION)
 
     override fun processOption(option: CliOption, value: String, configuration: CompilerConfiguration) {
         when (option) {
@@ -151,6 +161,7 @@ class Kapt3CommandLineProcessor : CommandLineProcessor {
             INCREMENTAL_DATA_OUTPUT_DIR_OPTION -> configuration.put(Kapt3ConfigurationKeys.INCREMENTAL_DATA_OUTPUT_DIR, value)
             VERBOSE_MODE_OPTION -> configuration.put(Kapt3ConfigurationKeys.VERBOSE_MODE, value)
             APT_ONLY_OPTION -> configuration.put(Kapt3ConfigurationKeys.APT_ONLY, value)
+            APT_MODE_OPTION -> configuration.put(Kapt3ConfigurationKeys.APT_MODE, value)
             USE_LIGHT_ANALYSIS_OPTION -> configuration.put(Kapt3ConfigurationKeys.USE_LIGHT_ANALYSIS, value)
             CORRECT_ERROR_TYPES_OPTION -> configuration.put(Kapt3ConfigurationKeys.CORRECT_ERROR_TYPES, value)
             else -> throw CliOptionProcessingException("Unknown option: ${option.name}")
@@ -182,7 +193,9 @@ class Kapt3ComponentRegistrar : ComponentRegistrar {
     }
 
     override fun registerProjectComponents(project: MockProject, configuration: CompilerConfiguration) {
-        val isAptOnly = configuration.get(Kapt3ConfigurationKeys.APT_ONLY) == "true"
+        val aptMode = AptMode.parse(configuration.get(Kapt3ConfigurationKeys.APT_MODE) ?:
+                                    configuration.get(Kapt3ConfigurationKeys.APT_ONLY))
+
         val isVerbose = configuration.get(Kapt3ConfigurationKeys.VERBOSE_MODE) == "true"
         val messageCollector = configuration.get(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY)
                                ?: PrintingMessageCollector(System.err, MessageRenderer.PLAIN_FULL_PATHS, isVerbose)
@@ -205,7 +218,7 @@ class Kapt3ComponentRegistrar : ComponentRegistrar {
         val apClasspath = configuration.get(ANNOTATION_PROCESSOR_CLASSPATH)?.map(::File)
 
         if (sourcesOutputDir == null || classFilesOutputDir == null || apClasspath == null || stubsOutputDir == null) {
-            if (isAptOnly) {
+            if (aptMode != AptMode.WITH_COMPILATION) {
                 val nonExistentOptionName = when {
                     sourcesOutputDir == null -> "Sources output directory"
                     classFilesOutputDir == null -> "Classes output directory"
@@ -237,7 +250,7 @@ class Kapt3ComponentRegistrar : ComponentRegistrar {
 
         if (isVerbose) {
             logger.info("Kapt3 is enabled.")
-            logger.info("Do annotation processing only: $isAptOnly")
+            logger.info("Annotation processing mode: $aptMode")
             logger.info("Use light analysis: $useLightAnalysis")
             logger.info("Correct error types: $correctErrorTypes")
             logger.info("Source output directory: $sourcesOutputDir")
@@ -254,7 +267,7 @@ class Kapt3ComponentRegistrar : ComponentRegistrar {
         val kapt3AnalysisCompletedHandlerExtension = ClasspathBasedKapt3Extension(
                 compileClasspath, apClasspath, javaSourceRoots, sourcesOutputDir, classFilesOutputDir,
                 stubsOutputDir, incrementalDataOutputDir, apOptions, javacCliOptions, annotationProcessors,
-                isAptOnly, useLightAnalysis, correctErrorTypes, System.currentTimeMillis(), logger)
+                aptMode, useLightAnalysis, correctErrorTypes, System.currentTimeMillis(), logger)
         AnalysisHandlerExtension.registerExtension(project, kapt3AnalysisCompletedHandlerExtension)
     }
 
@@ -282,4 +295,25 @@ class Kapt3ComponentRegistrar : ComponentRegistrar {
             return AnalysisResult.Companion.success(bindingTrace.bindingContext, module, shouldGenerateCode = false)
         }
     }
+}
+
+enum class AptMode {
+    WITH_COMPILATION, STUBS_AND_APT, STUBS_ONLY, APT_ONLY;
+
+    val runAnnotationProcessing
+        get() = this != STUBS_ONLY
+
+    val generateStubs
+        get() = this != APT_ONLY
+
+    companion object {
+        // Supports both deprecated APT_ONLY and new APT_MODE options
+        fun parse(mode: String?): AptMode = when (mode) {
+            "true", "stubsAndApt" -> STUBS_AND_APT
+            "false", "compile" -> WITH_COMPILATION
+            "apt" -> APT_ONLY
+            "stubs" -> STUBS_ONLY
+            else -> WITH_COMPILATION
+        }
+    }
 }

plugins/kapt3/src/org/jetbrains/kotlin/kapt3/KaptContext.kt

@@ -21,6 +21,7 @@ import com.sun.tools.javac.jvm.ClassReader
 import com.sun.tools.javac.main.JavaCompiler
 import com.sun.tools.javac.util.Context
 import com.sun.tools.javac.util.Options
+import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.kapt3.javac.KaptJavaCompiler
 import org.jetbrains.kotlin.kapt3.javac.KaptJavaLog
 import org.jetbrains.kotlin.kapt3.javac.KaptTreeMaker
@@ -31,11 +32,12 @@ import org.jetbrains.kotlin.utils.keysToMap
 import org.jetbrains.org.objectweb.asm.tree.ClassNode
 import javax.tools.JavaFileManager
 
-class KaptContext(
+class KaptContext<out GState : GenerationState?>(
         val logger: KaptLogger,
         val bindingContext: BindingContext,
         val compiledClasses: List<ClassNode>,
         val origins: Map<Any, JvmDeclarationOrigin>,
+        val generationState: GState,
         processorOptions: Map<String, String>,
         javacOptions: Map<String, String> = emptyMap()
 ) : AutoCloseable {
@@ -80,5 +82,6 @@ class KaptContext(
     override fun close() {
         compiler.close()
         fileManager.close()
+        generationState?.destroy()
     }
 }

plugins/kapt3/src/org/jetbrains/kotlin/kapt3/annotationProcessing.kt

@@ -30,7 +30,7 @@ import javax.annotation.processing.Processor
 import javax.tools.JavaFileManager
 import com.sun.tools.javac.util.List as JavacList
 
-fun KaptContext.doAnnotationProcessing(
+fun KaptContext<*>.doAnnotationProcessing(
         javaSourceFiles: List<File>,
         processors: List<Processor>,
         compileClasspath: List<File>,

plugins/kapt3/src/org/jetbrains/kotlin/kapt3/stubs/ClassFileToSourceStubConverter.kt

@@ -23,6 +23,7 @@ import com.sun.tools.javac.parser.Tokens
 import com.sun.tools.javac.tree.JCTree
 import com.sun.tools.javac.tree.JCTree.*
 import com.sun.tools.javac.tree.TreeMaker
+import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.kapt3.*
@@ -45,8 +46,7 @@ import javax.tools.JavaFileManager
 import com.sun.tools.javac.util.List as JavacList
 
 class ClassFileToSourceStubConverter(
-        val kaptContext: KaptContext,
+        val kaptContext: KaptContext<GenerationState>,
-        val typeMapper: KotlinTypeMapper,
         val generateNonExistentClass: Boolean,
         val correctErrorTypes: Boolean
 ) {
@@ -463,7 +463,7 @@ class ClassFileToSourceStubConverter(
 
             val superClassConstructorCall = if (superClassConstructor != null) {
                 val args = mapJList(superClassConstructor.valueParameters) { param ->
-                    convertLiteralExpression(getDefaultValue(typeMapper.mapType(param.type)))
+                    convertLiteralExpression(getDefaultValue(kaptContext.generationState.typeMapper.mapType(param.type)))
                 }
                 val call = treeMaker.Apply(JavacList.nil(), treeMaker.SimpleName("super"), args)
                 JavacList.of<JCStatement>(treeMaker.Exec(call))
@@ -584,6 +584,8 @@ class ClassFileToSourceStubConverter(
 
     private fun getMostSuitableSuperTypeForAnonymousType(typeDescriptor: ClassDescriptor): JCExpression {
         val superClass = typeDescriptor.getSuperClassNotAny()
+        val typeMapper = kaptContext.generationState.typeMapper
+
         if (superClass != null) {
             return treeMaker.Type(typeMapper.mapType(superClass))
         } else {

plugins/kapt3/src/org/jetbrains/kotlin/kapt3/stubs/errorTypeConversion.kt

@@ -43,7 +43,7 @@ internal fun convertKtType(
         val kotlinType = converter.kaptContext.bindingContext[BindingContext.TYPE, reference]
         if (kotlinType != null && !kotlinType.containsErrorTypes()) {
             val signatureWriter = BothSignatureWriter(BothSignatureWriter.Mode.TYPE)
-            converter.typeMapper.mapType(kotlinType, signatureWriter,
+            converter.kaptContext.generationState.typeMapper.mapType(kotlinType, signatureWriter,
                     if (shouldBeBoxed) TypeMappingMode.GENERIC_ARGUMENT else TypeMappingMode.DEFAULT)
 
             return SignatureParser(converter.treeMaker).parseFieldSignature(
@@ -74,7 +74,7 @@ private fun convertUserType(type: KtUserType, converter: ClassFileToSourceStubCo
         // This could be List<SomeErrorType> or similar. List should be converted to java.util.List in this case.
         val referenceTarget = converter.kaptContext.bindingContext[BindingContext.REFERENCE_TARGET, type.referenceExpression]
         if (referenceTarget is ClassDescriptor) {
-            treeMaker.FqName(converter.typeMapper.mapType(referenceTarget.defaultType).className)
+            treeMaker.FqName(converter.kaptContext.generationState.typeMapper.mapType(referenceTarget.defaultType).className)
         }
         else {
             treeMaker.SimpleName(referencedName)

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

@@ -25,7 +25,10 @@ import org.jetbrains.kotlin.codegen.GenerationUtils
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.diagnostics.rendering.DefaultErrorMessages
 import org.jetbrains.kotlin.kapt3.AbstractKapt3Extension
+import org.jetbrains.kotlin.kapt3.AptMode
+import org.jetbrains.kotlin.kapt3.AptMode.STUBS_AND_APT
 import org.jetbrains.kotlin.kapt3.Kapt3BuilderFactory
+import org.jetbrains.kotlin.kapt3.KaptContext
 import org.jetbrains.kotlin.kapt3.diagnostic.DefaultErrorMessagesKapt3
 import org.jetbrains.kotlin.kapt3.javac.KaptJavaFileObject
 import org.jetbrains.kotlin.kapt3.stubs.ClassFileToSourceStubConverter
@@ -152,7 +155,7 @@ abstract class AbstractKotlinKapt3IntegrationTest : CodegenTestCase() {
             incrementalDataOutputDir: File
     ) : AbstractKapt3Extension(PathUtil.getJdkClassesRootsFromCurrentJre() + PathUtil.getKotlinPathsForIdeaPlugin().stdlibPath,
                                emptyList(), javaSourceRoots, outputDir, outputDir,
-                               stubsOutputDir, incrementalDataOutputDir, options, emptyMap(), "", true, System.currentTimeMillis(),
+                               stubsOutputDir, incrementalDataOutputDir, options, emptyMap(), "", STUBS_AND_APT, System.currentTimeMillis(),
                                KaptLogger(true), correctErrorTypes = true
     ) {
         internal var savedStubs: String? = null
@@ -174,7 +177,7 @@ abstract class AbstractKotlinKapt3IntegrationTest : CodegenTestCase() {
         }
 
         override fun saveIncrementalData(
-                generationState: GenerationState,
+                kaptContext: KaptContext<GenerationState>,
                 messageCollector: MessageCollector,
                 converter: ClassFileToSourceStubConverter
         ) {
@@ -184,7 +187,7 @@ abstract class AbstractKotlinKapt3IntegrationTest : CodegenTestCase() {
 
             this.savedBindings = converter.bindings
 
-            super.saveIncrementalData(generationState, messageCollector, converter)
+            super.saveIncrementalData(kaptContext, messageCollector, converter)
         }
     }
 }

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

@@ -24,7 +24,7 @@ import org.jetbrains.kotlin.cli.common.messages.MessageRenderer
 import org.jetbrains.kotlin.cli.common.messages.PrintingMessageCollector
 import org.jetbrains.kotlin.codegen.CodegenTestCase
 import org.jetbrains.kotlin.codegen.GenerationUtils
-import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper
+import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.kapt3.Kapt3BuilderFactory
 import org.jetbrains.kotlin.kapt3.KaptContext
 import org.jetbrains.kotlin.kapt3.doAnnotationProcessing
@@ -58,54 +58,51 @@ abstract class AbstractKotlinKapt3Test : CodegenTestCase() {
 
         val txtFile = File(wholeFile.parentFile, wholeFile.nameWithoutExtension + ".txt")
         val classBuilderFactory = Kapt3BuilderFactory()
-        val state = GenerationUtils.compileFiles(myFiles.psiFiles, myEnvironment, classBuilderFactory)
+        val generationState = GenerationUtils.compileFiles(myFiles.psiFiles, myEnvironment, classBuilderFactory)
 
         val logger = KaptLogger(isVerbose = true, messageCollector = messageCollector)
-        val kaptContext = KaptContext(logger, state.bindingContext, classBuilderFactory.compiledClasses,
+        val kaptContext = KaptContext(logger, generationState.bindingContext, classBuilderFactory.compiledClasses,
-                                      classBuilderFactory.origins, processorOptions = emptyMap())
+                                      classBuilderFactory.origins, generationState, processorOptions = emptyMap())
         try {
-            check(kaptContext, state.typeMapper, txtFile, wholeFile)
+            check(kaptContext, txtFile, wholeFile)
-        }
+        } finally {
-        finally {
             kaptContext.close()
         }
     }
 
     protected fun convert(
-            kaptRunner: KaptContext,
+            kaptContext: KaptContext<GenerationState>,
-            typeMapper: KotlinTypeMapper,
             generateNonExistentClass: Boolean,
             correctErrorTypes: Boolean
     ): JavacList<JCCompilationUnit> {
-        val converter = ClassFileToSourceStubConverter(kaptRunner, typeMapper, generateNonExistentClass, correctErrorTypes)
+        val converter = ClassFileToSourceStubConverter(kaptContext, generateNonExistentClass, correctErrorTypes)
         return converter.convert()
     }
 
     protected abstract fun check(
-            kaptRunner: KaptContext,
+            kaptContext: KaptContext<GenerationState>,
-            typeMapper: KotlinTypeMapper,
             txtFile: File,
             wholeFile: File)
 }
 
 abstract class AbstractClassFileToSourceStubConverterTest : AbstractKotlinKapt3Test() {
-    override fun check(kaptRunner: KaptContext, typeMapper: KotlinTypeMapper, txtFile: File, wholeFile: File) {
+    override fun check(kaptContext: KaptContext<GenerationState>, txtFile: File, wholeFile: File) {
         fun isOptionSet(name: String) = wholeFile.useLines { lines -> lines.any { it.trim() == "// $name" } }
 
         val generateNonExistentClass = isOptionSet("NON_EXISTENT_CLASS")
         val correctErrorTypes = isOptionSet("CORRECT_ERROR_TYPES")
         val validate = !isOptionSet("NO_VALIDATION")
 
-        val javaFiles = convert(kaptRunner, typeMapper, generateNonExistentClass, correctErrorTypes)
+        val javaFiles = convert(kaptContext, generateNonExistentClass, correctErrorTypes)
 
-        kaptRunner.javaLog.interceptorData.files = javaFiles.map { it.sourceFile to it }.toMap()
+        kaptContext.javaLog.interceptorData.files = javaFiles.map { it.sourceFile to it }.toMap()
-        if (validate) kaptRunner.compiler.enterTrees(javaFiles)
+        if (validate) kaptContext.compiler.enterTrees(javaFiles)
 
         val actualRaw = javaFiles.sortedBy { it.sourceFile.name }.joinToString (FILE_SEPARATOR)
         val actual = StringUtil.convertLineSeparators(actualRaw.trim({ it <= ' ' })).trimTrailingWhitespacesAndAddNewlineAtEOF()
 
-        if (kaptRunner.compiler.shouldStop(CompileStates.CompileState.ENTER)) {
+        if (kaptContext.compiler.shouldStop(CompileStates.CompileState.ENTER)) {
-            Log.instance(kaptRunner.context).flush()
+            Log.instance(kaptContext.context).flush()
             error("There were errors during analysis. See errors above. Stubs:\n\n$actual")
         }
         KotlinTestUtils.assertEqualsToFile(txtFile, actual)
@@ -113,11 +110,11 @@ abstract class AbstractClassFileToSourceStubConverterTest : AbstractKotlinKapt3T
 }
 
 abstract class AbstractKotlinKaptContextTest : AbstractKotlinKapt3Test() {
-    override fun check(kaptRunner: KaptContext, typeMapper: KotlinTypeMapper, txtFile: File, wholeFile: File) {
+    override fun check(kaptContext: KaptContext<GenerationState>, txtFile: File, wholeFile: File) {
-        val compilationUnits = convert(kaptRunner, typeMapper, generateNonExistentClass = false, correctErrorTypes = true)
+        val compilationUnits = convert(kaptContext, generateNonExistentClass = false, correctErrorTypes = true)
         val sourceOutputDir = Files.createTempDirectory("kaptRunner").toFile()
         try {
-            kaptRunner.doAnnotationProcessing(emptyList(), listOf(JavaKaptContextTest.simpleProcessor()),
+            kaptContext.doAnnotationProcessing(emptyList(), listOf(JavaKaptContextTest.simpleProcessor()),
                                               compileClasspath = PathUtil.getJdkClassesRootsFromCurrentJre() + PathUtil.getKotlinPathsForIdeaPlugin().stdlibPath,
                                               annotationProcessingClasspath = emptyList(), annotationProcessors = "",
                                               sourcesOutputDir = sourceOutputDir, classesOutputDir = sourceOutputDir,

plugins/kapt3/test/org/jetbrains/kotlin/kapt3/test/JavaKaptContextTest.kt

@@ -83,6 +83,7 @@ class JavaKaptContextTest {
                     bindingContext = BindingContext.EMPTY,
                     compiledClasses = emptyList(),
                     origins = emptyMap(),
+                    generationState = null,
                     processorOptions = emptyMap()
         ).doAnnotationProcessing(
                 listOf(javaSourceFile),