Rework check of suspension point inside critical section

There is a trade-off between robustness of check and accuracy of the
diagnostic: the previous version, which works on generation, was too
fragile and lead to false-positives. Now we check on state machine
generation. However, since we do not have PSI for call, we can only
report diagnostic on whole suspend function or suspend lambda.
Additionally, the state machine is generated on crossinline suspend
lambdas regeneration and thus we do not have the PSI for the lambda as
well!

 #KT-27130 Fixed
 #KT-27258 Open

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

@@ -448,7 +448,8 @@ class CoroutineCodegenForLambda private constructor(
                     return CoroutineTransformerMethodVisitor(
                         mv, access, name, desc, null, null,
                         obtainClassBuilderForCoroutineState = { v },
-                        lineNumber = CodegenUtil.getLineNumberForElement(element, false) ?: 0,
+                        element = element,
+                        diagnostics = state.diagnostics,
                         shouldPreserveClassInitialization = constructorCallNormalizationMode.shouldPreserveClassInitialization,
                         containingClassInternalName = v.thisName,
                         isForNamedFunction = false,

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.codegen.coroutines
 
 import com.intellij.util.containers.Stack
+import org.jetbrains.kotlin.backend.common.CodegenUtil
 import org.jetbrains.kotlin.codegen.AsmUtil
 import org.jetbrains.kotlin.codegen.ClassBuilder
 import org.jetbrains.kotlin.codegen.StackValue
@@ -19,7 +20,10 @@ import org.jetbrains.kotlin.codegen.optimization.fixStack.top
 import org.jetbrains.kotlin.codegen.optimization.transformer.MethodTransformer
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.config.isReleaseCoroutines
+import org.jetbrains.kotlin.diagnostics.DiagnosticSink
+import org.jetbrains.kotlin.psi.KtElement
 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.utils.addToStdlib.cast
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
@@ -56,9 +60,11 @@ class CoroutineTransformerMethodVisitor(
     obtainClassBuilderForCoroutineState: () -> ClassBuilder,
     private val isForNamedFunction: Boolean,
     private val shouldPreserveClassInitialization: Boolean,
-    private val lineNumber: Int,
     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 diagnostics: DiagnosticSink,
     // 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
@@ -68,6 +74,7 @@ class CoroutineTransformerMethodVisitor(
 ) : TransformationMethodVisitor(delegate, access, name, desc, signature, exceptions) {
 
     private val classBuilderForCoroutineState: ClassBuilder by lazy(obtainClassBuilderForCoroutineState)
+    private val lineNumber = element?.let { CodegenUtil.getLineNumberForElement(it, false) } ?: 0
 
     private var continuationIndex = if (isForNamedFunction) -1 else 0
     private var dataIndex = if (isForNamedFunction) -1 else 1
@@ -90,6 +97,8 @@ class CoroutineTransformerMethodVisitor(
 
         val suspensionPoints = collectSuspensionPoints(methodNode)
 
+        checkForSuspensionPointInsideMonitor(methodNode, suspensionPoints)
+
         // First instruction in the method node may change in case of named function
         val actualCoroutineStart = methodNode.instructions.first
 
@@ -191,6 +200,37 @@ class CoroutineTransformerMethodVisitor(
         }
     }
 
+    private fun checkForSuspensionPointInsideMonitor(methodNode: MethodNode, suspensionPoints: List<SuspensionPoint>) {
+        if (methodNode.instructions.asSequence().none { it.opcode == Opcodes.MONITORENTER }) return
+
+        val cfg = ControlFlowGraph.build(methodNode)
+        val monitorDepthMap = hashMapOf<AbstractInsnNode, Int>()
+        fun addMonitorDepthToSuccs(index: Int, depth: Int) {
+            val insn = methodNode.instructions[index]
+            monitorDepthMap[insn] = depth
+            val newDepth = when (insn.opcode) {
+                Opcodes.MONITORENTER -> depth + 1
+                Opcodes.MONITOREXIT -> depth - 1
+                else -> depth
+            }
+            for (succIndex in cfg.getSuccessorsIndices(index)) {
+                if (monitorDepthMap[methodNode.instructions[succIndex]] == null) {
+                    addMonitorDepthToSuccs(succIndex, newDepth)
+                }
+            }
+        }
+
+        addMonitorDepthToSuccs(0, 0)
+
+        for (suspensionPoint in suspensionPoints) {
+            if (monitorDepthMap[suspensionPoint.suspensionCallBegin]?.let { it > 0 } == true) {
+                // TODO: Support crossinline suspend lambdas
+                element?.let { diagnostics.report(ErrorsJvm.SUSPENSION_POINT_INSIDE_MONITOR.on(it)) }
+                return
+            }
+        }
+    }
+
     private fun fixLvtForParameters(methodNode: MethodNode, startLabel: LabelNode, endLabel: LabelNode) {
         // We need to skip continuation, since the inliner likes to remap variables there.
         // But this is not a problem, since we have separate $continuation LVT entry

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

@@ -66,7 +66,8 @@ open class SuspendFunctionGenerationStrategy(
         return CoroutineTransformerMethodVisitor(
             mv, access, name, desc, null, null, containingClassInternalName, this::classBuilderForCoroutineState,
             isForNamedFunction = true,
-            lineNumber = CodegenUtil.getLineNumberForElement(declaration, false) ?: 0,
+            element = declaration,
+            diagnostics = state.diagnostics,
             shouldPreserveClassInitialization = constructorCallNormalizationMode.shouldPreserveClassInitialization,
             needDispatchReceiver = originalSuspendDescriptor.dispatchReceiverParameter != null,
             internalNameForDispatchReceiver = containingClassInternalNameOrNull(),

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

@@ -448,7 +448,8 @@ class AnonymousObjectTransformer(
                     ArrayUtil.toStringArray(original.exceptions)
                 ), original.access, original.name, original.desc, null, null,
                 obtainClassBuilderForCoroutineState = { builder },
-                lineNumber = 0, // <- TODO
+                element = null,
+                diagnostics = state.diagnostics,
                 languageVersionSettings = languageVersionSettings,
                 shouldPreserveClassInitialization = state.constructorCallNormalizationMode.shouldPreserveClassInitialization,
                 containingClassInternalName = builder.thisName,
@@ -477,7 +478,8 @@ class AnonymousObjectTransformer(
                     ArrayUtil.toStringArray(original.exceptions)
                 ), original.access, original.name, original.desc, null, null,
                 obtainClassBuilderForCoroutineState = { (inliningContext as RegeneratedClassContext).continuationBuilders[continuationClassName]!! },
-                lineNumber = 0, // <- TODO
+                element = null,
+                diagnostics = state.diagnostics,
                 languageVersionSettings = languageVersionSettings,
                 shouldPreserveClassInitialization = state.constructorCallNormalizationMode.shouldPreserveClassInitialization,
                 containingClassInternalName = builder.thisName,