Provide access to 'coroutineContext' inside suspend functions/lambdas (KT-24829)

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/checkers/coroutineCallChecker.kt

@@ -12,8 +12,10 @@ import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.diagnostics.DiagnosticSink
 import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtCodeFragment
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtThisExpression
+import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.calls.callUtil.isCallableReference
@@ -74,7 +76,20 @@ object CoroutineSuspendCallChecker : CallChecker {
                 val callElement = resolvedCall.call.callElement as KtExpression
 
                 if (!InlineUtil.checkNonLocalReturnUsage(enclosingSuspendFunction, callElement, context.resolutionContext)) {
-                    context.trace.report(Errors.NON_LOCAL_SUSPENSION_POINT.on(reportOn))
+                    var shouldReport = true
+
+                    // Do not report for KtCodeFragment in a suspend function context
+                    val containingFile = callElement.containingFile
+                    if (containingFile is KtCodeFragment) {
+                        val c = containingFile.context?.getParentOfType<KtExpression>(false)
+                        if (c != null && InlineUtil.checkNonLocalReturnUsage(enclosingSuspendFunction, c, context.resolutionContext)) {
+                            shouldReport = false
+                        }
+                    }
+
+                    if (shouldReport) {
+                        context.trace.report(Errors.NON_LOCAL_SUSPENSION_POINT.on(reportOn))
+                    }
                 } else if (context.scope.parentsWithSelf.any { it.isScopeForDefaultParameterValuesOf(enclosingSuspendFunction) }) {
                     context.trace.report(Errors.UNSUPPORTED.on(reportOn, "suspend function calls in a context of default parameter value"))
                 }

idea/idea-jvm/src/org/jetbrains/kotlin/idea/debugger/evaluate/VariableFinder.kt

@@ -21,16 +21,26 @@ import org.jetbrains.kotlin.codegen.topLevelClassAsmType
 import org.jetbrains.kotlin.idea.debugger.*
 import org.jetbrains.kotlin.idea.util.application.runReadAction
 import org.jetbrains.kotlin.idea.util.attachment.mergeAttachments
+import org.jetbrains.kotlin.resolve.calls.checkers.COROUTINE_CONTEXT_1_3_FQ_NAME
 import org.jetbrains.kotlin.load.java.JvmAbi.LOCAL_VARIABLE_NAME_PREFIX_INLINE_ARGUMENT
 import org.jetbrains.kotlin.load.java.JvmAbi.LOCAL_VARIABLE_NAME_PREFIX_INLINE_FUNCTION
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.kotlin.resolve.jvm.JvmClassName
 import org.jetbrains.kotlin.resolve.jvm.JvmPrimitiveType
+import kotlin.coroutines.Continuation
 import org.jetbrains.org.objectweb.asm.Type as AsmType
 import com.sun.jdi.Type as JdiType
 
 class VariableFinder private constructor(private val context: EvaluationContextImpl, private val frameProxy: StackFrameProxyImpl) {
     companion object {
+        private val COROUTINE_CONTEXT_SIMPLE_NAME = COROUTINE_CONTEXT_1_3_FQ_NAME.shortName().asString()
+        private val CONTINUATION_TYPE = AsmType.getType(Continuation::class.java)
+
+        val SUSPEND_LAMBDA_CLASSES = listOf(
+            "kotlin.coroutines.jvm.internal.SuspendLambda",
+            "kotlin.coroutines.jvm.internal.RestrictedSuspendLambda"
+        )
+
         fun instance(context: EvaluationContextImpl): VariableFinder? {
             val frameProxy = context.frameProxy ?: return null
             return VariableFinder(context, frameProxy)
@@ -248,9 +258,60 @@ class VariableFinder private constructor(private val context: EvaluationContextI
             .firstOrNull()
             ?.let { return Result(it.value()) }
 
+        if (name == COROUTINE_CONTEXT_SIMPLE_NAME) {
+            val coroutineContext = findCoroutineContext()?.takeIf { kind.typeMatches(it.type()) }
+            if (coroutineContext != null) {
+                return Result(coroutineContext)
+            }
+        }
+
         return null
     }
 
+    private fun findCoroutineContext(): ObjectReference? {
+        val method = frameProxy.location().safeMethod() ?: return null
+        return findCoroutineContextForLambda(method) ?: findCoroutineContextForMethod(method)
+    }
+
+    private fun findCoroutineContextForLambda(method: Method): ObjectReference? {
+        if (method.name() != "invokeSuspend" || method.signature() != "(Ljava/lang/Object;)Ljava/lang/Object;") {
+            return null
+        }
+
+        val thisObject = frameProxy.thisObject() ?: return null
+        val thisType = thisObject.referenceType()
+
+        if (SUSPEND_LAMBDA_CLASSES.none { thisType.isSubtype(it) }) {
+            return null
+        }
+
+        return findCoroutineContextForContinuation(thisObject)
+    }
+
+    private fun findCoroutineContextForMethod(method: Method): ObjectReference? {
+        if (CONTINUATION_TYPE.descriptor + ")" !in method.signature()) {
+            return null
+        }
+
+        val continuationVariable = frameProxy.visibleVariableByName("\$continuation") ?: return null
+        val continuation = frameProxy.getValue(continuationVariable) as? ObjectReference ?: return null
+        return findCoroutineContextForContinuation(continuation)
+    }
+
+    private fun findCoroutineContextForContinuation(continuation: ObjectReference): ObjectReference? {
+        val continuationType = (continuation.referenceType() as? ClassType)
+            ?.allInterfaces()?.firstOrNull { it.name() == Continuation::class.java.name }
+            ?: return null
+
+        val getContextMethod = continuationType
+            .methodsByName("getContext", "()Lkotlin/coroutines/CoroutineContext;").firstOrNull()
+            ?: return null
+
+        val threadReference = frameProxy.threadProxy().threadReference.takeIf { it.isSuspended } ?: return null
+        val invokePolicy = context.suspendContext.getInvokePolicy()
+        return continuation.invokeMethod(threadReference, getContextMethod, emptyList(), invokePolicy) as? ObjectReference
+    }
+
     private fun findCapturedVariableInReceiver(variables: List<LocalVariableProxyImpl>, kind: VariableKind): Result? {
         fun isReceiverOrPassedThis(name: String) =
             name.startsWith(AsmUtil.LABELED_THIS_PARAMETER)

idea/idea-jvm/src/org/jetbrains/kotlin/idea/debugger/safeUtil.kt

@@ -30,6 +30,14 @@ fun ReferenceType.safeAllLineLocations(): List<Location> {
     return DebuggerUtilsEx.allLineLocations(this) ?: emptyList()
 }
 
+fun ReferenceType.safeSourceName(): String? {
+    return try {
+        sourceName()
+    } catch (e: AbsentInformationException) {
+        null
+    }
+}
+
 fun Method.safeLocationsOfLine(line: Int): List<Location> {
     return try {
         locationsOfLine(line)

idea/testData/debugger/tinyApp/src/evaluate/singleBreakpoint/frame/coroutineContextFun.kt

@@ -0,0 +1,15 @@
+package coroutineContextFun
+
+import kotlin.coroutines.coroutineContext
+
+suspend fun main() {
+    val a = 5
+    foo()
+    //Breakpoint!
+    foo()
+}
+
+private suspend fun foo() {}
+
+// EXPRESSION: coroutineContext
+// RESULT: instance of kotlin.coroutines.EmptyCoroutineContext(id=ID): Lkotlin/coroutines/EmptyCoroutineContext;

idea/testData/debugger/tinyApp/src/evaluate/singleBreakpoint/frame/coroutineContextFun.out

@@ -0,0 +1,8 @@
+LineBreakpoint created at coroutineContextFun.kt:9
+Run Java
+Connected to the target VM
+coroutineContextFun.kt:9
+Compile bytecode for coroutineContext
+Disconnected from the target VM
+
+Process finished with exit code 0

idea/testData/debugger/tinyApp/src/evaluate/singleBreakpoint/frame/coroutineContextLambda.kt

@@ -0,0 +1,19 @@
+package coroutineContextLambda
+
+import kotlin.coroutines.coroutineContext
+
+suspend fun main() {
+    foo()
+}
+
+private var foo: suspend () -> Unit = {
+    bar()
+}
+
+private var bar: suspend () -> Unit = {
+    //Breakpoint!
+    val a = 5
+}
+
+// EXPRESSION: coroutineContext
+// RESULT: instance of kotlin.coroutines.EmptyCoroutineContext(id=ID): Lkotlin/coroutines/EmptyCoroutineContext;

idea/testData/debugger/tinyApp/src/evaluate/singleBreakpoint/frame/coroutineContextLambda.out

@@ -0,0 +1,8 @@
+LineBreakpoint created at coroutineContextLambda.kt:15
+Run Java
+Connected to the target VM
+coroutineContextLambda.kt:15
+Compile bytecode for coroutineContext
+Disconnected from the target VM
+
+Process finished with exit code 0

idea/tests/org/jetbrains/kotlin/idea/debugger/evaluate/KotlinEvaluateExpressionTestGenerated.java

@@ -600,6 +600,16 @@ public class KotlinEvaluateExpressionTestGenerated extends AbstractKotlinEvaluat
                 runTest("idea/testData/debugger/tinyApp/src/evaluate/singleBreakpoint/frame/catchVariable.kt");
             }
 
+            @TestMetadata("coroutineContextFun.kt")
+            public void testCoroutineContextFun() throws Exception {
+                runTest("idea/testData/debugger/tinyApp/src/evaluate/singleBreakpoint/frame/coroutineContextFun.kt");
+            }
+
+            @TestMetadata("coroutineContextLambda.kt")
+            public void testCoroutineContextLambda() throws Exception {
+                runTest("idea/testData/debugger/tinyApp/src/evaluate/singleBreakpoint/frame/coroutineContextLambda.kt");
+            }
+
             @TestMetadata("delegatedPropertyInClass.kt")
             public void testDelegatedPropertyInClass() throws Exception {
                 runTest("idea/testData/debugger/tinyApp/src/evaluate/singleBreakpoint/frame/delegatedPropertyInClass.kt");