fix: extended coe notation and delaborator (#3295)

src/Lean/Elab/BuiltinNotation.lean

@@ -22,14 +22,14 @@ open Meta
   ensureHasType expectedType? e
 
 @[builtin_term_elab coeFunNotation] def elabCoeFunNotation : TermElab := fun stx _ => do
-  let x ← elabTerm stx none
+  let x ← elabTerm stx[1] none
   if let some ty ← coerceToFunction? x then
     return ty
   else
     throwError "cannot coerce to function{indentExpr x}"
 
 @[builtin_term_elab coeSortNotation] def elabCoeSortNotation : TermElab := fun stx _ => do
-  let x ← elabTerm stx none
+  let x ← elabTerm stx[1] none
   if let some ty ← coerceToSort? x then
     return ty
   else

src/Lean/PrettyPrinter/Delaborator/Builtins.lean

@@ -341,27 +341,6 @@ def withOverApp (arity : Nat) (x : Delab) : Delab := do
   else
     delabAppCore (n - arity) x
 
-/--
-This delaborator tries to elide functions which are known coercions.
-For example, `Int.ofNat` is a coercion, so instead of printing `ofNat n` we just print `↑n`,
-and when re-parsing this we can (usually) recover the specific coercion being used.
--/
-@[builtin_delab app]
-def coeDelaborator : Delab := whenPPOption getPPCoercions do
-  let e ← getExpr
-  let .const declName _ := e.getAppFn | failure
-  let some info ← Meta.getCoeFnInfo? declName | failure
-  let n := e.getAppNumArgs
-  withOverApp info.numArgs do
-    match info.type with
-    | .coe => `(↑$(← withNaryArg info.coercee delab))
-    | .coeFun =>
-      if n = info.numArgs then
-        `(⇑$(← withNaryArg info.coercee delab))
-      else
-        withNaryArg info.coercee delab
-    | .coeSort => `(↥$(← withNaryArg info.coercee delab))
-
 /-- State for `delabAppMatch` and helpers. -/
 structure AppMatchState where
   info        : MatcherInfo
@@ -832,6 +811,27 @@ def delabProjectionApp : Delab := whenPPOption getPPStructureProjections $ do
     let appStx ← withAppArg delab
     `($(appStx).$(mkIdent f):ident)
 
+/--
+This delaborator tries to elide functions which are known coercions.
+For example, `Int.ofNat` is a coercion, so instead of printing `ofNat n` we just print `↑n`,
+and when re-parsing this we can (usually) recover the specific coercion being used.
+-/
+@[builtin_delab app]
+def coeDelaborator : Delab := whenPPOption getPPCoercions do
+  let e ← getExpr
+  let .const declName _ := e.getAppFn | failure
+  let some info ← Meta.getCoeFnInfo? declName | failure
+  let n := e.getAppNumArgs
+  withOverApp info.numArgs do
+    match info.type with
+    | .coe => `(↑$(← withNaryArg info.coercee delab))
+    | .coeFun =>
+      if n = info.numArgs then
+        `(⇑$(← withNaryArg info.coercee delab))
+      else
+        withNaryArg info.coercee delab
+    | .coeSort => `(↥$(← withNaryArg info.coercee delab))
+
 @[builtin_delab app.dite]
 def delabDIte : Delab := whenPPOption getPPNotation <| withOverApp 5 do
   -- Note: we keep this as a delaborator for now because it actually accesses the expression.

tests/lean/coe.lean

@@ -0,0 +1,53 @@
+import Lean
+
+structure WrappedNat where
+  val : Nat
+
+structure WrappedFun (α) where
+  fn : Nat → α
+
+structure WrappedType where
+  typ : Type
+
+attribute [coe] WrappedNat.val
+instance : Coe WrappedNat Nat where coe := WrappedNat.val
+
+#eval Lean.Meta.registerCoercion ``WrappedFun.fn (some ⟨2, 1, .coeFun⟩)
+instance : CoeFun (WrappedFun α) (fun _ => Nat → α) where coe := WrappedFun.fn
+
+#eval Lean.Meta.registerCoercion ``WrappedType.typ (some ⟨1, 0, .coeSort⟩)
+instance : CoeSort WrappedType Type where coe := WrappedType.typ
+
+section coe
+variable (n : WrappedNat)
+
+#check (↑n : Nat)
+
+#check n.val
+
+end coe
+
+section coeFun
+variable (f : WrappedFun Nat) (g : Nat → WrappedFun Nat) (h : WrappedFun (WrappedFun Nat))
+
+#check f.fn
+
+#check ⇑f
+
+#check ⇑f 1
+
+#check ⇑(g 1)
+
+#check g 1 2
+
+#check ⇑h
+
+end coeFun
+
+section coeSort
+variable (t : WrappedType)
+
+#check t.typ
+#check ↥t
+
+end coeSort

tests/lean/coe.lean.expected.out

@@ -0,0 +1,12 @@
+
+
+↑n : Nat
+↑n : Nat
+⇑f : Nat → Nat
+⇑f : Nat → Nat
+f 1 : Nat
+⇑(g 1) : Nat → Nat
+(g 1) 2 : Nat
+⇑h : Nat → WrappedFun Nat
+↥t : Type
+↥t : Type