JVM_IR: make PropertyReferenceLowering less quadratic

ClassLoweringPass creates a visitor and calls lower() on each class,
which then creates another visitor that does not override `visitClass`
to ignore the nested classes. This is generally not a good idea.

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

@@ -5,7 +5,7 @@
 
 package org.jetbrains.kotlin.backend.jvm.lower
 
-import org.jetbrains.kotlin.backend.common.ClassLoweringPass
+import org.jetbrains.kotlin.backend.common.FileLoweringPass
 import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
 import org.jetbrains.kotlin.backend.common.ir.addSimpleDelegatingConstructor
 import org.jetbrains.kotlin.backend.common.ir.copyTo
@@ -14,7 +14,6 @@ import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
-import org.jetbrains.kotlin.backend.jvm.ir.JvmIrBuilder
 import org.jetbrains.kotlin.backend.jvm.ir.createJvmIrBuilder
 import org.jetbrains.kotlin.backend.jvm.ir.irArrayOf
 import org.jetbrains.kotlin.backend.jvm.ir.needsAccessor
@@ -38,7 +37,6 @@ import org.jetbrains.kotlin.ir.types.createType
 import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl
 import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection
 import org.jetbrains.kotlin.ir.util.*
-import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.load.java.JavaDescriptorVisibilities
 import org.jetbrains.kotlin.load.java.JvmAbi
 import org.jetbrains.kotlin.name.Name
@@ -54,7 +52,7 @@ internal val propertyReferencePhase = makeIrFilePhase(
     prerequisite = setOf(functionReferencePhase, suspendLambdaPhase)
 )
 
-internal class PropertyReferenceLowering(val context: JvmBackendContext) : ClassLoweringPass {
+private class PropertyReferenceLowering(val context: JvmBackendContext) : IrElementTransformerVoidWithContext(), FileLoweringPass {
     // Reflection metadata for local properties is serialized under the signature "<v#$N>" attached to the containing class.
     // This maps properties to values of N.
     private val localPropertyIndices = mutableMapOf<IrSymbol, Int>()
@@ -69,11 +67,6 @@ internal class PropertyReferenceLowering(val context: JvmBackendContext) : Class
     private val IrMemberAccessExpression<*>.field: IrFieldSymbol?
         get() = (this as? IrPropertyReference)?.field
 
-    // Plain Java fields do not have a getter, but can be referenced nonetheless. The signature should be the one
-    // that a getter would have, if it existed.
-    private val IrField.signature: String
-        get() = "${JvmAbi.getterName(name.asString())}()${context.methodSignatureMapper.mapReturnType(this)}"
-
     private val arrayItemGetter =
         context.ir.symbols.array.owner.functions.single { it.name.asString() == "get" }
 
@@ -100,16 +93,18 @@ internal class PropertyReferenceLowering(val context: JvmBackendContext) : Class
             return current.parent
         }
 
-    private fun IrBuilderWithScope.buildReflectedContainerReference(expression: IrMemberAccessExpression<*>): IrExpression =
+    // Plain Java fields do not have a getter, but can be referenced nonetheless. The signature should be the one
-        calculateOwner(expression.propertyContainer, this@PropertyReferenceLowering.context)
+    // that a getter would have, if it existed.
-
+    private val IrField.fakeGetterSignature: String
-    private fun JvmIrBuilder.buildReflectedContainerReferenceKClass(expression: IrMemberAccessExpression<*>): IrExpression =
+        get() = "${JvmAbi.getterName(name.asString())}()${context.methodSignatureMapper.mapReturnType(this)}"
-        calculateOwnerKClass(expression.propertyContainer, backendContext)
 
     private fun IrBuilderWithScope.computeSignatureString(expression: IrMemberAccessExpression<*>): IrExpression {
-        return expression.getter?.let { getter ->
+        if (expression is IrLocalDelegatedPropertyReference) {
-            localPropertyIndices[getter]?.let { irString("<v#$it>") }
+            // Local delegated properties are stored as a plain list, and the runtime library extracts the index from this string:
-                ?: irCall(signatureStringIntrinsic).apply {
+            val index = localPropertyIndices[expression.getter] ?: throw AssertionError("no index for ${expression.render()}")
+            return irString("<v#$index>")
+        }
+        val getter = expression.getter ?: return irString(expression.field!!.owner.fakeGetterSignature)
         // Work around for differences between `RuntimeTypeMapper.KotlinProperty` and the real Kotlin type mapper.
         // Most notably, the runtime type mapper does not perform inline class name mangling. This is usually not
         // a problem, since we will produce a getter signature as part of the Kotlin metadata, except when there
@@ -118,20 +113,12 @@ internal class PropertyReferenceLowering(val context: JvmBackendContext) : Class
         //
         // Note that we cannot compute the signature at this point, since we still need to mangle the names of
         // private properties in multifile-part classes.
-                    val needsDummySignature =
+        val needsDummySignature = getter.owner.correspondingPropertySymbol?.owner?.needsAccessor(getter.owner) == false ||
-                        getter.owner.correspondingPropertySymbol?.owner?.needsAccessor(getter.owner) == false ||
                 // Internal underlying vals of inline classes have no getter method
                 getter.owner.isInlineClassFieldGetter && getter.owner.visibility == DescriptorVisibilities.INTERNAL
-
+        val origin = if (needsDummySignature) InlineClassAbi.UNMANGLED_FUNCTION_REFERENCE else null
-                    putValueArgument(
+        val reference = IrFunctionReferenceImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, expression.type, getter, 0, getter, origin)
-                        0,
+        return irCall(signatureStringIntrinsic).apply { putValueArgument(0, reference) }
-                        IrFunctionReferenceImpl(
-                            UNDEFINED_OFFSET, UNDEFINED_OFFSET, expression.type, getter, 0, getter,
-                            if (needsDummySignature) InlineClassAbi.UNMANGLED_FUNCTION_REFERENCE else null
-                        )
-                    )
-                }
-        } ?: irString(expression.field!!.owner.signature)
     }
 
     private fun IrClass.addOverride(method: IrSimpleFunction, buildBody: IrBuilderWithScope.(List<IrValueParameter>) -> IrExpression) =
@@ -191,7 +178,7 @@ internal class PropertyReferenceLowering(val context: JvmBackendContext) : Class
 
     private data class PropertyInstance(val initializer: IrExpression, val index: Int)
 
-    override fun lower(irClass: IrClass) {
+    private inner class ClassData {
         val kProperties = mutableMapOf<IrSymbol, PropertyInstance>()
         val kPropertiesField = context.irFactory.buildField {
             name = Name.identifier(JvmAbi.DELEGATED_PROPERTIES_ARRAY_NAME)
@@ -202,10 +189,39 @@ internal class PropertyReferenceLowering(val context: JvmBackendContext) : Class
             visibility = JavaDescriptorVisibilities.PACKAGE_VISIBILITY
         }
         var localPropertiesInClass = 0
+    }
+
+    private var currentClassData: ClassData? = null
+
+    override fun lower(irFile: IrFile) =
+        irFile.transformChildrenVoid()
+
+    override fun visitClassNew(declaration: IrClass): IrStatement {
+        val data = ClassData()
+        val parentClassData = currentClassData
+        currentClassData = data
+        declaration.transformChildrenVoid()
+        currentClassData = parentClassData
+
+        // Put the new field at the beginning so that static delegated properties with initializers work correctly.
+        // Since we do not cache property references, the new field does not reference anything else.
+        if (data.kProperties.isNotEmpty()) {
+            declaration.declarations.add(0, data.kPropertiesField.apply {
+                parent = declaration
+                initializer = context.createJvmIrBuilder(declaration.symbol).run {
+                    val initializers = data.kProperties.values.sortedBy { it.index }.map { it.initializer }
+                    irExprBody(irArrayOf(kPropertiesFieldType, initializers))
+                }
+            })
+
+            context.localDelegatedProperties[declaration.attributeOwnerId] =
+                data.kProperties.keys.filterIsInstance<IrLocalDelegatedPropertySymbol>()
+        }
+        return declaration
+    }
 
-        irClass.transformChildrenVoid(object : IrElementTransformerVoidWithContext() {
     override fun visitLocalDelegatedProperty(declaration: IrLocalDelegatedProperty): IrStatement {
-                localPropertyIndices[declaration.getter.symbol] = localPropertiesInClass++
+        localPropertyIndices[declaration.getter.symbol] = currentClassData!!.localPropertiesInClass++
         return super.visitLocalDelegatedProperty(declaration)
     }
 
@@ -220,14 +236,15 @@ internal class PropertyReferenceLowering(val context: JvmBackendContext) : Class
         if (expression.origin != IrStatementOrigin.PROPERTY_REFERENCE_FOR_DELEGATE)
             return createSpecializedKProperty(expression)
 
+        val data = currentClassData ?: throw AssertionError("property reference not in class: ${expression.render()}")
         // For delegated properties, the getter and setter contain a reference each as the second argument to getValue
         // and setValue. Since it's highly unlikely that anyone will call get/set on these, optimize for space.
         return context.createIrBuilder(currentScope!!.scope.scopeOwnerSymbol, expression.startOffset, expression.endOffset).run {
-                    val (_, index) = kProperties.getOrPut(expression.symbol) {
+            val (_, index) = data.kProperties.getOrPut(expression.symbol) {
-                        PropertyInstance(createReflectedKProperty(expression), kProperties.size)
+                PropertyInstance(createReflectedKProperty(expression), data.kProperties.size)
             }
             irCall(arrayItemGetter).apply {
-                        dispatchReceiver = irGetField(null, kPropertiesField)
+                dispatchReceiver = irGetField(null, data.kPropertiesField)
                 putValueArgument(0, irInt(index))
             }
         }
@@ -242,7 +259,7 @@ internal class PropertyReferenceLowering(val context: JvmBackendContext) : Class
             irCall(referenceKind.wrapper).apply {
                 val constructor = referenceKind.implSymbol.constructors.single { it.owner.valueParameters.size == 3 }
                 putValueArgument(0, irCall(constructor).apply {
-                            putValueArgument(0, buildReflectedContainerReference(expression))
+                    putValueArgument(0, calculateOwner(expression.propertyContainer, this@PropertyReferenceLowering.context))
                     putValueArgument(1, irString(expression.referencedName.asString()))
                     putValueArgument(2, computeSignatureString(expression))
                 })
@@ -263,11 +280,7 @@ internal class PropertyReferenceLowering(val context: JvmBackendContext) : Class
     //
     private fun createSpecializedKProperty(expression: IrCallableReference<*>): IrExpression {
         val referenceClass = createKPropertySubclass(expression)
-                return context.createIrBuilder(
+        return context.createIrBuilder(currentScope!!.scope.scopeOwnerSymbol, expression.startOffset, expression.endOffset).irBlock {
-                        currentScope?.scope?.scopeOwnerSymbol ?: irClass.symbol, expression.startOffset, expression.endOffset
-                    )
-                    .irBlock {
-                        // TODO: Move this to the enclosing class, right now the parent field is wrong!
             +referenceClass
             +irCall(referenceClass.constructors.single()).apply {
                 var index = 0
@@ -286,7 +299,7 @@ internal class PropertyReferenceLowering(val context: JvmBackendContext) : Class
             origin = JvmLoweredDeclarationOrigin.GENERATED_PROPERTY_REFERENCE
             visibility = DescriptorVisibilities.LOCAL
         }.apply {
-                    parent = irClass
+            parent = currentDeclarationParent!!
             superTypes = listOf(superClass.defaultType)
             createImplicitParameterDeclarationWithWrappedDescriptor()
         }.copyAttributes(expression)
@@ -298,7 +311,7 @@ internal class PropertyReferenceLowering(val context: JvmBackendContext) : Class
             val getOwner = superClass.functions.single { it.name.asString() == "getOwner" }
             val getSignature = superClass.functions.single { it.name.asString() == "getSignature" }
             referenceClass.addOverride(getName) { irString(expression.referencedName.asString()) }
-                    referenceClass.addOverride(getOwner) { buildReflectedContainerReference(expression) }
+            referenceClass.addOverride(getOwner) { calculateOwner(expression.propertyContainer, this@PropertyReferenceLowering.context) }
             referenceClass.addOverride(getSignature) { computeSignatureString(expression) }
         }
 
@@ -392,7 +405,7 @@ internal class PropertyReferenceLowering(val context: JvmBackendContext) : Class
                             putValueArgument(index++, irGet(valueParameters.first()))
                         }
                         val callee = expression.symbol.owner as IrDeclaration
-                                val owner = buildReflectedContainerReferenceKClass(expression)
+                        val owner = calculateOwnerKClass(expression.propertyContainer, backendContext)
                         putValueArgument(index++, kClassToJavaClass(owner, backendContext))
                         putValueArgument(index++, irString(expression.referencedName.asString()))
                         putValueArgument(index++, computeSignatureString(expression))
@@ -403,21 +416,4 @@ internal class PropertyReferenceLowering(val context: JvmBackendContext) : Class
             }
         }
     }
-        })
-
-        // Put the new field at the beginning so that static delegated properties with initializers work correctly.
-        // Since we do not cache property references, the new field does not reference anything else.
-        if (kProperties.isNotEmpty()) {
-            irClass.declarations.add(0, kPropertiesField.apply {
-                parent = irClass
-                initializer = context.createJvmIrBuilder(irClass.symbol).run {
-                    val initializers = kProperties.values.sortedBy { it.index }.map { it.initializer }
-                    irExprBody(irArrayOf(kPropertiesFieldType, initializers))
-                }
-            })
-
-            context.localDelegatedProperties[irClass.attributeOwnerId] =
-                kProperties.keys.filterIsInstance<IrLocalDelegatedPropertySymbol>()
-        }
-    }
 }