Report suspension point in monitor error on crossinline lambdas

#KT-27258: Fixed

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

@@ -63,7 +63,7 @@ class CoroutineTransformerMethodVisitor(
     private val languageVersionSettings: LanguageVersionSettings,
     private val sourceFile: String,
     // These two are needed to report diagnostics about suspension points inside critical section
-    private val element: KtElement?,
+    private val element: KtElement,
     private val diagnostics: DiagnosticSink,
     // It's only matters for named functions, may differ from '!isStatic(access)' in case of DefaultImpls
     private val needDispatchReceiver: Boolean = false,
@@ -231,7 +231,7 @@ class CoroutineTransformerMethodVisitor(
                     sourceFile,
                     findSuspensionPointLineNumber(suspensionPoint)?.line ?: -1
                 )
-                element?.let { diagnostics.report(ErrorsJvm.SUSPENSION_POINT_INSIDE_MONITOR.on(it, "$stackTraceElement")) }
+                diagnostics.report(ErrorsJvm.SUSPENSION_POINT_INSIDE_MONITOR.on(element, "$stackTraceElement"))
                 return
             }
         }

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

@@ -21,8 +21,10 @@ import org.jetbrains.kotlin.metadata.jvm.JvmProtoBuf
 import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmProtoBufUtil
 import org.jetbrains.kotlin.metadata.jvm.serialization.JvmStringTable
 import org.jetbrains.kotlin.protobuf.MessageLite
+import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin.Companion.NO_ORIGIN
+import org.jetbrains.kotlin.utils.addToStdlib.cast
 import org.jetbrains.org.objectweb.asm.*
 import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter
 import org.jetbrains.org.objectweb.asm.tree.*
@@ -131,10 +133,9 @@ class AnonymousObjectTransformer(
                 constructor!!, allCapturedParamBuilder, constructorParamBuilder,transformationInfo, parentRemapper
         )
 
+        val crossinlineSuspendElement = capturedCrossinlineSuspendElement()
         val capturesCrossinlineSuspend = (!inliningContext.isInliningLambda || inliningContext.isContinuation) &&
-                inliningContext.expressionMap.values.any { lambda ->
-                    lambda is PsiExpressionLambda && lambda.isCrossInline && lambda.invokeMethodDescriptor.isSuspend
-                }
+                crossinlineSuspendElement != null
 
         val deferringMethods = ArrayList<DeferredMethodVisitor>()
 
@@ -158,11 +159,18 @@ class AnonymousObjectTransformer(
                 capturesCrossinlineSuspend && !inliningContext.isContinuation && continuationClassName != null
 
             val deferringVisitor =
-                when {
-                    generateStateMachineForLambda -> newStateMachineForLambda(classBuilder, next)
-                    generateStateMachineForNamedFunction -> newStateMachineForNamedFunction(classBuilder, next, continuationClassName!!)
-                    else -> newMethod(classBuilder, next)
-                }
+                if (crossinlineSuspendElement != null) {
+                    when {
+                        generateStateMachineForLambda -> newStateMachineForLambda(classBuilder, next, crossinlineSuspendElement)
+                        generateStateMachineForNamedFunction -> newStateMachineForNamedFunction(
+                            classBuilder,
+                            next,
+                            continuationClassName!!,
+                            crossinlineSuspendElement
+                        )
+                        else -> newMethod(classBuilder, next)
+                    }
+                } else newMethod(classBuilder, next)
             val funResult = inlineMethodAndUpdateGlobalResult(parentRemapper, deferringVisitor, next, allCapturedParamBuilder, false)
 
             val returnType = Type.getReturnType(next.desc)
@@ -227,6 +235,10 @@ class AnonymousObjectTransformer(
         return transformationResult
     }
 
+    private fun capturedCrossinlineSuspendElement(): KtExpression? = inliningContext.expressionMap.values.find { lambda ->
+        lambda is PsiExpressionLambda && lambda.isCrossInline && lambda.invokeMethodDescriptor.isSuspend
+    }?.cast<PsiExpressionLambda>()?.functionWithBodyOrCallableReference
+
     private fun writeTransformedMetadata(header: KotlinClassHeader, classBuilder: ClassBuilder) {
         writeKotlinMetadata(classBuilder, state, header.kind, header.extraInt) action@ { av ->
             val (newProto, newStringTable) = transformMetadata(header) ?: run {
@@ -435,7 +447,7 @@ class AnonymousObjectTransformer(
         }
     }
 
-    private fun newStateMachineForLambda(builder: ClassBuilder, original: MethodNode): DeferredMethodVisitor {
+    private fun newStateMachineForLambda(builder: ClassBuilder, original: MethodNode, element: KtExpression): DeferredMethodVisitor {
         return DeferredMethodVisitor(
             MethodNode(
                 original.access, original.name, original.desc, original.signature,
@@ -448,7 +460,7 @@ class AnonymousObjectTransformer(
                     ArrayUtil.toStringArray(original.exceptions)
                 ), original.access, original.name, original.desc, null, null,
                 obtainClassBuilderForCoroutineState = { builder },
-                element = null,
+                element = element,
                 diagnostics = state.diagnostics,
                 languageVersionSettings = languageVersionSettings,
                 shouldPreserveClassInitialization = state.constructorCallNormalizationMode.shouldPreserveClassInitialization,
@@ -463,7 +475,8 @@ class AnonymousObjectTransformer(
     private fun newStateMachineForNamedFunction(
         builder: ClassBuilder,
         original: MethodNode,
-        continuationClassName: String
+        continuationClassName: String,
+        element: KtExpression
     ): DeferredMethodVisitor {
         assert(inliningContext is RegeneratedClassContext)
         return DeferredMethodVisitor(
@@ -478,7 +491,7 @@ class AnonymousObjectTransformer(
                     ArrayUtil.toStringArray(original.exceptions)
                 ), original.access, original.name, original.desc, null, null,
                 obtainClassBuilderForCoroutineState = { (inliningContext as RegeneratedClassContext).continuationBuilders[continuationClassName]!! },
-                element = null,
+                element = element,
                 diagnostics = state.diagnostics,
                 languageVersionSettings = languageVersionSettings,
                 shouldPreserveClassInitialization = state.constructorCallNormalizationMode.shouldPreserveClassInitialization,

compiler/testData/compileKotlinAgainstCustomBinaries/suspensionPointInMonitor/library/a.kt

@@ -15,3 +15,12 @@ inline fun monitorInFinally(a: () -> Unit, b: () -> Unit) {
         }
     }
 }
+
+inline fun withCrossinline(crossinline a: suspend () -> Unit): suspend () -> Unit {
+    val c : suspend () -> Unit = {
+        inlineMe {
+            a()
+        }
+    }
+    return c
+}

compiler/testData/compileKotlinAgainstCustomBinaries/suspensionPointInMonitor/output.txt

@@ -4,4 +4,7 @@ compiler/testData/compileKotlinAgainstCustomBinaries/suspensionPointInMonitor/so
 compiler/testData/compileKotlinAgainstCustomBinaries/suspensionPointInMonitor/source.kt:14:13: error: a suspension point at SourceKt$test$2.invokeSuspend(source.kt:17) is inside a critical section
     builder {
             ^
+compiler/testData/compileKotlinAgainstCustomBinaries/suspensionPointInMonitor/source.kt:24:25: error: a suspension point at SourceKt$test$3$invokeSuspend$$inlined$withCrossinline$2.invokeSuspend(a.kt:31) is inside a critical section
+        withCrossinline {
+                        ^
 COMPILATION_ERROR

compiler/testData/compileKotlinAgainstCustomBinaries/suspensionPointInMonitor/source.kt

@@ -18,6 +18,14 @@ fun test() {
         )
     }
 
+    builder {
+        withCrossinline {}
+
+        withCrossinline {
+            suspensionPoint()
+        }
+    }
+
     synchronized(lock) {
         builder {
             suspensionPoint()