From c46608ced50a660c2caa96e09aba1d49113165fb Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Sat, 11 May 2019 16:28:11 -0700
Subject: [PATCH] feat(library/init/io): implement `unsafeIO` in Lean

Motivation: avoid closure allocation.
---
 library/init/io.lean                 | 11 ++--
 src/runtime/io.cpp                   |  6 ---
 src/stage0/init/control/estate.cpp   | 80 ++++++++++++++++++++++++++++
 src/stage0/init/io.cpp               | 30 +++++++++--
 src/stage0/init/lean/environment.cpp | 23 ++++----
 tests/playground/unsafe.lean         |  5 ++
 6 files changed, 130 insertions(+), 25 deletions(-)
 create mode 100644 tests/playground/unsafe.lean

diff --git a/library/init/io.lean b/library/init/io.lean
index a967251e77..815c271f31 100644
--- a/library/init/io.lean
+++ b/library/init/io.lean
@@ -7,8 +7,10 @@ prelude
 import init.control.estate init.data.string.basic
 
 /-- Like https://hackage.haskell.org/package/ghc-Prim-0.5.2.0/docs/GHC-Prim.html#t:RealWorld.
-    Makes sure we never reorder `IO` operations. -/
-constant IO.RealWorld : Type := Unit
+    Makes sure we never reorder `IO` operations.
+
+    TODO: mark opaque -/
+def IO.RealWorld : Type := Unit
 
 /- TODO(Leo): mark it as an opaque definition. Reason: prevent
    functions defined in other modules from accessing `IO.RealWorld`.
@@ -35,8 +37,9 @@ s
 
 abbrev IO : Type → Type := EIO IO.error
 
-@[extern "lean_io_unsafe"]
-unsafe constant unsafeIO {α : Type} (fn : IO α) : Option α := default _
+@[inline]
+unsafe def unsafeIO {α : Type} (fn : IO α) : Option α :=
+fn.run' ()
 
 @[extern 4 "lean_io_timeit"]
 constant timeit {α : Type} (msg : @& String) (fn : IO α) : IO α := default _
diff --git a/src/runtime/io.cpp b/src/runtime/io.cpp
index 327adb2b21..1497447dce 100644
--- a/src/runtime/io.cpp
+++ b/src/runtime/io.cpp
@@ -102,12 +102,6 @@ extern "C" obj_res lean_io_prim_handle_get_line(b_obj_arg /* h */, obj_arg /* w
     lean_unreachable();
 }
 
-/* constant unsafe_io {α : Type} (fn : io α) : option α */
-extern "C" obj_res lean_io_unsafe(obj_arg, obj_arg fn) {
-    object * r     = io_mk_world();
-    return option_of_io_result(apply_1(fn, r));
-}
-
 /* timeit {α : Type} (msg : @& string) (fn : io α) : io α */
 extern "C" obj_res lean_io_timeit(obj_arg, b_obj_arg msg, obj_arg fn, obj_arg r) {
     auto start = std::chrono::steady_clock::now();
diff --git a/src/stage0/init/control/estate.cpp b/src/stage0/init/control/estate.cpp
index c0c70c39b8..58b21ad7bd 100644
--- a/src/stage0/init/control/estate.cpp
+++ b/src/stage0/init/control/estate.cpp
@@ -31,11 +31,13 @@ obj* l_EState_map___rarg(obj*, obj*, obj*);
 obj* l_EState_Monad___boxed(obj*, obj*);
 obj* l_EState_orelse(obj*, obj*, obj*);
 obj* l_EState_Result_toString___main___boxed(obj*, obj*, obj*);
+obj* l_EState_run_x_27(obj*, obj*, obj*);
 obj* l_EState_Inhabited___boxed(obj*, obj*, obj*);
 obj* l_EState_orelse_x_27(obj*, obj*, obj*);
 obj* l_EState_unreachableError(obj*, obj*, obj*, obj*, obj*, obj*, obj*);
 obj* l_EState_orelse_x_27___rarg___boxed(obj*, obj*, obj*, obj*);
 obj* l_EState_Result_toString___main(obj*, obj*, obj*);
+obj* l_EState_run___rarg(obj*, obj*);
 obj* l_EState_throw___boxed(obj*, obj*, obj*);
 obj* l_EState_Result_toString___main___rarg(obj*, obj*, obj*);
 obj* l_EState_HasToString___boxed(obj*, obj*, obj*);
@@ -52,10 +54,12 @@ obj* l_EState_Result_toString___boxed(obj*, obj*, obj*);
 namespace lean {
 obj* string_append(obj*, obj*);
 }
+obj* l_EState_run_x_27___boxed(obj*, obj*, obj*);
 obj* l_EState_Result_repr___main___rarg(obj*, obj*, obj*);
 extern obj* l_Option_HasRepr___rarg___closed__3;
 obj* l_EState_HasOrelse___closed__1;
 obj* l_EState_pure(obj*, obj*, obj*);
+obj* l_EState_run___boxed(obj*, obj*, obj*);
 extern obj* l_Except_toString___main___rarg___closed__2;
 obj* l_EState_HasOrelse(obj*, obj*, obj*);
 obj* l_EState_Monad___lambda__1(obj*, obj*, obj*, obj*, obj*);
@@ -76,6 +80,7 @@ obj* l_EState_bind___rarg(obj*, obj*, obj*);
 obj* l_EState_Monad___lambda__3(obj*, obj*, obj*, obj*, obj*);
 extern obj* l_Except_toString___main___rarg___closed__1;
 obj* l_EState_Result_repr___main___boxed(obj*, obj*, obj*);
+obj* l_EState_run_x_27___rarg(obj*, obj*);
 obj* l_EState_Result_repr___rarg(obj*, obj*, obj*);
 obj* l_EState_get___boxed(obj*, obj*);
 obj* l_EState_orelse___rarg(obj*, obj*, obj*);
@@ -97,6 +102,7 @@ obj* l_EState_MonadState___closed__1;
 obj* l_EState_map___boxed(obj*, obj*, obj*, obj*);
 obj* l_EState_resultOk_mk___boxed(obj*, obj*, obj*);
 obj* l_EState_Result_repr(obj*, obj*, obj*);
+obj* l_EState_run(obj*, obj*, obj*);
 obj* l_EState_set___rarg(obj*, obj*);
 obj* l_EState_Inhabited___rarg(obj*, obj*);
 obj* l_EState_MonadState___boxed(obj*, obj*);
@@ -1581,6 +1587,80 @@ lean::dec(x_1);
 return x_2;
 }
 }
+obj* l_EState_run___rarg(obj* x_0, obj* x_1) {
+_start:
+{
+obj* x_2; obj* x_3; obj* x_4; 
+x_2 = lean::box(0);
+x_3 = lean::alloc_cnstr(0, 2, 0);
+lean::cnstr_set(x_3, 0, x_2);
+lean::cnstr_set(x_3, 1, x_1);
+x_4 = lean::apply_1(x_0, x_3);
+return x_4;
+}
+}
+obj* l_EState_run(obj* x_0, obj* x_1, obj* x_2) {
+_start:
+{
+obj* x_3; 
+x_3 = lean::alloc_closure(reinterpret_cast<void*>(l_EState_run___rarg), 2, 0);
+return x_3;
+}
+}
+obj* l_EState_run___boxed(obj* x_0, obj* x_1, obj* x_2) {
+_start:
+{
+obj* x_3; 
+x_3 = l_EState_run(x_0, x_1, x_2);
+lean::dec(x_0);
+lean::dec(x_1);
+lean::dec(x_2);
+return x_3;
+}
+}
+obj* l_EState_run_x_27___rarg(obj* x_0, obj* x_1) {
+_start:
+{
+obj* x_2; 
+x_2 = l_EState_run___rarg(x_0, x_1);
+if (lean::obj_tag(x_2) == 0)
+{
+obj* x_3; obj* x_6; 
+x_3 = lean::cnstr_get(x_2, 0);
+lean::inc(x_3);
+lean::dec(x_2);
+x_6 = lean::alloc_cnstr(1, 1, 0);
+lean::cnstr_set(x_6, 0, x_3);
+return x_6;
+}
+else
+{
+obj* x_8; 
+lean::dec(x_2);
+x_8 = lean::box(0);
+return x_8;
+}
+}
+}
+obj* l_EState_run_x_27(obj* x_0, obj* x_1, obj* x_2) {
+_start:
+{
+obj* x_3; 
+x_3 = lean::alloc_closure(reinterpret_cast<void*>(l_EState_run_x_27___rarg), 2, 0);
+return x_3;
+}
+}
+obj* l_EState_run_x_27___boxed(obj* x_0, obj* x_1, obj* x_2) {
+_start:
+{
+obj* x_3; 
+x_3 = l_EState_run_x_27(x_0, x_1, x_2);
+lean::dec(x_0);
+lean::dec(x_1);
+lean::dec(x_2);
+return x_3;
+}
+}
 obj* initialize_init_control_state(obj*);
 obj* initialize_init_control_except(obj*);
 static bool _G_initialized = false;
diff --git a/src/stage0/init/io.cpp b/src/stage0/init/io.cpp
index 6ed9f675ce..1d232ee237 100644
--- a/src/stage0/init/io.cpp
+++ b/src/stage0/init/io.cpp
@@ -44,7 +44,7 @@ obj* l_IO_print___at_HasRepr_HasEval___spec__2___boxed(obj*, obj*);
 obj* l_IO_Prim_iterate___at_IO_Fs_handle_readToEnd___spec__3___boxed(obj*, obj*, obj*);
 obj* l_IO_Ref_reset___rarg(obj*, obj*, obj*);
 obj* l_EIO_Inhabited(obj*, obj*);
-extern "C" obj* lean_io_unsafe(obj*, obj*);
+obj* l_unsafeIO(obj*);
 obj* l_IO_Prim_iterate___at_IO_Fs_handle_readToEnd___spec__3(obj*, obj*, obj*);
 obj* l_EIO_Monad(obj*);
 obj* l_IO_HasEval___boxed(obj*);
@@ -87,7 +87,7 @@ obj* l_IO_mkRef___rarg(obj*, obj*, obj*);
 obj* l_IO_Fs_handle_mk___rarg___boxed(obj*, obj*, obj*, obj*);
 obj* l_IO_Prim_iterate(obj*, obj*);
 obj* l_IO_Fs_readFile___rarg___lambda__2(obj*, obj*, obj*, obj*);
-obj* l_unsafeIO___boxed(obj*, obj*);
+obj* l_unsafeIO___boxed(obj*);
 obj* l_String_HasToString___boxed(obj*);
 obj* l_IO_Prim_liftIO(obj*, obj*);
 obj* l_IO_Inhabited(obj*);
@@ -121,6 +121,7 @@ obj* l_IO_Fs_handle_getLine___at_IO_Fs_handle_readToEnd___spec__2___boxed(obj*,
 obj* l_IO_Fs_handle_getLine___boxed(obj*, obj*);
 extern "C" obj* lean_io_prim_handle_flush(obj*, obj*);
 obj* l___private_init_io_1__putStr(obj*, obj*);
+obj* l_EState_run_x_27___rarg(obj*, obj*);
 obj* l_IO_print(obj*, obj*);
 obj* l_IO_Fs_handle_getLine___rarg(obj*, obj*);
 obj* l_IO_Prim_handle_getLine___boxed(obj*, obj*);
@@ -145,6 +146,7 @@ obj* l_EState_MonadExcept(obj*, obj*);
 obj* l_timeit___boxed(obj*, obj*, obj*, obj*);
 obj* l_IO_Fs_handle_isEof___boxed(obj*, obj*);
 obj* l_IO_Prim_iterate___main(obj*, obj*);
+obj* l_unsafeIO___rarg(obj*);
 obj* l_IO_Prim_iterate___main___at_IO_Fs_handle_readToEnd___spec__4(obj*, obj*, obj*);
 obj* l_IO_Fs_handle_flush___rarg(obj*, obj*);
 extern "C" obj* lean_io_timeit(obj*, obj*, obj*, obj*);
@@ -277,14 +279,32 @@ lean::dec(x_0);
 return x_1;
 }
 }
-obj* l_unsafeIO___boxed(obj* x_0, obj* x_1) {
+obj* l_unsafeIO___rarg(obj* x_0) {
 _start:
 {
-obj* x_2; 
-x_2 = lean_io_unsafe(x_0, x_1);
+obj* x_1; obj* x_2; 
+x_1 = lean::box(0);
+x_2 = l_EState_run_x_27___rarg(x_0, x_1);
 return x_2;
 }
 }
+obj* l_unsafeIO(obj* x_0) {
+_start:
+{
+obj* x_1; 
+x_1 = lean::alloc_closure(reinterpret_cast<void*>(l_unsafeIO___rarg), 1, 0);
+return x_1;
+}
+}
+obj* l_unsafeIO___boxed(obj* x_0) {
+_start:
+{
+obj* x_1; 
+x_1 = l_unsafeIO(x_0);
+lean::dec(x_0);
+return x_1;
+}
+}
 obj* l_timeit___boxed(obj* x_0, obj* x_1, obj* x_2, obj* x_3) {
 _start:
 {
diff --git a/src/stage0/init/lean/environment.cpp b/src/stage0/init/lean/environment.cpp
index 92dbb5590e..73be137d01 100644
--- a/src/stage0/init/lean/environment.cpp
+++ b/src/stage0/init/lean/environment.cpp
@@ -46,7 +46,6 @@ obj* l_Lean_EnvExtension_modifyStateUnsafe(obj*);
 obj* l___private_init_lean_environment_4__getTrustLevel___boxed(obj*);
 obj* l_Lean_SMap_find___main___at_Lean_Environment_find___spec__1___boxed(obj*, obj*);
 uint8 l_AssocList_contains___main___at_Lean_Environment_add___spec__5(obj*, obj*);
-extern "C" obj* lean_io_unsafe(obj*, obj*);
 obj* l_Lean_SMap_insert___main___at_Lean_Environment_add___spec__1___closed__1;
 obj* l_Lean_SMap_insert___main___at_Lean_Environment_add___spec__1(obj*, obj*, obj*);
 obj* l_HashMapImp_moveEntries___main___at_Lean_Environment_add___spec__7(obj*, obj*, obj*);
@@ -154,6 +153,7 @@ obj* l_Lean_getCPPExtensionState___lambda__1(obj*, obj*, obj*, obj*);
 obj* l_Lean_EnvExtension_setStateUnsafe___boxed(obj*);
 obj* l_RBNode_insert___at_Lean_Environment_add___spec__2(obj*, obj*, obj*);
 obj* l_Lean_PersistentEnvExtension_addEntry___rarg(obj*, obj*, obj*);
+obj* l_EState_run_x_27___rarg(obj*, obj*);
 namespace lean {
 obj* environment_switch_core(obj*);
 }
@@ -4077,15 +4077,16 @@ namespace lean {
 obj* register_extension_core(obj* x_0) {
 _start:
 {
-obj* x_1; obj* x_2; obj* x_3; obj* x_4; 
+obj* x_1; obj* x_2; obj* x_3; obj* x_4; obj* x_5; 
 x_1 = lean::alloc_closure(reinterpret_cast<void*>(l_Lean_registerEnvExtensionUnsafe___rarg), 2, 1);
 lean::closure_set(x_1, 0, x_0);
 x_2 = l_Lean_registerCPPExtension___closed__1;
 x_3 = lean::alloc_closure(reinterpret_cast<void*>(l_EState_bind___rarg), 3, 2);
 lean::closure_set(x_3, 0, x_1);
 lean::closure_set(x_3, 1, x_2);
-x_4 = lean_io_unsafe(lean::box(0), x_3);
-return x_4;
+x_4 = lean::box(0);
+x_5 = l_EState_run_x_27___rarg(x_3, x_4);
+return x_5;
 }
 }
 }
@@ -4131,7 +4132,7 @@ namespace lean {
 obj* set_extension_core(obj* x_0, obj* x_1, obj* x_2) {
 _start:
 {
-obj* x_3; obj* x_4; obj* x_5; obj* x_6; 
+obj* x_3; obj* x_4; obj* x_5; obj* x_6; obj* x_7; 
 x_3 = lean::alloc_closure(reinterpret_cast<void*>(l_Lean_setCPPExtensionState___lambda__1___boxed), 5, 3);
 lean::closure_set(x_3, 0, x_1);
 lean::closure_set(x_3, 1, x_0);
@@ -4140,8 +4141,9 @@ x_4 = l_Lean_setCPPExtensionState___closed__1;
 x_5 = lean::alloc_closure(reinterpret_cast<void*>(l_EState_bind___rarg), 3, 2);
 lean::closure_set(x_5, 0, x_4);
 lean::closure_set(x_5, 1, x_3);
-x_6 = lean_io_unsafe(lean::box(0), x_5);
-return x_6;
+x_6 = lean::box(0);
+x_7 = l_EState_run_x_27___rarg(x_5, x_6);
+return x_7;
 }
 }
 }
@@ -4185,7 +4187,7 @@ namespace lean {
 obj* get_extension_core(obj* x_0, obj* x_1) {
 _start:
 {
-obj* x_2; obj* x_3; obj* x_4; obj* x_5; 
+obj* x_2; obj* x_3; obj* x_4; obj* x_5; obj* x_6; 
 x_2 = lean::alloc_closure(reinterpret_cast<void*>(l_Lean_getCPPExtensionState___lambda__1___boxed), 4, 2);
 lean::closure_set(x_2, 0, x_1);
 lean::closure_set(x_2, 1, x_0);
@@ -4193,8 +4195,9 @@ x_3 = l_Lean_setCPPExtensionState___closed__1;
 x_4 = lean::alloc_closure(reinterpret_cast<void*>(l_EState_bind___rarg), 3, 2);
 lean::closure_set(x_4, 0, x_3);
 lean::closure_set(x_4, 1, x_2);
-x_5 = lean_io_unsafe(lean::box(0), x_4);
-return x_5;
+x_5 = lean::box(0);
+x_6 = l_EState_run_x_27___rarg(x_4, x_5);
+return x_6;
 }
 }
 }
diff --git a/tests/playground/unsafe.lean b/tests/playground/unsafe.lean
new file mode 100644
index 0000000000..9e1373eacf
--- /dev/null
+++ b/tests/playground/unsafe.lean
@@ -0,0 +1,5 @@
+unsafe def test (x : Nat) : Option Nat :=
+unsafeIO (IO.println x *> pure (x+1))
+
+unsafe def main (xs : List String) : IO Unit :=
+IO.println $ test xs.head.toNat