fix(frontends/lean/pp): pretty printing coercions to functions, fixes #151

src/frontends/lean/pp.cpp

@@ -151,18 +151,35 @@ bool pretty_fn::is_prop(expr const & e) {
     }
 }
 
+auto pretty_fn::pp_coercion(expr const & e, unsigned bp) -> result {
+    buffer<expr> args;
+    expr const & f = get_app_args(e, args);
+    optional<pair<name, unsigned>> r = is_coercion(m_env, f);
+    lean_assert(r);
+    if (r->second >= args.size())
+        return pp_child_core(e, bp);
+    else if (r->second == args.size() - 1)
+        return pp_child(args.back(), bp);
+    else
+        return pp_child(mk_app(args.size() - r->second, args.data() + r->second), bp);
+}
+
+auto pretty_fn::pp_child_core(expr const & e, unsigned bp) -> result {
+    result r = pp(e);
+    if (r.second < bp) {
+        return mk_result(paren(r.first));
+    } else {
+        return r;
+    }
+}
+
 auto pretty_fn::pp_child(expr const & e, unsigned bp) -> result {
     if (is_app(e) && is_implicit(app_fn(e))) {
         return pp_child(app_fn(e), bp);
     } else if (is_app(e) && !m_coercion && is_coercion(m_env, get_app_fn(e))) {
-        return pp_child(app_arg(e), bp); // TODO(Fix): this is not correct for coercions to function-class
+        return pp_coercion(e, bp);
     } else {
-        result r = pp(e);
-        if (r.second < bp) {
-            return mk_result(paren(r.first));
-        } else {
-            return r;
-        }
+        return pp_child_core(e, bp);
     }
 }
 

src/frontends/lean/pp.h

@@ -58,6 +58,8 @@ private:
     format pp_binders(buffer<expr> const & locals);
     format pp_level(level const & l);
 
+    result pp_coercion(expr const & e, unsigned bp);
+    result pp_child_core(expr const & e, unsigned bp);
     result pp_child(expr const & e, unsigned bp);
     result pp_var(expr const & e);
     result pp_sort(expr const & e);

src/library/coercion.cpp

@@ -71,7 +71,7 @@ struct coercion_state {
     // m_from and m_to contain "direct" coercions
     rb_map<name, list<pair<coercion_class, expr>>, name_quick_cmp>  m_from; // map user-class -> list of (class, coercion-fun)
     rb_map<coercion_class, list<name>, coercion_class_cmp_fn>            m_to;
-    rb_tree<name, name_quick_cmp>                                        m_coercions;
+    rb_map<name, pair<name, unsigned>, name_quick_cmp>                   m_coercions; // map coercion -> (from-class, num-args)
     coercion_info get_info(name const & from, coercion_class const & to) {
         auto it = m_coercion_info.find(from);
         lean_assert(it);
@@ -254,7 +254,7 @@ struct add_coercion_fn {
             m_state.m_coercion_info.insert(C, infos);
         }
         if (is_constant(f))
-            m_state.m_coercions.insert(const_name(f));
+            m_state.m_coercions.insert(const_name(f), mk_pair(C, num_args));
     }
 
     void add_coercion(name const & C, expr const & f, expr const & f_type,
@@ -352,13 +352,18 @@ environment add_coercion(environment const & env, name const & f, io_state const
     return add_coercion(env, f, const_name(C), ios);
 }
 
-bool is_coercion(environment const & env, name const & f) {
+optional<pair<name, unsigned>> is_coercion(environment const & env, name const & f) {
     coercion_state const & ext = coercion_ext::get_state(env);
-    return ext.m_coercions.contains(f);
+    if (auto it = ext.m_coercions.find(f))
+        return optional<pair<name, unsigned>>(*it);
+    else
+        return optional<pair<name, unsigned>>();
 }
 
-bool is_coercion(environment const & env, expr const & f) {
-    return is_constant(f) && is_coercion(env, const_name(f));
+optional<pair<name, unsigned>> is_coercion(environment const & env, expr const & f) {
+    if (!is_constant(f))
+        return optional<pair<name, unsigned>>();
+    return is_coercion(env, const_name(f));
 }
 
 bool has_coercions_to(environment const & env, name const & D) {
@@ -485,10 +490,18 @@ static int add_coercion(lua_State * L) {
         return push_environment(L, add_coercion(to_environment(L, 1), to_name_ext(L, 2), to_name_ext(L, 3), to_io_state(L, 4)));
 }
 static int is_coercion(lua_State * L) {
+    optional<pair<name, unsigned>> r;
     if (is_expr(L, 2))
-        return push_boolean(L, is_coercion(to_environment(L, 1), to_expr(L, 2)));
+        r = is_coercion(to_environment(L, 1), to_expr(L, 2));
     else
-        return push_boolean(L, is_coercion(to_environment(L, 1), to_name_ext(L, 2)));
+        r = is_coercion(to_environment(L, 1), to_name_ext(L, 2));
+    if (r) {
+        push_name(L, r->first);
+        push_integer(L, r->second);
+        return 2;
+    } else {
+        return 0;
+    }
 }
 static int has_coercions_from(lua_State * L) {
     if (is_expr(L, 2))

src/library/coercion.h

@@ -38,8 +38,11 @@ namespace lean {
 */
 environment add_coercion(environment const & env, name const & f, io_state const & ios);
 environment add_coercion(environment const & env, name const & f, name const & C, io_state const & ios);
-bool is_coercion(environment const & env, name const & f);
-bool is_coercion(environment const & env, expr const & f);
+/** \brief If \c f is a coercion, then return the name of the 'from-class' and the number of
+    class parameters.
+*/
+optional<pair<name, unsigned>> is_coercion(environment const & env, name const & f);
+optional<pair<name, unsigned>> is_coercion(environment const & env, expr const & f);
 /** \brief Return true iff the given environment has coercions from a user-class named \c C. */
 bool has_coercions_from(environment const & env, name const & C);
 bool has_coercions_from(environment const & env, expr const & C);

tests/lean/interactive/coe.input

@@ -0,0 +1,3 @@
+VISIT coe.lean
+WAIT
+INFO 5 33

tests/lean/interactive/coe.input.expected.out

@@ -0,0 +1,13 @@
+-- BEGINWAIT
+-- ENDWAIT
+-- BEGININFO
+-- TYPE|5|33
+decidable (eq a b)
+-- ACK
+-- SYNTH|5|33
+int.has_decidable_eq a b
+-- ACK
+-- SYMBOL|5|33
+_
+-- ACK
+-- ENDINFO

tests/lean/interactive/coe.lean

@@ -0,0 +1,9 @@
+import data.int
+open int
+
+theorem has_decidable_eq [instance] [protected] : decidable_eq ℤ :=
+decidable_eq.intro (λ (a b : ℤ), _)
+
+variable n : nat
+variable i : int
+check n + i