From 2a8cd373ca323b01b37105b529cf04cfe19092bc Mon Sep 17 00:00:00 2001 From: Mac Malone Date: Sat, 14 Jun 2025 21:11:58 -0400 Subject: [PATCH] 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 | 1 + src/Lean/Elab/Frontend.lean | 12 ++-- src/Lean/Shell.lean | 84 +++++++++++++++++++++++++ src/library/compiler/ir_interpreter.cpp | 29 ++++----- src/library/compiler/ir_interpreter.h | 2 +- src/library/compiler/llvm.cpp | 14 +++++ src/library/time_task.cpp | 6 ++ src/runtime/platform.cpp | 14 ++++- src/util/shell.cpp | 78 ++++++++--------------- 9 files changed, 163 insertions(+), 77 deletions(-) create mode 100644 src/Lean/Shell.lean diff --git a/src/Lean.lean b/src/Lean.lean index 118f074592..39d1cef22d 100644 --- a/src/Lean.lean +++ b/src/Lean.lean @@ -43,3 +43,4 @@ import Lean.Namespace import Lean.EnvExtension import Lean.ErrorExplanation import Lean.DefEqAttrib +import Lean.Shell diff --git a/src/Lean/Elab/Frontend.lean b/src/Lean/Elab/Frontend.lean index 6d97c67563..b3b676924f 100644 --- a/src/Lean/Elab/Frontend.lean +++ b/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 diff --git a/src/Lean/Shell.lean b/src/Lean/Shell.lean new file mode 100644 index 0000000000..d6c609b5f4 --- /dev/null +++ b/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 diff --git a/src/library/compiler/ir_interpreter.cpp b/src/library/compiler/ir_interpreter.cpp index f1cf5f4937..de26f2c4bc 100644 --- a/src/library/compiler/ir_interpreter.cpp +++ b/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 const & args) { decl d = get_decl("main"); array_ref const & params = decl_params(d); - buffer args; + buffer 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(env, opts, "main", [&](interpreter & interp) { return interp.run_main(argv, argc); }); +uint32 run_main(elab_environment const & env, options const & opts, list_ref const & args) { + return interpreter::with_interpreter(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, args)); + return io_result_mk_ok(box(ret)); } extern "C" LEAN_EXPORT object * lean_eval_const(object * env, object * opts, object * c) { diff --git a/src/library/compiler/ir_interpreter.h b/src/library/compiler/ir_interpreter.h index c3a0022656..9da08495c1 100644 --- a/src/library/compiler/ir_interpreter.h +++ b/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 const & args); } void initialize_ir_interpreter(); void finalize_ir_interpreter(); diff --git a/src/library/compiler/llvm.cpp b/src/library/compiler/llvm.cpp index 892502a546..8b53c4a02d 100644 --- a/src/library/compiler/llvm.cpp +++ b/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 diff --git a/src/library/time_task.cpp b/src/library/time_task.cpp index 51796825f2..5f6c15b051 100644 --- a/src/library/time_task.cpp +++ b/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; diff --git a/src/runtime/platform.cpp b/src/runtime/platform.cpp index 0546abb19e..07859a3040 100644 --- a/src/runtime/platform.cpp +++ b/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 +} } diff --git a/src/util/shell.cpp b/src/util/shell.cpp index 308ecc0430..153159f173 100644 --- a/src/util/shell.cpp +++ b/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 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 const & olean_file_name, optional const & ilean_file_name, + optional const & c_file_name, + optional const & bc_file_name, uint8_t json_output, array_ref const & error_kinds, bool print_stats, - optional const & setup_file_name + optional const & setup_file_name, + bool run ) { - return get_io_result>(lean_run_frontend( + list_ref args; + while (argc > 0) { + argc--; + args = list_ref(string_ref(argv[argc]), args); + } + return get_io_scalar_result(lean_shell_main( + args.steal(), mk_string(input), opts.to_obj_arg(), mk_string(file_name), @@ -365,11 +373,13 @@ option_ref 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 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) {