From 53653c352643751c4b62ff63ec5e555f11dae8eb Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Tue, 21 Apr 2015 22:40:20 -0700
Subject: [PATCH] fix(frontends/lean): pretty printing in sections with
 parameters

fix #530
---
 src/frontends/lean/CMakeLists.txt          |  2 +-
 src/frontends/lean/builtin_cmds.cpp        | 21 +++---
 src/frontends/lean/decl_cmds.cpp           |  4 +-
 src/frontends/lean/init_module.cpp         |  3 +
 src/frontends/lean/local_ref_info.cpp      | 79 ++++++++++++++++++++++
 src/frontends/lean/local_ref_info.h        | 14 ++++
 src/frontends/lean/parser.cpp              | 15 ++++
 src/frontends/lean/parser.h                |  1 +
 src/frontends/lean/pp.cpp                  | 22 +++++-
 src/frontends/lean/type_util.cpp           |  2 +-
 tests/lean/errors.lean.expected.out        |  2 +-
 tests/lean/sec_param_pp.lean               | 20 ++++++
 tests/lean/sec_param_pp.lean.expected.out  |  4 ++
 tests/lean/sec_param_pp2.lean              | 17 +++++
 tests/lean/sec_param_pp2.lean.expected.out |  4 ++
 15 files changed, 193 insertions(+), 17 deletions(-)
 create mode 100644 src/frontends/lean/local_ref_info.cpp
 create mode 100644 src/frontends/lean/local_ref_info.h
 create mode 100644 tests/lean/sec_param_pp.lean
 create mode 100644 tests/lean/sec_param_pp.lean.expected.out
 create mode 100644 tests/lean/sec_param_pp2.lean
 create mode 100644 tests/lean/sec_param_pp2.lean.expected.out

diff --git a/src/frontends/lean/CMakeLists.txt b/src/frontends/lean/CMakeLists.txt
index e992f11ef8..210fd8f252 100644
--- a/src/frontends/lean/CMakeLists.txt
+++ b/src/frontends/lean/CMakeLists.txt
@@ -8,6 +8,6 @@ structure_cmd.cpp info_manager.cpp info_annotation.cpp find_cmd.cpp
 coercion_elaborator.cpp info_tactic.cpp
 init_module.cpp elaborator_context.cpp calc_proof_elaborator.cpp
 parse_tactic_location.cpp parse_rewrite_tactic.cpp
-type_util.cpp elaborator_exception.cpp migrate_cmd.cpp)
+type_util.cpp elaborator_exception.cpp migrate_cmd.cpp local_ref_info.cpp)
 
 target_link_libraries(lean_frontend ${LEAN_LIBS})
diff --git a/src/frontends/lean/builtin_cmds.cpp b/src/frontends/lean/builtin_cmds.cpp
index 4771179efe..b09f8eaa30 100644
--- a/src/frontends/lean/builtin_cmds.cpp
+++ b/src/frontends/lean/builtin_cmds.cpp
@@ -270,12 +270,12 @@ environment namespace_cmd(parser & p) {
     return push_scope(p.env(), p.ios(), scope_kind::Namespace, n);
 }
 
-static void redeclare_aliases(parser & p,
-                              list<pair<name, level>> old_level_entries,
-                              list<pair<name, expr>> old_entries) {
-    environment const & env = p.env();
-    if (!in_context(env) && !in_section(env))
-        return;
+static environment redeclare_aliases(environment env, parser & p,
+                                     list<pair<name, level>> old_level_entries,
+                                     list<pair<name, expr>> old_entries) {
+    environment const & old_env = p.env();
+    if (!in_context(old_env) && !in_section(old_env))
+        return env;
     list<pair<name, expr>> new_entries = p.get_local_entries();
     buffer<pair<name, expr>> to_redeclare;
     name_set popped_locals;
@@ -299,8 +299,9 @@ static void redeclare_aliases(parser & p,
     for (auto const & entry : to_redeclare) {
         expr new_ref = update_local_ref(entry.second, popped_levels, popped_locals);
         if (!is_constant(new_ref))
-            p.add_local_expr(entry.first, new_ref);
+            env = p.add_local_ref(env, entry.first, new_ref);
     }
+    return env;
 }
 
 environment end_scoped_cmd(parser & p) {
@@ -310,12 +311,10 @@ environment end_scoped_cmd(parser & p) {
     if (p.curr_is_identifier()) {
         name n = p.check_atomic_id_next("invalid end of scope, atomic identifier expected");
         environment env = pop_scope(p.env(), p.ios(), n);
-        redeclare_aliases(p, level_entries, entries);
-        return env;
+        return redeclare_aliases(env, p, level_entries, entries);
     } else {
         environment env = pop_scope(p.env(), p.ios());
-        redeclare_aliases(p, level_entries, entries);
-        return env;
+        return redeclare_aliases(env, p, level_entries, entries);
     }
 }
 
diff --git a/src/frontends/lean/decl_cmds.cpp b/src/frontends/lean/decl_cmds.cpp
index 5c00834cf6..96fa8fe0ca 100644
--- a/src/frontends/lean/decl_cmds.cpp
+++ b/src/frontends/lean/decl_cmds.cpp
@@ -952,10 +952,10 @@ class definition_cmd_fn {
             local_ls = remove_local_vars(m_p, local_ls);
             if (!locals.empty()) {
                 expr ref = mk_local_ref(m_real_name, local_ls, locals);
-                m_p.add_local_expr(m_name, ref);
+                m_env = m_p.add_local_ref(m_env, m_name, ref);
             } else if (local_ls) {
                 expr ref = mk_constant(m_real_name, local_ls);
-                m_p.add_local_expr(m_name, ref);
+                m_env = m_p.add_local_ref(m_env, m_name, ref);
             }
         } else {
             update_univ_parameters(m_ls_buffer, collect_univ_params(m_value, collect_univ_params(m_type)), m_p);
diff --git a/src/frontends/lean/init_module.cpp b/src/frontends/lean/init_module.cpp
index fb4f7d3723..21c2384a93 100644
--- a/src/frontends/lean/init_module.cpp
+++ b/src/frontends/lean/init_module.cpp
@@ -27,6 +27,7 @@ Author: Leonardo de Moura
 #include "frontends/lean/scanner.h"
 #include "frontends/lean/pp.h"
 #include "frontends/lean/server.h"
+#include "frontends/lean/local_ref_info.h"
 
 namespace lean {
 void initialize_frontend_lean_module() {
@@ -53,8 +54,10 @@ void initialize_frontend_lean_module() {
     initialize_calc_proof_elaborator();
     initialize_server();
     initialize_find_cmd();
+    initialize_local_ref_info();
 }
 void finalize_frontend_lean_module() {
+    finalize_local_ref_info();
     finalize_find_cmd();
     finalize_server();
     finalize_calc_proof_elaborator();
diff --git a/src/frontends/lean/local_ref_info.cpp b/src/frontends/lean/local_ref_info.cpp
new file mode 100644
index 0000000000..87a61f7376
--- /dev/null
+++ b/src/frontends/lean/local_ref_info.cpp
@@ -0,0 +1,79 @@
+/*
+Copyright (c) 2015 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+
+Author: Leonardo de Moura
+*/
+#include <string>
+#include "library/scoped_ext.h"
+
+namespace lean {
+/*
+When we declare a definition in a section, we create an alias for it that fixes the parameters in
+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;
+    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) {}
+};
+
+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;
+
+    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));
+    }
+    static name const & get_class_name() {
+        return *g_local_ref_name;
+    }
+    static std::string const & get_serialization_key() {
+        return *g_key;
+    }
+    static void  write_entry(serializer &, entry const &) {
+        lean_unreachable();
+    }
+    static entry read_entry(deserializer &) {
+        lean_unreachable();
+    }
+    static optional<unsigned> get_fingerprint(entry const &) {
+        return optional<unsigned>();
+    }
+};
+
+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) {
+    bool persistent = false;
+    return local_ref_ext::add_entry(env, get_dummy_ios(), local_ref_entry(n, num_fix_univs, num_fix_args), persistent);
+}
+
+optional<pair<unsigned, unsigned>> 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);
+    else
+        return optional<pair<unsigned, unsigned>>();
+}
+
+void initialize_local_ref_info() {
+    g_local_ref_name = new name("localrefinfo");
+    g_key            = new std::string("localrefinfo");
+    local_ref_ext::initialize();
+}
+
+void finalize_local_ref_info() {
+    local_ref_ext::finalize();
+    delete g_local_ref_name;
+    delete g_key;
+}
+}
diff --git a/src/frontends/lean/local_ref_info.h b/src/frontends/lean/local_ref_info.h
new file mode 100644
index 0000000000..c7441c3e3f
--- /dev/null
+++ b/src/frontends/lean/local_ref_info.h
@@ -0,0 +1,14 @@
+/*
+Copyright (c) 2015 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+
+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);
+void initialize_local_ref_info();
+void finalize_local_ref_info();
+}
diff --git a/src/frontends/lean/parser.cpp b/src/frontends/lean/parser.cpp
index 4500f94b99..63eb36a65f 100644
--- a/src/frontends/lean/parser.cpp
+++ b/src/frontends/lean/parser.cpp
@@ -47,6 +47,7 @@ Author: Leonardo de Moura
 #include "frontends/lean/info_annotation.h"
 #include "frontends/lean/parse_rewrite_tactic.h"
 #include "frontends/lean/update_environment_exception.h"
+#include "frontends/lean/local_ref_info.h"
 
 #ifndef LEAN_DEFAULT_PARSER_SHOW_ERRORS
 #define LEAN_DEFAULT_PARSER_SHOW_ERRORS true
@@ -494,6 +495,20 @@ void parser::add_local_expr(name const & n, expr const & p, bool is_variable) {
     }
 }
 
+environment parser::add_local_ref(environment const & env, name const & n, expr const & ref) {
+    add_local_expr(n, ref, false);
+    if (is_as_atomic(ref)) {
+        buffer<expr> args;
+        expr f = get_app_args(get_as_atomic_arg(ref), args);
+        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 env;
+}
+
 void parser::add_parameter(name const & n, expr const & p) {
     lean_assert(is_local(p));
     add_local_expr(n, p, false);
diff --git a/src/frontends/lean/parser.h b/src/frontends/lean/parser.h
index 3d2e33b746..31f95bee06 100644
--- a/src/frontends/lean/parser.h
+++ b/src/frontends/lean/parser.h
@@ -400,6 +400,7 @@ public:
     bool has_locals() const { return !m_local_decls.empty() || !m_local_level_decls.empty(); }
     void add_local_level(name const & n, level const & l, bool is_variable = false);
     void add_local_expr(name const & n, expr const & p, bool is_variable = false);
+    environment add_local_ref(environment const & env, name const & n, expr const & ref);
     void add_parameter(name const & n, expr const & p);
     void add_local(expr const & p) { return add_local_expr(local_pp_name(p), p); }
     bool has_params() const { return m_has_params; }
diff --git a/src/frontends/lean/pp.cpp b/src/frontends/lean/pp.cpp
index 9a598318a1..9fedaddf36 100644
--- a/src/frontends/lean/pp.cpp
+++ b/src/frontends/lean/pp.cpp
@@ -33,6 +33,7 @@ Author: Leonardo de Moura
 #include "frontends/lean/token_table.h"
 #include "frontends/lean/builtin_exprs.h"
 #include "frontends/lean/parser_config.h"
+#include "frontends/lean/local_ref_info.h"
 
 namespace lean {
 static format * g_ellipsis_n_fmt  = nullptr;
@@ -438,9 +439,19 @@ auto pretty_fn::pp_const(expr const & e) -> result {
             n = *n1;
     }
     if (m_universes && !empty(const_levels(e))) {
+        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)
+                return result(format(n));
+            else
+                first_idx = info->first;
+        }
         format r = compose(format(n), format(".{"));
         bool first = true;
-        for (auto const & l : const_levels(e)) {
+        for (unsigned i = first_idx; i < ls.size(); i++) {
+            level const & l = ls[i];
             format l_fmt = pp_level(l);
             if (is_max(l) || is_imax(l))
                 l_fmt = paren(l_fmt);
@@ -490,6 +501,15 @@ bool pretty_fn::has_implicit_args(expr const & f) {
 }
 
 auto pretty_fn::pp_app(expr const & e) -> result {
+    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)
+                return pp_const(rfn);
+        }
+    }
     expr const & fn = app_fn(e);
     if (auto it = is_abbreviated(fn))
         return pp_abbreviation(e, *it, true);
diff --git a/src/frontends/lean/type_util.cpp b/src/frontends/lean/type_util.cpp
index bd3f377d13..cd72949ee2 100644
--- a/src/frontends/lean/type_util.cpp
+++ b/src/frontends/lean/type_util.cpp
@@ -31,7 +31,7 @@ environment add_alias(parser & p, environment env, bool composite,
         id = name(full_id.get_string());
     if (!empty(ctx_levels) || !ctx_params.empty()) {
         expr r = mk_local_ref(full_id, ctx_levels, ctx_params);
-        p.add_local_expr(id, r);
+        env = p.add_local_ref(env, id, r);
     }
     if (full_id != id)
         env = add_expr_alias_rec(env, id, full_id);
diff --git a/tests/lean/errors.lean.expected.out b/tests/lean/errors.lean.expected.out
index 9d7966f81d..92616c4489 100644
--- a/tests/lean/errors.lean.expected.out
+++ b/tests/lean/errors.lean.expected.out
@@ -3,7 +3,7 @@ tst1 : nat → nat → nat
 errors.lean:12:2: error: function expected at
   tactic.cases [tactic.expr add] tactic.expr_list.nil
 errors.lean:22:12: error: unknown identifier 'b'
-tst3 A : A → A → A
+tst3 : A → A → A
 foo.tst1 : ℕ → ℕ → ℕ
 foo.tst2 : ℕ → ℕ → ℕ
 foo.tst3 : Π (A : Type), A → A → A
diff --git a/tests/lean/sec_param_pp.lean b/tests/lean/sec_param_pp.lean
new file mode 100644
index 0000000000..9e1a41360f
--- /dev/null
+++ b/tests/lean/sec_param_pp.lean
@@ -0,0 +1,20 @@
+section
+ parameters {A : Type} (a : A)
+ variable f : A → A → A
+
+ definition id : A := a
+ check id
+
+
+ definition pr (b : A) : A := f a b
+
+ check pr f id
+
+ set_option pp.universes true
+
+ check pr f id
+
+ definition pr2 (B : Type) (b : B) : A := a
+
+ check pr2 num 10
+end
diff --git a/tests/lean/sec_param_pp.lean.expected.out b/tests/lean/sec_param_pp.lean.expected.out
new file mode 100644
index 0000000000..817589c1e1
--- /dev/null
+++ b/tests/lean/sec_param_pp.lean.expected.out
@@ -0,0 +1,4 @@
+id : A
+pr f id : A
+pr f id : A
+pr2.{1} num 10 : A
diff --git a/tests/lean/sec_param_pp2.lean b/tests/lean/sec_param_pp2.lean
new file mode 100644
index 0000000000..71828df578
--- /dev/null
+++ b/tests/lean/sec_param_pp2.lean
@@ -0,0 +1,17 @@
+section
+ parameters {A : Type} (a : A)
+
+ section
+   parameters {B : Type} (b : B)
+
+   variable f : A → B → A
+
+   definition id := f a b
+
+   check id
+   set_option pp.universes true
+   check id
+ end
+ check id
+end
+check id
diff --git a/tests/lean/sec_param_pp2.lean.expected.out b/tests/lean/sec_param_pp2.lean.expected.out
new file mode 100644
index 0000000000..b98a28e5a5
--- /dev/null
+++ b/tests/lean/sec_param_pp2.lean.expected.out
@@ -0,0 +1,4 @@
+id : (A → B → A) → A
+id : (A → B → A) → A
+id.{l_2} : ?B a → (A → ?B a → A) → A
+id.{l_1 l_2} : ?A → (Π {B : Type.{l_2}}, B → (?A → B → ?A) → ?A)