Add `at_cleanup()` function

8 files changed, 116 insertions, 13 deletions

diff --git a/CHANGES.md b/CHANGES.md
index fb4ce929..42db4715 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -12,8 +12,9 @@
 - Syntax for text literals and inline C code has been simplified somewhat.
 - Syntax for tables has changed to use colons (`{k: v}`) instead of equals
   (`{k=v}`).
-- Added `Path.lines()`
-- Added `Text.find(text, target, start=1)`
+- Added `Path.lines()`.
+- Added `Text.find(text, target, start=1)`.
+- Added `at_cleanup()` to register cleanup functions.
 - Deprecated:
   - Sets are no longer a separate type with separate methods.
     - Instead of sets, use tables with a value type of `{KeyType:Empty}`.
diff --git a/api/api.md b/api/api.md
index 02ad054e..5c9dc9c4 100644
--- a/api/api.md
+++ b/api/api.md
@@ -33,10 +33,34 @@ force_tty | `Bool` | Whether or not to force the use of /dev/tty.  | `yes`
 assert ask("What's your name? ") == "Arthur Dent"
 
 ```
+## at_cleanup
+
+```tomo
+at_cleanup : func(fn: func() -> Void)
+```
+
+Register a function that runs at cleanup time for Tomo programs. Cleanup time happens when a program exits (see `atexit()` in C), or immediately before printing error messages in a call to `fail()`. This allows for terminal cleanup so error messages can be visible as the program shuts down.
+
+Use this API very carefully, because errors that occur during cleanup functions may make it extremely hard to figure out what's going on. Cleanup functions should be designed to not error under any circumstances.
+
+Argument | Type | Description | Default
+---------|------|-------------|---------
+fn | `func()` | A function to run at cleanup time.  | -
+
+**Return:** Nothing.
+
+
+**Example:**
+```tomo
+at_cleanup(func()
+    (/tmp/file.txt).remove(ignore_missing=yes)
+)
+
+```
 ## exit
 
 ```tomo
-exit : func(message: Text? = none, status: Int32 = Int32(1) -> Void)
+exit : func(message: Text? = none, status: Int32 = Int32(1) -> Abort)
 ```
 
 Exits the program with a given status and optionally prints a message.
diff --git a/api/builtins.md b/api/builtins.md
index 0b06a41b..6d042741 100644
--- a/api/builtins.md
+++ b/api/builtins.md
@@ -33,10 +33,34 @@ force_tty | `Bool` | Whether or not to force the use of /dev/tty.  | `yes`
 assert ask("What's your name? ") == "Arthur Dent"
 
 ```
+## at_cleanup
+
+```tomo
+at_cleanup : func(fn: func() -> Void)
+```
+
+Register a function that runs at cleanup time for Tomo programs. Cleanup time happens when a program exits (see `atexit()` in C), or immediately before printing error messages in a call to `fail()`. This allows for terminal cleanup so error messages can be visible as the program shuts down.
+
+Use this API very carefully, because errors that occur during cleanup functions may make it extremely hard to figure out what's going on. Cleanup functions should be designed to not error under any circumstances.
+
+Argument | Type | Description | Default
+---------|------|-------------|---------
+fn | `func()` | A function to run at cleanup time.  | -
+
+**Return:** Nothing.
+
+
+**Example:**
+```tomo
+at_cleanup(func()
+    (/tmp/file.txt).remove(ignore_missing=yes)
+)
+
+```
 ## exit
 
 ```tomo
-exit : func(message: Text? = none, status: Int32 = Int32(1) -> Void)
+exit : func(message: Text? = none, status: Int32 = Int32(1) -> Abort)
 ```
 
 Exits the program with a given status and optionally prints a message.
diff --git a/api/builtins.yaml b/api/builtins.yaml
index 2eae5340..764a1bd4 100644
--- a/api/builtins.yaml
+++ b/api/builtins.yaml
@@ -38,7 +38,7 @@ exit:
   description: >
     Exits the program with a given status and optionally prints a message.
   return:
-    type: 'Void'
+    type: 'Abort'
     description: >
       This function never returns.
   args:
@@ -56,6 +56,32 @@ exit:
   example: |
     exit(status=1, "Goodbye forever!")
 
+at_cleanup:
+  short: register a cleanup function
+  description: >
+    Register a function that runs at cleanup time for Tomo programs. Cleanup
+    time happens when a program exits (see `atexit()` in C), or immediately
+    before printing error messages in a call to `fail()`. This allows for
+    terminal cleanup so error messages can be visible as the program shuts
+    down.
+  note: >
+    Use this API very carefully, because errors that occur during cleanup
+    functions may make it extremely hard to figure out what's going on. Cleanup
+    functions should be designed to not error under any circumstances.
+  args:
+    fn:
+      type: 'func()'
+      description: >
+        A function to run at cleanup time.
+  return:
+    type: 'Void'
+    description: >
+      Nothing.
+  example: |
+    at_cleanup(func()
+        (/tmp/file.txt).remove(ignore_missing=yes)
+    )
+
 getenv:
   short: get an environment variable
   description: >
diff --git a/man/man3/tomo-exit.3 b/man/man3/tomo-exit.3
index 48e0bd79..50c022b8 100644
--- a/man/man3/tomo-exit.3
+++ b/man/man3/tomo-exit.3
@@ -2,14 +2,14 @@
 .\" Copyright (c) 2025 Bruce Hill
 .\" All rights reserved.
 .\"
-.TH exit 3 2025-05-17 "Tomo man-pages"
+.TH exit 3 2025-11-23 "Tomo man-pages"
 .SH NAME
 exit \- exit the program
 .SH LIBRARY
 Tomo Standard Library
 .SH SYNOPSIS
 .nf
-.BI exit\ :\ func(message:\ Text?\ =\ none,\ status:\ Int32\ =\ Int32(1)\ ->\ Void)
+.BI exit\ :\ func(message:\ Text?\ =\ none,\ status:\ Int32\ =\ Int32(1)\ ->\ Abort)
 .fi
 .SH DESCRIPTION
 Exits the program with a given status and optionally prints a message.
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);                                     \