fix: run shell main on a thread with the configured stack size (#14343)

This PR fixes the new 1GB stack size not being used for the main `lean`
thread itself, e.g. for serialization or `--run`.

Co-authored-by: Claude <noreply@anthropic.com>

Author: Kha

src/runtime/thread.cpp

@@ -167,8 +167,7 @@ extern "C" LEAN_EXPORT lean_obj_res lean_internal_set_thread_stack_size(size_t s
     return lean_box(0);
 }
 
-extern "C" LEAN_EXPORT lean_object * lean_run_main(lean_object * (*main_fn)(int, char **), int argc, char ** argv) {
-#ifdef LEAN_MULTI_THREAD
+LEAN_EXPORT void set_thread_stack_size_from_env() {
     const char * stack_size_env = std::getenv("LEAN_STACK_SIZE_KB");
     if (stack_size_env) {
         size_t sz = std::strtoull(stack_size_env, nullptr, 10);
@@ -177,18 +176,24 @@ extern "C" LEAN_EXPORT lean_object * lean_run_main(lean_object * (*main_fn)(int,
             lthread::set_thread_stack_size(sz);
         }
     }
+}
+
+LEAN_EXPORT void run_with_thread_stack(std::function<void()> const & fn) {
     const char * use_thread_env = std::getenv("LEAN_MAIN_USE_THREAD");
     if (use_thread_env && std::strcmp(use_thread_env, "0") == 0) {
-        return main_fn(argc, argv);
+        fn();
+    } else {
+        // Start new thread to use given/default stack size
+        lthread t(fn);
+        t.join();
     }
-    // Start new thread to use given/default stack size
+}
+
+extern "C" LEAN_EXPORT lean_object * lean_run_main(lean_object * (*main_fn)(int, char **), int argc, char ** argv) {
+    set_thread_stack_size_from_env();
     lean_object * res = nullptr;
-    lthread t([&]() { res = main_fn(argc, argv); });
-    t.join();
+    run_with_thread_stack([&]() { res = main_fn(argc, argv); });
     return res;
-#else
-    return main_fn(argc, argv);
-#endif
 }
 
 LEAN_THREAD_VALUE(bool, g_finalizing, false);

src/runtime/thread.h

@@ -267,4 +267,6 @@ LEAN_EXPORT void register_thread_local_reset_fn(std::function<void()> fn);
    We invoke this function before processing a command
    and before executing a task. */
 LEAN_EXPORT void reset_thread_local();
+LEAN_EXPORT void run_with_thread_stack(std::function<void()> const & fn);
+LEAN_EXPORT void set_thread_stack_size_from_env();
 }

src/util/shell.cpp

@@ -321,6 +321,10 @@ extern "C" LEAN_EXPORT int lean_main(int argc, char ** argv) {
     }
     consume_io_result(lean_enable_initializer_execution());
 
+    // Default the configured thread stack size from the environment as in `lean_run_main`;
+    // `--tstack` below overrides it.
+    set_thread_stack_size_from_env();
+
     int rc;
     object_ref shell_opts;
     try {
@@ -348,12 +352,16 @@ extern "C" LEAN_EXPORT int lean_main(int argc, char ** argv) {
 
     scoped_task_manager scope_task_man(get_shell_num_threads(shell_opts));
 
-    try {
-        return run_shell_main(argc - optind, argv + optind, shell_opts);
-    } catch (lean::throwable & ex) {
-        std::cerr << ex.what() << "\n";
-    } catch (std::bad_alloc & ex) {
-        std::cerr << "out of memory" << std::endl;
-    }
-    return 1;
+    int shell_rc = 1;
+    // Do not rely on OS thread stack size, as for Lean executables
+    run_with_thread_stack([&]() {
+        try {
+            shell_rc = run_shell_main(argc - optind, argv + optind, shell_opts);
+        } catch (lean::throwable & ex) {
+            std::cerr << ex.what() << "\n";
+        } catch (std::bad_alloc & ex) {
+            std::cerr << "out of memory" << std::endl;
+        }
+    });
+    return shell_rc;
 }

tests/compile/shell_thread_stack.lean

@@ -0,0 +1,36 @@
+import Lean.CompactedRegion
+
+open Lean
+
+/-!
+Regression test that the `lean` driver runs its main logic on a thread with the configured
+Lean thread stack size instead of directly on the OS main thread, whose fixed default stack
+(often 8 MiB) is too small for the runtime's recursive algorithms. Exercises two stack-hungry
+paths that run on the driver's main thread: interpreting a non-tail-recursive function at a
+depth beyond the OS default stack's interpreter headroom, and compacting a deeply nested
+object graph with `CompactedRegion.save`, whose compactor recurses once per nesting level.
+-/
+
+def depth : Nat := 100000
+
+def deepRec : Nat → Nat
+  | 0 => 0
+  | n+1 => 1 + deepRec n
+
+def mkDeepList (n : Nat) : List Nat := Id.run do
+  let mut xs := []
+  for i in [0:n] do
+    xs := i :: xs
+  return xs
+
+unsafe def main : IO UInt32 := do
+  unless deepRec depth == depth do
+    throw <| IO.userError "deepRec mismatch"
+  let tmpFile : System.FilePath := "./_tmp_shell_thread_stack.olean"
+  let xs := mkDeepList (4 * depth)
+  let _ ← CompactedRegion.save tmpFile `ShellThreadStack xs #[] none
+  let (ys, _region) ← CompactedRegion.read (α := List Nat) tmpFile #[]
+  unless ys.length == 4 * depth do
+    throw <| IO.userError s!"round-trip length mismatch: expected {4 * depth}, got {ys.length}"
+  IO.FS.removeFile tmpFile
+  return 0