fix(library/kernel_bindings): set_environment and set_io_state objects

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

src/frontends/lean/parser.h

@@ -87,8 +87,8 @@ class parser {
     template<typename F>
     typename std::result_of<F(lua_State * L)>::type using_script(F && f) {
         script_state S = get_thread_script_state();
-        set_io_state    set1(S.get_state(), m_ios);
-        set_environment set2(S.get_state(), m_env);
+        set_io_state    set1(S, m_ios);
+        set_environment set2(S, m_env);
         return f(S.get_state());
     }
 

src/library/kernel_bindings.cpp

@@ -1286,34 +1286,34 @@ static const struct luaL_Reg environment_m[] = {
 };
 
 static char g_set_environment_key;
-
-void set_global_environment(lua_State * L, environment const & env) {
+static environment * get_global_environment_ptr(lua_State * L) {
     lua_pushlightuserdata(L, static_cast<void *>(&g_set_environment_key));
-    push_environment(L, env);
+    lua_gettable(L, LUA_REGISTRYINDEX);
+    if (!lua_islightuserdata(L, -1))
+        return nullptr;
+    environment * r = static_cast<environment*>(const_cast<void*>(lua_topointer(L, -1)));
+    lua_pop(L, 1);
+    return r;
+}
+static void set_global_environment_ptr(lua_State * L, environment * env) {
+    lua_pushlightuserdata(L, static_cast<void *>(&g_set_environment_key));
+    lua_pushlightuserdata(L, static_cast<void *>(env));
     lua_settable(L, LUA_REGISTRYINDEX);
 }
+static environment get_global_environment(lua_State * L) {
+    environment * env = get_global_environment_ptr(L);
+    if (env == nullptr)
+        throw exception("Lua state does not have an environment object");
+    return *env;
+}
 
-set_environment::set_environment(lua_State * L, environment & env):
-    m_env(env) {
-    m_state = L;
-    set_global_environment(L, env);
+set_environment::set_environment(lua_State * L, environment & env):m_state(L) {
+    m_old_env = get_global_environment_ptr(L);
+    set_global_environment_ptr(L, &env);
 }
 
 set_environment::~set_environment() {
-    m_env = get_global_environment(m_state);
-    lua_pushlightuserdata(m_state, static_cast<void *>(&g_set_environment_key));
-    lua_pushnil(m_state);
-    lua_settable(m_state, LUA_REGISTRYINDEX);
-}
-
-static environment get_global_environment(lua_State * L) {
-    lua_pushlightuserdata(L, static_cast<void *>(&g_set_environment_key));
-    lua_gettable(L, LUA_REGISTRYINDEX);
-    if (!is_environment(L, -1))
-        return environment(); // return empty environment
-    environment r = to_environment(L, -1);
-    lua_pop(L, 1);
-    return r;
+    set_global_environment_ptr(m_state, m_old_env);
 }
 
 int get_environment(lua_State * L) {
@@ -1321,7 +1321,10 @@ int get_environment(lua_State * L) {
 }
 
 int set_environment(lua_State * L) {
-    set_global_environment(L, to_environment(L, 1));
+    environment * env = get_global_environment_ptr(L);
+    if (env == nullptr)
+        throw exception("Lua state does not have an environment object");
+    *env = to_environment(L, 1);
     return 0;
 }
 

src/library/kernel_bindings.h

@@ -5,7 +5,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Author: Leonardo de Moura
 */
 #pragma once
-#include "util/lua.h"
+#include "util/script_state.h"
 #include "kernel/environment.h"
 
 namespace lean {
@@ -52,10 +52,11 @@ void set_global_environment(lua_State * L, environment const & env);
    with an environment object.
 */
 class set_environment {
-    environment & m_env;
     lua_State *   m_state;
+    environment * m_old_env;
 public:
     set_environment(lua_State * L, environment & env);
+    set_environment(script_state & S, environment & env):set_environment(S.get_state(), env) {}
     ~set_environment();
 };
 
@@ -70,7 +71,8 @@ class set_io_state {
     io_state *  m_prev;
     options     m_prev_options;
 public:
-    set_io_state(lua_State * L, io_state & st);
+    set_io_state(lua_State * L, io_state & ios);
+    set_io_state(script_state & S, io_state & ios):set_io_state(S.get_state(), ios) {}
     ~set_io_state();
 };
 /**

src/shell/lean.cpp

@@ -39,6 +39,8 @@ using lean::io_state_stream;
 using lean::regular;
 using lean::mk_environment;
 using lean::mk_hott_environment;
+using lean::set_environment;
+using lean::set_io_state;
 
 enum class input_kind { Unspecified, Lean, Lua };
 
@@ -184,11 +186,13 @@ int main(int argc, char ** argv) {
         }
     }
 
-    script_state S = lean::get_thread_script_state();
     environment env = mode == lean_mode::Standard ? mk_environment(trust_lvl) : mk_hott_environment(trust_lvl);
     io_state ios(lean::mk_simple_formatter());
     if (quiet)
         ios.set_option("verbose", false);
+    script_state S = lean::get_thread_script_state();
+    set_environment set1(S, env);
+    set_io_state    set2(S, ios);
 
     try {
         bool ok = true;