JVM_IR: Support crossinline suspend lambdas

The main idea is the following: since we need to generate (fake)continuations before inlining, we move IrClasses of suspend lambdas and continuation classes of named functions into the functions. Thus, it allows the codegen to generate them prior to inlining and the inliner will happily transform them for us. Because of that, lowerings which transform call-site function are likely to change reference to lowered suspend lambdas or functions. Hence, do not rely on references to lowered suspend lambdas or functions, instead, rely on attributes. Do not generate continuation for inline suspend lambdas. Previously, inline suspend lambdas were treated like suspend functions, thus we generated continuations for them. Now we just do not treat them as suspend functions or lambdas during AddContinuationLowering. We should add continuation parameter to them, however. Do not generate secondary constructor for suspend lambdas, otherwise, the inliner is unable to transform them (it requires only one constructor to be present). Generate continuation classes for suspend functions as first statement inside the function. This enables suspend functions in local object inside inline functions. Since we already have attributes inside suspend named functions, we just reuse them to generate continuation class names. This allows us to close the gap between code generated by old back-end and the new one. If a suspend named function captures crossinline lambda, we should generate a template for inliner: a copy of the function without state-machine and a continuation constructor call. The call is needed so the inliner transforms the continuation as well. Refactor CoroutineTransformerMethodVisitor, so it no longer depends on PSI.
2019-10-22 20:02:34 +03:00
parent 3b37f6bd32
commit a1448ebb37
55 changed files with 680 additions and 366 deletions
@@ -6,6 +6,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
@@ -17,6 +18,7 @@ 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
 import org.jetbrains.kotlin.codegen.serialization.JvmSerializerExtension
+import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.config.isReleaseCoroutines
 import org.jetbrains.kotlin.descriptors.*
@@ -32,6 +34,7 @@ import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
 import org.jetbrains.kotlin.resolve.descriptorUtil.module
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes
+import org.jetbrains.kotlin.resolve.jvm.diagnostics.ErrorsJvm
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.OtherOrigin
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature
@@ -470,8 +473,9 @@ class CoroutineCodegenForLambda private constructor(
                    val stateMachineBuilder = CoroutineTransformerMethodVisitor(
                        mv, access, name, desc, null, null,
                        obtainClassBuilderForCoroutineState = { v },
-                        element = element,
-                        diagnostics = state.diagnostics,
+                        reportSuspensionPointInsideMonitor = { reportSuspensionPointInsideMonitor(element, state, it) },
+                        lineNumber = CodegenUtil.getLineNumberForElement(element, false) ?: 0,
+                        sourceFile = element.containingKtFile.name,
                        shouldPreserveClassInitialization = constructorCallNormalizationMode.shouldPreserveClassInitialization,
                        containingClassInternalName = v.thisName,
                        isForNamedFunction = false,
@@ -768,3 +772,7 @@ private object FailingFunctionGenerationStrategy : FunctionGenerationStrategy()
        error("This functions must not be called")
    }
 }
+
+fun reportSuspensionPointInsideMonitor(element: KtElement, state: GenerationState, stackTraceElement: String) {
+    state.diagnostics.report(ErrorsJvm.SUSPENSION_POINT_INSIDE_MONITOR.on(element, stackTraceElement))
+}
@@ -60,13 +60,13 @@ class CoroutineTransformerMethodVisitor(
    private val isForNamedFunction: Boolean,
    private val shouldPreserveClassInitialization: Boolean,
    private val languageVersionSettings: LanguageVersionSettings,
-    // These two are needed to report diagnostics about suspension points inside critical section
-    private val element: KtElement,
-    private val diagnostics: DiagnosticSink,
    // Since tail-call optimization of functions with Unit return type relies on ability of call-site to recognize them,
    // in order to ignore return value and push Unit, when we cannot ensure this ability, for example, when the function overrides function,
    // returning Any, we need to disable tail-call optimization for these functions.
    private val disableTailCallOptimizationForFunctionReturningUnit: Boolean,
+    private val reportSuspensionPointInsideMonitor: (String) -> Unit,
+    private val lineNumber: Int,
+    private val sourceFile: String,
    // It's only matters for named functions, may differ from '!isStatic(access)' in case of DefaultImpls
    private val needDispatchReceiver: Boolean = false,
    // May differ from containingClassInternalName in case of DefaultImpls
@@ -76,8 +76,6 @@ class CoroutineTransformerMethodVisitor(
 ) : TransformationMethodVisitor(delegate, access, name, desc, signature, exceptions) {

    private val classBuilderForCoroutineState: ClassBuilder by lazy(obtainClassBuilderForCoroutineState)
-    private val lineNumber = CodegenUtil.getLineNumberForElement(element, false) ?: 0
-    private val sourceFile = element.containingKtFile.name

    private var continuationIndex = if (isForNamedFunction) -1 else 0
    private var dataIndex = if (isForNamedFunction) -1 else 1
@@ -315,7 +313,7 @@ class CoroutineTransformerMethodVisitor(
                    sourceFile,
                    findSuspensionPointLineNumber(suspensionPoint)?.line ?: -1
                )
-                diagnostics.report(ErrorsJvm.SUSPENSION_POINT_INSIDE_MONITOR.on(element, "$stackTraceElement"))
+                reportSuspensionPointInsideMonitor("$stackTraceElement")
                return
            }
        }
@@ -5,6 +5,7 @@

 package org.jetbrains.kotlin.codegen.coroutines

+import org.jetbrains.kotlin.backend.common.CodegenUtil
 import org.jetbrains.kotlin.codegen.*
 import org.jetbrains.kotlin.codegen.binding.CodegenBinding
 import org.jetbrains.kotlin.codegen.inline.addFakeContinuationConstructorCallMarker
@@ -96,8 +97,9 @@ open class SuspendFunctionGenerationStrategy(
        return CoroutineTransformerMethodVisitor(
            mv, access, name, desc, null, null, containingClassInternalName, this::classBuilderForCoroutineState,
            isForNamedFunction = true,
-            element = declaration,
-            diagnostics = state.diagnostics,
+            reportSuspensionPointInsideMonitor = { reportSuspensionPointInsideMonitor(declaration, state, it) },
+            lineNumber = CodegenUtil.getLineNumberForElement(declaration, false) ?: 0,
+            sourceFile = declaration.containingKtFile.name,
            shouldPreserveClassInitialization = constructorCallNormalizationMode.shouldPreserveClassInitialization,
            needDispatchReceiver = originalSuspendDescriptor.dispatchReceiverParameter != null,
            internalNameForDispatchReceiver = containingClassInternalNameOrNull(),
@@ -302,7 +302,8 @@ class AnonymousObjectTransformer(
                sourceNode.name,
                if (isConstructor) transformationInfo.newConstructorDescriptor else sourceNode.desc,
                inliningContext.callSiteInfo.isInlineOrInsideInline,
-                isSuspendFunctionOrLambda(sourceNode)
+                isSuspendFunctionOrLambda(sourceNode),
+                inliningContext.root.sourceCompilerForInline.inlineCallSiteInfo.lineNumber
            ), null
        )

@@ -10,5 +10,6 @@ class InlineCallSiteInfo(
    val functionName: String?,
    val functionDesc: String?,
    val isInlineOrInsideInline: Boolean,
-    val isSuspend: Boolean
+    val isSuspend: Boolean,
+    val lineNumber: Int
 )
@@ -157,6 +157,7 @@ class MethodInliner(
            ), AsmTypeRemapper(remapper, result)
        )

+        val fakeContinuationName = CoroutineTransformer.findFakeContinuationConstructorClassName(node)
        val markerShift = calcMarkerShift(parameters, node)
        val lambdaInliner = object : InlineAdapter(remappingMethodAdapter, parameters.argsSizeOnStack, sourceMapper) {
            private var transformationInfo: TransformationInfo? = null
@@ -176,9 +177,7 @@ class MethodInliner(
                        inlineCallSiteInfo
                    )
                    val transformer = transformationInfo!!.createTransformer(
-                        childInliningContext,
-                        isSameModule,
-                        CoroutineTransformer.findFakeContinuationConstructorClassName(node)
+                        childInliningContext, isSameModule, fakeContinuationName
                    )

                    val transformResult = transformer.doTransform(nodeRemapper)
@@ -75,6 +75,8 @@ interface SourceCompilerForInline {
    val compilationContextFunctionDescriptor: FunctionDescriptor

    fun getContextLabels(): Set<String>
+
+    fun reportSuspensionPointInsideMonitor(stackTraceElement: String)
 }


@@ -120,7 +122,8 @@ class PsiSourceCompilerForInline(private val codegen: ExpressionCodegen, overrid
                signature.asmMethod.name,
                signature.asmMethod.descriptor,
                compilationContextFunctionDescriptor.isInlineOrInsideInline(),
-                compilationContextFunctionDescriptor.isSuspend
+                compilationContextFunctionDescriptor.isSuspend,
+                CodegenUtil.getLineNumberForElement(callElement, false) ?: 0
            )
        }

@@ -365,6 +368,10 @@ class PsiSourceCompilerForInline(private val codegen: ExpressionCodegen, overrid
        )
    }

+    override fun reportSuspensionPointInsideMonitor(stackTraceElement: String) {
+        org.jetbrains.kotlin.codegen.coroutines.reportSuspensionPointInsideMonitor(callElement, state, stackTraceElement)
+    }
+
    companion object {
        fun getContext(
            descriptor: DeclarationDescriptor,
@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.codegen.inline.coroutines

 import com.intellij.util.ArrayUtil
+import org.jetbrains.kotlin.backend.common.CodegenUtil
 import org.jetbrains.kotlin.codegen.AsmUtil.CAPTURED_THIS_FIELD
 import org.jetbrains.kotlin.codegen.ClassBuilder
 import org.jetbrains.kotlin.codegen.TransformationMethodVisitor
@@ -66,29 +67,20 @@ class CoroutineTransformer(
            if (state.languageVersionSettings.isReleaseCoroutines()) superClassName.endsWith("ContinuationImpl")
            else methods.any { it.name == "getLabel" }

-    private fun crossinlineLambda(): PsiExpressionLambda? = inliningContext.expressionMap.values.find {
-        it is PsiExpressionLambda && it.isCrossInline
-    }?.cast()
-
    private fun isStateMachine(node: MethodNode): Boolean =
        node.instructions.asSequence().any { insn -> insn is LdcInsnNode && insn.cst == ILLEGAL_STATE_ERROR_MESSAGE }

    private fun isSuspendLambda(node: MethodNode) = isResumeImpl(node)

    fun newMethod(node: MethodNode): DeferredMethodVisitor {
-        // Find ANY element to report error about suspension point in monitor on.
-        val element = crossinlineLambda()?.functionWithBodyOrCallableReference
-            ?: inliningContext.root.sourceCompilerForInline.callElement as? KtElement
-            ?: error("crossinline lambda should have element")
-
        return when {
            isResumeImpl(node) -> {
                assert(!isStateMachine(node)) {
                    "Inlining/transforming state-machine"
                }
-                newStateMachineForLambda(node, element)
+                newStateMachineForLambda(node)
            }
-            isSuspendFunctionWithFakeConstructorCall(node) -> newStateMachineForNamedFunction(node, element)
+            isSuspendFunctionWithFakeConstructorCall(node) -> newStateMachineForNamedFunction(node)
            else -> error("no need to generate state maching for ${node.name}")
        }
    }
@@ -99,7 +91,7 @@ class CoroutineTransformer(

    private fun isSuspendFunctionWithFakeConstructorCall(node: MethodNode): Boolean = findFakeContinuationConstructorClassName(node) != null

-    private fun newStateMachineForLambda(node: MethodNode, element: KtElement): DeferredMethodVisitor {
+    private fun newStateMachineForLambda(node: MethodNode): DeferredMethodVisitor {
        val name = node.name.removeSuffix(FOR_INLINE_SUFFIX)
        return DeferredMethodVisitor(
            MethodNode(
@@ -107,13 +99,16 @@ class CoroutineTransformer(
                ArrayUtil.toStringArray(node.exceptions)
            )
        ) {
+            val sourceCompilerForInline = inliningContext.root.sourceCompilerForInline
            val stateMachineBuilder = surroundNoinlineCallsWithMarkers(
                node,
                CoroutineTransformerMethodVisitor(
                    createNewMethodFrom(node, name), node.access, name, node.desc, null, null,
                    obtainClassBuilderForCoroutineState = { classBuilder },
-                    element = element,
-                    diagnostics = state.diagnostics,
+                    reportSuspensionPointInsideMonitor = { sourceCompilerForInline.reportSuspensionPointInsideMonitor(it) },
+                    // TODO: this linenumbers might not be correct and since they are used only for step-over, check them.
+                    lineNumber = sourceCompilerForInline.inlineCallSiteInfo.lineNumber,
+                    sourceFile = sourceCompilerForInline.callsiteFile?.name ?: "",
                    languageVersionSettings = state.languageVersionSettings,
                    shouldPreserveClassInitialization = state.constructorCallNormalizationMode.shouldPreserveClassInitialization,
                    containingClassInternalName = classBuilder.thisName,
@@ -137,7 +132,7 @@ class CoroutineTransformer(
        }
    }

-    private fun newStateMachineForNamedFunction(node: MethodNode, element: KtElement): DeferredMethodVisitor {
+    private fun newStateMachineForNamedFunction(node: MethodNode): DeferredMethodVisitor {
        val name = node.name.removeSuffix(FOR_INLINE_SUFFIX)
        val continuationClassName = findFakeContinuationConstructorClassName(node)
        assert(inliningContext is RegeneratedClassContext)
@@ -149,20 +144,23 @@ class CoroutineTransformer(
        ) {
            // If the node already has state-machine, it is safer to generate state-machine.
            val disableTailCallOptimization = methods.find { it.name == name && it.desc == node.desc }?.let { isStateMachine(it) } ?: false
+            val sourceCompilerForInline = inliningContext.root.sourceCompilerForInline
            val stateMachineBuilder = surroundNoinlineCallsWithMarkers(
                node,
                CoroutineTransformerMethodVisitor(
                    createNewMethodFrom(node, name), node.access, name, node.desc, null, null,
                    obtainClassBuilderForCoroutineState = { (inliningContext as RegeneratedClassContext).continuationBuilders[continuationClassName]!! },
-                    element = element,
-                    diagnostics = state.diagnostics,
+                    reportSuspensionPointInsideMonitor = { sourceCompilerForInline.reportSuspensionPointInsideMonitor(it) },
+                    lineNumber = sourceCompilerForInline.inlineCallSiteInfo.lineNumber,
+                    sourceFile = sourceCompilerForInline.callsiteFile?.name ?: "",
                    languageVersionSettings = state.languageVersionSettings,
                    shouldPreserveClassInitialization = state.constructorCallNormalizationMode.shouldPreserveClassInitialization,
                    containingClassInternalName = classBuilder.thisName,
                    isForNamedFunction = true,
                    needDispatchReceiver = true,
                    internalNameForDispatchReceiver = classBuilder.thisName,
-                    disableTailCallOptimizationForFunctionReturningUnit = disableTailCallOptimization
+                    disableTailCallOptimizationForFunctionReturningUnit = disableTailCallOptimization,
+                    putContinuationParameterToLvt = !state.isIrBackend
                )
            )

@@ -233,7 +231,7 @@ class CoroutineTransformer(

    fun oldContinuationFrom(method: MethodNode): String? =
        methods.find { it.name == method.name + FOR_INLINE_SUFFIX && it.desc == method.desc }
-                ?.let { findFakeContinuationConstructorClassName(it) }
+            ?.let { findFakeContinuationConstructorClassName(it) }

    companion object {
        fun findFakeContinuationConstructorClassName(node: MethodNode): String? {
@@ -431,7 +431,7 @@ fun addSuspendMarker(v: InstructionAdapter, isStartNotEnd: Boolean) {
    v.emitInlineMarker(if (isStartNotEnd) INLINE_MARKER_BEFORE_SUSPEND_ID else INLINE_MARKER_AFTER_SUSPEND_ID)
 }

-internal fun addFakeContinuationConstructorCallMarker(v: InstructionAdapter, isStartNotEnd: Boolean) {
+fun addFakeContinuationConstructorCallMarker(v: InstructionAdapter, isStartNotEnd: Boolean) {
    v.emitInlineMarker(if (isStartNotEnd) INLINE_MARKER_BEFORE_FAKE_CONTINUATION_CONSTRUCTOR_CALL else INLINE_MARKER_AFTER_FAKE_CONTINUATION_CONSTRUCTOR_CALL)
 }