From fd2a5dd45e9ff7cc6944fd407320ffa00085bfee Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Fri, 10 May 2019 11:26:49 -0700
Subject: [PATCH] feat(library/init/io): add `IO.initializing`

---
 library/init/io.lean              | 8 ++++++++
 src/init/init.cpp                 | 1 +
 src/library/compiler/emit_cpp.cpp | 1 +
 src/runtime/io.cpp                | 7 +++++++
 src/runtime/object.h              | 1 +
 5 files changed, 18 insertions(+)

diff --git a/library/init/io.lean b/library/init/io.lean
index 714df896a2..5f843f9e28 100644
--- a/library/init/io.lean
+++ b/library/init/io.lean
@@ -44,6 +44,14 @@ constant timeit {α : Type} (msg : @& String) (fn : IO α) : IO α := default _
 @[extern 4 "lean_io_allocprof"]
 constant allocprof {α : Type} (msg : @& String) (fn : IO α) : IO α := default _
 
+/- Programs can execute IO actions during initialization that occurs before
+   the `main` function is executed. The attribute `[init <action>]` specifies
+   which IO action is executed to set the value of an opaque constant.
+
+   The action `initializing` returns `true` iff it is invoked during initialization. -/
+@[extern 1 "lean_io_initializing"]
+constant IO.initializing : IO Bool := default _
+
 abbrev monadIO (m : Type → Type) := HasMonadLiftT IO m
 
 def ioOfExcept {ε α : Type} [HasToString ε] (e : Except ε α) : IO α :=
diff --git a/src/init/init.cpp b/src/init/init.cpp
index 893244ee98..c945814e95 100644
--- a/src/init/init.cpp
+++ b/src/init/init.cpp
@@ -38,6 +38,7 @@ void initialize() {
     initialize_frontend_lean_module();
     object * w = initialize_init_default(io_mk_world());
     w = initialize_init_lean_default(w);
+    lean::io_mark_end_initialization();
     if (io_result_is_error(w)) {
         io_result_show_error(w);
         dec(w);
diff --git a/src/library/compiler/emit_cpp.cpp b/src/library/compiler/emit_cpp.cpp
index 5604f95ba1..e484363aab 100644
--- a/src/library/compiler/emit_cpp.cpp
+++ b/src/library/compiler/emit_cpp.cpp
@@ -1154,6 +1154,7 @@ static void emit_main_fn(std::ostream & out, environment const & env, module_nam
         out << "lean::initialize_runtime_module();\n";
     out << "obj * w = lean::io_mk_world();\n";
     out << "w = initialize_" << mangle(m, false) << "(w);\n";
+    out << "lean::io_mark_end_initialization();\n";
     out << "if (io_result_is_ok(w)) {\n";
     out << "lean::scoped_task_manager tmanager(lean::hardware_concurrency());\n";
     if (arity == 2) {
diff --git a/src/runtime/io.cpp b/src/runtime/io.cpp
index f62a14ee31..327adb2b21 100644
--- a/src/runtime/io.cpp
+++ b/src/runtime/io.cpp
@@ -55,6 +55,13 @@ static obj_res option_of_io_result(obj_arg r) {
     }
 }
 
+static bool g_initializing = true;
+void io_mark_end_initialization() { g_initializing = false; }
+
+extern "C" obj_res lean_io_initializing(obj_arg r) {
+    return set_io_result(r, box(g_initializing));
+}
+
 extern "C" obj_res lean_io_prim_put_str(b_obj_arg s, obj_arg r) {
     std::cout << string_to_std(s); // TODO(Leo): use out handle
     return set_io_result(r, box(0));
diff --git a/src/runtime/object.h b/src/runtime/object.h
index 46317b8cd5..579c063e40 100644
--- a/src/runtime/object.h
+++ b/src/runtime/object.h
@@ -1507,6 +1507,7 @@ inline bool io_result_is_error(b_obj_arg r) { return cnstr_tag(r) == 1; }
 inline b_obj_res io_result_get_value(b_obj_arg r) { lean_assert(io_result_is_ok(r)); return cnstr_get(r, 0); }
 inline b_obj_arg io_result_get_error(b_obj_arg r) { lean_assert(io_result_is_error(r)); return cnstr_get(r, 0); }
 void io_result_show_error(b_obj_arg r);
+void io_mark_end_initialization();
 
 // =======================================
 // IO ref primitives