feat(lua): allow lua scripts (embedded in Lean files) to access the environment

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

src/bindings/lua/environment.cpp

@@ -12,6 +12,7 @@ Author: Leonardo de Moura
 #include "bindings/lua/level.h"
 #include "bindings/lua/expr.h"
 #include "bindings/lua/context.h"
+#include "bindings/lua/environment.h"
 
 namespace lean {
 constexpr char const * environment_mt = "environment.mt";
@@ -29,14 +30,18 @@ static int environment_gc(lua_State * L) {
     return 0;
 }
 
-static int mk_environment(lua_State * L) {
+int push_environment(lua_State * L, environment const & env) {
     void * mem = lua_newuserdata(L, sizeof(environment));
-    new (mem) environment();
+    new (mem) environment(env);
     luaL_getmetatable(L, environment_mt);
     lua_setmetatable(L, -2);
     return 1;
 }
 
+static int mk_environment(lua_State * L) {
+    return push_environment(L, environment());
+}
+
 static int environment_add_uvar(lua_State * L) {
     int nargs = lua_gettop(L);
     if (nargs == 2)
@@ -120,6 +125,29 @@ static const struct luaL_Reg environment_m[] = {
     {0, 0}
 };
 
+static char g_set_environment_key;
+
+set_environment::set_environment(lua_State * L, environment & env) {
+    m_state = L;
+    lua_pushlightuserdata(m_state, (void *)&g_set_environment_key);
+    push_environment(m_state, env);
+    lua_settable(m_state, LUA_REGISTRYINDEX);
+}
+
+set_environment::~set_environment() {
+    lua_pushlightuserdata(m_state, (void *)&g_set_environment_key);
+    lua_pushnil(m_state);
+    lua_settable(m_state, LUA_REGISTRYINDEX);
+}
+
+int get_environment(lua_State * L) {
+    lua_pushlightuserdata(L, (void *)&g_set_environment_key);
+    lua_gettable(L, LUA_REGISTRYINDEX);
+    if (!is_environment(L, -1))
+        luaL_error(L, "Lua registry does not contain a Lean environment");
+    return push_environment(L, to_environment(L, -1));
+}
+
 void open_environment(lua_State * L) {
     luaL_newmetatable(L, environment_mt);
     lua_pushvalue(L, -1);
@@ -128,5 +156,7 @@ void open_environment(lua_State * L) {
 
     set_global_function<mk_environment>(L, "environment");
     set_global_function<environment_pred>(L, "is_environment");
+    set_global_function<get_environment>(L, "get_environment");
+    set_global_function<get_environment>(L, "env");
 }
 }

src/bindings/lua/environment.h

@@ -11,5 +11,15 @@ class environment;
 void open_environment(lua_State * L);
 bool is_environment(lua_State * L, int idx);
 environment & to_environment(lua_State * L, int idx);
-int push_environment(lua_State * L, environment const & o);
+int push_environment(lua_State * L, environment const & env);
+/**
+   \brief Auxiliary class for setting the Lua registry of a Lua state
+   with an environment object.
+*/
+class set_environment {
+    lua_State * m_state;
+public:
+    set_environment(lua_State * L, environment & env);
+    ~set_environment();
+};
 }

src/bindings/lua/leanlua_state.cpp

@@ -68,6 +68,11 @@ struct leanlua_state::imp {
         std::lock_guard<std::mutex> lock(m_mutex);
         ::lean::dostring(m_state, str);
     }
+
+    void dostring(char const * str, environment & env) {
+        set_environment set(m_state, env);
+        dostring(str);
+    }
 };
 
 leanlua_state::leanlua_state():
@@ -84,4 +89,8 @@ void leanlua_state::dofile(char const * fname) {
 void leanlua_state::dostring(char const * str) {
     m_ptr->dostring(str);
 }
+
+void leanlua_state::dostring(char const * str, environment & env) {
+    m_ptr->dostring(str, env);
+}
 }

src/bindings/lua/leanlua_state.h

@@ -9,6 +9,7 @@ Author: Leonardo de Moura
 #include "bindings/lua/lua_exception.h"
 
 namespace lean {
+class environment;
 /**
    \brief Wrapper for lua_State objects which contains all Lean bindings.
 */
@@ -29,5 +30,6 @@ public:
        This method throws an exception if an error occurs.
     */
     void dostring(char const * str);
+    void dostring(char const * str, environment & env);
 };
 }

src/frontends/lean/parser.cpp

@@ -1593,7 +1593,7 @@ class parser::imp {
         m_last_script_pos = mk_pair(m_scanner.get_script_block_line(), m_scanner.get_script_block_pos());
         if (!m_leanlua_state)
             throw exception("failed to execute Lua script, parser does not have a Lua interpreter");
-        m_leanlua_state->dostring(m_scanner.get_str_val().c_str());
+        m_leanlua_state->dostring(m_scanner.get_str_val().c_str(), m_frontend.get_environment());
         next();
     }
 

tests/lean/lua4.lean

@@ -0,0 +1,12 @@
+Variable x : Int
+
+{{
+-- Add a variable to the environment using Lua
+-- The type of the new variable is equal to the type
+-- of x
+typeofx = env():check_type(Const("x"))
+print("type of x is " .. tostring(typeofx))
+env():add_var("y", typeofx)
+}}
+
+Check x + y

tests/lean/lua4.lean.expected.out

@@ -0,0 +1,5 @@
+  Set: pp::colors
+  Set: pp::unicode
+  Assumed: x
+type of x is Int
+x + y : ℤ

tests/lean/lua5.lean

@@ -0,0 +1,13 @@
+Variable x : Int
+
+{{
+local N = 100
+-- Create N variables with the same type of x
+typeofx = env():check_type(Const("x"))
+for i = 1, N do
+    env():add_var("y_" .. i, typeofx)
+end
+}}
+
+Show Environment 101
+Check x + y_1 + y_2

tests/lean/lua5.lean.expected.out

@@ -0,0 +1,105 @@
+  Set: pp::colors
+  Set: pp::unicode
+  Assumed: x
+Variable x : ℤ
+Variable y_1 : ℤ
+Variable y_2 : ℤ
+Variable y_3 : ℤ
+Variable y_4 : ℤ
+Variable y_5 : ℤ
+Variable y_6 : ℤ
+Variable y_7 : ℤ
+Variable y_8 : ℤ
+Variable y_9 : ℤ
+Variable y_10 : ℤ
+Variable y_11 : ℤ
+Variable y_12 : ℤ
+Variable y_13 : ℤ
+Variable y_14 : ℤ
+Variable y_15 : ℤ
+Variable y_16 : ℤ
+Variable y_17 : ℤ
+Variable y_18 : ℤ
+Variable y_19 : ℤ
+Variable y_20 : ℤ
+Variable y_21 : ℤ
+Variable y_22 : ℤ
+Variable y_23 : ℤ
+Variable y_24 : ℤ
+Variable y_25 : ℤ
+Variable y_26 : ℤ
+Variable y_27 : ℤ
+Variable y_28 : ℤ
+Variable y_29 : ℤ
+Variable y_30 : ℤ
+Variable y_31 : ℤ
+Variable y_32 : ℤ
+Variable y_33 : ℤ
+Variable y_34 : ℤ
+Variable y_35 : ℤ
+Variable y_36 : ℤ
+Variable y_37 : ℤ
+Variable y_38 : ℤ
+Variable y_39 : ℤ
+Variable y_40 : ℤ
+Variable y_41 : ℤ
+Variable y_42 : ℤ
+Variable y_43 : ℤ
+Variable y_44 : ℤ
+Variable y_45 : ℤ
+Variable y_46 : ℤ
+Variable y_47 : ℤ
+Variable y_48 : ℤ
+Variable y_49 : ℤ
+Variable y_50 : ℤ
+Variable y_51 : ℤ
+Variable y_52 : ℤ
+Variable y_53 : ℤ
+Variable y_54 : ℤ
+Variable y_55 : ℤ
+Variable y_56 : ℤ
+Variable y_57 : ℤ
+Variable y_58 : ℤ
+Variable y_59 : ℤ
+Variable y_60 : ℤ
+Variable y_61 : ℤ
+Variable y_62 : ℤ
+Variable y_63 : ℤ
+Variable y_64 : ℤ
+Variable y_65 : ℤ
+Variable y_66 : ℤ
+Variable y_67 : ℤ
+Variable y_68 : ℤ
+Variable y_69 : ℤ
+Variable y_70 : ℤ
+Variable y_71 : ℤ
+Variable y_72 : ℤ
+Variable y_73 : ℤ
+Variable y_74 : ℤ
+Variable y_75 : ℤ
+Variable y_76 : ℤ
+Variable y_77 : ℤ
+Variable y_78 : ℤ
+Variable y_79 : ℤ
+Variable y_80 : ℤ
+Variable y_81 : ℤ
+Variable y_82 : ℤ
+Variable y_83 : ℤ
+Variable y_84 : ℤ
+Variable y_85 : ℤ
+Variable y_86 : ℤ
+Variable y_87 : ℤ
+Variable y_88 : ℤ
+Variable y_89 : ℤ
+Variable y_90 : ℤ
+Variable y_91 : ℤ
+Variable y_92 : ℤ
+Variable y_93 : ℤ
+Variable y_94 : ℤ
+Variable y_95 : ℤ
+Variable y_96 : ℤ
+Variable y_97 : ℤ
+Variable y_98 : ℤ
+Variable y_99 : ℤ
+Variable y_100 : ℤ
+x + y_1 + y_2 : ℤ