From eed0fb6635cf7543842e88f1c7aacfe241e28a1f Mon Sep 17 00:00:00 2001
From: Daniel Selsam <daniel.selsam@protonmail.com>
Date: Wed, 28 Jul 2021 20:22:38 -0700
Subject: [PATCH] feat: special support for 'fun x => x'

---
 .../PrettyPrinter/Delaborator/TopDownAnalyze.lean  | 14 ++++++++++++++
 tests/lean/run/PPTopDownAnalyze.lean               |  8 ++++++++
 2 files changed, 22 insertions(+)

diff --git a/src/Lean/PrettyPrinter/Delaborator/TopDownAnalyze.lean b/src/Lean/PrettyPrinter/Delaborator/TopDownAnalyze.lean
index 0e336ffe47..24d7d8245b 100644
--- a/src/Lean/PrettyPrinter/Delaborator/TopDownAnalyze.lean
+++ b/src/Lean/PrettyPrinter/Delaborator/TopDownAnalyze.lean
@@ -54,6 +54,12 @@ register_builtin_option pp.analyze.trustOfNat : Bool := {
   descr    := "(pretty printer analyzer) always 'pretend' `OfNat.ofNat` applications can elab bottom-up"
 }
 
+register_builtin_option pp.analyze.trustId : Bool := {
+  defValue := true
+  group    := "pp.analyze"
+  descr    := "(pretty printer analyzer) always assume an implicit `fun x => x` can be inferred"
+}
+
 register_builtin_option pp.analyze.trustKnownFOType2TypeHOFuns : Bool := {
   defValue := true
   group    := "pp.analyze"
@@ -65,6 +71,7 @@ def getPPAnalyzeCheckInstances              (o : Options) : Bool := o.get pp.ana
 def getPPAnalyzeTypeAscriptions             (o : Options) : Bool := o.get pp.analyze.typeAscriptions.name pp.analyze.typeAscriptions.defValue
 def getPPAnalyzeTrustSubst                  (o : Options) : Bool := o.get pp.analyze.trustSubst.name pp.analyze.trustSubst.defValue
 def getPPAnalyzeTrustOfNat                  (o : Options) : Bool := o.get pp.analyze.trustOfNat.name pp.analyze.trustOfNat.defValue
+def getPPAnalyzeTrustId                     (o : Options) : Bool := o.get pp.analyze.trustId.name pp.analyze.trustId.defValue
 def getPPAnalyzeTrustKnownFOType2TypeHOFuns (o : Options) : Bool := o.get pp.analyze.trustKnownFOType2TypeHOFuns.name pp.analyze.trustKnownFOType2TypeHOFuns.defValue
 
 def getPPAnalysisSkip          (o : Options) : Bool := o.get `pp.analysis.skip false
@@ -96,6 +103,12 @@ def isFOLike (motive : Expr) : MetaM Bool := do
   let f := motive.getAppFn
   f.isFVar || f.isConst
 
+def isIdLike (arg : Expr) : Bool := do
+  -- TODO: allow `id` constant as well?
+  match arg with
+  | Expr.lam _ _ (Expr.bvar ..) .. => true
+  | _ => false
+
 namespace TopDownAnalyze
 
 def replaceLPsWithVars (e : Expr) : MetaM Expr := do
@@ -290,6 +303,7 @@ where
     for i in [:args.size] do
       if not (← bInfos[i] == BinderInfo.implicit) then continue
       if not (← isHigherOrder (← inferType args[i])) then continue
+      if getPPAnalyzeTrustId (← getOptions) && isIdLike args[i] then continue
 
       if getPPAnalyzeTrustKnownFOType2TypeHOFuns (← getOptions) && not (← valUnknown mvars[i])
         && (← isType2Type (args[i])) && (← isFOLike (args[i])) then continue
diff --git a/tests/lean/run/PPTopDownAnalyze.lean b/tests/lean/run/PPTopDownAnalyze.lean
index 8ddde3bed1..15c4a5f20f 100644
--- a/tests/lean/run/PPTopDownAnalyze.lean
+++ b/tests/lean/run/PPTopDownAnalyze.lean
@@ -185,6 +185,14 @@ set_option pp.analyze.trustSubst true in
 #testDelab (fun (x y z : Nat) (hxy : x = y) (hyz : x = z) => hxy ▸ hyz : ∀ (x y z : Nat), x = y → x = z → y = z)
   expecting fun x y z hxy hyz => hxy ▸ hyz
 
+set_option pp.analyze.trustId true in
+#testDelab Sigma.mk (β := fun α => α) Bool true
+  expecting { fst := Bool, snd := true : Sigma fun α => α }
+
+set_option pp.analyze.trustId false in
+#testDelab Sigma.mk (β := fun α => α) Bool true
+  expecting Sigma.mk (β := fun α => α) Bool true
+
 #testDelab let xs := #[true]; xs
   expecting let xs := #[true]; xs