[JS IR][Wasm] Preload (Suspend)FunctionN needed in lowerings

This reverts commit 15d178d850d15675ee39cbd833d13f4c8aed83ff.

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

@@ -57,14 +57,14 @@ class JsCommonCoroutineSymbols(
     val coroutineImpl =
         symbolTable.referenceClass(findClass(coroutinePackage.memberScope, COROUTINE_IMPL_NAME))
 
-    val coroutineImplLabelPropertyGetter by context.lazy2 { coroutineImpl.getPropertyGetter("state")!!.owner }
+    val coroutineImplLabelPropertyGetter by lazy { coroutineImpl.getPropertyGetter("state")!!.owner }
-    val coroutineImplLabelPropertySetter by context.lazy2 { coroutineImpl.getPropertySetter("state")!!.owner }
+    val coroutineImplLabelPropertySetter by lazy { coroutineImpl.getPropertySetter("state")!!.owner }
-    val coroutineImplResultSymbolGetter by context.lazy2 { coroutineImpl.getPropertyGetter("result")!!.owner }
+    val coroutineImplResultSymbolGetter by lazy { coroutineImpl.getPropertyGetter("result")!!.owner }
-    val coroutineImplResultSymbolSetter by context.lazy2 { coroutineImpl.getPropertySetter("result")!!.owner }
+    val coroutineImplResultSymbolSetter by lazy { coroutineImpl.getPropertySetter("result")!!.owner }
-    val coroutineImplExceptionPropertyGetter by context.lazy2 { coroutineImpl.getPropertyGetter("exception")!!.owner }
+    val coroutineImplExceptionPropertyGetter by lazy { coroutineImpl.getPropertyGetter("exception")!!.owner }
-    val coroutineImplExceptionPropertySetter by context.lazy2 { coroutineImpl.getPropertySetter("exception")!!.owner }
+    val coroutineImplExceptionPropertySetter by lazy { coroutineImpl.getPropertySetter("exception")!!.owner }
-    val coroutineImplExceptionStatePropertyGetter by context.lazy2 { coroutineImpl.getPropertyGetter("exceptionState")!!.owner }
+    val coroutineImplExceptionStatePropertyGetter by lazy { coroutineImpl.getPropertyGetter("exceptionState")!!.owner }
-    val coroutineImplExceptionStatePropertySetter by context.lazy2 { coroutineImpl.getPropertySetter("exceptionState")!!.owner }
+    val coroutineImplExceptionStatePropertySetter by lazy { coroutineImpl.getPropertySetter("exceptionState")!!.owner }
 
     val continuationClass = symbolTable.referenceClass(
         coroutinePackage.memberScope.getContributedClassifier(

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

@@ -78,7 +78,8 @@ fun compile(
     }
 
     val (moduleFragment: IrModuleFragment, dependencyModules, irBuiltIns, symbolTable, deserializer, moduleToName) =
-        loadIr(depsDescriptors, irFactory, verifySignatures, filesToLower)
+        loadIr(depsDescriptors, irFactory, verifySignatures, filesToLower, loadFunctionInterfacesIntoStdlib = true)
+
     val mainModule = depsDescriptors.mainModule
     val configuration = depsDescriptors.compilerConfiguration
 

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

@@ -91,7 +91,7 @@ private fun IrSimpleFunction.createSuspendFunctionStub(context: JsCommonBackendC
         updateFrom(this@createSuspendFunctionStub)
         isSuspend = false
         name = this@createSuspendFunctionStub.name
-        origin = this@createSuspendFunctionStub.origin
+        origin = IrDeclarationOrigin.LOWERED_SUSPEND_FUNCTION
         returnType = loweredSuspendFunctionReturnType(this@createSuspendFunctionStub, context.irBuiltIns)
     }.also { function ->
         function.parent = parent

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

@@ -11,7 +11,6 @@ import org.jetbrains.kotlin.backend.common.lower.FinallyBlocksLowering
 import org.jetbrains.kotlin.ir.backend.js.JsStatementOrigins
 import org.jetbrains.kotlin.backend.common.lower.ReturnableBlockTransformer
 import org.jetbrains.kotlin.ir.backend.js.JsCommonBackendContext
-import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
 import org.jetbrains.kotlin.ir.backend.js.ir.JsIrBuilder
 import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.declarations.*

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

@@ -121,10 +121,6 @@ class JsIntrinsicTransformers(backendContext: JsIrBackendContext) {
 
             addIfNotNull(intrinsics.jsCode) { _, _ -> error("Should not be called") }
 
-            add(intrinsics.jsGetContinuation) { _, _ -> error("getContinuation should be lowered") }
-
-            add(intrinsics.jsCoroutineContext) { _, _ -> error("coroutineContext should be lowered") }
-
             add(intrinsics.jsArrayLength) { call, context ->
                 val args = translateCallArguments(call, context)
                 JsNameRef("length", args[0])

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

@@ -52,15 +52,6 @@ fun compileWasm(
         ExternalDependenciesGenerator(symbolTable, listOf(deserializer)).generateUnboundSymbolsAsDependencies()
     }
 
-    // Preloading function interfaces that will potentially be referenced by IR lowering.
-    // TODO: Do a smart preload based on what references we have in IR and support big arity
-    repeat(22) {
-        irBuiltIns.functionN(it)
-        irBuiltIns.kFunctionN(it)
-        irBuiltIns.kSuspendFunctionN(it)
-        irBuiltIns.suspendFunctionN(it)
-    }
-
     val irFiles = allModules.flatMap { it.files }
     moduleFragment.files.clear()
     moduleFragment.files += irFiles

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/IrDeclarationOrigin.kt

@@ -79,6 +79,7 @@ interface IrDeclarationOrigin {
 
     object FIELD_FOR_OUTER_THIS : IrDeclarationOriginImpl("FIELD_FOR_OUTER_THIS", isSynthetic = true)
     object CONTINUATION : IrDeclarationOriginImpl("CONTINUATION", isSynthetic = true)
+    object LOWERED_SUSPEND_FUNCTION : IrDeclarationOriginImpl("LOWERED_SUSPEND_FUNCTION", isSynthetic = true)
 
     val isSynthetic: Boolean get() = false
 }

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/descriptors/IrDescriptorBasedFunctionFactory.kt

@@ -85,7 +85,9 @@ abstract class IrAbstractDescriptorBasedFunctionFactory {
 class IrDescriptorBasedFunctionFactory(
     private val irBuiltIns: IrBuiltInsOverDescriptors,
     private val symbolTable: SymbolTable,
-    getPackageFragment: ((PackageFragmentDescriptor) -> IrPackageFragment)? = null
+    getPackageFragment: ((PackageFragmentDescriptor) -> IrPackageFragment)? = null,
+    // Needed for JS and Wasm backends to "preload" interfaces that can referenced during lowerings
+    private val referenceFunctionsWhenKFunctionAreReferenced: Boolean = false,
 ) : IrAbstractDescriptorBasedFunctionFactory() {
     val getPackageFragment =
         getPackageFragment ?: symbolTable::declareExternalPackageFragmentIfNotExists
@@ -139,6 +141,9 @@ class IrDescriptorBasedFunctionFactory(
     }
 
     override fun kFunctionN(arity: Int, declarator: SymbolTable.((IrClassSymbol) -> IrClass) -> IrClass): IrClass {
+        if (referenceFunctionsWhenKFunctionAreReferenced)
+            functionN(arity)
+
         return kFunctionNMap.getOrPut(arity) {
             symbolTable.declarator { symbol ->
                 val descriptor = symbol.descriptor as FunctionClassDescriptor
@@ -149,6 +154,9 @@ class IrDescriptorBasedFunctionFactory(
     }
 
     override fun kSuspendFunctionN(arity: Int, declarator: SymbolTable.((IrClassSymbol) -> IrClass) -> IrClass): IrClass {
+        if (referenceFunctionsWhenKFunctionAreReferenced)
+            suspendFunctionN(arity)
+
         return kSuspendFunctionNMap.getOrPut(arity) {
             symbolTable.declarator { symbol ->
                 val descriptor = symbol.descriptor as FunctionClassDescriptor

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

@@ -274,8 +274,11 @@ fun loadIr(
     val createFunctionFactoryCallback =
         if (loadFunctionInterfacesIntoStdlib) {
             { packageFragmentDescriptor: PackageFragmentDescriptor ->
-                IrFileImpl(NaiveSourceBasedFileEntryImpl("${packageFragmentDescriptor.fqName}-[K][Suspend]Functions"), packageFragmentDescriptor, stdlibModule)
+                IrFileImpl(
-                    .also { stdlibModule.files += it }
+                    NaiveSourceBasedFileEntryImpl("${packageFragmentDescriptor.fqName}-[K][Suspend]Functions"),
+                    packageFragmentDescriptor,
+                    stdlibModule
+                ).also { stdlibModule.files += it }
             }
         } else {
             null
@@ -287,7 +290,7 @@ fun loadIr(
             val (psi2IrContext, _) = preparePsi2Ir(depsDescriptors, errorPolicy, symbolTable)
             val irBuiltIns = psi2IrContext.irBuiltIns
             (irBuiltIns as IrBuiltInsOverDescriptors).functionFactory =
-                IrDescriptorBasedFunctionFactory(irBuiltIns, symbolTable, createFunctionFactoryCallback)
+                IrDescriptorBasedFunctionFactory(irBuiltIns, symbolTable, createFunctionFactoryCallback, true)
 
             val feContext = psi2IrContext.run {
                 JsIrLinker.JsFePluginContext(moduleDescriptor, symbolTable, typeTranslator, irBuiltIns)
@@ -349,7 +352,7 @@ fun loadIr(
             val typeTranslator =
                 TypeTranslatorImpl(symbolTable, depsDescriptors.compilerConfiguration.languageVersionSettings, moduleDescriptor)
             val irBuiltIns = IrBuiltInsOverDescriptors(moduleDescriptor.builtIns, typeTranslator, symbolTable)
-            irBuiltIns.functionFactory = IrDescriptorBasedFunctionFactory(irBuiltIns, symbolTable, createFunctionFactoryCallback)
+            irBuiltIns.functionFactory = IrDescriptorBasedFunctionFactory(irBuiltIns, symbolTable, createFunctionFactoryCallback, true)
 
             val loweredIcData = if (!depsDescriptors.icUseStdlibCache && !depsDescriptors.icUseStdlibCache) emptyMap() else {
                 val result = mutableMapOf<ModuleDescriptor, SerializedIcData>()

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

@@ -15,7 +15,7 @@ internal fun <T> getContinuation(): Continuation<T> { throw Exception("Implement
 
 @PublishedApi
 @Suppress("UNCHECKED_CAST")
-internal suspend fun <T> returnIfSuspended(@Suppress("UNUSED_PARAMETER") argument: Any?): T {
+internal suspend fun <T> returnIfSuspended(argument: Any?): T {
     return argument as T
 }