feat(bindings/lua): expose io_state object in the Lua API

Signed-off-by: Leonardo de Moura <leonardo@microsoft.com>

src/bindings/lua/cex_builder.cpp

@@ -60,5 +60,4 @@ void open_cex_builder(lua_State * L) {
     SET_GLOBAL_FUN(cex_builder_pred, "is_cex_builder");
     SET_GLOBAL_FUN(mk_cex_builder, "cex_builder");
 }
-
 }

src/bindings/lua/io_state.cpp

@@ -6,13 +6,74 @@ Author: Leonardo de Moura
 */
 #include <lua.hpp>
 #include "library/io_state.h"
+#include "bindings/lua/util.h"
 #include "bindings/lua/io_state.h"
+#include "bindings/lua/options.h"
+#include "bindings/lua/formatter.h"
 
 namespace lean {
+DECL_UDATA(io_state)
+
+int mk_io_state(lua_State * L) {
+    int nargs = lua_gettop(L);
+    if (nargs == 0)
+        return push_io_state(L, io_state());
+    else if (nargs == 1)
+        return push_io_state(L, io_state(to_io_state(L, 1)));
+    else
+        return push_io_state(L, io_state(to_options(L, 1), to_formatter(L, 2)));
+}
+
+int io_state_get_options(lua_State * L) {
+    return push_options(L, to_io_state(L, 1).get_options());
+}
+
+int io_state_get_formatter(lua_State * L) {
+    return push_formatter(L, to_io_state(L, 1).get_formatter());
+}
+
+int io_state_set_options(lua_State * L) {
+    to_io_state(L, 1).set_options(to_options(L, 2));
+    return 0;
+}
+
+int print(lua_State * L, io_state & ios, int start, bool reg);
+
+int io_state_print_regular(lua_State * L) {
+    return print(L, to_io_state(L, 1), 2, true);
+}
+
+int io_state_print_diagnostic(lua_State * L) {
+    return print(L, to_io_state(L, 1), 2, false);
+}
+
+static const struct luaL_Reg io_state_m[] = {
+    {"__gc",             io_state_gc}, // never throws
+    {"get_options",      safe_function<io_state_get_options>},
+    {"set_options",      safe_function<io_state_set_options>},
+    {"get_formatter",    safe_function<io_state_get_formatter>},
+    {"print_diagnostic", safe_function<io_state_print_diagnostic>},
+    {"print_regular",    safe_function<io_state_print_regular>},
+    {"print",            safe_function<io_state_print_regular>},
+    {"diagnostic",       safe_function<io_state_print_diagnostic>},
+    {0, 0}
+};
+
+void open_io_state(lua_State * L) {
+    luaL_newmetatable(L, io_state_mt);
+    lua_pushvalue(L, -1);
+    lua_setfield(L, -2, "__index");
+    setfuncs(L, io_state_m, 0);
+
+    SET_GLOBAL_FUN(io_state_pred, "is_io_state");
+    SET_GLOBAL_FUN(mk_io_state, "io_state");
+}
+
 static char g_set_state_key;
 
 set_io_state::set_io_state(lua_State * L, io_state & st) {
     m_state = L;
+    m_prev  = get_io_state(L);
     lua_pushlightuserdata(m_state, static_cast<void *>(&g_set_state_key));
     lua_pushlightuserdata(m_state, &st);
     lua_settable(m_state, LUA_REGISTRYINDEX);
@@ -20,7 +81,7 @@ set_io_state::set_io_state(lua_State * L, io_state & st) {
 
 set_io_state::~set_io_state() {
     lua_pushlightuserdata(m_state, static_cast<void *>(&g_set_state_key));
-    lua_pushnil(m_state);
+    lua_pushlightuserdata(m_state, m_prev);
     lua_settable(m_state, LUA_REGISTRYINDEX);
 }
 

src/bindings/lua/io_state.h

@@ -9,12 +9,15 @@ Author: Leonardo de Moura
 #include "library/io_state.h"
 
 namespace lean {
+UDATA_DEFS(io_state)
+void open_io_state(lua_State * L);
 /**
    \brief Auxiliary class for temporarily setting the Lua registry of a Lua state
    with a Lean io_state object.
 */
 class set_io_state {
     lua_State * m_state;
+    io_state *  m_prev;
 public:
     set_io_state(lua_State * L, io_state & st);
     ~set_io_state();

src/bindings/lua/leanlua_state.cpp

@@ -196,6 +196,7 @@ struct leanlua_state::imp {
         open_frontend_lean(m_state);
         open_thread(m_state);
         open_interrupt(m_state);
+        open_io_state(m_state);
         dostring(g_leanlua_extra);
     }
 
@@ -253,14 +254,25 @@ void leanlua_state::dostring(char const * str, environment & env, io_state & st)
 
 static std::mutex g_print_mutex;
 
+static void print(io_state * ios, bool reg, char const * msg) {
+    if (ios) {
+        if (reg)
+            regular(*ios) << msg;
+        else
+            diagnostic(*ios) << msg;
+    } else {
+        std::cout << msg;
+    }
+}
+
 /** \brief Thread safe version of print function */
-static int print(lua_State * L) {
+static int print(lua_State * L, int start, bool reg = false) {
-    io_state * io = get_io_state(L);
+    std::lock_guard<std::mutex> lock(g_print_mutex);
+    io_state * ios = get_io_state(L);
     int n = lua_gettop(L);
     int i;
     lua_getglobal(L, "tostring");
-    std::lock_guard<std::mutex> lock(g_print_mutex);
+    for (i = start; i <= n; i++) {
-    for (i = 1; i <= n; i++) {
         char const * s;
         size_t l;
         lua_pushvalue(L, -1);
@@ -269,25 +281,25 @@ static int print(lua_State * L) {
         s = lua_tolstring(L, -1, &l);
         if (s == NULL)
             throw exception("'to_string' must return a string to 'print'");
-        if (i > 1) {
+        if (i > start) {
-            if (io)
+            print(ios, reg, "\t");
-                regular(*io) << "\t";
-            else
-                std::cout << "\t";
         }
-        if (io)
+        print(ios, reg, s);
-            regular(*io) << s;
-        else
-            std::cout << s;
         lua_pop(L, 1);
     }
-    if (io)
+    print(ios, reg, "\n");
-        regular(*io) << endl;
-    else
-        std::cout << std::endl;
     return 0;
 }
 
+int print(lua_State * L, io_state & ios, int start, bool reg) {
+    set_io_state set(L, ios);
+    return print(L, start, reg);
+}
+
+static int print(lua_State * L) {
+    return print(L, 1, true);
+}
+
 /** \brief Redefine some functions from the Lua library */
 static void open_patch(lua_State * L) {
     SET_GLOBAL_FUN(print, "print");

src/bindings/lua/options.cpp

@@ -31,7 +31,7 @@ static int mk_options(lua_State * L) {
         } else {
             option_declaration const & d = it->second;
             switch (d.kind()) {
-            case BoolOption:      r = r.update(k, lua_toboolean(L, i+1)); break;
+            case BoolOption:      r = r.update(k, static_cast<bool>(lua_toboolean(L, i+1))); break;
             case IntOption:       r = r.update(k, static_cast<int>(lua_tointeger(L, i+1))); break;
             case UnsignedOption:  r = r.update(k, static_cast<unsigned>(lua_tointeger(L, i+1))); break;
             case DoubleOption:    r = r.update(k, static_cast<double>(lua_tonumber(L, i+1))); break;

src/bindings/lua/proof_state.cpp

@@ -180,11 +180,11 @@ static const struct luaL_Reg proof_state_m[] = {
     {"get_menv",             safe_function<proof_state_get_menv>},
     {"get_proof_builder",    safe_function<proof_state_get_proof_builder>},
     {"get_cex_builder",      safe_function<proof_state_get_cex_builder>},
-    {"precision",        safe_function<proof_state_get_precision>},
+    {"precision",            safe_function<proof_state_get_precision>},
-    {"goals",            safe_function<proof_state_get_goals>},
+    {"goals",                safe_function<proof_state_get_goals>},
-    {"menv",             safe_function<proof_state_get_menv>},
+    {"menv",                 safe_function<proof_state_get_menv>},
-    {"proof_builder",    safe_function<proof_state_get_proof_builder>},
+    {"proof_builder",        safe_function<proof_state_get_proof_builder>},
-    {"cex_builder",      safe_function<proof_state_get_cex_builder>},
+    {"cex_builder",          safe_function<proof_state_get_cex_builder>},
     {"is_proof_final_state", safe_function<proof_state_is_proof_final_state>},
     {"is_cex_final_state",   safe_function<proof_state_is_cex_final_state>},
     {0, 0}

tests/lean/lua16.lean

@@ -0,0 +1,8 @@
+
+Variables a b : Int
+
+(**
+  local ios = io_state()
+  ios:print(parse_lean("a + b"))
+  print(parse_lean("a + b"))
+**)

tests/lean/lua16.lean.expected.out

@@ -0,0 +1,6 @@
+  Set: pp::colors
+  Set: pp::unicode
+  Assumed: a
+  Assumed: b
+Int::add a b
+a + b

tests/lua/io_state1.lua

@@ -0,0 +1,9 @@
+local ios = io_state()
+assert(is_io_state(ios))
+ios:print(Const("a"), Const("b"))
+print(ios:get_options())
+assert(is_formatter(ios:get_formatter()))
+ios:set_options(options({"pp", "unicode"}, false))
+print(ios:get_options())
+ios:print_diagnostic("warning")
+ios:print_diagnostic(ios:get_options())