Add experimental support for bytecode instrumentation with JPS

#KT-13563 Fixed

idea/idea-jvm/src/org/jetbrains/kotlin/idea/compiler/configuration/KotlinBuildProcessParametersProvider.kt

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.idea.compiler.configuration
 
 import com.intellij.compiler.server.BuildProcessParametersProvider
+import com.intellij.openapi.util.registry.Registry
 import org.jetbrains.kotlin.config.IncrementalCompilation
 import org.jetbrains.kotlin.idea.PluginStartupComponent
 
@@ -35,6 +36,9 @@ class KotlinBuildProcessParametersProvider(
         if (compilerWorkspaceSettings.enableDaemon) {
             res.add("-Dkotlin.daemon.enabled")
         }
+        if (Registry.`is`("kotlin.jps.instrument.bytecode", false)) {
+            res.add("-Dkotlin.jps.instrument.bytecode=true")
+        }
         kotlinPluginStartupComponent.aliveFlagPath.let {
             if (!it.isBlank()) {
                 // TODO: consider taking the property name from compiler/daemon/common (check whether dependency will be not too heavy)

idea/resources/META-INF/jvm.xml

@@ -177,6 +177,11 @@
                    description="Whether to convert `KtStringTemplateExpression` to `KotlinStringTemplateUPolyadicExpression` in all cases"
                    defaultValue="false"
                    restartRequired="false"/>
+
+      <registryKey key="kotlin.jps.instrument.bytecode"
+                   description="Enable bytecode instrumentation for Kotlin classes"
+                   defaultValue="false"
+                   restartRequired="false"/>
     </extensions>
 
     <extensions defaultExtensionNs="org.jetbrains.uast">

idea/resources/META-INF/jvm.xml.182

@@ -171,6 +171,11 @@
                    description="Whether to use multi resolve for UAST in Kotlin provided by `Call.resolveCandidates`, otherwise PsiPolyVariantReference-based multiResolve will be performed"
                    defaultValue="true"
                    restartRequired="false"/>
+
+      <registryKey key="kotlin.jps.instrument.bytecode"
+                   description="Enable bytecode instrumentation for Kotlin classes"
+                   defaultValue="false"
+                   restartRequired="false"/>
     </extensions>
 
     <extensions defaultExtensionNs="org.jetbrains.uast">

idea/resources/META-INF/jvm.xml.183

@@ -172,6 +172,11 @@
                    description="Whether to use multi resolve for UAST in Kotlin provided by `Call.resolveCandidates`, otherwise PsiPolyVariantReference-based multiResolve will be performed"
                    defaultValue="true"
                    restartRequired="false"/>
+
+      <registryKey key="kotlin.jps.instrument.bytecode"
+                   description="Enable bytecode instrumentation for Kotlin classes"
+                   defaultValue="false"
+                   restartRequired="false"/>
     </extensions>
 
     <extensions defaultExtensionNs="org.jetbrains.uast">

jps-plugin/src/org/jetbrains/kotlin/jps/build/KotlinBuilder.kt

@@ -35,6 +35,7 @@ import org.jetbrains.kotlin.cli.common.arguments.CommonCompilerArguments
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity.*
 import org.jetbrains.kotlin.cli.common.messages.MessageCollectorUtil
+import org.jetbrains.kotlin.cli.common.toBooleanLenient
 import org.jetbrains.kotlin.compilerRunner.*
 import org.jetbrains.kotlin.config.IncrementalCompilation
 import org.jetbrains.kotlin.config.KotlinModuleKind
@@ -54,6 +55,7 @@ import org.jetbrains.kotlin.preloading.ClassCondition
 import org.jetbrains.kotlin.utils.KotlinPaths
 import org.jetbrains.kotlin.utils.KotlinPathsFromHomeDir
 import org.jetbrains.kotlin.utils.PathUtil
+import org.jetbrains.org.objectweb.asm.ClassReader
 import java.io.File
 import java.util.*
 import kotlin.collections.HashSet
@@ -424,8 +426,11 @@ class KotlinBuilder : ModuleLevelBuilder(BuilderCategory.SOURCE_PROCESSOR) {
         }
 
         val generatedFiles = getGeneratedFiles(context, chunk, environment.outputItemsCollector)
-
-        registerOutputItems(outputConsumer, generatedFiles)
+        val kotlinTargets = kotlinContext.targetsBinding
+        for ((target, outputItems) in generatedFiles) {
+            val kotlinTarget = kotlinTargets[target] ?: error("Could not find Kotlin target for JPS target $target")
+            kotlinTarget.registerOutputItems(outputConsumer, outputItems)
+        }
         kotlinChunk.saveVersions()
 
         if (targets.any { kotlinContext.hasKotlinMarker[it] == null }) {
@@ -633,14 +638,6 @@ class KotlinBuilder : ModuleLevelBuilder(BuilderCategory.SOURCE_PROCESSOR) {
         return outputItemCollector.outputs.groupBy(SimpleOutputItem::target, SimpleOutputItem::toGeneratedFile)
     }
 
-    private fun registerOutputItems(outputConsumer: OutputConsumer, outputItems: Map<ModuleBuildTarget, List<GeneratedFile>>) {
-        for ((target, outputs) in outputItems) {
-            for (output in outputs) {
-                outputConsumer.registerOutputFile(target, output.outputFile, output.sourceFiles.map { it.path })
-            }
-        }
-    }
-
     private fun updateLookupStorage(
         lookupTracker: LookupTracker,
         lookupStorageManager: JpsLookupStorageManager,

jps-plugin/src/org/jetbrains/kotlin/jps/build/KotlinCompileContext.kt

@@ -13,6 +13,7 @@ import org.jetbrains.jps.incremental.fs.CompilationRound
 import org.jetbrains.jps.incremental.messages.BuildMessage
 import org.jetbrains.jps.incremental.messages.CompilerMessage
 import org.jetbrains.kotlin.config.CompilerRunnerConstants
+import org.jetbrains.kotlin.config.CompilerRunnerConstants.*
 import org.jetbrains.kotlin.incremental.LookupSymbol
 import org.jetbrains.kotlin.incremental.storage.FileToPathConverter
 import org.jetbrains.kotlin.jps.incremental.*
@@ -67,6 +68,15 @@ class KotlinCompileContext(val jpsContext: CompileContext) {
 
     val hasKotlinMarker = HasKotlinMarker(dataManager)
 
+    val isInstrumentationEnabled: Boolean by lazy {
+        val value = System.getProperty("kotlin.jps.instrument.bytecode")?.toBoolean() ?: false
+        if (value) {
+            val message = "Experimental bytecode instrumentation for Kotlin classes is enabled"
+            jpsContext.processMessage(CompilerMessage(KOTLIN_COMPILER_NAME, BuildMessage.Kind.INFO, message))
+        }
+        value
+    }
+
     val fileToPathConverter: FileToPathConverter =
         JpsFileToPathConverter(jpsContext.projectDescriptor.project)
 
@@ -287,7 +297,7 @@ class KotlinCompileContext(val jpsContext: CompileContext) {
             testingLogger?.addCustomMessage(msg)
             jpsContext.processMessage(
                 CompilerMessage(
-                    CompilerRunnerConstants.KOTLIN_COMPILER_NAME,
+                    KOTLIN_COMPILER_NAME,
                     BuildMessage.Kind.WARNING,
                     msg
                 )

jps-plugin/src/org/jetbrains/kotlin/jps/targets/KotlinJvmModuleBuildTarget.kt

@@ -14,8 +14,7 @@ import gnu.trove.THashSet
 import org.jetbrains.jps.ModuleChunk
 import org.jetbrains.jps.builders.java.JavaBuilderUtil
 import org.jetbrains.jps.builders.storage.BuildDataPaths
-import org.jetbrains.jps.incremental.CompileContext
-import org.jetbrains.jps.incremental.ModuleBuildTarget
+import org.jetbrains.jps.incremental.*
 import org.jetbrains.jps.model.java.JpsJavaExtensionService
 import org.jetbrains.jps.model.module.JpsSdkDependency
 import org.jetbrains.kotlin.build.GeneratedFile
@@ -149,6 +148,22 @@ class KotlinJvmModuleBuildTarget(kotlinContext: KotlinCompileContext, jpsModuleB
         return true
     }
 
+    override fun registerOutputItems(outputConsumer: ModuleLevelBuilder.OutputConsumer, outputItems: List<GeneratedFile>) {
+        if (kotlinContext.isInstrumentationEnabled) {
+            val (classFiles, nonClassFiles) = outputItems.partition { it is GeneratedJvmClass }
+            super.registerOutputItems(outputConsumer, nonClassFiles)
+
+            for (output in classFiles) {
+                val bytes = output.outputFile.readBytes()
+                val binaryContent = BinaryContent(bytes)
+                val compiledClass = CompiledClass(output.outputFile, output.sourceFiles, ClassReader(bytes).className, binaryContent)
+                outputConsumer.registerCompiledClass(jpsModuleBuildTarget, compiledClass)
+            }
+        } else {
+            super.registerOutputItems(outputConsumer, outputItems)
+        }
+    }
+
     private fun generateChunkModuleDescription(dirtyFilesHolder: KotlinDirtySourceFilesHolder): File? {
         val builder = KotlinModuleXmlBuilder()
 

jps-plugin/src/org/jetbrains/kotlin/jps/targets/KotlinModuleBuildTarget.kt

@@ -9,6 +9,7 @@ import org.jetbrains.jps.ModuleChunk
 import org.jetbrains.jps.builders.storage.BuildDataPaths
 import org.jetbrains.jps.incremental.CompileContext
 import org.jetbrains.jps.incremental.ModuleBuildTarget
+import org.jetbrains.jps.incremental.ModuleLevelBuilder
 import org.jetbrains.jps.incremental.ProjectBuildException
 import org.jetbrains.jps.model.java.JpsJavaClasspathKind
 import org.jetbrains.jps.model.java.JpsJavaExtensionService
@@ -208,6 +209,12 @@ abstract class KotlinModuleBuildTarget<BuildMetaInfoType : BuildMetaInfo> intern
         environment: JpsCompilerEnvironment
     ): Boolean
 
+    open fun registerOutputItems(outputConsumer: ModuleLevelBuilder.OutputConsumer, outputItems: List<GeneratedFile>) {
+        for (output in outputItems) {
+            outputConsumer.registerOutputFile(jpsModuleBuildTarget, output.outputFile, output.sourceFiles.map { it.path })
+        }
+    }
+
     protected fun reportAndSkipCircular(environment: JpsCompilerEnvironment): Boolean {
         if (chunk.targets.size > 1) {
             // We do not support circular dependencies, but if they are present, we do our best should not break the build,