From 5952a857cd70e7bd71d7b60cf0784bbeb76d7a84 Mon Sep 17 00:00:00 2001
From: Daniel Selsam <daniel.selsam@protonmail.com>
Date: Wed, 11 Aug 2021 12:56:58 -0700
Subject: [PATCH] feat: pp.analyze improve heuristics for fun binders

---
 .../Delaborator/TopDownAnalyze.lean           | 34 +++++++++++++++++--
 tests/lean/run/PPTopDownAnalyze.lean          | 22 ++++++++++++
 2 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/src/Lean/PrettyPrinter/Delaborator/TopDownAnalyze.lean b/src/Lean/PrettyPrinter/Delaborator/TopDownAnalyze.lean
index 54371e5167..1495c5ab19 100644
--- a/src/Lean/PrettyPrinter/Delaborator/TopDownAnalyze.lean
+++ b/src/Lean/PrettyPrinter/Delaborator/TopDownAnalyze.lean
@@ -334,6 +334,7 @@ structure App.Context where
 structure App.State where
   bottomUps       : Array Bool
   higherOrders    : Array Bool
+  funBinders      : Array Bool
   provideds       : Array Bool
   namedArgs       : Array Name := #[]
 
@@ -397,7 +398,8 @@ mutual
       analyzeAppStagedCore ⟨f, fType, args, mvars, bInfos, forceRegularApp⟩ |>.run' {
         bottomUps    := mkArray args.size false,
         higherOrders := mkArray args.size false,
-        provideds    := mkArray args.size false
+        provideds    := mkArray args.size false,
+        funBinders   := mkArray args.size false
       }
 
       if not rest.isEmpty then
@@ -453,6 +455,7 @@ mutual
     hBinOpHeuristic
     collectTrivialBottomUps
     discard <| processPostponed (mayPostpone := true)
+    applyFunBinderHeuristic
     analyzeFn
     for i in [:(← read).args.size] do analyzeArg i
     maybeSetExplicit
@@ -500,6 +503,31 @@ mutual
             tryUnify args[i] mvars[i]
             modify fun s => { s with bottomUps := s.bottomUps.set! i true }
 
+    applyFunBinderHeuristic := do
+      let ⟨f, _, args, mvars, bInfos, _⟩ ← read
+
+      let rec core (argIdx : Nat) (mvarType : Expr) : AnalyzeAppM Bool := do
+        match ← getExpr, mvarType with
+        | Expr.lam _ _ _ .., Expr.forallE n t b .. =>
+          let mut annotated := false
+          for i in [:argIdx] do
+            if ← bInfos[i] == BinderInfo.implicit <&&> valUnknown mvars[i] <&&> withNewMCtxDepth (checkpointDefEq t mvars[i]) then
+              annotateBool `pp.funBinderTypes
+              tryUnify args[i] mvars[i]
+              -- Note: currently we always analyze the lambda binding domains in `analyzeLam`
+              -- (so we don't need to analyze it again here)
+              annotated := true
+              break
+          let annotatedBody ← withBindingBody Name.anonymous (core argIdx b)
+          return annotated || annotatedBody
+
+        | _, _ => return false
+
+      for i in [:args.size] do
+        if ← bInfos[i] == BinderInfo.default then
+          let b ← withNaryArg i (core i (← inferType mvars[i]))
+          if b then modify fun s => { s with funBinders := s.funBinders.set! i true }
+
     analyzeFn := do
       -- Now, if this is the first staging, analyze the n-ary function without expected type
       let ⟨f, fType, _, _, _, forceRegularApp⟩ ← read
@@ -511,7 +539,7 @@ mutual
 
     analyzeArg (i : Nat) := do
       let ⟨f, _, args, mvars, bInfos, forceRegularApp⟩ ← read
-      let ⟨bottomUps, higherOrders,_, _⟩ ← get
+      let ⟨bottomUps, higherOrders, funBinders, _, _⟩ ← get
       let arg := args[i]
       let argType ← inferType arg
 
@@ -527,7 +555,7 @@ mutual
 
           match bInfos[i] with
           | BinderInfo.default =>
-            if ← !(← valUnknown mvars[i]) <&&> !(← readThe Context).inBottomUp <&&> !(← isFunLike arg) <&&> checkpointDefEq mvars[i] arg then
+            if ← !(← valUnknown mvars[i]) <&&> !(← readThe Context).inBottomUp <&&> !(← isFunLike arg) <&&> !funBinders[i] <&&> checkpointDefEq mvars[i] arg then
               annotateBool `pp.analysis.hole
             else
               modify fun s => { s with provideds := s.provideds.set! i true }
diff --git a/tests/lean/run/PPTopDownAnalyze.lean b/tests/lean/run/PPTopDownAnalyze.lean
index 7b2caf87a7..99a7eea13b 100644
--- a/tests/lean/run/PPTopDownAnalyze.lean
+++ b/tests/lean/run/PPTopDownAnalyze.lean
@@ -278,6 +278,28 @@ set_option pp.proofs.withType false in
     let ctxCore ← readThe Core.Context
     pure ctxCore.currNamespace
 
+structure SubtypeLike1 {α : Sort u} (p : α → Prop) where
+
+#testDelab SubtypeLike1 fun (x : Nat) => x < 10
+  expecting SubtypeLike1 fun (x : Nat) => x < 10
+
+#eval "prevent comment from parsing as part of previous expression"
+
+--Note: currently we do not try "bottom-upping" inside lambdas
+--(so we will always annotate the binder type)
+#testDelab SubtypeLike1 fun (x : Nat) => Nat.succ x = x
+  expecting SubtypeLike1 fun (x : Nat) => Nat.succ x = x
+
+structure SubtypeLike3 {α β γ : Sort u} (p : α → β → γ → Prop) where
+
+#testDelab SubtypeLike3 fun (x y z : Nat) => x + y < z
+  expecting SubtypeLike3 fun (x y z : Nat) => x + y < z
+
+structure SubtypeLike3Double {α β γ : Sort u} (p₁ : α → β → Prop) (p₂ : β → γ → Prop) where
+
+#testDelab SubtypeLike3Double (fun (x y : Nat) => x = y) (fun (y z : Nat) => y = z)
+  expecting SubtypeLike3Double (fun (x y : Nat) => x = y) fun y (z : Nat) => y = z
+
 def takesStricts ⦃α : Type⦄ {β : Type} ⦃γ : Type⦄ : Unit := ()
 #testDelab takesStricts expecting takesStricts
 #testDelab @takesStricts expecting takesStricts