From f017408612921c1d43dec3e04f14c28d69e4a472 Mon Sep 17 00:00:00 2001
From: Laurens Valk <laurens@pybricks.com>
Date: Mon, 26 Aug 2024 17:25:31 +0200
Subject: [PATCH] pybricks.tools.wait: Fix awaiting on wait(0).

---
 CHANGELOG.md                       |  2 ++
 pybricks/tools/pb_module_tools.c   |  2 +-
 pybricks/tools/pb_type_awaitable.c | 44 +++++++++++++++++++++---------
 pybricks/tools/pb_type_awaitable.h |  2 ++
 4 files changed, 36 insertions(+), 14 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4ca58ce46..fa391552f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -28,9 +28,11 @@
 ### Fixed
 - Fixed not able to connect to new Technic Move hub with `LWP3Device()`.
 - Removed `gc_collect()` from `tools.run_task()` loop to fix unwanted delays.
+- Fixed `await wait(0)` never yielding, so parallel tasks could lock up ([support#1429]).
 
 [pybricks-micropython#250]: https://github.com/pybricks/pybricks-micropython/pull/250
 [pybricks-micropython#253]: https://github.com/pybricks/pybricks-micropython/pull/253
+[support#1429]: https://github.com/pybricks/support/issues/1429
 [support#1615]: https://github.com/pybricks/support/issues/1615
 [support#1622]: https://github.com/pybricks/support/issues/1622
 [support#1678]: https://github.com/pybricks/support/issues/1678
diff --git a/pybricks/tools/pb_module_tools.c b/pybricks/tools/pb_module_tools.c
index 1976d52f6..03f298d70 100644
--- a/pybricks/tools/pb_module_tools.c
+++ b/pybricks/tools/pb_module_tools.c
@@ -75,7 +75,7 @@ static mp_obj_t pb_module_tools_wait(size_t n_args, const mp_obj_t *pos_args, mp
         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,
+        time > 0 ? pb_module_tools_wait_test_completion : pb_type_awaitable_test_completion_yield_once,
         pb_type_awaitable_return_none,
         pb_type_awaitable_cancel_none,
         PB_TYPE_AWAITABLE_OPT_NONE);
diff --git a/pybricks/tools/pb_type_awaitable.c b/pybricks/tools/pb_type_awaitable.c
index 4cc89a6b6..420915b65 100644
--- a/pybricks/tools/pb_type_awaitable.c
+++ b/pybricks/tools/pb_type_awaitable.c
@@ -54,6 +54,27 @@ static mp_obj_t pb_type_awaitable_close(mp_obj_t self_in) {
 }
 static MP_DEFINE_CONST_FUN_OBJ_1(pb_type_awaitable_close_obj, pb_type_awaitable_close);
 
+/**
+ * Completion checker that is always true.
+ *
+ * Linked awaitables are gracefully cancelled by setting this as the completion
+ * checker. This allows MicroPython to handle completion during the next call
+ * to iternext.
+ */
+static bool pb_type_awaitable_test_completion_completed(mp_obj_t self_in, uint32_t start_time) {
+    return true;
+}
+
+/**
+ * Special completion test for awaitable that should yield exactly once.
+ *
+ * It will return false to indicate that it is not done. Then the iternext will
+ * replace this test with one that is always done, thus completing next time.
+ */
+bool pb_type_awaitable_test_completion_yield_once(mp_obj_t obj, uint32_t end_time) {
+    return false;
+}
+
 static mp_obj_t pb_type_awaitable_iternext(mp_obj_t self_in) {
     pb_type_awaitable_obj_t *self = MP_OBJ_TO_PTR(self_in);
 
@@ -62,8 +83,16 @@ static mp_obj_t pb_type_awaitable_iternext(mp_obj_t self_in) {
         return MP_OBJ_STOP_ITERATION;
     }
 
+    bool complete = self->test_completion(self->obj, self->end_time);
+
+    // If this was a special awaitable that was supposed to yield exactly once,
+    // it will now be yielding by being not complete, but complete the next time.
+    if (self->test_completion == pb_type_awaitable_test_completion_yield_once) {
+        self->test_completion = pb_type_awaitable_test_completion_completed;
+    }
+
     // Keep going if not completed by returning None.
-    if (!self->test_completion(self->obj, self->end_time)) {
+    if (!complete) {
         return mp_const_none;
     }
 
@@ -120,17 +149,6 @@ static pb_type_awaitable_obj_t *pb_type_awaitable_get(mp_obj_t awaitables_in) {
     return awaitable;
 }
 
-/**
- * Completion checker that is always true.
- *
- * Linked awaitables are gracefully cancelled by setting this as the completion
- * 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) {
-    return true;
-}
-
 /**
  * Checks and updates all awaitables associated with an object.
  *
@@ -166,7 +184,7 @@ void pb_type_awaitable_update_all(mp_obj_t awaitables_in, pb_type_awaitable_opt_
         // Set awaitable to done so it gets cancelled it gracefully on the
         // next iteration.
         if (options & PB_TYPE_AWAITABLE_OPT_CANCEL_ALL) {
-            awaitable->test_completion = pb_type_awaitable_completed;
+            awaitable->test_completion = pb_type_awaitable_test_completion_completed;
         }
 
     }
diff --git a/pybricks/tools/pb_type_awaitable.h b/pybricks/tools/pb_type_awaitable.h
index 9dea9898a..b63f434b8 100644
--- a/pybricks/tools/pb_type_awaitable.h
+++ b/pybricks/tools/pb_type_awaitable.h
@@ -82,6 +82,8 @@ typedef void (*pb_type_awaitable_cancel_t)(mp_obj_t obj);
 
 #define pb_type_awaitable_cancel_none (NULL)
 
+bool pb_type_awaitable_test_completion_yield_once(mp_obj_t obj, uint32_t end_time);
+
 void pb_type_awaitable_update_all(mp_obj_t awaitables_in, pb_type_awaitable_opt_t options);
 
 mp_obj_t pb_type_awaitable_await_or_wait(