[K/JS] Rework ES modules part with squashed JsImport and right renaming strategy inside import/export statements

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

@@ -90,6 +90,7 @@ public class K2JSCompiler extends CLICompiler<K2JSCompilerArguments> {
         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 esClassWhichNeedBoxParameters = DefaultDelegateFactory.newDeclarationToValueMapping<IrClass, Boolean>()
-    val esClassToPossibilityForOptimization = mutableMapOf<IrClass, MutableReference<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

@@ -21,6 +21,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) {
@@ -164,6 +165,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,8 +177,8 @@ 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 IrClass.needsOfBoxParameter by context.mapping.esClassWhichNeedBoxParameters
-    private val esClassToPossibilityForOptimization = context.mapping.esClassToPossibilityForOptimization
+    private var IrClass.possibilityToOptimizeForEsClass by context.mapping.esClassToPossibilityForOptimization
 
     override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
         if (
@@ -186,7 +186,7 @@ class ES6CollectPrimaryConstructorsWhichCouldBeOptimizedLowering(private val con
             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)) {
+        while (currentClass != null && !currentClass.isExternal && currentClass.possibilityToOptimizeForEsClass == null) {
-            esClassToPossibilityForOptimization[currentClass] = nearestOptimizationDecision
+            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 headerToBuilder = headers.associateWith { JsIrModuleCrossModuleReferenceBuilder(moduleKind, it, relativeRequirePath) }
-        val definitionModule = mutableMapOf<String, JsIrModuleCrossModuleReferecenceBuilder>()
+        val definitionModule = mutableMapOf<String, JsIrModuleCrossModuleReferenceBuilder>()
 
-        val mainModuleHeader = headers.last()
+        if (moduleKind != ModuleKind.ES) {
-        val otherModuleHeaders = headers.dropLast(1)
+            val mainModuleHeader = headers.last()
-        headerToBuilder[mainModuleHeader]!!.transitiveJsExportFrom = otherModuleHeaders
+            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 {
+        fun import(moduleHeader: JsIrModuleHeader): JsImportedModule {
-            return importedModules.getOrPut(moduleHeader) {
+            return if (isImportOptional) {
-                val jsModuleName = JsName(moduleHeader.moduleName, false)
+                moduleHeader.toJsImportedModule()
-                val relativeRequirePath = relativeRequirePath(moduleHeader)
+            } else {
-
+                importedModules.getOrPut(moduleHeader) { moduleHeader.toJsImportedModule() }
-                JsImportedModule(
+            }
-                    moduleHeader.externalModuleName,
-                    jsModuleName,
-                    null,
-                    relativeRequirePath
-                )
-            }.internalName
         }
 
         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
+        val parentMain = File(header.externalModuleName).parentFile ?: return "./${moduleHeader.externalModuleName}"
-
-        if (parentMain == null) 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.key to it.value.generateCrossModuleImportStatement(importedAs)
-                it.value.exportedAs,
-                it.value.moduleExporter.let {
-                    if (moduleKind == ModuleKind.ES) {
-                        it.makeRef()
-                    } else {
-                        ReservedJsNames.makeCrossModuleNameRef(it)
-                    }
-                }
-            )
-            it.key to JsVars.JsVar(importedAs, exportRef)
         }
     }
 
+    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}")
+    }
+}

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

@@ -58,10 +58,10 @@ class Merger(
             }
         }
 
-        for ((tag, jsVar) in crossModuleReferences.jsImports) {
+        for ((tag, crossModuleJsImport) in crossModuleReferences.jsImports) {
             val importName = nameMap[tag] ?: error("Missing name for declaration '$tag'")
 
-            importStatements.putIfAbsent(tag, JsVars(JsVars.JsVar(importName, jsVar.initExpression)))
+            importStatements.putIfAbsent(tag, crossModuleJsImport.renameImportedSymbolInternalName(importName))
         }
 
         if (crossModuleReferences.exports.isNotEmpty()) {
@@ -70,7 +70,7 @@ class Merger(
             if (isEsModules) {
                 val exportedElements = crossModuleReferences.exports.entries.map { (tag, hash) ->
                     val internalName = nameMap[tag] ?: error("Missing name for declaration '$tag'")
-                    JsExport.Element(internalName, JsName(hash, false))
+                    JsExport.Element(internalName.makeRef(), JsName(hash, false))
                 }
 
                 additionalExports += JsExport(JsExport.Subject.Elements(exportedElements))
@@ -143,7 +143,7 @@ class Merger(
             val exportedElements = currentModuleExportStatements.takeIf { it.isNotEmpty() }
                 ?.asSequence()
                 ?.flatMap { (it.subject as JsExport.Subject.Elements).elements }
-                ?.distinctBy { (it.alias ?: it.name).ident }
+                ?.distinctBy { it.alias?.ident ?: it.name.ident }
                 ?.map { if (it.name.ident == it.alias?.ident) JsExport.Element(it.name, null) else it }
                 ?.toList()
 
@@ -176,7 +176,7 @@ class Merger(
     private fun transitiveJsExport(): List<JsStatement> {
         return if (isEsModules) {
             crossModuleReferences.transitiveJsExportFrom.map {
-                JsExport(JsExport.Subject.All, it.externalName)
+                JsExport(JsExport.Subject.All, it.getRequireEsmName())
             }
         } else {
             val internalModuleName = ReservedJsNames.makeInternalModuleName()

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.ir.backend.js.transformers.irToJs
 import org.jetbrains.kotlin.js.backend.ast.*
 import org.jetbrains.kotlin.js.common.isValidES5Identifier
 import org.jetbrains.kotlin.serialization.js.ModuleKind
+import org.jetbrains.kotlin.utils.addToStdlib.partitionIsInstance
 
 object ModuleWrapperTranslation {
     object Namer {
@@ -105,7 +106,7 @@ object ModuleWrapperTranslation {
         val scope = program.scope
         val defineName = scope.declareName("define")
         val invocationArgs = listOf(
-            JsArrayLiteral(listOf(JsStringLiteral("exports")) + importedModules.map { JsStringLiteral(it.requireName) }),
+            JsArrayLiteral(listOf(JsStringLiteral("exports")) + importedModules.map { JsStringLiteral(it.getRequireName()) }),
             function
         )
 
@@ -122,23 +123,32 @@ object ModuleWrapperTranslation {
         val moduleName = scope.declareName("module")
         val requireName = scope.declareName("require")
 
-        val invocationArgs = importedModules.map { JsInvocation(requireName.makeRef(), JsStringLiteral(it.requireName)) }
+        val invocationArgs = importedModules.map { JsInvocation(requireName.makeRef(), JsStringLiteral(it.getRequireName())) }
         val invocation = JsInvocation(function, listOf(JsNameRef("exports", moduleName.makeRef())) + invocationArgs)
         return listOf(invocation.makeStmt())
     }
 
     private fun wrapEsModule(function: JsFunction, importedModules: List<JsImportedModule>): List<JsStatement> {
+        val (alreadyPresentedImportStatements, restStatements) = function.body.statements.partitionIsInstance<JsStatement, JsImport>()
+
         val importStatements = importedModules.zip(function.parameters.drop(1)).map {
             JsImport(
-                it.first.externalName,
+                it.first.getRequireName(true),
                 if (it.first.plainReference == null) {
-                    JsImport.Target.All(alias = it.second.name)
+                    JsImport.Target.All(alias = it.second.name.makeRef())
                 } else {
-                    JsImport.Target.Default(name = it.second.name)
+                    JsImport.Target.Default(name = it.second.name.makeRef())
                 }
             )
         }
-       return importStatements + function.body.statements.dropLast(1)
+
+        val alreadyPresentedImportStatementsWithoutDuplicates = alreadyPresentedImportStatements
+            .groupBy { it.module }
+            .map { (module, import) ->
+                JsImport(module, *import.flatMap { it.elements }.toTypedArray())
+            }
+
+       return  importStatements + alreadyPresentedImportStatementsWithoutDuplicates + restStatements.dropLast(1)
     }
 
 

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

@@ -116,6 +116,21 @@ private fun JsNode.computeScopes(): Scope {
             super.visitNameRef(nameRef)
         }
 
+        override fun visitImport(import: JsImport) {
+            when (val target = import.target) {
+                is JsImport.Target.All -> target.alias.name?.let { currentScope.declaredNames += it }
+                is JsImport.Target.Default -> target.name.name?.let { currentScope.declaredNames += it }
+                is JsImport.Target.Elements -> target.elements.forEach { element ->
+                    if (element.alias != null) {
+                        element.alias?.name?.let { currentScope.declaredNames += it }
+                    } else {
+                        currentScope.declaredNames += element.name
+                    }
+                }
+            }
+            super.visitImport(import)
+        }
+
         override fun visitBreak(x: JsBreak) {}
 
         override fun visitContinue(x: JsContinue) {}

js/js.ast/src/org/jetbrains/kotlin/js/backend/JsToStringGenerationVisitor.java

@@ -1342,7 +1342,7 @@ public class JsToStringGenerationVisitor extends JsVisitor {
             blockOpen();
             List<JsExport.Element> elements = ((JsExport.Subject.Elements) subject).getElements();
             for (JsExport.Element element : elements) {
-                nameDef(element.getName());
+                visitNameRef(element.getName());
                 JsName alias = element.getAlias();
                 if (alias != null) {
                     p.print(" as ");
@@ -1369,10 +1369,10 @@ public class JsToStringGenerationVisitor extends JsVisitor {
         p.print("import ");
 
         if (target instanceof JsImport.Target.Default) {
-            nameDef(((JsImport.Target.Default) target).getName());
+            visitNameRef(((JsImport.Target.Default) target).getName());
         } else if (target instanceof JsImport.Target.All) {
             p.print("* as ");
-            nameDef(((JsImport.Target.All) target).getAlias());
+            visitNameRef(((JsImport.Target.All) target).getAlias());
         } else if (target instanceof JsImport.Target.Elements) {
             List<JsImport.Element> elements = ((JsImport.Target.Elements) target).getElements();
 
@@ -1386,10 +1386,10 @@ public class JsToStringGenerationVisitor extends JsVisitor {
 
             for (JsImport.Element element : elements) {
                 nameDef(element.getName());
-                JsName alias = element.getAlias();
+                JsNameRef alias = element.getAlias();
                 if (alias != null) {
                     p.print(" as ");
-                    nameDef(alias);
+                    visitNameRef(alias);
                 }
 
                 if (isMultiline) {

js/js.ast/src/org/jetbrains/kotlin/js/backend/ast/JsExport.kt

@@ -11,7 +11,7 @@ class JsExport(
 ) : SourceInfoAwareJsNode(), JsStatement {
 
     constructor(element: Element) : this(Subject.Elements(listOf(element)))
-    constructor(name: JsName, alias: JsName? = null) : this(Element(name, alias))
+    constructor(name: JsNameRef, alias: JsName? = null) : this(Element(name, alias))
 
     sealed class Subject {
         class Elements(val elements: List<Element>) : Subject()
@@ -19,7 +19,7 @@ class JsExport(
     }
 
     class Element(
-        val name: JsName,
+        val name: JsNameRef,
         val alias: JsName?
     )
 
@@ -28,6 +28,11 @@ class JsExport(
     }
 
     override fun acceptChildren(visitor: JsVisitor) {
+        if (subject is Subject.Elements) {
+            subject.elements.forEach {
+                visitor.accept(it.name)
+            }
+        }
     }
 
     override fun deepCopy(): JsStatement =

js/js.ast/src/org/jetbrains/kotlin/js/backend/ast/JsImport.kt

@@ -9,34 +9,39 @@ class JsImport(
     val module: String,
     val target: Target,
 ) : SourceInfoAwareJsNode(), JsStatement {
-    constructor(module: String, elements: MutableList<Element> = mutableListOf()) : this(module, Target.Elements(elements))
+    constructor(module: String, vararg elements: Element) : this(module, Target.Elements(elements.toMutableList()))
 
     val elements: MutableList<Element>
         get() = (target as Target.Elements).elements
 
     sealed class Target {
         class Elements(val elements: MutableList<Element>) : Target()
-        class Default(val name: JsName) : Target() {
+        class Default(val name: JsNameRef) : Target() {
-            constructor(name: String) : this(JsName(name, false))
+            constructor(name: String) : this(JsNameRef(name))
         }
 
-        class All(val alias: JsName) : Target() {
+        class All(val alias: JsNameRef) : Target() {
-            constructor(alias: String) : this(JsName(alias, false))
+            constructor(alias: String) : this(JsNameRef(alias))
         }
     }
 
     class Element(
         val name: JsName,
-        val alias: JsName?
+        val alias: JsNameRef?
-    ) {
+    )
-        constructor(name: String, alias: String?) : this(JsName(name, false), alias?.let { JsName(it, false) })
-    }
 
     override fun accept(visitor: JsVisitor) {
         visitor.visitImport(this)
     }
 
     override fun acceptChildren(visitor: JsVisitor) {
+        when (target) {
+            is Target.All -> visitor.accept(target.alias)
+            is Target.Default -> visitor.accept(target.name)
+            is Target.Elements -> target.elements.forEach {
+                it.alias?.let(visitor::accept)
+            }
+        }
     }
 
     override fun deepCopy(): JsStatement =

js/js.ast/src/org/jetbrains/kotlin/js/backend/ast/JsImportedModule.kt

@@ -26,7 +26,14 @@ class JsImportedModule @JvmOverloads constructor(
     val key = JsImportedModuleKey(externalName, plainReference?.toString())
 }
 
-val JsImportedModule.requireName: String
+const val REGULAR_EXTENSION = ".js"
-    get() = relativeRequirePath?.let { "$it.js" } ?: externalName
+const val ESM_EXTENSION = ".mjs"
+
+fun JsImportedModule.getRequireName(isEsm: Boolean = false): String {
+    return relativeRequirePath?.let {
+        val extension = if (isEsm) ESM_EXTENSION else REGULAR_EXTENSION
+        "$it$extension"
+    } ?: externalName
+}
 
 data class JsImportedModuleKey(val baseName: String, val plainName: String?)

js/js.tests/test/org/jetbrains/kotlin/js/test/converters/JsIrBackendFacade.kt

@@ -15,14 +15,14 @@ import org.jetbrains.kotlin.ir.backend.js.*
 import org.jetbrains.kotlin.ir.backend.js.codegen.JsGenerationGranularity
 import org.jetbrains.kotlin.ir.backend.js.ic.JsExecutableProducer
 import org.jetbrains.kotlin.ir.backend.js.lower.serialization.ir.JsManglerDesc
-import org.jetbrains.kotlin.ir.backend.js.transformers.irToJs.IrModuleToJsTransformer
 import org.jetbrains.kotlin.ir.backend.js.SourceMapsInfo
-import org.jetbrains.kotlin.ir.backend.js.transformers.irToJs.CompilationOutputs
+import org.jetbrains.kotlin.ir.backend.js.transformers.irToJs.*
-import org.jetbrains.kotlin.ir.backend.js.transformers.irToJs.TranslationMode
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.ir.declarations.impl.IrFactoryImplForJsIC
 import org.jetbrains.kotlin.ir.util.SymbolTable
 import org.jetbrains.kotlin.ir.util.irMessageLogger
+import org.jetbrains.kotlin.js.backend.ast.ESM_EXTENSION
+import org.jetbrains.kotlin.js.backend.ast.REGULAR_EXTENSION
 import org.jetbrains.kotlin.js.config.JSConfigurationKeys
 import org.jetbrains.kotlin.js.test.handlers.JsBoxRunner.Companion.TEST_FUNCTION
 import org.jetbrains.kotlin.js.test.utils.extractTestPackage
@@ -37,13 +37,12 @@ import org.jetbrains.kotlin.test.frontend.classic.moduleDescriptorProvider
 import org.jetbrains.kotlin.test.model.*
 import org.jetbrains.kotlin.test.services.*
 import org.jetbrains.kotlin.test.services.configuration.JsEnvironmentConfigurator
+import org.jetbrains.kotlin.test.services.configuration.JsEnvironmentConfigurator.Companion.getJsArtifactSimpleName
 import org.jetbrains.kotlin.utils.addToStdlib.ifTrue
+import org.jetbrains.kotlin.utils.addToStdlib.runIf
 import org.jetbrains.kotlin.utils.fileUtils.withReplacedExtensionOrNull
 import java.io.File
 
-const val REGULAR_EXTENSION = ".js"
-const val ESM_EXTENSION = ".mjs"
-
 class JsIrBackendFacade(
     val testServices: TestServices,
     private val firstTimeCompilation: Boolean
@@ -132,10 +131,10 @@ class JsIrBackendFacade(
 
 
         val loweredIr = compileIr(
-            irModuleFragment.apply { resolveTestPathes() },
+            irModuleFragment.apply { resolveTestPaths() },
             MainModule.Klib(inputArtifact.outputFile.absolutePath),
             configuration,
-            dependencyModules.apply { forEach { it.resolveTestPathes() } },
+            dependencyModules.onEach { it.resolveTestPaths() },
             emptyMap(),
             irModuleFragment.irBuiltins,
             symbolTable,
@@ -150,14 +149,10 @@ class JsIrBackendFacade(
             granularity = granularity,
         )
 
-        return loweredIr2JsArtifact(module, loweredIr, granularity)
+        return loweredIr2JsArtifact(module, loweredIr)
     }
 
-    private fun loweredIr2JsArtifact(
+    private fun loweredIr2JsArtifact(module: TestModule, loweredIr: LoweredIr): BinaryArtifacts.Js {
-        module: TestModule,
-        loweredIr: LoweredIr,
-        granularity: JsGenerationGranularity,
-    ): BinaryArtifacts.Js {
         val mainArguments = JsEnvironmentConfigurator.getMainCallParametersForModule(module)
             .run { if (shouldBeGenerated()) arguments() else null }
         val runIrDce = JsEnvironmentConfigurationDirectives.RUN_IR_DCE in module.directives
@@ -172,10 +167,13 @@ class JsIrBackendFacade(
         val transformer = IrModuleToJsTransformer(
             loweredIr.context,
             mainArguments,
-            moduleToName = JsIrModuleToPath(
+            moduleToName = runIf(isEsModules) {
-                testServices,
+                loweredIr.allModules.associateWith {
-                isEsModules && granularity != JsGenerationGranularity.WHOLE_PROGRAM
+                    "./${getJsArtifactSimpleName(testServices, it.safeName)}_v5.mjs".run {
-            )
+                        if (isWindows) minify() else this
+                    }
+                }
+            } ?: emptyMap()
         )
         // If runIrDce then include DCE results
         // If perModuleOnly then skip whole program
@@ -188,7 +186,7 @@ class JsIrBackendFacade(
         return BinaryArtifacts.Js.JsIrArtifact(outputFile, compilationOut).dump(module)
     }
 
-    private fun IrModuleFragment.resolveTestPathes() {
+    private fun IrModuleFragment.resolveTestPaths() {
         JsIrPathReplacer(testServices).let {
             files.forEach(it::lower)
         }
@@ -274,7 +272,7 @@ class JsIrBackendFacade(
 
     private fun CompilationOutputs.writeTo(outputFile: File, moduleId: String, moduleKind: ModuleKind) {
         val allJsFiles = writeAll(outputFile.parentFile, outputFile.nameWithoutExtension, false, moduleId, moduleKind).filter {
-            it.extension == "js"
+            it.extension == "js" || it.extension == "mjs"
         }
 
         val mainModuleFile = allJsFiles.last()
@@ -296,12 +294,6 @@ class JsIrBackendFacade(
     }
 }
 
-val ModuleKind.extension: String
-    get() = when (this) {
-        ModuleKind.ES -> ESM_EXTENSION
-        else -> REGULAR_EXTENSION
-    }
-
 val RegisteredDirectives.moduleKind: ModuleKind
     get() = get(JsEnvironmentConfigurationDirectives.MODULE_KIND).singleOrNull()
         ?: if (contains(JsEnvironmentConfigurationDirectives.ES_MODULES)) ModuleKind.ES else ModuleKind.PLAIN
@@ -311,15 +303,15 @@ val TestModule.kind: ModuleKind
 
 fun String.augmentWithModuleName(moduleName: String): String {
     val normalizedName = moduleName.run { if (isWindows) minify() else this }
+    val suffix = when {
+        endsWith(ESM_EXTENSION) -> ESM_EXTENSION
+        endsWith(REGULAR_EXTENSION) -> REGULAR_EXTENSION
+        else -> error("Unexpected file '$this' extension")
+    }
 
-    return if (normalizedName.isPath()) {
+    return if (suffix == ESM_EXTENSION) {
-        replaceAfterLast(File.separator, normalizedName.replace("./", ""))
+        replaceAfterLast(File.separator, normalizedName.replace("./", "")).removeSuffix(suffix) + suffix
     } else {
-        val suffix = when {
-            endsWith(ESM_EXTENSION) -> ESM_EXTENSION
-            endsWith(REGULAR_EXTENSION) -> REGULAR_EXTENSION
-            else -> error("Unexpected file '$this' extension")
-        }
         return removeSuffix("_v5$suffix") + "-${normalizedName}_v5$suffix"
     }
 }

js/js.tests/test/org/jetbrains/kotlin/js/test/converters/JsIrModuleToPath.kt

@@ -1,35 +0,0 @@
-/*
- * 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.js.test.converters
-
-import org.jetbrains.kotlin.cli.common.isWindows
-import org.jetbrains.kotlin.ir.backend.js.transformers.irToJs.safeName
-import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
-import org.jetbrains.kotlin.test.services.TestServices
-import org.jetbrains.kotlin.test.services.configuration.JsEnvironmentConfigurator.Companion.getJsArtifactSimpleName
-import org.jetbrains.kotlin.utils.addToStdlib.runIf
-
-private typealias K = IrModuleFragment
-private typealias V = String
-
-class JsIrModuleToPath(val testServices: TestServices, shouldProvidePaths: Boolean) : Map<K, V> {
-    override val size = if (!shouldProvidePaths) 0 else 1
-    override val entries = emptySet<Map.Entry<K, V>>()
-    override val keys = emptySet<K>()
-    override val values = emptyList<V>()
-
-    override fun isEmpty() = size == 0
-    override fun containsKey(key: K): Boolean = !isEmpty()
-    override fun containsValue(value: V): Boolean = !isEmpty()
-
-    override operator fun get(key: K): V? {
-        return runIf(!isEmpty()) {
-            "./${getJsArtifactSimpleName(testServices, key.safeName)}_v5.mjs".run {
-                if (isWindows) minify() else this
-            }
-        }
-    }
-}

js/js.tests/test/org/jetbrains/kotlin/js/test/utils/RunnerUtils.kt

@@ -6,11 +6,11 @@
 package org.jetbrains.kotlin.js.test.utils
 
 import org.jetbrains.kotlin.ir.backend.js.transformers.irToJs.TranslationMode
+import org.jetbrains.kotlin.ir.backend.js.transformers.irToJs.extension
 import org.jetbrains.kotlin.js.JavaScript
 import org.jetbrains.kotlin.js.config.JSConfigurationKeys
 import org.jetbrains.kotlin.js.test.JsAdditionalSourceProvider
 import org.jetbrains.kotlin.js.test.converters.augmentWithModuleName
-import org.jetbrains.kotlin.js.test.converters.extension
 import org.jetbrains.kotlin.js.test.converters.kind
 import org.jetbrains.kotlin.js.test.handlers.JsBoxRunner.Companion.TEST_FUNCTION
 import org.jetbrains.kotlin.js.testOld.*
@@ -131,7 +131,7 @@ fun testWithModuleSystem(testServices: TestServices): Boolean {
     val globalDirectives = testServices.moduleStructure.allDirectives
     val configuration = testServices.compilerConfigurationProvider.getCompilerConfiguration(getMainModule(testServices))
     val mainModuleKind = configuration[JSConfigurationKeys.MODULE_KIND]
-    return mainModuleKind != ModuleKind.PLAIN && NO_JS_MODULE_SYSTEM !in globalDirectives
+    return mainModuleKind != ModuleKind.PLAIN && mainModuleKind != ModuleKind.ES && NO_JS_MODULE_SYSTEM !in globalDirectives
 }
 
 fun getModeOutputFilePath(testServices: TestServices, module: TestModule, mode: TranslationMode): String {

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/BoxJsTestGenerated.java

@@ -8934,6 +8934,12 @@ public class BoxJsTestGenerated extends AbstractBoxJsTest {
             runTest("js/js.translator/testData/box/reflection/findAssociatedObject.kt");
         }
 
+        @Test
+        @TestMetadata("findAssociatedObjectInSeparatedFile.kt")
+        public void testFindAssociatedObjectInSeparatedFile() throws Exception {
+            runTest("js/js.translator/testData/box/reflection/findAssociatedObjectInSeparatedFile.kt");
+        }
+
         @Test
         @TestMetadata("findAssociatedObject_oldBE.kt")
         public void testFindAssociatedObject_oldBE() throws Exception {

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/fir/FirJsBoxTestGenerated.java

@@ -2297,12 +2297,6 @@ public class FirJsBoxTestGenerated extends AbstractFirJsBoxTest {
                 runTest("js/js.translator/testData/box/esModules/export/overriddenExternalMethodWithSameStableNameMethodInExportedFile.kt");
             }
 
-            @Test
-            @TestMetadata("reexport.kt")
-            public void testReexport() throws Exception {
-                runTest("js/js.translator/testData/box/esModules/export/reexport.kt");
-            }
-
             @Test
             @TestMetadata("reservedModuleName.kt")
             public void testReservedModuleName() throws Exception {
@@ -9844,6 +9838,12 @@ public class FirJsBoxTestGenerated extends AbstractFirJsBoxTest {
             runTest("js/js.translator/testData/box/reflection/findAssociatedObject.kt");
         }
 
+        @Test
+        @TestMetadata("findAssociatedObjectInSeparatedFile.kt")
+        public void testFindAssociatedObjectInSeparatedFile() throws Exception {
+            runTest("js/js.translator/testData/box/reflection/findAssociatedObjectInSeparatedFile.kt");
+        }
+
         @Test
         @TestMetadata("findAssociatedObject_oldBE.kt")
         public void testFindAssociatedObject_oldBE() throws Exception {

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/ir/IrBoxJsES6TestGenerated.java

@@ -2403,12 +2403,6 @@ public class IrBoxJsES6TestGenerated extends AbstractIrBoxJsES6Test {
                 runTest("js/js.translator/testData/box/esModules/export/overriddenExternalMethodWithSameStableNameMethodInExportedFile.kt");
             }
 
-            @Test
-            @TestMetadata("reexport.kt")
-            public void testReexport() throws Exception {
-                runTest("js/js.translator/testData/box/esModules/export/reexport.kt");
-            }
-
             @Test
             @TestMetadata("reservedModuleName.kt")
             public void testReservedModuleName() throws Exception {
@@ -9950,6 +9944,12 @@ public class IrBoxJsES6TestGenerated extends AbstractIrBoxJsES6Test {
             runTest("js/js.translator/testData/box/reflection/findAssociatedObject.kt");
         }
 
+        @Test
+        @TestMetadata("findAssociatedObjectInSeparatedFile.kt")
+        public void testFindAssociatedObjectInSeparatedFile() throws Exception {
+            runTest("js/js.translator/testData/box/reflection/findAssociatedObjectInSeparatedFile.kt");
+        }
+
         @Test
         @TestMetadata("findAssociatedObject_oldBE.kt")
         public void testFindAssociatedObject_oldBE() throws Exception {

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/ir/IrBoxJsTestGenerated.java

@@ -2297,12 +2297,6 @@ public class IrBoxJsTestGenerated extends AbstractIrBoxJsTest {
                 runTest("js/js.translator/testData/box/esModules/export/overriddenExternalMethodWithSameStableNameMethodInExportedFile.kt");
             }
 
-            @Test
-            @TestMetadata("reexport.kt")
-            public void testReexport() throws Exception {
-                runTest("js/js.translator/testData/box/esModules/export/reexport.kt");
-            }
-
             @Test
             @TestMetadata("reservedModuleName.kt")
             public void testReservedModuleName() throws Exception {
@@ -9844,6 +9838,12 @@ public class IrBoxJsTestGenerated extends AbstractIrBoxJsTest {
             runTest("js/js.translator/testData/box/reflection/findAssociatedObject.kt");
         }
 
+        @Test
+        @TestMetadata("findAssociatedObjectInSeparatedFile.kt")
+        public void testFindAssociatedObjectInSeparatedFile() throws Exception {
+            runTest("js/js.translator/testData/box/reflection/findAssociatedObjectInSeparatedFile.kt");
+        }
+
         @Test
         @TestMetadata("findAssociatedObject_oldBE.kt")
         public void testFindAssociatedObject_oldBE() throws Exception {

js/js.translator/testData/box/esModules/export/reexport.kt

@@ -1,31 +0,0 @@
-// IGNORE_FIR
-// DONT_TARGET_EXACT_BACKEND: JS
-// ES_MODULES
-// EXPECTED_REACHABLE_NODES: 1283
-
-// MODULE: lib1
-// FILE: lib1.kt
-@JsExport
-fun bar() = "O"
-
-// MODULE: lib2(lib1)
-// FILE: lib2.kt
-
-@JsExport
-fun baz() = "K"
-
-// MODULE: main(lib2)
-// FILE: main.kt
-
-@JsExport
-fun result(o: String, k: String) = o + k
-
-// FILE: entry.mjs
-// ENTRY_ES_MODULE
-import { bar, baz, result } from "./reexport_v5.mjs";
-
-export function box() {
-    const o = bar();
-    const k = baz();
-    return result(o, k);
-}

js/js.translator/testData/box/reflection/findAssociatedObjectInSeparatedFile.kt

@@ -0,0 +1,110 @@
+// IGNORE_BACKEND: JS, JS_IR, JS_IR_ES6
+// KJS_WITH_FULL_RUNTIME
+
+// FILE: annotations.kt
+import kotlin.reflect.*
+
+@OptIn(ExperimentalAssociatedObjects::class)
+@AssociatedObjectKey
+@Retention(AnnotationRetention.BINARY)
+annotation class Associated3(val kClass: KClass<*>)
+
+// FILE: foo.kt
+@Associated1(Bar::class)
+@Associated2(Baz::class)
+class Foo
+
+// FILE: bar.kt
+import kotlin.reflect.*
+
+@OptIn(ExperimentalAssociatedObjects::class)
+@AssociatedObjectKey
+@Retention(AnnotationRetention.BINARY)
+annotation class Associated1(val kClass: KClass<*>)
+
+object Bar
+
+// FILE: baz.kt
+import kotlin.reflect.*
+
+@OptIn(ExperimentalAssociatedObjects::class)
+@AssociatedObjectKey
+@Retention(AnnotationRetention.BINARY)
+annotation class Associated2(val kClass: KClass<*>)
+
+object Baz
+
+private class C(var list: List<String>?)
+
+private interface I1 {
+    fun foo(): Int
+    fun bar(c: C)
+}
+
+private object I1Impl : I1 {
+    override fun foo() = 42
+    override fun bar(c: C) {
+        c.list = mutableListOf("zzz")
+    }
+}
+
+@Associated1(I1Impl::class)
+private class I1ImplHolder
+
+private interface I2 {
+    fun foo(): Int
+}
+
+private object I2Impl : I2 {
+    override fun foo() = 17
+}
+
+@Associated1(I2Impl::class)
+private class I2ImplHolder
+
+@Associated2(A.Companion::class)
+class A {
+    companion object : I2 {
+        override fun foo() = 20
+    }
+}
+
+@OptIn(ExperimentalAssociatedObjects::class)
+fun KClass<*>.getAssociatedObjectByAssociated2(): Any? {
+    return this.findAssociatedObject<Associated2>()
+}
+
+@OptIn(ExperimentalAssociatedObjects::class)
+fun box(): String {
+
+    if (Foo::class.findAssociatedObject<Associated1>() != Bar) return "fail 1"
+
+    if (Foo::class.findAssociatedObject<Associated2>() != Baz) return "fail 2"
+
+    if (Foo::class.findAssociatedObject<Associated3>() != null) return "fail 3"
+
+    if (Bar::class.findAssociatedObject<Associated1>() != null) return "fail 4"
+
+    val i1 = I1ImplHolder::class.findAssociatedObject<Associated1>() as I1
+    if (i1.foo() != 42) return "fail 5"
+
+    val c = C(null)
+    i1.bar(c)
+    if (c.list!![0] != "zzz") return "fail 6"
+
+    val i2 = I2ImplHolder()::class.findAssociatedObject<Associated1>() as I2
+    if (i2.foo() != 17) return "fail 7"
+
+    val a = A::class.findAssociatedObject<Associated2>() as I2
+    if (a.foo() != 20) return "fail 8"
+
+    if (Foo::class.getAssociatedObjectByAssociated2() != Baz) return "fail 9"
+
+    if ((A::class.getAssociatedObjectByAssociated2() as I2).foo() != 20) return "fail 10"
+
+    if (Int::class.findAssociatedObject<Associated1>() != null) return "fail 11"
+
+    if (10::class.findAssociatedObject<Associated2>() != null) return "fail 12"
+
+    return "OK"
+}

libraries/stdlib/js-ir/runtime/BitMask.kt

@@ -5,7 +5,38 @@
 
 package kotlin.js
 
-internal fun implement(vararg interfaces: dynamic): BitMask {
+internal typealias BitMask = IntArray
+
+private fun bitMaskWith(activeBit: Int): BitMask {
+    val intArray = IntArray((activeBit shr 5) + 1)
+    val numberIndex = activeBit shr 5
+    val positionInNumber = activeBit and 31
+    val numberWithSettledBit = 1 shl positionInNumber
+    intArray[numberIndex] = intArray[numberIndex] or numberWithSettledBit
+    return intArray
+}
+
+internal fun BitMask.isBitSet(possibleActiveBit: Int): Boolean {
+    val numberIndex = possibleActiveBit shr 5
+    if (numberIndex > size) return false
+    val positionInNumber = possibleActiveBit and 31
+    val numberWithSettledBit = 1 shl positionInNumber
+    return get(numberIndex) and numberWithSettledBit != 0
+}
+
+private fun compositeBitMask(capacity: Int, masks: Array<BitMask>): BitMask {
+    return IntArray(capacity) { i ->
+        var result = 0
+        for (mask in masks) {
+            if (i < mask.size) {
+                result = result or mask[i]
+            }
+        }
+        result
+    }
+}
+
+internal fun implement(interfaces: Array<dynamic>): BitMask {
     var maxSize = 1
     val masks = js("[]")
 
@@ -15,15 +46,15 @@ internal fun implement(vararg interfaces: dynamic): BitMask {
 
         if (imask != null) {
             masks.push(imask)
-            currentSize = imask.intArray.size
+            currentSize = imask.size
         }
 
         val iid: Int? = i.`$metadata$`.iid
-        val iidImask: BitMask? = iid?.let { BitMask(arrayOf(it)) }
+        val iidImask: BitMask? = iid?.let { bitMaskWith(it) }
 
         if (iidImask != null) {
             masks.push(iidImask)
-            currentSize = JsMath.max(currentSize, iidImask.intArray.size)
+            currentSize = JsMath.max(currentSize, iidImask.size)
         }
 
         if (currentSize > maxSize) {
@@ -31,42 +62,5 @@ internal fun implement(vararg interfaces: dynamic): BitMask {
         }
     }
 
-    val resultIntArray = IntArray(maxSize) { i ->
+    return compositeBitMask(maxSize, masks)
-        masks.reduce({ acc: Int, it: BitMask ->
-            if (i >= it.intArray.size)
-                acc
-            else
-                acc or it.intArray[i]
-        }, 0)
-    }
-
-    val result = BitMask(emptyArray())
-    result.intArray = resultIntArray
-    return result
 }
-
-internal class BitMask(activeBits: Array<Int>) {
-    var intArray: IntArray = run {
-        if (activeBits.size == 0) {
-            IntArray(0)
-        } else {
-            val max: Int = JsMath.asDynamic().max.apply(null, activeBits)
-            val intArray = IntArray((max shr 5) + 1)
-            for (activeBit in activeBits) {
-                val numberIndex = activeBit shr 5
-                val positionInNumber = activeBit and 31
-                val numberWithSettledBit = 1 shl positionInNumber
-                intArray[numberIndex] = intArray[numberIndex] or numberWithSettledBit
-            }
-            intArray
-        }
-    }
-
-    fun isBitSet(possibleActiveBit: Int): Boolean {
-        val numberIndex = possibleActiveBit shr 5
-        if (numberIndex > intArray.size) return false
-        val positionInNumber = possibleActiveBit and 31
-        val numberWithSettledBit = 1 shl positionInNumber
-        return intArray[numberIndex] and numberWithSettledBit != 0
-    }
-}

libraries/stdlib/js-ir/runtime/reflectRuntime.kt

@@ -36,7 +36,7 @@ private fun getPropertyRefClass(obj: Ctor, metadata: Metadata, imask: BitMask):
 }
 
 private fun getInterfaceMaskFor(obj: Ctor, superType: dynamic): BitMask =
-    obj.`$imask$` ?: implement(superType)
+    obj.`$imask$` ?: implement(arrayOf(superType))
 
 @Suppress("UNUSED_PARAMETER")
 private fun getKPropMetadata(paramCount: Int, setter: Any?): dynamic {

libraries/stdlib/js-ir/runtime/typeCheckUtils.kt

@@ -27,24 +27,26 @@ internal fun setMetadataFor(
 
     if (interfaces != null) {
         val receiver = if (metadata.iid != null) ctor else ctor.prototype
-        receiver.`$imask$` = implement(*interfaces)
+        receiver.`$imask$` = implement(interfaces)
     }
 }
 
 // There was a problem with per-module compilation (KT-55758) when the top-level state (iid) was reinitialized during stdlib module initialization
 // As a result we miss already incremented iid and had the same iids in two different modules
-// So, to keep the state consistent it was moved into the object
+// So, to keep the state consistent it was moved into the next lateinit variable and function
-private object InterfaceIdService {
+private lateinit var iid: Any
-    var iid: Int = 0
+
+private fun generateInterfaceId(): Int {
+    if (!::iid.isInitialized) {
+        iid = 0
+    }
+    iid = iid.unsafeCast<Int>() + 1
+    return iid.unsafeCast<Int>()
 }
 
-private fun InterfaceIdService.generateInterfaceId(): Int {
-    iid += 1
-    return iid
-}
 
 internal fun interfaceMeta(name: String?, associatedObjectKey: Number?, associatedObjects: dynamic, suspendArity: Array<Int>?): Metadata {
-    return createMetadata("interface", name, associatedObjectKey, associatedObjects, suspendArity, InterfaceIdService.generateInterfaceId())
+    return createMetadata("interface", name, associatedObjectKey, associatedObjects, suspendArity, generateInterfaceId())
 }
 
 internal fun objectMeta(name: String?, associatedObjectKey: Number?, associatedObjects: dynamic, suspendArity: Array<Int>?): Metadata {