feat(frontends/lean/pp): improve pretty printer for prefix and postfix notation

closes #491

src/frontends/lean/pp.cpp

@@ -279,6 +279,7 @@ void pretty_fn::set_options_core(options const & o) {
     m_beta            = get_pp_beta(o);
     m_numerals        = get_pp_numerals(o);
     m_abbreviations   = get_pp_abbreviations(o);
+    m_extra_spaces    = get_pp_extra_spaces(o);
     m_hide_full_terms = get_formatter_hide_full_terms(o);
     m_num_nat_coe     = m_numerals && !m_coercion && has_coercion_num_nat(m_env);
 }
@@ -880,6 +881,7 @@ auto pretty_fn::pp_notation(notation_entry const & entry, buffer<optional<expr>>
         unsigned last_rbp  = inf_bp()-1;
         unsigned token_lbp = 0;
         bool extra_space   = false;
+        bool last_is_skip  = false;
         bool last          = true;
         while (i > 0) {
             --i;
@@ -889,8 +891,10 @@ auto pretty_fn::pp_notation(notation_entry const & entry, buffer<optional<expr>>
             switch (a.kind()) {
             case notation::action_kind::Skip:
                 curr = format(tk);
-                if (last)
+                if (last) {
-                    last_rbp = inf_bp();
+                    last_rbp     = inf_bp();
+                    last_is_skip = true;
+                }
                 break;
             case notation::action_kind::Expr:
                 if (args.empty() || !args.back()) {
@@ -901,7 +905,10 @@ auto pretty_fn::pp_notation(notation_entry const & entry, buffer<optional<expr>>
                     result e_r   = pp_notation_child(e, token_lbp, a.rbp());
                     format e_fmt = e_r.fmt();
                     curr = format(tk);
-                    if (add_extra_space(tk))
+                    // we add space after the token only when
+                    // 1- add_extra_space(tk) is true AND
+                    // 2- tk is the first token in a nud notation
+                    if (add_extra_space(tk) && (!entry.is_nud() || i != 0 || m_extra_spaces))
                         curr = curr + space();
                     curr = curr + e_fmt;
                     if (last)
@@ -959,7 +966,7 @@ auto pretty_fn::pp_notation(notation_entry const & entry, buffer<optional<expr>>
                         --i;
                         result arg_res = pp_notation_child(rec_args[i], curr_lbp, a.rbp());
                         if (i == 0) {
-                            if (add_extra_space_first(tk))
+                            if (m_extra_spaces && add_extra_space_first(tk))
                                 curr = format(tk) + space() + arg_res.fmt() + curr;
                             else
                                 curr = format(tk) + arg_res.fmt() + curr;
@@ -1044,7 +1051,10 @@ auto pretty_fn::pp_notation(notation_entry const & entry, buffer<optional<expr>>
             expr e = *args.back();
             args.pop_back();
             format e_fmt = pp_notation_child(e, token_lbp, 0).fmt();
-            fmt = e_fmt + space() + fmt;
+            if (m_extra_spaces || !last_is_skip)
+                fmt = e_fmt + space() + fmt;
+            else
+                fmt = e_fmt + fmt;
         }
         return optional<result>(result(first_lbp, last_rbp, fmt));
     }

src/frontends/lean/pp.h

@@ -68,6 +68,7 @@ private:
     bool                m_numerals;
     bool                m_abbreviations;
     bool                m_hide_full_terms;
+    bool                m_extra_spaces;
 
     name mk_metavar_name(name const & m);
     name mk_local_name(name const & n, name const & suggested);

src/library/pp_options.cpp

@@ -63,6 +63,10 @@ Author: Leonardo de Moura
 #define LEAN_DEFAULT_PP_ABBREVIATIONS true
 #endif
 
+#ifndef LEAN_DEFAULT_PP_EXTRA_SPACES
+#define LEAN_DEFAULT_PP_EXTRA_SPACES false
+#endif
+
 namespace lean {
 static name * g_pp_max_depth       = nullptr;
 static name * g_pp_max_steps       = nullptr;
@@ -78,6 +82,7 @@ static name * g_pp_purify_locals   = nullptr;
 static name * g_pp_beta            = nullptr;
 static name * g_pp_numerals        = nullptr;
 static name * g_pp_abbreviations   = nullptr;
+static name * g_pp_extra_spaces    = nullptr;
 static list<options> * g_distinguishing_pp_options = nullptr;
 
 void initialize_pp_options() {
@@ -95,6 +100,7 @@ void initialize_pp_options() {
     g_pp_beta            = new name{"pp", "beta"};
     g_pp_numerals        = new name{"pp", "numerals"};
     g_pp_abbreviations   = new name{"pp", "abbreviations"};
+    g_pp_extra_spaces    = new name{"pp", "extra_spaces"};
     register_unsigned_option(*g_pp_max_depth, LEAN_DEFAULT_PP_MAX_DEPTH,
                              "(pretty printer) maximum expression depth, after that it will use ellipsis");
     register_unsigned_option(*g_pp_max_steps, LEAN_DEFAULT_PP_MAX_STEPS,
@@ -125,6 +131,8 @@ void initialize_pp_options() {
                          "(pretty printer) display nat/num numerals in decimal notation");
     register_bool_option(*g_pp_abbreviations, LEAN_DEFAULT_PP_ABBREVIATIONS,
                          "(pretty printer) display abbreviations (i.e., revert abbreviation expansion when pretty printing)");
+    register_bool_option(*g_pp_extra_spaces, LEAN_DEFAULT_PP_EXTRA_SPACES,
+                         "(pretty printer) add space after prefix operators and before postfix ones");
 
     options universes_true(*g_pp_universes, true);
     options full_names_true(*g_pp_full_names, true);
@@ -139,6 +147,7 @@ void initialize_pp_options() {
 }
 
 void finalize_pp_options() {
+    delete g_pp_extra_spaces;
     delete g_pp_abbreviations;
     delete g_pp_numerals;
     delete g_pp_max_depth;
@@ -180,5 +189,6 @@ bool     get_pp_purify_locals(options const & opts)   { return opts.get_bool(*g_
 bool     get_pp_beta(options const & opts)            { return opts.get_bool(*g_pp_beta, LEAN_DEFAULT_PP_BETA); }
 bool     get_pp_numerals(options const & opts)        { return opts.get_bool(*g_pp_numerals, LEAN_DEFAULT_PP_NUMERALS); }
 bool     get_pp_abbreviations(options const & opts)   { return opts.get_bool(*g_pp_abbreviations, LEAN_DEFAULT_PP_ABBREVIATIONS); }
+bool     get_pp_extra_spaces(options const & opts)    { return opts.get_bool(*g_pp_extra_spaces, LEAN_DEFAULT_PP_EXTRA_SPACES); }
 list<options> const & get_distinguishing_pp_options() { return *g_distinguishing_pp_options; }
 }

src/library/pp_options.h

@@ -31,6 +31,7 @@ bool     get_pp_purify_metavars(options const & opts);
 bool     get_pp_purify_locals(options const & opts);
 bool     get_pp_numerals(options const & opts);
 bool     get_pp_abbreviations(options const & opts);
+bool     get_pp_extra_spaces(options const & opts);
 list<options> const & get_distinguishing_pp_options();
 
 void initialize_pp_options();

tests/lean/491.lean

@@ -0,0 +1,8 @@
+import data.int
+open eq.ops int
+variable A : Type
+variables a b : A
+variable H : a = b
+check H⁻¹
+check -1
+check 1 + -2

tests/lean/491.lean.expected.out

@@ -0,0 +1,3 @@
+H⁻¹ : b = a
+-1 : ℤ
+1 + -2 : ℤ