feat(frontends/lean): add Alias command

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

src/frontends/lean/frontend.cpp

@@ -54,6 +54,22 @@ static void read_mark_implicit(environment const & env, io_state const &, deseri
 }
 static object_cell::register_deserializer_fn mark_implicit_ds("MarkImplicit", read_mark_implicit);
 
+class alias_command : public neutral_object_cell {
+    name m_name;
+    expr m_expr;
+public:
+    alias_command(name const & n, expr const & e):m_name(n), m_expr(e) {}
+    virtual ~alias_command() {}
+    virtual char const * keyword() const { return "Alias"; }
+    virtual void write(serializer & s) const { s << "Alias" << m_name << m_expr; }
+};
+static void read_alias(environment const & env, io_state const &, deserializer & d) {
+    name n = read_name(d);
+    expr e = read_expr(d);
+    add_alias(env, n, e);
+}
+static object_cell::register_deserializer_fn add_alias_ds("Alias", read_alias);
+
 static std::vector<bool> g_empty_vector;
 /**
    \brief Environment extension object for the Lean default frontend.
@@ -68,6 +84,7 @@ struct lean_extension : public environment_extension {
     typedef expr_pair_struct_map<expr>           coercion_map;
     typedef expr_struct_map<list<expr_pair>>     expr_to_coercions;
     typedef expr_struct_set coercion_set;
+    typedef expr_struct_map<list<name>>          inv_aliases;
 
     operator_table        m_nud; // nud table for Pratt's parser
     operator_table        m_led; // led table for Pratt's parser
@@ -80,6 +97,8 @@ struct lean_extension : public environment_extension {
     coercion_set          m_coercion_set; // Set of coercions
     expr_to_coercions     m_type_coercions; // mapping type -> list (to-type, function)
     name_set              m_explicit_names; // set of explicit version of constants with implicit parameters
+    name_map<expr>        m_aliases;
+    inv_aliases           m_inv_aliases;  // inverse map for m_aliases
 
     lean_extension() {
     }
@@ -498,6 +517,40 @@ struct lean_extension : public environment_extension {
         lean_extension const * parent = get_parent();
         return parent && parent->is_coercion(f);
     }
+
+    optional<expr> get_alias(name const & n) const {
+        auto it = m_aliases.find(n);
+        if (it != m_aliases.end())
+            return some_expr(it->second);
+        lean_extension const * parent = get_parent();
+        if (parent)
+            return parent->get_alias(n);
+        else
+            return none_expr();
+    }
+
+    optional<list<name>> get_aliased(expr const & e) const {
+        auto it = m_inv_aliases.find(e);
+        if (it != m_inv_aliases.end())
+            return optional<list<name>>(it->second);
+        lean_extension const * parent = get_parent();
+        if (parent)
+            return parent->get_aliased(e);
+        else
+            return optional<list<name>>();
+    }
+
+    void add_alias(name const & n, expr const & e, environment const & env) {
+        if (get_alias(n))
+            throw exception(sstream() << "alias '" << n << "' was already defined");
+        m_aliases[n]     = e;
+        auto l = get_aliased(e);
+        if (l)
+            m_inv_aliases[e] = list<name>(n, *l);
+        else
+            m_inv_aliases[e] = list<name>(n);
+        env->add_neutral_object(new alias_command(n, e));
+    }
 };
 
 struct lean_extension_initializer {
@@ -637,4 +690,13 @@ list<expr_pair> get_coercions(ro_environment const & env, expr const & from_type
 bool is_coercion(ro_environment const & env, expr const & f) {
     return to_ext(env).is_coercion(f);
 }
+optional<expr> get_alias(ro_environment const & env, name const & n) {
+    return to_ext(env).get_alias(n);
+}
+optional<list<name>> get_aliased(ro_environment const & env, expr const & e) {
+    return to_ext(env).get_aliased(e);
+}
+void add_alias(environment const & env, name const & n, expr const & e) {
+    to_ext(env).add_alias(n, e, env);
+}
 }

src/frontends/lean/frontend.h

@@ -159,4 +159,13 @@ optional<expr> get_coercion(ro_environment const & env, expr const & from_type,
 */
 list<expr_pair> get_coercions(ro_environment const & env, expr const & from_type);
 /*@}*/
+
+/**
+   @name Aliases
+*/
+/*@{*/
+optional<expr> get_alias(ro_environment const & env, name const & n);
+optional<list<name>> get_aliased(ro_environment const & env, expr const & e);
+void add_alias(environment const & env, name const & n, expr const & e);
+/*@}*/
 }

src/frontends/lean/parser.cpp

@@ -79,6 +79,7 @@ bool     get_parser_show_errors(options const & opts)  { return opts.get_bool(g_
 
 // ==========================================
 // Builtin commands
+static name g_alias_kwd("Alias");
 static name g_definition_kwd("Definition");
 static name g_variable_kwd("Variable");
 static name g_variables_kwd("Variables");
@@ -115,7 +116,7 @@ static list<name> g_tactic_cmds = { g_apply, g_done, g_back, g_abort, g_assumpti
 static list<name> g_command_keywords = {g_definition_kwd, g_variable_kwd, g_variables_kwd, g_theorem_kwd, g_axiom_kwd, g_universe_kwd, g_eval_kwd,
                                         g_show_kwd, g_check_kwd, g_infix_kwd, g_infixl_kwd, g_infixr_kwd, g_notation_kwd, g_echo_kwd,
                                         g_set_option_kwd, g_set_opaque_kwd, g_env_kwd, g_options_kwd, g_import_kwd, g_help_kwd, g_coercion_kwd,
-                                        g_exit_kwd, g_push_kwd, g_pop_kwd, g_scope_kwd, g_end_scope_kwd};
+                                        g_exit_kwd, g_push_kwd, g_pop_kwd, g_scope_kwd, g_end_scope_kwd, g_alias_kwd};
 // ==========================================
 
 // ==========================================
@@ -1162,6 +1163,8 @@ class parser::imp {
             return save(mk_var(m_num_local_decls - it->second - 1), p);
         } else if (m_expr_macros && m_expr_macros->find(id) != m_expr_macros->end()) {
             return parse_expr_macro(id, p);
+        } else if (auto alias = get_alias(m_env, id)) {
+            return save(*alias, p);
         } else {
             operator_info op = find_nud(m_env, id);
             if (op) {
@@ -2468,6 +2471,7 @@ class parser::imp {
             }
         } else {
             regular(m_io_state) << "Available commands:" << endl
+                                << "  Alias [id] : [expr]    define an alias for the given expression" << endl
                                 << "  Axiom [id] : [type]    assert/postulate a new axiom" << endl
                                 << "  Check [expr]           type check the given expression" << endl
                                 << "  Definition [id] : [type] := [expr]   define a new element" << endl
@@ -2552,6 +2556,14 @@ class parser::imp {
         m_env->add_uvar(id, lvl);
     }
 
+    void parse_alias() {
+        next();
+        name id   = check_identifier_next("invalid alias declaration, identifier expected");
+        check_colon_next("invalid alias declaration, ':' expected");
+        expr e    = parse_expr();
+        add_alias(m_env, id, e);
+    }
+
     /** \brief Parse a Lean command. */
     bool parse_command() {
         m_elaborator.clear();
@@ -2606,6 +2618,8 @@ class parser::imp {
             parse_end_scope();
         } else if (cmd_id == g_universe_kwd) {
             parse_universe();
+        } else if (cmd_id == g_alias_kwd) {
+            parse_alias();
         } else if (m_cmd_macros && m_cmd_macros->find(cmd_id) != m_cmd_macros->end()) {
             parse_cmd_macro(cmd_id, m_last_cmd_pos);
         } else {

src/frontends/lean/pp.cpp

@@ -88,7 +88,7 @@ static name g_pp_max_steps       {"lean", "pp", "max_steps"};
 static name g_pp_implicit         {"lean", "pp", "implicit"};
 static name g_pp_notation         {"lean", "pp", "notation"};
 static name g_pp_extra_lets       {"lean", "pp", "extra_lets"};
-static name g_pp_alias_min_weight{"lean", "pp", "alias_min_weight"};
+static name g_pp_alias_min_weight {"lean", "pp", "alias_min_weight"};
 static name g_pp_coercion         {"lean", "pp", "coercion"};
 static name g_pp_def_value        {"lean", "pp", "definition_value"};
 
@@ -111,7 +111,7 @@ unsigned get_pp_alias_min_weight(options const & opts) { return opts.get_unsigne
 bool     get_pp_def_value(options const & opts)        { return opts.get_bool(g_pp_def_value, LEAN_DEFAULT_PP_DEFINITION_VALUE); }
 
 // =======================================
-// Prefixes for naming aliases (auxiliary local decls)
+// Prefixes for naming local aliases (auxiliary local decls)
 static name g_kappa("\u03BA");
 static name g_pho("\u03C1");
 static name g_nu("\u03BD");
@@ -168,13 +168,13 @@ bool supported_by_pp(object const & obj) {
 
 /** \brief Functional object for pretty printing expressions */
 class pp_fn {
-    typedef scoped_map<expr, name, expr_hash_alloc, expr_eqp> aliases;
+    typedef scoped_map<expr, name, expr_hash_alloc, expr_eqp> local_aliases;
-    typedef std::vector<std::pair<name, format>>              aliases_defs;
+    typedef std::vector<std::pair<name, format>>              local_aliases_defs;
     typedef scoped_set<name, name_hash, name_eq>              local_names;
     ro_environment   m_env;
     // State
-    aliases          m_aliases;
+    local_aliases      m_local_aliases;
-    aliases_defs     m_aliases_defs;
+    local_aliases_defs m_local_aliases_defs;
     local_names        m_local_names;
     unsigned           m_num_steps;
     name               m_aux;
@@ -193,13 +193,13 @@ class pp_fn {
     struct mk_scope {
         pp_fn &  m_fn;
         unsigned m_old_size;
-        mk_scope(pp_fn & fn):m_fn(fn), m_old_size(fn.m_aliases_defs.size()) {
+        mk_scope(pp_fn & fn):m_fn(fn), m_old_size(fn.m_local_aliases_defs.size()) {
-            m_fn.m_aliases.push();
+            m_fn.m_local_aliases.push();
         }
         ~mk_scope() {
-            lean_assert(m_old_size <= m_fn.m_aliases_defs.size());
+            lean_assert(m_old_size <= m_fn.m_local_aliases_defs.size());
-            m_fn.m_aliases.pop();
+            m_fn.m_local_aliases.pop();
-            m_fn.m_aliases_defs.resize(m_old_size);
+            m_fn.m_local_aliases_defs.resize(m_old_size);
         }
     };
 
@@ -215,6 +215,10 @@ class pp_fn {
        \brief Return true iff \c e is an atomic operation.
     */
     bool is_atomic(expr const & e) {
+        if (auto aliased_list = get_aliased(m_env, e)) {
+            if (m_unicode || std::any_of(aliased_list->begin(), aliased_list->end(), [](name const & a) { return a.is_safe_ascii(); }))
+                return true;
+        }
         switch (e.kind()) {
         case expr_kind::Var: case expr_kind::Constant: case expr_kind::Type:
             return true;
@@ -774,7 +778,7 @@ class pp_fn {
         } else {
             mk_scope s(*this);
             result r = pp(e, depth + 1, true);
-            if (m_aliases_defs.size() == s.m_old_size) {
+            if (m_local_aliases_defs.size() == s.m_old_size) {
                 if (prec <= get_operator_precedence(e))
                     return r;
                 else
@@ -783,9 +787,9 @@ class pp_fn {
                 format r_format   = g_let_fmt;
                 unsigned r_weight = 2;
                 unsigned begin    = s.m_old_size;
-                unsigned end      = m_aliases_defs.size();
+                unsigned end      = m_local_aliases_defs.size();
                 for (unsigned i = begin; i < end; i++) {
-                    auto b = m_aliases_defs[i];
+                    auto b = m_local_aliases_defs[i];
                     name const & n = b.first;
                     format beg = i == begin ? space() : line();
                     format sep = i < end - 1 ? comma() : format();
@@ -1105,9 +1109,19 @@ class pp_fn {
             return pp_ellipsis();
         } else {
             m_num_steps++;
+            if (auto aliased_list = get_aliased(m_env, e)) {
+                if (m_unicode) {
+                    return mk_result(format(head(*aliased_list)), 1);
+                } else {
+                    for (auto n : *aliased_list) {
+                        if (n.is_safe_ascii())
+                            return mk_result(format(n), 1);
+                    }
+                }
+            }
             if (m_extra_lets && is_shared(e)) {
-                auto it = m_aliases.find(e);
+                auto it = m_local_aliases.find(e);
-                if (it != m_aliases.end())
+                if (it != m_local_aliases.end())
                     return mk_result(format(it->second), 1);
             }
             result r;
@@ -1128,9 +1142,9 @@ class pp_fn {
                 }
             }
             if (!main && m_extra_lets && is_shared(e) && r.second > m_alias_min_weight) {
-                name new_aux = name(m_aux, m_aliases_defs.size()+1);
+                name new_aux = name(m_aux, m_local_aliases_defs.size()+1);
-                m_aliases.insert(e, new_aux);
+                m_local_aliases.insert(e, new_aux);
-                m_aliases_defs.emplace_back(new_aux, r.first);
+                m_local_aliases_defs.emplace_back(new_aux, r.first);
                 return mk_result(format(new_aux), 1);
             }
             return r;
@@ -1175,8 +1189,8 @@ class pp_fn {
     }
 
     void init(expr const & e) {
-        m_aliases.clear();
+        m_local_aliases.clear();
-        m_aliases_defs.clear();
+        m_local_aliases_defs.clear();
         m_num_steps = 0;
         m_aux = find_unused_prefix(e);
     }

tests/lean/alias1.lean

@@ -0,0 +1,3 @@
+Alias BB : Bool.
+Variable x : BB.
+Show Environment 1.

tests/lean/alias1.lean.expected.out

@@ -0,0 +1,4 @@
+  Set: pp::colors
+  Set: pp::unicode
+  Assumed: x
+Variable x : BB

tests/lean/alias2.lean

@@ -0,0 +1,14 @@
+Variable Natural : Type.
+Alias ℕℕ : Natural.
+Variable x : Natural.
+Show Environment 1.
+SetOption pp::unicode false.
+Show Environment 1.
+SetOption pp::unicode true.
+Show Environment 1.
+Alias NN : Natural.
+Show Environment 1.
+Alias ℕℕℕ : Natural.
+Show Environment 1.
+SetOption pp::unicode false.
+Show Environment 1.

tests/lean/alias2.lean.expected.out

@@ -0,0 +1,13 @@
+  Set: pp::colors
+  Set: pp::unicode
+  Assumed: Natural
+  Assumed: x
+Variable x : ℕℕ
+  Set: pp::unicode
+Variable x : Natural
+  Set: pp::unicode
+Variable x : ℕℕ
+Variable x : NN
+Variable x : ℕℕℕ
+  Set: pp::unicode
+Variable x : NN