JVM: refactor inline ExpressionLambda initialization

2021-05-25 12:38:34 +02:00
parent d023966054
commit 117fad2018
7 changed files with 56 additions and 97 deletions
@@ -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,
@@ -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(
@@ -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
 }
@@ -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()

@@ -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(
@@ -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
@@ -61,7 +61,7 @@ class IrInlineCodegen(

        when (next) {
            is IrExpressionLambdaImpl -> next.reference.getArgumentsWithIr().forEachIndexed { index, (_, ir) ->
-                putCapturedValueOnStack(ir, next.capturedParamsInDesc[index], index)
+                putCapturedValueOnStack(ir, next.capturedVars[index].type, index)
            }
            is IrDefaultLambda -> rememberCapturedForDefaultLambda(next)
            else -> throw RuntimeException("Unknown lambda: $next")
@@ -87,7 +87,7 @@ class IrInlineCodegen(
            val boundReceiver = irReference.extensionReceiver
            if (boundReceiver != null) {
                activeLambda = lambdaInfo
-                putCapturedValueOnStack(boundReceiver, lambdaInfo.capturedParamsInDesc.single(), 0)
+                putCapturedValueOnStack(boundReceiver, lambdaInfo.capturedVars.single().type, 0)
                activeLambda = null
            }
        } else {
@@ -185,7 +185,11 @@ class IrExpressionLambdaImpl(
    override val isBoundCallableReference: Boolean
        get() = reference.extensionReceiver != null

-    override val isSuspend: Boolean = function.isSuspend
+    override val isSuspend: Boolean
+        get() = reference.symbol.owner.isSuspend
+
+    override val hasDispatchReceiver: Boolean
+        get() = false

    // This name doesn't actually matter: it is used internally to tell this lambda's captured
    // arguments apart from any other scope's. So long as it's unique, any value is fine.
@@ -194,44 +198,29 @@ class IrExpressionLambdaImpl(
    override val lambdaClassType: Type = codegen.context.getLocalClassType(reference)
        ?: throw AssertionError("callable reference ${reference.dump()} has no name in context")

-    private val capturedParameters: Map<CapturedParamDesc, IrValueParameter> =
-        reference.getArgumentsWithIr().associate { (param, _) ->
-            capturedParamDesc(param.name.asString(), codegen.typeMapper.mapType(param.type)) to param
-        }
-
-    override val capturedVars: List<CapturedParamDesc> = capturedParameters.keys.toList()
-
-    private val loweredMethod = codegen.methodSignatureMapper.mapAsmMethod(function)
-
-    private val captureParameterIndices: Pair<Int, Int>
-        get() = when {
-            isBoundCallableReference -> 0 to 1 // (bound receiver, real parameters...)
-            isExtensionLambda -> 1 to capturedVars.size + 1 // (unbound receiver, captures..., real parameters...)
-            else -> 0 to capturedVars.size // (captures..., real parameters...)
-        }
-
-    val capturedParamsInDesc: List<Type> =
-        captureParameterIndices.let { (from, to) -> loweredMethod.argumentTypes.take(to).drop(from) }
-
-    override val invokeMethod: Method = loweredMethod.let {
-        val (startCapture, endCapture) = captureParameterIndices
-        Method(it.name, it.returnType, (it.argumentTypes.take(startCapture) + it.argumentTypes.drop(endCapture)).toTypedArray())
-    }
-
+    override val capturedVars: List<CapturedParamDesc>
+    override val invokeMethod: Method
    override val invokeMethodParameters: List<KotlinType?>
-        get() {
-            val allParameters = function.explicitParameters.map { it.type.toIrBasedKotlinType() }
-            val (startCapture, endCapture) = captureParameterIndices
-            return allParameters.take(startCapture) + allParameters.drop(endCapture)
-        }
-
    override val invokeMethodReturnType: KotlinType
-        get() = function.returnType.toIrBasedKotlinType()

-    override val hasDispatchReceiver: Boolean = false
-
-    override fun isCapturedSuspend(desc: CapturedParamDesc): Boolean =
-        capturedParameters[desc]?.let { it.isInlineParameter() && it.type.isSuspendFunctionTypeOrSubtype() } == true
+    init {
+        val asmMethod = codegen.methodSignatureMapper.mapAsmMethod(function)
+        val capturedParameters = reference.getArgumentsWithIr()
+        val (startCapture, endCapture) = when {
+            isBoundCallableReference -> 0 to 1 // (bound receiver, real parameters...)
+            isExtensionLambda -> 1 to capturedParameters.size + 1 // (unbound receiver, captures..., real parameters...)
+            else -> 0 to capturedParameters.size // (captures..., real parameters...)
+        }
+        capturedVars = capturedParameters.mapIndexed { index, (parameter, _) ->
+            val isSuspend = parameter.isInlineParameter() && parameter.type.isSuspendFunctionTypeOrSubtype()
+            capturedParamDesc(parameter.name.asString(), asmMethod.argumentTypes[startCapture + index], isSuspend)
+        }
+        val freeParameters = function.explicitParameters.let { it.take(startCapture) + it.drop(endCapture) }
+        val freeAsmParameters = asmMethod.argumentTypes.let { it.take(startCapture) + it.drop(endCapture) }
+        invokeMethod = Method(asmMethod.name, asmMethod.returnType, freeAsmParameters.toTypedArray())
+        invokeMethodParameters = freeParameters.map { it.type.toIrBasedKotlinType() }
+        invokeMethodReturnType = function.returnType.toIrBasedKotlinType()
+    }
 }

 class IrDefaultLambda(