refactor: port more of shell.cpp to Lean (#10086)

This PR ports more of the post-initialization C++ shell code to Lean.

All that remains is the initialization of the profiler and task manager.
As initialization tasks rather than main shell code, they were left in
C++ (where the rest of the initialization code currently is).

The `max_memory` and `timeout` Lean options used by the the `--memory`
and `--timeout` command-line options are now properly registered. The
server defaults for max memory and max heartbeats (timeout) were removed
as they were not actually used (because the `server` option that was
checked was neither set nor exists).

This PR also makes better use of the module system in `Shell.lean` and
fixes a minor bug in a previous port where the file name check was
dependent on building the `.ilean` rather than the `.c` file (as was
originally the case).

Fixes #9879.
This commit is contained in:
Mac Malone 2025-08-26 16:02:42 -04:00 committed by GitHub
parent c83674bdff
commit db3fb47109
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 123 additions and 89 deletions

View file

@ -1063,7 +1063,6 @@ where
severity? := DiagnosticSeverity.error
message := err.toString }]
@[export lean_server_worker_main]
def workerMain (opts : Options) : IO UInt32 := do
let i ← IO.getStdin
let o ← IO.getStdout

View file

@ -1663,7 +1663,6 @@ def initAndRunWatchdog (args : List String) (i o e : FS.Stream) : IO Unit := do
: ServerContext
}
@[export lean_server_watchdog_main]
def watchdogMain (args : List String) : IO UInt32 := do
let i ← IO.getStdin
let o ← IO.getStdout

View file

@ -6,11 +6,11 @@ Authors: Leonardo de Moura, Mac Malone
module
prelude
public import Lean.Elab.Frontend
public import Lean.Elab.ParseImportsFast
public import Lean.Compiler.IR.EmitC
public section
import Lean.Elab.Frontend
import Lean.Elab.ParseImportsFast
import Lean.Server.Watchdog
import Lean.Server.FileWorker
import Lean.Compiler.IR.EmitC
/- Lean companion to `shell.cpp` -/
@ -27,48 +27,70 @@ This is used instead of the standard `String` functions because Lean is expected
abort on files with invalid UTF-8.
-/
@[extern "lean_decode_lossy_utf8"]
private opaque decodeLossyUTF8 (a : @& ByteArray) : String
opaque decodeLossyUTF8 (a : @& ByteArray) : String
/- 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
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
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
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
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
opaque Internal.hasAddressSanitizer (_ : Unit) : Bool
/-- Whether Lean was built with multithread support. -/
@[extern "lean_internal_is_multi_thread"]
private opaque Internal.isMultiThread (_ : Unit) : Bool
opaque Internal.isMultiThread (_ : Unit) : Bool
/-- Whether Lean was built in debug mode. -/
@[extern "lean_internal_is_debug"]
private opaque Internal.isDebug (_ : Unit) : Bool
opaque Internal.isDebug (_ : Unit) : Bool
/-- The mode Lean was built in. -/
/-- Returns the mode Lean was built in. -/
@[extern "lean_internal_get_build_type"]
private opaque Internal.getBuildType (_ : Unit) : String
opaque Internal.getBuildType (_ : Unit) : String
/--
Returns the default max memory (in megabytes) Lean was built with
(i.e., `LEAN_DEFAULT_MAX_MEMORY`).
-/
@[extern "lean_internal_get_default_max_memory"]
opaque Internal.getDefaultMaxMemory (_ : Unit) : Nat
/-- Sets Lean's internal maximum memory (in bytes) for the C runtime. -/
@[extern "lean_internal_set_max_memory"]
opaque Internal.setMaxMemory (max : USize) : BaseIO Unit
/--
Returns the default max heartbeats (in thousands) Lean was built with
(i.e., `LEAN_DEFAULT_MAX_HEARTBEAT`).
-/
@[extern "lean_internal_get_default_max_heartbeat"]
opaque Internal.getDefaultMaxHeartbeat (_ : Unit) : Nat
/-- Sets Lean's internal maximum heartbeats for the C runtime. -/
@[extern "lean_internal_set_max_heartbeat"]
opaque Internal.setMaxHeartbeat (max : USize) : BaseIO Unit
/-- Lean equivalent of `LEAN_VERSION_STRING` / `g_short_version_string` -/
private def shortVersionString : String :=
def shortVersionString : String :=
if version.specialDesc ≠ "" then
s!"{versionStringCore}-{version.specialDesc}"
else if !version.isRelease then
@ -77,7 +99,7 @@ private def shortVersionString : String :=
versionStringCore
/-- The full Lean version header (i.e, what is printed by `lean --version`). -/
private def versionHeader : String := Id.run do
def versionHeader : String := Id.run do
let mut ver := shortVersionString
if Platform.target ≠ "" then
ver := s!"{ver}, {Platform.target}"
@ -87,12 +109,12 @@ private def versionHeader : String := Id.run do
/-- Print the Lean version header. -/
@[export lean_display_header]
private def displayHeader : IO Unit := do
def displayHeader : IO Unit := do
IO.println versionHeader
/-- Print the Lean CLI help message. -/
@[export lean_display_help]
private def displayHelp (useStderr : Bool) : IO Unit := do
def displayHelp (useStderr : Bool) : IO Unit := do
let out ← if useStderr then IO.getStderr else IO.getStdout
out.putStrLn versionHeader
out.putStrLn "Miscellaneous"
@ -137,9 +159,24 @@ private def displayHelp (useStderr : Bool) : IO Unit := do
out.putStrLn " --debug=tag enable assertions with the given tag"
out.putStrLn " -D name=value set a configuration option (see set_option command)"
inductive ShellComponent
| frontend
| watchdog
| worker
private builtin_initialize maxMemory : Lean.Option Nat ←
Lean.Option.register `max_memory {defValue := Internal.getDefaultMaxMemory ()}
private builtin_initialize timeout : Lean.Option Nat ←
Lean.Option.register `timeout {defValue := Internal.getDefaultMaxHeartbeat ()}
@[export lean_shell_main]
private def shellMain
def shellMain
(args : List String)
(forwardedArgs : List String)
(component : ShellComponent := .frontend)
(printPrefix : Bool := false)
(printLibDir : Bool := false)
(useStdin : Bool := false)
(onlyDeps : Bool := false)
(onlySrcDeps : Bool := false)
@ -157,6 +194,25 @@ private def shellMain
(printStats : Bool := false)
(run : Bool := false)
: IO UInt32 := do
if printPrefix then
IO.println (← getBuildDir)
return 0
if printLibDir then
IO.println (← getLibDir (← getBuildDir))
return 0
let maxMemory := maxMemory.get opts
if maxMemory != 0 then
Internal.setMaxMemory (maxMemory.toUSize * 1024 * 1024)
let timeout := timeout.get opts
if timeout != 0 then
Internal.setMaxHeartbeat (timeout.toUSize * 1000)
match component with
| .frontend =>
pure ()
| .watchdog =>
return ← Server.Watchdog.watchdogMain forwardedArgs
| .worker =>
return ← Server.FileWorker.workerMain opts
if onlyDeps && depsJson then
let fns ←
if useStdin then
@ -214,7 +270,7 @@ private def shellMain
pure setup.name
else if let some fileName := fileName? then
try moduleNameOfFileName fileName rootDir? catch e =>
if oleanFileName?.isNone && ileanFileName?.isNone then
if oleanFileName?.isNone && cFileName?.isNone then
pure `_stdin
else
throw e

View file

@ -75,11 +75,9 @@ end SearchPath
builtin_initialize searchPathRef : IO.Ref SearchPath ← IO.mkRef {}
@[export lean_get_prefix]
def getBuildDir : IO FilePath := do
return (← IO.appDir).parent |>.get!
@[export lean_get_libdir]
def getLibDir (leanSysroot : FilePath) : IO FilePath := do
let mut buildDir := leanSysroot
-- use stage1 stdlib with stage0 executable (which should never be distributed outside of the build directory)

View file

@ -16,12 +16,25 @@ namespace lean {
LEAN_THREAD_VALUE(size_t, g_max_heartbeat, 0);
LEAN_THREAD_VALUE(size_t, g_heartbeat, 0);
extern "C" LEAN_EXPORT obj_res lean_internal_get_default_max_heartbeat(obj_arg) {
#ifdef LEAN_DEFAULT_MAX_HEARTBEAT
return lean_box(LEAN_DEFAULT_MAX_HEARTBEAT);
#else
return lean_box(0);
#endif
}
void inc_heartbeat() { g_heartbeat++; }
void reset_heartbeat() { g_heartbeat = 0; }
void set_max_heartbeat(size_t max) { g_max_heartbeat = max; }
extern "C" LEAN_EXPORT obj_res lean_internal_set_max_heartbeat(usize max) {
set_max_heartbeat(max);
return lean_io_result_mk_ok(lean_box(0));
}
size_t get_max_heartbeat() { return g_max_heartbeat; }
void set_max_heartbeat_thousands(unsigned max) { g_max_heartbeat = static_cast<size_t>(max) * 1000; }

View file

@ -114,10 +114,23 @@ namespace lean {
static size_t g_max_memory = 0;
LEAN_THREAD_VALUE(size_t, g_counter, 0);
extern "C" LEAN_EXPORT lean_obj_res lean_internal_get_default_max_memory(lean_obj_arg) {
#ifdef LEAN_DEFAULT_MAX_MEMORY
return lean_box(LEAN_DEFAULT_MAX_MEMORY);
#else
return lean_box(0);
#endif
}
void set_max_memory(size_t max) {
g_max_memory = max;
}
extern "C" LEAN_EXPORT lean_obj_res lean_internal_set_max_memory(size_t max) {
set_max_memory(max);
return lean_io_result_mk_ok(lean_box(0));
}
void set_max_memory_megabyte(unsigned max) {
size_t m = max;
m *= 1024 * 1024;

View file

@ -157,19 +157,6 @@ int getopt_long(int argc, char *in_argv[], const char *optstring, const option *
using namespace lean; // NOLINT
#ifndef LEAN_SERVER_DEFAULT_MAX_MEMORY
#define LEAN_SERVER_DEFAULT_MAX_MEMORY 1024
#endif
#ifndef LEAN_DEFAULT_MAX_MEMORY
#define LEAN_DEFAULT_MAX_MEMORY 0
#endif
#ifndef LEAN_DEFAULT_MAX_HEARTBEAT
#define LEAN_DEFAULT_MAX_HEARTBEAT 0
#endif
#ifndef LEAN_SERVER_DEFAULT_MAX_HEARTBEAT
#define LEAN_SERVER_DEFAULT_MAX_HEARTBEAT 100000
#endif
extern "C" obj_res lean_display_header(obj_arg);
static void display_header() {
consume_io_result(lean_display_header(io_mk_world()));
@ -286,6 +273,10 @@ options set_config_option(options const & opts, char const * in) {
namespace lean {
extern "C" obj_res lean_shell_main(
obj_arg args,
obj_arg forwarded_args,
uint8 component,
uint8 print_prefix,
uint8 print_libdir,
uint8 use_stdin,
uint8 only_deps,
uint8 only_src_deps,
@ -306,6 +297,10 @@ extern "C" obj_res lean_shell_main(
);
uint32 run_shell_main(
int argc, char* argv[],
buffer<string_ref> forwarded_args,
int run_server,
bool print_prefix,
bool print_libdir,
bool use_stdin,
bool only_deps,
bool only_src_deps,
@ -330,6 +325,10 @@ uint32 run_shell_main(
}
return get_io_scalar_result<uint32>(lean_shell_main(
args.steal(),
to_list_ref(forwarded_args).to_obj_arg(),
run_server,
print_prefix,
print_libdir,
use_stdin,
only_deps, only_src_deps, deps_json,
opts.to_obj_arg(),
@ -348,19 +347,6 @@ uint32 run_shell_main(
));
}
/* def workerMain : Options → IO UInt32 */
extern "C" object * lean_server_worker_main(object * opts, object * w);
uint32_t run_server_worker(options const & opts) {
return get_io_scalar_result<uint32_t>(lean_server_worker_main(opts.to_obj_arg(), io_mk_world()));
}
/* def watchdogMain (args : List String) : IO Uint32 */
extern "C" object* lean_server_watchdog_main(object* args, object* w);
uint32_t run_server_watchdog(buffer<string_ref> const & args) {
list_ref<string_ref> arglist = to_list_ref(args);
return get_io_scalar_result<uint32_t>(lean_server_watchdog_main(arglist.to_obj_arg(), io_mk_world()));
}
extern "C" object* lean_init_search_path(object* w);
void init_search_path() {
get_io_scalar_result<unsigned>(lean_init_search_path(io_mk_world()));
@ -372,9 +358,6 @@ void environment_free_regions(elab_environment && env) {
}
}
extern "C" object * lean_get_prefix(object * w);
extern "C" object * lean_get_libdir(object * sysroot, object * w);
void check_optarg(char const * option_name) {
if (!optarg) {
std::cerr << "error: argument missing for option '-" << option_name << "'" << std::endl;
@ -601,29 +584,6 @@ extern "C" LEAN_EXPORT int lean_main(int argc, char ** argv) {
lean::io_mark_end_initialization();
if (print_prefix) {
std::cout << get_io_result<string_ref>(lean_get_prefix(io_mk_world())).data() << std::endl;
return 0;
}
if (print_libdir) {
string_ref prefix = get_io_result<string_ref>(lean_get_prefix(io_mk_world()));
std::cout << get_io_result<string_ref>(lean_get_libdir(prefix.to_obj_arg(), io_mk_world())).data() << std::endl;
return 0;
}
if (auto max_memory = opts.get_unsigned(get_max_memory_opt_name(),
opts.get_bool("server") ? LEAN_SERVER_DEFAULT_MAX_MEMORY
: LEAN_DEFAULT_MAX_MEMORY)) {
set_max_memory_megabyte(max_memory);
}
if (auto timeout = opts.get_unsigned(get_timeout_opt_name(),
opts.get_bool("server") ? LEAN_SERVER_DEFAULT_MAX_HEARTBEAT
: LEAN_DEFAULT_MAX_HEARTBEAT)) {
set_max_heartbeat_thousands(timeout);
}
if (get_profiler(opts)) {
g_lean_report_task_get_blocked_time = report_task_get_blocked_time;
report_profiling_time("initialization", init_time);
@ -632,18 +592,14 @@ extern "C" LEAN_EXPORT int lean_main(int argc, char ** argv) {
scoped_task_manager scope_task_man(num_threads);
try {
if (run_server == 1)
return run_server_watchdog(forwarded_args);
else if (run_server == 2)
return run_server_worker(opts);
else
return run_shell_main(
argc - optind, argv + optind,
use_stdin, only_deps, only_src_deps, deps_json,
opts, trust_lvl, root_dir, setup_fn,
olean_fn, ilean_fn, c_output, llvm_output,
json_output, error_kinds, stats, run
);
return run_shell_main(
argc - optind, argv + optind,
forwarded_args, run_server, print_prefix, print_libdir,
use_stdin, only_deps, only_src_deps, deps_json,
opts, trust_lvl, root_dir, setup_fn,
olean_fn, ilean_fn, c_output, llvm_output,
json_output, error_kinds, stats, run
);
} catch (lean::throwable & ex) {
std::cerr << ex.what() << "\n";
} catch (std::bad_alloc & ex) {