fix(frontends/lean/pp): fixes #634

trying again...

src/frontends/lean/decl_cmds.cpp

@@ -41,6 +41,8 @@ namespace lean {
 static environment declare_universe(parser & p, environment env, name const & n, bool local) {
     if (local) {
         p.add_local_level(n, mk_param_univ(n), local);
+    } else if (in_section(env)) {
+        p.add_local_level(n, mk_param_univ(n), false);
     } else {
         name const & ns = get_namespace(env);
         name full_n  = ns + n;

src/frontends/lean/local_ref_info.cpp

@@ -13,25 +13,23 @@ When we declare a definition in a section, we create an alias for it that fixes
 universe parameters. We have to store the number of parameters and universes that have been fixed
 to be able to correctly pretty print terms.
 */
-
 struct local_ref_entry {
     name     m_name;
-    unsigned m_num_fix_univs;
-    unsigned m_num_fix_args;
+    expr     m_ref;
     local_ref_entry() {}
-    local_ref_entry(name const & n, unsigned u, unsigned p):
-        m_name(n), m_num_fix_univs(u), m_num_fix_args(p) {}
+    local_ref_entry(name const & n, expr const & ref):
+    m_name(n), m_ref(ref) {}
 };
 
 static name * g_local_ref_name  = nullptr;
 static std::string * g_key      = nullptr;
 
 struct local_ref_config {
-    typedef name_map<pair<unsigned, unsigned>> state;
-    typedef local_ref_entry                    entry;
+    typedef name_map<expr>   state;
+    typedef local_ref_entry  entry;
 
     static void add_entry(environment const &, io_state const &, state & s, entry const & e) {
-        s.insert(e.m_name, mk_pair(e.m_num_fix_univs, e.m_num_fix_args));
+        s.insert(e.m_name, e.m_ref);
     }
     static name const & get_class_name() {
         return *g_local_ref_name;
@@ -53,16 +51,16 @@ struct local_ref_config {
 template class scoped_ext<local_ref_config>;
 typedef scoped_ext<local_ref_config> local_ref_ext;
 
-environment save_local_ref_info(environment const & env, name const & n, unsigned num_fix_univs, unsigned num_fix_args) {
+environment save_local_ref_info(environment const & env, name const & n, expr const & ref) {
     bool persistent = false;
-    return local_ref_ext::add_entry(env, get_dummy_ios(), local_ref_entry(n, num_fix_univs, num_fix_args), persistent);
+    return local_ref_ext::add_entry(env, get_dummy_ios(), local_ref_entry(n, ref), persistent);
 }
 
-optional<pair<unsigned, unsigned>> get_local_ref_info(environment const & env, name const & n) {
+optional<expr> get_local_ref_info(environment const & env, name const & n) {
     if (auto r = local_ref_ext::get_state(env).find(n))
-        return optional<pair<unsigned, unsigned>>(*r);
+        return some_expr(*r);
     else
-        return optional<pair<unsigned, unsigned>>();
+        return none_expr();
 }
 
 void initialize_local_ref_info() {

src/frontends/lean/local_ref_info.h

@@ -7,8 +7,8 @@ Author: Leonardo de Moura
 #pragma once
 
 namespace lean {
-environment save_local_ref_info(environment const & env, name const & n, unsigned num_fix_univs, unsigned num_fix_args);
-optional<pair<unsigned, unsigned>> get_local_ref_info(environment const & env, name const & n);
+environment save_local_ref_info(environment const & env, name const & n, expr const & ref);
+optional<expr> get_local_ref_info(environment const & env, name const & n);
 void initialize_local_ref_info();
 void finalize_local_ref_info();
 }

src/frontends/lean/parser.cpp

@@ -510,10 +510,15 @@ environment parser::add_local_ref(environment const & env, name const & n, expr
         if (is_explicit(f))
             f = get_explicit_arg(f);
         if (is_constant(f)) {
-            return save_local_ref_info(env, const_name(f), length(const_levels(f)), args.size());
+            return save_local_ref_info(env, const_name(f), ref);
+        } else {
+            return env;
         }
+    } else if (is_constant(ref) && const_levels(ref)) {
+        return save_local_ref_info(env, const_name(ref), ref);
+    } else {
+        return env;
     }
-    return env;
 }
 
 void parser::add_parameter(name const & n, expr const & p) {

src/frontends/lean/pp.cpp

@@ -365,8 +365,7 @@ auto pretty_fn::pp_coercion(expr const & e, unsigned bp, bool ignore_hide) -> re
     }
 }
 
-auto pretty_fn::pp_child_core(expr const & e, unsigned bp, bool ignore_hide) -> result {
-    result r = pp(e, ignore_hide);
+auto pretty_fn::add_paren_if_needed(result const & r, unsigned bp) -> result {
     if (r.rbp() < bp) {
         return result(paren(r.fmt()));
     } else {
@@ -374,24 +373,110 @@ auto pretty_fn::pp_child_core(expr const & e, unsigned bp, bool ignore_hide) ->
     }
 }
 
+auto pretty_fn::pp_child_core(expr const & e, unsigned bp, bool ignore_hide) -> result {
+    return add_paren_if_needed(pp(e, ignore_hide), bp);
+}
+
+static expr consume_ref_annotations(expr const & e) {
+    if (is_explicit(e))
+        return consume_ref_annotations(get_explicit_arg(e));
+    else
+        return e;
+}
+
+enum local_ref_kind { LocalRef, OverridenLocalRef, NotLocalRef };
+
+static local_ref_kind check_local_ref(environment const & env, expr const & e, unsigned & num_ref_univ_params) {
+    expr const & rfn = get_app_fn(e);
+    if (!is_constant(rfn))
+        return NotLocalRef;
+    auto ref = get_local_ref_info(env, const_name(rfn));
+    if (!ref)
+        return NotLocalRef;
+    if (is_as_atomic(*ref))
+        ref = get_as_atomic_arg(*ref);
+    buffer<expr> ref_args, e_args;
+    expr ref_fn = consume_ref_annotations(get_app_args(*ref, ref_args));
+    get_app_args(e, e_args);
+    if (!is_constant(ref_fn) || e_args.size() != ref_args.size()) {
+        return NotLocalRef;
+    }
+    for (unsigned i = 0; i < e_args.size(); i++) {
+        expr e_arg   = e_args[i];
+        expr ref_arg = consume_ref_annotations(ref_args[i]);
+        if (!is_local(e_arg) || !is_local(ref_arg) || local_pp_name(e_arg) != local_pp_name(ref_arg))
+            return OverridenLocalRef;
+    }
+    num_ref_univ_params = length(const_levels(ref_fn));
+    buffer<level> lvls;
+    to_buffer(const_levels(rfn), lvls);
+    if (lvls.size() < num_ref_univ_params) {
+        return NotLocalRef;
+    }
+    for (unsigned i = 0; i < num_ref_univ_params; i++) {
+        level const & l = lvls[i];
+        if (!is_param(l)) {
+            return OverridenLocalRef;
+        }
+        for (unsigned j = 0; j < i; j++)
+            if (lvls[i] == lvls[j]) {
+                return OverridenLocalRef;
+            }
+    }
+    return LocalRef;
+}
+
+static bool is_local_ref(environment const & env, expr const & e) {
+    unsigned num_ref_univ_params;
+    return check_local_ref(env, e, num_ref_univ_params) == LocalRef;
+}
+
+auto pretty_fn::pp_overriden_local_ref(expr const & e) -> result {
+    buffer<expr> args;
+    expr const & fn = get_app_args(e, args);
+    result res_fn;
+    {
+        flet<bool> set1(m_full_names, true);
+        res_fn = pp_const(fn, some(0u));
+    }
+    format fn_fmt    = res_fn.fmt();
+    if (const_name(fn).is_atomic())
+        fn_fmt = compose(format("_root_."), fn_fmt);
+    if (m_implict && has_implicit_args(fn))
+        fn_fmt = compose(*g_explicit_fmt, fn_fmt);
+    format r_fmt = fn_fmt;
+    expr curr_fn = fn;
+    for (unsigned i = 0; i < args.size(); i++) {
+        expr const & arg = args[i];
+        if (m_implict || !is_implicit(curr_fn)) {
+            result res_arg   = pp_child(arg, max_bp());
+            r_fmt = group(compose(r_fmt, nest(m_indent, compose(line(), res_arg.fmt()))));
+        }
+        curr_fn = mk_app(curr_fn, arg);
+    }
+    return result(max_bp()-1, r_fmt);
+}
+
+bool pretty_fn::ignore_local_ref(expr const & e) {
+    expr const & fn = get_app_fn(e);
+    return m_full_names && (!is_constant(fn) || !const_name(fn).is_atomic());
+}
+
 // Return some result if \c e is of the form (c p_1 ... p_n) where
 // c is a constant, and p_i's are parameters fixed in a section.
 auto pretty_fn::pp_local_ref(expr const & e) -> optional<result> {
-    if (m_full_names)
+    if (ignore_local_ref(e))
         return optional<result>();
-    lean_assert(is_app(e));
-    expr const & rfn = get_app_fn(e);
-    if (is_constant(rfn)) {
-        if (auto info = get_local_ref_info(m_env, const_name(rfn))) {
-            buffer<expr> args;
-            get_app_args(e, args);
-            if (args.size() == info->second) {
-                // TODO(Leo): must check if the arguments are really the fixed parameters.
-                return some(pp_const(rfn));
-            }
-        }
+    unsigned num_ref_univ_params;
+    switch (check_local_ref(m_env, e, num_ref_univ_params)) {
+    case NotLocalRef:
+        return optional<result>();
+    case LocalRef:
+        return some(pp_const(get_app_fn(e), optional<unsigned>(num_ref_univ_params)));
+    case OverridenLocalRef:
+        return some(pp_overriden_local_ref(e));
     }
-    return optional<result>();
+    lean_unreachable();
 }
 
 auto pretty_fn::pp_child(expr const & e, unsigned bp, bool ignore_hide) -> result {
@@ -399,7 +484,7 @@ auto pretty_fn::pp_child(expr const & e, unsigned bp, bool ignore_hide) -> resul
         return pp_abbreviation(e, *it, false, bp, ignore_hide);
     if (is_app(e)) {
         if (auto r = pp_local_ref(e))
-            return *r;
+            return add_paren_if_needed(*r, bp);
         expr const & f = app_fn(e);
         if (auto it = is_abbreviated(f)) {
             return pp_abbreviation(e, *it, true, bp, ignore_hide);
@@ -446,7 +531,14 @@ optional<name> pretty_fn::is_abbreviated(expr const & e) const {
     return optional<name>();
 }
 
-auto pretty_fn::pp_const(expr const & e) -> result {
+auto pretty_fn::pp_const(expr const & e, optional<unsigned> const & num_ref_univ_params) -> result {
+    if (!num_ref_univ_params) {
+        if (auto r = pp_local_ref(e))
+            return *r;
+    }
+    // Remark: if num_ref_univ_params is "some", then it contains the number of
+    // universe levels that are fixed in a section. That is, \c e corresponds to
+    // a constant in a section which has fixed levels.
     name n = const_name(e);
     if (!m_full_names) {
         if (auto it = is_aliased(n)) {
@@ -474,11 +566,11 @@ auto pretty_fn::pp_const(expr const & e) -> result {
         unsigned first_idx = 0;
         buffer<level> ls;
         to_buffer(const_levels(e), ls);
-        if (auto info = get_local_ref_info(m_env, n)) {
-            if (ls.size() <= info->first)
+        if (num_ref_univ_params) {
+            if (ls.size() <= *num_ref_univ_params)
                 return result(format(n));
             else
-                first_idx = info->first;
+                first_idx = *num_ref_univ_params;
         }
         format r = compose(format(n), format(".{"));
         bool first = true;
@@ -547,7 +639,7 @@ auto pretty_fn::pp_app(expr const & e) -> result {
     bool ignore_hide = true;
     result res_fn    = pp_child(fn, max_bp()-1, ignore_hide);
     format fn_fmt    = res_fn.fmt();
-    if (m_implict && !is_app(fn) && has_implicit_args(fn))
+    if (m_implict && (!is_app(fn) || (!ignore_local_ref(fn) && is_local_ref(m_env, fn))) && has_implicit_args(fn))
         fn_fmt = compose(*g_explicit_fmt, fn_fmt);
     result res_arg = pp_child(app_arg(e), max_bp());
     return result(max_bp()-1, group(compose(fn_fmt, nest(m_indent, compose(line(), res_arg.fmt())))));

src/frontends/lean/pp.h

@@ -91,6 +91,10 @@ private:
     optional<result> pp_notation(notation_entry const & entry, buffer<optional<expr>> & args);
     optional<result> pp_notation(expr const & e);
 
+    result add_paren_if_needed(result const & r, unsigned bp);
+
+    result pp_overriden_local_ref(expr const & e);
+    bool ignore_local_ref(expr const & e);
     optional<result> pp_local_ref(expr const & e);
 
     result pp_coercion_fn(expr const & e, unsigned sz, bool ignore_hide = false);
@@ -99,7 +103,7 @@ private:
     result pp_child(expr const & e, unsigned bp, bool ignore_hide = false);
     result pp_var(expr const & e);
     result pp_sort(expr const & e);
-    result pp_const(expr const & e);
+    result pp_const(expr const & e, optional<unsigned> const & num_ref_univs = optional<unsigned>());
     result pp_meta(expr const & e);
     result pp_local(expr const & e);
     result pp_app(expr const & e);

tests/lean/634.lean.expected.out

@@ -1,5 +1,5 @@
-A n : Type₁
+@A n : Type₁
 @foo.A X n : Type₁
 @foo.A X n : Type₁
-A n : Type₁
-A n : Type₁
+@A n : Type₁
+@A n : Type₁

tests/lean/634b.lean

@@ -0,0 +1,41 @@
+open nat
+namespace foo
+section
+  parameter (X : Type₁)
+  definition A {n : ℕ} : Type₁ := X
+  definition B : Type₁ := X
+  variable {n : ℕ}
+  check @A n
+  check foo.A nat
+  check foo.A (X × B)
+  check @foo.A (X × B) 10
+  check @foo.A (@foo.B (@A n)) n
+  check @foo.A (@foo.B (@foo.A X n)) n
+  check @foo.A (@foo.B (@foo.A nat n)) n
+
+  set_option pp.full_names true
+  check A
+  check foo.A nat
+  check @foo.A (X × B) 10
+  check @foo.A (@foo.B (@foo.A X n)) n
+  check @foo.A (@foo.B (@foo.A nat n)) n
+
+  set_option pp.full_names false
+
+  set_option pp.implicit true
+  check @A n
+  check @foo.A nat 10
+  check @foo.A X n
+  set_option pp.full_names true
+  check @foo.A X n
+  check @A n
+
+  set_option pp.full_names false
+  check @foo.A X n
+  check @foo.A B n
+  check @foo.A (@foo.B (@A n)) n
+  check @foo.A (@foo.B (@foo.A X n)) n
+  check @foo.A (@foo.B (@foo.A nat n)) n
+  check @A n
+end
+end foo

tests/lean/634b.lean.expected.out

@@ -0,0 +1,23 @@
+A : Type₁
+foo.A ℕ : Type₁
+foo.A (X × B) : Type₁
+foo.A (X × B) : Type₁
+foo.A (foo.B A) : Type₁
+foo.A (foo.B A) : Type₁
+foo.A (foo.B (foo.A ℕ)) : Type₁
+foo.A X : Type₁
+foo.A ℕ : Type₁
+foo.A (X × foo.B X) : Type₁
+foo.A (foo.B (foo.A X)) : Type₁
+foo.A (foo.B (foo.A ℕ)) : Type₁
+@A n : Type₁
+@foo.A ℕ 10 : Type₁
+@A n : Type₁
+@foo.A X n : Type₁
+@foo.A X n : Type₁
+@A n : Type₁
+@foo.A B n : Type₁
+@foo.A (foo.B (@A n)) n : Type₁
+@foo.A (foo.B (@A n)) n : Type₁
+@foo.A (foo.B (@foo.A ℕ n)) n : Type₁
+@A n : Type₁

tests/lean/634c.lean

@@ -0,0 +1,38 @@
+open nat
+section
+  parameter (X : Type₁)
+  definition A {n : ℕ} : Type₁ := X
+  definition B : Type₁ := X
+  variable {n : ℕ}
+  check @A n
+  check _root_.A nat
+  check _root_.A (X × B)
+  check @_root_.A (X × B) 10
+  check @_root_.A (_root_.B (@_root_.A X n)) n
+  check @_root_.A (@_root_.B (@_root_.A nat n)) n
+
+  set_option pp.full_names true
+  check A
+  check _root_.A nat
+  check @_root_.A (X × B) 10
+  check @_root_.A (@_root_.B (@_root_.A X n)) n
+  check @_root_.A (@_root_.B (@_root_.A nat n)) n
+
+  set_option pp.full_names false
+
+  set_option pp.implicit true
+  check @A n
+  check @_root_.A nat 10
+  check @_root_.A X n
+  set_option pp.full_names true
+  check @_root_.A X n
+  check @_root_.A B n
+
+  set_option pp.full_names false
+  check @_root_.A X n
+  check @_root_.A B n
+  check @_root_.A (@_root_.B (@A n)) n
+  check @_root_.A (@_root_.B (@_root_.A X n)) n
+  check @_root_.A (@_root_.B (@_root_.A nat n)) n
+  check @A n
+end

tests/lean/634c.lean.expected.out

@@ -0,0 +1,22 @@
+A : Type₁
+_root_.A ℕ : Type₁
+_root_.A (X × B) : Type₁
+_root_.A (X × B) : Type₁
+_root_.A (_root_.B A) : Type₁
+_root_.A (_root_.B (_root_.A ℕ)) : Type₁
+A : Type₁
+_root_.A ℕ : Type₁
+_root_.A (X × B) : Type₁
+_root_.A (_root_.B A) : Type₁
+_root_.A (_root_.B (_root_.A ℕ)) : Type₁
+@A n : Type₁
+@_root_.A ℕ 10 : Type₁
+@A n : Type₁
+@A n : Type₁
+@_root_.A B n : Type₁
+@A n : Type₁
+@_root_.A B n : Type₁
+@_root_.A (_root_.B (@A n)) n : Type₁
+@_root_.A (_root_.B (@A n)) n : Type₁
+@_root_.A (_root_.B (@_root_.A ℕ n)) n : Type₁
+@A n : Type₁