py/modatexit: Switch to using list.

Signed-off-by: Andrew Leech <[email protected]>
andrewleech · Nov 24, 2024 · 0bd99fd · 0bd99fd
1 parent f80c6c7
commit 0bd99fd
Showing 1 changed file with 19 additions and 38 deletions.
diff --git a/py/modatexit.c b/py/modatexit.c
@@ -34,45 +34,26 @@
 
 #if MICROPY_PY_ATEXIT
 
-typedef struct _m_atexit_node_t {
-    struct _m_atexit_node_t *prev;
-    struct _m_atexit_node_t *next;
-    mp_obj_t fn;
-} m_atexit_node_t;
-
-// atexit.register(function): Functions are called LIFO when soft-reset / exit is called.
 mp_obj_t mp_atexit_register(mp_obj_t function) {
     if (!mp_obj_is_callable(function)) {
         mp_raise_ValueError(MP_ERROR_TEXT("function not callable"));
     }
-    m_atexit_node_t *node = m_malloc(sizeof(m_atexit_node_t));
-    if (MP_STATE_VM(atexit) != NULL) {
-        MP_STATE_VM(atexit)->prev = node;
+    if (MP_STATE_VM(atexit) == NULL) {
+        MP_STATE_VM(atexit) = mp_obj_new_list(0, mp_const_none);
     }
-    node->fn = function;
-    node->prev = NULL;
-    node->next = MP_STATE_VM(atexit);
-    MP_STATE_VM(atexit) = node;
+    mp_obj_list_append(MP_STATE_VM(atexit), function);
     // return the passed in function so this can be used as a decorator
     return function;
 }
 static MP_DEFINE_CONST_FUN_OBJ_1(mp_atexit_register_obj, mp_atexit_register);
 
 #if MICROPY_PY_ATEXIT_UNREGISTER
 mp_obj_t mp_atexit_unregister(mp_obj_t function) {
-    m_atexit_node_t *node = MP_STATE_VM(atexit);
-    while (node != NULL) {
-        if (mp_obj_equal(node->fn, function)) {
-            if (node->next != NULL) {
-                node->next->prev = node->prev;
-            }
-            if (node->prev != NULL) {
-                node->prev->next = node->next;
-            } else {
-                MP_STATE_VM(atexit) = node->next;
-            }
+    nlr_buf_t nlr;
+    if (nlr_push(&nlr) == 0) {
+        while (MP_STATE_VM(atexit) != NULL) {
+            mp_obj_list_remove(MP_STATE_VM(atexit), function);
         }
-        node = node->next;
     }
     return mp_const_none;
 }
@@ -86,30 +67,30 @@ static const mp_rom_map_elem_t mp_module_atexit_globals_table[] = {
     { MP_ROM_QSTR(MP_QSTR_unregister), MP_ROM_PTR(&mp_atexit_unregister_obj) },
     #endif
 };
-
 static MP_DEFINE_CONST_DICT(mp_module_atexit_globals, mp_module_atexit_globals_table);
 
 const mp_obj_module_t mp_module_atexit = {
     .base = { &mp_type_module },
     .globals = (mp_obj_dict_t *)&mp_module_atexit_globals,
 };
 
-MP_REGISTER_ROOT_POINTER(struct _m_atexit_node_t *atexit);
+MP_REGISTER_ROOT_POINTER(mp_obj_list_t * atexit);
 MP_REGISTER_MODULE(MP_QSTR_atexit, mp_module_atexit);
 
 int mp_atexit_execute(void) {
     int exit_code = 0;
-    // This function is intended to be run by a port during its soft-reset / exit.
-    // walk down linked list last in / first out and execute each function.
-    // Beware, the sys.settrace function should be disabled before running sys.atexit.
-    while (MP_STATE_VM(atexit) != NULL) {
-        nlr_buf_t nlr;
-        if (nlr_push(&nlr) == 0) {
-            mp_call_function_0(MP_STATE_VM(atexit)->fn);
-        } else {
-            exit_code = pyexec_handle_uncaught_exception(nlr.ret_val);
+    if (MP_STATE_VM(atexit) != NULL) {
+        mp_obj_list_t *list = MP_STATE_VM(atexit);
+        for (size_t i = 0; i < list->len; i++) {
+            mp_obj_t function = list->items[i];
+
+            nlr_buf_t nlr;
+            if (nlr_push(&nlr) == 0) {
+                mp_call_function_0(function);
+            } else {
+                exit_code = pyexec_handle_uncaught_exception(nlr.ret_val);
+            }
         }
-        MP_STATE_VM(atexit) = MP_STATE_VM(atexit)->next;
     }
     return exit_code;
 }