py/nlr: Implement jump callbacks.

NLR buffers are usually quite large (use lots of C stack) and expensive to
push and pop.  Some of the time they are only needed to perform clean up if
an exception happens, and then they re-raise the exception.

This commit allows optimizing that scenario by introducing a linked-list of
NLR callbacks that are called automatically when an exception is raised.
They are essentially a light-weight NLR handler that can implement a
"finally" block, i.e. clean-up when an exception is raised, or (by passing
`true` to nlr_pop_jump_callback) when execution leaves the scope.

Signed-off-by: Damien George <damien@micropython.org>

3 files changed, 56 insertions, 2 deletions

diff --git a/py/mpstate.h b/py/mpstate.h
index f6b911af56..80b49cb6b6 100644
--- a/py/mpstate.h
+++ b/py/mpstate.h
@@ -271,6 +271,7 @@ typedef struct _mp_state_thread_t {
     mp_obj_dict_t *dict_globals;
 
     nlr_buf_t *nlr_top;
+    nlr_jump_callback_node_t *nlr_jump_callback_top;
 
     // pending exception object (MP_OBJ_NULL if not pending)
     volatile mp_obj_t mp_pending_exception;
diff --git a/py/nlr.c b/py/nlr.c
index 92db8ffb19..7ab0c0955a 100644
--- a/py/nlr.c
+++ b/py/nlr.c
@@ -3,7 +3,7 @@
  *
  * The MIT License (MIT)
  *
- * Copyright (c) 2013-2017 Damien P. George
+ * Copyright (c) 2013-2023 Damien P. George
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -50,6 +50,36 @@ void nlr_pop(void) {
     *top = (*top)->prev;
 }
 
+void nlr_push_jump_callback(nlr_jump_callback_node_t *node, nlr_jump_callback_fun_t fun) {
+    nlr_jump_callback_node_t **top = &MP_STATE_THREAD(nlr_jump_callback_top);
+    node->prev = *top;
+    node->fun = fun;
+    *top = node;
+}
+
+void nlr_pop_jump_callback(bool run_callback) {
+    nlr_jump_callback_node_t **top = &MP_STATE_THREAD(nlr_jump_callback_top);
+    nlr_jump_callback_node_t *cur = *top;
+    *top = (*top)->prev;
+    if (run_callback) {
+        cur->fun(cur);
+    }
+}
+
+// This function pops and runs all callbacks that were registered after `nlr`
+// was pushed (via nlr_push).  It assumes:
+//  - a descending C stack,
+//  - that all nlr_jump_callback_node_t's in the linked-list pointed to by
+//    nlr_jump_callback_top are on the C stack
+// It works by popping each node in turn until the next node is NULL or above
+// the `nlr` pointer on the C stack (and so pushed before `nlr` was pushed).
+void nlr_call_jump_callbacks(nlr_buf_t *nlr) {
+    nlr_jump_callback_node_t **top = &MP_STATE_THREAD(nlr_jump_callback_top);
+    while (*top != NULL && (void *)*top < (void *)nlr) {
+        nlr_pop_jump_callback(true);
+    }
+}
+
 #if MICROPY_ENABLE_VM_ABORT
 NORETURN void nlr_jump_abort(void) {
     MP_STATE_THREAD(nlr_top) = MP_STATE_VM(nlr_abort);
diff --git a/py/nlr.h b/py/nlr.h
index b37370b224..62972dba62 100644
--- a/py/nlr.h
+++ b/py/nlr.h
@@ -3,7 +3,7 @@
  *
  * The MIT License (MIT)
  *
- * Copyright (c) 2013, 2014 Damien P. George
+ * Copyright (c) 2013-2023 Damien P. George
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -31,6 +31,7 @@
 
 #include <limits.h>
 #include <assert.h>
+#include <stdbool.h>
 
 #include "py/mpconfig.h"
 
@@ -123,6 +124,15 @@ struct _nlr_buf_t {
     #endif
 };
 
+typedef void (*nlr_jump_callback_fun_t)(void *ctx);
+
+typedef struct _nlr_jump_callback_node_t nlr_jump_callback_node_t;
+
+struct _nlr_jump_callback_node_t {
+    nlr_jump_callback_node_t *prev;
+    nlr_jump_callback_fun_t fun;
+};
+
 // Helper macros to save/restore the pystack state
 #if MICROPY_ENABLE_PYSTACK
 #define MP_NLR_SAVE_PYSTACK(nlr_buf) (nlr_buf)->pystack = MP_STATE_THREAD(pystack_cur)
@@ -140,6 +150,7 @@ struct _nlr_buf_t {
         nlr_jump_fail(val); \
     } \
     top->ret_val = val; \
+    nlr_call_jump_callbacks(top); \
     MP_NLR_RESTORE_PYSTACK(top); \
     *_top_ptr = top->prev; \
 
@@ -187,4 +198,16 @@ NORETURN void nlr_jump_fail(void *val);
 
 #endif
 
+// Push a callback on to the linked-list of NLR jump callbacks.  The `node` pointer must
+// be on the C stack.  The `fun` callback will be executed if an NLR jump is taken which
+// unwinds the C stack through this `node`.
+void nlr_push_jump_callback(nlr_jump_callback_node_t *node, nlr_jump_callback_fun_t fun);
+
+// Pop a callback from the linked-list of NLR jump callbacks.  The corresponding function
+// will be called if `run_callback` is true.
+void nlr_pop_jump_callback(bool run_callback);
+
+// Pop and call all NLR jump callbacks that were registered after `nlr` buffer was pushed.
+void nlr_call_jump_callbacks(nlr_buf_t *nlr);
+
 #endif // MICROPY_INCLUDED_PY_NLR_H