JVM: refactor inline ExpressionLambda initialization

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

@@ -7,9 +7,7 @@ package org.jetbrains.kotlin.codegen.coroutines
 
 import com.intellij.util.ArrayUtil
 import org.jetbrains.kotlin.backend.common.CodegenUtil
-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
@@ -601,21 +599,6 @@ class CoroutineCodegenForLambda private constructor(
     }
 }
 
-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
-    }
-    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(
     delegate: MethodVisitor,
     access: Int,

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

@@ -547,7 +547,7 @@ class AnonymousObjectTransformer(
                         capturedParamBuilder.addCapturedParam(desc, desc.fieldName, !capturedOuterThisTypes.add(desc.type.className))
                     else
                         capturedParamBuilder.addCapturedParam(desc, addUniqueField(desc.fieldName + INLINE_TRANSFORMATION_SUFFIX), false)
-                    if (info is ExpressionLambda && info.isCapturedSuspend(desc)) {
+                    if (desc.isSuspend) {
                         recapturedParamInfo.functionalArgument = NonInlineableArgumentForInlineableParameterCalledInSuspend
                     }
                     val composed = StackValue.field(

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

@@ -18,8 +18,6 @@ package org.jetbrains.kotlin.codegen.inline
 
 import org.jetbrains.org.objectweb.asm.Type
 
-class CapturedParamDesc(private val containingLambdaType: Type, val fieldName: String, val type: Type) {
-
-    val containingLambdaName: String
-        get() = containingLambdaType.internalName
+class CapturedParamDesc(containingLambdaType: Type, val fieldName: String, val type: Type, val isSuspend: Boolean = false) {
+    val containingLambdaName: String = containingLambdaType.internalName
 }

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

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.codegen.inline
 
 import org.jetbrains.kotlin.codegen.AsmUtil
-import org.jetbrains.kotlin.codegen.context.EnclosedValueDescriptor
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes.*
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.util.OperatorNameConventions
@@ -53,7 +52,7 @@ abstract class LambdaInfo(@JvmField val isCrossInline: Boolean) : FunctionalArgu
             val field = remapper.findField(FieldInsnNode(0, info.containingLambdaName, info.fieldName, ""))
                 ?: error("Captured field not found: " + info.containingLambdaName + "." + info.fieldName)
             val recapturedParamInfo = builder.addCapturedParam(field, info.fieldName)
-            if (this is ExpressionLambda && isCapturedSuspend(info)) {
+            if (info.isSuspend) {
                 recapturedParamInfo.functionalArgument = NonInlineableArgumentForInlineableParameterCalledInSuspend
             }
         }
@@ -61,14 +60,9 @@ abstract class LambdaInfo(@JvmField val isCrossInline: Boolean) : FunctionalArgu
         return builder.buildParameters()
     }
 
-
     companion object {
-        fun LambdaInfo.getCapturedParamInfo(descriptor: EnclosedValueDescriptor): CapturedParamDesc {
-            return capturedParamDesc(descriptor.fieldName, descriptor.type)
-        }
-
-        fun LambdaInfo.capturedParamDesc(fieldName: String, fieldType: Type): CapturedParamDesc {
-            return CapturedParamDesc(lambdaClassType, fieldName, fieldType)
+        fun LambdaInfo.capturedParamDesc(fieldName: String, fieldType: Type, isSuspend: Boolean): CapturedParamDesc {
+            return CapturedParamDesc(lambdaClassType, fieldName, fieldType, isSuspend)
         }
     }
 }
@@ -81,8 +75,6 @@ abstract class ExpressionLambda(isCrossInline: Boolean) : LambdaInfo(isCrossInli
         node = sourceCompiler.generateLambdaBody(this, reifiedTypeParametersUsages)
         node.node.preprocessSuspendMarkers(forInline = true, keepFakeContinuation = false)
     }
-
-    abstract fun isCapturedSuspend(desc: CapturedParamDesc): Boolean
 }
 
 abstract class DefaultLambda(
@@ -123,11 +115,11 @@ abstract class DefaultLambda(
             if (isFunctionReference || isPropertyReference)
                 constructor?.desc?.let { Type.getArgumentTypes(it) }?.singleOrNull()?.let {
                     originalBoundReceiverType = it
-                    listOf(capturedParamDesc(AsmUtil.RECEIVER_PARAMETER_NAME, AsmUtil.boxType(it)))
+                    listOf(capturedParamDesc(AsmUtil.RECEIVER_PARAMETER_NAME, AsmUtil.boxType(it), isSuspend = false))
                 } ?: emptyList()
             else
                 constructor?.findCapturedFieldAssignmentInstructions()?.map { fieldNode ->
-                    capturedParamDesc(fieldNode.name, Type.getType(fieldNode.desc))
+                    capturedParamDesc(fieldNode.name, Type.getType(fieldNode.desc), isSuspend = false)
                 }?.toList() ?: emptyList()
         isBoundCallableReference = (isFunctionReference || isPropertyReference) && capturedVars.isNotEmpty()
 

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

@@ -11,9 +11,7 @@ import org.jetbrains.kotlin.codegen.*
 import org.jetbrains.kotlin.codegen.DescriptorAsmUtil.getMethodAsmFlags
 import org.jetbrains.kotlin.codegen.binding.CalculatedClosure
 import org.jetbrains.kotlin.codegen.binding.CodegenBinding
-import org.jetbrains.kotlin.codegen.context.EnclosedValueDescriptor
 import org.jetbrains.kotlin.codegen.coroutines.getOrCreateJvmSuspendFunctionView
-import org.jetbrains.kotlin.codegen.coroutines.isCapturedSuspendLambda
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper
 import org.jetbrains.kotlin.config.LanguageVersionSettings
@@ -299,42 +297,27 @@ class PsiExpressionLambda(
         arrayListOf<CapturedParamDesc>().apply {
             val captureThis = closure.capturedOuterClassDescriptor
             if (captureThis != null) {
-                val kotlinType = captureThis.defaultType
-                val type = typeMapper.mapType(kotlinType)
-                val descriptor = EnclosedValueDescriptor(
-                    AsmUtil.CAPTURED_THIS_FIELD, null,
-                    StackValue.field(type, lambdaClassType, AsmUtil.CAPTURED_THIS_FIELD, false, StackValue.LOCAL_0),
-                    type, kotlinType
-                )
-                add(getCapturedParamInfo(descriptor))
+                add(capturedParamDesc(AsmUtil.CAPTURED_THIS_FIELD, typeMapper.mapType(captureThis.defaultType), isSuspend = false))
             }
 
             val capturedReceiver = closure.capturedReceiverFromOuterContext
             if (capturedReceiver != null) {
+                val fieldName = closure.getCapturedReceiverFieldName(typeMapper.bindingContext, languageVersionSettings)
                 val type = typeMapper.mapType(capturedReceiver).let {
                     if (isBoundCallableReference) AsmUtil.boxType(it) else it
                 }
-
-                val fieldName = closure.getCapturedReceiverFieldName(typeMapper.bindingContext, languageVersionSettings)
-                val descriptor = EnclosedValueDescriptor(
-                    fieldName, null,
-                    StackValue.field(type, capturedReceiver, lambdaClassType, fieldName, false, StackValue.LOCAL_0),
-                    type, capturedReceiver
-                )
-                add(getCapturedParamInfo(descriptor))
+                add(capturedParamDesc(fieldName, type, isSuspend = false))
             }
 
-            closure.captureVariables.values.forEach { descriptor ->
-                add(getCapturedParamInfo(descriptor))
+            closure.captureVariables.forEach { (parameter, value) ->
+                val isSuspend = parameter is ValueParameterDescriptor && parameter.type.isSuspendFunctionTypeOrSubtype
+                add(capturedParamDesc(value.fieldName, value.type, isSuspend))
             }
         }
     }
 
     val isPropertyReference: Boolean
         get() = propertyReferenceInfo != null
-
-    override fun isCapturedSuspend(desc: CapturedParamDesc): Boolean =
-        isCapturedSuspendLambda(closure, desc.fieldName, typeMapper.bindingContext)
 }
 
 class PsiDefaultLambda(

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

@@ -6,7 +6,9 @@
 package org.jetbrains.kotlin.codegen.inline.coroutines
 
 import com.intellij.util.ArrayUtil
+import org.jetbrains.kotlin.builtins.isSuspendFunctionTypeOrSubtype
 import org.jetbrains.kotlin.codegen.ClassBuilder
+import org.jetbrains.kotlin.codegen.binding.CodegenBinding
 import org.jetbrains.kotlin.codegen.coroutines.*
 import org.jetbrains.kotlin.codegen.inline.*
 import org.jetbrains.kotlin.codegen.optimization.common.asSequence
@@ -15,6 +17,7 @@ import org.jetbrains.kotlin.config.JVMConfigurationKeys
 import org.jetbrains.kotlin.config.isReleaseCoroutines
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin
 import org.jetbrains.org.objectweb.asm.MethodVisitor
@@ -250,7 +253,18 @@ fun FieldInsnNode.isSuspendLambdaCapturedByOuterObjectOrLambda(inliningContext:
     while (container !is ClassDescriptor) {
         container = container.containingDeclaration ?: return false
     }
-    return isCapturedSuspendLambda(container, name, inliningContext.state.bindingContext)
+    val bindingContext = inliningContext.state.bindingContext
+    var classDescriptor: ClassDescriptor? = container
+    while (classDescriptor != null) {
+        val closure = bindingContext[CodegenBinding.CLOSURE, classDescriptor] ?: return false
+        for ((param, value) in closure.captureVariables) {
+            if (param is ValueParameterDescriptor && value.fieldName == name) {
+                return param.type.isSuspendFunctionTypeOrSubtype
+            }
+        }
+        classDescriptor = closure.capturedOuterClassDescriptor
+    }
+    return false
 }
 
 // Interpreter, that keeps track of captured functional arguments