Check variables, captured by outer object when searching for suspend crossinline parameters

compiler/backend/src/org/jetbrains/kotlin/codegen/FunctionCodegen.java

@@ -14,6 +14,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.backend.common.CodegenUtil;
 import org.jetbrains.kotlin.backend.common.bridges.Bridge;
+import org.jetbrains.kotlin.codegen.binding.CalculatedClosure;
 import org.jetbrains.kotlin.codegen.binding.CodegenBinding;
 import org.jetbrains.kotlin.codegen.context.*;
 import org.jetbrains.kotlin.codegen.coroutines.CoroutineCodegenUtilKt;
@@ -1645,7 +1646,7 @@ public class FunctionCodegen {
     }
 
     @Nullable
-    public Map<DeclarationDescriptor, EnclosedValueDescriptor> getCaptureVariables() {
-        return owner.closure == null ? null : owner.closure.getCaptureVariables();
+    public CalculatedClosure getClosure() {
+        return owner.closure;
     }
 }

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

@@ -8,10 +8,11 @@ package org.jetbrains.kotlin.codegen.coroutines
 import com.intellij.util.ArrayUtil
 import org.jetbrains.kotlin.builtins.isSuspendFunctionTypeOrSubtype
 import org.jetbrains.kotlin.codegen.*
+import org.jetbrains.kotlin.codegen.binding.CalculatedClosure
 import org.jetbrains.kotlin.codegen.binding.CodegenBinding
 import org.jetbrains.kotlin.codegen.binding.CodegenBinding.CAPTURES_CROSSINLINE_LAMBDA
+import org.jetbrains.kotlin.codegen.binding.CodegenBinding.CLOSURE
 import org.jetbrains.kotlin.codegen.context.ClosureContext
-import org.jetbrains.kotlin.codegen.context.EnclosedValueDescriptor
 import org.jetbrains.kotlin.codegen.context.MethodContext
 import org.jetbrains.kotlin.codegen.inline.coroutines.SurroundSuspendLambdaCallsWithSuspendMarkersMethodVisitor
 import org.jetbrains.kotlin.codegen.serialization.JvmSerializationBindings.METHOD_FOR_FUNCTION
@@ -26,6 +27,7 @@ import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.load.kotlin.header.KotlinClassHeader
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
 import org.jetbrains.kotlin.resolve.descriptorUtil.module
@@ -479,7 +481,7 @@ class CoroutineCodegenForLambda private constructor(
                         MethodNodeCopyingMethodVisitor(
                             SurroundSuspendLambdaCallsWithSuspendMarkersMethodVisitor(
                                 stateMachineBuilder, access, name, desc, v.thisName,
-                                isCapturedSuspendLambda = { isCapturedSuspendLambda(closure.captureVariables, it.name) }
+                                isCapturedSuspendLambda = { isCapturedSuspendLambda(closure, it.name, state.bindingContext) }
                             ), access, name, desc,
                             newMethod = { origin, newAccess, newName, newDesc ->
                                 functionCodegen.newMethod(origin, newAccess, newName, newDesc, null, null)
@@ -528,13 +530,19 @@ class CoroutineCodegenForLambda private constructor(
     }
 }
 
-fun isCapturedSuspendLambda(captureVariables: Map<DeclarationDescriptor, EnclosedValueDescriptor>, name: String): Boolean {
-    for ((param, value) in captureVariables) {
+fun isCapturedSuspendLambda(closure: CalculatedClosure, name: String, bindingContext: BindingContext): Boolean {
+    for ((param, value) in closure.captureVariables) {
         if (param !is ValueParameterDescriptor) continue
         if (value.fieldName != name) continue
         return param.type.isSuspendFunctionTypeOrSubtype
     }
-    return false
+    val classDescriptor = closure.capturedOuterClassDescriptor ?: return false
+    return isCapturedSuspendLambda(classDescriptor, name, bindingContext)
+}
+
+fun isCapturedSuspendLambda(classDescriptor: ClassDescriptor, name: String, bindingContext: BindingContext): Boolean {
+    val closure = bindingContext[CLOSURE, classDescriptor] ?: return false
+    return isCapturedSuspendLambda(closure, name, bindingContext)
 }
 
 private class AddEndLabelMethodVisitor(

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

@@ -76,10 +76,11 @@ open class SuspendFunctionGenerationStrategy(
                         access, name, desc, containingClassInternalName,
                         isCapturedSuspendLambda = {
                             isCapturedSuspendLambda(
-                                functionCodegen.captureVariables.sure {
+                                functionCodegen.closure.sure {
                                     "Anonymous object should have closure"
                                 },
-                                it.name
+                                it.name,
+                                state.bindingContext
                             )
                         }
                     ), access, name, desc,

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/AnonymousObjectTransformer.kt

@@ -484,9 +484,9 @@ class AnonymousObjectTransformer(
                         alreadyAddedParam?.newFieldName ?: getNewFieldName(desc.fieldName, false),
                         alreadyAddedParam != null
                     )
-                    if (info is PsiExpressionLambda && info.captureVariables.any { it.value.fieldName == desc.fieldName }) {
+                    if (info is PsiExpressionLambda && info.closure.captureVariables.any { it.value.fieldName == desc.fieldName }) {
                         recapturedParamInfo.functionalArgument = NonInlineableArgumentForInlineableParameterCalledInSuspend(
-                            isCapturedSuspendLambda(info.captureVariables, desc.fieldName)
+                            isCapturedSuspendLambda(info.closure, desc.fieldName, inliningContext.state.bindingContext)
                         )
                     }
                     val composed = StackValue.field(

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/LambdaInfo.kt

@@ -230,7 +230,8 @@ class PsiExpressionLambda(
 
     private val labels: Set<String>
 
-    private var closure: CalculatedClosure
+    var closure: CalculatedClosure
+        private set
 
     init {
         val bindingContext = typeMapper.bindingContext
@@ -307,6 +308,4 @@ class PsiExpressionLambda(
 
     val isPropertyReference: Boolean
         get() = propertyReferenceInfo != null
-
-    val captureVariables = closure.captureVariables
 }

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/coroutines/CoroutineTransformer.kt

@@ -10,13 +10,13 @@ import org.jetbrains.kotlin.codegen.AsmUtil.CAPTURED_THIS_FIELD
 import org.jetbrains.kotlin.codegen.ClassBuilder
 import org.jetbrains.kotlin.codegen.TransformationMethodVisitor
 import org.jetbrains.kotlin.codegen.coroutines.*
-import org.jetbrains.kotlin.codegen.coroutines.getLastParameterIndex
-import org.jetbrains.kotlin.codegen.coroutines.replaceFakeContinuationsWithRealOnes
 import org.jetbrains.kotlin.codegen.inline.*
 import org.jetbrains.kotlin.codegen.optimization.common.asSequence
 import org.jetbrains.kotlin.codegen.optimization.common.findPreviousOrNull
 import org.jetbrains.kotlin.codegen.optimization.transformer.MethodTransformer
 import org.jetbrains.kotlin.config.isReleaseCoroutines
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin
 import org.jetbrains.kotlin.utils.addToStdlib.cast
@@ -98,7 +98,7 @@ class CoroutineTransformer(
                 ArrayUtil.toStringArray(node.exceptions)
             )
         ) {
-            val stateMachineBuilder = surroundNoinlineCallsWithMarkersIfNeeded(
+            val stateMachineBuilder = surroundNoinlineCallsWithMarkers(
                 node,
                 CoroutineTransformerMethodVisitor(
                     createNewMethodFrom(node, name), node.access, name, node.desc, null, null,
@@ -137,7 +137,7 @@ class CoroutineTransformer(
                 ArrayUtil.toStringArray(node.exceptions)
             )
         ) {
-            val stateMachineBuilder = surroundNoinlineCallsWithMarkersIfNeeded(
+            val stateMachineBuilder = surroundNoinlineCallsWithMarkers(
                 node,
                 CoroutineTransformerMethodVisitor(
                     createNewMethodFrom(node, name), node.access, name, node.desc, null, null,
@@ -165,16 +165,34 @@ class CoroutineTransformer(
         }
     }
 
-    private fun surroundNoinlineCallsWithMarkersIfNeeded(node: MethodNode, delegate: MethodVisitor): MethodVisitor =
-        if (capturedParams.any { it.functionalArgument?.isSuspendLambda() == true })
-            SurroundSuspendLambdaCallsWithSuspendMarkersMethodVisitor(
-                delegate, node.access, node.name, node.desc, classBuilder.thisName, this::isCapturedSuspendLambda
-            )
-        else
-            delegate
+    private fun surroundNoinlineCallsWithMarkers(node: MethodNode, delegate: MethodVisitor): MethodVisitor =
+        SurroundSuspendLambdaCallsWithSuspendMarkersMethodVisitor(
+            delegate, node.access, node.name, node.desc, classBuilder.thisName, this::fieldIsCapturedSuspendLambda
+        )
 
-    private fun isCapturedSuspendLambda(field: FieldInsnNode): Boolean =
-        capturedParams.find { it.newFieldName == field.name }?.functionalArgument?.isSuspendLambda() == true
+    private fun fieldIsCapturedSuspendLambda(field: FieldInsnNode): Boolean =
+        capturedParams.find { it.newFieldName == field.name }?.let { it.functionalArgument?.isSuspendLambda() == true }
+            ?: isSuspendLambdaCapturedByOuterObjectOrLambda(field)
+
+    // We cannot find the lambda in captured parameters: it came from object outside of the our reach:
+    // this can happen when the lambda capture by non-transformed closure:
+    //   inline fun inlineMe(crossinline c: suspend() -> Unit) = suspend { c() }
+    //   inline fun inlineMe2(crossinline c: suspend() -> Unit) = suspend { inlineMe { c() }() }
+    // Suppose, we inline inlineMe into inlineMe2: the only knowledge we have about inlineMe$1 is captured receiver (this$0)
+    // Thus, transformed lambda from inlineMe, inlineMe3$$inlined$inlineMe2$1 contains the following bytecode
+    //   ALOAD 0
+    //   GETFIELD inlineMe2$1$invokeSuspend$$inlined$inlineMe$1.this$0 : LScratchKt$inlineMe2$1;
+    //   GETFIELD inlineMe2$1.$c : Lkotlin/jvm/functions/Function1;
+    // Since inlineMe2's lambda is outside of reach of the inliner, find crossinline parameter from compilation context:
+    private fun isSuspendLambdaCapturedByOuterObjectOrLambda(field: FieldInsnNode): Boolean {
+        val functionDescriptor = inliningContext.root.sourceCompilerForInline.compilationContextFunctionDescriptor
+        val classDescriptor = functionDescriptor.findContainingClassOrLambda() ?: return false
+        return isCapturedSuspendLambda(classDescriptor, field.name, inliningContext.state.bindingContext)
+    }
+
+    private tailrec fun DeclarationDescriptor.findContainingClassOrLambda(): ClassDescriptor? =
+        if (containingDeclaration is ClassDescriptor) containingDeclaration as ClassDescriptor
+        else containingDeclaration?.findContainingClassOrLambda()
 
     private fun createNewMethodFrom(node: MethodNode, name: String): MethodVisitor {
         return classBuilder.newMethod(