feat: support coercions in the differential geometry elaborators (leanprover-community#30307)

Generalise the `CMDiff`, `MDiff` and `mfderiv` elaborators (and its friends) to apply to any object which coerces to a function. We can such speak about the smoothness of a `PartialEquiv`, `OpenPartialHomeomorph` (or the future `ClosedPartialHomeomorph` and `PartialHomeomorph`), smooth section or bundled smooth map without the need for a coercion.

As a by-product, add new functions `Lean.Meta.ensure{HasType,IsFunction,IsSort}` which apply the respective coercion if necessary.

While at it, add more sections (and section headers) to the test file for the elaborators.

Co-authored-by: pre-commit-ci-lite[bot] <117423508+pre-commit-ci-lite[bot]@users.noreply.github.com>

Author: 2 people

Mathlib/Geometry/Manifold/Notation.lean

@@ -67,7 +67,7 @@ variable {s : E → E'} in
 These elaborators can be combined: `CMDiffAt[u] n (T% s) x`
 
 **Warning.** These elaborators are a proof of concept; the implementation should be considered a
-prototype. Don't rewrite all of mathlib to use it just yet. Notable bugs and limitations include
+prototype. Don't rewrite all of mathlib to use it just yet. Notable limitations include
 the following.
 
 ## TODO
@@ -76,10 +76,8 @@ the following.
   is correct 90% of the time).
   For products of vector spaces `E × F`, this could print a warning about making a choice between
   the model in `E × F` and the product of the models on `E` and `F`.
-- extend the elaborators to support `OpenPartialHomeomorph`s and `PartialEquiv`s
-- better error messages (as needed)
-- further testing and fixing of edge cases
-- add tests for all of the above
+- better error messages (as needed), with tests
+- further testing and fixing of edge cases (with tests)
 - add delaborators for these elaborators
 
 -/
@@ -102,7 +100,7 @@ private def findSomeLocalInstanceOf? (c : Name) {α} (p : Expr → Expr → Meta
     MetaM (Option α) := do
   (← getLocalInstances).findSomeM? fun inst ↦ do
     if inst.className == c then
-      let type ← whnfR <|← instantiateMVars <|← inferType inst.fvar
+      let type ← whnfR <| ← instantiateMVars <| ← inferType inst.fvar
       p inst.fvar type
     else return none
 
@@ -112,7 +110,7 @@ to `p`. -/
 private def findSomeLocalHyp? {α} (p : Expr → Expr → MetaM (Option α)) : MetaM (Option α) := do
   (← getLCtx).findDeclRevM? fun decl ↦ do
     if decl.isImplementationDetail then return none
-    let type ← whnfR <|← instantiateMVars decl.type
+    let type ← whnfR <| ← instantiateMVars decl.type
     p decl.toExpr type
 
 end Elab
@@ -134,7 +132,7 @@ run.
 -- TODO: factor out `MetaM` component for reuse
 scoped elab:max "T% " t:term:arg : term => do
   let e ← Term.elabTerm t none
-  let etype ← whnf <|← instantiateMVars <|← inferType e
+  let etype ← whnf <| ← instantiateMVars <| ← inferType e
   match etype with
   | .forallE x base tgt _ => withLocalDeclD x base fun x ↦ do
     let tgtHasLooseBVars := tgt.hasLooseBVars
@@ -319,7 +317,7 @@ We pass `e` instead of just its type for better diagnostics.
 
 If `es` is `some`, we verify that `src` and the type of `es` are definitionally equal. -/
 def findModels (e : Expr) (es : Option Expr) : TermElabM (Expr × Expr) := do
-  let etype ← whnf <|← instantiateMVars <|← inferType e
+  let etype ← whnf <| ← instantiateMVars <| ← inferType e
   match etype with
   | .forallE _ src tgt _ =>
     if tgt.hasLooseBVars then
@@ -332,7 +330,7 @@ def findModels (e : Expr) (es : Option Expr) : TermElabM (Expr × Expr) := do
       /- Note: we use `isDefEq` here since persistent metavariable assignments in `src` and
       `estype` are acceptable.
       TODO: consider attempting to coerce `es` to a `Set`. -/
-      if !(← isDefEq estype <|← mkAppM ``Set #[src]) then
+      if !(← isDefEq estype <| ← mkAppM ``Set #[src]) then
         throwError "The domain {src} of {e} is not definitionally equal to the carrier type of \
           the set {es} : {estype}"
     let tgtI ← findModel tgt (src, srcI)
@@ -348,15 +346,15 @@ trying to determine `I` and `J` from the local context.
 The argument `x` can be omitted. -/
 scoped elab:max "MDiffAt[" s:term "]" ppSpace f:term:arg : term => do
   let es ← Term.elabTerm s none
-  let ef ← Term.elabTerm f none
+  let ef ← ensureIsFunction <| ← Term.elabTerm f none
   let (srcI, tgtI) ← findModels ef es
   mkAppM ``MDifferentiableWithinAt #[srcI, tgtI, ef, es]
 
 /-- `MDiffAt f x` elaborates to `MDifferentiableAt I J f x`,
 trying to determine `I` and `J` from the local context.
 The argument `x` can be omitted. -/
 scoped elab:max "MDiffAt" ppSpace t:term:arg : term => do
-  let e ← Term.elabTerm t none
+  let e ← ensureIsFunction <| ← Term.elabTerm t none
   let (srcI, tgtI) ← findModels e none
   mkAppM ``MDifferentiableAt #[srcI, tgtI, e]
 
@@ -366,7 +364,7 @@ scoped elab:max "MDiffAt" ppSpace t:term:arg : term => do
 -- The argument `x` can be omitted. -/
 -- scoped elab:max "MDiffAt2" ppSpace t:term:arg : term => do
 --   let e ← Term.elabTerm t none
---   let etype ← whnfR <|← instantiateMVars <|← inferType e
+--   let etype ← whnfR <| ← instantiateMVars <| ← inferType e
 --   forallBoundedTelescope etype (some 1) fun src tgt ↦ do
 --     if let some src := src[0]? then
 --       let srcI ← findModel (← inferType src)
@@ -383,25 +381,30 @@ scoped elab:max "MDiffAt" ppSpace t:term:arg : term => do
 trying to determine `I` and `J` from the local context. -/
 scoped elab:max "MDiff[" s:term "]" ppSpace t:term:arg : term => do
   let es ← Term.elabTerm s none
-  let et ← Term.elabTerm t none
+  let et ← ensureIsFunction <| ← Term.elabTerm t none
   let (srcI, tgtI) ← findModels et es
   mkAppM ``MDifferentiableOn #[srcI, tgtI, et, es]
 
 /-- `MDiff f` elaborates to `MDifferentiable I J f`,
 trying to determine `I` and `J` from the local context. -/
 scoped elab:max "MDiff" ppSpace t:term:arg : term => do
-  let e ← Term.elabTerm t none
+  let e ← ensureIsFunction <| ← Term.elabTerm t none
   let (srcI, tgtI) ← findModels e none
   mkAppM ``MDifferentiable #[srcI, tgtI, e]
 
+-- We ensure the type of `n` before checking `f` is a function to provide better error messages
+-- in case e.g. `f` and `n` are swapped.
+-- TODO: provide better error messages if just `n` is forgotten (say, by making `n` optional in
+-- the parser and erroring later in the elaborator); currently, this yields just a parser error.
+
 /-- `CMDiffAt[s] n f x` elaborates to `ContMDiffWithinAt I J n f s x`,
 trying to determine `I` and `J` from the local context.
 `n` is coerced to `WithTop ℕ∞` if necessary (so passing a `ℕ`, `∞` or `ω` are all supported).
 The argument `x` can be omitted. -/
 scoped elab:max "CMDiffAt[" s:term "]" ppSpace nt:term:arg ppSpace f:term:arg : term => do
   let es ← Term.elabTerm s none
-  let ef ← Term.elabTerm f none
   let ne ← Term.elabTermEnsuringType nt q(WithTop ℕ∞)
+  let ef ← ensureIsFunction <| ← Term.elabTerm f none
   let (srcI, tgtI) ← findModels ef es
   mkAppM ``ContMDiffWithinAt #[srcI, tgtI, ne, ef, es]
 
@@ -410,7 +413,7 @@ trying to determine `I` and `J` from the local context.
 `n` is coerced to `WithTop ℕ∞` if necessary (so passing a `ℕ`, `∞` or `ω` are all supported).
 The argument `x` can be omitted. -/
 scoped elab:max "CMDiffAt" ppSpace nt:term:arg ppSpace t:term:arg : term => do
-  let e ← Term.elabTerm t none
+  let e ← ensureIsFunction <| ← Term.elabTerm t none
   let ne ← Term.elabTermEnsuringType nt q(WithTop ℕ∞)
   let (srcI, tgtI) ← findModels e none
   mkAppM ``ContMDiffAt #[srcI, tgtI, ne, e]
@@ -420,32 +423,32 @@ trying to determine `I` and `J` from the local context.
 `n` is coerced to `WithTop ℕ∞` if necessary (so passing a `ℕ`, `∞` or `ω` are all supported). -/
 scoped elab:max "CMDiff[" s:term "]" ppSpace nt:term:arg ppSpace f:term:arg : term => do
   let es ← Term.elabTerm s none
-  let ef ← Term.elabTerm f none
   let ne ← Term.elabTermEnsuringType nt q(WithTop ℕ∞)
+  let ef ← ensureIsFunction <| ← Term.elabTerm f none
   let (srcI, tgtI) ← findModels ef es
   mkAppM ``ContMDiffOn #[srcI, tgtI, ne, ef, es]
 
 /-- `CMDiff n f` elaborates to `ContMDiff I J n f`,
 trying to determine `I` and `J` from the local context.
 `n` is coerced to `WithTop ℕ∞` if necessary (so passing a `ℕ`, `∞` or `ω` are all supported). -/
 scoped elab:max "CMDiff" ppSpace nt:term:arg ppSpace f:term:arg : term => do
-  let e ← Term.elabTerm f none
   let ne ← Term.elabTermEnsuringType nt q(WithTop ℕ∞)
+  let e ← ensureIsFunction <| ← Term.elabTerm f none
   let (srcI, tgtI) ← findModels e none
   mkAppM ``ContMDiff #[srcI, tgtI, ne, e]
 
 /-- `mfderiv[u] f x` elaborates to `mfderivWithin I J f u x`,
 trying to determine `I` and `J` from the local context. -/
 scoped elab:max "mfderiv[" s:term "]" ppSpace t:term:arg : term => do
   let es ← Term.elabTerm s none
-  let e ← Term.elabTerm t none
+  let e ← ensureIsFunction <| ← Term.elabTerm t none
   let (srcI, tgtI) ← findModels e es
   mkAppM ``mfderivWithin #[srcI, tgtI, e, es]
 
 /-- `mfderiv% f x` elaborates to `mfderiv I J f x`,
 trying to determine `I` and `J` from the local context. -/
 scoped elab:max "mfderiv%" ppSpace t:term:arg : term => do
-  let e ← Term.elabTerm t none
+  let e ← ensureIsFunction <| ← Term.elabTerm t none
   let (srcI, tgtI) ← findModels e none
   mkAppM ``mfderiv #[srcI, tgtI, e]
 

Mathlib/Lean/Meta/Basic.lean

@@ -94,3 +94,28 @@ def Lean.Meta.withEnsuringLocalInstance {α : Type} (inst : MVarId) (k : MetaM (
       let (e, v) ← k
       let e' := (← e.abstractM #[inst']).instantiate1 instE
       return (e', v)
+
+/-- Checks that `e` has type `expectedType` (i.e. that its type is defeq to `expectedType`
+at the current transparency). If not, coerces `e` to this type
+or fails with a descriptive error. -/
+def Lean.Meta.ensureHasType (e expectedType : Expr) : MetaM Expr := do
+  let ty ← inferType e
+  if ← withNewMCtxDepth (isDefEq ty expectedType) then return e else
+    (← coerceSimple? e expectedType).toOption.getDM <|
+      throwError "Expected{indentD e}\nto have type{indentD ty}\n or to be coercible to it"
+
+/-- Checks that `e` is a function (i.e. that its type is a `.forallE` after `instantiateMVars` and
+`whnf`). If not, coerces `e` to a function or fails with a descriptive error. -/
+def Lean.Meta.ensureIsFunction (e : Expr) : MetaM Expr := do
+  let ty ← whnf <| ← instantiateMVars <| ← inferType e
+  if ty.isForall then return e else (← coerceToFunction? e).getDM <|
+    throwError "Expected{indentD e}\nof type{indentD ty}\nto be a function, or to be coercible to \
+      a function"
+
+/-- Checks that `e` is a type (i.e. that its type is a `Sort` after `instantiateMVars` and
+`whnf`). If not, coerces `e` to a Sort or fails with a descriptive error. -/
+def Lean.Meta.ensureIsSort (e : Expr) : MetaM Expr := do
+  let ty ← whnf <| ← instantiateMVars <| ← inferType e
+  if ty.isSort then return e else (← coerceToSort? e).getDM <|
+    throwError "Expected{indentD e}\nof type{indentD ty}\nto be a Sort, or to be coercible to \
+      a Sort"

MathlibTest/DifferentialGeometry/Notation.lean

@@ -92,11 +92,13 @@ end TotalSpace
 /-! Tests for the elaborators for `MDifferentiable{WithinAt,At,On}`. -/
 section differentiability
 
--- Start with some basic tests: a simple function, both in applied and unapplied form.
 variable {EM' : Type*} [NormedAddCommGroup EM']
   [NormedSpace 𝕜 EM'] {H' : Type*} [TopologicalSpace H'] (I' : ModelWithCorners 𝕜 EM' H')
   {M' : Type*} [TopologicalSpace M'] [ChartedSpace H' M']
 
+/-! Some basic tests: a simple function, both in applied and unapplied form -/
+section basic
+
 -- General case: a function between two manifolds.
 variable {f : M → M'} {s : Set M} {m : M}
 
@@ -306,12 +308,98 @@ variable {f : 𝕜 → 𝕜} {u : Set 𝕜} {a : 𝕜}
 #guard_msgs in
 #check MDiff[u] f
 
--- This elaborator can be combined with the total space elaborator.
--- XXX: these tests might be incomplete; extend as needed!
+end basic
+
+/-! A partial homeomorphism or partial equivalence. More generally, this works for any type
+with a coercion to (possibly dependent) functions. -/
+section coercion
+
+variable {s : Set M} {m : M}
+
+variable {φ : OpenPartialHomeomorph M E} {ψ : PartialEquiv M E}
+
+/-- info: MDifferentiableWithinAt I 𝓘(𝕜, E) (↑φ) s : M → Prop -/
+#guard_msgs in
+#check MDiffAt[s] φ
+
+/-- info: MDifferentiableWithinAt I 𝓘(𝕜, E) (↑ψ) s : M → Prop -/
+#guard_msgs in
+#check MDiffAt[s] ψ
+
+/-- info: MDifferentiableAt I 𝓘(𝕜, E) ↑φ : M → Prop -/
+#guard_msgs in
+#check MDiffAt φ
+
+/-- info: MDifferentiableAt I 𝓘(𝕜, E) ↑ψ : M → Prop -/
+#guard_msgs in
+#check MDiffAt ψ
+
+/-- info: MDifferentiableOn I 𝓘(𝕜, E) (↑φ) s : Prop -/
+#guard_msgs in
+#check MDiff[s] φ
+
+/-- info: MDifferentiableOn I 𝓘(𝕜, E) (↑ψ) s : Prop -/
+#guard_msgs in
+#check MDiff[s] ψ
+
+/-- info: MDifferentiable I 𝓘(𝕜, E) ↑φ : Prop -/
+#guard_msgs in
+#check MDiff φ
+
+/-- info: ContMDiffWithinAt I 𝓘(𝕜, E) 2 (↑ψ) s : M → Prop -/
+#guard_msgs in
+#check CMDiffAt[s] 2 ψ
+
+/-- info: ContMDiffOn I 𝓘(𝕜, E) 2 (↑φ) s : Prop -/
+#guard_msgs in
+#check CMDiff[s] 2 φ
+
+/-- info: ContMDiffAt I 𝓘(𝕜, E) 2 ↑φ : M → Prop -/
+#guard_msgs in
+#check CMDiffAt 2 φ
+
+/-- info: ContMDiff I 𝓘(𝕜, E) 2 ↑ψ : Prop -/
+#guard_msgs in
+#check CMDiff 2 ψ
+
+/-- info: mfderiv I 𝓘(𝕜, E) ↑φ : (x : M) → TangentSpace I x →L[𝕜] TangentSpace 𝓘(𝕜, E) (↑φ x) -/
+#guard_msgs in
+#check mfderiv% φ
+
+/--
+info: mfderivWithin I 𝓘(𝕜, E) (↑ψ) s : (x : M) → TangentSpace I x →L[𝕜] TangentSpace 𝓘(𝕜, E) (↑ψ x)
+-/
+#guard_msgs in
+#check mfderiv[s] ψ
+
+/--
+info: mfderivWithin I 𝓘(𝕜, E) (↑ψ) s : (x : M) → TangentSpace I x →L[𝕜] TangentSpace 𝓘(𝕜, E) (↑ψ x)
+-/
+#guard_msgs in
+variable {f : ContMDiffSection I F n V} in
+#check mfderiv[s] ψ
+
+/-- info: mfderiv I I' ⇑g : (x : M) → TangentSpace I x →L[𝕜] TangentSpace I' (g x) -/
+#guard_msgs in
+variable {g : ContMDiffMap I I' M M' n} in
+#check mfderiv% g
+
+-- An example of "any type" which coerces to functions.
+/-- info: mfderiv I I' ⇑g : (x : M) → TangentSpace I x →L[𝕜] TangentSpace I' (g x) -/
+#guard_msgs in
+variable {g : Equiv M M'} in
+#check mfderiv% g
+
+end coercion
 
 variable {σ : Π x : M, V x} {σ' : (x : E) → Trivial E E' x} {s : E → E'}
 variable (X : (m : M) → TangentSpace I m) [IsManifold I 1 M]
 
+/-! These elaborators can be combined with the total space elaborator. -/
+section interaction
+
+-- Note: these tests might be incomplete; extend as needed!
+
 /-- info: MDifferentiableAt I (I.prod 𝓘(𝕜, E)) fun m ↦ TotalSpace.mk' E m (X m) : M → Prop -/
 #guard_msgs in
 #check MDiffAt (T% X)
@@ -326,6 +414,8 @@ info: MDifferentiableAt 𝓘(𝕜, E) (𝓘(𝕜, E).prod 𝓘(𝕜, E')) fun x
 #guard_msgs in
 #check MDiffAt (T% σ')
 
+end interaction
+
 section
 
 variable [IsManifold I 2 M]
@@ -479,10 +569,11 @@ variable {f : M → M'} {s : Set M} {m : M}
 
 variable [IsManifold I 1 M] [IsManifold I' 1 M']
 
--- TODO: can there be better error messages when forgetting the smoothness exponent?
+-- Testing error messages when forgetting the smoothness exponent or swapping arguments.
 section error
 
--- yields a parse error, "unexpected toekn '/--'; expected term"
+-- yields a parse error, "unexpected token '/--'; expected term"
+-- TODO: make this parse, but error in the elaborator
 -- #check CMDiffAt[s] f
 
 /--
@@ -498,7 +589,7 @@ error: Expected
   m
 of type
   M
-to be a function
+to be a function, or to be coercible to a function
 -/
 #guard_msgs in
 #check CMDiffAt[s] f m
@@ -516,14 +607,32 @@ error: Expected
   m
 of type
   M
-to be a function
+to be a function, or to be coercible to a function
 -/
 #guard_msgs in
-#check CMDiffAt[s] f m
+#check CMDiff[s] f m
 
--- yields a parse error, "unexpected toekn '/--'; expected term"
+-- yields a parse error, "unexpected token '/--'; expected term"
 -- #check CMDiffAt f
 
+/--
+error: Type mismatch
+  f
+has type
+  M → M'
+of sort `Type (max u_10 u_4)` but is expected to have type
+  WithTop ℕ∞
+of sort `Type`
+---
+error: Expected
+  n
+of type
+  Option ℕ∞
+to be a function, or to be coercible to a function
+-/
+#guard_msgs in
+#check CMDiff f n
+
 end error
 
 /-- info: ContMDiffWithinAt I I' 1 f s : M → Prop -/