From 2e3e15fadeb2a6e4743d3927b9120ff6195ac64a Mon Sep 17 00:00:00 2001
From: Kyle Miller <kmill31415@gmail.com>
Date: Wed, 23 Oct 2024 18:50:53 -0700
Subject: [PATCH] feat: preserve binder names in signatures

Part of the interface of a function are the argument names. This makes it so that names after the colon become hygienic if they are not able to be used as named arguments.

Closes #5810
---
 src/Lean/PrettyPrinter/Delaborator/Basic.lean |  9 ++++-
 .../PrettyPrinter/Delaborator/Builtins.lean   | 38 ++++++++++++++++---
 .../PrettyPrinter/Delaborator/SubExpr.lean    |  8 ++++
 ...toImplicitChainNameIssue.lean.expected.out |  2 +-
 tests/lean/match1.lean.expected.out           |  2 +-
 tests/lean/run/2846.lean                      | 12 +++---
 6 files changed, 56 insertions(+), 15 deletions(-)

diff --git a/src/Lean/PrettyPrinter/Delaborator/Basic.lean b/src/Lean/PrettyPrinter/Delaborator/Basic.lean
index bf0de8759ebd..5409e373a705 100644
--- a/src/Lean/PrettyPrinter/Delaborator/Basic.lean
+++ b/src/Lean/PrettyPrinter/Delaborator/Basic.lean
@@ -262,9 +262,14 @@ def mkAnnotatedIdent (n : Name) (e : Expr) : DelabM Ident := do
 Enters the body of the current expression, which must be a lambda or forall.
 The binding variable is passed to `d` as `Syntax`, and it is an identifier that has been annotated with the fvar expression
 for the variable.
+
+If `preserveName` is `false` (the default), then gives the binder an unused name.
+Otherwise, it preserves the binder name (contrary to the name of this function).
 -/
-def withBindingBodyUnusedName {α} (d : Syntax → DelabM α) : DelabM α := do
-  let n ← getUnusedName (← getExpr).bindingName! (← getExpr).bindingBody!
+def withBindingBodyUnusedName {α} (d : Syntax → DelabM α) (preserveName := false) : DelabM α := do
+  let mut n := (← getExpr).bindingName!
+  unless preserveName do
+    n ← getUnusedName n (← getExpr).bindingBody!
   withBindingBody' n (mkAnnotatedIdent n) (d ·)
 
 inductive OmissionReason
diff --git a/src/Lean/PrettyPrinter/Delaborator/Builtins.lean b/src/Lean/PrettyPrinter/Delaborator/Builtins.lean
index aa13f8f2c5cb..0a39b839185c 100644
--- a/src/Lean/PrettyPrinter/Delaborator/Builtins.lean
+++ b/src/Lean/PrettyPrinter/Delaborator/Builtins.lean
@@ -887,6 +887,11 @@ def delabLam : Delab :=
       else
         `(fun $binders* => $stxBody)
 
+/-- Don't do any renaming for forall binders. -/
+def ppPiPreserveNames := `pp.piPreserveNames
+/-- Causes non-dependent foralls to print with binder names. -/
+def ppPiBinderNames := `pp.piBinderNames
+
 /--
 Similar to `delabBinders`, but tracking whether `forallE` is dependent or not.
 
@@ -898,13 +903,14 @@ private partial def delabForallBinders (delabGroup : Array Syntax → Bool → S
     -- don't group
     delabGroup curNames curDep (← delab)
   else
+    let preserve := (← getOptionsAtCurrPos).get ppPiPreserveNames false
     let curDep := dep
     if ← shouldGroupWithNext then
       -- group with nested binder => recurse immediately
-      withBindingBodyUnusedName fun stxN => delabForallBinders delabGroup (curNames.push stxN) curDep
+      withBindingBodyUnusedName (preserveName := preserve) fun stxN => delabForallBinders delabGroup (curNames.push stxN) curDep
     else
       -- don't group => delab body and prepend current binder group
-      let (stx, stxN) ← withBindingBodyUnusedName fun stxN => return (← delab, stxN)
+      let (stx, stxN) ← withBindingBodyUnusedName (preserveName := preserve) fun stxN => return (← delab, stxN)
       delabGroup (curNames.push stxN) curDep stx
 
 @[builtin_delab forallE]
@@ -920,7 +926,7 @@ def delabForall : Delab := do
     | BinderInfo.instImplicit   => `(bracketedBinderF|[$curNames.back : $stxT])
     | _                         =>
       -- NOTE: non-dependent arrows are available only for the default binder info
-      if dependent then
+      if dependent || (← getOptionsAtCurrPos).get ppPiBinderNames false then
         if prop && !(← getPPOption getPPPiBinderTypes) then
           return ← `(∀ $curNames:ident*, $stxBody)
         else
@@ -1276,8 +1282,10 @@ where
     else if e.isForall && !e.bindingName!.hasMacroScopes && !bindingNames.contains e.bindingName! then
       delabParamsAux bindingNames idStx groups #[]
     else
-      let type ← delabTy
-      `(declSigWithId| $idStx:ident $groups* : $type)
+      let (opts', e') ← processSpine {} (← readThe SubExpr)
+      withReader (fun ctx => {ctx with optionsPerPos := opts', subExpr := { ctx.subExpr with expr := e' }}) do
+        let type ← delabTy
+        `(declSigWithId| $idStx:ident $groups* : $type)
   /--
   Inner loop for `delabParams`, collecting binders.
   Invariants:
@@ -1320,5 +1328,25 @@ where
     e.bindingDomain! == e'.bindingDomain! &&
     -- Inst implicits can't be grouped:
     e'.binderInfo != BinderInfo.instImplicit
+  /--
+  Go through rest of type, alpha renaming and setting options along the spine.
+  -/
+  processSpine (opts : OptionsPerPos) (subExpr : SubExpr) : MetaM (OptionsPerPos × Expr) := do
+    if let .forallE n t b bi := subExpr.expr then
+      let used := (← getLCtx).usesUserName n
+      withLocalDecl n bi t fun fvar => do
+        let (opts, b') ← processSpine opts { expr := b.instantiate1 fvar, pos := subExpr.pos.pushBindingBody }
+        let b' := b'.abstract #[fvar]
+        let opts := opts.insertAt subExpr.pos ppPiPreserveNames true
+        if n.hasMacroScopes then
+          return (opts, .forallE n t b' bi)
+        else if !used then
+          let opts := opts.insertAt subExpr.pos ppPiBinderNames true
+          return (opts, .forallE n t b' bi)
+        else
+          let n' ← withFreshMacroScope <| MonadQuotation.addMacroScope n
+          return (opts, .forallE n' t b' bi)
+    else
+      return (opts, subExpr.expr)
 
 end Lean.PrettyPrinter.Delaborator
diff --git a/src/Lean/PrettyPrinter/Delaborator/SubExpr.lean b/src/Lean/PrettyPrinter/Delaborator/SubExpr.lean
index dfa8311f3ed0..30909660e47c 100644
--- a/src/Lean/PrettyPrinter/Delaborator/SubExpr.lean
+++ b/src/Lean/PrettyPrinter/Delaborator/SubExpr.lean
@@ -18,6 +18,14 @@ namespace Lean.PrettyPrinter.Delaborator
 
 abbrev OptionsPerPos := RBMap SubExpr.Pos Options compare
 
+def OptionsPerPos.insertAt (optionsPerPos : OptionsPerPos) (pos : SubExpr.Pos) (name : Name) (value : DataValue) : OptionsPerPos :=
+  let opts := optionsPerPos.find? pos |>.getD {}
+  optionsPerPos.insert pos <| opts.insert name value
+
+/-- Merges two collections of options, where the second overrides the first. -/
+def OptionsPerPos.merge : OptionsPerPos → OptionsPerPos → OptionsPerPos :=
+  RBMap.mergeBy (fun _ => KVMap.mergeBy (fun _ _ dv => dv))
+
 namespace SubExpr
 
 open Lean.SubExpr
diff --git a/tests/lean/autoImplicitChainNameIssue.lean.expected.out b/tests/lean/autoImplicitChainNameIssue.lean.expected.out
index 749246019cfd..f5430942390a 100644
--- a/tests/lean/autoImplicitChainNameIssue.lean.expected.out
+++ b/tests/lean/autoImplicitChainNameIssue.lean.expected.out
@@ -3,4 +3,4 @@ case nil
 α✝ : Type u_1
 as : List α✝
 ⊢ Palindrome [].reverse
-palindrome_reverse.{u_1} : ∀ {α : Type u_1} {as : List α}, Palindrome as → Palindrome as.reverse
+palindrome_reverse.{u_1} : ∀ {α✝ : Type u_1} {as : List α✝} (h : Palindrome as), Palindrome as.reverse
diff --git a/tests/lean/match1.lean.expected.out b/tests/lean/match1.lean.expected.out
index 14f9bcfa50e7..185733340dab 100644
--- a/tests/lean/match1.lean.expected.out
+++ b/tests/lean/match1.lean.expected.out
@@ -47,5 +47,5 @@ fun x =>
   | #[3, 4, 5] => 3
   | x => 4 : Array Nat → Nat
 g.match_1.{u_1, u_2} {α : Type u_1} (motive : List α → Sort u_2) :
-  (x : List α) → ((a : α) → motive [a]) → ((x : List α) → motive x) → motive x
+  (x✝ : List α) → (h_1 : (a : α) → motive [a]) → (h_2 : (x : List α) → motive x) → motive x✝
 fun e => nomatch e : Empty → False
diff --git a/tests/lean/run/2846.lean b/tests/lean/run/2846.lean
index 88dcf2f153f9..aabdbc357c48 100644
--- a/tests/lean/run/2846.lean
+++ b/tests/lean/run/2846.lean
@@ -30,20 +30,20 @@ def Nat.pos_pow_of_pos' {n : Nat} (m : Nat) : 0 < n → 0 < n ^ m := Nat.pos_pow
 
 /-!
 Repetition of a named argument, only the first is printed as a named argument.
+The second is made hygienic.
 -/
 def foo (n n : Nat) : Fin (n + 1) := 0
 
-/-- info: foo (n : Nat) : (n : Nat) → Fin (n + 1) -/
+/-- info: foo (n : Nat) : (n✝ : Nat) → Fin (n✝ + 1) -/
 #guard_msgs in #check foo
 
 /-!
-Repetition of a named argument, only the first is printed as a named argument.
-This is checking that shadowing is handled correctly (that's not the responsibility of
-`delabConstWithSignature`, but it assumes that the main delaborator will handle this correctly).
+Same, but a named argument still follows, and its name is preserved.
 -/
+
 def foo' (n n : Nat) (a : Fin ((by clear n; exact n) + 1)) : Fin (n + 1) := 0
 
-/-- info: foo' (n : Nat) : (n_1 : Nat) → Fin (n + 1) → Fin (n_1 + 1) -/
+/-- info: foo' (n : Nat) : (n✝ : Nat) → (a : Fin (n + 1)) → Fin (n✝ + 1) -/
 #guard_msgs in #check foo'
 
 /-!
@@ -52,7 +52,7 @@ But, it does not print using named pi notation since this is not a dependent typ
 -/
 def foo'' : String → (needle : String) → String := fun _ yo => yo
 
-/-- info: foo'' : String → String → String -/
+/-- info: foo'' : (a✝needle : String) → String -/
 #guard_msgs in #check foo''
 
 /-!