Rewrite inner class lowerings

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

@@ -14,5 +14,6 @@ interface DeclarationFactory {
     fun getFieldForEnumEntry(enumEntry: IrEnumEntry): IrField
     fun getOuterThisField(innerClass: IrClass): IrField
     fun getInnerClassConstructorWithOuterThisParameter(innerClassConstructor: IrConstructor): IrConstructor
+    fun getInnerClassOriginalPrimaryConstructorOrNull(innerClass: IrClass): IrConstructor?
     fun getFieldForObjectInstance(singleton: IrClass): IrField
 }

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

@@ -17,15 +17,95 @@ import org.jetbrains.kotlin.ir.symbols.IrConstructorSymbol
 import org.jetbrains.kotlin.ir.symbols.IrValueParameterSymbol
 import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
 import org.jetbrains.kotlin.ir.types.classOrNull
+import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
 import org.jetbrains.kotlin.ir.util.dump
 import org.jetbrains.kotlin.ir.util.parentAsClass
 import org.jetbrains.kotlin.ir.util.patchDeclarationParents
-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.addToStdlib.cast
 
-class InnerClassesLowering(val context: BackendContext) : ClassLoweringPass {
+class InnerClassesLowering(val context: BackendContext) : DeclarationTransformer {
+
+    override fun lower(irFile: IrFile) {
+        runPostfix(true).toFileLoweringPass().lower(irFile)
+    }
+
+    override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
+
+        if (declaration is IrClass && declaration.isInner) {
+            declaration.declarations += context.declarationFactory.getOuterThisField(declaration)
+        } else if (declaration is IrConstructor) {
+            val irClass = declaration.parentAsClass
+            if (!irClass.isInner) return null
+
+            val newConstructor = lowerConstructor(declaration)
+            val oldConstructorParameterToNew = context.primaryConstructorParameterMap(declaration)
+            for ((oldParam, newParam) in oldConstructorParameterToNew.entries) {
+                newParam.defaultValue = oldParam.defaultValue?.let { oldDefault ->
+                    IrExpressionBodyImpl(oldDefault.startOffset, oldDefault.endOffset) {
+                        expression = oldDefault.expression.patchDeclarationParents(newConstructor)
+                    }
+                }
+            }
+
+            return listOf(newConstructor)
+        }
+
+        return null
+    }
+
+    private fun lowerConstructor(irConstructor: IrConstructor): IrConstructor {
+        val loweredConstructor = context.declarationFactory.getInnerClassConstructorWithOuterThisParameter(irConstructor)
+        val outerThisParameter = loweredConstructor.valueParameters[0]
+
+        val irClass = irConstructor.parentAsClass
+        val parentThisField = context.declarationFactory.getOuterThisField(irClass)
+
+        val blockBody = irConstructor.body as? IrBlockBody ?: throw AssertionError("Unexpected constructor body: ${irConstructor.body}")
+
+        loweredConstructor.body = IrBlockBodyImpl(blockBody.startOffset, blockBody.endOffset) {
+            context.createIrBuilder(irConstructor.symbol, irConstructor.startOffset, irConstructor.endOffset).apply {
+                statements.add(0, irSetField(irGet(irClass.thisReceiver!!), parentThisField, irGet(outerThisParameter)))
+            }
+
+            statements.addAll(blockBody.statements)
+
+            if (statements.find { it is IrInstanceInitializerCall } == null) {
+                val delegatingConstructorCall =
+                    statements.find { it is IrDelegatingConstructorCall } as IrDelegatingConstructorCall?
+                        ?: throw AssertionError("Delegating constructor call expected: ${irConstructor.dump()}")
+                delegatingConstructorCall.apply { dispatchReceiver = IrGetValueImpl(startOffset, endOffset, outerThisParameter.symbol) }
+            }
+            patchDeclarationParents(loweredConstructor)
+
+            val oldConstructorParameterToNew = context.primaryConstructorParameterMap(irConstructor)
+            transformChildrenVoid(VariableRemapper(oldConstructorParameterToNew))
+        }
+
+        return loweredConstructor
+    }
+
+}
+
+private fun BackendContext.primaryConstructorParameterMap(originalConstructor: IrConstructor): Map<IrValueParameter, IrValueParameter> {
+    val oldConstructorParameterToNew = HashMap<IrValueParameter, IrValueParameter>()
+
+    val loweredConstructor = declarationFactory.getInnerClassConstructorWithOuterThisParameter(originalConstructor)
+
+    originalConstructor.valueParameters.forEach { old ->
+        oldConstructorParameterToNew[old] = loweredConstructor.valueParameters[old.index + 1]
+    }
+
+    return oldConstructorParameterToNew
+}
+
+
+class InnerClassesMemberBodyLowering(val context: BackendContext) : BodyLoweringPass {
+    override fun lower(irFile: IrFile) {
+        runOnFilePostfix(irFile, true)
+    }
+
     private val IrValueSymbol.classForImplicitThis: IrClass?
         // TODO: is this the correct way to get the class?
         get() =
@@ -37,40 +117,31 @@ class InnerClassesLowering(val context: BackendContext) : ClassLoweringPass {
             } else
                 null
 
-    override fun lower(irClass: IrClass) {
+    override fun lower(irBody: IrBody, container: IrDeclaration) {
+        val irClass = container.parent as? IrClass ?: return
+
         if (!irClass.isInner) return
 
-        val parentThisField = context.declarationFactory.getOuterThisField(irClass)
+        if (container is IrField || container is IrAnonymousInitializer || container is IrValueParameter) {
-        val oldConstructorParameterToNew = HashMap<IrValueParameter, IrValueParameter>()
+            val primaryConstructor = context.declarationFactory.getInnerClassOriginalPrimaryConstructorOrNull(irClass)
-
+            if (primaryConstructor != null) {
-        fun lowerConstructor(irConstructor: IrConstructor): IrConstructor {
+                val oldConstructorParameterToNew = context.primaryConstructorParameterMap(primaryConstructor)
-            val loweredConstructor = context.declarationFactory.getInnerClassConstructorWithOuterThisParameter(irConstructor)
+                irBody.transformChildrenVoid(VariableRemapper(oldConstructorParameterToNew))
-            val outerThisParameter = loweredConstructor.valueParameters[0]
-
-            irConstructor.valueParameters.forEach { old ->
-                oldConstructorParameterToNew[old] = loweredConstructor.valueParameters[old.index + 1]
             }
-
-            val blockBody = irConstructor.body as? IrBlockBody ?: throw AssertionError("Unexpected constructor body: ${irConstructor.body}")
-            context.createIrBuilder(irConstructor.symbol, irConstructor.startOffset, irConstructor.endOffset).apply {
-                blockBody.statements.add(0, irSetField(irGet(irClass.thisReceiver!!), parentThisField, irGet(outerThisParameter)))
-            }
-            if (blockBody.statements.find { it is IrInstanceInitializerCall } == null) {
-                val delegatingConstructorCall =
-                    blockBody.statements.find { it is IrDelegatingConstructorCall } as IrDelegatingConstructorCall?
-                        ?: throw AssertionError("Delegating constructor call expected: ${irConstructor.dump()}")
-                delegatingConstructorCall.apply { dispatchReceiver = IrGetValueImpl(startOffset, endOffset, outerThisParameter.symbol) }
-            }
-            blockBody.patchDeclarationParents(loweredConstructor)
-            loweredConstructor.body = blockBody
-            return loweredConstructor
         }
 
-        irClass.declarations += parentThisField
+        irBody.fixThisReference(irClass, container)
-        irClass.transformDeclarationsFlat { irMember -> (irMember as? IrConstructor)?.let { listOf(lowerConstructor(it)) } }
+    }
-        irClass.transformChildrenVoid(VariableRemapper(oldConstructorParameterToNew))
 
-        irClass.transformChildrenVoid(object : IrElementTransformerVoidWithContext() {
+    fun IrBody.fixThisReference(irClass: IrClass, container: IrDeclaration) {
+        val enclosingFunction: IrDeclaration? = run {
+            var current: IrDeclaration? = container
+            while (current != null && current !is IrFunction) {
+                current = current.parent as? IrDeclaration
+            }
+            current
+        }
+        transformChildrenVoid(object : IrElementTransformerVoidWithContext() {
             override fun visitClassNew(declaration: IrClass): IrStatement =
                 declaration
 
@@ -83,7 +154,7 @@ class InnerClassesLowering(val context: BackendContext) : ClassLoweringPass {
                 val startOffset = expression.startOffset
                 val endOffset = expression.endOffset
                 val origin = expression.origin
-                val function = currentFunction?.irElement as? IrFunction
+                val function = (currentFunction?.irElement ?: enclosingFunction) as? IrFunction
                 val enclosingThisReceiver = function?.dispatchReceiverParameter ?: irClass.thisReceiver!!
 
                 var irThis: IrExpression = IrGetValueImpl(startOffset, endOffset, enclosingThisReceiver.symbol, origin)
@@ -117,9 +188,9 @@ val innerClassConstructorCallsPhase = makeIrFilePhase(
     description = "Handle constructor calls for inner classes"
 )
 
-class InnerClassConstructorCallsLowering(val context: BackendContext) : FileLoweringPass {
+class InnerClassConstructorCallsLowering(val context: BackendContext) : BodyLoweringPass {
-    override fun lower(irFile: IrFile) {
+    override fun lower(irBody: IrBody, container: IrDeclaration) {
-        irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
+        irBody.transformChildrenVoid(object : IrElementTransformerVoid() {
             override fun visitConstructorCall(expression: IrConstructorCall): IrExpression {
                 expression.transformChildrenVoid(this)
 

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

@@ -32,6 +32,7 @@ class JsDeclarationFactory(mapping: JsMapping) : DeclarationFactory {
     private val singletonFieldDescriptors = mapping.singletonFieldDescriptors
     private val outerThisFieldSymbols = mapping.outerThisFieldSymbols
     private val innerClassConstructors = mapping.innerClassConstructors
+    private val originalInnerClassPrimaryConstructorByClass = mapping.originalInnerClassPrimaryConstructorByClass
 
     override fun getFieldForEnumEntry(enumEntry: IrEnumEntry): IrField = TODO()
 
@@ -80,9 +81,19 @@ class JsDeclarationFactory(mapping: JsMapping) : DeclarationFactory {
 
         return innerClassConstructors.getOrPut(innerClassConstructor) {
             createInnerClassConstructorWithOuterThisParameter(innerClassConstructor)
+        }.also {
+            if (innerClassConstructor.isPrimary) {
+                originalInnerClassPrimaryConstructorByClass[innerClass] = innerClassConstructor
+            }
         }
     }
 
+    override fun getInnerClassOriginalPrimaryConstructorOrNull(innerClass: IrClass): IrConstructor? {
+        assert(innerClass.isInner) { "Class is not inner: $innerClass" }
+
+        return originalInnerClassPrimaryConstructorByClass[innerClass]
+    }
+
     private fun createInnerClassConstructorWithOuterThisParameter(oldConstructor: IrConstructor): IrConstructor {
         val irClass = oldConstructor.parent as IrClass
         val outerThisType = (irClass.parent as IrClass).defaultType

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

@@ -217,6 +217,13 @@ private val innerClassesLoweringPhase = makeJsModulePhase(
     description = "Capture outer this reference to inner class"
 )
 
+private val innerClassesMemberBodyLoweringPhase = makeJsModulePhase(
+    ::InnerClassesMemberBodyLowering,
+    name = "InnerClassesMemberBody",
+    description = "Replace `this` with 'outer this' field references",
+    prerequisite = setOf(innerClassesLoweringPhase)
+)
+
 private val innerClassConstructorCallsLoweringPhase = makeJsModulePhase(
     ::InnerClassConstructorCallsLowering,
     name = "InnerClassConstructorCallsLowering",
@@ -444,6 +451,7 @@ val jsPhases = namedIrModulePhase(
             localDeclarationsLoweringPhase then
             localClassExtractionPhase then
             innerClassesLoweringPhase then
+            innerClassesMemberBodyLoweringPhase then
             innerClassConstructorCallsLoweringPhase then
             propertiesLoweringPhase then
             primaryConstructorLoweringPhase then

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

@@ -14,4 +14,5 @@ class JsMapping : DefaultMapping() {
     val singletonFieldDescriptors = newMapping<IrClass, IrField>()
     val outerThisFieldSymbols = newMapping<IrClass, IrField>()
     val innerClassConstructors = newMapping<IrConstructor, IrConstructor>()
+    val originalInnerClassPrimaryConstructorByClass = newMapping<IrClass, IrConstructor>()
 }

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmLower.kt

@@ -183,6 +183,13 @@ private val innerClassesPhase = makeIrFilePhase(
     prerequisite = setOf(localDeclarationsPhase)
 )
 
+private val innerClassesMemberBodyPhase = makeIrFilePhase(
+    ::InnerClassesMemberBodyLowering,
+    name = "InnerClassesMemberBody",
+    description = "Replace `this` with 'outer this' field references",
+    prerequisite = setOf(innerClassesPhase)
+)
+
 private val staticInitializersPhase = makeIrFilePhase(
     ::StaticInitializersLowering,
     name = "StaticInitializers",
@@ -305,6 +312,7 @@ private val jvmFilePhases =
         addContinuationPhase then
 
         innerClassesPhase then
+        innerClassesMemberBodyPhase then
         innerClassConstructorCallsPhase then
 
         makePatchParentsPhase(2) then

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/descriptors/JvmDeclarationFactory.kt

@@ -39,6 +39,7 @@ class JvmDeclarationFactory(
     private val interfaceCompanionFieldDeclarations = HashMap<IrSymbolOwner, IrField>()
     private val outerThisDeclarations = HashMap<IrClass, IrField>()
     private val innerClassConstructors = HashMap<IrConstructor, IrConstructor>()
+    private val originalInnerClassPrimaryConstructorByClass = HashMap<IrClass, IrConstructor>()
     private val staticBackingFields = HashMap<IrProperty, IrField>()
 
     private val defaultImplsMethods = HashMap<IrSimpleFunction, IrSimpleFunction>()
@@ -74,13 +75,24 @@ class JvmDeclarationFactory(
         }
 
     override fun getInnerClassConstructorWithOuterThisParameter(innerClassConstructor: IrConstructor): IrConstructor {
-        assert((innerClassConstructor.parent as IrClass).isInner) { "Class is not inner: ${(innerClassConstructor.parent as IrClass).dump()}" }
+        val innerClass = innerClassConstructor.parent as IrClass
+        assert(innerClass.isInner) { "Class is not inner: ${(innerClassConstructor.parent as IrClass).dump()}" }
 
         return innerClassConstructors.getOrPut(innerClassConstructor) {
             createInnerClassConstructorWithOuterThisParameter(innerClassConstructor)
+        }.also {
+            if (innerClassConstructor.isPrimary) {
+                originalInnerClassPrimaryConstructorByClass[innerClass] = innerClassConstructor
+            }
         }
     }
 
+    override fun getInnerClassOriginalPrimaryConstructorOrNull(innerClass: IrClass): IrConstructor? {
+        assert(innerClass.isInner) { "Class is not inner: $innerClass" }
+
+        return originalInnerClassPrimaryConstructorByClass[innerClass]
+    }
+
     private fun createInnerClassConstructorWithOuterThisParameter(oldConstructor: IrConstructor): IrConstructor {
         val newDescriptor = WrappedClassConstructorDescriptor(oldConstructor.descriptor.annotations)
         return IrConstructorImpl(

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

@@ -165,6 +165,13 @@ private val innerClassesLoweringPhase = makeWasmModulePhase(
     description = "Capture outer this reference to inner class"
 )
 
+private val innerClassesMemberBodyLoweringPhase = makeWasmModulePhase(
+    ::InnerClassesMemberBodyLowering,
+    name = "InnerClassesMemberBody",
+    description = "Replace `this` with 'outer this' field references",
+    prerequisite = setOf(innerClassesLoweringPhase)
+)
+
 private val innerClassConstructorCallsLoweringPhase = makeWasmModulePhase(
     ::InnerClassConstructorCallsLowering,
     name = "InnerClassConstructorCallsLowering",
@@ -354,6 +361,7 @@ val wasmPhases = namedIrModulePhase<WasmBackendContext>(
             localDeclarationsLoweringPhase then
             localClassExtractionPhase then
             innerClassesLoweringPhase then
+            innerClassesMemberBodyLoweringPhase then
             innerClassConstructorCallsLoweringPhase then
             propertiesLoweringPhase then
             primaryConstructorLoweringPhase then