Kapt: Analyze sources again if the received analysis result is 'RetryWithAdditionalJavaRoots'. ->

Add additional Java roots to Java source roots and clear diagnostic messages before the next analysis round.
(cherry picked from commit 17ad807)

compiler/cli/src/org/jetbrains/kotlin/cli/common/messages/PrintingMessageCollector.java

@@ -32,6 +32,11 @@ public class PrintingMessageCollector implements MessageCollector {
         this.messageRenderer = messageRenderer;
     }
 
+    @Override
+    public void clear() {
+        // Do nothing, messages are already reported
+    }
+
     @Override
     public void report(
             @NotNull CompilerMessageSeverity severity,

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

@@ -64,7 +64,7 @@ class CliLightClassGenerationSupport(project: Project) : LightClassGenerationSup
     private var bindingContext: BindingContext by Delegates.notNull()
     private var module: ModuleDescriptor by Delegates.notNull()
 
-
+    
     override fun initialize(trace: BindingTrace, module: ModuleDescriptor, codeAnalyzer: KotlinCodeAnalyzer) {
         this.bindingContext = trace.bindingContext
         this.module = module

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

@@ -58,16 +58,16 @@ interface JvmDependenciesIndex {
     ): Set<String>
 }
 
-interface JvmDependenciesIndexFactory {
-    fun makeIndexFor(roots: List<JavaRoot>): JvmDependenciesIndex
+interface JvmDependenciesIndexFactory<out T : JvmDependenciesIndex> {
+    fun makeIndexFor(roots: List<JavaRoot>): T
 }
 
-class JvmStaticDependenciesIndexFactory : JvmDependenciesIndexFactory {
-    override fun makeIndexFor(roots: List<JavaRoot>): JvmDependenciesIndex = JvmDependenciesIndexImpl(roots)
+class JvmStaticDependenciesIndexFactory : JvmDependenciesIndexFactory<JvmDependenciesIndex> {
+    override fun makeIndexFor(roots: List<JavaRoot>) = JvmDependenciesIndexImpl(roots)
 }
 
-class JvmUpdatableDependenciesIndexFactory : JvmDependenciesIndexFactory {
-    override fun makeIndexFor(roots: List<JavaRoot>): JvmDependenciesIndex = JvmDependenciesDynamicCompoundIndex().apply {
+class JvmUpdatableDependenciesIndexFactory : JvmDependenciesIndexFactory<JvmDependenciesDynamicCompoundIndex> {
+    override fun makeIndexFor(roots: List<JavaRoot>) = JvmDependenciesDynamicCompoundIndex().apply {
         addIndex(JvmDependenciesIndexImpl(roots))
     }
 }

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

@@ -107,9 +107,9 @@ class KotlinCoreEnvironment private constructor(
         }
     }
     private val sourceFiles = ArrayList<KtFile>()
-    private val rootsIndex: JvmDependenciesIndex
+    private val rootsIndex: JvmDependenciesDynamicCompoundIndex
 
-    val configuration: CompilerConfiguration = configuration.copy().apply { setReadOnly(true) }
+    val configuration: CompilerConfiguration = configuration.copy()
 
     init {
         PersistentFSConstants.setMaxIntellisenseFileSize(FileUtilRt.LARGE_FOR_CONTENT_LOADING)
@@ -154,11 +154,9 @@ class KotlinCoreEnvironment private constructor(
         }
 
         val initialRoots = configuration.getList(JVMConfigurationKeys.CONTENT_ROOTS).classpathRoots()
-
-        // TODO: pass index factory in the configuration to avoid selection logic here
-        // for a moment using updatable index for REPL mode, so we can add roots to classpath then compiling subsequent lines in REPL
-        val indexFactory = if (configuration.getBoolean(CommonConfigurationKeys.REPL_MODE)) JvmUpdatableDependenciesIndexFactory()
-                           else JvmStaticDependenciesIndexFactory()
+        
+        // REPL and kapt2 update classpath dynamically
+        val indexFactory = JvmUpdatableDependenciesIndexFactory()
 
         rootsIndex = indexFactory.makeIndexFor(initialRoots)
         updateClasspathFromRootsIndex(rootsIndex)
@@ -226,19 +224,17 @@ class KotlinCoreEnvironment private constructor(
             projectEnvironment.addSourcesToClasspath(it.file)
         }
     }
+    
+    fun updateClasspath(roots: List<ContentRoot>): List<File>? {
+        return rootsIndex.addNewIndexForRoots(roots.classpathRoots())?.let {
+            updateClasspathFromRootsIndex(it)
+            it.indexedRoots.mapNotNull { File(it.file.path.substringBefore(URLUtil.JAR_SEPARATOR)) }.toList()
+        } ?: emptyList()
+    }
 
     @Suppress("unused") // used externally
-    fun tryUpdateClasspath(files: Iterable<File>): List<File>? =
-        if (rootsIndex !is JvmDependenciesDynamicCompoundIndex) {
-            report(WARNING, "Unable to update classpath after initialization, it is only allowed in REPL")
-            null
-        }
-        else {
-            rootsIndex.addNewIndexForRoots(files.map { JvmClasspathRoot(it) }.classpathRoots())?.let {
-                updateClasspathFromRootsIndex(it)
-                it.indexedRoots.mapNotNull { File(it.file.path.substringBefore(URLUtil.JAR_SEPARATOR)) }.toList()
-            } ?: emptyList()
-        }
+    @Deprecated("Use updateClasspath() instead.", ReplaceWith("updateClasspath(files)"))
+    fun tryUpdateClasspath(files: Iterable<File>): List<File>? = updateClasspath(files.map(::JvmClasspathRoot))
 
     fun contentRootToVirtualFile(root: JvmContentRoot): VirtualFile? {
         when (root) {

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

@@ -18,6 +18,8 @@ package org.jetbrains.kotlin.cli.jvm.compiler
 
 import com.intellij.openapi.util.io.JarUtil
 import org.jetbrains.annotations.TestOnly
+import com.intellij.psi.PsiManager
+import com.intellij.psi.impl.PsiModificationTrackerImpl
 import org.jetbrains.kotlin.analyzer.AnalysisResult
 import org.jetbrains.kotlin.asJava.FilteredJvmDiagnostics
 import org.jetbrains.kotlin.backend.common.output.OutputFileCollection
@@ -34,6 +36,8 @@ import org.jetbrains.kotlin.codegen.ClassBuilderFactories
 import org.jetbrains.kotlin.codegen.CompilationErrorHandler
 import org.jetbrains.kotlin.codegen.GeneratedClassLoader
 import org.jetbrains.kotlin.codegen.KotlinCodegenFacade
+import org.jetbrains.kotlin.cli.jvm.config.*
+import org.jetbrains.kotlin.codegen.*
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.codegen.state.GenerationStateEventCallback
 import org.jetbrains.kotlin.config.CompilerConfiguration
@@ -124,7 +128,26 @@ object KotlinToJVMBytecodeCompiler {
         }
 
         val targetDescription = "in targets [" + chunk.joinToString { input -> input.getModuleName() + "-" + input.getModuleType() } + "]"
-        val result = analyze(environment, targetDescription)
+        var result = analyze(environment, targetDescription)
+        
+        if (result is AnalysisResult.RetryWithAdditionalJavaRoots) {
+            val oldReadOnlyValue = projectConfiguration.isReadOnly
+            projectConfiguration.isReadOnly = false
+            projectConfiguration.addJavaSourceRoots(result.additionalJavaRoots)
+            projectConfiguration.isReadOnly = oldReadOnlyValue
+
+            environment.updateClasspath(result.additionalJavaRoots.map { JavaSourceRoot(it, null) })
+ 
+            // Clear package caches (see KotlinJavaPsiFacade)
+            (PsiManager.getInstance(environment.project).modificationTracker as? PsiModificationTrackerImpl)?.incCounter()
+            
+            // Clear all diagnostic messages
+            projectConfiguration[CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY]?.clear()
+            
+            // Repeat analysis with additional Java roots (kapt generated sources)
+            result = analyze(environment, targetDescription)
+        }
+        
         if (result == null || !result.shouldGenerateCode) return false
 
         ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()
@@ -402,7 +425,12 @@ object KotlinToJVMBytecodeCompiler {
 
         K2JVMCompiler.reportPerf(environment.configuration, message)
 
-        return if (analyzerWithCompilerReport.hasErrors()) null else analyzerWithCompilerReport.analysisResult
+        val analysisResult = analyzerWithCompilerReport.analysisResult
+        
+        return if (!analyzerWithCompilerReport.hasErrors() || analysisResult is AnalysisResult.RetryWithAdditionalJavaRoots)
+            analysisResult
+        else
+            null
     }
 
     private fun generate(