Add awaitable C functions

CIRCUITPY_DEFINE_ASYNC_FUN_OBJ_KW(obj, n_args_min, fun) defines a method
that returns an awaitable. As with a CPython coroutine, its arguments are
stored and only parsed when it is awaited. While an operation runs, the
task waits on asyncio's I/O queue: the awaitable is a stream that polls
readable once the operation's done flag is set, possibly from an
interrupt. CIRCUITPY_DEFINE_ASYNC_POLL_FUN_OBJ_KW is for operations with
no interrupt: the task tries again on each pass of the event loop.
Based on the awaitable in adafruit#11040.

Author: lynt-smitka

supervisor/shared/awaitable.c

@@ -0,0 +1,187 @@
+// This file is part of the CircuitPython project: https://circuitpython.org
+//
+// SPDX-FileCopyrightText: Copyright (c) 2025 Scott Shawcroft for Adafruit Industries
+//
+// SPDX-License-Identifier: MIT
+
+#include <string.h>
+
+#include "supervisor/shared/awaitable.h"
+#include "py/mperrno.h"
+#include "py/objmodule.h"
+#include "py/runtime.h"
+#include "py/stream.h"
+
+#if MICROPY_PY_ASYNC_AWAIT
+
+#if MICROPY_PY_ASYNCIO
+// The globals of asyncio.core, if it has been imported. Found through sys.modules, because
+// mp_asyncio_context in modasyncio.c is left pointing at the old heap after a soft reload.
+static mp_map_t *asyncio_core_globals(void) {
+    mp_map_elem_t *elem = mp_map_lookup(&MP_STATE_VM(mp_loaded_modules_dict).map,
+        MP_OBJ_NEW_QSTR(MP_QSTR_asyncio_dot_core), MP_MAP_LOOKUP);
+    return elem == NULL ? NULL : &mp_obj_module_get_globals(elem->value)->map;
+}
+
+static mp_obj_t map_get(mp_map_t *map, qstr name) {
+    mp_map_elem_t *elem = mp_map_lookup(map, MP_OBJ_NEW_QSTR(name), MP_MAP_LOOKUP);
+    return elem == NULL ? mp_const_none : elem->value;
+}
+
+// Make the running asyncio task wait for this awaitable: on the I/O queue until it polls
+// readable, or with poll, on the task queue so that it runs again on the next pass.
+static void asyncio_wait(mp_obj_t awaitable, bool poll) {
+    mp_map_t *core = asyncio_core_globals();
+    if (core == NULL) {
+        return;
+    }
+    mp_obj_t cur_task = map_get(core, MP_QSTR_cur_task);
+    if (cur_task == mp_const_none) {
+        return;
+    }
+    mp_obj_t dest[4];
+    if (poll) {
+        mp_load_method(map_get(core, MP_QSTR__task_queue), MP_QSTR_push, dest);
+        dest[2] = cur_task;
+        mp_call_method_n_kw(1, 0, dest);
+    } else {
+        mp_load_method(map_get(core, MP_QSTR__io_queue), MP_QSTR__enqueue, dest);
+        dest[2] = awaitable;
+        dest[3] = MP_OBJ_NEW_SMALL_INT(0);  // wait to read
+        mp_call_method_n_kw(2, 0, dest);
+    }
+}
+#endif
+
+// ---- the async function ------------------------------------------------------
+
+// Store the arguments; they are parsed by start() when the awaitable is first resumed.
+static mp_obj_t async_fun_call(mp_obj_t self_in, size_t n_args, size_t n_kw, const mp_obj_t *args) {
+    const circuitpy_async_fun_obj_t *fun = MP_OBJ_TO_PTR(self_in);
+    mp_arg_check_num(n_args, n_kw, fun->n_args_min, MP_OBJ_FUN_ARGS_MAX, true);
+    size_t n = n_args + 2 * n_kw;
+    circuitpy_awaitable_obj_t *aw = mp_obj_malloc_var_with_finaliser(circuitpy_awaitable_obj_t,
+        args, mp_obj_t, n, &circuitpy_awaitable_type);
+    aw->fun = fun;
+    aw->context = NULL;
+    aw->started = false;
+    aw->finished = false;
+    aw->n_args = n_args;
+    aw->n_kw = n_kw;
+    memcpy(aw->args, args, n * sizeof(mp_obj_t));
+    return MP_OBJ_FROM_PTR(aw);
+}
+
+MP_DEFINE_CONST_OBJ_TYPE(
+    circuitpy_async_fun_type,
+    MP_QSTR_function,
+    MP_TYPE_FLAG_BINDS_SELF,
+    call, async_fun_call
+    );
+
+// ---- the awaitable -----------------------------------------------------------
+
+static void awaitable_cancel(circuitpy_awaitable_obj_t *self) {
+    if (self->started && !self->finished) {
+        self->finished = true;
+        self->fun->cancel(self->context, &self->done);
+    }
+}
+
+static mp_obj_t awaitable_iternext(mp_obj_t self_in) {
+    circuitpy_awaitable_obj_t *self = MP_OBJ_TO_PTR(self_in);
+    if (self->finished) {
+        // Awaited again: like a finished generator, return None at once.
+        MP_STATE_THREAD(stop_iteration_arg) = mp_const_none;
+        return MP_OBJ_STOP_ITERATION;
+    }
+    if (!self->started) {
+        mp_map_t kw_args;
+        mp_map_init_fixed_table(&kw_args, self->n_kw, self->args + self->n_args);
+        CIRCUITPY_ASYNC_FLAG_INIT(&self->done);
+        // An exception from start() (bad arguments) propagates out of the await.
+        self->context = self->fun->start(&self->done, self->n_args, self->args, &kw_args);
+        self->started = true;
+    }
+    if (!CIRCUITPY_ASYNC_FLAG_IS_SET(&self->done) && self->fun->poll != NULL && self->fun->poll(self->context)) {
+        CIRCUITPY_ASYNC_FLAG_SET(&self->done);
+    }
+    if (!CIRCUITPY_ASYNC_FLAG_IS_SET(&self->done)) {
+        #if MICROPY_PY_ASYNCIO
+        asyncio_wait(self_in, self->fun->poll != NULL);
+        #endif
+        return mp_const_none;
+    }
+    self->finished = true;
+    MP_STATE_THREAD(stop_iteration_arg) = self->fun->end(self->context, &self->done);
+    return MP_OBJ_STOP_ITERATION;
+}
+
+static mp_obj_t awaitable_await(mp_obj_t self_in) {
+    return self_in;
+}
+static MP_DEFINE_CONST_FUN_OBJ_1(awaitable_await_obj, awaitable_await);
+
+static mp_obj_t awaitable_send(mp_obj_t self_in, mp_obj_t value) {
+    if (value != mp_const_none) {
+        mp_raise_TypeError(MP_ERROR_TEXT("can't send non-None value to a just-started generator"));
+    }
+    mp_obj_t ret = awaitable_iternext(self_in);
+    if (ret == MP_OBJ_STOP_ITERATION) {
+        mp_raise_StopIteration(MP_STATE_THREAD(stop_iteration_arg));
+    }
+    return ret;
+}
+static MP_DEFINE_CONST_FUN_OBJ_2(awaitable_send_obj, awaitable_send);
+
+// throw() is how asyncio cancels a task: stop the operation and raise the exception.
+static mp_obj_t awaitable_throw(size_t n_args, const mp_obj_t *args) {
+    awaitable_cancel(MP_OBJ_TO_PTR(args[0]));
+    mp_obj_t exc = args[1];
+    if (n_args > 2 && args[2] != mp_const_none) {
+        exc = args[2];
+    }
+    nlr_raise(mp_make_raise_obj(exc));
+}
+static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(awaitable_throw_obj, 2, 4, awaitable_throw);
+
+static mp_obj_t awaitable_close(mp_obj_t self_in) {
+    awaitable_cancel(MP_OBJ_TO_PTR(self_in));
+    return mp_const_none;
+}
+static MP_DEFINE_CONST_FUN_OBJ_1(awaitable_close_obj, awaitable_close);
+
+// Readable once the operation has finished.
+static mp_uint_t awaitable_ioctl(mp_obj_t self_in, mp_uint_t request, uintptr_t arg, int *errcode) {
+    circuitpy_awaitable_obj_t *self = MP_OBJ_TO_PTR(self_in);
+    if (request != MP_STREAM_POLL) {
+        *errcode = MP_EINVAL;
+        return MP_STREAM_ERROR;
+    }
+    return CIRCUITPY_ASYNC_FLAG_IS_SET(&self->done) ? (arg & MP_STREAM_POLL_RD) : 0;
+}
+
+static const mp_stream_p_t awaitable_stream_p = {
+    .ioctl = awaitable_ioctl,
+};
+
+static const mp_rom_map_elem_t awaitable_locals_dict_table[] = {
+    { MP_ROM_QSTR(MP_QSTR___await__), MP_ROM_PTR(&awaitable_await_obj) },
+    { MP_ROM_QSTR(MP_QSTR_send), MP_ROM_PTR(&awaitable_send_obj) },
+    { MP_ROM_QSTR(MP_QSTR_throw), MP_ROM_PTR(&awaitable_throw_obj) },
+    { MP_ROM_QSTR(MP_QSTR_close), MP_ROM_PTR(&awaitable_close_obj) },
+    // Runs during GC: an operation still in flight must not write into freed memory.
+    { MP_ROM_QSTR(MP_QSTR___del__), MP_ROM_PTR(&awaitable_close_obj) },
+};
+static MP_DEFINE_CONST_DICT(awaitable_locals_dict, awaitable_locals_dict_table);
+
+MP_DEFINE_CONST_OBJ_TYPE(
+    circuitpy_awaitable_type,
+    MP_QSTR_awaitable,
+    MP_TYPE_FLAG_ITER_IS_ITERNEXT,
+    iter, awaitable_iternext,
+    protocol, &awaitable_stream_p,
+    locals_dict, &awaitable_locals_dict
+    );
+
+#endif // MICROPY_PY_ASYNC_AWAIT

supervisor/shared/awaitable.h

@@ -0,0 +1,84 @@
+// This file is part of the CircuitPython project: https://circuitpython.org
+//
+// SPDX-FileCopyrightText: Copyright (c) 2025 Scott Shawcroft for Adafruit Industries
+//
+// SPDX-License-Identifier: MIT
+
+// Awaitable functions implemented in C.
+//
+// CIRCUITPY_DEFINE_ASYNC_FUN_OBJ_KW(obj_name, n_args_min, fun_name) defines obj_name, a callable
+// that takes the same arguments as an MP_DEFINE_CONST_FUN_OBJ_KW function but returns an
+// awaitable. A call only checks the number of positional arguments; as with a coroutine in
+// CPython, nothing runs and the arguments are parsed only when the awaitable is awaited. Then:
+//
+//   void *fun_name(circuitpy_async_flag_t *done, size_t n_args, const mp_obj_t *pos_args,
+//       mp_map_t *kw_args)
+//       parses the arguments, starts the operation and returns a context. The operation sets
+//       *done, possibly from an interrupt, once it has finished.
+//   mp_obj_t common_hal_<fun_name>_end(void *context, circuitpy_async_flag_t *done)
+//       finishes the operation once *done is set and returns the result of the await.
+//   void common_hal_<fun_name>_cancel(void *context, circuitpy_async_flag_t *done)
+//       stops the operation, or finishes it if *done is already set. It must not allocate: it
+//       may run during GC. done tells which operation is meant, if the context has started
+//       another one since.
+//
+// CIRCUITPY_DEFINE_ASYNC_POLL_FUN_OBJ_KW also uses, for an operation with no interrupt to set *done:
+//   bool common_hal_<fun_name>_poll(void *context)
+//       tries to finish the operation and returns whether it has. The awaiting task tries again
+//       on each pass of the asyncio event loop.
+//
+// While an operation with an interrupt runs, the awaiting asyncio task waits on asyncio's I/O
+// queue: the awaitable is a stream that polls readable once *done is set.
+
+#pragma once
+
+#include "py/obj.h"
+#include "supervisor/shared/async_flag.h"
+
+#if MICROPY_PY_ASYNC_AWAIT
+
+typedef void *(*circuitpy_awaitable_start_fn)(circuitpy_async_flag_t *done, size_t n_args,
+    const mp_obj_t *pos_args, mp_map_t *kw_args);
+typedef mp_obj_t (*circuitpy_awaitable_end_fn)(void *context, circuitpy_async_flag_t *done);
+typedef void (*circuitpy_awaitable_cancel_fn)(void *context, circuitpy_async_flag_t *done);
+typedef bool (*circuitpy_awaitable_poll_fn)(void *context);
+
+// An async function: calling it returns a circuitpy_awaitable_obj_t.
+typedef struct {
+    mp_obj_base_t base;
+    uint16_t n_args_min;
+    circuitpy_awaitable_start_fn start;
+    circuitpy_awaitable_end_fn end;
+    circuitpy_awaitable_cancel_fn cancel;
+    circuitpy_awaitable_poll_fn poll;   // NULL when an interrupt sets done
+} circuitpy_async_fun_obj_t;
+
+extern const mp_obj_type_t circuitpy_async_fun_type;
+
+typedef struct {
+    mp_obj_base_t base;
+    const circuitpy_async_fun_obj_t *fun;
+    void *context;                  // returned by start()
+    circuitpy_async_flag_t done;
+    bool started;
+    bool finished;
+    uint16_t n_args;
+    uint16_t n_kw;
+    mp_obj_t args[];                // positional arguments, then keyword name/value pairs
+} circuitpy_awaitable_obj_t;
+
+extern const mp_obj_type_t circuitpy_awaitable_type;
+
+#define CIRCUITPY_DEFINE_ASYNC_FUN_OBJ_KW(obj_name, n_args_min, fun_name) \
+    const circuitpy_async_fun_obj_t obj_name = { \
+        {&circuitpy_async_fun_type}, n_args_min, fun_name, \
+        common_hal_##fun_name##_end, common_hal_##fun_name##_cancel, NULL \
+    }
+
+#define CIRCUITPY_DEFINE_ASYNC_POLL_FUN_OBJ_KW(obj_name, n_args_min, fun_name) \
+    const circuitpy_async_fun_obj_t obj_name = { \
+        {&circuitpy_async_fun_type}, n_args_min, fun_name, \
+        common_hal_##fun_name##_end, common_hal_##fun_name##_cancel, common_hal_##fun_name##_poll \
+    }
+
+#endif // MICROPY_PY_ASYNC_AWAIT

supervisor/supervisor.mk

@@ -27,6 +27,11 @@ ifeq ($(CIRCUITPY_LIB_TLSF),1)
 SRC_SUPERVISOR += lib/tlsf/tlsf.c
 endif
 
+# Awaitable C functions, for the async_* modules.
+ifeq ($(CIRCUITPY_ASYNC_SPI),1)
+SRC_SUPERVISOR += supervisor/shared/awaitable.c
+endif
+
 # For tlsf
 CFLAGS += -D_DEBUG=0