[Gradle, Wasm] Binaryen to wasm targets

^KT-62024 fixed ^KT-62025 fixed
2023-09-18 14:42:20 +00:00
parent bb381e7aef
commit 6a514c3209
8 changed files with 201 additions and 58 deletions
@@ -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)
+                }
+            }
+        }
+    }
 }
@@ -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()
@@ -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 {
@@ -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)
+        }
+    }
 }
@@ -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,
@@ -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
    }

@@ -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)
-                binary.linkTask.flatMap { linkTask ->
-                    linkTask.outputFileProperty
+                if (project.locateTask<BinaryenExec>((binary as ExecutableWasm).optimizeTaskName) != null) {
+                    testJs.dependsOn(binary.optimizeTask)
+                    binary.optimizeTask.flatMap { optimizeTask ->
+                        optimizeTask.outputFileProperty.asFile
+                    }
+                } else {
+                    testJs.dependsOn(binary.linkTask)
+                    binary.linkTask.flatMap { linkTask ->
+                        linkTask.outputFileProperty
+                    }
                }
            }

@@ -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.registerTask
+import org.jetbrains.kotlin.gradle.tasks.*
+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) {
-        val binaryenApplied = binaryenApplied
-        if (binaryenApplied != null) {
-            body(binaryenApplied)
-        } else {
-            applyBinaryenHandlers += body
+                    if (
+                        wasmTargetType == KotlinWasmTargetType.WASI && nodejsLazyDelegate.isInitialized() ||
+                        project.locateTask<IncrementalSyncTask>(binary.linkSyncTaskName) != null
+                    ) {
+                        configureBinaryen(binary, binaryenExec)
+                    } else {
+                        binaryenReplaceInput[binary] = {
+                            configureBinaryen(binary, binaryenExec)
+                        }
+                    }
+                }
+            }
        }
    }

-    override fun applyBinaryen(body: BinaryenExec.() -> Unit) {
-        binaryenApplied = body
-        applyBinaryenHandlers.forEach { handler ->
-            handler(body)
+    private fun createBinaryen(binary: ExecutableWasm, binaryenDsl: BinaryenExec.() -> Unit): TaskProvider<BinaryenExec> {
+        val linkTask = binary.linkTask
+
+        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()