Kotlin/JS codegen compiler extension

compiler/cli/build.gradle.kts

@@ -12,6 +12,7 @@ dependencies {
     api(project(":compiler:backend-common"))
     api(project(":compiler:backend"))
     api(project(":compiler:backend.jvm"))
+    api(project(":compiler:backend.js"))
     implementation(project(":compiler:backend.jvm.entrypoint"))
     api(project(":compiler:ir.backend.common"))
     api(project(":compiler:light-classes"))

compiler/cli/cli-common/resources/META-INF/extensions/compiler.xml

@@ -32,6 +32,9 @@
         <extensionPoint qualifiedName="org.jetbrains.kotlin.irGenerationExtension"
                         interface="org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension"
                         area="IDEA_PROJECT" dynamic="true"/>
+        <extensionPoint qualifiedName="org.jetbrains.kotlin.jsGenerationExtension"
+                        interface="org.jetbrains.kotlin.ir.backend.js.extensions.JsGenerationExtension"
+                        area="IDEA_PROJECT" dynamic="true"/>
         <extensionPoint qualifiedName="org.jetbrains.kotlin.simpleNameReferenceExtension"
                         interface="org.jetbrains.kotlin.plugin.references.SimpleNameReferenceExtension"
                         area="IDEA_PROJECT" dynamic="true"/>

compiler/cli/cli-js/src/org/jetbrains/kotlin/cli/js/K2JsIrCompiler.kt

@@ -42,6 +42,7 @@ import org.jetbrains.kotlin.incremental.js.IncrementalResultsConsumer
 import org.jetbrains.kotlin.backend.wasm.dce.eliminateDeadDeclarations
 import org.jetbrains.kotlin.ir.backend.js.*
 import org.jetbrains.kotlin.ir.backend.js.codegen.JsGenerationGranularity
+import org.jetbrains.kotlin.ir.backend.js.extensions.IrToJsTransformationExtension
 import org.jetbrains.kotlin.ir.backend.js.ic.*
 import org.jetbrains.kotlin.ir.backend.js.transformers.irToJs.IrModuleToJsTransformer
 import org.jetbrains.kotlin.ir.backend.js.transformers.irToJs.IrModuleToJsTransformerTmp
@@ -55,7 +56,6 @@ import org.jetbrains.kotlin.library.KLIB_FILE_EXTENSION
 import org.jetbrains.kotlin.metadata.deserialization.BinaryVersion
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.resolve.CompilerEnvironment
 import org.jetbrains.kotlin.serialization.js.ModuleKind
 import org.jetbrains.kotlin.utils.KotlinPaths
 import org.jetbrains.kotlin.utils.PathUtil
@@ -402,7 +402,8 @@ class K2JsIrCompiler : CLICompiler<K2JSCompilerArguments>() {
                         ir.context,
                         mainCallArguments,
                         relativeRequirePath = true,
-                        moduleToName = ir.moduleFragmentToUniqueName
+                        moduleToName = ir.moduleFragmentToUniqueName,
+                        irToJsTransformationExtensions = IrToJsTransformationExtension.getInstances(module.project)
                     )
 
                     transformer.generateModule(

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/KotlinCoreEnvironment.kt

@@ -74,6 +74,7 @@ import org.jetbrains.kotlin.extensions.internal.CandidateInterceptor
 import org.jetbrains.kotlin.extensions.internal.TypeResolutionInterceptor
 import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrarAdapter
 import org.jetbrains.kotlin.idea.KotlinFileType
+import org.jetbrains.kotlin.ir.backend.js.extensions.IrToJsTransformationExtension
 import org.jetbrains.kotlin.js.translate.extensions.JsSyntheticTranslateExtension
 import org.jetbrains.kotlin.load.kotlin.KotlinBinaryClassCache
 import org.jetbrains.kotlin.load.kotlin.MetadataFinderFactory
@@ -633,6 +634,7 @@ class KotlinCoreEnvironment private constructor(
             ProcessSourcesBeforeCompilingExtension.registerExtensionPoint(project)
             ExtraImportsProviderExtension.registerExtensionPoint(project)
             IrGenerationExtension.registerExtensionPoint(project)
+            IrToJsTransformationExtension.registerExtensionPoint(project)
             ScriptEvaluationExtension.registerExtensionPoint(project)
             ShellExtension.registerExtensionPoint(project)
             TypeResolutionInterceptor.registerExtensionPoint(project)

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/compiler.kt

@@ -165,7 +165,12 @@ fun generateJsCode(
     moveBodilessDeclarationsToSeparatePlace(context, moduleFragment)
     jsPhases.invokeToplevel(PhaseConfig(jsPhases), context, listOf(moduleFragment))
 
-    val transformer = IrModuleToJsTransformer(context, null, true, nameTables)
+    val transformer = IrModuleToJsTransformer(
+        context,
+        null,
+        true,
+        nameTables,
+    )
     return transformer.generateModule(listOf(moduleFragment)).outputs[TranslationMode.FULL]!!.jsCode
 }
 

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/dce/Dce.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.ir.backend.js.dce
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
 import org.jetbrains.kotlin.ir.backend.js.export.isExported
+import org.jetbrains.kotlin.ir.backend.js.extensions.IrToJsTransformationExtension
 import org.jetbrains.kotlin.ir.backend.js.utils.*
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.IrBody
@@ -22,8 +23,9 @@ fun eliminateDeadDeclarations(
     modules: Iterable<IrModuleFragment>,
     context: JsIrBackendContext,
     removeUnusedAssociatedObjects: Boolean = true,
+    irToJsTransformationExtensions: List<IrToJsTransformationExtension> = emptyList(),
 ) {
-    val allRoots = buildRoots(modules, context)
+    val allRoots = buildRoots(modules, context, irToJsTransformationExtensions)
 
     val printReachabilityInfo =
         context.configuration.getBoolean(JSConfigurationKeys.PRINT_REACHABILITY_INFO) ||
@@ -85,7 +87,11 @@ private fun IrDeclaration.addRootsTo(
     }
 }
 
-private fun buildRoots(modules: Iterable<IrModuleFragment>, context: JsIrBackendContext): List<IrDeclaration> = buildList {
+private fun buildRoots(
+    modules: Iterable<IrModuleFragment>,
+    context: JsIrBackendContext,
+    irToJsTransformationExtensions: List<IrToJsTransformationExtension>
+): List<IrDeclaration> = buildList {
     val declarationsCollector = object : IrElementVisitorVoid {
         override fun visitElement(element: IrElement): Unit = element.acceptChildrenVoid(this)
         override fun visitBody(body: IrBody): Unit = Unit // Skip
@@ -103,6 +109,14 @@ private fun buildRoots(modules: Iterable<IrModuleFragment>, context: JsIrBackend
         }
     }
 
+    if (irToJsTransformationExtensions.isNotEmpty()) {
+        for (extension in irToJsTransformationExtensions) {
+            for (module in modules) {
+                addAll(extension.getAdditionalDceRoots(module))
+            }
+        }
+    }
+
     val dceRuntimeDiagnostic = context.dceRuntimeDiagnostic
     if (dceRuntimeDiagnostic != null) {
         dceRuntimeDiagnostic.unreachableDeclarationMethod(context).owner.acceptVoid(declarationsCollector)

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/extensions/IrToJsExtensionKey.kt

@@ -0,0 +1,8 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.ir.backend.js.extensions
+
+abstract class IrToJsExtensionKey

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/extensions/IrToJsTransformationExtension.kt

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.ir.backend.js.extensions
+
+import org.jetbrains.kotlin.extensions.ProjectExtensionDescriptor
+import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
+import org.jetbrains.kotlin.ir.backend.js.transformers.irToJs.JsIrModule
+import org.jetbrains.kotlin.ir.backend.js.transformers.irToJs.JsModuleOrigin
+import org.jetbrains.kotlin.ir.declarations.IrDeclaration
+import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
+import org.jetbrains.kotlin.js.backend.ast.JsImportedModule
+import org.jetbrains.kotlin.js.backend.ast.JsProgram
+import org.jetbrains.kotlin.js.backend.ast.JsStatement
+import org.jetbrains.kotlin.serialization.js.ModuleKind
+
+interface IrToJsTransformationExtension {
+    companion object : ProjectExtensionDescriptor<IrToJsTransformationExtension>(
+        "org.jetbrains.kotlin.jsGenerationExtension", IrToJsTransformationExtension::class.java
+    )
+
+    fun generateAdditionalJsIrModules(
+        module: IrModuleFragment,
+        context: JsIrBackendContext,
+        minimizedMemberNames: Boolean
+    ): List<JsIrModule> = emptyList()
+
+    fun getAdditionalDceRoots(module: IrModuleFragment): List<IrDeclaration> = emptyList()
+
+    fun transformMainFunction(mainInvocation: JsStatement): JsStatement = mainInvocation
+
+    fun postprocessJsAst(
+        jsProgram: JsProgram,
+        moduleKind: ModuleKind,
+        moduleOrigin: JsModuleOrigin,
+        importedJsModules: List<JsImportedModule>
+    ) {}
+
+    val extensionKey: IrToJsExtensionKey
+}

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/IrModuleToJsTransformerTmp.kt

@@ -7,9 +7,12 @@ package org.jetbrains.kotlin.ir.backend.js.transformers.irToJs
 
 import org.jetbrains.kotlin.config.CommonConfigurationKeys
 import org.jetbrains.kotlin.config.CompilerConfiguration
-import org.jetbrains.kotlin.ir.backend.js.*
+import org.jetbrains.kotlin.ir.backend.js.CompilationOutputs
+import org.jetbrains.kotlin.ir.backend.js.CompilerResult
+import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
 import org.jetbrains.kotlin.ir.backend.js.dce.eliminateDeadDeclarations
 import org.jetbrains.kotlin.ir.backend.js.export.*
+import org.jetbrains.kotlin.ir.backend.js.extensions.IrToJsTransformationExtension
 import org.jetbrains.kotlin.ir.backend.js.lower.StaticMembersLowering
 import org.jetbrains.kotlin.ir.backend.js.utils.*
 import org.jetbrains.kotlin.ir.backend.js.utils.serialization.JsIrAstSerializer
@@ -36,7 +39,7 @@ import java.util.*
 enum class TranslationMode(
     val dce: Boolean,
     val perModule: Boolean,
-    val minimizedMemberNames: Boolean,
+    val minimizedMemberNames: Boolean
 ) {
     FULL(dce = false, perModule = false, minimizedMemberNames = false),
     FULL_DCE(dce = true, perModule = false, minimizedMemberNames = false),
@@ -71,6 +74,7 @@ class IrModuleToJsTransformerTmp(
     private val relativeRequirePath: Boolean = false,
     private val moduleToName: Map<IrModuleFragment, String> = emptyMap(),
     private val removeUnusedAssociatedObjects: Boolean = true,
+    private val irToJsTransformationExtensions: List<IrToJsTransformationExtension> = emptyList()
 ) {
     private val generateRegionComments = backendContext.configuration.getBoolean(JSConfigurationKeys.GENERATE_REGION_COMMENTS)
 
@@ -100,6 +104,7 @@ class IrModuleToJsTransformerTmp(
             SourceMapsInfo.from(backendContext.configuration),
             relativeRequirePath,
             generateScriptModule,
+            irToJsTransformationExtensions
         )
 
         val result = EnumMap<TranslationMode, CompilationOutputs>(TranslationMode::class.java)
@@ -113,7 +118,7 @@ class IrModuleToJsTransformerTmp(
         }
 
         if (modes.any { it.dce }) {
-            eliminateDeadDeclarations(modules, backendContext, removeUnusedAssociatedObjects)
+            eliminateDeadDeclarations(modules, backendContext, removeUnusedAssociatedObjects, irToJsTransformationExtensions)
         }
 
         modes.filter { it.dce }.forEach {
@@ -164,23 +169,37 @@ class IrModuleToJsTransformerTmp(
         minimizedMemberNames: Boolean
     ): JsIrProgram {
         return JsIrProgram(
-            modules.map { m ->
-                JsIrModule(
-                    m.safeName,
-                    m.externalModuleName(),
-                    m.files.map {
-                        val exports = exportData[m]!![it]!!
-                        generateProgramFragment(it, exports, minimizedMemberNames)
-                    },
-                )
+            buildList {
+                modules.forEach { m ->
+                    for (extension in irToJsTransformationExtensions) {
+                        val generatedModules = extension.generateAdditionalJsIrModules(m, backendContext, minimizedMemberNames)
+                        generatedModules.forEach { it.origin = JsModuleOrigin.Extension(extension.extensionKey) }
+                        addAll(generatedModules)
+                    }
+
+                    add(
+                        JsIrModule(
+                            m.safeName,
+                            m.externalModuleName(),
+                            m.files.map {
+                                val exports = exportData[m]!![it]!!
+                                generateProgramFragment(it, exports, minimizedMemberNames)
+                            },
+                        )
+                    )
+                }
             }
         )
     }
 
     private val generateFilePaths = backendContext.configuration.getBoolean(JSConfigurationKeys.GENERATE_COMMENTS_WITH_FILE_PATH)
     private val pathPrefixMap = backendContext.configuration.getMap(JSConfigurationKeys.FILE_PATHS_PREFIX_MAP)
 
-    private fun generateProgramFragment(file: IrFile, exports: List<ExportedDeclaration>, minimizedMemberNames: Boolean): JsIrProgramFragment {
+    private fun generateProgramFragment(
+        file: IrFile,
+        exports: List<ExportedDeclaration>,
+        minimizedMemberNames: Boolean
+    ): JsIrProgramFragment {
         val nameGenerator = JsNameLinkingNamer(backendContext, minimizedMemberNames)
 
         val globalNameScope = NameTable<IrDeclaration>()
@@ -254,7 +273,9 @@ class IrModuleToJsTransformerTmp(
                 val jsName = staticContext.getNameForStaticFunction(it)
                 val generateArgv = it.valueParameters.firstOrNull()?.isStringArrayParameter() ?: false
                 val generateContinuation = it.isLoweredSuspendFunction(backendContext)
-                result.mainFunction = JsInvocation(jsName.makeRef(), generateMainArguments(generateArgv, generateContinuation, staticContext)).makeStmt()
+                val mainInvocation =
+                    JsInvocation(jsName.makeRef(), generateMainArguments(generateArgv, generateContinuation, staticContext)).makeStmt()
+                result.mainFunction = irToJsTransformationExtensions.fold(mainInvocation) { acc, ext -> ext.transformMainFunction(acc) }
             }
         }
 
@@ -332,50 +353,59 @@ private fun generateWrappedModuleBody(
     program: JsIrProgram,
     sourceMapsInfo: SourceMapsInfo?,
     relativeRequirePath: Boolean,
-    generateScriptModule: Boolean
+    generateScriptModule: Boolean,
+    irToJsTransformationExtensions: List<IrToJsTransformationExtension>
 ): CompilationOutputs {
-    if (multiModule) {
-
-        val moduleToRef = program.crossModuleDependencies(relativeRequirePath)
+    val moduleToRef = program.crossModuleDependencies(relativeRequirePath)
 
+    val (compiledMain, others) = if (multiModule) {
         val main = program.mainModule
-        val others = program.otherModules
-
-        val mainModule = generateSingleWrappedModuleBody(
+        val module = generateSingleWrappedModuleBody(
             mainModuleName,
             moduleKind,
             main.fragments,
             sourceMapsInfo,
             generateScriptModule,
             generateCallToMain = true,
             moduleToRef[main]!!,
+            irToJsTransformationExtensions = irToJsTransformationExtensions
         )
-
-        val dependencies = others.map { module ->
-            val moduleName = module.externalModuleName
-
-            moduleName to generateSingleWrappedModuleBody(
-                moduleName,
-                moduleKind,
-                module.fragments,
-                sourceMapsInfo,
-                generateScriptModule,
-                generateCallToMain = false,
-                moduleToRef[module]!!,
-            )
+        Pair(module, program.otherModules)
+    } else {
+        val (generatedModules, otherModules) = program.modules.partition { it.origin is JsModuleOrigin.Extension }
+        if (generatedModules.isNotEmpty()) {
+            error("Additional js modules generation supported only in multi module builds")
         }
 
-        return CompilationOutputs(mainModule.jsCode, mainModule.jsProgram, mainModule.sourceMap, dependencies)
-    } else {
-        return generateSingleWrappedModuleBody(
+        val module = generateSingleWrappedModuleBody(
             mainModuleName,
             moduleKind,
-            program.modules.flatMap { it.fragments },
+            otherModules.flatMap { it.fragments },
             sourceMapsInfo,
             generateScriptModule,
             generateCallToMain = true,
+            irToJsTransformationExtensions = irToJsTransformationExtensions
+        )
+        Pair(module, generatedModules)
+    }
+
+    val dependencies = others.map { module ->
+        val moduleName = module.externalModuleName
+
+        moduleName to generateSingleWrappedModuleBody(
+            moduleName,
+            moduleKind,
+            module.fragments,
+            sourceMapsInfo,
+            generateScriptModule,
+            generateCallToMain = false,
+            moduleToRef[module]!!,
+            module.origin,
+            irToJsTransformationExtensions = irToJsTransformationExtensions
         )
     }
+
+    return CompilationOutputs(compiledMain.jsCode, compiledMain.jsProgram, compiledMain.sourceMap, dependencies)
 }
 
 fun generateSingleWrappedModuleBody(
@@ -386,6 +416,8 @@ fun generateSingleWrappedModuleBody(
     generateScriptModule: Boolean,
     generateCallToMain: Boolean,
     crossModuleReferences: CrossModuleReferences = CrossModuleReferences.Empty,
+    moduleOrigin: JsModuleOrigin = JsModuleOrigin.Source,
+    irToJsTransformationExtensions: List<IrToJsTransformationExtension> = emptyList()
 ): CompilationOutputs {
     val program = Merger(
         moduleName,
@@ -395,6 +427,8 @@ fun generateSingleWrappedModuleBody(
         generateScriptModule,
         generateRegionComments = true,
         generateCallToMain,
+        moduleOrigin,
+        irToJsTransformationExtensions = irToJsTransformationExtensions
     ).merge()
 
     program.resolveTemporaryNames()
@@ -425,11 +459,10 @@ fun generateSingleWrappedModuleBody(
     }
 
     program.accept(JsToStringGenerationVisitor(jsCode, sourceMapBuilderConsumer))
-
     return CompilationOutputs(
         jsCode.toString(),
         program,
-        sourceMapBuilder?.build()
+        sourceMapBuilder?.build(),
     )
 }