Worklets. Add extension

Author: turansky

seskar/seskar-compiler-plugin/src/main/kotlin/seskar/compiler/worklets/backend/WorkletModule.kt

@@ -0,0 +1,16 @@
+package seskar.compiler.worklets.backend
+
+import org.jetbrains.kotlin.ir.declarations.IrProperty
+import org.jetbrains.kotlin.ir.types.classFqName
+import org.jetbrains.kotlin.ir.util.isTopLevel
+import org.jetbrains.kotlin.name.FqName
+
+internal const val WORKLET_MODULE_DELIMITER = "$$"
+
+private val WORKLET_MODULES = setOf(
+    FqName("web.audio.AudioWorkletModule"),
+    FqName("web.css.painting.PaintWorkletModule"),
+)
+
+internal fun IrProperty.isWorkletModule(): Boolean =
+    isTopLevel && getter?.returnType?.classFqName in WORKLET_MODULES

seskar/seskar-compiler-plugin/src/main/kotlin/seskar/compiler/worklets/backend/WorkletTransformer.kt

@@ -0,0 +1,35 @@
+package seskar.compiler.worklets.backend
+
+import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
+import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.declarations.IrProperty
+import org.jetbrains.kotlin.ir.declarations.name
+import org.jetbrains.kotlin.ir.util.file
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
+import seskar.compiler.common.backend.JsFileName
+import seskar.compiler.common.backend.JsName
+
+internal class WorkletTransformer(
+    private val context: IrPluginContext,
+) : IrElementTransformerVoid() {
+    override fun visitProperty(declaration: IrProperty): IrStatement {
+        if (!declaration.isWorkletModule())
+            return declaration
+
+        val fileName = declaration.file.name.removeSuffix(".kt")
+        declaration.file.annotations += JsFileName(context, "${fileName}__worklet__module")
+
+        val jsName = sequenceOf(
+            fileName,
+            declaration.name.identifier,
+        ).joinToString(
+            separator = WORKLET_MODULE_DELIMITER,
+            prefix = WORKLET_MODULE_DELIMITER,
+            postfix = WORKLET_MODULE_DELIMITER,
+        )
+
+        declaration.annotations += JsName(context, jsName)
+
+        return super.visitProperty(declaration)
+    }
+}

seskar/seskar-compiler-plugin/src/main/kotlin/seskar/compiler/worklets/extensions/WorkletExtension.kt

@@ -0,0 +1,20 @@
+package seskar.compiler.worklets.extensions
+
+import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
+import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
+import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
+import org.jetbrains.kotlin.platform.isJs
+import seskar.compiler.worklets.backend.WorkletTransformer
+
+internal class WorkletExtension : IrGenerationExtension {
+    override fun generate(
+        moduleFragment: IrModuleFragment,
+        pluginContext: IrPluginContext,
+    ) {
+        pluginContext.platform
+            ?.takeIf { it.isJs() }
+            ?: return
+
+        moduleFragment.transformChildren(WorkletTransformer(pluginContext), null)
+    }
+}

seskar/seskar-compiler-plugin/src/main/kotlin/seskar/compiler/worklets/extensions/WorkletRegistrar.kt

@@ -0,0 +1,13 @@
+package seskar.compiler.worklets.extensions
+
+import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
+import org.jetbrains.kotlin.config.CompilerConfiguration
+import seskar.compiler.common.extensions.JsCompilerPluginRegistrar
+
+class WorkletRegistrar : JsCompilerPluginRegistrar() {
+    override fun ExtensionStorage.registerExtensions(
+        configuration: CompilerConfiguration,
+    ) {
+        IrGenerationExtension.registerExtension(WorkletExtension())
+    }
+}