JVM_IR: refine the condition for caching method nodes

Don't cache non-inline `suspend fun`s if there is no `$$forInline`;
do cache `invokeSuspend` if there *is* `invokeSuspend$$forInline`.

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

@@ -312,10 +312,12 @@ class ClassCodegen private constructor(
     private val generatedInlineMethods = mutableMapOf<IrFunction, SMAPAndMethodNode>()
 
     fun generateMethodNode(method: IrFunction): SMAPAndMethodNode {
-        if (!method.isInline && !method.isSuspend) {
+        if (!method.isInline && !method.isSuspendCapturingCrossinline()) {
             // Inline methods can be used multiple times by `IrSourceCompilerForInline`, suspend methods
-            // could be used twice if they capture crossinline lambdas, and everything else is only
-            // generated by `generateMethod` below so does not need caching.
+            // are used twice (`f` and `f$$forInline`) if they capture crossinline lambdas, and everything
+            // else is only generated by `generateMethod` below so does not need caching.
+            // TODO: inline lambdas are not marked `isInline`, and are generally used once, but may be needed
+            //       multiple times if declared in a `finally` block - should they be cached?
             return FunctionCodegen(method, this).generate()
         }
         val (node, smap) = generatedInlineMethods.getOrPut(method) { FunctionCodegen(method, this).generate() }
@@ -348,14 +350,13 @@ class ClassCodegen private constructor(
             override fun visitLineNumber(line: Int, start: Label) =
                 super.visitLineNumber(smapCopier.mapLineNumber(line), start)
         }
-        if (method.hasContinuation() || method.isInvokeSuspendOfLambda()) {
+        if (method.hasContinuation()) {
             // Generate a state machine within this method. The continuation class for it should be generated
             // lazily so that if tail call optimization kicks in, the unused class will not be written to the output.
-            val continuationClassCodegen = lazy { getOrCreate(method.continuationClass()!!, context, method) }
-            node.acceptWithStateMachine(method, this, smapCopyingVisitor) {
-                if (method.isSuspend) continuationClassCodegen.value.visitor else visitor
-            }
-            if (continuationClassCodegen.isInitialized() || method.alwaysNeedsContinuation()) {
+            val continuationClass = method.continuationClass() // null if `SuspendLambda.invokeSuspend` - `this` is continuation itself
+            val continuationClassCodegen = lazy { if (continuationClass != null) getOrCreate(continuationClass, context, method) else this }
+            node.acceptWithStateMachine(method, this, smapCopyingVisitor) { continuationClassCodegen.value.visitor }
+            if (continuationClass != null && (continuationClassCodegen.isInitialized() || method.isSuspendCapturingCrossinline())) {
                 continuationClassCodegen.value.generate()
             }
         } else {

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

@@ -100,7 +100,7 @@ internal fun IrFunction.suspendForInlineToOriginal(): IrSimpleFunction? {
     } as IrSimpleFunction?
 }
 
-internal fun IrFunction.alwaysNeedsContinuation(): Boolean =
+internal fun IrFunction.isSuspendCapturingCrossinline(): Boolean =
     this is IrSimpleFunction && hasContinuation() && parentAsClass.declarations.any {
         it is IrSimpleFunction && it.attributeOwnerId == attributeOwnerId &&
                 it.origin == JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE
@@ -157,10 +157,10 @@ internal fun IrFunction.shouldContainSuspendMarkers(): Boolean = !isInvokeSuspen
         !isBridgeToSuspendImplMethod() &&
         !isStaticInlineClassReplacementDelegatingCall()
 
-internal fun IrFunction.hasContinuation(): Boolean = isSuspend && shouldContainSuspendMarkers() &&
-        // This is inline-only function
+internal fun IrFunction.hasContinuation(): Boolean = isInvokeSuspendOfLambda() ||
+        isSuspend && shouldContainSuspendMarkers() &&
+        // These are templates for the inliner; the continuation is borrowed from the caller method.
         !isEffectivelyInlineOnly() &&
-        // These are templates for the inliner; the continuation will be generated after it runs.
         origin != IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA &&
         origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE &&
         origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE

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

@@ -248,6 +248,7 @@ class ExpressionCodegen(
     }
 
     private fun generateFakeContinuationConstructorIfNeeded() {
+        if (!irFunction.isSuspendCapturingCrossinline()) return
         val continuationClass = irFunction.continuationClass() ?: return
         val continuationType = typeMapper.mapClass(continuationClass)
         val continuationIndex = frameMap.getIndex(irFunction.continuationParameter()!!.symbol)

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

@@ -214,7 +214,7 @@ private class SuspendLambdaLowering(context: JvmBackendContext) : SuspendLowerin
             INVOKE_SUSPEND_METHOD_NAME + FOR_INLINE_SUFFIX,
             context.irBuiltIns.anyNType,
             Modality.FINAL,
-            origin = JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE
+            origin = JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE
         ).apply {
             copyAttributes(invokeSuspend)
             generateErrorForInlineBody()