[Gradle, Wasm] Binaryen to wasm targets

^KT-62024 fixed
^KT-62025 fixed

libraries/tools/kotlin-gradle-plugin-integration-tests/src/test/kotlin/org/jetbrains/kotlin/gradle/KotlinWasmGradlePluginIT.kt

@@ -61,4 +61,48 @@ class KotlinWasmGradlePluginIT : KGPBaseTest() {
             }
         }
     }
+
+    @DisplayName("Check wasi target")
+    @GradleTest
+    fun wasiTargetWithBinaryen(gradleVersion: GradleVersion) {
+        project("new-mpp-wasm-wasi-test", gradleVersion) {
+            buildGradleKts.modify(::transformBuildScriptWithPluginsDsl)
+            buildGradleKts.modify {
+                it.replace("wasmWasi {", "wasmWasi {\napplyBinaryen()\nbinaries.executable()")
+            }
+
+            build("assemble") {
+                assertTasksExecuted(":compileProductionExecutableKotlinWasmWasi")
+                assertTasksExecuted(":compileProductionExecutableKotlinWasmWasiOptimize")
+
+                val original = projectPath.resolve("build/compileSync/wasmWasi/main/productionExecutable/kotlin/new-mpp-wasm-wasi-test-wasm-wasi.wasm")
+                val optimized = projectPath.resolve("build/compileSync/wasmWasi/main/productionExecutable/optimized/new-mpp-wasm-wasi-test-wasm-wasi.wasm")
+                assertTrue {
+                    Files.size(original) > Files.size(optimized)
+                }
+            }
+        }
+    }
+
+    @DisplayName("Check js target")
+    @GradleTest
+    fun jsTargetWithBinaryen(gradleVersion: GradleVersion) {
+        project("new-mpp-wasm-js", gradleVersion) {
+            buildGradleKts.modify(::transformBuildScriptWithPluginsDsl)
+            buildGradleKts.modify {
+                it.replace("wasmJs {", "wasmJs {\napplyBinaryen()")
+            }
+
+            build("assemble") {
+                assertTasksExecuted(":compileProductionExecutableKotlinWasmJs")
+                assertTasksExecuted(":compileProductionExecutableKotlinWasmJsOptimize")
+
+                val original = projectPath.resolve("build/compileSync/wasmJs/main/productionExecutable/kotlin/redefined-wasm-module-name.wasm")
+                val optimized = projectPath.resolve("build/compileSync/wasmJs/main/productionExecutable/optimized/redefined-wasm-module-name.wasm")
+                assertTrue {
+                    Files.size(original) > Files.size(optimized)
+                }
+            }
+        }
+    }
 }

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/targets/js/binaryen/BinaryenExec.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.gradle.targets.js.binaryen
 
+import org.gradle.api.file.FileSystemOperations
 import org.gradle.api.file.RegularFileProperty
 import org.gradle.api.tasks.*
 import org.gradle.work.DisableCachingByDefault
@@ -15,13 +16,16 @@ import org.jetbrains.kotlin.gradle.utils.newFileProperty
 import javax.inject.Inject
 
 @DisableCachingByDefault
-open class BinaryenExec
+abstract class BinaryenExec
 @Inject
 constructor() : AbstractExecTask<BinaryenExec>(BinaryenExec::class.java) {
     @Transient
     @get:Internal
     lateinit var binaryen: BinaryenRootExtension
 
+    @get:Inject
+    abstract val fs: FileSystemOperations
+
     init {
         onlyIf {
             !inputFileProperty.isPresent || inputFileProperty.asFile.map { it.exists() }.get()

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/targets/js/dsl/KotlinWasmJsTargetDsl.kt

@@ -16,8 +16,6 @@ interface KotlinWasmSubTargetContainerDsl : KotlinTarget {
     val isD8Configured: Boolean
 
     fun whenD8Configured(body: KotlinWasmD8Dsl.() -> Unit)
-
-    fun whenBinaryenApplied(body: (BinaryenExec.() -> Unit) -> Unit)
 }
 
 interface KotlinWasmJsTargetDsl : KotlinWasmTargetDsl, KotlinJsTargetDsl {
@@ -28,14 +26,6 @@ interface KotlinWasmJsTargetDsl : KotlinWasmTargetDsl, KotlinJsTargetDsl {
             fn.execute(this)
         }
     }
-
-    fun applyBinaryen() = applyBinaryen { }
-    fun applyBinaryen(body: BinaryenExec.() -> Unit)
-    fun applyBinaryen(fn: Action<BinaryenExec>) {
-        applyBinaryen {
-            fn.execute(this)
-        }
-    }
 }
 
 interface KotlinWasmD8Dsl : KotlinJsSubTargetDsl {

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/targets/js/dsl/KotlinWasmTargetDsl.kt

@@ -5,9 +5,19 @@
 
 package org.jetbrains.kotlin.gradle.targets.js.dsl
 
+import org.gradle.api.Action
 import org.jetbrains.kotlin.gradle.plugin.KotlinTarget
 import org.jetbrains.kotlin.gradle.targets.js.KotlinWasmTargetType
+import org.jetbrains.kotlin.gradle.targets.js.binaryen.BinaryenExec
 
 interface KotlinWasmTargetDsl : KotlinTarget {
     val wasmTargetType: KotlinWasmTargetType?
+
+    fun applyBinaryen() = applyBinaryen { }
+    fun applyBinaryen(body: BinaryenExec.() -> Unit)
+    fun applyBinaryen(fn: Action<BinaryenExec>) {
+        applyBinaryen {
+            fn.execute(this)
+        }
+    }
 }

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/targets/js/ir/JsBinaries.kt

@@ -10,6 +10,7 @@ import org.gradle.api.tasks.TaskProvider
 import org.jetbrains.kotlin.gradle.plugin.KotlinCompilation
 import org.jetbrains.kotlin.gradle.plugin.KotlinTarget
 import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinJsCompilation
+import org.jetbrains.kotlin.gradle.targets.js.binaryen.BinaryenExec
 import org.jetbrains.kotlin.gradle.targets.js.dsl.Distribution
 import org.jetbrains.kotlin.gradle.targets.js.dsl.KotlinJsBinaryMode
 import org.jetbrains.kotlin.gradle.targets.js.ir.KotlinJsBinaryContainer.Companion.generateBinaryName
@@ -83,7 +84,7 @@ sealed class JsIrBinary(
         get() = target.project
 }
 
-class Executable(
+open class Executable(
     compilation: KotlinJsCompilation,
     name: String,
     mode: KotlinJsBinaryMode
@@ -107,6 +108,26 @@ class Executable(
         )
 }
 
+open class ExecutableWasm(
+    compilation: KotlinJsCompilation,
+    name: String,
+    mode: KotlinJsBinaryMode
+) : Executable(
+    compilation,
+    name,
+    mode
+) {
+    val optimizeTaskName: String = optimizeTaskName()
+
+    val optimizeTask: TaskProvider<BinaryenExec>
+        get() = target.project.tasks
+            .withType<BinaryenExec>()
+            .named(optimizeTaskName)
+
+    private fun optimizeTaskName(): String =
+        "${linkTaskName}Optimize"
+}
+
 class Library(
     compilation: KotlinJsCompilation,
     name: String,

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/targets/js/ir/KotlinJsBinaryContainer.kt

@@ -9,12 +9,14 @@ import org.gradle.api.DomainObjectSet
 import org.gradle.api.GradleException
 import org.gradle.api.Project
 import org.gradle.api.plugins.ExtensionAware
+import org.gradle.language.base.plugins.LifecycleBasePlugin
 import org.jetbrains.kotlin.gradle.plugin.KotlinCompilation
 import org.jetbrains.kotlin.gradle.plugin.KotlinPlatformType
 import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinJsCompilation
 import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinTargetWithBinaries
 import org.jetbrains.kotlin.gradle.plugin.mpp.isMain
 import org.jetbrains.kotlin.gradle.targets.js.KotlinJsTarget
+import org.jetbrains.kotlin.gradle.targets.js.KotlinWasmTargetType
 import org.jetbrains.kotlin.gradle.targets.js.binaryen.BinaryenExec
 import org.jetbrains.kotlin.gradle.targets.js.dsl.Distribution
 import org.jetbrains.kotlin.gradle.targets.js.dsl.KotlinJsBinaryMode
@@ -23,6 +25,7 @@ import org.jetbrains.kotlin.gradle.targets.js.dsl.KotlinJsBinaryMode.PRODUCTION
 import org.jetbrains.kotlin.gradle.targets.js.subtargets.DefaultDistribution
 import org.jetbrains.kotlin.gradle.targets.js.subtargets.KotlinJsSubTarget
 import org.jetbrains.kotlin.gradle.targets.js.subtargets.createDefaultDistribution
+import org.jetbrains.kotlin.gradle.tasks.dependsOn
 import org.jetbrains.kotlin.gradle.utils.lowerCamelCaseName
 import org.jetbrains.kotlin.util.capitalizeDecapitalize.toLowerCaseAsciiOnly
 import javax.inject.Inject
@@ -41,26 +44,6 @@ constructor(
     private val defaultCompilation: KotlinJsCompilation
         get() = target.compilations.getByName(KotlinCompilation.MAIN_COMPILATION_NAME)
 
-    private fun configureBinaryen(binary: JsIrBinary, binaryenDsl: BinaryenExec.() -> Unit) {
-        val linkTask = binary.linkTask
-
-        val compiledWasmFile = linkTask.map { link ->
-            link.destinationDirectory.asFile.get().resolve(link.compilerOptions.moduleName.get() + ".wasm")
-        }
-
-        //TODO This is temporary solution that overrides compiled files that triggers recompile and reoptimize wasm every time (when binaryen is enabled)
-        val binaryenTask = BinaryenExec.create(binary.compilation, "${linkTask.name}Optimize") {
-            dependsOn(linkTask)
-            inputFileProperty.fileProvider(compiledWasmFile)
-            outputFileProperty.fileProvider(compiledWasmFile)
-            binaryenDsl()
-        }
-
-        binary.linkSyncTask.configure {
-            it.dependsOn(binaryenTask)
-        }
-    }
-
     // For Groovy DSL
     @JvmOverloads
     fun executable(
@@ -103,7 +86,13 @@ constructor(
     internal fun executableIrInternal(compilation: KotlinJsCompilation): List<JsBinary> = createBinaries(
         compilation = compilation,
         jsBinaryType = KotlinJsBinaryType.EXECUTABLE,
-        create = ::Executable
+        create = { compilation, name, mode ->
+            if (target.platformType == KotlinPlatformType.wasm) {
+                ExecutableWasm(compilation, name, mode)
+            } else {
+                Executable(compilation, name, mode)
+            }
+        }
     )
 
     private fun executableLegacyInternal(compilation: KotlinJsCompilation) = createBinaries(
@@ -198,12 +187,6 @@ constructor(
             extensions.add(binary.name, binary)
         }
 
-        if (compilation.platformType == KotlinPlatformType.wasm && target is KotlinJsIrTarget && binary is JsIrBinary) {
-            target.whenBinaryenApplied {
-                configureBinaryen(binary, it)
-            }
-        }
-
         return binary
     }
 

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/targets/js/ir/KotlinJsIrSubTarget.kt

@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.gradle.plugin.mpp.isMain
 import org.jetbrains.kotlin.gradle.plugin.whenEvaluated
 import org.jetbrains.kotlin.gradle.targets.js.KotlinJsPlatformTestRun
 import org.jetbrains.kotlin.gradle.targets.js.KotlinWasmTargetType
+import org.jetbrains.kotlin.gradle.targets.js.binaryen.BinaryenExec
 import org.jetbrains.kotlin.gradle.targets.js.dsl.Distribution
 import org.jetbrains.kotlin.gradle.targets.js.dsl.ExperimentalDistributionDsl
 import org.jetbrains.kotlin.gradle.targets.js.dsl.KotlinJsBinaryMode
@@ -26,6 +27,7 @@ import org.jetbrains.kotlin.gradle.targets.js.dsl.KotlinJsSubTargetDsl
 import org.jetbrains.kotlin.gradle.targets.js.npm.npmProject
 import org.jetbrains.kotlin.gradle.targets.js.testing.KotlinJsTest
 import org.jetbrains.kotlin.gradle.tasks.dependsOn
+import org.jetbrains.kotlin.gradle.tasks.locateTask
 import org.jetbrains.kotlin.gradle.tasks.registerTask
 import org.jetbrains.kotlin.gradle.testing.internal.configureConventions
 import org.jetbrains.kotlin.gradle.testing.internal.kotlinTestRegistry
@@ -130,9 +132,16 @@ abstract class KotlinJsIrSubTarget(
                     }
                 }
             } else {
-                testJs.dependsOn(binary.linkTask)
+                if (project.locateTask<BinaryenExec>((binary as ExecutableWasm).optimizeTaskName) != null) {
-                binary.linkTask.flatMap { linkTask ->
+                    testJs.dependsOn(binary.optimizeTask)
-                    linkTask.outputFileProperty
+                    binary.optimizeTask.flatMap { optimizeTask ->
+                        optimizeTask.outputFileProperty.asFile
+                    }
+                } else {
+                    testJs.dependsOn(binary.linkTask)
+                    binary.linkTask.flatMap { linkTask ->
+                        linkTask.outputFileProperty
+                    }
                 }
             }
 

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/targets/js/ir/KotlinJsIrTarget.kt

@@ -8,7 +8,9 @@ package org.jetbrains.kotlin.gradle.targets.js.ir
 import org.gradle.api.NamedDomainObjectContainer
 import org.gradle.api.Project
 import org.gradle.api.Task
+import org.gradle.api.provider.Provider
 import org.gradle.api.tasks.TaskProvider
+import org.gradle.language.base.plugins.LifecycleBasePlugin
 import org.jetbrains.kotlin.gradle.dsl.KotlinJsOptions
 import org.jetbrains.kotlin.gradle.dsl.KotlinMultiplatformExtension
 import org.jetbrains.kotlin.gradle.dsl.kotlinExtension
@@ -31,8 +33,8 @@ import org.jetbrains.kotlin.gradle.targets.js.nodejs.NodeJsRootPlugin
 import org.jetbrains.kotlin.gradle.targets.js.npm.NpmResolverPlugin
 import org.jetbrains.kotlin.gradle.targets.js.npm.npmProject
 import org.jetbrains.kotlin.gradle.targets.js.typescript.TypeScriptValidationTask
-import org.jetbrains.kotlin.gradle.tasks.locateOrRegisterTask
+import org.jetbrains.kotlin.gradle.tasks.*
-import org.jetbrains.kotlin.gradle.tasks.registerTask
+import org.jetbrains.kotlin.gradle.utils.*
 import org.jetbrains.kotlin.gradle.utils.dashSeparatedName
 import org.jetbrains.kotlin.gradle.utils.decamelize
 import org.jetbrains.kotlin.gradle.utils.lowerCamelCaseName
@@ -41,6 +43,7 @@ import org.jetbrains.kotlin.gradle.utils.setProperty
 import org.jetbrains.kotlin.util.capitalizeDecapitalize.capitalizeAsciiOnly
 import org.jetbrains.kotlin.util.capitalizeDecapitalize.toLowerCaseAsciiOnly
 import org.jetbrains.kotlin.utils.addIfNotNull
+import java.io.File
 import javax.inject.Inject
 
 abstract class KotlinJsIrTarget
@@ -160,13 +163,14 @@ constructor(
             }
     }
 
-    private val commonLazy by lazy {
+    private val commonLazyDelegate = lazy {
         NpmResolverPlugin.apply(project)
         compilations.all { compilation ->
             compilation.binaries
                 .withType(JsIrBinary::class.java)
                 .all { binary ->
                     val syncTask = registerCompileSync(binary)
+                    binaryenReplaceInput[binary]?.invoke()
                     val tsValidationTask = registerTypeScriptCheckTask(binary)
 
                     binary.linkTask.configure {
@@ -181,6 +185,8 @@ constructor(
         }
     }
 
+    private val commonLazy by commonLazyDelegate
+
     private fun registerCompileSync(binary: JsIrBinary): TaskProvider<DefaultIncrementalSyncTask> {
         val compilation = binary.compilation
         val npmProject = compilation.npmProject
@@ -215,25 +221,96 @@ constructor(
     }
 
     //Binaryen
-    private val applyBinaryenHandlers = mutableListOf<(BinaryenExec.() -> Unit) -> Unit>()
+    private var binaryenReplaceInput: MutableMap<JsIrBinary, () -> Unit> = mutableMapOf()
 
-    private var binaryenApplied: (BinaryenExec.() -> Unit)? = null
+    override fun applyBinaryen(body: BinaryenExec.() -> Unit) {
+        compilations.all {
+            it.binaries.all { binary ->
+                if (binary is ExecutableWasm) {
+                    val binaryenExec = createBinaryen(binary, body)
 
-    override fun whenBinaryenApplied(body: (BinaryenExec.() -> Unit) -> Unit) {
+                    if (
-        val binaryenApplied = binaryenApplied
+                        wasmTargetType == KotlinWasmTargetType.WASI && nodejsLazyDelegate.isInitialized() ||
-        if (binaryenApplied != null) {
+                        project.locateTask<IncrementalSyncTask>(binary.linkSyncTaskName) != null
-            body(binaryenApplied)
+                    ) {
-        } else {
+                        configureBinaryen(binary, binaryenExec)
-            applyBinaryenHandlers += body
+                    } else {
+                        binaryenReplaceInput[binary] = {
+                            configureBinaryen(binary, binaryenExec)
+                        }
+                    }
+                }
+            }
         }
     }
 
-    override fun applyBinaryen(body: BinaryenExec.() -> Unit) {
+    private fun createBinaryen(binary: ExecutableWasm, binaryenDsl: BinaryenExec.() -> Unit): TaskProvider<BinaryenExec> {
-        binaryenApplied = body
+        val linkTask = binary.linkTask
-        applyBinaryenHandlers.forEach { handler ->
+
-            handler(body)
+        val compileWasmDestDir = linkTask.map {
+            it.destinationDirectory
+        }
+
+        val compiledWasmFile = linkTask.map { link ->
+            link.destinationDirectory.asFile.get().resolve(link.compilerOptions.moduleName.get() + ".wasm")
+        }
+
+        return BinaryenExec.create(binary.compilation, binary.optimizeTaskName) {
+            val compilation = binary.compilation
+            dependsOn(linkTask)
+            inputFileProperty.fileProvider(compiledWasmFile)
+
+            val outputDirectory: Provider<File> = binary.target.project.layout.buildDirectory
+                .dir(COMPILE_SYNC)
+                .map { it.dir(compilation.target.targetName) }
+                .map { it.dir(compilation.name) }
+                .map { it.dir(binary.name) }
+                .map { it.dir("optimized") }
+                .map { it.asFile }
+
+            val outputFile = outputDirectory.map { it.resolve(compiledWasmFile.get().name) }
+
+            outputFileProperty.fileProvider(
+                outputFile
+            )
+
+            doLast {
+                fs.copy {
+                    it.from(compileWasmDestDir)
+                    it.into(outputDirectory)
+                    it.eachFile {
+                        if (it.relativePath.getFile(outputDirectory.get()).exists()) {
+                            it.exclude()
+                        }
+                    }
+                }
+            }
+
+            binaryenDsl()
+        }
+    }
+
+    private fun configureBinaryen(binary: JsIrBinary, binaryenExec: TaskProvider<BinaryenExec>) {
+        if (wasmTargetType == KotlinWasmTargetType.WASI) {
+            if (binary.compilation.isMain() && binary.mode == KotlinJsBinaryMode.PRODUCTION) {
+                project.tasks.named(LifecycleBasePlugin.ASSEMBLE_TASK_NAME).dependsOn(binaryenExec)
+            }
+
+            whenNodejsConfigured {
+                testTask {
+                    val name = binary.linkTask.flatMap { it.outputFileProperty.map { it.name } }
+                    it.inputFileProperty.fileProvider(
+                        binaryenExec.flatMap { it.outputFileProperty.map { it.asFile.parentFile.resolve(name.get()) } }
+                    )
+                    it.dependsOn(binaryenExec)
+                }
+            }
+        } else {
+            binary.linkSyncTask.configure {
+                it.from.setFrom(binaryenExec.flatMap { it.outputFileProperty.map { it.asFile.parentFile } })
+                it.dependsOn(binaryenExec)
+            }
         }
-        browserConfiguredHandlers.clear()
     }
 
     //Browser
@@ -265,6 +342,11 @@ constructor(
             commonLazy
         } else {
             NodeJsRootPlugin.apply(project.rootProject)
+            compilations.all { compilation ->
+                compilation.binaries.all { binary ->
+                    binaryenReplaceInput[binary]?.invoke()
+                }
+            }
         }
         project.objects.newInstance(KotlinNodeJsIr::class.java, this).also {
             it.configureSubTarget()