Refactor InnerClassesLowering

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

@@ -10,9 +10,10 @@ import org.jetbrains.kotlin.backend.common.BodyLoweringPass
 import org.jetbrains.kotlin.backend.common.ClassLoweringPass
 import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
 import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.builders.irGet
+import org.jetbrains.kotlin.ir.builders.irSetField
 import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.declarations.IrConstructor
-import org.jetbrains.kotlin.ir.declarations.IrField
 import org.jetbrains.kotlin.ir.declarations.IrValueParameter
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.*
@@ -20,111 +21,71 @@ 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.classifierOrNull
-import org.jetbrains.kotlin.ir.util.defaultType
 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 java.util.*
 
 val innerClassesPhase = makeIrFilePhase(
     ::InnerClassesLowering,
     name = "InnerClasses",
-    description = "Move inner classes to toplevel"
+    description = "Add 'outer this' fields to inner classes"
 )
 
 class InnerClassesLowering(val context: BackendContext) : ClassLoweringPass {
-    override fun lower(irClass: IrClass) {
+    private val IrValueSymbol.classForImplicitThis: IrClass?
-        InnerClassTransformer(irClass).lowerInnerClass()
+        // TODO: is this the correct way to get the class?
-    }
+        // -1 means value is either IMPLICIT or EXTENSION receiver
-
+        get() = if (this is IrValueParameterSymbol && owner.index == -1 && owner.name.isSpecial /* <this> */)
-    private inner class InnerClassTransformer(val irClass: IrClass) {
+            owner.type.classifierOrNull?.owner as IrClass
-        lateinit var outerThisField: IrField
-
-        val oldConstructorParameterToNew = HashMap<IrValueParameter, IrValueParameter>()
-
-        fun lowerInnerClass() {
-            if (!irClass.isInner) return
-
-            createOuterThisField()
-            lowerConstructors()
-            lowerConstructorParameterUsages()
-            lowerOuterThisReferences()
-        }
-
-        private fun createOuterThisField() {
-            val field = context.declarationFactory.getOuterThisField(irClass)
-            outerThisField = field
-            irClass.declarations += field
-        }
-
-        private fun lowerConstructors() {
-            irClass.transformDeclarationsFlat { irMember ->
-                if (irMember is IrConstructor)
-                    listOf(lowerConstructor(irMember))
         else
             null
-            }
-        }
 
-        private fun lowerConstructor(irConstructor: IrConstructor): IrConstructor {
+    override fun lower(irClass: IrClass) {
-            val startOffset = irConstructor.startOffset
+        if (!irClass.isInner) return
-            val endOffset = irConstructor.endOffset
 
+        val parentThisField = context.declarationFactory.getOuterThisField(irClass)
+        val oldConstructorParameterToNew = HashMap<IrValueParameter, IrValueParameter>()
+
+        fun lowerConstructor(irConstructor: IrConstructor): IrConstructor {
             val loweredConstructor = context.declarationFactory.getInnerClassConstructorWithOuterThisParameter(irConstructor)
-            val outerThisValueParameter = loweredConstructor.valueParameters[0].symbol
+            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 {
-            val instanceInitializerIndex = blockBody.statements.indexOfFirst { it is IrInstanceInitializerCall }
+                blockBody.statements.add(0, irSetField(irGet(irClass.thisReceiver!!), parentThisField, irGet(outerThisParameter)))
-
+            }
-            // Initializing constructor: initialize 'this.this$0' with '$outer'
+            if (blockBody.statements.find { it is IrInstanceInitializerCall } == null) {
-            blockBody.statements.add(
+                val delegatingConstructorCall =
-                0,
+                    blockBody.statements.find { it is IrDelegatingConstructorCall } as IrDelegatingConstructorCall?
-                IrSetFieldImpl(
-                    startOffset, endOffset, outerThisField.symbol,
-                    IrGetValueImpl(startOffset, endOffset, irClass.thisReceiver!!.symbol),
-                    IrGetValueImpl(startOffset, endOffset, outerThisValueParameter),
-                    context.irBuiltIns.unitType
-                )
-            )
-            if (instanceInitializerIndex < 0) {
-                // Delegating constructor: invoke old constructor with dispatch receiver '$outer'
-                val delegatingConstructorCall = (blockBody.statements.find { it is IrDelegatingConstructorCall }
                         ?: throw AssertionError("Delegating constructor call expected: ${irConstructor.dump()}")
-                        ) as IrDelegatingConstructorCall
+                delegatingConstructorCall.apply { dispatchReceiver = IrGetValueImpl(startOffset, endOffset, outerThisParameter.symbol) }
-                delegatingConstructorCall.dispatchReceiver = IrGetValueImpl(
-                    delegatingConstructorCall.startOffset, delegatingConstructorCall.endOffset, outerThisValueParameter
-                )
             }
             blockBody.patchDeclarationParents(loweredConstructor)
             loweredConstructor.body = blockBody
             return loweredConstructor
         }
 
-        private fun lowerConstructorParameterUsages() {
+        irClass.declarations += parentThisField
+        irClass.transformDeclarationsFlat { irMember -> (irMember as? IrConstructor)?.let { listOf(lowerConstructor(it)) } }
         irClass.transformChildrenVoid(VariableRemapper(oldConstructorParameterToNew))
-        }
 
-        private fun lowerOuterThisReferences() {
         irClass.transformChildrenVoid(object : IrElementTransformerVoid() {
-
-
-                override fun visitClass(declaration: IrClass): IrStatement =
             // TODO: maybe add another transformer that skips specified elements
+            override fun visitClass(declaration: IrClass): IrStatement =
                 declaration
 
             override fun visitGetValue(expression: IrGetValue): IrExpression {
                 expression.transformChildrenVoid(this)
 
-                    val implicitThisClass = expression.symbol.getClassForImplicitThis() ?: return expression
+                val implicitThisClass = expression.symbol.classForImplicitThis
-
+                if (implicitThisClass == null || implicitThisClass == irClass) return expression
-                    if (implicitThisClass == irClass) return expression
 
                 val startOffset = expression.startOffset
                 val endOffset = expression.endOffset
@@ -132,7 +93,6 @@ class InnerClassesLowering(val context: BackendContext) : ClassLoweringPass {
 
                 var irThis: IrExpression = IrGetValueImpl(startOffset, endOffset, irClass.thisReceiver!!.symbol, origin)
                 var innerClass = irClass
-
                 while (innerClass != implicitThisClass) {
                     if (!innerClass.isInner) {
                         // Captured 'this' unrelated to inner classes nesting hierarchy, leave it as is -
@@ -142,30 +102,12 @@ class InnerClassesLowering(val context: BackendContext) : ClassLoweringPass {
 
                     val outerThisField = context.declarationFactory.getOuterThisField(innerClass)
                     irThis = IrGetFieldImpl(startOffset, endOffset, outerThisField.symbol, outerThisField.type, irThis, origin)
-
+                    innerClass = innerClass.parentAsClass
-                        val outer = innerClass.parent
-                        innerClass = outer as? IrClass ?:
-                                throw AssertionError("Unexpected containing declaration for inner class $innerClass: $outer")
                 }
-
                 return irThis
             }
         })
     }
-
-        private fun IrValueSymbol.getClassForImplicitThis(): IrClass? {
-            //TODO: is it correct way to get class
-            if (this is IrValueParameterSymbol) {
-                val declaration = owner
-                if (declaration.index == -1) { // means value is either IMPLICIT or EXTENSION receiver
-                    if (declaration.name.isSpecial) { // whether name is <this>
-                        return owner.type.classifierOrNull?.owner as IrClass
-                    }
-                }
-            }
-            return null
-        }
-    }
 }
 
 val innerClassConstructorCallsPhase = makeIrFilePhase(