[JS IR] A bunch of fixes to make PIR mode work

- provide required mappings
 - fix fake overrides
 - fix local declaration lowering
 - fix DCE

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

@@ -15,6 +15,9 @@ interface Mapping {
     val suspendFunctionToCoroutineConstructor: Delegate<IrFunction, IrConstructor>
     val lateInitFieldToNullableField: Delegate<IrField, IrField>
     val inlineClassMemberToStatic: Delegate<IrFunction, IrSimpleFunction>
+    val capturedFields: Delegate<IrClass, Collection<IrField>>
+    val capturedConstructors: Delegate<IrConstructor, IrConstructor>
+    val reflectedNameAccessor: Delegate<IrClass, IrSimpleFunction>
 
     abstract class Delegate<K : IrDeclaration, V> {
         abstract operator fun get(key: K): V?
@@ -36,6 +39,9 @@ open class DefaultMapping : Mapping {
     override val suspendFunctionToCoroutineConstructor: Mapping.Delegate<IrFunction, IrConstructor> = newMapping()
     override val lateInitFieldToNullableField: Mapping.Delegate<IrField, IrField> = newMapping()
     override val inlineClassMemberToStatic: Mapping.Delegate<IrFunction, IrSimpleFunction> = newMapping()
+    override val capturedFields: Mapping.Delegate<IrClass, Collection<IrField>> = newMapping()
+    override val capturedConstructors: Mapping.Delegate<IrConstructor, IrConstructor> = newMapping()
+    override val reflectedNameAccessor: Mapping.Delegate<IrClass, IrSimpleFunction> = newMapping()
 
     protected open fun <K : IrDeclaration, V> newMapping() = object : Mapping.Delegate<K, V>() {
         private val map: MutableMap<K, V> = mutableMapOf()

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

@@ -33,9 +33,7 @@ import org.jetbrains.kotlin.ir.symbols.impl.IrSimpleFunctionSymbolImpl
 import org.jetbrains.kotlin.ir.symbols.impl.IrTypeParameterSymbolImpl
 import org.jetbrains.kotlin.ir.symbols.impl.IrValueParameterSymbolImpl
 import org.jetbrains.kotlin.ir.types.*
-import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeBuilder
 import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl
-import org.jetbrains.kotlin.ir.types.impl.buildSimpleType
 import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
@@ -479,14 +477,14 @@ val IrFunction.allParameters: List<IrValueParameter>
         explicitParameters
     }
 
-fun IrClass.addFakeOverrides() {
+fun IrClass.addFakeOverrides(implementedMembers: List<IrSimpleFunction> = emptyList()) {
     fun IrDeclaration.toList() = when (this) {
         is IrSimpleFunction -> listOf(this)
         is IrProperty -> listOfNotNull(getter, setter)
         else -> emptyList()
     }
 
-    val overriddenFunctions = declarations
+    val overriddenFunctions = (declarations + implementedMembers)
         .flatMap { it.toList() }
         .flatMap { it.overriddenSymbols.map { it.owner } }
         .toSet()

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/ClosureAnnotator.kt

@@ -37,9 +37,6 @@ class ClosureAnnotator(body: IrBody, declaration: IrDeclaration) {
 
     init {
         // Collect all closures for classes and functions. Collect call graph
-        if (declaration is IrSimpleFunction && declaration.name.asString() == "arrayIterator") {
-            1
-        }
         body.accept(ClosureCollectorVisitor(), declaration.closureBuilderOrNull ?: declaration.parentClosureBuilder)
     }
 

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/LocalClassPopupLowering.kt

@@ -7,10 +7,14 @@ package org.jetbrains.kotlin.backend.common.lower
 
 import org.jetbrains.kotlin.backend.common.*
 import org.jetbrains.kotlin.backend.common.ir.addChild
+import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.IrBody
 import org.jetbrains.kotlin.ir.expressions.impl.IrCompositeImpl
+import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
+import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
+import org.jetbrains.kotlin.ir.visitors.acceptVoid
 
 //This lower takes part of old LocalDeclarationLowering job to pop up local classes from functions
 open class LocalClassPopupLowering(val context: BackendContext) : BodyLoweringPass {
@@ -44,7 +48,20 @@ open class LocalClassPopupLowering(val context: BackendContext) : BodyLoweringPa
 
         for ((local, newContainer) in extractedLocalClasses) {
             newContainer.addChild(local)
-            context.extractedLocalClasses += local
+
+            local.acceptVoid(object : IrElementVisitorVoid {
+                override fun visitElement(element: IrElement) {
+                    element.acceptChildrenVoid(this)
+                }
+
+                override fun visitBody(body: IrBody) {
+                }
+
+                override fun visitClass(declaration: IrClass) {
+                    super.visitClass(declaration)
+                    context.extractedLocalClasses += declaration
+                }
+            })
         }
     }
 

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/LocalDeclarationsLowering.kt

@@ -90,7 +90,7 @@ object BOUND_RECEIVER_PARAMETER : IrDeclarationOriginImpl("BOUND_RECEIVER_PARAME
   to proceed nevertheless.
 */
 class LocalDeclarationsLowering(
-    val context: BackendContext,
+    val context: CommonBackendContext,
     val localNameProvider: LocalNameProvider = LocalNameProvider.DEFAULT,
     val visibilityPolicy: VisibilityPolicy = VisibilityPolicy.DEFAULT
 ) :
@@ -463,6 +463,9 @@ class LocalDeclarationsLowering(
 
             irClass.declarations += localClassContext.capturedValueToField.values
 
+            context.mapping.capturedFields[irClass] =
+                (context.mapping.capturedFields[irClass] ?: emptyList()) + localClassContext.capturedValueToField.values
+
             for (constructorContext in constructorsCallingSuper) {
                 val blockBody = constructorContext.declaration.body as? IrBlockBody
                     ?: throw AssertionError("Unexpected constructor body: ${constructorContext.declaration.body}")
@@ -733,6 +736,7 @@ class LocalDeclarationsLowering(
             newDeclaration.metadata = oldDeclaration.metadata
 
             transformedDeclarations[oldDeclaration] = newDeclaration
+            context.mapping.capturedConstructors[oldDeclaration] = newDeclaration
         }
 
         private fun createFieldForCapturedValue(

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

@@ -311,9 +311,10 @@ fun usefulDeclarations(roots: Iterable<IrDeclaration>, context: JsIrBackendConte
             // TODO find out how `doResume` gets removed
             if (klass.symbol == context.ir.symbols.coroutineImpl) {
                 ArrayList(klass.declarations).forEach {
-                    if (it is IrSimpleFunction && it.name.asString() == "doResume") {
+                    // TODO: fix the heck
+//                    if (it is IrSimpleFunction && it.name.asString() == "doResume") {
                         it.enqueue()
-                    }
+//                    }
                 }
             }
         }

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

@@ -11,10 +11,13 @@ import org.jetbrains.kotlin.backend.common.ir.copyTo
 import org.jetbrains.kotlin.backend.common.ir.createImplicitParameterDeclarationWithWrappedDescriptor
 import org.jetbrains.kotlin.backend.common.ir.moveBodyTo
 import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
+import org.jetbrains.kotlin.backend.common.runOnFilePostfix
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
-import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.builders.declarations.*
+import org.jetbrains.kotlin.ir.builders.irBlockBody
+import org.jetbrains.kotlin.ir.builders.irDelegatingConstructorCall
+import org.jetbrains.kotlin.ir.builders.setSourceRange
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.declarations.impl.IrValueParameterImpl
 import org.jetbrains.kotlin.ir.descriptors.WrappedReceiverParameterDescriptor
@@ -24,13 +27,21 @@ import org.jetbrains.kotlin.ir.symbols.impl.IrValueParameterSymbolImpl
 import org.jetbrains.kotlin.ir.types.IrSimpleType
 import org.jetbrains.kotlin.ir.types.IrTypeProjection
 import org.jetbrains.kotlin.ir.types.classOrNull
-import org.jetbrains.kotlin.ir.util.*
+import org.jetbrains.kotlin.ir.util.defaultType
+import org.jetbrains.kotlin.ir.util.explicitParameters
+import org.jetbrains.kotlin.ir.util.isSuspend
+import org.jetbrains.kotlin.ir.util.parentAsClass
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.SpecialNames
 
 class CallableReferenceLowering(private val context: CommonBackendContext) : BodyLoweringPass {
+
+    override fun lower(irFile: IrFile) {
+        runOnFilePostfix(irFile, withLocalDeclarations = true)
+    }
+
     override fun lower(irBody: IrBody, container: IrDeclaration) {
         val realContainer = container as? IrDeclarationParent ?: container.parent
         irBody.transformChildrenVoid(ReferenceTransformer(realContainer))
@@ -40,6 +51,11 @@ class CallableReferenceLowering(private val context: CommonBackendContext) : Bod
     private val stringType = context.irBuiltIns.stringType
 
     private inner class ReferenceTransformer(private val container: IrDeclarationParent) : IrElementTransformerVoid() {
+
+        override fun visitBody(body: IrBody): IrBody {
+            return body
+        }
+
         override fun visitFunctionExpression(expression: IrFunctionExpression): IrExpression {
             expression.transformChildrenVoid(this)
 
@@ -316,6 +332,8 @@ class CallableReferenceLowering(private val context: CommonBackendContext) : Bod
                     )
                 )
             )
+
+            context.mapping.reflectedNameAccessor[clazz] = getter
         }
 
         fun build(): Pair<IrClass, IrConstructor> {

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

@@ -202,7 +202,7 @@ class InteropCallableReferenceLowering(val context: JsIrBackendContext) : BodyLo
         val functionExpression =
             expression.run { IrFunctionExpressionImpl(startOffset, endOffset, type, lambdaDeclaration, expression.origin!!) }
 
-        val nameGetter = lambdaClass.declarations.filterIsInstance<IrSimpleFunction>().singleOrNull { it.correspondingPropertySymbol != null }
+        val nameGetter = context.mapping.reflectedNameAccessor[lambdaClass]
 
         if (nameGetter != null || lambdaDeclaration.isSuspend) {
             val tmpVar = JsIrBuilder.buildVar(functionExpression.type, factoryFunction, "l", initializer = functionExpression)

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

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.ir.backend.js.lower.coroutines
 
 import org.jetbrains.kotlin.backend.common.*
 import org.jetbrains.kotlin.backend.common.ir.*
-import org.jetbrains.kotlin.backend.common.lower.LocalDeclarationsLowering
 import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Modality
@@ -258,7 +257,9 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
 
         private fun buildConstructor(): IrConstructor {
             if (isSuspendLambda) {
-                return coroutineClass.declarations.filterIsInstance<IrConstructor>().single()
+                return coroutineClass.declarations.filterIsInstance<IrConstructor>().single().let {
+                    context.mapping.capturedConstructors[it] ?: it
+                }
             }
 
             return WrappedClassConstructorDescriptor().let { d ->
@@ -401,9 +402,8 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
 
                     val thisReceiver = this.dispatchReceiverParameter!!
 
-                    val boundFields = coroutineClass.declarations.filterIsInstance<IrField>().filter { f ->
+                    val boundFields =
-                        f.origin == LocalDeclarationsLowering.DECLARATION_ORIGIN_FIELD_FOR_CAPTURED_VALUE || f.origin == LocalDeclarationsLowering.DECLARATION_ORIGIN_FIELD_FOR_CROSSINLINE_CAPTURED_VALUE
+                        context.mapping.capturedFields[coroutineClass] ?: error("No captured values for class ${coroutineClass.render()}")
-                    }
 
                     val irBuilder = context.createIrBuilder(symbol, startOffset, endOffset)
                     body = irBuilder.irBlockBody(startOffset, endOffset) {
@@ -422,9 +422,7 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
 
                         assert(createValueParameters.size - 1 == argumentToPropertiesMap.size)
 
-                        val ptofMap = createValueParameters.zip(argumentToPropertiesMap.values)
+                        for ((p, f) in createValueParameters.zip(argumentToPropertiesMap.values)) {
-
-                        for ((p, f) in ptofMap) {
                             +irSetField(irGet(instanceVal), f, irGet(p))
                         }
 
@@ -460,9 +458,13 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
         fun build(): BuiltCoroutine {
             val coroutineConstructor = buildConstructor()
 
+            val implementedMembers = ArrayList<IrSimpleFunction>(2)
+
             val superInvokeSuspendFunction = coroutineBaseClass.owner.simpleFunctions().single { it.name == stateMachineMethodName }
             val invokeSuspendMethod = buildInvokeSuspendMethod(superInvokeSuspendFunction)
 
+            implementedMembers.add(invokeSuspendMethod)
+
             if (isSuspendLambda) {
                 // Suspend lambda - create factory methods.
                 val createFunction = coroutineBaseClass.owner.simpleFunctions()
@@ -471,13 +473,14 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
                     }
 
                 val createMethod = buildCreateMethod(createFunction, coroutineConstructor)
+                implementedMembers.add(createMethod)
 
                 transformInvokeMethod(createMethod, invokeSuspendMethod)
             } else {
                 coroutineClass.superTypes += coroutineBaseClass.defaultType
             }
 
-            coroutineClass.addFakeOverrides()
+            coroutineClass.addFakeOverrides(implementedMembers)
 
             initializeStateMachine(listOf(coroutineConstructor), coroutineClassThis)