[K/JS] Rework ES modules part with squashed JsImport and right renami…

…ng strategy inside import/export statements

compiler/cli/cli-js/src/org/jetbrains/kotlin/cli/js/K2JSCompiler.java

@@ -90,6 +90,7 @@ protected void addPlatformOptions(@NotNull List<String> $self, @NotNull K2JSComp
         moduleKindMap.put(K2JsArgumentConstants.MODULE_COMMONJS, ModuleKind.COMMON_JS);
         moduleKindMap.put(K2JsArgumentConstants.MODULE_AMD, ModuleKind.AMD);
         moduleKindMap.put(K2JsArgumentConstants.MODULE_UMD, ModuleKind.UMD);
+        moduleKindMap.put(K2JsArgumentConstants.MODULE_ES, ModuleKind.ES);
 
         sourceMapContentEmbeddingMap.put(K2JsArgumentConstants.SOURCE_MAP_SOURCE_CONTENT_ALWAYS, SourceMapSourceEmbedding.ALWAYS);
         sourceMapContentEmbeddingMap.put(K2JsArgumentConstants.SOURCE_MAP_SOURCE_CONTENT_NEVER, SourceMapSourceEmbedding.NEVER);

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/Mappings.kt

@@ -44,6 +44,8 @@ interface DelegateFactory {
 }
 
 object DefaultDelegateFactory : DelegateFactory {
+    fun <K : IrDeclaration, V> newDeclarationToValueMapping(): Mapping.Delegate<K, V> = newMappingImpl()
+
     override fun <K : IrDeclaration, V : IrDeclaration> newDeclarationToDeclarationMapping(): Mapping.Delegate<K, V> = newMappingImpl()
 
     override fun <K : IrDeclaration, V : Collection<IrDeclaration>> newDeclarationToDeclarationCollectionMapping(): Mapping.Delegate<K, V> = newMappingImpl()

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

@@ -11,8 +11,8 @@ import org.jetbrains.kotlin.ir.backend.js.utils.MutableReference
 import org.jetbrains.kotlin.ir.declarations.*
 
 class JsMapping : DefaultMapping() {
-    val esClassWhichNeedBoxParameters = mutableSetOf<IrClass>()
-    val esClassToPossibilityForOptimization = mutableMapOf<IrClass, MutableReference<Boolean>>()
+    val esClassWhichNeedBoxParameters = DefaultDelegateFactory.newDeclarationToValueMapping<IrClass, Boolean>()
+    val esClassToPossibilityForOptimization = DefaultDelegateFactory.newDeclarationToValueMapping<IrClass, MutableReference<Boolean>>()
 
     val outerThisFieldSymbols = DefaultDelegateFactory.newDeclarationToDeclarationMapping<IrClass, IrField>()
     val innerClassConstructors = DefaultDelegateFactory.newDeclarationToDeclarationMapping<IrConstructor, IrConstructor>()

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

@@ -83,7 +83,7 @@ class ExportModelToJsStatements(
                     namespace != null ->
                         listOf(jsAssignment(jsElementAccess(declaration.name, namespace), JsNameRef(name)).makeStmt())
 
-                    esModules -> listOf(JsExport(name, alias = JsName(declaration.name, false)))
+                    esModules -> listOf(JsExport(name.makeRef(), alias = JsName(declaration.name, false)))
                     else -> emptyList()
                 }
             }
@@ -96,7 +96,7 @@ class ExportModelToJsStatements(
                 when {
                     namespace == null -> {
                         val property = declaration.generateTopLevelGetters()
-                        listOf(JsVars(property), JsExport(property.name, JsName(declaration.name, false)))
+                        listOf(JsVars(property), JsExport(property.name.makeRef(), JsName(declaration.name, false)))
                     }
                     es6mode && declaration.isMember -> {
                         val jsClass = parentClass?.getCorrespondingJsClass() ?: error("Expect to have parentClass at this point")
@@ -168,7 +168,7 @@ class ExportModelToJsStatements(
                 }
                 val klassExport = when {
                     namespace != null -> jsAssignment(newNameSpace, JsNameRef(name)).makeStmt()
-                    esModules -> JsExport(name, alias = JsName(declaration.name, false))
+                    esModules -> JsExport(name.makeRef(), alias = JsName(declaration.name, false))
                     else -> null
                 }
 

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

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.js.config.JSConfigurationKeys
 import org.jetbrains.kotlin.library.impl.buffer
 import org.jetbrains.kotlin.protobuf.CodedInputStream
 import org.jetbrains.kotlin.protobuf.CodedOutputStream
+import org.jetbrains.kotlin.serialization.js.ModuleKind
 import java.security.MessageDigest
 
 internal fun Hash128Bits.toProtoStream(out: CodedOutputStream) {
@@ -158,6 +159,13 @@ internal fun CrossModuleReferences.crossModuleReferencesHashForIC() = HashCalcul
         val import = imports[tag]!!
         update(tag)
         update(import.exportedAs)
-        update(import.moduleExporter.toString())
+
+        if (moduleKind == ModuleKind.ES) {
+            update(import.moduleExporter.internalName.toString())
+            update(import.moduleExporter.externalName)
+            update(import.moduleExporter.relativeRequirePath ?: "")
+        } else {
+            update(import.moduleExporter.internalName.toString())
+        }
     }
 }.finalize()

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

@@ -25,7 +25,7 @@ import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.util.collectionUtils.filterIsInstanceAnd
 
 class ES6ConstructorBoxParameterOptimizationLowering(private val context: JsIrBackendContext) : BodyLoweringPass {
-    private val esClassWhichNeedBoxParameters = context.mapping.esClassWhichNeedBoxParameters
+    private val IrClass.needsOfBoxParameter by context.mapping.esClassWhichNeedBoxParameters
 
     override fun lower(irBody: IrBody, container: IrDeclaration) {
         if (!context.es6mode) return
@@ -85,12 +85,12 @@ class ES6ConstructorBoxParameterOptimizationLowering(private val context: JsIrBa
     }
 
     private fun IrClass.requiredToHaveBoxParameter(): Boolean {
-        return esClassWhichNeedBoxParameters.contains(this)
+        return needsOfBoxParameter == true
     }
 }
 
 class ES6CollectConstructorsWhichNeedBoxParameters(private val context: JsIrBackendContext) : DeclarationTransformer {
-    private val esClassWhichNeedBoxParameters = context.mapping.esClassWhichNeedBoxParameters
+    private var IrClass.needsOfBoxParameter by context.mapping.esClassWhichNeedBoxParameters
 
     override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
         if (!context.es6mode || declaration !is IrClass) return null
@@ -134,7 +134,7 @@ class ES6CollectConstructorsWhichNeedBoxParameters(private val context: JsIrBack
 
     private fun IrClass.addToClassListWhichNeedBoxParameter() {
         if (isExternal) return
-        esClassWhichNeedBoxParameters.add(this)
+        needsOfBoxParameter = true
         superClass?.addToClassListWhichNeedBoxParameter()
     }
 }

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

@@ -177,16 +177,16 @@ class ES6PrimaryConstructorUsageOptimizationLowering(private val context: JsIrBa
  * Otherwise, we can generate a simple ES-class constructor in each class of the hierarchy
  */
 class ES6CollectPrimaryConstructorsWhichCouldBeOptimizedLowering(private val context: JsIrBackendContext) : DeclarationTransformer {
-    private val esClassWhichNeedBoxParameters = context.mapping.esClassWhichNeedBoxParameters
-    private val esClassToPossibilityForOptimization = context.mapping.esClassToPossibilityForOptimization
+    private val IrClass.needsOfBoxParameter by context.mapping.esClassWhichNeedBoxParameters
+    private var IrClass.possibilityToOptimizeForEsClass by context.mapping.esClassToPossibilityForOptimization
 
     override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
         if (
             context.es6mode &&
             declaration is IrClass &&
             !declaration.isExternal &&
             !context.inlineClassesUtils.isClassInlineLike(declaration) &&
-            !esClassToPossibilityForOptimization.contains(declaration)
+            declaration.possibilityToOptimizeForEsClass == null
         ) {
             declaration.checkIfCanBeOptimized()
         }
@@ -199,7 +199,7 @@ class ES6CollectPrimaryConstructorsWhichCouldBeOptimizedLowering(private val con
         var nearestOptimizationDecision: MutableReference<Boolean>? = null
 
         while (currentClass != null && !currentClass.isExternal) {
-            val currentClassOptimizationDecision = esClassToPossibilityForOptimization[currentClass]
+            val currentClassOptimizationDecision = currentClass.possibilityToOptimizeForEsClass
 
             if (currentClassOptimizationDecision != null) {
                 nearestOptimizationDecision = currentClassOptimizationDecision
@@ -214,8 +214,8 @@ class ES6CollectPrimaryConstructorsWhichCouldBeOptimizedLowering(private val con
         }
 
         currentClass = this
-        while (currentClass != null && !currentClass.isExternal && !esClassToPossibilityForOptimization.contains(currentClass)) {
-            esClassToPossibilityForOptimization[currentClass] = nearestOptimizationDecision
+        while (currentClass != null && !currentClass.isExternal && currentClass.possibilityToOptimizeForEsClass == null) {
+            currentClass.possibilityToOptimizeForEsClass = nearestOptimizationDecision
 
             if (nearestOptimizationDecision.value && !currentClass.canBeOptimized()) {
                 nearestOptimizationDecision.value = false
@@ -249,7 +249,7 @@ class ES6CollectPrimaryConstructorsWhichCouldBeOptimizedLowering(private val con
     }
 
     private fun IrClass.isSubclassOfExternalClassWithRequiredBoxParameter(): Boolean {
-        return superClass?.isExternal == true && esClassWhichNeedBoxParameters.contains(this)
+        return superClass?.isExternal == true && needsOfBoxParameter == true
     }
 }
 

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

@@ -7,11 +7,19 @@ package org.jetbrains.kotlin.ir.backend.js.transformers.irToJs
 
 import org.jetbrains.kotlin.ir.backend.js.export.TypeScriptFragment
 import org.jetbrains.kotlin.ir.backend.js.export.toTypeScript
+import org.jetbrains.kotlin.js.backend.ast.ESM_EXTENSION
 import org.jetbrains.kotlin.js.backend.ast.JsProgram
+import org.jetbrains.kotlin.js.backend.ast.REGULAR_EXTENSION
 import org.jetbrains.kotlin.serialization.js.ModuleKind
 import java.io.File
 import java.nio.file.Files
 
+val ModuleKind.extension: String
+    get() = when (this) {
+        ModuleKind.ES -> ESM_EXTENSION
+        else -> REGULAR_EXTENSION
+    }
+
 abstract class CompilationOutputs {
     var dependencies: Collection<Pair<String, CompilationOutputs>> = emptyList()
 
@@ -35,10 +43,10 @@ abstract class CompilationOutputs {
         }
 
         dependencies.forEach { (name, content) ->
-            outputDir.resolve("$name.js").writeAsJsFile(content)
+            outputDir.resolve("$name${moduleKind.extension}").writeAsJsFile(content)
         }
 
-        val outputJsFile = outputDir.resolve("$outputName.js")
+        val outputJsFile = outputDir.resolve("$outputName${moduleKind.extension}")
         outputJsFile.writeAsJsFile(this)
 
         if (genDTS) {

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

@@ -94,12 +94,14 @@ class CrossModuleDependenciesResolver(
     private val headers: List<JsIrModuleHeader>
 ) {
     fun resolveCrossModuleDependencies(relativeRequirePath: Boolean): Map<JsIrModuleHeader, CrossModuleReferences> {
-        val headerToBuilder = headers.associateWith { JsIrModuleCrossModuleReferecenceBuilder(moduleKind, it, relativeRequirePath) }
-        val definitionModule = mutableMapOf<String, JsIrModuleCrossModuleReferecenceBuilder>()
+        val headerToBuilder = headers.associateWith { JsIrModuleCrossModuleReferenceBuilder(moduleKind, it, relativeRequirePath) }
+        val definitionModule = mutableMapOf<String, JsIrModuleCrossModuleReferenceBuilder>()
 
-        val mainModuleHeader = headers.last()
-        val otherModuleHeaders = headers.dropLast(1)
-        headerToBuilder[mainModuleHeader]!!.transitiveJsExportFrom = otherModuleHeaders
+        if (moduleKind != ModuleKind.ES) {
+            val mainModuleHeader = headers.last()
+            val otherModuleHeaders = headers.dropLast(1)
+            headerToBuilder[mainModuleHeader]!!.transitiveJsExportFrom = otherModuleHeaders
+        }
 
         for (header in headers) {
             val builder = headerToBuilder[header]!!
@@ -130,9 +132,9 @@ class CrossModuleDependenciesResolver(
     }
 }
 
-private class CrossModuleRef(val module: JsIrModuleCrossModuleReferecenceBuilder, val tag: String)
+private class CrossModuleRef(val module: JsIrModuleCrossModuleReferenceBuilder, val tag: String)
 
-private class JsIrModuleCrossModuleReferecenceBuilder(
+private class JsIrModuleCrossModuleReferenceBuilder(
     val moduleKind: ModuleKind,
     val header: JsIrModuleHeader,
     val relativeRequirePath: Boolean
@@ -150,20 +152,15 @@ private class JsIrModuleCrossModuleReferecenceBuilder(
 
     fun buildCrossModuleRefs(): CrossModuleReferences {
         buildExportNames()
+        val isImportOptional = moduleKind == ModuleKind.ES
         val importedModules = mutableMapOf<JsIrModuleHeader, JsImportedModule>()
 
-        fun import(moduleHeader: JsIrModuleHeader): JsName {
-            return importedModules.getOrPut(moduleHeader) {
-                val jsModuleName = JsName(moduleHeader.moduleName, false)
-                val relativeRequirePath = relativeRequirePath(moduleHeader)
-
-                JsImportedModule(
-                    moduleHeader.externalModuleName,
-                    jsModuleName,
-                    null,
-                    relativeRequirePath
-                )
-            }.internalName
+        fun import(moduleHeader: JsIrModuleHeader): JsImportedModule {
+            return if (isImportOptional) {
+                moduleHeader.toJsImportedModule()
+            } else {
+                importedModules.getOrPut(moduleHeader) { moduleHeader.toJsImportedModule() }
+            }
         }
 
         val resultImports = imports.associate { crossModuleRef ->
@@ -173,13 +170,13 @@ private class JsIrModuleCrossModuleReferecenceBuilder(
                 "Cross module dependency resolution failed due to signature '$tag' redefinition"
             }
             val exportedAs = crossModuleRef.module.exportNames[tag]!!
-            val moduleName = import(crossModuleRef.module.header)
+            val importedModule = import(crossModuleRef.module.header)
 
-            tag to CrossModuleImport(exportedAs, moduleName)
+            tag to CrossModuleImport(exportedAs, importedModule)
         }
 
         val transitiveExport = transitiveJsExportFrom.mapNotNull {
-            if (!it.hasJsExports) null else CrossModuleTransitiveExport(import(it), it.externalModuleName)
+            if (!it.hasJsExports) null else CrossModuleTransitiveExport(import(it).internalName, it.externalModuleName)
         }
         return CrossModuleReferences(
             moduleKind,
@@ -190,12 +187,22 @@ private class JsIrModuleCrossModuleReferecenceBuilder(
         )
     }
 
+    private fun JsIrModuleHeader.toJsImportedModule(): JsImportedModule {
+        val jsModuleName = JsName(moduleName, false)
+        val relativeRequirePath = relativeRequirePath(this)
+
+        return JsImportedModule(
+            externalModuleName,
+            jsModuleName,
+            null,
+            relativeRequirePath
+        )
+    }
+
     private fun relativeRequirePath(moduleHeader: JsIrModuleHeader): String? {
         if (!this.relativeRequirePath) return null
 
-        val parentMain = File(header.externalModuleName).parentFile
-
-        if (parentMain == null) return "./${moduleHeader.externalModuleName}"
+        val parentMain = File(header.externalModuleName).parentFile ?: return "./${moduleHeader.externalModuleName}"
 
         val relativePath = File(moduleHeader.externalModuleName)
             .toRelativeString(parentMain)
@@ -206,10 +213,12 @@ private class JsIrModuleCrossModuleReferecenceBuilder(
     }
 }
 
-class CrossModuleImport(val exportedAs: String, val moduleExporter: JsName)
+class CrossModuleImport(val exportedAs: String, val moduleExporter: JsImportedModule)
 
 class CrossModuleTransitiveExport(val internalName: JsName, val externalName: String)
 
+fun CrossModuleTransitiveExport.getRequireEsmName() = "$externalName$ESM_EXTENSION"
+
 class CrossModuleReferences(
     val moduleKind: ModuleKind,
     val importedModules: List<JsImportedModule>, // additional Kotlin imported modules
@@ -218,28 +227,45 @@ class CrossModuleReferences(
     val imports: Map<String, CrossModuleImport>, // tag -> import statement
 ) {
     // built from imports
-    var jsImports = emptyMap<String, JsVars.JsVar>() // tag -> import statement
+    var jsImports = emptyMap<String, JsStatement>() // tag -> import statement
         private set
 
     fun initJsImportsForModule(module: JsIrModule) {
         val tagToName = module.fragments.flatMap { it.nameBindings.entries }.associate { it.key to it.value }
         jsImports = imports.entries.associate {
             val importedAs = tagToName[it.key] ?: error("Internal error: cannot find imported name for signature ${it.key}")
-            val exportRef = JsNameRef(
-                it.value.exportedAs,
-                it.value.moduleExporter.let {
-                    if (moduleKind == ModuleKind.ES) {
-                        it.makeRef()
-                    } else {
-                        ReservedJsNames.makeCrossModuleNameRef(it)
-                    }
-                }
-            )
-            it.key to JsVars.JsVar(importedAs, exportRef)
+            it.key to it.value.generateCrossModuleImportStatement(importedAs)
         }
     }
 
+    private fun CrossModuleImport.generateCrossModuleImportStatement(importedAs: JsName): JsStatement {
+        return when (moduleKind) {
+            ModuleKind.ES -> generateJsImportStatement(importedAs)
+            else -> generateImportVariableDeclaration(importedAs)
+        }
+    }
+
+    private fun CrossModuleImport.generateImportVariableDeclaration(importedAs: JsName): JsStatement {
+        val exportRef = JsNameRef(exportedAs, ReservedJsNames.makeCrossModuleNameRef(moduleExporter.internalName))
+        return JsVars(JsVars.JsVar(importedAs, exportRef))
+    }
+
+    private fun CrossModuleImport.generateJsImportStatement(importedAs: JsName): JsStatement {
+        return JsImport(
+            moduleExporter.getRequireName(true),
+            JsImport.Element(JsName(exportedAs, false), importedAs.makeRef())
+        )
+    }
+
     companion object {
         fun Empty(moduleKind: ModuleKind) = CrossModuleReferences(moduleKind, listOf(), emptyList(), emptyMap(), emptyMap())
     }
 }
+
+fun JsStatement.renameImportedSymbolInternalName(newName: JsName): JsStatement {
+    return when (this) {
+        is JsImport -> JsImport(module, JsImport.Element((target as JsImport.Target.Elements).elements.single().name, newName.makeRef()))
+        is JsVars -> JsVars(JsVars.JsVar(newName, vars.single().initExpression))
+        else -> error("Unexpected cross-module import statement ${this::class.qualifiedName}")
+    }
+}