Atomicfu Js/Ir compiler plugin supported in the gradle plugin (#215)

* Atomicfu Js/Ir compiler plugin supported in the gradle plugin

* Fix: apply atomicfu compiler plugin by default

* Add gradle.properties option to turn on IR transformations

* Fix README.md

Co-authored-by: Vsevolod Tolstopyatov <qwwdfsad@gmail.com>

* Fix README.md

Co-authored-by: Vsevolod Tolstopyatov <qwwdfsad@gmail.com>

* Property value toBooleanStrict

Co-authored-by: Vsevolod Tolstopyatov <qwwdfsad@gmail.com>

5 files changed, 103 insertions, 41 deletions

diff --git a/README.md b/README.md
index 87b54d5..9944900 100644
--- a/README.md
+++ b/README.md
@@ -119,7 +119,7 @@ apply plugin: 'kotlinx-atomicfu'
 
 ### JS 
 
-Configure add apply plugin just like for [JVM](#jvm). 
+Configure add apply plugin just like for [JVM](#jvm).
 
 ### Native
 
@@ -149,16 +149,49 @@ dependencies {
 }
 ```
 
-### Additional configuration
+## IR transformation for Kotlin/JS
 
-There are the following additional parameters (with their defaults):
+There is a new option to turn on IR transformation for Kotlin/JS backend.
+You can add `kotlinx.atomicfu.enableIrTransformation=true` to your `gradle.properties` file in order to enable it.
 
+Here is how transformation is performed for different [JS compiler modes](https://kotlinlang.org/docs/js-ir-compiler.html) with this option enabled:
+
+- `kotlin.js.compiler=legacy`: JavaScript transformer from the library is applied to the final compiled *.js files.
+- `kotlin.js.compiler=ir`: compiler plugin transformations are appiled to the generated IR.
+- `kotlin.js.compiler=both`: compiler plugin transformations are appiled to all compilations of IR targets, while compilations of legacy targets are transformed by the library.
+
+## Additional configuration
+
+To set configuration options you should create `atomicfu` section in a `build.gradle` file, 
+like this:
+```groovy
+atomicfu {
+  dependenciesVersion = '0.17.1'
+}
+```
+
+### JVM transformation options
+
+To turn off transformation for Kotlin/JVM set option `transformJvm` to `false`.
+
+Configuration option `jvmVariant` defines the Java class that replaces atomics during bytecode transformation.
+Here are the valid options:
+- `FU` – atomics are replaced with [AtomicXxxFieldUpdater](https://docs.oracle.com/javase/10/docs/api/java/util/concurrent/atomic/AtomicIntegerFieldUpdater.html).
+- `VH` – atomics are replaced with [VarHandle](https://docs.oracle.com/javase/9/docs/api/java/lang/invoke/VarHandle.html), 
+  this option is supported for JDK 9+.
+- `BOTH` – [multi-release jar file](https://openjdk.java.net/jeps/238) will be created with both `AtomicXxxFieldUpdater` for JDK <= 8 and `VarHandle` for JDK 9+.
+
+### JS transformation options
+
+To turn off transformation for Kotlin/JS set option `transformJs` to `false`.
+
+Here are all available configuration options (with their defaults):
 ```groovy
 atomicfu {
   dependenciesVersion = '0.17.1' // set to null to turn-off auto dependencies
   transformJvm = true // set to false to turn off JVM transformation
-  transformJs = true // set to false to turn off JS transformation
-  variant = "FU" // JVM transformation variant: FU,VH, or BOTH 
+  jvmVariant = "FU" // JVM transformation variant: FU,VH, or BOTH 
+  jsVariant = "JS" // JS transformation variant: JS or IR
   verbose = false // set to true to be more verbose  
 }
 ```
@@ -237,22 +270,6 @@ which is then transformed to a regular `classes` directory to be used later by t
 
 AtomicFU provides some additional features that you can optionally use.
 
-### VarHandles with Java 9
-
-AtomicFU can produce code that uses Java 9 
-[VarHandle](https://docs.oracle.com/javase/9/docs/api/java/lang/invoke/VarHandle.html)
-instead of `AtomicXxxFieldUpdater`. Configure transformation `variant` in Gradle build file:
- 
-```groovy
-atomicfu {
-    variant = "VH"
-}
-``` 
- 
-It can also create [JEP 238](https://openjdk.java.net/jeps/238) multi-release jar file with both
-`AtomicXxxFieldUpdater` for JDK<=8 and `VarHandle` for for JDK9+ if you 
-set `variant` to `"BOTH"`.
-
 ### Arrays of atomic values
 
 You can declare arrays of all supported atomic value types. 
diff --git a/atomicfu-gradle-plugin/build.gradle b/atomicfu-gradle-plugin/build.gradle
index 58ee3a1..79ba241 100644
--- a/atomicfu-gradle-plugin/build.gradle
+++ b/atomicfu-gradle-plugin/build.gradle
@@ -25,6 +25,8 @@ dependencies {
     compile 'org.jetbrains.kotlin:kotlin-stdlib'
 
     compileOnly "org.jetbrains.kotlin:kotlin-gradle-plugin:$kotlin_version"
+    // atomicfu compiler plugin dependency will be loaded to kotlinCompilerPluginClasspath
+    implementation "org.jetbrains.kotlin:atomicfu:$kotlin_version"
 
     testCompile gradleTestKit()
     testCompile 'org.jetbrains.kotlin:kotlin-test'
diff --git a/atomicfu-gradle-plugin/src/main/kotlin/kotlinx/atomicfu/plugin/gradle/AtomicFUGradlePlugin.kt b/atomicfu-gradle-plugin/src/main/kotlin/kotlinx/atomicfu/plugin/gradle/AtomicFUGradlePlugin.kt
index 16953d1..b1dd098 100644
--- a/atomicfu-gradle-plugin/src/main/kotlin/kotlinx/atomicfu/plugin/gradle/AtomicFUGradlePlugin.kt
+++ b/atomicfu-gradle-plugin/src/main/kotlin/kotlinx/atomicfu/plugin/gradle/AtomicFUGradlePlugin.kt
@@ -18,18 +18,25 @@ import org.jetbrains.kotlin.gradle.plugin.*
 import java.io.*
 import java.util.*
 import java.util.concurrent.*
+import org.jetbrains.kotlin.gradle.targets.js.*
+import org.jetbrains.kotlin.gradle.targets.js.ir.KotlinJsIrTarget
+import org.jetbrains.kotlinx.atomicfu.gradle.*
 
 private const val EXTENSION_NAME = "atomicfu"
 private const val ORIGINAL_DIR_NAME = "originalClassesDir"
 private const val COMPILE_ONLY_CONFIGURATION = "compileOnly"
 private const val IMPLEMENTATION_CONFIGURATION = "implementation"
 private const val TEST_IMPLEMENTATION_CONFIGURATION = "testImplementation"
+private const val ENABLE_IR_TRANSFORMATION = "kotlinx.atomicfu.enableIrTransformation"
 
 open class AtomicFUGradlePlugin : Plugin<Project> {
     override fun apply(project: Project) = project.run {
         val pluginVersion = rootProject.buildscript.configurations.findByName("classpath")
             ?.allDependencies?.find { it.name == "atomicfu-gradle-plugin" }?.version
         extensions.add(EXTENSION_NAME, AtomicFUPluginExtension(pluginVersion))
+        if (rootProject.getBooleanProperty(ENABLE_IR_TRANSFORMATION)) {
+            plugins.apply(AtomicfuKotlinGradleSubplugin::class.java)
+        }
         configureDependencies()
         configureTasks()
     }
@@ -49,6 +56,7 @@ private fun Project.configureDependencies() {
             getAtomicfuDependencyNotation(Platform.JS, version)
         )
         dependencies.add(TEST_IMPLEMENTATION_CONFIGURATION, getAtomicfuDependencyNotation(Platform.JS, version))
+        addCompilerPluginDependency()
     }
     withPluginWhenEvaluatedDependencies("kotlin-multiplatform") { version ->
         configureMultiplatformPluginDependencies(version)
@@ -78,6 +86,36 @@ private fun Project.configureTasks() {
     }
 }
 
+private fun Project.getBooleanProperty(name: String) =
+    rootProject.findProperty(name)?.toString()?.toBooleanStrict() ?: false
+
+private fun String.toBooleanStrict(): Boolean = when (this) {
+    "true" -> true
+    "false" -> false
+    else -> throw IllegalArgumentException("The string doesn't represent a boolean value: $this")
+}
+
+private fun Project.needsJsIrTransformation(target: KotlinTarget): Boolean =
+    config.transformJs && target.isJsIrTarget()
+
+private fun KotlinTarget.isJsIrTarget() = (this is KotlinJsTarget && this.irTarget != null) || this is KotlinJsIrTarget
+
+private fun Project.addCompilerPluginDependency() {
+    val kotlinVersion = rootProject.buildscript.configurations.findByName("classpath")
+        ?.allDependencies?.find { it.name == "kotlin-gradle-plugin" }?.version
+    withKotlinTargets { target ->
+        if (needsJsIrTransformation(target)) {
+            target.compilations.forEach { kotlinCompilation ->
+                kotlinCompilation.dependencies {
+                    // add atomicfu compiler plugin dependency
+                    // to provide the `kotlinx-atomicfu-runtime` library used during compiler plugin transformation
+                    compileOnly("org.jetbrains.kotlin:atomicfu:$kotlinVersion")
+                }
+            }
+        }
+    }
+}
+
 private enum class Platform(val suffix: String) {
     JVM("-jvm"),
     JS("-js"),
@@ -129,9 +167,9 @@ fun Project.withKotlinTargets(fn: (KotlinTarget) -> Unit) {
     extensions.findByType(KotlinProjectExtension::class.java)?.let { kotlinExtension ->
         val targetsExtension = (kotlinExtension as? ExtensionAware)?.extensions?.findByName("targets")
         @Suppress("UNCHECKED_CAST")
-        val targets = targetsExtension as NamedDomainObjectContainer<KotlinTarget>
+        val targets = targetsExtension as? NamedDomainObjectContainer<KotlinTarget>
         // find all compilations given sourceSet belongs to
-        targets.all { target -> fn(target) }
+        targets?.all { target -> fn(target) }
     }
 }
 
@@ -180,7 +218,10 @@ private fun Project.configureTransformationForTarget(target: KotlinTarget) {
                 )
             }
             KotlinPlatformType.js -> {
-                if (!config.transformJs) return@compilations // skip when transformation is turned off
+                // skip when js transformation is not needed or when IR is transformed
+                if (!config.transformJs || (needsJsIrTransformation(target))) {
+                    return@compilations
+                }
                 project.createJsTransformTask(compilation).configureJsTask(
                     compilation.compileAllTaskName,
                     transformedClassesDir,
@@ -194,7 +235,7 @@ private fun Project.configureTransformationForTarget(target: KotlinTarget) {
         classesDirs.setFrom(transformedClassesDir)
         classesDirs.builtBy(transformTask)
         (tasks.findByName(target.artifactsTaskName) as? Jar)?.apply {
-            setupJarManifest(multiRelease = config.variant.toVariant() == Variant.BOTH)
+            setupJarManifest(multiRelease = config.jvmVariant.toJvmVariant() == JvmVariant.BOTH)
         }
         // test should compile and run against original production binaries
         if (compilationType == CompilationType.TEST) {
@@ -232,7 +273,8 @@ fun Project.sourceSetsByCompilation(): Map<KotlinSourceSet, List<KotlinCompilati
 }
 
 fun Project.configureMultiplatformPluginDependencies(version: String) {
-    if (rootProject.findProperty("kotlin.mpp.enableGranularSourceSetsMetadata").toString().toBoolean()) {
+    if (rootProject.getBooleanProperty("kotlin.mpp.enableGranularSourceSetsMetadata")) {
+        addCompilerPluginDependency()
         val mainConfigurationName = project.extensions.getByType(KotlinMultiplatformExtension::class.java).sourceSets
                 .getByName(KotlinSourceSet.COMMON_MAIN_SOURCE_SET_NAME)
                 .compileOnlyConfigurationName
@@ -256,6 +298,7 @@ fun Project.configureMultiplatformPluginDependencies(version: String) {
         }
     } else {
         sourceSetsByCompilation().forEach { (sourceSet, compilations) ->
+            addCompilerPluginDependency()
             val platformTypes = compilations.map { it.platformType }.toSet()
             val compilationNames = compilations.map { it.compilationName }.toSet()
             if (compilationNames.size != 1)
@@ -303,7 +346,7 @@ fun Project.configureJvmTransformation(
         //now transformTask is responsible for compiling this source set into the classes directory
         sourceSet.compiledBy(transformTask)
         (tasks.findByName(sourceSet.jarTaskName) as? Jar)?.apply {
-            setupJarManifest(multiRelease = config.variant.toVariant() == Variant.BOTH)
+            setupJarManifest(multiRelease = config.jvmVariant.toJvmVariant() == JvmVariant.BOTH)
         }
         // test should compile and run against original production binaries
         if (compilationType == CompilationType.TEST) {
@@ -329,7 +372,7 @@ fun Project.configureJvmTransformation(
     }
 }
 
-fun String.toVariant(): Variant = enumValueOf(toUpperCase(Locale.US))
+fun String.toJvmVariant(): JvmVariant = enumValueOf(toUpperCase(Locale.US))
 
 fun Project.createJvmTransformTask(compilation: KotlinCompilation<*>): AtomicFUTransformTask =
     tasks.create(
@@ -358,7 +401,7 @@ fun AtomicFUTransformTask.configureJvmTask(
         classPath = classpath
         inputFiles = originalClassesDir
         outputDir = transformedClassesDir
-        variant = config.variant
+        jvmVariant = config.jvmVariant
         verbose = config.verbose
     }
 
@@ -390,7 +433,7 @@ class AtomicFUPluginExtension(pluginVersion: String?) {
     var dependenciesVersion = pluginVersion
     var transformJvm = true
     var transformJs = true
-    var variant: String = "FU"
+    var jvmVariant: String = "FU"
     var verbose: Boolean = false
 }
 
@@ -408,7 +451,7 @@ open class AtomicFUTransformTask : ConventionTask() {
     lateinit var classPath: FileCollection
 
     @Input
-    var variant = "FU"
+    var jvmVariant = "FU"
     @Input
     var verbose = false
 
@@ -417,7 +460,7 @@ open class AtomicFUTransformTask : ConventionTask() {
         val cp = classPath.files.map { it.absolutePath }
         inputFiles.files.forEach { inputDir ->
             AtomicFUTransformer(cp, inputDir, outputDir).let { t ->
-                t.variant = variant.toVariant()
+                t.jvmVariant = jvmVariant.toJvmVariant()
                 t.verbose = verbose
                 t.transform()
             }
diff --git a/atomicfu-maven-plugin/src/main/kotlin/kotlinx/atomicfu/plugin/TransformMojo.kt b/atomicfu-maven-plugin/src/main/kotlin/kotlinx/atomicfu/plugin/TransformMojo.kt
index 28580d8..5f45f3e 100644
--- a/atomicfu-maven-plugin/src/main/kotlin/kotlinx/atomicfu/plugin/TransformMojo.kt
+++ b/atomicfu-maven-plugin/src/main/kotlin/kotlinx/atomicfu/plugin/TransformMojo.kt
@@ -17,7 +17,7 @@
 package kotlinx.atomicfu.plugin
 
 import kotlinx.atomicfu.transformer.AtomicFUTransformer
-import kotlinx.atomicfu.transformer.Variant
+import kotlinx.atomicfu.transformer.JvmVariant
 import org.apache.maven.plugin.AbstractMojo
 import org.apache.maven.plugins.annotations.LifecyclePhase
 import org.apache.maven.plugins.annotations.Mojo
@@ -50,8 +50,8 @@ class TransformMojo : AbstractMojo() {
     /**
      * Transformation variant: "FU", "VH", or "BOTH".
      */
-    @Parameter(defaultValue = "FU", property = "variant", required = true)
-    lateinit var variant: Variant
+    @Parameter(defaultValue = "FU", property = "jvmVariant", required = true)
+    lateinit var jvmVariant: JvmVariant
 
     /**
      * Verbose debug info.
@@ -60,7 +60,7 @@ class TransformMojo : AbstractMojo() {
     var verbose: Boolean = false
 
     override fun execute() {
-        val t = AtomicFUTransformer(classpath, input, output, variant)
+        val t = AtomicFUTransformer(classpath, input, output, jvmVariant)
         t.verbose = verbose
         t.transform()
     }
diff --git a/atomicfu-transformer/src/main/kotlin/kotlinx/atomicfu/transformer/AtomicFUTransformer.kt b/atomicfu-transformer/src/main/kotlin/kotlinx/atomicfu/transformer/AtomicFUTransformer.kt
index ca6c143..a138422 100644
--- a/atomicfu-transformer/src/main/kotlin/kotlinx/atomicfu/transformer/AtomicFUTransformer.kt
+++ b/atomicfu-transformer/src/main/kotlin/kotlinx/atomicfu/transformer/AtomicFUTransformer.kt
@@ -168,13 +168,13 @@ class FieldInfo(
     override fun toString(): String = "${owner.prettyStr()}::$name"
 }
 
-enum class Variant { FU, VH, BOTH }
+enum class JvmVariant { FU, VH, BOTH }
 
 class AtomicFUTransformer(
     classpath: List<String>,
     inputDir: File,
     outputDir: File = inputDir,
-    var variant: Variant = Variant.FU
+    var jvmVariant: JvmVariant = JvmVariant.FU
 ) : AtomicFUTransformerBase(inputDir, outputDir) {
 
     private val classPathLoader = URLClassLoader(
@@ -195,7 +195,7 @@ class AtomicFUTransformer(
         val files = inputDir.walk().filter { it.isFile }.toList()
         val needTransform = analyzeFilesForFields(files)
         if (needTransform || outputDir == inputDir) {
-            val vh = variant == Variant.VH
+            val vh = jvmVariant == JvmVariant.VH
             // visit method bodies for external references to fields, runs all logic, fails if anything is wrong
             val needsTransform = analyzeFilesForRefs(files, vh)
             // perform transformation
@@ -205,7 +205,7 @@ class AtomicFUTransformer(
                 val outBytes = if (file.isClassFile() && file in needsTransform) transformFile(file, bytes, vh) else bytes
                 val outFile = file.toOutputFile()
                 outFile.mkdirsAndWrite(outBytes)
-                if (variant == Variant.BOTH && outBytes !== bytes) {
+                if (jvmVariant == JvmVariant.BOTH && outBytes !== bytes) {
                     val vhBytes = transformFile(file, bytes, true)
                     val vhFile = outputDir / "META-INF" / "versions" / "9" / file.relativeTo(inputDir).toString()
                     vhFile.mkdirsAndWrite(vhBytes)
@@ -1512,7 +1512,7 @@ fun main(args: Array<String>) {
     }
     val t = AtomicFUTransformer(emptyList(), File(args[0]))
     if (args.size > 1) t.outputDir = File(args[1])
-    if (args.size > 2) t.variant = enumValueOf(args[2].toUpperCase(Locale.US))
+    if (args.size > 2) t.jvmVariant = enumValueOf(args[2].toUpperCase(Locale.US))
     t.verbose = true
     t.transform()
 }