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, " +
-                        "IrTypeOperatorCall for SAM conversions) invokes `copyAttributes` on the transformed element."
+                        "Ensure that any lowering that transforms elements with local class info (classes, function references) " +
+                        "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.IrClass
-import org.jetbrains.kotlin.ir.declarations.copyAttributes
+import org.jetbrains.kotlin.ir.declarations.*
 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,104 +43,128 @@ internal val singleAbstractMethodPhase = makeIrFilePhase(
     prerequisite = setOf(localDeclarationsPhase)
 )
 
-class SingleAbstractMethodLowering(val context: CommonBackendContext) : ClassLoweringPass {
-    override fun lower(irClass: IrClass) {
-        val cachedImplementations = mutableMapOf<Pair<IrType, IrType>, IrClass>()
-        irClass.transformChildrenVoid(object : IrElementTransformerVoidWithContext() {
-            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.
-                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
+class SingleAbstractMethodLowering(val context: CommonBackendContext) : FileLoweringPass, IrElementTransformerVoidWithContext() {
+    private val cachedImplementations = mutableMapOf<IrType, IrClass>()
+    private var enclosingClass: IrClass? = null
 
-                    val implementation = cachedImplementations.getOrPut(superType to invokable.type) {
-                        createObjectProxy(superType, invokable.type, expression)
-                    }
+    // 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
+    }
 
-                    return if (superType.isNullable() && invokable.type.isNullable()) {
-                        irBlock(invokable, null, superType) {
-                            val invokableVariable = irTemporary(invokable)
-                            val instance = irCall(implementation.constructors.single()).apply {
-                                putValueArgument(0, irGet(invokableVariable))
-                            }
-                            irIfNull(superType, irGet(invokableVariable), irNull(), instance)
-                        }
-                    } else {
-                        irCall(implementation.constructors.single()).apply { putValueArgument(0, invokable) }
-                    }
-                }
+    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)
+        context.createIrBuilder(currentScope!!.scope.scopeOwnerSymbol).apply {
+            // Do not generate a wrapper class for null, it has no invoke() anyway.
+            if (invokable.isNullConst())
+                return invokable
+
+            // Coming from the frontend, every SAM interface is associated with exactly one function type
+            // (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)
             }
 
-            // 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 {
-                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" }
-
-                val subclass = buildClass {
-                    name = Name.special("<sam adapter for ${superClass.fqNameWhenAvailable}>")
-                    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")
-                    type = invokableType
-                    origin = subclass.origin
-                    visibility = Visibilities.PRIVATE
-                }
-
-                subclass.addConstructor {
-                    origin = subclass.origin
-                    isPrimary = true
-                }.apply {
-                    val parameter = addValueParameter {
-                        name = field.name
-                        type = field.type
-                        origin = subclass.origin
-                    }
-
-                    body = context.createIrBuilder(symbol).irBlockBody(startOffset, endOffset) {
-                        +irDelegatingConstructorCall(context.irBuiltIns.anyClass.owner.constructors.single())
-                        +irSetField(irGet(subclass.thisReceiver!!), field, irGet(parameter))
-                        +IrInstanceInitializerCallImpl(startOffset, endOffset, subclass.symbol, context.irBuiltIns.unitType)
+            return if (superType.isNullable() && invokable.type.isNullable()) {
+                irBlock(invokable, null, superType) {
+                    val invokableVariable = irTemporary(invokable)
+                    val instance = irCall(implementation.constructors.single()).apply {
+                        putValueArgument(0, irGet(invokableVariable))
                     }
+                    irIfNull(superType, irGet(invokableVariable), irNull(), instance)
                 }
-
-                val superMethod = superClass.functions.single { it.modality == Modality.ABSTRACT }
-                subclass.addFunction {
-                    name = superMethod.name
-                    returnType = superMethod.returnType
-                    visibility = superMethod.visibility
-                    origin = subclass.origin
-                }.apply {
-                    overriddenSymbols += superMethod.symbol
-                    dispatchReceiverParameter = subclass.thisReceiver!!.copyTo(this)
-                    superMethod.valueParameters.mapTo(valueParameters) { it.copyTo(this) }
-                    val invokableClass = invokableType.classifierOrFail.owner as IrClass
-                    body = context.createIrBuilder(symbol).run {
-                        irExprBody(irCall(invokableClass.functions.single { it.name == OperatorNameConventions.INVOKE }).apply {
-                            dispatchReceiver = irGetField(irGet(dispatchReceiverParameter!!), field)
-                            valueParameters.forEachIndexed { i, parameter -> putValueArgument(i, irGet(parameter)) }
-                        })
-                    }
-                }
-
-                return subclass
+            } else {
+                irCall(implementation.constructors.single()).apply { putValueArgument(0, invokable) }
             }
-        })
-        irClass.declarations += cachedImplementations.values
+        }
+    }
+
+    // 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): 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 = wrapperName
+            origin = JvmLoweredDeclarationOrigin.GENERATED_SAM_IMPLEMENTATION
+        }.apply {
+            createImplicitParameterDeclarationWithWrappedDescriptor()
+            // TODO convert all type parameters to upper bounds? See the kt11696 test.
+            superTypes += superType
+        }
+
+        val field = subclass.addField {
+            name = Name.identifier(SamWrapperCodegen.FUNCTION_FIELD_NAME)
+            type = invokableType
+            origin = subclass.origin
+            visibility = Visibilities.PRIVATE
+        }
+
+        subclass.addConstructor {
+            origin = subclass.origin
+            isPrimary = true
+        }.apply {
+            val parameter = addValueParameter {
+                name = field.name
+                type = field.type
+                origin = subclass.origin
+            }
+
+            body = context.createIrBuilder(symbol).irBlockBody(startOffset, endOffset) {
+                +irDelegatingConstructorCall(context.irBuiltIns.anyClass.owner.constructors.single())
+                +irSetField(irGet(subclass.thisReceiver!!), field, irGet(parameter))
+                +IrInstanceInitializerCallImpl(startOffset, endOffset, subclass.symbol, context.irBuiltIns.unitType)
+            }
+        }
+
+        val superMethod = superClass.functions.single { it.modality == Modality.ABSTRACT }
+        subclass.addFunction {
+            name = superMethod.name
+            returnType = superMethod.returnType
+            visibility = superMethod.visibility
+            origin = subclass.origin
+        }.apply {
+            overriddenSymbols += superMethod.symbol
+            dispatchReceiverParameter = subclass.thisReceiver!!.copyTo(this)
+            superMethod.valueParameters.mapTo(valueParameters) { it.copyTo(this) }
+            val invokableClass = invokableType.classifierOrFail.owner as IrClass
+            body = context.createIrBuilder(symbol).run {
+                irExprBody(irCall(invokableClass.functions.single { it.name == OperatorNameConventions.INVOKE }).apply {
+                    dispatchReceiver = irGetField(irGet(dispatchReceiverParameter!!), field)
+                    valueParameters.forEachIndexed { i, parameter -> putValueArgument(i, irGet(parameter)) }
+                })
+            }
+        }
+
+        return subclass
     }
 }

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