Fix Java interop of inline suspend functions with suspend parameters

In 1.3.31 I fixed Java interop for inline function with coroutines
(TL;DR: when we need a state machine, generate two methods: one with
normal name, and the other one with $$forInline suffix, for the inliner
to use, just like inline suspend functions), however, I forgot a case
with inline suspend function with inline suspend function parameter.
In this case, the compiler a generated two functions, as needed, but,
neither of them had a state-machine. This change adds the state-machine
for the method with normal name. Note, that suspend inline functions
with crossinline parameter, which are also supported by the change,
did not cause incorrect behaviour, since until now they were generated
as synthetic.

 #KT-31354 Fixed

compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/SuspendFunctionGenerationStrategy.kt

@@ -30,7 +30,7 @@ open class SuspendFunctionGenerationStrategy(
     state: GenerationState,
     protected val originalSuspendDescriptor: FunctionDescriptor,
     protected val declaration: KtFunction,
-    private val containingClassInternalName: String,
+    protected val containingClassInternalName: String,
     private val constructorCallNormalizationMode: JVMConstructorCallNormalizationMode,
     protected val functionCodegen: FunctionCodegen
 ) : FunctionGenerationStrategy.CodegenBased(state) {
@@ -53,16 +53,7 @@ open class SuspendFunctionGenerationStrategy(
     override fun wrapMethodVisitor(mv: MethodVisitor, access: Int, name: String, desc: String): MethodVisitor {
         if (access and Opcodes.ACC_ABSTRACT != 0) return mv
 
-        val stateMachineBuilder = CoroutineTransformerMethodVisitor(
-            mv, access, name, desc, null, null, containingClassInternalName, this::classBuilderForCoroutineState,
-            isForNamedFunction = true,
-            element = declaration,
-            diagnostics = state.diagnostics,
-            shouldPreserveClassInitialization = constructorCallNormalizationMode.shouldPreserveClassInitialization,
-            needDispatchReceiver = originalSuspendDescriptor.dispatchReceiverParameter != null,
-            internalNameForDispatchReceiver = containingClassInternalNameOrNull(),
-            languageVersionSettings = languageVersionSettings
-        )
+        val stateMachineBuilder = createStateMachineBuilder(mv, access, name, desc)
 
         val forInline = state.bindingContext[CodegenBinding.CAPTURES_CROSSINLINE_LAMBDA, originalSuspendDescriptor] == true
         // Both capturing and inline functions share the same suffix, however, inline functions can also be capturing
@@ -95,6 +86,24 @@ open class SuspendFunctionGenerationStrategy(
             ) else stateMachineBuilder
     }
 
+    protected fun createStateMachineBuilder(
+        mv: MethodVisitor,
+        access: Int,
+        name: String,
+        desc: String
+    ): CoroutineTransformerMethodVisitor {
+        return CoroutineTransformerMethodVisitor(
+            mv, access, name, desc, null, null, containingClassInternalName, this::classBuilderForCoroutineState,
+            isForNamedFunction = true,
+            element = declaration,
+            diagnostics = state.diagnostics,
+            shouldPreserveClassInitialization = constructorCallNormalizationMode.shouldPreserveClassInitialization,
+            needDispatchReceiver = originalSuspendDescriptor.dispatchReceiverParameter != null,
+            internalNameForDispatchReceiver = containingClassInternalNameOrNull(),
+            languageVersionSettings = languageVersionSettings
+        )
+    }
+
     private fun containingClassInternalNameOrNull() =
         originalSuspendDescriptor.containingDeclaration.safeAs<ClassDescriptor>()?.let(state.typeMapper::mapClass)?.internalName
 

compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/SuspendInlineFunctionGenerationStrategy.kt

@@ -5,20 +5,30 @@
 
 package org.jetbrains.kotlin.codegen.coroutines
 
+import org.jetbrains.kotlin.builtins.isSuspendFunctionType
 import org.jetbrains.kotlin.codegen.ExpressionCodegen
 import org.jetbrains.kotlin.codegen.FunctionCodegen
 import org.jetbrains.kotlin.codegen.FunctionGenerationStrategy
 import org.jetbrains.kotlin.codegen.TransformationMethodVisitor
 import org.jetbrains.kotlin.codegen.inline.coroutines.FOR_INLINE_SUFFIX
+import org.jetbrains.kotlin.codegen.inline.coroutines.findReceiverOfInvoke
+import org.jetbrains.kotlin.codegen.inline.coroutines.surroundInvokesWithSuspendMarkers
+import org.jetbrains.kotlin.codegen.inline.isInvokeOnLambda
+import org.jetbrains.kotlin.codegen.optimization.common.asSequence
+import org.jetbrains.kotlin.codegen.optimization.fixStack.FixStackMethodTransformer
+import org.jetbrains.kotlin.codegen.optimization.transformer.MethodTransformer
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.config.JVMConstructorCallNormalizationMode
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
 import org.jetbrains.kotlin.psi.KtFunction
+import org.jetbrains.kotlin.resolve.inline.InlineUtil
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature
 import org.jetbrains.org.objectweb.asm.MethodVisitor
 import org.jetbrains.org.objectweb.asm.Opcodes
-import org.jetbrains.org.objectweb.asm.tree.MethodNode
+import org.jetbrains.org.objectweb.asm.tree.*
+import org.jetbrains.org.objectweb.asm.tree.analysis.SourceInterpreter
 
 // For named suspend function we generate two methods:
 // 1) to use as noinline function, which have state machine
@@ -44,7 +54,10 @@ class SuspendInlineFunctionGenerationStrategy(
         if (access and Opcodes.ACC_ABSTRACT != 0) return mv
 
         return MethodNodeCopyingMethodVisitor(
-            super.wrapMethodVisitor(mv, access, name, desc), access, name, desc,
+            SurroundSuspendParameterCallsWithSuspendMarkersMethodVisitor(
+                createStateMachineBuilder(mv, access, name, desc),
+                access, name, desc, containingClassInternalName, originalSuspendDescriptor.valueParameters
+            ), access, name, desc,
             newMethod = { origin, newAccess, newName, newDesc ->
                 functionCodegen.newMethod(origin, newAccess, newName, newDesc, null, null)
             }, keepAccess = false
@@ -78,3 +91,39 @@ class MethodNodeCopyingMethodVisitor(
             else access or Opcodes.ACC_PRIVATE and Opcodes.ACC_PUBLIC.inv() and Opcodes.ACC_PROTECTED.inv()
     }
 }
+
+private class SurroundSuspendParameterCallsWithSuspendMarkersMethodVisitor(
+    delegate: MethodVisitor,
+    access: Int,
+    name: String,
+    desc: String,
+    private val thisName: String,
+    private val valueParameters: List<ValueParameterDescriptor>
+): TransformationMethodVisitor(delegate, access, name, desc, null, null) {
+    override fun performTransformations(methodNode: MethodNode) {
+        fun AbstractInsnNode.index() = methodNode.instructions.indexOf(this)
+        fun AbstractInsnNode.isInlineSuspendParameter(): Boolean {
+            if (this !is VarInsnNode) return false
+            val index = `var` - (if (methodNode.access and Opcodes.ACC_STATIC != 0) 0 else 1)
+            return opcode == Opcodes.ALOAD && index < valueParameters.size && InlineUtil.isInlineParameter(valueParameters[index]) &&
+                    valueParameters[index].type.isSuspendFunctionType
+        }
+
+        FixStackMethodTransformer().transform(thisName, methodNode)
+
+        val sourceFrames = MethodTransformer.analyze(thisName, methodNode, SourceInterpreter())
+
+        val noinlineInvokes = arrayListOf<Pair<AbstractInsnNode, AbstractInsnNode>>()
+
+        for (insn in methodNode.instructions.asSequence()) {
+            if (insn.opcode != Opcodes.INVOKEINTERFACE) continue
+            insn as MethodInsnNode
+            if (!isInvokeOnLambda(insn.owner, insn.name)) continue
+            val frame = sourceFrames[insn.index()] ?: continue
+            val aload = findReceiverOfInvoke(frame, insn).takeIf { it?.isInlineSuspendParameter() == true } as? VarInsnNode ?: continue
+            noinlineInvokes.add(insn to aload)
+        }
+
+        surroundInvokesWithSuspendMarkers(methodNode, noinlineInvokes)
+    }
+}