[IR] Update parents when transforming declaration containers

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

@@ -34,6 +34,7 @@ import org.jetbrains.kotlin.ir.symbols.impl.IrValueParameterSymbolImpl
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
 import org.jetbrains.kotlin.ir.util.defaultType
+import org.jetbrains.kotlin.ir.util.transformDeclarationsFlat
 import org.jetbrains.kotlin.ir.util.transformFlat
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
@@ -44,7 +45,7 @@ import org.jetbrains.kotlin.name.Name
 open class DefaultArgumentStubGenerator constructor(val context: CommonBackendContext, private val skipInlineMethods: Boolean = true) :
     DeclarationContainerLoweringPass {
     override fun lower(irDeclarationContainer: IrDeclarationContainer) {
-        irDeclarationContainer.declarations.transformFlat { memberDeclaration ->
+        irDeclarationContainer.transformDeclarationsFlat { memberDeclaration ->
             if (memberDeclaration is IrFunction)
                 lower(memberDeclaration)
             else
@@ -371,7 +372,7 @@ private fun IrFunction.generateDefaultsFunctionImpl(context: CommonBackendContex
     val newFunction = buildFunctionDeclaration(this, origin)
 
     val syntheticParameters = MutableList((valueParameters.size + 31) / 32) { i ->
-        valueParameter(valueParameters.size + i, parameterMaskName(i), context.irBuiltIns.intType)
+        newFunction.valueParameter(valueParameters.size + i, parameterMaskName(i), context.irBuiltIns.intType)
     }
 
     if (this is IrConstructor) {

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

@@ -21,7 +21,7 @@ import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
 import org.jetbrains.kotlin.ir.types.classifierOrNull
 import org.jetbrains.kotlin.ir.util.defaultType
 import org.jetbrains.kotlin.ir.util.dump
-import org.jetbrains.kotlin.ir.util.transformFlat
+import org.jetbrains.kotlin.ir.util.transformDeclarationsFlat
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import java.util.*
@@ -52,7 +52,7 @@ class InnerClassesLowering(val context: BackendContext) : ClassLoweringPass {
         }
 
         private fun lowerConstructors() {
-            irClass.declarations.transformFlat { irMember ->
+            irClass.transformDeclarationsFlat { irMember ->
                 if (irMember is IrConstructor)
                     listOf(lowerConstructor(irMember))
                 else

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

@@ -28,6 +28,7 @@ import org.jetbrains.kotlin.ir.symbols.impl.IrFieldSymbolImpl
 import org.jetbrains.kotlin.ir.symbols.impl.IrSimpleFunctionSymbolImpl
 import org.jetbrains.kotlin.ir.symbols.impl.IrValueParameterSymbolImpl
 import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.util.transformDeclarationsFlat
 import org.jetbrains.kotlin.ir.util.transformFlat
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
@@ -86,7 +87,7 @@ class LocalDeclarationsLowering(
         // Continuous numbering across all declarations in the container.
         lambdasCount = 0
 
-        irDeclarationContainer.declarations.transformFlat { memberDeclaration ->
+        irDeclarationContainer.transformDeclarationsFlat { memberDeclaration ->
             // TODO: may be do the opposite - specify the list of IR elements which need not to be transformed
             when (memberDeclaration) {
                 is IrFunction -> LocalDeclarationsTransformer(memberDeclaration).lowerLocalDeclarations()

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

@@ -10,7 +10,7 @@ import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.impl.IrBlockImpl
-import org.jetbrains.kotlin.ir.util.transformFlat
+import org.jetbrains.kotlin.ir.util.transformDeclarationsFlat
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.utils.addIfNotNull
@@ -23,13 +23,13 @@ class PropertiesLowering : IrElementTransformerVoid(), FileLoweringPass {
 
     override fun visitFile(declaration: IrFile): IrFile {
         declaration.transformChildrenVoid(this)
-        declaration.declarations.transformFlat { lowerProperty(it, ClassKind.CLASS) }
+        declaration.transformDeclarationsFlat { lowerProperty(it, ClassKind.CLASS) }
         return declaration
     }
 
     override fun visitClass(declaration: IrClass): IrStatement {
         declaration.transformChildrenVoid(this)
-        declaration.declarations.transformFlat { lowerProperty(it, declaration.kind) }
+        declaration.transformDeclarationsFlat { lowerProperty(it, declaration.kind) }
         return declaration
     }
 

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

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.*
 import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.util.transformDeclarationsFlat
 import org.jetbrains.kotlin.ir.util.transformFlat
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformer
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
@@ -26,7 +27,7 @@ class BlockDecomposerLowering(context: JsIrBackendContext) : DeclarationContaine
     private val nothingType = context.irBuiltIns.nothingType
 
     override fun lower(irDeclarationContainer: IrDeclarationContainer) {
-        irDeclarationContainer.declarations.transformFlat { declaration ->
+        irDeclarationContainer.transformDeclarationsFlat { declaration ->
             when (declaration) {
                 is IrFunction -> {
                     lower(declaration)

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

@@ -81,15 +81,16 @@ class EnumUsageLowering(val context: JsIrBackendContext) : FileLoweringPass {
 
 class EnumClassLowering(val context: JsIrBackendContext) : DeclarationContainerLoweringPass {
     override fun lower(irDeclarationContainer: IrDeclarationContainer) {
-        irDeclarationContainer.declarations.transformFlat { declaration ->
+        irDeclarationContainer.transformDeclarationsFlat { declaration ->
             if (declaration is IrClass && declaration.isEnumClass) {
                 if (declaration.descriptor.isExpect) {
                     emptyList()
                 } else {
                     EnumClassTransformer(context, declaration).transform()
                 }
-            } else
+            } else {
                 listOf(declaration)
+            }
         }
     }
 }
@@ -291,7 +292,7 @@ class EnumClassTransformer(val context: JsIrBackendContext, private val irClass:
     }
 
     private fun replaceIrEntriesWithCorrespondingClasses() {
-        irClass.declarations.transformFlat {
+        irClass.transformDeclarationsFlat {
             listOfNotNull(if (it is IrEnumEntry) it.correspondingClass else it)
         }
     }

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

@@ -25,7 +25,7 @@ import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
 import org.jetbrains.kotlin.ir.util.defaultType
-import org.jetbrains.kotlin.ir.util.transformFlat
+import org.jetbrains.kotlin.ir.util.transformDeclarationsFlat
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformer
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
@@ -54,7 +54,7 @@ class SecondaryCtorLowering(val context: JsIrBackendContext) {
     }
 
     private fun updateConstructorDeclarations(irClass: IrClass) {
-        irClass.declarations.transformFlat {
+        irClass.transformDeclarationsFlat {
             if (it is IrConstructor) {
                 oldCtorToNewMap[it.symbol]?.let { (newInit, newCreate) ->
                     listOf(newInit.owner, newCreate.owner)

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

@@ -58,7 +58,7 @@ internal class SuspendFunctionsLowering(val context: JsIrBackendContext): FileLo
     }
 
     private fun buildCoroutines(irFile: IrFile) {
-        irFile.declarations.transformFlat(::tryTransformSuspendFunction)
+        irFile.transformDeclarationsFlat(::tryTransformSuspendFunction)
         irFile.acceptVoid(object : IrElementVisitorVoid {
             override fun visitElement(element: IrElement) {
                 element.acceptChildrenVoid(this)
@@ -66,7 +66,7 @@ internal class SuspendFunctionsLowering(val context: JsIrBackendContext): FileLo
 
             override fun visitClass(declaration: IrClass) {
                 declaration.acceptChildrenVoid(this)
-                declaration.declarations.transformFlat(::tryTransformSuspendFunction)
+                declaration.transformDeclarationsFlat(::tryTransformSuspendFunction)
             }
         })
     }
@@ -343,6 +343,7 @@ internal class SuspendFunctionsLowering(val context: JsIrBackendContext): FileLo
             coroutineClassThis =
                     JsIrBuilder.buildValueParameter(Name.special("<this>"), -1, thisType, IrDeclarationOrigin.INSTANCE_RECEIVER)
             coroutineClass.thisReceiver = coroutineClassThis
+            coroutineClassThis.parent = coroutineClass
 
             val overriddenMap = mutableMapOf<IrSimpleFunction, IrSimpleFunctionSymbol>()
             val constructors = mutableSetOf<IrConstructor>()

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

@@ -8,12 +8,12 @@ package org.jetbrains.kotlin.ir.backend.js.lower.inline
 import org.jetbrains.kotlin.backend.common.DeclarationContainerLoweringPass
 import org.jetbrains.kotlin.ir.declarations.IrDeclarationContainer
 import org.jetbrains.kotlin.ir.declarations.IrFunction
-import org.jetbrains.kotlin.ir.util.transformFlat
+import org.jetbrains.kotlin.ir.util.transformDeclarationsFlat
 
 
 object RemoveInlineFunctionsWithReifiedTypeParametersLowering: DeclarationContainerLoweringPass {
     override fun lower(irDeclarationContainer: IrDeclarationContainer) {
-        irDeclarationContainer.declarations.transformFlat {
+        irDeclarationContainer.transformDeclarationsFlat {
             if (it is IrFunction && it.isInline && it.typeParameters.any { it.isReified }) listOf() else null
         }
     }

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/EnumClassLowering.kt

@@ -170,7 +170,7 @@ class EnumClassLowering(val context: JvmBackendContext) : ClassLoweringPass {
         }
 
         private fun lowerEnumEntries() {
-            irClass.declarations.transformFlat { declaration ->
+            irClass.transformDeclarationsFlat { declaration ->
                 if (declaration is IrEnumEntry) {
                     listOfNotNull(
                         createFieldForEnumEntry(declaration),

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/transform.kt

@@ -17,6 +17,8 @@
 package org.jetbrains.kotlin.ir.util
 
 import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.declarations.IrDeclaration
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationContainer
 
 inline fun <reified T : IrElement> MutableList<T>.transform(transformation: (T) -> IrElement) {
     forEachIndexed { i, item ->
@@ -46,3 +48,16 @@ inline fun <T> MutableList<T>.transformFlat(transformation: (T) -> List<T>?) {
         removeAt(i)
     }
 }
+
+/**
+ * Transforms declarations in declaration container.
+ * Behaves similar to like MutableList<T>.transformFlat but also updates
+ * parent property for transformed declarations.
+ */
+fun IrDeclarationContainer.transformDeclarationsFlat(transformation: (IrDeclaration) -> List<IrDeclaration>?) {
+    declarations.transformFlat { declaration ->
+        val transformed = transformation(declaration)
+        transformed?.forEach { it.parent = this }
+        transformed
+    }
+}