refactor(frontends/lean/pp): replace weak_ref with a strong reference, add new function (lean_formatter) for creating a Lean object formatter in the Lua API

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

src/frontends/lean/pp.cpp

@@ -153,7 +153,7 @@ class pp_fn {
     typedef scoped_map<expr, name, expr_hash_alloc, expr_eqp> aliases;
     typedef std::vector<std::pair<name, format>>              aliases_defs;
     typedef scoped_set<name, name_hash, name_eq>              local_names;
-    environment::weak_ref     m_env;
+    environment      m_env;
     // State
     aliases          m_aliases;
     aliases_defs     m_aliases_defs;
@@ -171,9 +171,7 @@ class pp_fn {
     bool             m_extra_lets;       //!< introduce extra let-expression to cope with sharing.
     unsigned         m_alias_min_weight; //!< minimal weight for creating an alias
 
-    environment env() const {
+    environment const & env() const { return m_env; }
-        return environment(m_env);
-    }
 
     // Create a scope for local definitions
     struct mk_scope {
@@ -1157,7 +1155,7 @@ class pp_fn {
     }
 
 public:
-    pp_fn(environment::weak_ref const & env, options const & opts):
+    pp_fn(environment const & env, options const & opts):
         m_env(env) {
         set_options(opts);
         m_num_steps   = 0;
@@ -1184,19 +1182,17 @@ public:
 };
 
 class pp_formatter_cell : public formatter_cell {
-    environment::weak_ref     m_env;
+    environment m_env;
 
-    environment env() const {
+    environment const & env() const { return m_env; }
-        return environment(m_env);
-    }
 
     format pp(expr const & e, options const & opts) {
-        pp_fn fn(m_env, opts);
+        pp_fn fn(env(), opts);
         return fn(e);
     }
 
     format pp(context const & c, expr const & e, bool include_e, options const & opts) {
-        pp_fn fn(m_env, opts);
+        pp_fn fn(env(), opts);
         unsigned indent = get_pp_indent(opts);
         format r;
         bool first = true;
@@ -1255,7 +1251,7 @@ class pp_formatter_cell : public formatter_cell {
             std::vector<bool> const * implicit_args = nullptr;
             if (has_implicit_arguments(env(), n))
                 implicit_args = &(get_implicit_arguments(env(), n));
-            pp_fn fn(m_env, opts);
+            pp_fn fn(env(), opts);
             format def = fn.pp_definition(v, t, implicit_args);
             return format{highlight_command(format(kwd)), space(), format(n), def};
         }
@@ -1271,7 +1267,7 @@ class pp_formatter_cell : public formatter_cell {
         name const & n = obj.get_name();
         format r = format{highlight_command(format(kwd)), space(), format(n)};
         if (has_implicit_arguments(env(), n)) {
-            pp_fn fn(m_env, opts);
+            pp_fn fn(env(), opts);
             r += fn.pp_pi_with_implicit_args(obj.get_type(), get_implicit_arguments(env(), n));
         } else {
             r += format{space(), colon(), space(), pp(obj.get_type(), opts)};
@@ -1313,10 +1309,7 @@ class pp_formatter_cell : public formatter_cell {
 
 public:
     pp_formatter_cell(environment const & env):
-        m_env(env.to_weak_ref()) {
+        m_env(env) {
-    }
-
-    virtual ~pp_formatter_cell() {
     }
 
     virtual format operator()(expr const & e, options const & opts) {

src/frontends/lean/register_module.cpp

@@ -12,6 +12,7 @@ Author: Leonardo de Moura
 #include "library/kernel_bindings.h"
 #include "frontends/lean/parser.h"
 #include "frontends/lean/frontend.h"
+#include "frontends/lean/pp.h"
 
 namespace lean {
 /** \see parse_lean_expr */
@@ -30,7 +31,7 @@ static int parse_lean_expr_core(lua_State * L, ro_environment const & env, io_st
 static int parse_lean_expr_core(lua_State * L, ro_environment const & env) {
     io_state * io = get_io_state(L);
     if (io == nullptr) {
-        io_state s(get_global_options(L), get_global_formatter(L));
+        io_state s(get_global_options(L), mk_pp_formatter(env));
         return parse_lean_expr_core(L, env, s);
     } else {
         return parse_lean_expr_core(L, env, *io);
@@ -66,7 +67,7 @@ static int parse_lean_expr(lua_State * L) {
             return parse_lean_expr_core(L, env);
         } else {
             options opts    = to_options(L, 3);
-            formatter fmt   = nargs == 3 ? get_global_formatter(L) : to_formatter(L, 4);
+            formatter fmt   = nargs == 3 ? mk_pp_formatter(env) : to_formatter(L, 4);
             io_state st(opts, fmt);
             return parse_lean_expr_core(L, env, st);
         }
@@ -87,7 +88,7 @@ static void parse_lean_cmds_core(lua_State * L, rw_environment & env, io_state &
 static void parse_lean_cmds_core(lua_State * L, rw_environment & env) {
     io_state * io = get_io_state(L);
     if (io == nullptr) {
-        io_state s(get_global_options(L), get_global_formatter(L));
+        io_state s(get_global_options(L), mk_pp_formatter(env));
         parse_lean_cmds_core(L, env, s);
         set_global_options(L, s.get_options());
     } else {
@@ -114,7 +115,7 @@ static int parse_lean_cmds(lua_State * L) {
             return 0;
         } else {
             options opts    = to_options(L, 3);
-            formatter fmt   = nargs == 3 ? get_global_formatter(L) : to_formatter(L, 4);
+            formatter fmt   = nargs == 3 ? mk_pp_formatter(env) : to_formatter(L, 4);
             io_state st(opts, fmt);
             parse_lean_cmds_core(L, env, st);
             push_options(L, st.get_options());
@@ -128,8 +129,13 @@ static int mk_environment(lua_State * L) {
     return push_environment(L, f.get_environment());
 }
 
+static int mk_lean_formatter(lua_State * L) {
+    return push_formatter(L, mk_pp_formatter(to_environment(L, 1)));
+}
+
 void open_frontend_lean(lua_State * L) {
     SET_GLOBAL_FUN(mk_environment,    "environment");
+    SET_GLOBAL_FUN(mk_lean_formatter, "lean_formatter");
     SET_GLOBAL_FUN(parse_lean_expr,   "parse_lean");
     SET_GLOBAL_FUN(parse_lean_cmds,   "parse_lean_cmds");
 }

src/library/kernel_bindings.cpp

@@ -869,23 +869,31 @@ static const struct luaL_Reg formatter_m[] = {
 static char g_formatter_key;
 static formatter g_simple_formatter = mk_simple_formatter();
 
-formatter get_global_formatter(lua_State * L) {
+optional<formatter> get_global_formatter_core(lua_State * L) {
     io_state * io = get_io_state(L);
     if (io != nullptr) {
-        return io->get_formatter();
+        return optional<formatter>(io->get_formatter());
     } else {
         lua_pushlightuserdata(L, static_cast<void *>(&g_formatter_key));
         lua_gettable(L, LUA_REGISTRYINDEX);
         if (is_formatter(L, -1)) {
             formatter r = to_formatter(L, -1);
             lua_pop(L, 1);
-            return r;
+            return optional<formatter>(r);
         } else {
             lua_pop(L, 1);
+            return optional<formatter>();
+        }
+    }
+}
+
+formatter get_global_formatter(lua_State * L) {
+    auto r = get_global_formatter_core(L);
+    if (r)
+        return *r;
+    else
         return g_simple_formatter;
 }
-    }
-}
 
 void set_global_formatter(lua_State * L, formatter const & fmt) {
     io_state * io = get_io_state(L);
@@ -898,10 +906,15 @@ void set_global_formatter(lua_State * L, formatter const & fmt) {
     }
 }
 
-int get_formatter(lua_State * L) {
+static int get_formatter(lua_State * L) {
     return push_formatter(L, get_global_formatter(L));
 }
 
+static int set_formatter(lua_State * L) {
+    set_global_formatter(L, to_formatter(L, 1));
+    return 0;
+}
+
 static void open_formatter(lua_State * L) {
     luaL_newmetatable(L, formatter_mt);
     lua_pushvalue(L, -1);
@@ -910,6 +923,7 @@ static void open_formatter(lua_State * L) {
 
     SET_GLOBAL_FUN(formatter_pred, "is_formatter");
     SET_GLOBAL_FUN(get_formatter,  "get_formatter");
+    SET_GLOBAL_FUN(set_formatter,  "set_formatter");
 }
 
 DECL_UDATA(environment)

src/library/kernel_bindings.h

@@ -30,6 +30,11 @@ int push_optional_object(lua_State * L, optional<object> const & o);
    1- Lean state object associated with \c L
    2- Lua Registry associated with \c L
 */
+optional<formatter> get_global_formatter_core(lua_State * L);
+/**
+   \brief Similar to \c get_global_formatter_core, but returns
+   the simple_formatter if a formatter can't be found.
+*/
 formatter get_global_formatter(lua_State * L);
 /**
    \brief Update the formatter object associated with the given Lua State.

tests/lua/front.lua

@@ -0,0 +1,52 @@
+local env = environment()
+print(get_options())
+parse_lean_cmds([[
+Variable f : Int -> Int -> Int
+Variable g : Bool -> Bool -> Bool
+Variables a b : Int
+Variables i j : Int
+Variables p q : Bool
+Notation 100 _ ++ _ : f
+Notation 100 _ ++ _ : g
+Set pp::colors true
+Set pp::width  300
+]], env)
+print(get_options())
+assert(get_options():get{"pp", "colors"})
+assert(get_options():get{"pp", "width"} == 300)
+parse_lean_cmds([[
+  Show i ++ j
+  Show f i j
+]], env)
+
+local env2 = environment()
+parse_lean_cmds([[
+Variable f : Int -> Int -> Int
+Variables a b : Int
+Show f a b
+Notation 100 _ -- _ : f
+]], env2)
+
+local f, a, b = Consts("f, a, b")
+assert(tostring(f(a, b)) == "f a b")
+set_formatter(lean_formatter(env))
+assert(tostring(f(a, b)) == "a ++ b")
+set_formatter(lean_formatter(env2))
+assert(tostring(f(a, b)) == "a -- b")
+
+local fmt = lean_formatter(env)
+-- We can parse commands with respect to environment env2,
+-- but using a formatter based on env.
+parse_lean_cmds([[
+Show f a b
+]], env2, options(), fmt)
+
+set_formatter(fmt)
+env  = nil
+env2 = nil
+fmt  = nil
+collectgarbage()
+-- The references to env, env2 and fmt were removed, but The global
+-- formatter (set with set_formatter) still has a reference to its
+-- environment.
+assert(tostring(f(a, b)) == "a ++ b")