JVM IR: fix enclosing constructor for lambdas in inner classes

Before this change, we stored the enclosing constructor in a map in
JvmBackendContext before moving lambdas and local classes out of
initializer blocks. However, in case the lambda was declared in an inner
class, we stored a reference to the unlowered constructor of the inner
class, whose JVM signature is "()V" instead of the correct "(LOuter;)V".
Java reflection then threw exception if we tried to call
`getEnclosingConstructor()` on such class at runtime. Proguard finished
with errors for the same reason.

It turns out that we can just store the fact that the class has been
moved, and load the matching constructor in codegen, where everything is
already lowered and guaranteed to match with the signatures of the
actual generated declarations.

 #KT-41668 Fixed

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

@@ -91,7 +91,7 @@ class JvmBackendContext(
         localClassType[container.attributeOwnerId] = value
     }
 
-    internal val customEnclosingFunction = mutableMapOf<IrAttributeContainer, IrFunction>()
+    internal val isEnclosedInConstructor = mutableSetOf<IrAttributeContainer>()
 
     internal val classCodegens = mutableMapOf<IrClass, ClassCodegen>()
 

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

@@ -36,6 +36,7 @@ import org.jetbrains.kotlin.resolve.jvm.annotations.VOLATILE_ANNOTATION_FQ_NAME
 import org.jetbrains.kotlin.resolve.jvm.checkers.JvmSimpleNameBacktickChecker
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.*
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmClassSignature
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import org.jetbrains.org.objectweb.asm.*
 import org.jetbrains.org.objectweb.asm.commons.Method
@@ -360,7 +361,13 @@ abstract class ClassCodegen protected constructor(
         // the current class. If the current class is immediately enclosed by a method
         // or constructor, the name and type of the function is recorded as well.
         if (parentClassCodegen != null) {
-            val enclosingFunction = context.customEnclosingFunction[irClass.attributeOwnerId] ?: parentFunction
+            // In case there's no primary constructor, it's unclear which constructor should be the enclosing one, so we select the first.
+            val enclosingFunction = if (irClass.attributeOwnerId in context.isEnclosedInConstructor) {
+                val containerClass = parentClassCodegen.irClass
+                containerClass.primaryConstructor
+                    ?: containerClass.declarations.firstIsInstanceOrNull<IrConstructor>()
+                    ?: error("Class in a non-static initializer found, but container has no constructors: ${containerClass.render()}")
+            } else parentFunction
             if (enclosingFunction != null || irClass.isAnonymousObject) {
                 val method = enclosingFunction?.let(context.methodSignatureMapper::mapAsmMethod)
                 visitor.visitOuterClass(parentClassCodegen.type.internalName, method?.name, method?.descriptor)

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

@@ -113,7 +113,7 @@ val IrType.isExtensionFunctionType: Boolean
 
 fun writeInnerClass(innerClass: IrClass, typeMapper: IrTypeMapper, context: JvmBackendContext, v: ClassBuilder) {
     val outerClassInternalName =
-        if (context.customEnclosingFunction[innerClass.attributeOwnerId] != null) null
+        if (innerClass.attributeOwnerId in context.isEnclosedInConstructor) null
         else innerClass.parent.safeAs<IrClass>()?.let(typeMapper::classInternalName)
     val innerName = innerClass.name.takeUnless { it.isSpecial }?.asString()
     val innerClassInternalName = typeMapper.classInternalName(innerClass)

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

@@ -12,10 +12,6 @@ import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
 import org.jetbrains.kotlin.backend.jvm.ir.IrInlineReferenceLocator
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.IrFunctionReference
-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) {
     private val inlineLambdaToScope = mutableMapOf<IrFunction, IrDeclaration>()
@@ -49,17 +45,13 @@ class JvmLocalClassPopupLowering(context: JvmBackendContext) : LocalClassPopupLo
         while (parent is IrFunction) {
             parent = inlineLambdaToScope[parent] ?: break
         }
-        val container = when (parent) {
-            is IrAnonymousInitializer -> parent.parentAsClass.takeUnless { parent.isStatic }
-            is IrField -> parent.parentAsClass.takeUnless { parent.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<IrConstructor>()
-                    ?: error("Class in a non-static initializer found, but container has no constructors: ${container.render()}")
-
-        return true
+        if (parent is IrAnonymousInitializer && !parent.isStatic ||
+            parent is IrField && !parent.isStatic
+        ) {
+            (context as JvmBackendContext).isEnclosedInConstructor.add(klass.attributeOwnerId)
+            return true
+        }
+        return false
     }
 }

compiler/testData/codegen/bytecodeListing/inline/enclosingInfo/lambdaInInnerClassConstructor.kt

@@ -0,0 +1,6 @@
+class Outer {
+    inner class Inner {
+        // This test checks the EnclosingMethod attribute of this lambda class.
+        val lambda = {}
+    }
+}

compiler/testData/codegen/bytecodeListing/inline/enclosingInfo/lambdaInInnerClassConstructor.txt

@@ -0,0 +1,30 @@
+@kotlin.Metadata
+final class Outer$Inner$lambda$1 {
+    // source: 'lambdaInInnerClassConstructor.kt'
+    enclosing method Outer$Inner.<init>(LOuter;)V
+    public final static field INSTANCE: Outer$Inner$lambda$1
+    inner (anonymous) class Outer$Inner$lambda$1
+    static method <clinit>(): void
+    method <init>(): void
+    public synthetic bridge method invoke(): java.lang.Object
+    public final method invoke(): void
+    public final inner class Outer$Inner
+}
+
+@kotlin.Metadata
+public final class Outer$Inner {
+    // source: 'lambdaInInnerClassConstructor.kt'
+    private final @org.jetbrains.annotations.NotNull field lambda: kotlin.jvm.functions.Function0
+    synthetic final field this$0: Outer
+    inner (anonymous) class Outer$Inner$lambda$1
+    public method <init>(p0: Outer): void
+    public final @org.jetbrains.annotations.NotNull method getLambda(): kotlin.jvm.functions.Function0
+    public final inner class Outer$Inner
+}
+
+@kotlin.Metadata
+public final class Outer {
+    // source: 'lambdaInInnerClassConstructor.kt'
+    public method <init>(): void
+    public final inner class Outer$Inner
+}

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

@@ -474,6 +474,11 @@ public class BytecodeListingTestGenerated extends AbstractBytecodeListingTest {
                 runTest("compiler/testData/codegen/bytecodeListing/inline/enclosingInfo/lambdaInInitBlockNoPrimaryConstructor.kt");
             }
 
+            @TestMetadata("lambdaInInnerClassConstructor.kt")
+            public void testLambdaInInnerClassConstructor() throws Exception {
+                runTest("compiler/testData/codegen/bytecodeListing/inline/enclosingInfo/lambdaInInnerClassConstructor.kt");
+            }
+
             @TestMetadata("transformedConstructor.kt")
             public void testTransformedConstructor() throws Exception {
                 runTest("compiler/testData/codegen/bytecodeListing/inline/enclosingInfo/transformedConstructor.kt");

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

@@ -444,6 +444,11 @@ public class IrBytecodeListingTestGenerated extends AbstractIrBytecodeListingTes
                 runTest("compiler/testData/codegen/bytecodeListing/inline/enclosingInfo/lambdaInInitBlockNoPrimaryConstructor.kt");
             }
 
+            @TestMetadata("lambdaInInnerClassConstructor.kt")
+            public void testLambdaInInnerClassConstructor() throws Exception {
+                runTest("compiler/testData/codegen/bytecodeListing/inline/enclosingInfo/lambdaInInnerClassConstructor.kt");
+            }
+
             @TestMetadata("transformedConstructor.kt")
             public void testTransformedConstructor() throws Exception {
                 runTest("compiler/testData/codegen/bytecodeListing/inline/enclosingInfo/transformedConstructor.kt");