JVM IR: add jvmLocalClassExtractionPhase to lift out local classes from initializers

Otherwise a local class in a field initializer or anonymous init block
is copied into each constructor of the containing class (because
InitializersLowering calls deepCopy).

Since the code structure no longer resembles the original source code
here, record a custom EnclosingMethod mapping before moving such
classes, and use it in codegen.

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

@@ -1,8 +1,14 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
 package org.jetbrains.kotlin.backend.common.lower
 
 import org.jetbrains.kotlin.backend.common.BackendContext
 import org.jetbrains.kotlin.backend.common.FileLoweringPass
 import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
+import org.jetbrains.kotlin.backend.common.ScopeWithIr
 import org.jetbrains.kotlin.backend.common.ir.addChild
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.declarations.IrClass
@@ -12,8 +18,7 @@ import org.jetbrains.kotlin.ir.declarations.IrScript
 import org.jetbrains.kotlin.ir.expressions.impl.IrCompositeImpl
 
 //This lower takes part of old LocalDeclarationLowering job to pop up local classes from functions
-class LocalClassPopupLowering(val context: BackendContext) : FileLoweringPass {
-
+open class LocalClassPopupLowering(val context: BackendContext) : FileLoweringPass {
     override fun lower(irFile: IrFile) {
         val extractedLocalClasses = arrayListOf<Pair<IrClass, IrDeclarationContainer>>()
 
@@ -21,9 +26,10 @@ class LocalClassPopupLowering(val context: BackendContext) : FileLoweringPass {
 
             override fun visitClassNew(declaration: IrClass): IrStatement {
                 val newDeclaration = super.visitClassNew(declaration)
-                if (newDeclaration !is IrClass || !newDeclaration.isLocalNotInner()) {
-                    return newDeclaration
-                }
+                if (newDeclaration !is IrClass) return newDeclaration
+
+                val currentScope = allScopes[allScopes.lastIndex - 1]
+                if (!shouldPopUp(declaration, currentScope)) return newDeclaration
 
                 val newContainer = allScopes.asReversed().drop(1/*skip self*/).firstOrNull {
                     //find first class local or not;
@@ -39,4 +45,7 @@ class LocalClassPopupLowering(val context: BackendContext) : FileLoweringPass {
             newContainer.addChild(local)
         }
     }
-}
+
+    protected open fun shouldPopUp(klass: IrClass, currentScope: ScopeWithIr?): Boolean =
+        klass.isLocalNotInner()
+}

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

@@ -27,12 +27,8 @@ import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.descriptors.IrBuiltIns
 import org.jetbrains.kotlin.ir.expressions.IrExpression
-import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
-import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
-import org.jetbrains.kotlin.ir.symbols.IrLocalDelegatedPropertySymbol
-import org.jetbrains.kotlin.ir.symbols.IrTypeParameterSymbol
+import org.jetbrains.kotlin.ir.symbols.*
 import org.jetbrains.kotlin.ir.types.IrType
-import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.util.ReferenceSymbolTable
 import org.jetbrains.kotlin.ir.util.SymbolTable
 import org.jetbrains.kotlin.name.FqName
@@ -65,7 +61,6 @@ class JvmBackendContext(
 
     val irIntrinsics = IrIntrinsicMethods(irBuiltIns, ir.symbols)
 
-    // TODO: also store info for EnclosingMethod
     internal class LocalClassInfo(val internalName: String)
 
     private val localClassInfo = mutableMapOf<IrAttributeContainer, LocalClassInfo>()
@@ -77,6 +72,8 @@ class JvmBackendContext(
         localClassInfo[container.attributeOwnerId] = value
     }
 
+    internal val customEnclosingFunction = mutableMapOf<IrAttributeContainer, IrFunction>()
+
     // TODO cache these at ClassCodegen level. Currently, sharing this map between classes in a module is required
     //      because IrSourceCompilerForInline constructs a new (Fake)ClassCodegen for every call to
     //      an inline function in the same module. Thus, if two inline functions happen to have the same name

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

@@ -145,6 +145,12 @@ internal val localDeclarationsPhase = makeIrFilePhase<CommonBackendContext>(
     prerequisite = setOf(callableReferencePhase, sharedVariablesPhase)
 )
 
+private val jvmLocalClassExtractionPhase = makeIrFilePhase(
+    ::JvmLocalClassPopupLowering,
+    name = "JvmLocalClassExtraction",
+    description = "Move local classes from field initializers and anonymous init blocks into the containing class"
+)
+
 private val defaultArgumentStubPhase = makeIrFilePhase(
     ::JvmDefaultArgumentStubGenerator,
     name = "DefaultArgumentsStubGenerator",
@@ -193,7 +199,9 @@ private val initializersPhase = makeIrFilePhase(
                 error("No anonymous initializers should remain at this stage")
             }
         })
-    })
+    }),
+    // Depends on local class extraction, because otherwise local classes in initializers will be copied into each constructor.
+    prerequisite = setOf(jvmLocalClassExtractionPhase)
 )
 
 private val returnableBlocksPhase = makeIrFilePhase(
@@ -248,6 +256,7 @@ private val jvmFilePhases =
         assertionPhase then
         returnableBlocksPhase then
         localDeclarationsPhase then
+        jvmLocalClassExtractionPhase then
 
         jvmOverloadsAnnotationPhase then
         jvmDefaultConstructorPhase then

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ClassCodegen.kt

@@ -122,7 +122,7 @@ open class ClassCodegen protected constructor(
 
         val nestedClasses = irClass.declarations.mapNotNull { declaration ->
             if (declaration is IrClass) {
-                ClassCodegen(declaration, context, this)
+                ClassCodegen(declaration, context, this, withinInline = withinInline)
             } else null
         }
 
@@ -411,9 +411,9 @@ open class ClassCodegen protected constructor(
         // or constructor, the name and type of the function is recorded as well.
         if (parentClassCodegen != null) {
             val outerClassName = parentClassCodegen.type.internalName
-            // TODO: LocalDeclarationsLowering could have moved this class out of its enclosing method.
-            if (parentFunction != null) {
-                val method = methodSignatureMapper.mapAsmMethod(parentFunction)
+            val enclosingFunction = context.customEnclosingFunction[irClass.attributeOwnerId] ?: parentFunction
+            if (enclosingFunction != null) {
+                val method = methodSignatureMapper.mapAsmMethod(enclosingFunction)
                 visitor.visitOuterClass(outerClassName, method.name, method.descriptor)
             } else if (irClass.isAnonymousObject) {
                 visitor.visitOuterClass(outerClassName, null, null)

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/irCodegenUtils.kt

@@ -35,7 +35,6 @@ import org.jetbrains.kotlin.load.java.JavaVisibilities
 import org.jetbrains.kotlin.load.java.JvmAbi
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.checkers.ExpectedActualDeclarationChecker
 import org.jetbrains.kotlin.resolve.inline.INLINE_ONLY_ANNOTATION_FQ_NAME
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmClassSignature
@@ -110,10 +109,12 @@ val IrType.isExtensionFunctionType: Boolean
     get() = isFunctionTypeOrSubtype() && hasAnnotation(KotlinBuiltIns.FQ_NAMES.extensionFunctionType)
 
 
-/* Borrowed from MemberCodegen.java */
+/* Borrowed with modifications from MemberCodegen.java */
 
 fun writeInnerClass(innerClass: IrClass, typeMapper: IrTypeMapper, context: JvmBackendContext, v: ClassBuilder) {
-    val outerClassInternalName = innerClass.parent.safeAs<IrClass>()?.let { typeMapper.classInternalName(it) }
+    val outerClassInternalName =
+        if (context.customEnclosingFunction[innerClass.attributeOwnerId] != null) null
+        else innerClass.parent.safeAs<IrClass>()?.let(typeMapper::classInternalName)
     val innerName = innerClass.name.takeUnless { it.isSpecial }?.asString()
     val innerClassInternalName = typeMapper.classInternalName(innerClass)
     v.visitInnerClass(innerClassInternalName, outerClassInternalName, innerName, innerClass.calculateInnerClassAccessFlags(context))

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

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.backend.jvm.lower
+
+import org.jetbrains.kotlin.backend.common.ScopeWithIr
+import org.jetbrains.kotlin.backend.common.lower.LocalClassPopupLowering
+import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
+import org.jetbrains.kotlin.ir.declarations.IrAnonymousInitializer
+import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrField
+import org.jetbrains.kotlin.ir.util.parentAsClass
+import org.jetbrains.kotlin.ir.util.primaryConstructor
+import org.jetbrains.kotlin.ir.util.render
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
+
+class JvmLocalClassPopupLowering(context: JvmBackendContext) : LocalClassPopupLowering(context) {
+    // On JVM, we only pop up local classes in field initializers and anonymous init blocks, so that InitializersLowering would not copy
+    // them to each constructor. (Moving all local classes is not possible because of cases where they use reified type parameters,
+    // or capture crossinline lambdas.)
+    // Upon moving such class, we record customEnclosingFunction for it to be the class constructor. This is needed because otherwise
+    // the class will not get any EnclosingMethod in the codegen later, since it won't be local anymore.
+    override fun shouldPopUp(klass: IrClass, currentScope: ScopeWithIr?): Boolean {
+        // On JVM, lambdas have package-private visibility after LocalDeclarationsLowering, so we have to check something else.
+        val isLocal = super.shouldPopUp(klass, currentScope) ||
+                klass.origin == JvmLoweredDeclarationOrigin.LAMBDA_IMPL ||
+                klass.origin == JvmLoweredDeclarationOrigin.FUNCTION_REFERENCE_IMPL ||
+                klass.origin == JvmLoweredDeclarationOrigin.GENERATED_PROPERTY_REFERENCE
+        if (!isLocal) return false
+
+        val container = when (val element = currentScope?.irElement) {
+            is IrAnonymousInitializer -> element.parentAsClass.takeUnless { element.isStatic }
+            is IrField -> element.parentAsClass.takeUnless { element.isStatic }
+            else -> null
+        } ?: return false
+
+        // In case there's no primary constructor, it's unclear which constructor should be the enclosing one, so we select the first.
+        (context as JvmBackendContext).customEnclosingFunction[klass.attributeOwnerId] =
+            container.primaryConstructor ?: container.declarations.firstIsInstanceOrNull()
+                    ?: error("Class in a non-static initializer found, but container has no constructors: ${container.render()}")
+
+        return true
+    }
+}

compiler/testData/codegen/box/reflection/enclosing/lambdaInPropertyDelegate.kt

@@ -0,0 +1,47 @@
+// TARGET_BACKEND: JVM
+// WITH_RUNTIME
+
+class C {
+    val f by foo {
+        {}
+    }
+}
+
+fun foo(f: () -> Any): Any = f()
+
+operator fun Any.getValue(thiz: Any?, metadata: Any?): Any = this
+
+fun box(): String {
+    // This is the class for lambda inside the `foo` call (`{}`)
+    val innerLambda = C().f.javaClass
+
+    val emInner = innerLambda.getEnclosingMethod()
+    if (emInner?.getName() != "invoke") return "Fail: incorrect enclosing method for inner lambda: $emInner"
+
+    val ecInner = innerLambda.getEnclosingClass()
+    if (ecInner?.getName() != "C\$f\$2") return "Fail: incorrect enclosing class for inner lambda: $ecInner"
+
+    val ectorInner = innerLambda.getEnclosingConstructor()
+    if (ectorInner != null) return "Fail: inner lambda should not have enclosing constructor: $ectorInner"
+
+    val dcInner = innerLambda.getDeclaringClass()
+    if (dcInner != null) return "Fail: inner lambda should not have declaring class: $dcInner"
+
+
+    // This is the class for lambda that is passed as an argument to `foo`
+    val outerLambda = ecInner
+
+    val emOuter = outerLambda.getEnclosingMethod()
+    if (emOuter != null) return "Fail: outer lambda should not have enclosing method: $emOuter"
+
+    val ecOuter = outerLambda.getEnclosingClass()
+    if (ecOuter?.getName() != "C") return "Fail: incorrect enclosing class for outer lambda: $ecOuter"
+
+    val ectorOuter = outerLambda.getEnclosingConstructor()
+    if (ectorOuter == null) return "Fail: outer lambda _should_ have enclosing constructor"
+
+    val dcOuter = outerLambda.getDeclaringClass()
+    if (dcOuter != null) return "Fail: outer lambda should not have declaring class: $dcOuter"
+
+    return "OK"
+}

compiler/testData/compileKotlinAgainstKotlin/metadataForMembersInLocalClassInInitializer.kt

@@ -0,0 +1,27 @@
+// TARGET_BACKEND: JVM
+// FILE: A.kt
+
+class A {
+    val o = object {
+        @JvmName("jvmGetO")
+        fun getO(): String = "O"
+    }
+
+    val k = object {
+        @get:JvmName("jvmGetK")
+        val k: String = "K"
+    }
+}
+
+// FILE: B.kt
+
+import kotlin.reflect.full.*
+
+fun box(): String {
+    val a = A()
+    val obj1 = a.o
+    val o = obj1::class.declaredMemberFunctions.single().call(obj1) as String
+    val obj2 = a.k
+    val k = obj2::class.declaredMemberProperties.single().call(obj2) as String
+    return o + k
+}

compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java

@@ -23105,6 +23105,11 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTest("compiler/testData/codegen/box/reflection/enclosing/lambdaInPackage.kt");
             }
 
+            @TestMetadata("lambdaInPropertyDelegate.kt")
+            public void testLambdaInPropertyDelegate() throws Exception {
+                runTest("compiler/testData/codegen/box/reflection/enclosing/lambdaInPropertyDelegate.kt");
+            }
+
             @TestMetadata("lambdaInPropertyGetter.kt")
             public void testLambdaInPropertyGetter() throws Exception {
                 runTest("compiler/testData/codegen/box/reflection/enclosing/lambdaInPropertyGetter.kt");

compiler/tests/org/jetbrains/kotlin/codegen/CompileKotlinAgainstKotlinTestGenerated.java

@@ -223,6 +223,11 @@ public class CompileKotlinAgainstKotlinTestGenerated extends AbstractCompileKotl
         runTest("compiler/testData/compileKotlinAgainstKotlin/kt21775.kt");
     }
 
+    @TestMetadata("metadataForMembersInLocalClassInInitializer.kt")
+    public void testMetadataForMembersInLocalClassInInitializer() throws Exception {
+        runTest("compiler/testData/compileKotlinAgainstKotlin/metadataForMembersInLocalClassInInitializer.kt");
+    }
+
     @TestMetadata("multifileClassInlineFunctionAccessingProperty.kt")
     public void testMultifileClassInlineFunctionAccessingProperty() throws Exception {
         runTest("compiler/testData/compileKotlinAgainstKotlin/multifileClassInlineFunctionAccessingProperty.kt");

compiler/tests/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -21922,6 +21922,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTest("compiler/testData/codegen/box/reflection/enclosing/lambdaInPackage.kt");
             }
 
+            @TestMetadata("lambdaInPropertyDelegate.kt")
+            public void testLambdaInPropertyDelegate() throws Exception {
+                runTest("compiler/testData/codegen/box/reflection/enclosing/lambdaInPropertyDelegate.kt");
+            }
+
             @TestMetadata("lambdaInPropertyGetter.kt")
             public void testLambdaInPropertyGetter() throws Exception {
                 runTest("compiler/testData/codegen/box/reflection/enclosing/lambdaInPropertyGetter.kt");

compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java

@@ -21670,6 +21670,11 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/reflection/enclosing/lambdaInPackage.kt");
             }
 
+            @TestMetadata("lambdaInPropertyDelegate.kt")
+            public void testLambdaInPropertyDelegate() throws Exception {
+                runTest("compiler/testData/codegen/box/reflection/enclosing/lambdaInPropertyDelegate.kt");
+            }
+
             @TestMetadata("lambdaInPropertyGetter.kt")
             public void testLambdaInPropertyGetter() throws Exception {
                 runTest("compiler/testData/codegen/box/reflection/enclosing/lambdaInPropertyGetter.kt");

compiler/tests/org/jetbrains/kotlin/codegen/ir/IrCompileKotlinAgainstKotlinTestGenerated.java

@@ -218,6 +218,11 @@ public class IrCompileKotlinAgainstKotlinTestGenerated extends AbstractIrCompile
         runTest("compiler/testData/compileKotlinAgainstKotlin/kt21775.kt");
     }
 
+    @TestMetadata("metadataForMembersInLocalClassInInitializer.kt")
+    public void testMetadataForMembersInLocalClassInInitializer() throws Exception {
+        runTest("compiler/testData/compileKotlinAgainstKotlin/metadataForMembersInLocalClassInInitializer.kt");
+    }
+
     @TestMetadata("multifileClassInlineFunctionAccessingProperty.kt")
     public void testMultifileClassInlineFunctionAccessingProperty() throws Exception {
         runTest("compiler/testData/compileKotlinAgainstKotlin/multifileClassInlineFunctionAccessingProperty.kt");