HACK: fix await wait(0)

This is what my late-night brain came up with. Likely a better way to
do this.

Fixes: pybricks/support#1775

pybricks/common/pb_type_device.c

@@ -69,7 +69,7 @@ void *pb_type_device_get_data_blocking(mp_obj_t self_in, uint8_t mode) {
  * @return                  True if operation is complete (device ready),
  *                          false otherwise.
  */
-static bool pb_pup_device_test_completion(mp_obj_t self_in, uint32_t end_time) {
+static bool pb_pup_device_test_completion(mp_obj_t self_in, uint32_t *end_time) {
     pb_type_device_obj_base_t *sensor = MP_OBJ_TO_PTR(self_in);
     pbio_error_t err = pbdrv_legodev_is_ready(sensor->legodev);
     if (err == PBIO_ERROR_AGAIN) {

pybricks/common/pb_type_motor.c

@@ -259,7 +259,7 @@ static mp_obj_t pb_type_Motor_hold(mp_obj_t self_in) {
 }
 static MP_DEFINE_CONST_FUN_OBJ_1(pb_type_Motor_hold_obj, pb_type_Motor_hold);
 
-static bool pb_type_Motor_test_completion(mp_obj_t self_in, uint32_t end_time) {
+static bool pb_type_Motor_test_completion(mp_obj_t self_in, uint32_t *end_time) {
     pb_type_Motor_obj_t *self = MP_OBJ_TO_PTR(self_in);
     // Handle I/O exceptions like port unplugged.
     if (!pbio_servo_update_loop_is_running(self->srv)) {

pybricks/common/pb_type_speaker.c

@@ -123,7 +123,7 @@ static mp_obj_t pb_type_Speaker_make_new(const mp_obj_type_t *type, size_t n_arg
     return MP_OBJ_FROM_PTR(self);
 }
 
-static bool pb_type_Speaker_beep_test_completion(mp_obj_t self_in, uint32_t end_time) {
+static bool pb_type_Speaker_beep_test_completion(mp_obj_t self_in, uint32_t *end_time) {
     pb_type_Speaker_obj_t *self = MP_OBJ_TO_PTR(self_in);
     if (mp_hal_ticks_ms() - self->beep_end_time < (uint32_t)INT32_MAX) {
         pb_type_Speaker_stop_beep();
@@ -337,7 +337,7 @@ static void pb_type_Speaker_play_note(pb_type_Speaker_obj_t *self, mp_obj_t obj,
     self->beep_end_time = release ? time_now + 7 * duration / 8 : time_now + duration;
 }
 
-static bool pb_type_Speaker_notes_test_completion(mp_obj_t self_in, uint32_t end_time) {
+static bool pb_type_Speaker_notes_test_completion(mp_obj_t self_in, uint32_t *end_time) {
     pb_type_Speaker_obj_t *self = MP_OBJ_TO_PTR(self_in);
 
     bool release_done = mp_hal_ticks_ms() - self->release_end_time < (uint32_t)INT32_MAX;

pybricks/robotics/pb_type_drivebase.c

@@ -90,7 +90,7 @@ static mp_obj_t pb_type_DriveBase_make_new(const mp_obj_type_t *type, size_t n_a
     return MP_OBJ_FROM_PTR(self);
 }
 
-static bool pb_type_DriveBase_test_completion(mp_obj_t self_in, uint32_t end_time) {
+static bool pb_type_DriveBase_test_completion(mp_obj_t self_in, uint32_t *end_time) {
 
     pb_type_DriveBase_obj_t *self = MP_OBJ_TO_PTR(self_in);
 

pybricks/tools/pb_module_tools.c

@@ -48,8 +48,18 @@ void pb_module_tools_assert_blocking(void) {
 // us share the same code with other awaitables. It also minimizes allocation.
 MP_REGISTER_ROOT_POINTER(mp_obj_t wait_awaitables);
 
-static bool pb_module_tools_wait_test_completion(mp_obj_t obj, uint32_t end_time) {
-    return mp_hal_ticks_ms() - end_time < UINT32_MAX / 2;
+static bool pb_module_tools_wait_one(mp_obj_t obj, uint32_t *end_time) {
+    // yield only on the first iteration.
+    if (*end_time == pb_type_awaitable_end_time_none) {
+        (*end_time)++;
+        return false;
+    }
+
+    return true;
+}
+
+static bool pb_module_tools_wait_test_completion(mp_obj_t obj, uint32_t *end_time) {
+    return mp_hal_ticks_ms() - *end_time < UINT32_MAX / 2;
 }
 
 static mp_obj_t pb_module_tools_wait(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
@@ -71,8 +81,22 @@ static mp_obj_t pb_module_tools_wait(size_t n_args, const mp_obj_t *pos_args, mp
     // test completion state in iteration loop.
     time = pbio_int_math_bind(time, 0, INT32_MAX >> 2);
 
+    // Special case to emulate CPython asyncio.sleep(0) where we ensure that
+    // the task yields exactly once.
+    if (time == 0) {
+        return pb_type_awaitable_await_or_wait(
+            MP_OBJ_NULL,
+            MP_STATE_PORT(wait_awaitables),
+            pb_type_awaitable_end_time_none,
+            pb_module_tools_wait_one,
+            pb_type_awaitable_return_none,
+            pb_type_awaitable_cancel_none,
+            PB_TYPE_AWAITABLE_OPT_NONE
+        );
+    }
+
     return pb_type_awaitable_await_or_wait(
-        NULL, // wait functions are not associated with an object
+        MP_OBJ_NULL, // wait functions are not associated with an object
         MP_STATE_PORT(wait_awaitables),
         mp_hal_ticks_ms() + time,
         pb_module_tools_wait_test_completion,
@@ -133,7 +157,7 @@ void pb_module_tools_pbio_task_do_blocking(pbio_task_t *task, mp_int_t timeout)
 // here instead of with each Bluetooth-related MicroPython object.
 MP_REGISTER_ROOT_POINTER(mp_obj_t pbio_task_awaitables);
 
-static bool pb_module_tools_pbio_task_test_completion(mp_obj_t obj, uint32_t end_time) {
+static bool pb_module_tools_pbio_task_test_completion(mp_obj_t obj, uint32_t *end_time) {
     pbio_task_t *task = MP_OBJ_TO_PTR(obj);
 
     // Keep going if not done yet.

pybricks/tools/pb_type_awaitable.c

@@ -63,7 +63,7 @@ static mp_obj_t pb_type_awaitable_iternext(mp_obj_t self_in) {
     }
 
     // Keep going if not completed by returning None.
-    if (!self->test_completion(self->obj, self->end_time)) {
+    if (!self->test_completion(self->obj, &self->end_time)) {
         return mp_const_none;
     }
 
@@ -127,7 +127,7 @@ static pb_type_awaitable_obj_t *pb_type_awaitable_get(mp_obj_t awaitables_in) {
  * checker. This allows MicroPython to handle completion during the next call
  * to iternext.
  */
-static bool pb_type_awaitable_completed(mp_obj_t self_in, uint32_t start_time) {
+static bool pb_type_awaitable_completed(mp_obj_t self_in, uint32_t *end_time) {
     return true;
 }
 
@@ -219,7 +219,7 @@ mp_obj_t pb_type_awaitable_await_or_wait(
     }
 
     // Outside run loop, block until the operation is complete.
-    while (test_completion_func && !test_completion_func(obj, end_time)) {
+    while (test_completion_func && !test_completion_func(obj, &end_time)) {
         mp_hal_delay_ms(1);
     }
     if (!return_value_func) {

pybricks/tools/pb_type_awaitable.h

@@ -54,10 +54,10 @@ typedef struct _pb_type_awaitable_obj_t pb_type_awaitable_obj_t;
  * below, which always stops the relevant hardware (i.e. always coast).
  *
  * @param [in]  obj            The object associated with this awaitable.
- * @param [in]  start_time     The time when the awaitable was created.
+ * @param [in]  end_time       Can be used by the callback for timeouts, etc.
  * @return                     True if operation is complete, False otherwise.
  */
-typedef bool (*pb_type_awaitable_test_completion_t)(mp_obj_t obj, uint32_t end_time);
+typedef bool (*pb_type_awaitable_test_completion_t)(mp_obj_t obj, uint32_t *end_time);
 
 /**
  * Gets the return value of the awaitable. If it always returns None, providing