JVM_IR: mark inlined callable references with a different origin

Makes them easier to find.

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

@@ -8,7 +8,7 @@ package org.jetbrains.kotlin.backend.jvm.lower
 import org.jetbrains.kotlin.backend.common.FileLoweringPass
 import org.jetbrains.kotlin.backend.common.phaser.makeIrModulePhase
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
-import org.jetbrains.kotlin.backend.jvm.codegen.isInlineIrExpression
+import org.jetbrains.kotlin.backend.jvm.JvmLoweredStatementOrigin
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.IrFunctionReference
@@ -41,8 +41,8 @@ private class BytecodeInliningPreparationLowering(val context: JvmBackendContext
 
             override fun visitFunctionReference(expression: IrFunctionReference, data: String) {
                 // Remove inline lambdas from their declaration parents. They should not appear in the output
-                // bytecode in non-inlined form.
-                if (expression.origin.isInlineIrExpression()) {
+                // bytecode in non-inlined form. TODO: this does not remove unused lambdas; need to mark function origin?
+                if (expression.origin == JvmLoweredStatementOrigin.INLINE_LAMBDA) {
                     loweredLambdasToDelete.add(expression.symbol.owner)
                 }
                 super.visitFunctionReference(expression, data)

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

@@ -37,44 +37,30 @@ internal class FakeInliningLocalVariablesLowering(val context: JvmBackendContext
     override fun visitFunction(declaration: IrFunction, data: IrDeclaration?) {
         super.visitFunction(declaration, data)
         if (declaration.isInline && !declaration.origin.isSynthetic && declaration.body != null && !declaration.isInlineOnly()) {
-            declaration.addFakeInliningLocalVariables()
+            val currentFunctionName = context.methodSignatureMapper.mapFunctionName(declaration)
+            val localName = "${JvmAbi.LOCAL_VARIABLE_NAME_PREFIX_INLINE_FUNCTION}$currentFunctionName"
+            declaration.addFakeLocalVariable(localName)
         }
     }
 
     override fun visitInlineLambda(argument: IrFunctionReference, callee: IrFunction, parameter: IrValueParameter, scope: IrDeclaration) {
         val lambda = argument.symbol.owner
-        if (lambda.origin == IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA) {
-            val argumentToFunctionName = context.methodSignatureMapper.mapFunctionName(callee)
-            val lambdaReferenceName = context.getLocalClassType(argument)!!.internalName.substringAfterLast("/")
-            val localName = "${JvmAbi.LOCAL_VARIABLE_NAME_PREFIX_INLINE_ARGUMENT}-$argumentToFunctionName-$lambdaReferenceName"
-            lambda.addFakeLocalVariable(localName)
-        }
-    }
-
-    private fun IrFunction.addFakeInliningLocalVariables() {
-        val currentFunctionName = context.methodSignatureMapper.mapFunctionName(this)
-        val localName = "${JvmAbi.LOCAL_VARIABLE_NAME_PREFIX_INLINE_FUNCTION}$currentFunctionName"
-        addFakeLocalVariable(localName)
+        val argumentToFunctionName = context.methodSignatureMapper.mapFunctionName(callee)
+        val lambdaReferenceName = context.getLocalClassType(argument)!!.internalName.substringAfterLast("/")
+        val localName = "${JvmAbi.LOCAL_VARIABLE_NAME_PREFIX_INLINE_ARGUMENT}-$argumentToFunctionName-$lambdaReferenceName"
+        lambda.addFakeLocalVariable(localName)
     }
 
     private fun IrFunction.addFakeLocalVariable(name: String) {
-        val oldBody = body
-        context.createIrBuilder(symbol).run {
-            body = irBlockBody {
-                // Create temporary variable, but make sure it's origin is `DEFINED` so that
-                // it will materialize in the code.
-                // Also, do not forget to remove $$forInline suffix, otherwise, IDE will not be able to navigate to inline function.
-                createTmpVariable(irInt(0), name.removeSuffix(FOR_INLINE_SUFFIX), origin = IrDeclarationOrigin.DEFINED)
-                when (oldBody) {
-                    is IrExpressionBody -> {
-                        +irReturn(oldBody.expression)
-                    }
-                    is IrBlockBody ->
-                        oldBody.statements.forEach { +it }
-                    else -> {
-                        throw AssertionError("Unexpected body:\n${this@addFakeLocalVariable.dump()}")
-                    }
-                }
+        body = context.createIrBuilder(symbol).irBlockBody {
+            // Create temporary variable, but make sure it's origin is `DEFINED` so that
+            // it will materialize in the code.
+            // Also, do not forget to remove $$forInline suffix, otherwise, IDE will not be able to navigate to inline function.
+            createTmpVariable(irInt(0), name.removeSuffix(FOR_INLINE_SUFFIX), origin = IrDeclarationOrigin.DEFINED)
+            when (val oldBody = body) {
+                is IrExpressionBody -> +irReturn(oldBody.expression)
+                is IrBlockBody -> oldBody.statements.forEach { +it }
+                else -> throw AssertionError("Unexpected body:\n${this@addFakeLocalVariable.dump()}")
             }
         }
     }

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

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.backend.common.lower.parents
 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.JvmLoweredStatementOrigin
 import org.jetbrains.kotlin.backend.jvm.JvmSymbols
 import org.jetbrains.kotlin.backend.jvm.ir.*
 import org.jetbrains.kotlin.backend.jvm.lower.indy.LambdaMetafactoryArguments
@@ -45,42 +46,26 @@ internal val functionReferencePhase = makeIrFilePhase(
 )
 
 internal class FunctionReferenceLowering(private val context: JvmBackendContext) : FileLoweringPass, IrElementTransformerVoidWithContext() {
-    // This pass ignores function references used as inline arguments. `InlineCallableReferenceToLambdaPhase`
-    // converts them into lambdas instead, so that after inlining there is only a direct call left, with no
-    // function reference classes needed.
-    private val ignoredFunctionReferences = mutableSetOf<IrCallableReference<*>>()
-
     private val crossinlineLambdas = HashSet<IrSimpleFunction>()
 
     private val IrFunctionReference.isIgnored: Boolean
-        get() = (!type.isFunctionOrKFunction() && !isSuspendFunctionReference()) || ignoredFunctionReferences.contains(this)
+        get() = (!type.isFunctionOrKFunction() && !isSuspendFunctionReference()) || origin == JvmLoweredStatementOrigin.INLINE_LAMBDA
 
-    // `suspend` function references are the same as non-`suspend` ones, just with a `suspend` invoke;
-    // however, suspending lambdas require different generation implemented in AddContinuationLowering
+    // `suspend` function references are the same as non-`suspend` ones, just with an extra continuation parameter;
+    // however, suspending lambdas require different generation implemented in SuspendLambdaLowering
     // because they are also their own continuation classes.
     // TODO: Currently, origin of callable references explicitly written in source code is null. Do we need to create one?
     private fun IrFunctionReference.isSuspendFunctionReference(): Boolean = isSuspend &&
             (origin == null || origin == IrStatementOrigin.ADAPTED_FUNCTION_REFERENCE || origin == IrStatementOrigin.SUSPEND_CONVERSION)
 
     override fun lower(irFile: IrFile) {
-        irFile.accept(
-            object : IrInlineReferenceLocator(context) {
-                override fun visitInlineLambda(
-                    argument: IrFunctionReference,
-                    callee: IrFunction,
-                    parameter: IrValueParameter,
-                    scope: IrDeclaration
-                ) {
-                    ignoredFunctionReferences.add(argument)
-                    val argumentFun = argument.symbol.owner
-                    if (parameter.isCrossinline && argumentFun is IrSimpleFunction) {
-                        crossinlineLambdas.add(argumentFun)
-                    }
-                }
-            },
-            null
-        )
+        irFile.findInlineLambdas(context) { argument, _, parameter, _ ->
+            if (parameter.isCrossinline) {
+                crossinlineLambdas.add(argument.symbol.owner as IrSimpleFunction)
+            }
+        }
         irFile.transformChildrenVoid(this)
+        crossinlineLambdas.clear()
     }
 
     private val shouldGenerateIndySamConversions =

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

@@ -6,15 +6,16 @@
 package org.jetbrains.kotlin.backend.jvm.lower
 
 import org.jetbrains.kotlin.backend.common.FileLoweringPass
-import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
 import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
-import org.jetbrains.kotlin.backend.common.lower.irBlock
 import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
-import org.jetbrains.kotlin.backend.jvm.ir.IrInlineReferenceLocator
+import org.jetbrains.kotlin.backend.jvm.JvmLoweredStatementOrigin
+import org.jetbrains.kotlin.backend.jvm.codegen.isInlineFunctionCall
 import org.jetbrains.kotlin.backend.jvm.ir.createJvmIrBuilder
 import org.jetbrains.kotlin.backend.jvm.ir.irArray
+import org.jetbrains.kotlin.backend.jvm.ir.isInlineParameter
 import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
+import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.builders.declarations.addExtensionReceiver
 import org.jetbrains.kotlin.ir.builders.declarations.addValueParameter
@@ -25,13 +26,13 @@ import org.jetbrains.kotlin.ir.expressions.impl.IrFunctionReferenceImpl
 import org.jetbrains.kotlin.ir.types.IrSimpleType
 import org.jetbrains.kotlin.ir.types.IrTypeProjection
 import org.jetbrains.kotlin.ir.util.*
-import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
+import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
 import org.jetbrains.kotlin.name.Name
 
 internal val inlineCallableReferenceToLambdaPhase = makeIrFilePhase(
     ::InlineCallableReferenceToLambdaPhase,
     name = "InlineCallableReferenceToLambdaPhase",
-    description = "Transform callable reference to inline lambda"
+    description = "Transform callable reference to inline lambdas, mark inline lambdas for later passes"
 )
 
 // This lowering transforms CR passed to inline function to lambda which would be inlined
@@ -43,158 +44,125 @@ internal val inlineCallableReferenceToLambdaPhase = makeIrFilePhase(
 //      foo(::smth) -> foo { a -> smth(a) }
 //
 internal class InlineCallableReferenceToLambdaPhase(val context: JvmBackendContext) : FileLoweringPass {
-    override fun lower(irFile: IrFile) {
-        val inlinableReferences = mutableSetOf<IrCallableReference<*>>()
-        irFile.accept(object : IrInlineReferenceLocator(context) {
-            override fun visitInlineReference(argument: IrCallableReference<*>) {
-                inlinableReferences.add(argument)
-            }
-
-            override fun visitInlineLambda(
-                argument: IrFunctionReference, callee: IrFunction, parameter: IrValueParameter, scope: IrDeclaration
-            ) {
-                // Obviously needs no extra wrapping.
-            }
-        }, null)
-        irFile.transformChildrenVoid(InlineCallableReferenceToLambdaTransformer(context, inlinableReferences))
-    }
+    override fun lower(irFile: IrFile) =
+        irFile.accept(InlineCallableReferenceToLambdaVisitor(context), null)
 }
 
 const val STUB_FOR_INLINING = "stub_for_inlining"
 
-private class InlineCallableReferenceToLambdaTransformer(
-    val context: JvmBackendContext,
-    val inlinableReferences: Set<IrCallableReference<*>>
-) : IrElementTransformerVoidWithContext() {
-    override fun visitFunctionReference(expression: IrFunctionReference): IrExpression {
-        expression.transformChildrenVoid(this)
-        if (expression !in inlinableReferences) return expression
-        return expandInlineFunctionReferenceToLambda(expression, expression.symbol.owner)
-    }
+private class InlineCallableReferenceToLambdaVisitor(val context: JvmBackendContext) : IrElementVisitor<Unit, IrDeclaration?> {
+    override fun visitElement(element: IrElement, data: IrDeclaration?) =
+        element.acceptChildren(this, element as? IrDeclaration ?: data)
 
-    override fun visitPropertyReference(expression: IrPropertyReference): IrExpression {
-        expression.transformChildrenVoid(this)
-        if (expression !in inlinableReferences) return expression
-
-        return if (expression.field?.owner == null) {
-            // Use getter if field is absent ...
-            expandInlineFunctionReferenceToLambda(expression, expression.getter!!.owner)
-        } else {
-            // ... else use field itself
-            expandInlineFieldReferenceToLambda(expression, expression.field!!.owner)
+    override fun visitFunctionAccess(expression: IrFunctionAccessExpression, data: IrDeclaration?) {
+        expression.acceptChildren(this, data)
+        val function = expression.symbol.owner
+        if (function.isInlineFunctionCall(context)) {
+            for (parameter in function.valueParameters) {
+                if (parameter.isInlineParameter()) {
+                    expression.putValueArgument(parameter.index, expression.getValueArgument(parameter.index)?.transform(data))
+                }
+            }
         }
     }
 
-    private fun expandInlineFieldReferenceToLambda(expression: IrPropertyReference, field: IrField): IrExpression {
-        val irBuilder = context.createJvmIrBuilder(currentScope!!.scope.scopeOwnerSymbol, expression.startOffset, expression.endOffset)
-        return irBuilder.irBlock(expression, IrStatementOrigin.LAMBDA) {
-            val boundReceiver = expression.dispatchReceiver ?: expression.extensionReceiver
-            val function = context.irFactory.buildFun {
-                setSourceRange(expression)
-                origin = IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA
-                name = Name.identifier("stub_for_inline")
-                visibility = DescriptorVisibilities.LOCAL
-                returnType = field.type
-                isSuspend = false
-            }.apply {
-                parent = currentDeclarationParent ?: error("No current declaration parent at ${expression.dump()}")
-                val receiver = when {
+    private fun IrExpression.transform(scope: IrDeclaration?) = when {
+        this is IrBlock && origin.isInlinable -> apply {
+            // Already a lambda or similar, just mark it with an origin.
+            statements[statements.size - 1] =
+                (statements[statements.size - 1] as IrFunctionReference).replaceOrigin(JvmLoweredStatementOrigin.INLINE_LAMBDA)
+        }
+
+        this is IrFunctionReference -> // ::function -> { args... -> function(args...) }
+            wrapFunction(symbol.owner).toLambda(this, scope!!)
+
+        this is IrPropertyReference -> // ::property -> { receiver -> receiver.property }; prefer direct field access if allowed.
+            (if (field != null) wrapField(field!!.owner) else wrapFunction(getter!!.owner)).toLambda(this, scope!!)
+
+        else -> this // not an inline argument
+    }
+
+    private fun IrPropertyReference.wrapField(field: IrField): IrSimpleFunction =
+        context.irFactory.buildFun {
+            setSourceRange(this@wrapField)
+            origin = IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA
+            name = Name.identifier(STUB_FOR_INLINING)
+            visibility = DescriptorVisibilities.LOCAL
+            returnType = field.type
+        }.apply {
+            body = context.createIrBuilder(symbol).run {
+                val boundReceiver = dispatchReceiver ?: extensionReceiver
+                val fieldReceiver = when {
                     field.isStatic -> null
                     boundReceiver != null -> irGet(addExtensionReceiver(boundReceiver.type))
                     else -> irGet(addValueParameter("receiver", field.parentAsClass.defaultType))
                 }
-                body = this@InlineCallableReferenceToLambdaTransformer.context.createIrBuilder(symbol).run {
-                    irExprBody(irGetField(receiver, field))
-                }
-            }
-
-            +function
-            +IrFunctionReferenceImpl.fromSymbolOwner(
-                expression.startOffset,
-                expression.endOffset,
-                expression.type,
-                function.symbol,
-                typeArgumentsCount = 0,
-                reflectionTarget = null,
-                origin = IrStatementOrigin.LAMBDA
-            ).apply {
-                copyAttributes(expression)
-                extensionReceiver = boundReceiver
+                irExprBody(irGetField(fieldReceiver, field))
             }
         }
-    }
 
-    private fun expandInlineFunctionReferenceToLambda(expression: IrCallableReference<*>, referencedFunction: IrFunction): IrExpression {
-        val irBuilder = context.createJvmIrBuilder(currentScope!!.scope.scopeOwnerSymbol, expression.startOffset, expression.endOffset)
-        return irBuilder.irBlock(expression, IrStatementOrigin.LAMBDA) {
-            val (receiverParameter, receiverValue) = expression.getArgumentsWithIr().singleOrNull() ?: (null to null)
-            val kFunctionArguments = (expression.type as IrSimpleType).arguments.map { (it as IrTypeProjection).type }
-            val argumentTypes = kFunctionArguments.dropLast(1)
-            val resultType = kFunctionArguments.last()
-
-            val function = context.irFactory.buildFun {
-                setSourceRange(expression)
-                origin = IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA
-                name = Name.identifier(STUB_FOR_INLINING)
-                visibility = DescriptorVisibilities.LOCAL
-                returnType = resultType
-                isSuspend = referencedFunction.isSuspend
-            }.apply {
-                parent = currentDeclarationParent!!
-                if (receiverValue != null) {
-                    addExtensionReceiver(receiverValue.type)
+    private fun IrCallableReference<*>.wrapFunction(referencedFunction: IrFunction): IrSimpleFunction =
+        context.irFactory.buildFun {
+            setSourceRange(this@wrapFunction)
+            origin = IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA
+            name = Name.identifier(STUB_FOR_INLINING)
+            visibility = DescriptorVisibilities.LOCAL
+            returnType = ((type as IrSimpleType).arguments.last() as IrTypeProjection).type
+            isSuspend = referencedFunction.isSuspend
+        }.apply {
+            body = context.createJvmIrBuilder(symbol, startOffset, endOffset).run {
+                // TODO: could there be a star projection here?
+                val argumentTypes = (type as IrSimpleType).arguments.dropLast(1).map { (it as IrTypeProjection).type }
+                val boundReceiver = dispatchReceiver ?: extensionReceiver
+                val boundReceiverParameter = when {
+                    dispatchReceiver != null -> referencedFunction.dispatchReceiverParameter
+                    extensionReceiver != null -> referencedFunction.extensionReceiverParameter
+                    else -> null
                 }
-                for ((index, argumentType) in argumentTypes.withIndex()) {
-                    addValueParameter {
-                        name = Name.identifier("p$index")
-                        type = argumentType
-                    }
-                }
-
-                body = this@InlineCallableReferenceToLambdaTransformer.context.createJvmIrBuilder(
-                    symbol,
-                    expression.startOffset,
-                    expression.endOffset
-                ).run {
-                    irExprBody(irCall(referencedFunction.symbol, resultType).apply {
-                        symbol.owner.allTypeParameters.forEach {
-                            putTypeArgument(it.index, expression.getTypeArgument(it.index))
-                        }
-
-                        var unboundIndex = 0
-                        for (parameter in referencedFunction.explicitParameters) {
-                            when {
-                                receiverParameter == parameter ->
-                                    irGet(extensionReceiverParameter!!)
-                                parameter.isVararg && unboundIndex < argumentTypes.size && parameter.type == valueParameters[unboundIndex].type ->
-                                    irGet(valueParameters[unboundIndex++])
-                                parameter.isVararg && (unboundIndex < argumentTypes.size || !parameter.hasDefaultValue()) ->
-                                    irArray(parameter.type) {
-                                        (unboundIndex until argumentTypes.size).forEach { +irGet(valueParameters[unboundIndex++]) }
+                irExprBody(irCall(referencedFunction.symbol, returnType).apply {
+                    copyTypeArgumentsFrom(this@wrapFunction)
+                    for (parameter in referencedFunction.explicitParameters) {
+                        val next = valueParameters.size
+                        when {
+                            boundReceiverParameter == parameter ->
+                                irGet(addExtensionReceiver(boundReceiver!!.type))
+                            parameter.isVararg && next < argumentTypes.size && parameter.type == argumentTypes[next] ->
+                                irGet(addValueParameter("p$next", argumentTypes[next]))
+                            parameter.isVararg && (next < argumentTypes.size || !parameter.hasDefaultValue()) ->
+                                irArray(parameter.type) {
+                                    for (i in next until argumentTypes.size) {
+                                        +irGet(addValueParameter("p$i", argumentTypes[i]))
                                     }
-                                unboundIndex >= argumentTypes.size ->
-                                    null
-                                else ->
-                                    irGet(valueParameters[unboundIndex++])
-                            }?.let { putArgument(referencedFunction, parameter, it) }
-                        }
-                    })
-                }
-            }
-
-            +function
-            +IrFunctionReferenceImpl.fromSymbolOwner(
-                expression.startOffset,
-                expression.endOffset,
-                expression.type,
-                function.symbol,
-                typeArgumentsCount = function.typeParameters.size,
-                reflectionTarget = null,
-                origin = IrStatementOrigin.LAMBDA
-            ).apply {
-                copyAttributes(expression)
-                extensionReceiver = receiverValue
+                                }
+                            next >= argumentTypes.size ->
+                                null
+                            else ->
+                                irGet(addValueParameter("p$next", argumentTypes[next]))
+                        }?.let { putArgument(referencedFunction, parameter, it) }
+                    }
+                })
+            }
+        }
+
+    private fun IrSimpleFunction.toLambda(original: IrCallableReference<*>, scope: IrDeclaration) =
+        context.createIrBuilder(scope.symbol).irBlock(startOffset, endOffset, IrStatementOrigin.LAMBDA) {
+            this@toLambda.parent = parent
+            +this@toLambda
+            +IrFunctionReferenceImpl.fromSymbolOwner(
+                startOffset, endOffset, original.type, symbol, typeArgumentsCount = 0, reflectionTarget = null,
+                origin = JvmLoweredStatementOrigin.INLINE_LAMBDA
+            ).apply {
+                copyAttributes(original)
+                extensionReceiver = original.dispatchReceiver ?: original.extensionReceiver
             }
         }
-    }
 }
+
+private val IrStatementOrigin?.isInlinable: Boolean
+    get() = isLambda || this == IrStatementOrigin.ADAPTED_FUNCTION_REFERENCE || this == IrStatementOrigin.SUSPEND_CONVERSION
+
+private fun IrFunctionReference.replaceOrigin(origin: IrStatementOrigin): IrFunctionReference =
+    IrFunctionReferenceImpl(startOffset, endOffset, type, symbol, typeArgumentsCount, valueArgumentsCount, reflectionTarget, origin).also {
+        it.copyAttributes(this)
+        it.copyTypeAndValueArgumentsFrom(this)
+    }

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

@@ -8,23 +8,17 @@ 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.backend.jvm.ir.IrInlineReferenceLocator
+import org.jetbrains.kotlin.backend.jvm.ir.findInlineLambdas
 import org.jetbrains.kotlin.backend.jvm.isGeneratedLambdaClass
 import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.expressions.IrFunctionReference
 
 class JvmLocalClassPopupLowering(context: JvmBackendContext) : LocalClassPopupLowering(context) {
     private val inlineLambdaToScope = mutableMapOf<IrFunction, IrDeclaration>()
 
     override fun lower(irFile: IrFile) {
-        irFile.accept(object : IrInlineReferenceLocator(context as JvmBackendContext) {
-            override fun visitInlineLambda(
-                argument: IrFunctionReference, callee: IrFunction, parameter: IrValueParameter, scope: IrDeclaration
-            ) {
-                inlineLambdaToScope[argument.symbol.owner] = scope
-            }
-        }, null)
+        irFile.findInlineLambdas(context as JvmBackendContext) { argument, _, _, scope ->
+            inlineLambdaToScope[argument.symbol.owner] = scope
+        }
         super.lower(irFile)
         inlineLambdaToScope.clear()
     }

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

@@ -15,7 +15,6 @@ 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.codegen.isReadOfCrossinline
-import org.jetbrains.kotlin.backend.jvm.ir.IrInlineReferenceLocator
 import org.jetbrains.kotlin.backend.jvm.ir.hasChild
 import org.jetbrains.kotlin.codegen.coroutines.COROUTINE_LABEL_FIELD_NAME
 import org.jetbrains.kotlin.codegen.coroutines.INVOKE_SUSPEND_METHOD_NAME
@@ -99,11 +98,10 @@ internal abstract class SuspendLoweringUtils(protected val context: JvmBackendCo
 
 private class SuspendLambdaLowering(context: JvmBackendContext) : SuspendLoweringUtils(context), FileLoweringPass {
     override fun lower(irFile: IrFile) {
-        val inlineReferences = IrInlineReferenceLocator.scan(context, irFile)
         irFile.transformChildrenVoid(object : IrElementTransformerVoidWithContext() {
             override fun visitBlock(expression: IrBlock): IrExpression {
                 val reference = expression.statements.lastOrNull() as? IrFunctionReference ?: return super.visitBlock(expression)
-                if (reference.isSuspend && reference.origin.isLambda && reference !in inlineReferences) {
+                if (reference.isSuspend && reference.origin.isLambda) {
                     assert(expression.statements.size == 2 && expression.statements[0] is IrFunction)
                     expression.transformChildrenVoid(this)
                     val parent = currentDeclarationParent ?: error("No current declaration parent at ${reference.dump()}")

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

@@ -81,6 +81,8 @@ internal class SyntheticAccessorLowering(val context: JvmBackendContext) : IrEle
         }, null)
 
         irFile.transformChildrenVoid(this)
+        inlineLambdaToCallSite.clear()
+        inlineFunctionToCallSites.clear()
 
         for (accessor in pendingAccessorsToAdd) {
             assert(accessor.fileOrNull == irFile) {
@@ -90,6 +92,7 @@ internal class SyntheticAccessorLowering(val context: JvmBackendContext) : IrEle
             }
             (accessor.parent as IrDeclarationContainer).declarations.add(accessor)
         }
+        pendingAccessorsToAdd.clear()
     }
 
     private data class FieldKey(val fieldSymbol: IrFieldSymbol, val parent: IrDeclarationParent, val superQualifierSymbol: IrClassSymbol?)

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

@@ -10,4 +10,5 @@ import org.jetbrains.kotlin.ir.expressions.IrStatementOriginImpl
 interface JvmLoweredStatementOrigin {
     object DEFAULT_STUB_CALL_TO_IMPLEMENTATION : IrStatementOriginImpl("DEFAULT_STUB_CALL_TO_IMPLEMENTATION")
     object DO_WHILE_COUNTER_LOOP: IrStatementOriginImpl("DO_WHILE_COUNTER_LOOP")
+    object INLINE_LAMBDA : IrStatementOriginImpl("INLINE_LAMBDA")
 }

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.backend.jvm.codegen
 
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.backend.jvm.JvmLoweredStatementOrigin
 import org.jetbrains.kotlin.backend.jvm.ir.isInlineParameter
 import org.jetbrains.kotlin.codegen.IrExpressionLambda
 import org.jetbrains.kotlin.codegen.JvmKotlinType
@@ -13,6 +14,7 @@ import org.jetbrains.kotlin.codegen.StackValue
 import org.jetbrains.kotlin.codegen.ValueKind
 import org.jetbrains.kotlin.codegen.inline.*
 import org.jetbrains.kotlin.codegen.state.GenerationState
+import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.declarations.IrConstructor
 import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
 import org.jetbrains.kotlin.ir.declarations.IrFunction
@@ -85,10 +87,9 @@ class IrInlineCodegen(
         codegen: ExpressionCodegen,
         blockInfo: BlockInfo
     ) {
-        val isInlineParameter = irValueParameter.isInlineParameter()
-        if (isInlineParameter && argumentExpression.isInlineIrExpression()) {
-            val irReference = (argumentExpression as IrBlock).statements.last() as IrFunctionReference
-            val lambdaInfo = IrExpressionLambdaImpl(codegen, irReference)
+        val inlineLambda = argumentExpression.unwrapInlineLambda()
+        if (inlineLambda != null) {
+            val lambdaInfo = IrExpressionLambdaImpl(codegen, inlineLambda)
             rememberClosure(parameterType, irValueParameter.index, lambdaInfo)
             lambdaInfo.generateLambdaBody(sourceCompiler)
             lambdaInfo.reference.getArgumentsWithIr().forEachIndexed { index, (_, ir) ->
@@ -97,6 +98,7 @@ class IrInlineCodegen(
                 putCapturedToLocalVal(onStack, param, ir.type.toIrBasedKotlinType())
             }
         } else {
+            val isInlineParameter = irValueParameter.isInlineParameter()
             val kind = when {
                 irValueParameter.origin == IrDeclarationOrigin.MASK_FOR_DEFAULT_FUNCTION ->
                     ValueKind.DEFAULT_MASK
@@ -210,19 +212,11 @@ class IrExpressionLambdaImpl(
     }
 }
 
-fun IrExpression.isInlineIrExpression() =
-    when (this) {
-        is IrBlock -> origin.isInlineIrExpression()
-        is IrCallableReference<*> -> true.also {
-            assert((0 until valueArgumentsCount).none { getValueArgument(it) != null }) {
-                "Expecting 0 value arguments for bound callable reference: ${dump()}"
-            }
-        }
-        else -> false
-    }
-
-fun IrStatementOrigin?.isInlineIrExpression() =
-    isLambda || this == IrStatementOrigin.ADAPTED_FUNCTION_REFERENCE || this == IrStatementOrigin.SUSPEND_CONVERSION
+fun IrStatement.unwrapInlineLambda(): IrFunctionReference? = when (this) {
+    is IrBlock -> statements.lastOrNull()?.unwrapInlineLambda()
+    is IrFunctionReference -> takeIf { it.origin == JvmLoweredStatementOrigin.INLINE_LAMBDA }
+    else -> null
+}
 
 fun IrFunction.isInlineFunctionCall(context: JvmBackendContext) =
     (!context.state.isInlineDisabled || typeParameters.any { it.isReified }) && (isInline || isInlineArrayConstructor(context))

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

@@ -7,59 +7,34 @@ package org.jetbrains.kotlin.backend.jvm.ir
 
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.codegen.isInlineFunctionCall
-import org.jetbrains.kotlin.backend.jvm.codegen.isInlineIrExpression
+import org.jetbrains.kotlin.backend.jvm.codegen.unwrapInlineLambda
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.expressions.IrBlock
-import org.jetbrains.kotlin.ir.expressions.IrCallableReference
 import org.jetbrains.kotlin.ir.expressions.IrFunctionAccessExpression
 import org.jetbrains.kotlin.ir.expressions.IrFunctionReference
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
 
-open class IrInlineReferenceLocator(private val context: JvmBackendContext) : IrElementVisitor<Unit, IrDeclaration?> {
-    override fun visitElement(element: IrElement, data: IrDeclaration?) {
-        element.acceptChildren(this, data)
-    }
-
-    override fun visitDeclaration(declaration: IrDeclarationBase, data: IrDeclaration?) {
-        val scope = if (declaration is IrVariable) data else declaration
-        declaration.acceptChildren(this, scope)
-    }
+abstract class IrInlineReferenceLocator(private val context: JvmBackendContext) : IrElementVisitor<Unit, IrDeclaration?> {
+    override fun visitElement(element: IrElement, data: IrDeclaration?) =
+        element.acceptChildren(this, if (element is IrDeclaration && element !is IrVariable) element else data)
 
     override fun visitFunctionAccess(expression: IrFunctionAccessExpression, data: IrDeclaration?) {
         val function = expression.symbol.owner
         if (function.isInlineFunctionCall(context)) {
             for (parameter in function.valueParameters) {
-                if (!parameter.isInlineParameter())
-                    continue
-
-                val valueArgument = expression.getValueArgument(parameter.index) ?: continue
-                if (!valueArgument.isInlineIrExpression())
-                    continue
-
-                if (valueArgument is IrBlock) {
-                    visitInlineLambda(valueArgument.statements.last() as IrFunctionReference, function, parameter, data!!)
-                } else if (valueArgument is IrCallableReference<*>) {
-                    visitInlineReference(valueArgument)
-                }
+                val lambda = expression.getValueArgument(parameter.index)?.unwrapInlineLambda() ?: continue
+                visitInlineLambda(lambda, function, parameter, data!!)
             }
         }
         return super.visitFunctionAccess(expression, data)
     }
 
-    open fun visitInlineReference(argument: IrCallableReference<*>) {}
-
-    open fun visitInlineLambda(argument: IrFunctionReference, callee: IrFunction, parameter: IrValueParameter, scope: IrDeclaration) =
-        visitInlineReference(argument)
-
-    companion object {
-        fun scan(context: JvmBackendContext, element: IrElement): Set<IrCallableReference<*>> =
-            mutableSetOf<IrCallableReference<*>>().apply {
-                element.accept(object : IrInlineReferenceLocator(context) {
-                    override fun visitInlineReference(argument: IrCallableReference<*>) {
-                        add(argument)
-                    }
-                }, null)
-            }
-    }
+    abstract fun visitInlineLambda(argument: IrFunctionReference, callee: IrFunction, parameter: IrValueParameter, scope: IrDeclaration)
 }
+
+inline fun IrFile.findInlineLambdas(
+    context: JvmBackendContext, crossinline onLambda: (IrFunctionReference, IrFunction, IrValueParameter, IrDeclaration) -> Unit
+) = accept(object : IrInlineReferenceLocator(context) {
+    override fun visitInlineLambda(argument: IrFunctionReference, callee: IrFunction, parameter: IrValueParameter, scope: IrDeclaration) =
+        onLambda(argument, callee, parameter, scope)
+}, null)