JVM IR: Cache SAM wrappers in top-level classes and use proper naming scheme

compiler/backend/src/org/jetbrains/kotlin/codegen/SamWrapperCodegen.java

@@ -45,7 +45,7 @@ import static org.jetbrains.kotlin.resolve.jvm.AsmTypes.OBJECT_TYPE;
 import static org.jetbrains.org.objectweb.asm.Opcodes.*;
 
 public class SamWrapperCodegen {
-    private static final String FUNCTION_FIELD_NAME = "function";
+    public static final String FUNCTION_FIELD_NAME = "function";
 
     private final GenerationState state;
     private final boolean isInsideInline;

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

@@ -146,8 +146,8 @@ class IrTypeMapper(private val context: JvmBackendContext) {
             }
             else -> error(
                 "Local class should have its name computed in InventNamesForLocalClasses: ${klass.fqNameWhenAvailable}\n" +
-                        "Ensure that any lowering that transforms elements with local class info (classes, function references, " +
+                        "Ensure that any lowering that transforms elements with local class info (classes, function references) " +
-                        "IrTypeOperatorCall for SAM conversions) invokes `copyAttributes` on the transformed element."
+                        "invokes `copyAttributes` on the transformed element."
             )
         }
     }

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

@@ -175,27 +175,6 @@ class InventNamesForLocalClasses(private val context: JvmBackendContext) : FileL
             declaration.acceptChildren(this, data.makeLocal())
         }
 
-        override fun visitTypeOperator(expression: IrTypeOperatorCall, data: Data) {
-            if (expression.operator == IrTypeOperator.SAM_CONVERSION) {
-                val invokable = expression.argument
-                // Function references (even those that are transformed into SAM wrappers) are handled in visitFunctionReference.
-                if (invokable !is IrFunctionReference && !(invokable is IrBlock && invokable.statements.last() is IrFunctionReference)) {
-                    val superClass = expression.typeOperandClassifier.owner as IrClass
-                    val superClassName = superClass.fqNameWhenAvailable!!.pathSegments().joinToString("_") { it.toString() }
-
-                    // TODO: consider dropping "sam\$$superClassName$0" from the name here.
-                    // It's only here to give resemblance of the name generated for SAM wrappers by the old backend. The exact logic
-                    // of the old backend is somewhat difficult to emulate consistently, see `SamWrapperCodegen.getWrapperName`.
-                    val internalName = inventName(Name.identifier("sam\$$superClassName$0"), data)
-                    context.putLocalClassInfo(expression, JvmBackendContext.LocalClassInfo(internalName))
-                    invokable.accept(this, data.withName(internalName))
-                    return
-                }
-            }
-
-            expression.acceptChildren(this, data)
-        }
-
         override fun visitElement(element: IrElement, data: Data) {
             element.acceptChildren(this, data)
         }

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

@@ -5,8 +5,8 @@
 
 package org.jetbrains.kotlin.backend.jvm.lower
 
-import org.jetbrains.kotlin.backend.common.ClassLoweringPass
 import org.jetbrains.kotlin.backend.common.CommonBackendContext
+import org.jetbrains.kotlin.backend.common.FileLoweringPass
 import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
 import org.jetbrains.kotlin.backend.common.ir.copyTo
 import org.jetbrains.kotlin.backend.common.ir.createImplicitParameterDeclarationWithWrappedDescriptor
@@ -15,14 +15,13 @@ import org.jetbrains.kotlin.backend.common.lower.irBlock
 import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
 import org.jetbrains.kotlin.backend.jvm.localDeclarationsPhase
+import org.jetbrains.kotlin.codegen.SamWrapperCodegen
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.builders.declarations.*
-import org.jetbrains.kotlin.ir.declarations.IrAttributeContainer
+import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.declarations.IrClass
-import org.jetbrains.kotlin.ir.declarations.copyAttributes
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrTypeOperator
 import org.jetbrains.kotlin.ir.expressions.IrTypeOperatorCall
@@ -34,7 +33,6 @@ import org.jetbrains.kotlin.ir.util.constructors
 import org.jetbrains.kotlin.ir.util.fqNameWhenAvailable
 import org.jetbrains.kotlin.ir.util.functions
 import org.jetbrains.kotlin.ir.util.isNullConst
-import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.util.OperatorNameConventions
 
@@ -45,24 +43,47 @@ internal val singleAbstractMethodPhase = makeIrFilePhase(
     prerequisite = setOf(localDeclarationsPhase)
 )
 
-class SingleAbstractMethodLowering(val context: CommonBackendContext) : ClassLoweringPass {
+class SingleAbstractMethodLowering(val context: CommonBackendContext) : FileLoweringPass, IrElementTransformerVoidWithContext() {
-    override fun lower(irClass: IrClass) {
+    private val cachedImplementations = mutableMapOf<IrType, IrClass>()
-        val cachedImplementations = mutableMapOf<Pair<IrType, IrType>, IrClass>()
+    private var enclosingClass: IrClass? = null
-        irClass.transformChildrenVoid(object : IrElementTransformerVoidWithContext() {
+
+    // SAM wrappers are cached, either in the file class (if it exists), or in the top-level enclosing class.
+    private inline fun <R> withCacheFor(irClass: IrClass?, block: () -> R): R {
+        enclosingClass = enclosingClass ?: irClass
+        val result = block()
+        if (irClass != null && enclosingClass === irClass) {
+            cachedImplementations.values.mapTo(irClass.declarations) {
+                it.parent = irClass
+                it
+            }
+            cachedImplementations.clear()
+        }
+        return result
+    }
+
+    override fun lower(irFile: IrFile) {
+        val fileClass = irFile.declarations.filterIsInstance<IrClass>().find { it.origin == IrDeclarationOrigin.FILE_CLASS }
+        withCacheFor(fileClass) { irFile.transformChildrenVoid() }
+    }
+
+    override fun visitClassNew(declaration: IrClass) =
+        withCacheFor(declaration) { super.visitClassNew(declaration) }
+
     override fun visitTypeOperator(expression: IrTypeOperatorCall): IrExpression {
         if (expression.operator != IrTypeOperator.SAM_CONVERSION)
             return super.visitTypeOperator(expression)
         val superType = expression.typeOperand
         val invokable = expression.argument.transform(this, null)
-                // Running in the class context, so empty scope means this is a field/anonymous initializer.
+        context.createIrBuilder(currentScope!!.scope.scopeOwnerSymbol).apply {
-                val scopeOwnerSymbol = currentScope?.scope?.scopeOwnerSymbol ?: irClass.symbol
-                context.createIrBuilder(scopeOwnerSymbol).apply {
             // Do not generate a wrapper class for null, it has no invoke() anyway.
             if (invokable.isNullConst())
                 return invokable
 
-                    val implementation = cachedImplementations.getOrPut(superType to invokable.type) {
+            // Coming from the frontend, every SAM interface is associated with exactly one function type
-                        createObjectProxy(superType, invokable.type, expression)
+            // (see SamType.getKotlinFunctionType). That's why we can cache implementations just based on
+            // the superType. Type parameters should have been erased before we get here.
+            val implementation = cachedImplementations.getOrPut(superType) {
+                createObjectProxy(superType, invokable.type)
             }
 
             return if (superType.isNullable() && invokable.type.isNullable()) {
@@ -81,24 +102,28 @@ class SingleAbstractMethodLowering(val context: CommonBackendContext) : ClassLow
 
     // Construct a class that wraps an invokable object into an implementation of an interface:
     //     class sam$n(private val invokable: F) : Interface { override fun method(...) = invokable(...) }
-            private fun createObjectProxy(superType: IrType, invokableType: IrType, attributeOwner: IrAttributeContainer): IrClass {
+    private fun createObjectProxy(superType: IrType, invokableType: IrType): IrClass {
         val superClass = superType.classifierOrFail.owner as IrClass
         // The language documentation prohibits casting lambdas to classes, but if it was allowed,
         // the `irDelegatingConstructorCall` in the constructor below would need to be modified.
         assert(superClass.kind == ClassKind.INTERFACE) { "SAM conversion to an abstract class not allowed" }
 
+        // TODO: In the scope of an inline function, we need to generate a *public* sam wrapper with a name of the form
+        //       sam$i$superClassFqName$0 (note the additiona $i compared to the name below)
+        val superFqName = superClass.fqNameWhenAvailable!!.asString().replace('.', '_')
+        val wrapperName = Name.identifier("sam\$$superFqName${SamWrapperCodegen.SAM_WRAPPER_SUFFIX}")
+
         val subclass = buildClass {
-                    name = Name.special("<sam adapter for ${superClass.fqNameWhenAvailable}>")
+            name = wrapperName
             origin = JvmLoweredDeclarationOrigin.GENERATED_SAM_IMPLEMENTATION
         }.apply {
             createImplicitParameterDeclarationWithWrappedDescriptor()
-                    parent = irClass
             // TODO convert all type parameters to upper bounds? See the kt11696 test.
             superTypes += superType
-                }.copyAttributes(attributeOwner)
+        }
 
         val field = subclass.addField {
-                    name = Name.identifier("arg0")
+            name = Name.identifier(SamWrapperCodegen.FUNCTION_FIELD_NAME)
             type = invokableType
             origin = subclass.origin
             visibility = Visibilities.PRIVATE
@@ -142,7 +167,4 @@ class SingleAbstractMethodLowering(val context: CommonBackendContext) : ClassLow
 
         return subclass
     }
-        })
-        irClass.declarations += cachedImplementations.values
-    }
 }

compiler/testData/codegen/box/sam/kt17091.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 // TARGET_BACKEND: JVM
 // WITH_RUNTIME
 // FILE: Foo.kt

compiler/testData/codegen/box/sam/kt17091_2.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 // TARGET_BACKEND: JVM
 // WITH_RUNTIME
 // FILE: Foo.kt

compiler/testData/codegen/box/sam/kt17091_3.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 // TARGET_BACKEND: JVM
 // WITH_RUNTIME
 // FILE: Foo.kt

compiler/testData/codegen/box/sam/kt17091_4.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 // TARGET_BACKEND: JVM
 // WITH_RUNTIME
 // FILE: Foo.kt

compiler/testData/codegen/box/sam/kt22906.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 // TARGET_BACKEND: JVM
 // WITH_RUNTIME
 // FILE: kt22906_1.kt

compiler/testData/codegen/box/sam/kt22906_2.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 // TARGET_BACKEND: JVM
 // WITH_RUNTIME
 // FILE: kt22906_1.kt