[JS IR BE] move exceptionState initialization from constructor to coroutine entryState

This fixes the DCE-driven mode. Before this fix the suspend lambda constructor was modified
by the `invoke` body lowering. Which is wrong, becuase the corresponding class was created
by the CallableReferenceLoiwering much earlier.

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

@@ -398,16 +398,6 @@ fun usefulDeclarations(roots: Iterable<IrDeclaration>, context: JsIrBackendConte
                     declaration.setter?.enqueue("(setter) overrides useful declaration")
                 }
             }
-
-            // TODO find out how `doResume` gets removed
-            if (klass.symbol == context.ir.symbols.coroutineImpl) {
-                ArrayList(klass.declarations).forEach {
-                    // TODO: fix the heck
-//                    if (it is IrSimpleFunction && it.name.asString() == "doResume") {
-                        it.enqueue(klass, "hack for CoroutineImpl::doResume")
-//                    }
-                }
-            }
         }
     }
 

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

@@ -59,7 +59,10 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
 
     protected abstract fun IrBlockBodyBuilder.generateCoroutineStart(invokeSuspendFunction: IrFunction, receiver: IrExpression)
 
-    protected abstract fun initializeStateMachine(coroutineConstructors: List<IrConstructor>, coroutineClassThis: IrValueDeclaration)
+    protected fun initializeStateMachine(coroutineConstructors: List<IrConstructor>, coroutineClassThis: IrValueDeclaration) {
+        // Do nothing by default
+        // TODO find out if Kotlin/Native needs this method.
+    }
 
     protected open fun IrBuilderWithScope.generateDelegatedCall(expectedType: IrType, delegatingCall: IrExpression): IrExpression =
         delegatingCall
@@ -470,7 +473,7 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
 
             coroutineClass.addFakeOverridesViaIncorrectHeuristic(implementedMembers)
 
-            // TODO: to meet PIR lower model constructor modification shouldn't be performed here
+            // TODO: find out whether Kotlin/Native needs this call
             initializeStateMachine(listOf(coroutineConstructor), coroutineClassThis)
 
             return BuiltCoroutine(

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

@@ -18,7 +18,6 @@ import org.jetbrains.kotlin.ir.expressions.impl.*
 import org.jetbrains.kotlin.ir.symbols.*
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.util.explicitParameters
-import org.jetbrains.kotlin.ir.util.patchDeclarationParents
 import org.jetbrains.kotlin.ir.visitors.*
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.utils.DFS
@@ -35,7 +34,6 @@ class JsSuspendFunctionsLowering(ctx: JsIrBackendContext) : AbstractSuspendFunct
     private val coroutineImplResultSymbolGetter = ctx.coroutineImplResultSymbolGetter
     private val coroutineImplResultSymbolSetter = ctx.coroutineImplResultSymbolSetter
 
-    private var exceptionTrapId = -1
     private var coroutineId = 0
 
     override val stateMachineMethodName = Name.identifier("doResume")
@@ -124,7 +122,17 @@ class JsSuspendFunctionsLowering(ctx: JsIrBackendContext) : AbstractSuspendFunct
 
         assignStateIds(stateMachineBuilder.entryState, stateVar.symbol, switch, rootLoop)
 
-        exceptionTrapId = stateMachineBuilder.rootExceptionTrap.id
+        // Set exceptionState to the global catch block
+        stateMachineBuilder.entryState.entryBlock.run {
+            val receiver = JsIrBuilder.buildGetValue(coroutineClass.thisReceiver!!.symbol)
+            val exceptionTrapId = stateMachineBuilder.rootExceptionTrap.id
+            assert(exceptionTrapId >= 0)
+            val id = JsIrBuilder.buildInt(context.irBuiltIns.intType, exceptionTrapId)
+            statements.add(0, JsIrBuilder.buildCall(coroutineImplExceptionStatePropertySetter.symbol).also { call ->
+                call.dispatchReceiver = receiver
+                call.putValueArgument(0, id)
+            })
+        }
 
         val functionBody =
             IrBlockBodyImpl(stateMachineFunction.startOffset, stateMachineFunction.endOffset, listOf(suspendResult, rootLoop))
@@ -210,20 +218,6 @@ class JsSuspendFunctionsLowering(ctx: JsIrBackendContext) : AbstractSuspendFunct
         return result
     }
 
-    override fun initializeStateMachine(coroutineConstructors: List<IrConstructor>, coroutineClassThis: IrValueDeclaration) {
-        for (it in coroutineConstructors) {
-            (it.body as? IrBlockBody)?.run {
-                val receiver = JsIrBuilder.buildGetValue(coroutineClassThis.symbol)
-                assert(exceptionTrapId >= 0)
-                val id = JsIrBuilder.buildInt(context.irBuiltIns.intType, exceptionTrapId)
-                statements += JsIrBuilder.buildCall(coroutineImplExceptionStatePropertySetter.symbol).also { call ->
-                    call.dispatchReceiver = receiver
-                    call.putValueArgument(0, id)
-                }
-            }
-        }
-    }
-
     override fun IrBuilderWithScope.generateDelegatedCall(expectedType: IrType, delegatingCall: IrExpression): IrExpression {
         val fromType = (delegatingCall as? IrCall)?.symbol?.owner?.returnType ?: delegatingCall.type
         if (!needUnboxingOrUnit(fromType, expectedType)) return delegatingCall

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

@@ -45,7 +45,9 @@ class JsGenerationContext(
 
     private fun isCoroutineDoResume(): Boolean {
         val overriddenSymbols = (currentFunction as? IrSimpleFunction)?.overriddenSymbols ?: return false
-        return staticContext.doResumeFunctionSymbol in overriddenSymbols
+        return overriddenSymbols.any {
+            it.owner.name.asString() == "doResume" && it.owner.parent == staticContext.coroutineImplDeclaration
+        }
     }
 
     fun checkIfJsCode(symbol: IrFunctionSymbol): Boolean = symbol == staticContext.backendContext.intrinsics.jsCode

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

@@ -21,8 +21,6 @@ class JsStaticContext(
     val intrinsics = JsIntrinsicTransformers(backendContext)
     val classModels = mutableMapOf<IrClassSymbol, JsIrClassModel>()
     val coroutineImplDeclaration = backendContext.ir.symbols.coroutineImpl.owner
-    val doResumeFunctionSymbol = coroutineImplDeclaration.declarations
-        .filterIsInstance<IrSimpleFunction>().single { it.name.asString() == "doResume" }.symbol
 
     val initializerBlock = JsGlobalBlock()
 }