Add `at_cleanup()` function

3 files changed, 34 insertions, 6 deletions

diff --git a/src/environment.c b/src/environment.c
index 3a2995f7..d8c82513 100644
--- a/src/environment.c
+++ b/src/environment.c
@@ -529,17 +529,18 @@ env_t *global_env(bool source_mapping) {
     struct {
         const char *name, *code, *type_str;
     } global_vars[] = {
-        {"USE_COLOR", "USE_COLOR", "Bool"},
+        {"EMPTY", "EMPTY", "Empty"},
         {"TOMO_VERSION", "TOMO_VERSION_TEXT", "Text"},
-        {"say", "say", "func(text:Text, newline=yes)"},
-        {"print", "say", "func(text:Text, newline=yes)"},
-        {"getenv", "getenv_text", "func(name:Text -> Text?)"},
-        {"setenv", "setenv_text", "func(name:Text, value:Text -> Text?)"},
+        {"USE_COLOR", "USE_COLOR", "Bool"},
         {"ask", "ask", "func(prompt:Text, bold=yes, force_tty=yes -> Text?)"},
+        {"at_cleanup", "tomo_at_cleanup", "func(fn:func())"},
         {"exit", "tomo_exit", "func(message:Text?=none, code=Int32(1) -> Abort)"},
         {"fail", "fail_text", "func(message:Text -> Abort)"},
+        {"getenv", "getenv_text", "func(name:Text -> Text?)"},
+        {"print", "say", "func(text:Text, newline=yes)"},
+        {"say", "say", "func(text:Text, newline=yes)"},
+        {"setenv", "setenv_text", "func(name:Text, value:Text -> Text?)"},
         {"sleep", "sleep_num", "func(seconds:Num)"},
-        {"EMPTY", "EMPTY", "Empty"},
     };
 
     for (size_t i = 0; i < sizeof(global_vars) / sizeof(global_vars[0]); i++) {
diff --git a/src/stdlib/stdlib.c b/src/stdlib/stdlib.c
index 157dbfd4..21547efe 100644
--- a/src/stdlib/stdlib.c
+++ b/src/stdlib/stdlib.c
@@ -72,6 +72,7 @@ void tomo_init(void) {
     sigemptyset(&sigact.sa_mask);
     sigact.sa_flags = 0;
     sigaction(SIGILL, &sigact, (struct sigaction *)NULL);
+    atexit(tomo_cleanup);
 }
 
 public
@@ -219,3 +220,25 @@ OptionalText_t getenv_text(Text_t name) {
 
 public
 void setenv_text(Text_t name, Text_t value) { setenv(Text$as_c_string(name), Text$as_c_string(value), 1); }
+
+typedef struct cleanup_s {
+    Closure_t cleanup_fn;
+    struct cleanup_s *next;
+} cleanup_t;
+
+static cleanup_t *cleanups = NULL;
+
+public
+void tomo_at_cleanup(Closure_t fn) { cleanups = new (cleanup_t, .cleanup_fn = fn, .next = cleanups); }
+
+public
+void tomo_cleanup(void) {
+    while (cleanups) {
+        // NOTE: we *must* remove the cleanup function from the stack before calling it,
+        // otherwise it will cause an infinite loop if the cleanup function fails or exits.
+        void (*run_cleanup)(void *) = cleanups->cleanup_fn.fn;
+        void *userdata = cleanups->cleanup_fn.userdata;
+        cleanups = cleanups->next;
+        run_cleanup(userdata);
+    }
+}
diff --git a/src/stdlib/stdlib.h b/src/stdlib/stdlib.h
index aadce27c..392b5f23 100644
--- a/src/stdlib/stdlib.h
+++ b/src/stdlib/stdlib.h
@@ -16,9 +16,12 @@ extern bool USE_COLOR;
 extern Text_t TOMO_VERSION_TEXT;
 
 void tomo_init(void);
+void tomo_at_cleanup(Closure_t fn);
+void tomo_cleanup(void);
 
 #define fail(...)                                                                                                      \
     ({                                                                                                                 \
+        tomo_cleanup();                                                                                                \
         fflush(stdout);                                                                                                \
         if (USE_COLOR) fputs("\x1b[31;7m ==================== ERROR ==================== \033[m\n\n", stderr);         \
         else fputs("==================== ERROR ====================\n\n", stderr);                                     \
@@ -35,6 +38,7 @@ void tomo_init(void);
 
 #define fail_source(filename, start, end, ...)                                                                         \
     ({                                                                                                                 \
+        tomo_cleanup();                                                                                                \
         fflush(stdout);                                                                                                \
         if (USE_COLOR) fputs("\x1b[31;7m ==================== ERROR ==================== \n\n\x1b[0;1m", stderr);      \
         else fputs("==================== ERROR ====================\n\n", stderr);                                     \