feat(library/export,checker): add basic support for mixfix notation

doc/export_format.md

@@ -228,3 +228,15 @@ Quotient declaration
 
 The declaration of the quotient type and its computational rule is exported as
 `#QUOT`.
+
+Notation
+--------
+
+The export contains information about prefix, postfix, and infix notation,
+where the head symbol is a constant.
+These are indicated using the `#PREFIX`, `#POSTFIX`, and `#INFIX` commands.
+They all follow the same syntax:
+
+```
+#PREFIX <name_idx> <precedence> <token>
+```

src/checker/checker.cpp

@@ -67,13 +67,14 @@ int main(int argc, char ** argv) {
 
         unsigned trust_lvl = 0;
         auto env = mk_environment(trust_lvl);
-        import_from_text(in, env);
+        lowlevel_notations notations;
+        import_from_text(in, env, notations);
 
         env.for_each_declaration([&] (declaration const & d) {
             if (d.is_axiom()) {
                 std::cout << compose_many(
                         {format("axiom"), space(), simple_pp(d.get_name()), space(), format(":"), space(),
-                         simple_pp(env, d.get_type()), line()});
+                         simple_pp(env, d.get_type(), notations), line()});
             }
         });
 
@@ -82,7 +83,7 @@ int main(int argc, char ** argv) {
             auto d = env.get(n);
             std::cout << compose_many(
                     {format("theorem"), space(), simple_pp(d.get_name()), space(), format(":"), space(),
-                     simple_pp(env, d.get_type()), line()});
+                     simple_pp(env, d.get_type(), notations), line()});
         }
 
         unsigned num_decls = 0;

src/checker/simple_pp.cpp

@@ -39,9 +39,11 @@ static format pp_name(name n) {
 
 struct simple_pp_fn {
     type_checker m_tc;
+    lowlevel_notations m_notations;
     unsigned m_indent = 0;
 
-    simple_pp_fn(environment const & env) : m_tc(env) {}
+    simple_pp_fn(environment const & env, lowlevel_notations const & notations) :
+            m_tc(env), m_notations(notations) {}
 
     static constexpr unsigned MAX_PRIO = 1024;
     struct result {
@@ -130,15 +132,45 @@ struct simple_pp_fn {
         }
     }
 
-    result pp_app(expr const & e) {
-        auto fn = get_app_fn(e);
-        format fmt = pp(fn).maybe_paren();
-        buffer<expr> args; get_app_args(e, args);
-        for (auto & arg : args) {
-            if (!is_implicit(fn)) {
-                fmt = compose_many({fmt, space(), pp(arg).maybe_paren()});
+    expr get_explicit_args(expr e, buffer<expr> & args) {
+        while (is_app(e)) {
+            if (!is_implicit(app_fn(e))) {
+                args.push_back(app_arg(e));
             }
-            fn = mk_app(fn, arg);
+            e = app_fn(e);
+        }
+        std::reverse(args.begin(), args.end());
+        return e;
+    }
+
+    result pp_app(expr const & e) {
+        buffer<expr> args;
+        auto fn = get_explicit_args(e, args);
+
+        if (is_constant(fn)) {
+            auto it = m_notations.find(const_name(fn));
+            if (it != m_notations.end()) {
+                format fmt;
+                auto prec = it->second.m_prec;
+                if (it->second.m_kind == lowlevel_notation_kind::Postfix && args.size() == 1) {
+                    fmt = compose_many({pp(args[0]).maybe_paren(prec), format(it->second.m_token)});
+                }
+                if (it->second.m_kind == lowlevel_notation_kind::Prefix && args.size() == 1) {
+                    fmt = compose_many({format(it->second.m_token), pp(args[0]).maybe_paren(prec)});
+                }
+                if (it->second.m_kind == lowlevel_notation_kind::Infix && args.size() == 2) {
+                    fmt = compose_many({pp(args[0]).maybe_paren(prec),
+                                        format(it->second.m_token),
+                                        pp(args[1]).maybe_paren(prec)});
+                }
+                if (!fmt.is_nil_fmt())
+                    return result(group(fmt), it->second.m_prec-1);
+            }
+        }
+
+        format fmt = pp(fn).maybe_paren();
+        for (auto & arg : args) {
+            fmt = compose_many({fmt, space(), pp(arg).maybe_paren()});
         }
         return result(group(fmt), MAX_PRIO-1);
     }
@@ -269,8 +301,8 @@ format simple_pp(name const & n) {
     return pp_name(n);
 }
 
-format simple_pp(environment const & env, expr const & e) {
-    simple_pp_fn fn(env);
+format simple_pp(environment const & env, expr const & e, lowlevel_notations const & notations) {
+    simple_pp_fn fn(env, notations);
     fn.mark_used_const_names(e);
     return fn.pp(e);
 }

src/checker/simple_pp.h

@@ -6,12 +6,13 @@ Author: Gabriel Ebner
 */
 #pragma once
 #include "kernel/environment.h"
+#include "checker/text_import.h"
 
 namespace lean {
 
 format compose_many(std::initializer_list<format> const & fmts);
 
 format simple_pp(name const & n);
-format simple_pp(environment const & env, expr const & e);
+format simple_pp(environment const & env, expr const & e, lowlevel_notations const & notations);
 
 }

src/checker/text_import.cpp

@@ -21,6 +21,8 @@ struct text_importer {
     std::unordered_map<unsigned, name> m_name;
     std::unordered_map<unsigned, level> m_level;
 
+    lowlevel_notations m_notations;
+
     environment m_env;
 
     text_importer(environment const & env) : m_env(env) {
@@ -78,6 +80,20 @@ struct text_importer {
         m_env = m_env.add(check(m_env, mk_axiom(m_name.at(name_idx), ls, m_expr.at(type_idx))));
     }
 
+    void handle_notation(std::istream & in, lowlevel_notation_kind kind) {
+        unsigned name_idx, prec;
+        in >> name_idx >> prec;
+
+        std::string token;
+        std::getline(in, token);
+        if (!token.empty() && token.front())
+            token.erase(token.begin());
+        if (!token.empty() && token.back() == '\n')
+            token.erase(token.end() - 1);
+
+        m_notations[m_name.at(name_idx)] = { kind, token, prec };
+    }
+
     binder_info read_binder_info(std::string const & tok) {
         if (tok == "#BI") {
             return mk_implicit_binder_info();
@@ -106,6 +122,12 @@ struct text_importer {
             handle_ax(in);
         } else if (cmd == "#QUOT") {
             m_env = declare_quotient(m_env);
+        } else if (cmd == "#PREFIX") {
+            handle_notation(in, lowlevel_notation_kind::Prefix);
+        } else if (cmd == "#POSTFIX") {
+            handle_notation(in, lowlevel_notation_kind::Postfix);
+        } else if (cmd == "#INFIX") {
+            handle_notation(in, lowlevel_notation_kind::Infix);
         } else if (std::istringstream(cmd) >> idx) {
             std::string kind;
             in >> kind;
@@ -162,7 +184,7 @@ struct text_importer {
     }
 };
 
-void import_from_text(std::istream & in, environment & env) {
+void import_from_text(std::istream & in, environment & env, lowlevel_notations & notations) {
     text_importer importer(env);
 
     std::string line;
@@ -179,6 +201,7 @@ void import_from_text(std::istream & in, environment & env) {
     }
 
     env = importer.m_env;
+    notations = std::move(importer.m_notations);
 }
 
 }

src/checker/text_import.h

@@ -7,9 +7,21 @@ Author: Gabriel Ebner
 #pragma once
 #include "kernel/environment.h"
 #include <iostream>
+#include <unordered_map>
 
 namespace lean {
 
-void import_from_text(std::istream & in, environment & env);
+enum class lowlevel_notation_kind {
+    Prefix, Postfix, Infix,
+};
+struct lowlevel_notation_info {
+    lowlevel_notation_kind m_kind;
+    std::string m_token;
+    unsigned m_prec;
+};
+
+using lowlevel_notations = std::unordered_map<name, lowlevel_notation_info, name_hash>;
+
+void import_from_text(std::istream & in, environment & env, lowlevel_notations & notations);
 
 }

src/frontends/lean/parser_config.cpp

@@ -357,6 +357,15 @@ list<notation_entry> get_notation_entries(environment const & env, head_index co
         return list<notation_entry>();
 }
 
+std::vector<notation_entry> get_notation_entries(environment const &env) {
+    std::vector<notation_entry> entries;
+    notation_ext::get_state(env).m_inv_map.for_each_entry(
+            [&] (head_index const &, notation_entry const & entry) {
+                entries.emplace_back(entry);
+            });
+    return entries;
+}
+
 struct cmd_ext : public environment_extension {
     cmd_table m_cmds;
     cmd_ext() {

src/frontends/lean/parser_config.h

@@ -90,6 +90,8 @@ list<expr> get_mpz_notation(environment const & env, mpz const & n);
 */
 list<notation_entry> get_notation_entries(environment const & env, head_index const & idx);
 
+std::vector<notation_entry> get_notation_entries(environment const & env);
+
 void initialize_parser_config();
 void finalize_parser_config();
 }

src/kernel/quotient/quotient.cpp

@@ -18,7 +18,6 @@ Quotient types for kernels with proof irrelevance.
 
 namespace lean {
 static name * g_quotient_extension = nullptr;
-static name * g_propext        = nullptr;
 static name * g_quotient       = nullptr;
 static name * g_quotient_lift  = nullptr;
 static name * g_quotient_ind   = nullptr;
@@ -211,12 +210,11 @@ bool has_quotient(environment const & env) {
 }
 
 std::vector<name> quotient_required_decls() {
-    return {"eq", *g_propext};
+    return {"eq"};
 }
 
 void initialize_quotient_module() {
     g_quotient_extension = new name("quotient_extension");
-    g_propext            = new name{"propext"};
     g_quotient           = new name{"quot"};
     g_quotient_lift      = new name{"quot", "lift"};
     g_quotient_ind       = new name{"quot", "ind"};
@@ -226,7 +224,6 @@ void initialize_quotient_module() {
 
 void finalize_quotient_module() {
     delete g_ext;
-    delete g_propext;
     delete g_quotient_extension;
     delete g_quotient;
     delete g_quotient_lift;

src/library/export.cpp

@@ -5,6 +5,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Author: Leonardo de Moura
 */
 #include <unordered_map>
+#include "frontends/lean/parser_config.h"
 #include "kernel/quotient/quotient.h"
 #include "kernel/expr_maps.h"
 #include "kernel/for_each_fn.h"
@@ -304,6 +305,39 @@ class exporter {
         m_out << "#QUOT\n";
     }
 
+    void export_notation(notation_entry const & entry) {
+        if (entry.parse_only()) return;
+        if (length(entry.get_transitions()) != 1) return;
+        auto & t = head(entry.get_transitions());
+
+        buffer<expr> args;
+        auto & fn = get_app_rev_args(entry.get_expr(), args);
+
+        char const * type = nullptr;
+        if (args.size() == 1 && args[0] == mk_var(0)) {
+            if (entry.is_nud()) {
+                type = "#PREFIX";
+            } else {
+                type = "#POSTFIX";
+            }
+        } else if (!entry.is_nud() && args.size() == 2 && args[0] == mk_var(0) && args[1] == mk_var(1)) {
+            type = "#INFIX";
+        }
+
+        if (type && is_constant(fn)) {
+            auto fni = export_name(const_name(fn));
+            auto prec_opt = get_expr_precedence(get_token_table(m_env), t.get_token().get_string());
+            auto prec = prec_opt ? *prec_opt : 0;
+            m_out << type << " " << fni << " " << prec << " " << t.get_pp_token().get_string() << "\n";
+        }
+    }
+
+    void export_notation() {
+        for (auto & entry : get_notation_entries(m_env)) {
+            export_notation(entry);
+        }
+    }
+
 public:
     exporter(std::ostream & out, environment const & env) : m_out(out), m_env(env) {}
 
@@ -314,6 +348,7 @@ public:
         if (has_quotient(m_env))
             export_quotient();
         export_declarations();
+        export_notation();
     }
 };