feat: respect lean --setup module name in code generation (#8780)

This PR makes Lean code generation respect the module name provided
through `lean --setup`.

This is accomplished by porting to Lean the portion of `shell.cpp` that
spans running the frontend to exiting the process. This makes it easier
to load the module setup and control how its name is passed to the code
generation functions. This port attempts to minimize the changes made to
Lean. It marks the new Lean functions `private` and tries to preserve as
faithfully as possible the behavior of the original C++ code. Exposing
the new Lean interface publicly and/or further improving the code now
that is written in Lean is left for the future.

src/Lean.lean

@@ -43,3 +43,4 @@ import Lean.Namespace
 import Lean.EnvExtension
 import Lean.ErrorExplanation
 import Lean.DefEqAttrib
+import Lean.Shell

src/Lean/Elab/Frontend.lean

@@ -132,20 +132,19 @@ def process (input : String) (env : Environment) (opts : Options) (fileName : Op
   let s ← IO.processCommands inputCtx { : Parser.ModuleParserState } (Command.mkState env {} opts)
   pure (s.commandState.env, s.commandState.messages)
 
-@[export lean_run_frontend]
 def runFrontend
     (input : String)
     (opts : Options)
     (fileName : String)
     (mainModuleName : Name)
     (trustLevel : UInt32 := 0)
-    (oleanFileName? : Option String := none)
-    (ileanFileName? : Option String := none)
+    (oleanFileName? : Option System.FilePath := none)
+    (ileanFileName? : Option System.FilePath := none)
     (jsonOutput : Bool := false)
     (errorOnKinds : Array Name := #[])
     (plugins : Array System.FilePath := #[])
     (printStats : Bool := false)
-    (setupFileName? : Option System.FilePath := none)
+    (setup? : Option ModuleSetup := none)
     : IO (Option Environment) := do
   let startTime := (← IO.monoNanosNow).toFloat / 1000000000
   let inputCtx := Parser.mkInputContext input fileName
@@ -154,8 +153,7 @@ def runFrontend
   let opts := Elab.async.setIfNotSet opts true
   let ctx := { inputCtx with }
   let setup stx := do
-    if let some file := setupFileName? then
-      let setup ← ModuleSetup.load file
+    if let some setup := setup? then
       liftM <| setup.dynlibs.forM Lean.loadDynlib
       return .ok {
         trustLevel
@@ -164,7 +162,7 @@ def runFrontend
         imports := setup.imports
         plugins := plugins ++ setup.plugins
         modules := setup.modules
-        -- override cmdline options with header options
+        -- override cmdline options with setup options
         opts := opts.mergeBy (fun _ _ hOpt => hOpt) setup.options.toOptions
       }
     else

src/Lean/Shell.lean

@@ -0,0 +1,84 @@
+/-
+Copyright (c) 2025 Lean FRO. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura, Mac Malone
+-/
+prelude
+import Lean.Elab.Frontend
+import Lean.Compiler.IR.EmitC
+
+/-  Lean companion to  `shell.cpp` -/
+
+open System
+
+namespace Lean
+
+/- Runs the `main` function of the module with `args` using the Lean interpreter. -/
+@[extern "lean_run_main"]
+private opaque runMain (env : @& Environment) (opts : @& Options) (args : @& List String) : BaseIO UInt32
+
+/--
+Initializes the LLVM subsystem.
+If Lean lacks LLVM support, this function will fail with an assertion violation.
+-/
+@[extern "lean_init_llvm"]
+private opaque initLLVM : IO Unit
+
+/--
+Emits LLVM bitcode for the module.
+Before calling this function, the LLVM subsystem must first be successfully initialized.
+-/
+@[extern "lean_emit_llvm"]
+private opaque emitLLVM (env : Environment) (modName : Name) (filepath : FilePath) : IO Unit
+
+/-- Print all profiling times (if any) to standard error. -/
+@[extern "lean_display_cumulative_profiling_times"]
+private opaque displayCumulativeProfilingTimes : BaseIO Unit
+
+/-- Whether Lean was built with an address sanitizer enabled. -/
+@[extern "lean_internal_has_address_sanitizer"]
+private opaque Internal.hasAddressSanitizer (_ : Unit) : Bool
+
+@[export lean_shell_main]
+private def shellMain
+    (args : List String)
+    (input : String)
+    (opts : Options)
+    (fileName : String)
+    (mainModuleName : Name)
+    (trustLevel : UInt32 := 0)
+    (oleanFileName? : Option System.FilePath := none)
+    (ileanFileName? : Option System.FilePath := none)
+    (cFileName? : Option System.FilePath := none)
+    (bcFileName? : Option System.FilePath := none)
+    (jsonOutput : Bool := false)
+    (errorOnKinds : Array Name := #[])
+    (printStats : Bool := false)
+    (setupFileName? : Option System.FilePath := none)
+    (run : Bool := false)
+    : IO UInt32 := do
+  let setup? ← setupFileName?.mapM ModuleSetup.load
+  let mainModuleName := setup?.map (·.name) |>.getD mainModuleName
+  let env? ← Elab.runFrontend input opts fileName mainModuleName trustLevel
+      oleanFileName? ileanFileName? jsonOutput errorOnKinds #[] printStats setup?
+  if let some env := env? then
+    if run then
+      return ← runMain env opts args
+    if let some c := cFileName? then
+      let .ok out ← IO.FS.Handle.mk c .write |>.toBaseIO
+        | IO.eprintln s!"failed to create '{c}'"
+          return 1
+      profileitIO "C code generation" opts do
+        let data ← IO.ofExcept <| IR.emitC env mainModuleName
+        out.write data.toUTF8
+    if let some bc := bcFileName? then
+      initLLVM
+      profileitIO "LLVM code generation" opts do
+        emitLLVM env mainModuleName bc
+  displayCumulativeProfilingTimes
+  if Internal.hasAddressSanitizer () then
+    return if env?.isSome then 0 else 1
+  else
+    -- When not using the address/leak sanitizer, we interrupt execution without garbage collecting.
+    -- This is useful when profiling improvements to Lean startup time.
+    IO.Process.exit <| if env?.isSome then 0 else 1

src/library/compiler/ir_interpreter.cpp

@@ -1050,27 +1050,18 @@ public:
         return r;
     }
 
-    uint32 run_main(int argc, char * argv[]) {
+    uint32 run_main(list_ref<string_ref> const & args) {
         decl d = get_decl("main");
         array_ref<param> const & params = decl_params(d);
-        buffer<object *> args;
+        buffer<object *> call_args;
         if (params.size() == 2) { // List String -> IO _
-            lean_object * in = lean_box(0);
-            int i = argc;
-            while (i > 0) {
-                i--;
-                lean_object * n = lean_alloc_ctor(1, 2, 0);
-                lean_ctor_set(n, 0, lean_mk_string(argv[i]));
-                lean_ctor_set(n, 1, in);
-                in = n;
-            }
-            args.push_back(in);
+            call_args.push_back(args.to_obj_arg());
         } else { // IO _
             lean_assert(params.size() == 1);
         }
         object * w = io_mk_world();
-        args.push_back(w);
-        w = call_boxed("main", args.size(), &args[0]);
+        call_args.push_back(w);
+        w = call_boxed("main", call_args.size(), &call_args[0]);
         if (io_result_is_ok(w)) {
             int ret = 0;
             lean::expr ret_ty = m_env.get("main").get_type();
@@ -1137,8 +1128,14 @@ object * run_boxed_kernel(environment const & env, options const & opts, name co
     return run_boxed(elab_environment_of_kernel_env(env), opts, fn, n, args);
 }
 
-uint32 run_main(elab_environment const & env, options const & opts, int argv, char * argc[]) {
-    return interpreter::with_interpreter<uint32>(env, opts, "main", [&](interpreter & interp) { return interp.run_main(argv, argc); });
+uint32 run_main(elab_environment const & env, options const & opts, list_ref<string_ref> const & args) {
+    return interpreter::with_interpreter<uint32>(env, opts, "main", [&](interpreter & interp) { return interp.run_main(args); });
+}
+
+/* runMain (env : Environment) (opts : Iptions) (args : List String) : BaseIO UInt32 */
+extern "C" LEAN_EXPORT obj_res lean_run_main(b_obj_arg env, b_obj_arg opts, b_obj_arg args, obj_arg) {
+    uint32 ret = run_main(TO_REF(elab_environment, env), TO_REF(options, opts), TO_REF(list_ref<string_ref>, args));
+    return io_result_mk_ok(box(ret));
 }
 
 extern "C" LEAN_EXPORT object * lean_eval_const(object * env, object * opts, object * c) {

src/library/compiler/ir_interpreter.h

@@ -12,7 +12,7 @@ namespace lean {
 namespace ir {
 /** \brief Run `n` using the "boxed" ABI, i.e. with all-owned parameters. */
 object * run_boxed(elab_environment const & env, options const & opts, name const & fn, unsigned n, object **args);
-LEAN_EXPORT uint32 run_main(elab_environment const & env, options const & opts, int argv, char * argc[]);
+LEAN_EXPORT uint32 run_main(elab_environment const & env, options const & opts, list_ref<string_ref> const & args);
 }
 void initialize_ir_interpreter();
 void finalize_ir_interpreter();

src/library/compiler/llvm.cpp

@@ -37,6 +37,20 @@ Lean's IR.
 #pragma GCC diagnostic ignored "-Wunused-parameter"
 #endif
 
+namespace lean {
+/*  initLLVM : IO Unit */
+extern "C" obj_res initialize_Lean_Compiler_IR_EmitLLVM(uint8_t builtin, obj_arg);
+extern "C" LEAN_EXPORT obj_res lean_init_llvm(obj_arg) {
+    return initialize_Lean_Compiler_IR_EmitLLVM(/*builtin*/ false, io_mk_world());
+}
+
+/*  emitLLVM (env : Environment) (modName : Name) (filepath : FilePath) : IO Unit */
+extern "C" obj_res lean_ir_emit_llvm(obj_arg env, obj_arg mod_name, obj_arg filepath, obj_arg);
+extern "C" LEAN_EXPORT obj_res lean_emit_llvm(obj_arg env, obj_arg mod_name, obj_arg filepath, obj_arg) {
+    return lean_ir_emit_llvm(env, mod_name, filepath, io_mk_world());
+}
+}
+
 extern "C" LEAN_EXPORT lean_object* lean_llvm_initialize_target_info(lean_object * /* w */) {
 
 #ifdef LEAN_LLVM

src/library/time_task.cpp

@@ -38,6 +38,12 @@ void display_cumulative_profiling_times(std::ostream & out) {
     out << ss.str();
 }
 
+/* displayCumulativeProfilingTimes : BaseIO Unit */
+extern "C" LEAN_EXPORT obj_res lean_display_cumulative_profiling_times(obj_arg) {
+   display_cumulative_profiling_times(std::cerr);
+   return lean_io_result_mk_ok(box(0));
+}
+
 void initialize_time_task() {
     g_cum_times_mutex = new mutex;
     g_cum_times = new std::map<std::string, second_duration>;

src/runtime/platform.cpp

@@ -42,11 +42,23 @@ extern "C" LEAN_EXPORT uint8 lean_system_platform_emscripten(obj_arg) {
 
 extern "C" object * lean_get_githash(obj_arg) { return lean_mk_string(LEAN_GITHASH); }
 
-extern "C" LEAN_EXPORT uint8_t lean_internal_has_llvm_backend(lean_obj_arg _unit) {
+extern "C" LEAN_EXPORT uint8 lean_internal_has_llvm_backend(obj_arg) {
 #ifdef LEAN_LLVM
     return 1;
 #else
     return 0;
 #endif
 }
+
+extern "C" LEAN_EXPORT uint8 lean_internal_has_address_sanitizer(obj_arg) {
+#if defined(__has_feature)
+#if __has_feature(address_sanitizer)
+    return 1;
+#else
+    return 0;
+#endif
+#else
+    return 0;
+#endif
+}
 }

src/util/shell.cpp

@@ -170,11 +170,6 @@ using namespace lean; // NOLINT
 #define LEAN_SERVER_DEFAULT_MAX_HEARTBEAT 100000
 #endif
 
-extern "C" void *initialize_Lean_Compiler_IR_EmitLLVM(uint8_t builtin,
-                                                      lean_object *);
-extern "C" object *lean_ir_emit_llvm(object *env, object *mod_name,
-                                     object *filepath, object *w);
-
 static void display_header(std::ostream & out) {
     out << "Lean (version " << get_version_string() << ", " << LEAN_STR(LEAN_BUILD_TYPE) << ")\n";
 }
@@ -330,7 +325,8 @@ options set_config_option(options const & opts, char const * in) {
 }
 
 namespace lean {
-extern "C" object * lean_run_frontend(
+extern "C" object * lean_shell_main(
+    object * args,
     object * input,
     object * opts,
     object * filename,
@@ -338,26 +334,38 @@ extern "C" object * lean_run_frontend(
     uint32_t trust_level,
     object * olean_filename,
     object * ilean_filename,
+    object * c_filename,
+    object * bc_filename,
     uint8_t  json_output,
     object * error_kinds,
-    object * plugins,
     bool     print_stats,
-    object * header_file_name,
+    object * setup_file_name,
+    bool     run,
     object * w
 );
-option_ref<elab_environment> run_new_frontend(
+uint32 run_shell_main(
+    int argc, char* argv[],
     std::string const & input,
     options const & opts, std::string const & file_name,
     name const & main_module_name,
     uint32_t trust_level,
     optional<std::string> const & olean_file_name,
     optional<std::string> const & ilean_file_name,
+    optional<std::string> const & c_file_name,
+    optional<std::string> const & bc_file_name,
     uint8_t json_output,
     array_ref<name> const & error_kinds,
     bool print_stats,
-    optional<std::string> const & setup_file_name
+    optional<std::string> const & setup_file_name,
+    bool run
 ) {
-    return get_io_result<option_ref<elab_environment>>(lean_run_frontend(
+    list_ref<string_ref> args;
+    while (argc > 0) {
+        argc--;
+        args = list_ref<string_ref>(string_ref(argv[argc]), args);
+    }
+    return get_io_scalar_result<uint32>(lean_shell_main(
+        args.steal(),
         mk_string(input),
         opts.to_obj_arg(),
         mk_string(file_name),
@@ -365,11 +373,13 @@ option_ref<elab_environment> run_new_frontend(
         trust_level,
         olean_file_name ? mk_option_some(mk_string(*olean_file_name)) : mk_option_none(),
         ilean_file_name ? mk_option_some(mk_string(*ilean_file_name)) : mk_option_none(),
+        c_file_name ? mk_option_some(mk_string(*c_file_name)) : mk_option_none(),
+        bc_file_name ? mk_option_some(mk_string(*bc_file_name)) : mk_option_none(),
         json_output,
         error_kinds.to_obj_arg(),
-        mk_empty_array(),
         print_stats,
         setup_file_name ? mk_option_some(mk_string(*setup_file_name)) : mk_option_none(),
+        run,
         io_mk_world()
     ));
 }
@@ -765,48 +775,12 @@ extern "C" LEAN_EXPORT int lean_main(int argc, char ** argv) {
 
         if (!main_module_name)
             main_module_name = name("_stdin");
-        option_ref<elab_environment> opt_env = run_new_frontend(
+        return run_shell_main(
+            argc - optind, argv + optind,
             contents, opts, mod_fn, *main_module_name, trust_lvl,
-            olean_fn, ilean_fn, json_output, error_kinds, stats, setup_fn
+            olean_fn, ilean_fn, c_output, llvm_output,
+            json_output, error_kinds, stats, setup_fn, run
         );
-
-        if (opt_env) {
-            elab_environment env = opt_env.get_val();
-            if (run) {
-                uint32 ret = ir::run_main(env, opts, argc - optind, argv + optind);
-                return ret;
-            }
-            if (c_output) {
-                std::ofstream out(*c_output, std::ios_base::binary);
-                if (out.fail()) {
-                    std::cerr << "failed to create '" << *c_output << "'\n";
-                    return 1;
-                }
-                time_task _("C code generation", opts);
-                out << lean::ir::emit_c(env, *main_module_name).data();
-                out.close();
-            }
-            if (llvm_output) {
-                initialize_Lean_Compiler_IR_EmitLLVM(/*builtin*/ false,
-                        lean_io_mk_world());
-                time_task _("LLVM code generation", opts);
-                lean::consume_io_result(lean_ir_emit_llvm(
-                            env.to_obj_arg(), (*main_module_name).to_obj_arg(),
-                            lean::string_ref(*llvm_output).to_obj_arg(),
-                            lean_io_mk_world()));
-            }
-        }
-
-        display_cumulative_profiling_times(std::cerr);
-
-#if defined(__has_feature)
-#if __has_feature(address_sanitizer)
-        return opt_env ? 0 : 1;
-#endif
-#endif
-        // When not using the address/leak sanitizer, we interrupt execution without garbage collecting.
-        // This is useful when profiling improvements to Lean startup time.
-        exit(opt_env ? 0 : 1);
     } catch (lean::throwable & ex) {
         std::cerr << ex.what() << "\n";
     } catch (std::bad_alloc & ex) {