Improve pretty printer for Pi's

Signed-off-by: Leonardo de Moura <leonardo@microsoft.com>

src/frontends/lean/lean_pp.cpp

@@ -686,14 +686,17 @@ class pp_fn {
         }
     }
 
-    result pp_scoped_child(expr const & e, unsigned depth) {
+    result pp_scoped_child(expr const & e, unsigned depth, unsigned prec = 0) {
         if (is_atomic(e)) {
             return pp(e, depth + 1, true);
         } else {
             mk_scope s(*this);
             result r = pp(e, depth + 1, true);
             if (m_aliases_defs.size() == s.m_old_size) {
+                if (prec <= get_operator_precedence(e))
                     return r;
+                else
+                    return mk_result(paren(r.first), r.second);
             } else {
                 format r_format   = g_let_fmt;
                 unsigned r_weight = 2;
@@ -715,12 +718,12 @@ class pp_fn {
         }
     }
 
-    result pp_arrow_body(expr const & e, unsigned depth) {
+    result pp_arrow_child(expr const & e, unsigned depth) {
-        if (is_atomic(e) || is_arrow(e)) {
+        return pp_scoped_child(e, depth, g_arrow_precedence + 1);
-            return pp(e, depth + 1);
-        } else {
-            return pp_child_with_paren(e, depth);
     }
+
+    result pp_arrow_body(expr const & e, unsigned depth) {
+        return pp_scoped_child(e, depth, g_arrow_precedence);
     }
 
     template<typename It>
@@ -738,6 +741,44 @@ class pp_fn {
         return implicit_args && (*implicit_args)[arg_pos];
     }
 
+    /**
+       \brief Auxiliary method for computing where Pi can be pretty printed as an arrow.
+       Examples:
+       Pi x : Int, Pi y : Int, Int        ===> return 0
+       Pi A : Type, Pi x : A, Pi y : A, A ===> return 1
+       Pi A : Type, Pi x : Int, A         ===> return 1
+       Pi A : Type, Pi x : Int, x > 0     ===> return UINT_MAX (there is no tail that can be printed as a arrow)
+
+       If \c e is not Pi, it returns UINT_MAX
+    */
+    unsigned get_arrow_starting_at(expr e) {
+        if (!is_pi(e))
+            return std::numeric_limits<unsigned>::max();
+        unsigned pos = 0;
+        while (is_pi(e)) {
+            expr e2 = abst_body(e);
+            unsigned num_vars = 1;
+            bool ok = true;
+            while (true) {
+                if (has_free_var(e2, 0, num_vars)) {
+                    ok = false;
+                    break;
+                }
+                if (is_pi(e2)) {
+                    e2 = abst_body(e2);
+                } else {
+                    break;
+                }
+            }
+            if (ok) {
+                return pos;
+            }
+            e = abst_body(e);
+            pos++;
+        }
+        return std::numeric_limits<unsigned>::max();
+    }
+
     /**
        \brief Pretty print Lambdas, Pis and compact definitions.
        When T != 0, it is a compact definition.
@@ -759,11 +800,12 @@ class pp_fn {
         local_names::mk_scope mk(m_local_names);
         if (is_arrow(e) && !implicit_args) {
             lean_assert(!T);
-            result p_lhs    = pp_child(abst_domain(e), depth);
+            result p_lhs    = pp_arrow_child(abst_domain(e), depth);
             result p_rhs    = pp_arrow_body(abst_body(e), depth);
             format r_format = group(format{p_lhs.first, space(), m_unicode ? g_arrow_n_fmt : g_arrow_fmt, line(), p_rhs.first});
             return mk_result(r_format, p_lhs.second + p_rhs.second + 1);
         } else {
+            unsigned arrow_starting_at = get_arrow_starting_at(e);
             buffer<std::pair<name, expr>> nested;
             auto p = collect_nested(e, T, e.kind(), nested);
             expr b = p.first;
@@ -814,6 +856,37 @@ class pp_fn {
                     ++it2;
                     bool implicit = is_implicit(implicit_args, arg_pos);
                     ++arg_pos;
+                    if (!implicit && !is_lambda(e) && arg_pos > arrow_starting_at) {
+                        // the rest is an arrow, but we must check if we are not missing implicit annotations.
+                        auto it2_aux = it2;
+                        unsigned arg_pos_aux = arg_pos;
+                        while (it2_aux != end && !is_implicit(implicit_args, arg_pos_aux)) {
+                            ++arg_pos_aux;
+                            ++it2_aux;
+                        }
+                        if (it2_aux == end) {
+                            // the rest is a sequence of arrows.
+                            format block;
+                            bool first_domain = true;
+                            for (; it != end; ++it) {
+                                result p_domain = pp_arrow_child(it->second, depth);
+                                r_weight += p_domain.second;
+                                if (first_domain) {
+                                    first_domain = false;
+                                    block = p_domain.first;
+                                } else {
+                                    block += format{space(), m_unicode ? g_arrow_n_fmt : g_arrow_fmt, line(), p_domain.first};
+                                }
+                            }
+                            result p_body = pp_arrow_child(b, depth);
+                            r_weight += p_body.second;
+                            block += format{space(), m_unicode ? g_arrow_n_fmt : g_arrow_fmt, line(), p_body.first};
+                            block = group(block);
+                            format r_format = group(nest(head_indent, format{head, space(), group(bindings), body_sep, line(), block}));
+                            return mk_result(r_format, r_weight);
+                        }
+                    }
+                    // Continue with standard encoding
                     while (it2 != end && it2->second == it->second && implicit == is_implicit(implicit_args, arg_pos)) {
                         ++it2;
                         ++arg_pos;

tests/lean/arrow.lean

@@ -0,0 +1,4 @@
+Show (Int -> Int) -> Int
+Show Int -> Int -> Int
+Show Int -> (Int -> Int)
+Show (Int -> Int) -> (Int -> Int) -> Int

tests/lean/arrow.lean.expected.out

@@ -0,0 +1,6 @@
+  Set: pp::colors
+  Set: pp::unicode
+(ℤ → ℤ) → ℤ
+ℤ → ℤ → ℤ
+ℤ → ℤ → ℤ
+(ℤ → ℤ) → (ℤ → ℤ) → ℤ

tests/lean/cast1.lean.expected.out

@@ -9,7 +9,7 @@
 Error (line: 12, pos: 6) type mismatch at application
     f m v1
 Function type:
-    Π (n : ℕ) (v : vector ℤ n), ℤ
+    Π (n : ℕ), (vector ℤ n) → ℤ
 Arguments types:
     m : ℕ
     v1 : vector ℤ (m + 0)

tests/lean/elab1.lean.expected.out

@@ -4,7 +4,7 @@
 Error (line: 4, pos: 6) type mismatch at application
     f 10 ⊤
 Function type:
-    Π (A : Type) (a b : A), A
+    Π (A : Type), A → A → A
 Arguments types:
     ℕ : Type
     10 : ℕ
@@ -16,7 +16,7 @@ Error (line: 7, pos: 6) unsolved placeholder at term
 Error (line: 11, pos: 27) application type mismatch during term elaboration
     h A x
 Function type:
-    Π (A : Type) (_ : A), A
+    Π (A : Type), A → A
 Arguments types:
     A : Type
     x : lift:0:2 ?M0
@@ -26,7 +26,7 @@ Elaborator state
 Error (line: 15, pos: 51) application type mismatch during term elaboration
     eq C a b
 Function type:
-    Π (A : Type) (_ _ : A), Bool
+    Π (A : Type), A → A → Bool
 Arguments types:
     C : Type
     a : lift:0:3 ?M0
@@ -52,7 +52,7 @@ Elaborator state
 Error (line: 22, pos: 22) type mismatch at application
     Trans (Refl a) (Refl b)
 Function type:
-    Π (A : Type U) (a b c : A) (H1 : a = b) (H2 : b = c), a = c
+    Π (A : Type U) (a b c : A), (a = b) → (b = c) → (a = c)
 Arguments types:
     Bool : Type
     a : Bool
@@ -63,7 +63,7 @@ Arguments types:
 Error (line: 24, pos: 6) type mismatch at application
     f Bool Bool
 Function type:
-    Π (A : Type) (a b : A), A
+    Π (A : Type), A → A → A
 Arguments types:
     Type : Type 1
     Bool : Type
@@ -71,7 +71,7 @@ Arguments types:
 Error (line: 27, pos: 21) type mismatch at application
     DisjCases (EM a) (λ H_a : a, H) (λ H_na : ¬ a, NotImp1 (MT H H_na))
 Function type:
-    Π (a b c : Bool) (H1 : a ∨ b) (H2 : a → c) (H3 : b → c), c
+    Π (a b c : Bool), (a ∨ b) → (a → c) → (b → c) → c
 Arguments types:
     a : Bool
     ¬ a : Bool

tests/lean/elab4.lean.expected.out

@@ -5,9 +5,9 @@
   Assumed: R
   Proved: R2
   Set: lean::pp::implicit
-Variable C {A B : Type} (H : A = B) (a : A) : B
+Variable C {A B : Type} : (A = B) → A → B
 Definition C::explicit (A B : Type) (H : A = B) (a : A) : B := C H a
-Variable D {A A' : Type} {B : A → Type} {B' : A' → Type} (H : (Π x : A, B x) = (Π x : A', B' x)) : A = A'
+Variable D {A A' : Type} {B : A → Type} {B' : A' → Type} : ((Π x : A, B x) = (Π x : A', B' x)) → (A = A')
 Definition D::explicit (A A' : Type) (B : A → Type) (B' : A' → Type) (H : (Π x : A, B x) = (Π x : A', B' x)) : A =
     A' :=
     D H

tests/lean/ex3.lean.expected.out

@@ -7,7 +7,7 @@ myeq Bool ⊤ ⊥
 Error (line: 5, pos: 6) type mismatch at application
     myeq Bool ⊤ a
 Function type:
-    Π (A : Type) (_ _ : A), Bool
+    Π (A : Type), A → A → Bool
 Arguments types:
     Bool : Type
     ⊤ : Bool
@@ -17,7 +17,7 @@ Arguments types:
 Error (line: 9, pos: 15) type mismatch at application
     myeq2::explicit Bool ⊤ a
 Function type:
-    Π (A : Type) (a b : A), Bool
+    Π (A : Type), A → A → Bool
 Arguments types:
     Bool : Type
     ⊤ : Bool

tests/lean/tst1.lean.expected.out

@@ -13,7 +13,7 @@
   Defined: select
   Defined: map
 Axiom two_lt_three : two < three
-Definition vector (A : Type) (n : N) : Type := Π (i : N) (H : i < n), A
+Definition vector (A : Type) (n : N) : Type := Π (i : N), (i < n) → A
 Definition const {A : Type} (n : N) (d : A) : vector A n := λ (i : N) (H : i < n), d
 Definition const::explicit (A : Type) (n : N) (d : A) : vector A n := const n d
 Definition update {A : Type} {n : N} (v : vector A n) (i : N) (d : A) : vector A n :=
@@ -30,20 +30,20 @@ Bool
 vector Bool three
 
 --------
-Π (A : Type) (n : N) (v : vector A n) (i : N) (H : i < n), A
+Π (A : Type) (n : N) (v : vector A n) (i : N), (i < n) → A
 
 map type ---> 
-Π (A B C : Type) (n : N) (f : A → B → C) (v1 : vector A n) (v2 : vector B n), vector C n
+Π (A B C : Type) (n : N), (A → B → C) → (vector A n) → (vector B n) → (vector C n)
 
 map normal form --> 
 λ (A B C : Type)
   (n : N)
   (f : A → B → C)
-  (v1 : Π (i : N) (H : i < n), A)
+  (v1 : Π (i : N), (i < n) → A)
-  (v2 : Π (i : N) (H : i < n), B)
+  (v2 : Π (i : N), (i < n) → B)
   (i : N)
   (H : i < n),
   f (v1 i H) (v2 i H)
 
 update normal form --> 
-λ (A : Type) (n : N) (v : Π (i : N) (H : i < n), A) (i : N) (d : A) (j : N) (H : j < n), if A (j = i) d (v j H)
+λ (A : Type) (n : N) (v : Π (i : N), (i < n) → A) (i : N) (d : A) (j : N) (H : j < n), if A (j = i) d (v j H)

tests/lean/tst10.lean.expected.out

@@ -19,6 +19,6 @@ if Bool a a ⊤
 a
   Assumed: H1
   Assumed: H2
-Π (a b : Bool) (H1 : a ⇒ b) (H2 : a), b
+Π (a b : Bool), (a ⇒ b) → a → b
 MP H2 H1
 b

tests/lean/tst11.lean.expected.out

@@ -6,7 +6,7 @@
 ⊤
   Assumed: a
 a ⊕ a ⊕ a
-Π (A : Type U) (a b : A) (P : A → Bool) (H1 : P a) (H2 : a = b), P b
+Π (A : Type U) (a b : A) (P : A → Bool), (P a) → (a = b) → (P b)
   Proved: EM2
 Π a : Bool, a ∨ ¬ a
 a ∨ ¬ a

tests/lean/tst4.lean.expected.out

@@ -10,7 +10,7 @@ f::explicit ((N → N) → N → N) (λ x : N → N, x) (λ y : N → N, y)
   Assumed: EqNice
 EqNice::explicit N n1 n2
 N
-Π (A : Type U) (B : A → Type U) (f g : Π x : A, B x) (a b : A) (H1 : f = g) (H2 : a = b), (f a) = (g b)
+Π (A : Type U) (B : A → Type U) (f g : Π x : A, B x) (a b : A), (f = g) → (a = b) → ((f a) = (g b))
 f::explicit N n1 n2
   Assumed: a
   Assumed: b

tests/lean/tst7.lean.expected.out

@@ -5,15 +5,15 @@
 Error (line: 4, pos: 40) application type mismatch during term elaboration
     f B a
 Function type:
-    Π (A : Type) (_ : A), Bool
+    Π (A : Type), A → Bool
 Arguments types:
     B : Type
     a : lift:0:2 ?M0
 Elaborator state
     #0 ≈ lift:0:2 ?M0
   Assumed: myeq
-myeq (Π (A : Type) (a : A), A) (λ (A : Type) (a : A), a) (λ (B : Type) (b : B), b)
+myeq (Π (A : Type), A → A) (λ (A : Type) (a : A), a) (λ (B : Type) (b : B), b)
 Bool
   Assumed: R
   Assumed: h
-Bool → (Π (f1 g1 : Π A : Type, R A) (G : Π A : Type, myeq (R A) (f1 A) (g1 A)), Bool)
+Bool → (Π (f1 g1 : Π A : Type, R A), (Π A : Type, myeq (R A) (f1 A) (g1 A)) → Bool)

tests/lean/tst8.lean.expected.out

@@ -1,6 +1,6 @@
   Set: pp::colors
   Set: pp::unicode
-Π (A : Type) (a : A), A
+Π (A : Type), A → A
   Assumed: g
   Defined: f
 f ℕ 10