[coroutine-debugger] Fix smart step into async coroutine builders (launch, async).

The original commit: 1776b92a

Instead of stepping from create to invokeSuspend method of the copied lambda, set a breakpoint into invokeSuspend and resume:
1. First set a breakpoint into `create` method of a given lambda.
2. Get the copied lambda instance.
3. Set a breakpoint into invokeSuspend of the new instance and resume.

IDEA-348340


Merge-request: IJ-MR-147058
Merged-by: Maria Sokolova <maria.sokolova@jetbrains.com>

(cherry picked from commit b190050f086880b02de10a29b5fb61175162348b)

IJ-CR-147164

GitOrigin-RevId: 151051226642470f8e1401d1d82a7a0ca2cd146d

plugins/kotlin/jvm-debugger/core/src/org/jetbrains/kotlin/idea/debugger/core/stepping/DebuggerSteppingHelper.kt

@@ -62,7 +62,7 @@ object DebuggerSteppingHelper {
                 if (context.frameProxy?.isOnSuspensionPoint() == true && nextLocationAfterResume != null) {
                     val filterThread = context.debugProcess.requestsManager.filterThread
                     // step till the next instruction after the resume location
-                    thisLogger().debug("Stepping to the resumeLocation in method ${context.location?.method()?.name()}, filterThread = $filterThread, resumeLocationCodeIndex = ${nextLocationAfterResume.codeIndex()}, currentIndex = ${context.location?.codeIndex()}")
+                    thisLogger().debug("Stepping to the resumeLocation in method ${context.location}, filterThread = $filterThread, resumeLocationCodeIndex = ${nextLocationAfterResume.codeIndex()}, currentIndex = ${context.location?.codeIndex()}")
                     val currentLocation = context.location ?: return super.getNextStepDepth(context)
                     // Make sure that we are stepping to the nextLocationAfterResume in the correct method.
                     if (nextLocationAfterResume.safeMethod() != currentLocation.safeMethod()) {

plugins/kotlin/jvm-debugger/core/src/org/jetbrains/kotlin/idea/debugger/stepping/smartStepInto/KotlinLambdaAsyncMethodFilter.kt

@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.idea.debugger.base.util.safeAllLineLocations
 import org.jetbrains.kotlin.idea.debugger.base.util.safeMethod
 import org.jetbrains.kotlin.idea.debugger.base.util.safeThisObject
 import org.jetbrains.kotlin.idea.debugger.core.DebuggerUtils.isGeneratedIrBackendLambdaMethodName
+import org.jetbrains.kotlin.idea.debugger.core.isInSuspendMethod
 import org.jetbrains.kotlin.idea.debugger.core.isInvokeSuspendMethod
 
 class KotlinLambdaAsyncMethodFilter(
@@ -79,7 +80,7 @@ class KotlinLambdaAsyncMethodFilter(
 
     override fun onReached(context: SuspendContextImpl, hint: RequestHint?): Int {
         try {
-            val breakpoint = createBreakpoint(context)
+            val breakpoint = createBreakpoint(context, hint)
             if (breakpoint != null) {
                 if (isSameCoroutineSuspendLambda) {
                     val filterThread = context.debugProcess.requestsManager.filterThread
@@ -96,11 +97,29 @@ class KotlinLambdaAsyncMethodFilter(
         return RequestHint.STOP
     }
 
-    private fun createBreakpoint(context: SuspendContextImpl): SteppingBreakpoint? {
+    private fun createBreakpoint(context: SuspendContextImpl, hint: RequestHint?): SteppingBreakpoint? {
         val lambdaReference = context.frameProxy?.getLambdaReference() ?: return null
         if (isAsyncSuspendLambda) {
+            /**
+             * In case of an async coroutine builder (launch or async) the `startCoroutine` method is invoked:
+             *
+             * public fun <T> (suspend () -> T).startCoroutine(completion: Continuation<T>) {
+             *     createCoroutineUnintercepted(completion).intercepted().resume(Unit)
+             * }
+             *
+             * It creates a coroutine, intercepts it and resumes with a dummy value (which calls invokeSuspend later).
+             * `createCoroutineUnintercepted` invokes suspending lambda's `create` function, it's generated by the compiler and
+             * creates a copy of the lambda by calling its constructor with captured variables.
+             *
+             * So, in case of an async builder, we cannot just set a breakpoint into invokeSuspend method of a lambda passed to a builder invocation,
+             * since invokeSuspend will already be invoked on the copied instance of the initial lambda.
+             * As a solution:
+             * 1. First set a breakpoint into `create` method of a given lambda
+             * 2. Get the copied lambda instance
+             * 3. Set a breakpoint into invokeSuspend of the new instance and resume.
+             */
             val lambdaMethod = lambdaReference.referenceType().methods().single { it.name() == CREATE }
-            return AsyncSuspendLambdaBreakpoint(context, lambdaReference, lambdaMethod)
+            return AsyncSuspendLambdaBreakpoint(context, hint, lambdaReference, lambdaMethod)
         } else {
             val position = lambdaFilter.breakpointPosition ?: return null
             return KotlinLambdaInstanceBreakpoint(
@@ -128,6 +147,7 @@ class KotlinLambdaAsyncMethodFilter(
 
     private inner class AsyncSuspendLambdaBreakpoint(
         private val context: SuspendContextImpl,
+        private val hint: RequestHint?,
         private val lambdaReference: ObjectReference,
         lambdaMethod: Method
     ) : StepIntoMethodBreakpoint(lambdaMethod.declaringType().name(), lambdaMethod.name(), lambdaMethod.signature(), context.debugProcess.project) {
@@ -136,10 +156,22 @@ class KotlinLambdaAsyncMethodFilter(
                 if (!stopped) return false
                 context.debugProcess.requestsManager.deleteRequest(this)
             }
-            scheduleStepsToInvokeSuspend(action.suspendContext!!).prepareSteppingRequestsAndHints(action.suspendContext!!)
+            val suspendContext = action.suspendContext ?: return false
+            val location = suspendContext.location ?: return false
+            if (isInvokeSuspendMethod(location.method())) {
+                thisLogger().debug("Hit AsyncSuspendLambdaBreakpoint and STOPPED at ${context.location?.method()}")
+                // Stop if we are inside invokeSuspend of the correct lambda instance
+                return true
+            }
+            scheduleStepsToInvokeSuspend(suspendContext).prepareSteppingRequestsAndHints(suspendContext)
             return false
         }
 
+        override fun evaluateCondition(context: EvaluationContextImpl, event: LocatableEvent): Boolean {
+            if (!super.evaluateCondition(context, event)) return false
+            return lambdaReference.checkLambdaBreakpoint(context, event.location())
+        }
+
         private fun scheduleStepsToInvokeSuspend(it: SuspendContextImpl): StepOverCommand =
             with(it.debugProcess) {
                 object : DebugProcessImpl.StepOverCommand(it, false, null, StepRequest.STEP_MIN) {
@@ -152,9 +184,16 @@ class KotlinLambdaAsyncMethodFilter(
                             object : RequestHint(stepThread, suspendContext, StepRequest.STEP_MIN, StepRequest.STEP_OVER, myMethodFilter, parentHint) {
                                 override fun getNextStepDepth(context: SuspendContextImpl): Int {
                                     val location = context.location
-                                    if (location != null && location.method().name() == CREATE) return StepRequest.STEP_OUT
-                                    if (location != null && isInvokeSuspendMethod(location.method()) && location.method().declaringType().name() == lambdaReference.referenceType().name()) {
-                                        return STOP
+                                    if (location != null && location.method().name() == "<init>" && location.method().declaringType() == lambdaReference.referenceType()) {
+                                        val lambdaReferenceCopy = context.thread?.frame(0)?.safeThisObject()
+                                        if (lambdaReferenceCopy == null) {
+                                            thisLogger().debug("Could not extract the copied instance of a lambda from the stack frame of <init> invocation, location = ${location.method()}.")
+                                            return RESUME
+                                        }
+                                        val invokeSuspendMethod = lambdaReferenceCopy.referenceType().methods().single { it.name() == INVOKE_SUSPEND }
+                                        val breakpoint = AsyncSuspendLambdaBreakpoint(context, hint, lambdaReferenceCopy, invokeSuspendMethod)
+                                        DebugProcessImpl.prepareAndSetSteppingBreakpoint(context, breakpoint, hint, true)
+                                        return RESUME
                                     }
                                     return StepRequest.STEP_INTO
                                 }
@@ -164,11 +203,6 @@ class KotlinLambdaAsyncMethodFilter(
                     }
                 }
             }
-
-        override fun evaluateCondition(context: EvaluationContextImpl, event: LocatableEvent): Boolean {
-            if (!super.evaluateCondition(context, event)) return false
-            return lambdaReference.checkLambdaBreakpoint(context, event.location())
-        }
     }
 
     private inner class KotlinLambdaInstanceBreakpoint(
@@ -191,7 +225,7 @@ class KotlinLambdaAsyncMethodFilter(
     }
 
     private fun isTargetLambdaName(name: String): Boolean {
-        if (isAsyncSuspendLambda) return name == CREATE
+        if (isAsyncSuspendLambda) return name == CREATE || name == INVOKE_SUSPEND
         return lambdaFilter.isTargetLambdaName(name)
     }
     
@@ -234,5 +268,6 @@ class KotlinLambdaAsyncMethodFilter(
 
     companion object {
         private const val CREATE = "create"
+        private const val INVOKE_SUSPEND = "invokeSuspend"
     }
 }

plugins/kotlin/jvm-debugger/test/testData/evaluation/singleBreakpoint/coroutines/stepInto/smartStepIntoAsyncBasic.kt

@@ -3,21 +3,34 @@
 
 import kotlinx.coroutines.*
 
+suspend fun foo(i: Int) {
+    coroutineScope {
+        if (i == 34) {
+            //Breakpoint!
+            "".toString()
+        }
+        val res = async(Dispatchers.Default) {
+            delay(10)
+            delay(10)
+            "After delay $i"
+        }
+        res.join()
+        println("Obtained result $res")
+    }
+    println("coroutineScope completed $i")
+}
+
+private suspend fun foo(lambda: suspend () -> Unit) {
+    lambda()
+    "".toString()
+}
+
 fun main() {
     runBlocking {
-        for (i in 0 .. 100) {
-            if (i == 34) {
-                //Breakpoint!
-                "".toString()
+        for (i in 0 .. 1000) {
+            launch {
+                foo(i)
             }
-            val res = async {
-                delay(10)
-                delay(10)
-                "After delay $i"
-            }
-            res.join()
-            println("Obtained result $res")
-            i.toString()
         }
     }
 }

plugins/kotlin/jvm-debugger/test/testData/evaluation/singleBreakpoint/coroutines/stepInto/smartStepIntoAsyncBasic.out

@@ -1,12 +1,12 @@
-LineBreakpoint created at smartStepIntoAsyncBasic.kt:11
+LineBreakpoint created at smartStepIntoAsyncBasic.kt:10
 Run Java
 Connected to the target VM
-smartStepIntoAsyncBasic.kt:11
-smartStepIntoAsyncBasic.kt:13
+smartStepIntoAsyncBasic.kt:10
+smartStepIntoAsyncBasic.kt:12
+smartStepIntoAsyncBasic.kt:12
 smartStepIntoAsyncBasic.kt:13
 smartStepIntoAsyncBasic.kt:14
 smartStepIntoAsyncBasic.kt:15
-smartStepIntoAsyncBasic.kt:16
 Disconnected from the target VM
 
 Process finished with exit code 0

plugins/kotlin/jvm-debugger/test/testData/evaluation/singleBreakpoint/coroutines/stepInto/smartStepIntoLaunchBasic.kt

@@ -3,9 +3,43 @@
 
 import kotlinx.coroutines.*
 
+suspend fun foo(i: Int) {
+    coroutineScope {
+        if (i == 1) {
+            //Breakpoint!
+            i.toString()
+        }
+        launch(Dispatchers.Default) {
+            delay(1)
+            launch(Dispatchers.Default) {
+                delay(1)
+                launch(Dispatchers.Default) {
+                    delay(1)
+                    launch(Dispatchers.Default) {
+                        delay(1)
+                        launch(Dispatchers.Default) {
+                            delay(1)
+                            launch(Dispatchers.Default) {
+                                startMethod(i)
+                                delay(1)
+                                println("After delay $i")
+                            }
+                            delay(1)
+                        }
+                        delay(1)
+                    }
+                }
+            }
+        }
+    }
+    println("coroutineScope completed $i")
+}
+
 suspend fun startMethod(i: Int) {
-    delay(1)
-    "".toString()
+    if (i == 5) {
+        delay(1)
+        "".toString()
+    }
 }
 
 suspend fun endMethod(i: Int) {
@@ -14,16 +48,9 @@ suspend fun endMethod(i: Int) {
 
 fun main() {
     runBlocking {
-        for (i in 0 .. 100) {
-            if (i == 25) {
-                //Breakpoint!
-                "".toString()
-            }
+        for (i in 0 .. 1000) {
             launch {
-                startMethod(i)
-                delay(10)
-                delay(1)
-                startMethod(i)
+                foo(i)
             }
         }
     }
@@ -31,4 +58,14 @@ fun main() {
 
 // STEP_OVER: 1
 // SMART_STEP_INTO_BY_INDEX: 1
+// STEP_OVER: 2
+// SMART_STEP_INTO_BY_INDEX: 1
+// STEP_OVER: 2
+// SMART_STEP_INTO_BY_INDEX: 1
+// STEP_OVER: 2
+// SMART_STEP_INTO_BY_INDEX: 1
+// STEP_OVER: 2
+// SMART_STEP_INTO_BY_INDEX: 1
+// STEP_OVER: 2
+// SMART_STEP_INTO_BY_INDEX: 1
 // STEP_OVER: 3

plugins/kotlin/jvm-debugger/test/testData/evaluation/singleBreakpoint/coroutines/stepInto/smartStepIntoLaunchBasic.out

@@ -1,7 +1,22 @@
-LineBreakpoint created at smartStepIntoLaunchBasic.kt:20
+LineBreakpoint created at smartStepIntoLaunchBasic.kt:10
 Run Java
 Connected to the target VM
+smartStepIntoLaunchBasic.kt:10
+smartStepIntoLaunchBasic.kt:12
+smartStepIntoLaunchBasic.kt:12
+smartStepIntoLaunchBasic.kt:13
+smartStepIntoLaunchBasic.kt:14
+smartStepIntoLaunchBasic.kt:14
+smartStepIntoLaunchBasic.kt:15
+smartStepIntoLaunchBasic.kt:16
+smartStepIntoLaunchBasic.kt:16
+smartStepIntoLaunchBasic.kt:17
+smartStepIntoLaunchBasic.kt:18
+smartStepIntoLaunchBasic.kt:18
+smartStepIntoLaunchBasic.kt:19
 smartStepIntoLaunchBasic.kt:20
+smartStepIntoLaunchBasic.kt:20
+smartStepIntoLaunchBasic.kt:21
 smartStepIntoLaunchBasic.kt:22
 smartStepIntoLaunchBasic.kt:22
 smartStepIntoLaunchBasic.kt:23

plugins/kotlin/jvm-debugger/test/testData/stepping/custom/coroutine.kt

@@ -13,7 +13,7 @@ suspend fun first() {
 
 fun main(args: Array<String>) {
     // SMART_STEP_INTO_BY_INDEX: 2
-    // STEP_OVER: 1
+    // STEP_OVER: 2
     //Breakpoint!
     builder {
         first()

plugins/kotlin/jvm-debugger/test/testData/stepping/custom/coroutine.out

@@ -3,6 +3,7 @@ Run Java
 Connected to the target VM
 coroutine.kt:18
 coroutine.kt:18
+coroutine.kt:18
 coroutine.kt:19
 Disconnected from the target VM