1 files changed, 22 insertions, 47 deletions

diff --git a/extmod/moduasyncio.c b/extmod/moduasyncio.c
index fe0f748ca..9717e3856 100644
--- a/extmod/moduasyncio.c
+++ b/extmod/moduasyncio.c
@@ -31,12 +31,19 @@
 
 #if MICROPY_PY_UASYNCIO
 
+#define TASK_STATE_RUNNING_NOT_WAITED_ON (mp_const_true)
+#define TASK_STATE_DONE_NOT_WAITED_ON (mp_const_none)
+#define TASK_STATE_DONE_WAS_WAITED_ON (mp_const_false)
+
+#define TASK_IS_DONE(task) ( \
+    (task)->state == TASK_STATE_DONE_NOT_WAITED_ON \
+    || (task)->state == TASK_STATE_DONE_WAS_WAITED_ON)
+
 typedef struct _mp_obj_task_t {
     mp_pairheap_t pairheap;
     mp_obj_t coro;
     mp_obj_t data;
-    mp_obj_t waiting;
-
+    mp_obj_t state;
     mp_obj_t ph_key;
 } mp_obj_task_t;
 
@@ -146,9 +153,6 @@ STATIC const mp_obj_type_t task_queue_type = {
 /******************************************************************************/
 // Task class
 
-// For efficiency, the task object is stored to the coro entry when the task is done.
-#define TASK_IS_DONE(task) ((task)->coro == MP_OBJ_FROM_PTR(task))
-
 // This is the core uasyncio context with cur_task, _task_queue and CancelledError.
 STATIC mp_obj_t uasyncio_context = MP_OBJ_NULL;
 
@@ -159,7 +163,7 @@ STATIC mp_obj_t task_make_new(const mp_obj_type_t *type, size_t n_args, size_t n
     mp_pairheap_init_node(task_lt, &self->pairheap);
     self->coro = args[0];
     self->data = mp_const_none;
-    self->waiting = mp_const_none;
+    self->state = TASK_STATE_RUNNING_NOT_WAITED_ON;
     self->ph_key = MP_OBJ_NEW_SMALL_INT(0);
     if (n_args == 2) {
         uasyncio_context = args[1];
@@ -218,24 +222,6 @@ STATIC mp_obj_t task_cancel(mp_obj_t self_in) {
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_1(task_cancel_obj, task_cancel);
 
-STATIC mp_obj_t task_throw(mp_obj_t self_in, mp_obj_t value_in) {
-    // This task raised an exception which was uncaught; handle that now.
-    mp_obj_task_t *self = MP_OBJ_TO_PTR(self_in);
-    // Set the data because it was cleared by the main scheduling loop.
-    self->data = value_in;
-    if (self->waiting == mp_const_none) {
-        // Nothing await'ed on the task so call the exception handler.
-        mp_obj_t _exc_context = mp_obj_dict_get(uasyncio_context, MP_OBJ_NEW_QSTR(MP_QSTR__exc_context));
-        mp_obj_dict_store(_exc_context, MP_OBJ_NEW_QSTR(MP_QSTR_exception), value_in);
-        mp_obj_dict_store(_exc_context, MP_OBJ_NEW_QSTR(MP_QSTR_future), self_in);
-        mp_obj_t Loop = mp_obj_dict_get(uasyncio_context, MP_OBJ_NEW_QSTR(MP_QSTR_Loop));
-        mp_obj_t call_exception_handler = mp_load_attr(Loop, MP_QSTR_call_exception_handler);
-        mp_call_function_1(call_exception_handler, _exc_context);
-    }
-    return mp_const_none;
-}
-STATIC MP_DEFINE_CONST_FUN_OBJ_2(task_throw_obj, task_throw);
-
 STATIC void task_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
     mp_obj_task_t *self = MP_OBJ_TO_PTR(self_in);
     if (dest[0] == MP_OBJ_NULL) {
@@ -244,32 +230,24 @@ STATIC void task_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
             dest[0] = self->coro;
         } else if (attr == MP_QSTR_data) {
             dest[0] = self->data;
-        } else if (attr == MP_QSTR_waiting) {
-            if (self->waiting != mp_const_none && self->waiting != mp_const_false) {
-                dest[0] = self->waiting;
-            }
+        } else if (attr == MP_QSTR_state) {
+            dest[0] = self->state;
         } else if (attr == MP_QSTR_done) {
             dest[0] = MP_OBJ_FROM_PTR(&task_done_obj);
             dest[1] = self_in;
         } else if (attr == MP_QSTR_cancel) {
             dest[0] = MP_OBJ_FROM_PTR(&task_cancel_obj);
             dest[1] = self_in;
-        } else if (attr == MP_QSTR_throw) {
-            dest[0] = MP_OBJ_FROM_PTR(&task_throw_obj);
-            dest[1] = self_in;
         } else if (attr == MP_QSTR_ph_key) {
             dest[0] = self->ph_key;
         }
     } else if (dest[1] != MP_OBJ_NULL) {
         // Store
-        if (attr == MP_QSTR_coro) {
-            self->coro = dest[1];
-            dest[0] = MP_OBJ_NULL;
-        } else if (attr == MP_QSTR_data) {
+        if (attr == MP_QSTR_data) {
             self->data = dest[1];
             dest[0] = MP_OBJ_NULL;
-        } else if (attr == MP_QSTR_waiting) {
-            self->waiting = dest[1];
+        } else if (attr == MP_QSTR_state) {
+            self->state = dest[1];
             dest[0] = MP_OBJ_NULL;
         }
     }
@@ -278,15 +256,12 @@ STATIC void task_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
 STATIC mp_obj_t task_getiter(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf) {
     (void)iter_buf;
     mp_obj_task_t *self = MP_OBJ_TO_PTR(self_in);
-    if (self->waiting == mp_const_none) {
-        // The is the first access of the "waiting" entry.
-        if (TASK_IS_DONE(self)) {
-            // Signal that the completed-task has been await'ed on.
-            self->waiting = mp_const_false;
-        } else {
-            // Lazily allocate the waiting queue.
-            self->waiting = task_queue_make_new(&task_queue_type, 0, 0, NULL);
-        }
+    if (TASK_IS_DONE(self)) {
+        // Signal that the completed-task has been await'ed on.
+        self->state = TASK_STATE_DONE_WAS_WAITED_ON;
+    } else if (self->state == TASK_STATE_RUNNING_NOT_WAITED_ON) {
+        // Allocate the waiting queue.
+        self->state = task_queue_make_new(&task_queue_type, 0, 0, NULL);
     }
     return self_in;
 }
@@ -299,7 +274,7 @@ STATIC mp_obj_t task_iternext(mp_obj_t self_in) {
     } else {
         // Put calling task on waiting queue.
         mp_obj_t cur_task = mp_obj_dict_get(uasyncio_context, MP_OBJ_NEW_QSTR(MP_QSTR_cur_task));
-        mp_obj_t args[2] = { self->waiting, cur_task };
+        mp_obj_t args[2] = { self->state, cur_task };
         task_queue_push_sorted(2, args);
         // Set calling task's data to this task that it waits on, to double-link it.
         ((mp_obj_task_t *)MP_OBJ_TO_PTR(cur_task))->data = self_in;