iterative + brent

Author: maurycy

Modules/_remote_debugging/_remote_debugging.h

@@ -733,8 +733,7 @@ extern int parse_task(
 extern int parse_coro_chain(
     RemoteUnwinderObject *unwinder,
     uintptr_t coro_address,
-    PyObject *render_to,
-    size_t depth
+    PyObject *render_to
 );
 
 extern int parse_async_frame_chain(

Modules/_remote_debugging/asyncio.c

@@ -248,13 +248,14 @@ parse_task_name(
  * ============================================================================ */
 
 static int
-handle_yield_from_frame(
+get_awaited_coro(
     RemoteUnwinderObject *unwinder,
     uintptr_t gi_iframe_addr,
     uintptr_t gen_type_addr,
-    PyObject *render_to,
-    size_t depth
+    uintptr_t *next_coro_addr
 ) {
+    *next_coro_addr = 0;
+
     // Read the entire interpreter frame at once
     char iframe[SIZEOF_INTERP_FRAME];
     int err = _Py_RemoteDebug_PagedReadRemoteMemory(
@@ -310,12 +311,7 @@ handle_yield_from_frame(
                    doesn't match the type of whatever it points to
                    in its cr_await.
                 */
-                err = parse_coro_chain(unwinder, gi_await_addr, render_to,
-                                       depth + 1);
-                if (err) {
-                    set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to parse coroutine chain in yield_from");
-                    return -1;
-                }
+                *next_coro_addr = gi_await_addr;
             }
         }
     }
@@ -327,63 +323,94 @@ int
 parse_coro_chain(
     RemoteUnwinderObject *unwinder,
     uintptr_t coro_address,
-    PyObject *render_to,
-    size_t depth
+    PyObject *render_to
 ) {
     assert((void*)coro_address != NULL);
 
-    if (depth >= MAX_FRAME_CHAIN_DEPTH) {
-        PyErr_SetString(PyExc_RuntimeError,
-            "Too many coroutine frames (possible infinite loop)");
-        set_exception_cause(unwinder, PyExc_RuntimeError,
-            "Coroutine chain depth limit exceeded");
-        return -1;
-    }
+    /* Brent's cycle detection algorithm */
+    uintptr_t checkpoint_addr = 0;
+    uintptr_t checkpoint_code = 0;
+    size_t window = 1;
+    size_t steps = 0;
 
-    // Read the entire generator object at once
-    char gen_object[SIZEOF_GEN_OBJ];
-    int err = _Py_RemoteDebug_PagedReadRemoteMemory(
-        &unwinder->handle,
-        coro_address,
-        SIZEOF_GEN_OBJ,
-        gen_object);
-    if (err < 0) {
-        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read generator object in coro chain");
-        return -1;
-    }
+    for (size_t depth = 0; (void*)coro_address != NULL; depth++) {
+        if (depth >= MAX_FRAME_CHAIN_DEPTH) {
+            PyErr_SetString(PyExc_RuntimeError,
+                "Too many coroutine frames (possible infinite loop)");
+            set_exception_cause(unwinder, PyExc_RuntimeError,
+                "Coroutine chain depth limit exceeded");
+            return -1;
+        }
 
-    int8_t frame_state = GET_MEMBER(int8_t, gen_object, unwinder->debug_offsets.gen_object.gi_frame_state);
-    if (frame_state == FRAME_CLEARED) {
-        return 0;
-    }
+        // Read the entire generator object at once
+        char gen_object[SIZEOF_GEN_OBJ];
+        int err = _Py_RemoteDebug_PagedReadRemoteMemory(
+            &unwinder->handle,
+            coro_address,
+            SIZEOF_GEN_OBJ,
+            gen_object);
+        if (err < 0) {
+            set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read generator object in coro chain");
+            return -1;
+        }
 
-    uintptr_t gen_type_addr = GET_MEMBER(uintptr_t, gen_object, unwinder->debug_offsets.pyobject.ob_type);
+        int8_t frame_state = GET_MEMBER(int8_t, gen_object, unwinder->debug_offsets.gen_object.gi_frame_state);
+        if (frame_state == FRAME_CLEARED) {
+            return 0;
+        }
 
-    PyObject* name = NULL;
+        uintptr_t gen_type_addr = GET_MEMBER(uintptr_t, gen_object, unwinder->debug_offsets.pyobject.ob_type);
 
-    // Parse the previous frame using the gi_iframe from local copy
-    uintptr_t prev_frame;
-    uintptr_t gi_iframe_addr = coro_address + (uintptr_t)unwinder->debug_offsets.gen_object.gi_iframe;
-    uintptr_t address_of_code_object = 0;
-    if (parse_frame_object(unwinder, &name, gi_iframe_addr, &address_of_code_object, &prev_frame) < 0) {
-        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to parse frame object in coro chain");
-        return -1;
-    }
+        PyObject* name = NULL;
 
-    if (!name) {
-        return 0;
-    }
+        // Parse the previous frame using the gi_iframe from local copy
+        uintptr_t prev_frame;
+        uintptr_t gi_iframe_addr = coro_address + (uintptr_t)unwinder->debug_offsets.gen_object.gi_iframe;
+        uintptr_t address_of_code_object = 0;
+        if (parse_frame_object(unwinder, &name, gi_iframe_addr, &address_of_code_object, &prev_frame) < 0) {
+            set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to parse frame object in coro chain");
+            return -1;
+        }
 
-    if (PyList_Append(render_to, name)) {
+        if (!name) {
+            return 0;
+        }
+
+        if (PyList_Append(render_to, name)) {
+            Py_DECREF(name);
+            set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to append frame to coro chain");
+            return -1;
+        }
         Py_DECREF(name);
-        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to append frame to coro chain");
-        return -1;
-    }
-    Py_DECREF(name);
 
-    if (frame_state == FRAME_SUSPENDED_YIELD_FROM) {
-        return handle_yield_from_frame(unwinder, gi_iframe_addr, gen_type_addr,
-                                       render_to, depth);
+        if (coro_address == checkpoint_addr
+            && address_of_code_object == checkpoint_code)
+        {
+            PyErr_SetString(PyExc_RuntimeError,
+                "Cycle detected in coroutine chain");
+            set_exception_cause(unwinder, PyExc_RuntimeError,
+                "Coroutine chain contains a cycle");
+            return -1;
+        }
+        if (steps == window) {
+            checkpoint_addr = coro_address;
+            checkpoint_code = address_of_code_object;
+            window *= 2;
+            steps = 0;
+        }
+        steps++;
+
+        if (frame_state != FRAME_SUSPENDED_YIELD_FROM) {
+            return 0;
+        }
+
+        uintptr_t next_coro_addr;
+        if (get_awaited_coro(unwinder, gi_iframe_addr, gen_type_addr,
+                             &next_coro_addr) < 0) {
+            set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to parse coroutine chain in yield_from");
+            return -1;
+        }
+        coro_address = next_coro_addr;
     }
 
     return 0;
@@ -429,7 +456,7 @@ create_task_result(
     coro_addr = GET_MEMBER_NO_TAG(uintptr_t, task_obj, unwinder->async_debug_offsets.asyncio_task_object.task_coro);
 
     if ((void*)coro_addr != NULL) {
-        if (parse_coro_chain(unwinder, coro_addr, call_stack, 0) < 0) {
+        if (parse_coro_chain(unwinder, coro_addr, call_stack) < 0) {
             set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to parse coroutine chain");
             goto error;
         }