JVM_IR: mark inline lambda functions with a special origin

Keeping the origin as LOCAL_FUNCTION_FOR_LAMBDA was a mistake as this
tells codegen nothing. Changing the origin in allows, for example,
removing the hack that detaches inline lambdas from the IR tree before
verification and codegen, or treating inline lambdas and inline
anonymous functions the same way.

This includes fake functions created for inline callable references.

 #KT-48319 Fixed
 #KT-47279 Fixed?

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

@@ -639,7 +639,8 @@ class LocalDeclarationsLowering(
             newDeclaration.copyAttributes(oldDeclaration)
 
             newDeclaration.valueParameters += createTransformedValueParameters(
-                capturedValues, localFunctionContext, oldDeclaration, newDeclaration, isExplicitLocalFunction = oldDeclaration.origin != IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA
+                capturedValues, localFunctionContext, oldDeclaration, newDeclaration,
+                isExplicitLocalFunction = oldDeclaration.origin == IrDeclarationOrigin.LOCAL_FUNCTION
             )
             newDeclaration.recordTransformedValueParameters(localFunctionContext)
 

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

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.backend.common.pop
 import org.jetbrains.kotlin.backend.common.push
 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.codegen.continuationParameter
 import org.jetbrains.kotlin.backend.jvm.codegen.hasContinuation
 import org.jetbrains.kotlin.backend.jvm.codegen.isInvokeSuspendOfContinuation
@@ -31,10 +32,7 @@ import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.builders.declarations.*
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
-import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
-import org.jetbrains.kotlin.ir.expressions.impl.IrFunctionReferenceImpl
-import org.jetbrains.kotlin.ir.expressions.impl.IrGetFieldImpl
-import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl
+import org.jetbrains.kotlin.ir.expressions.impl.*
 import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.types.IrType
@@ -458,20 +456,19 @@ private fun <T : IrMemberAccessExpression<IrFunctionSymbol>> T.retargetToSuspend
         it.copyTypeArgumentsFrom(this)
         it.dispatchReceiver = dispatchReceiver
         it.extensionReceiver = extensionReceiver
-        val continuationIndex = view.continuationParameter()!!.index
+        val continuationParameter = view.continuationParameter()!!
         for (i in 0 until valueArgumentsCount) {
-            it.putValueArgument(i + if (i >= continuationIndex) 1 else 0, getValueArgument(i))
+            it.putValueArgument(i + if (i >= continuationParameter.index) 1 else 0, getValueArgument(i))
         }
         if (caller != null) {
-            // At this point the only LOCAL_FUNCTION_FOR_LAMBDAs are inline and crossinline lambdas.
-            val continuation = if (caller.originalFunction.origin == IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA)
-                context.fakeContinuation
+            val continuation = if (caller.origin == JvmLoweredDeclarationOrigin.INLINE_LAMBDA)
+                IrCompositeImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, continuationParameter.type, JvmLoweredStatementOrigin.FAKE_CONTINUATION)
             else
                 IrGetValueImpl(
                     UNDEFINED_OFFSET, UNDEFINED_OFFSET, caller.continuationParameter()?.symbol
                         ?: throw AssertionError("${caller.render()} has no continuation; can't call ${owner.render()}")
                 )
-            it.putValueArgument(continuationIndex, continuation)
+            it.putValueArgument(continuationParameter.index, continuation)
         }
     }
 }

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

@@ -8,23 +8,19 @@ 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.JvmLoweredStatementOrigin
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.expressions.IrFunctionReference
 import org.jetbrains.kotlin.ir.expressions.IrLoop
-import org.jetbrains.kotlin.ir.util.parentAsClass
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
 
 internal val prepareForBytecodeInlining = makeIrModulePhase(
     ::BytecodeInliningPreparationLowering,
     name = "BytecodeInliningPreparation",
-    description = "Remove inline lambda declarations and label all loops"
+    description = "Label all loops for non-local break/continue"
 )
 
 private class BytecodeInliningPreparationLowering(val context: JvmBackendContext) : FileLoweringPass {
     override fun lower(irFile: IrFile) {
-        val loweredLambdasToDelete = mutableSetOf<IrFunction>()
         irFile.accept(object : IrElementVisitor<Unit, String> {
             // This counter is intentionally not local to every declaration because their names might clash.
             private var counter = 0
@@ -38,19 +34,6 @@ private class BytecodeInliningPreparationLowering(val context: JvmBackendContext
                 loop.label = "$data${++counter}"
                 super.visitLoop(loop, data)
             }
-
-            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. 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)
-            }
         }, "")
-
-        for (irClass in loweredLambdasToDelete.mapTo(mutableSetOf()) { it.parentAsClass }) {
-            irClass.declarations.removeAll(loweredLambdasToDelete)
-        }
     }
 }

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

@@ -220,6 +220,7 @@ private fun IrSimpleFunction.createMultifileDelegateIfNeeded(
     if (DescriptorVisibilities.isPrivate(originalVisibility) ||
         name == StaticInitializersLowering.clinitName ||
         origin == JvmLoweredDeclarationOrigin.SYNTHETIC_ACCESSOR ||
+        origin == JvmLoweredDeclarationOrigin.INLINE_LAMBDA ||
         origin == IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA ||
         origin == IrDeclarationOrigin.PROPERTY_DELEGATE ||
         // $annotations methods in the facade are only needed for const properties.

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

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.backend.common.FileLoweringPass
 import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 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.codegen.isInlineFunctionCall
 import org.jetbrains.kotlin.backend.jvm.ir.createJvmIrBuilder
@@ -69,8 +70,9 @@ private class InlineCallableReferenceToLambdaVisitor(val context: JvmBackendCont
     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)
+            val reference = statements.last() as IrFunctionReference
+            reference.symbol.owner.origin = JvmLoweredDeclarationOrigin.INLINE_LAMBDA
+            statements[statements.lastIndex] = reference.replaceOrigin(JvmLoweredStatementOrigin.INLINE_LAMBDA)
         }
 
         this is IrFunctionReference -> // ::function -> { args... -> function(args...) }
@@ -85,7 +87,7 @@ private class InlineCallableReferenceToLambdaVisitor(val context: JvmBackendCont
     private fun IrPropertyReference.wrapField(field: IrField): IrSimpleFunction =
         context.irFactory.buildFun {
             setSourceRange(this@wrapField)
-            origin = IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA
+            origin = JvmLoweredDeclarationOrigin.INLINE_LAMBDA
             name = Name.identifier(STUB_FOR_INLINING)
             visibility = DescriptorVisibilities.LOCAL
             returnType = field.type
@@ -104,7 +106,7 @@ private class InlineCallableReferenceToLambdaVisitor(val context: JvmBackendCont
     private fun IrCallableReference<*>.wrapFunction(referencedFunction: IrFunction): IrSimpleFunction =
         context.irFactory.buildFun {
             setSourceRange(this@wrapFunction)
-            origin = IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA
+            origin = JvmLoweredDeclarationOrigin.INLINE_LAMBDA
             name = Name.identifier(STUB_FOR_INLINING)
             visibility = DescriptorVisibilities.LOCAL
             returnType = ((type as IrSimpleType).arguments.last() as IrTypeProjection).type

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

@@ -132,7 +132,7 @@ internal class InterfaceLowering(val context: JvmBackendContext) : IrElementTran
                         (function.origin == JvmLoweredDeclarationOrigin.SYNTHETIC_METHOD_FOR_PROPERTY_OR_TYPEALIAS_ANNOTATIONS &&
                                 (isCompatibilityMode || jvmDefaultMode == JvmDefaultMode.ENABLE) &&
                                 function.isCompiledToJvmDefault(jvmDefaultMode)) -> {
-                    if (function.origin == IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA || function.origin == IrDeclarationOrigin.LOCAL_FUNCTION) {
+                    if (function.origin == JvmLoweredDeclarationOrigin.INLINE_LAMBDA) {
                         //move as is
                         val defaultImplsClass = context.cachedDeclarations.getDefaultImplsClass(irClass)
                         defaultImplsClass.declarations.add(function)

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

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.backend.common.lower.loops.isInductionVariable
 import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredStatementOrigin
+import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
 import org.jetbrains.kotlin.backend.jvm.ir.createJvmIrBuilder
 import org.jetbrains.kotlin.codegen.AsmUtil
 import org.jetbrains.kotlin.descriptors.Modality
@@ -125,12 +126,12 @@ class JvmOptimizationLowering(val context: JvmBackendContext) : FileLoweringPass
 
         // For some functions, we clear the current class field since the code could end up
         // in another class then the one it is nested under in the IR.
-        // TODO: Loosen this up for local functions for lambdas passed as an inline lambda
-        // argument to an inline function. In that case the code does end up in the current class.
+        // TODO: replace this with the code from SyntheticAccessorLowering that returns the current class
+        //   or package accounting for all inline functions and lambdas.
         override fun visitFunction(declaration: IrFunction, data: IrClass?): IrStatement {
             val codeMightBeGeneratedInDifferentClass = declaration.isSuspend ||
                     declaration.isInline ||
-                    declaration.origin == IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA
+                    declaration.origin == JvmLoweredDeclarationOrigin.INLINE_LAMBDA
             declaration.transformChildren(this, data.takeUnless { codeMightBeGeneratedInDifferentClass })
             return declaration
         }

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

@@ -16,7 +16,6 @@ import org.jetbrains.kotlin.backend.jvm.caches.CollectionStubComputer
 import org.jetbrains.kotlin.backend.jvm.codegen.ClassCodegen
 import org.jetbrains.kotlin.backend.jvm.codegen.IrTypeMapper
 import org.jetbrains.kotlin.backend.jvm.codegen.MethodSignatureMapper
-import org.jetbrains.kotlin.backend.jvm.codegen.createFakeContinuation
 import org.jetbrains.kotlin.backend.jvm.intrinsics.IrIntrinsicMethods
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
@@ -124,7 +123,6 @@ class JvmBackendContext(
 
     val suspendLambdaToOriginalFunctionMap = ConcurrentHashMap<IrFunctionReference, IrFunction>()
     val suspendFunctionOriginalToView = ConcurrentHashMap<IrSimpleFunction, IrSimpleFunction>()
-    val fakeContinuation: IrExpression = createFakeContinuation(this)
 
     val staticDefaultStubs = ConcurrentHashMap<IrSimpleFunctionSymbol, IrSimpleFunction>()
 

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

@@ -47,4 +47,5 @@ interface JvmLoweredDeclarationOrigin : IrDeclarationOrigin {
     object FIELD_FOR_STATIC_CALLABLE_REFERENCE_INSTANCE : IrDeclarationOriginImpl("FIELD_FOR_STATIC_CALLABLE_REFERENCE_INSTANCE")
     object ABSTRACT_BRIDGE_STUB : IrDeclarationOriginImpl("ABSTRACT_BRIDGE_STUB")
     object INVOVEDYNAMIC_CALL_TARGET : IrDeclarationOriginImpl("INVOVEDYNAMIC_CALL_TARGET")
+    object INLINE_LAMBDA : IrDeclarationOriginImpl("INLINE_LAMBDA")
 }

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

@@ -11,4 +11,5 @@ 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")
+    object FAKE_CONTINUATION : IrStatementOriginImpl("FAKE_CONTINUATION")
 }

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

@@ -136,7 +136,7 @@ class ClassCodegen private constructor(
         val smap = context.getSourceMapper(irClass)
         // 1. Any method other than `<clinit>` can add a field and a `<clinit>` statement:
         for (method in irClass.declarations.filterIsInstance<IrFunction>()) {
-            if (method.name.asString() != "<clinit>") {
+            if (method.name.asString() != "<clinit>" && method.origin != JvmLoweredDeclarationOrigin.INLINE_LAMBDA) {
                 generateMethod(method, smap)
             }
         }

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

@@ -18,18 +18,14 @@ import org.jetbrains.kotlin.codegen.coroutines.INVOKE_SUSPEND_METHOD_NAME
 import org.jetbrains.kotlin.codegen.coroutines.SUSPEND_IMPL_NAME_SUFFIX
 import org.jetbrains.kotlin.codegen.coroutines.reportSuspensionPointInsideMonitor
 import org.jetbrains.kotlin.config.LanguageFeature
-import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
-import org.jetbrains.kotlin.ir.expressions.impl.IrErrorExpressionImpl
 import org.jetbrains.kotlin.ir.types.*
-import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection
 import org.jetbrains.kotlin.ir.util.file
 import org.jetbrains.kotlin.ir.util.functions
 import org.jetbrains.kotlin.ir.util.isSuspend
 import org.jetbrains.kotlin.ir.util.parentAsClass
 import org.jetbrains.kotlin.psi.KtElement
-import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.org.objectweb.asm.MethodVisitor
 import org.jetbrains.org.objectweb.asm.Type
 import org.jetbrains.org.objectweb.asm.tree.MethodNode
@@ -167,7 +163,7 @@ fun IrFunction.hasContinuation(): Boolean = isInvokeSuspendOfLambda() ||
         isSuspend && shouldContainSuspendMarkers() &&
         // These are templates for the inliner; the continuation is borrowed from the caller method.
         !isEffectivelyInlineOnly() &&
-        origin != IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA &&
+        origin != JvmLoweredDeclarationOrigin.INLINE_LAMBDA &&
         origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE &&
         origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE
 
@@ -180,13 +176,6 @@ fun IrExpression?.isReadOfCrossinline(): Boolean = when (this) {
 internal fun IrExpression?.isReadOfInlineLambda(): Boolean = isReadOfCrossinline() ||
         (this is IrGetValue && origin == IrStatementOrigin.VARIABLE_AS_FUNCTION && (symbol.owner as? IrValueParameter)?.isNoinline == false)
 
-internal fun createFakeContinuation(context: JvmBackendContext): IrExpression = IrErrorExpressionImpl(
-    UNDEFINED_OFFSET,
-    UNDEFINED_OFFSET,
-    context.ir.symbols.continuationClass.createType(true, listOf(makeTypeProjection(context.irBuiltIns.anyNType, Variance.INVARIANT))),
-    "FAKE_CONTINUATION"
-)
-
 internal fun IrFunction.originalReturnTypeOfSuspendFunctionReturningUnboxedInlineClass(): IrType? {
     if (this !is IrSimpleFunction || !isSuspend) return null
     // Unlike `suspendFunctionOriginal()`, this also maps `$default` stubs to the original function.

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

@@ -208,11 +208,7 @@ class ExpressionCodegen(
 
     // TODO remove
     fun gen(expression: IrExpression, type: Type, irType: IrType, data: BlockInfo): StackValue {
-        if (expression.attributeOwnerId === context.fakeContinuation) {
-            addFakeContinuationMarker(mv)
-        } else {
-            expression.accept(this, data).materializeAt(type, irType)
-        }
+        expression.accept(this, data).materializeAt(type, irType)
         return StackValue.onStack(type, irType.toIrBasedKotlinType())
     }
 
@@ -444,7 +440,12 @@ class ExpressionCodegen(
     }
 
     override fun visitContainerExpression(expression: IrContainerExpression, data: BlockInfo) =
-        visitStatementContainer(expression, data)
+        if (expression.origin == JvmLoweredStatementOrigin.FAKE_CONTINUATION) {
+            addFakeContinuationMarker(mv)
+            expression.onStack
+        } else {
+            visitStatementContainer(expression, data)
+        }
 
     override fun visitCall(expression: IrCall, data: BlockInfo): PromisedValue {
         val intrinsic = classCodegen.context.irIntrinsics.getIntrinsic(expression.symbol)

compiler/testData/codegen/box/coroutines/suspendFromInlineLambda.kt

@@ -26,13 +26,12 @@ fun box(): String {
 
     builder {
         result += "-"
-        foo {
-            result += suspendHere(it).toString()
-        }
+        foo { result += suspendHere(it).toString() }
+        foo(fun(it: Int) { result += suspendHere(it).toString() })
         result += "+"
     }
 
-    if (result != "-24+") return "fail: $result"
+    if (result != "-2424+") return "fail: $result"
 
     return "OK"
 }

compiler/testData/codegen/bytecodeText/kt9603.kt

@@ -3,7 +3,8 @@ class A {
         private set
 
     fun test() {
-        { prop }()
+        val f = { prop }
+        f()
     }
 }