From c935a3ff6a7f36ed3e82ec05fe2e1903fb94e999 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Tue, 10 May 2022 12:48:14 -0700
Subject: [PATCH] feat: improve error location for universe level errors at
 `inductive` command

---
 src/Lean/Elab/Inductive.lean               | 39 +++++++++++-----------
 tests/lean/univInference.lean.expected.out |  2 +-
 2 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/src/Lean/Elab/Inductive.lean b/src/Lean/Elab/Inductive.lean
index 55faeb937c..5764bbb7f1 100644
--- a/src/Lean/Elab/Inductive.lean
+++ b/src/Lean/Elab/Inductive.lean
@@ -443,26 +443,37 @@ def accLevelAtCtor (u : Level) (r : Level) (rOffset : Nat) : StateRefT (Array Le
     else
       modify fun us => us.push u
 
+/--
+  Execute `k` using the `Syntax` reference associated with constructor `ctorName`.
+-/
+def withCtorRef [Monad m] [MonadRef m] (views : Array InductiveView) (ctorName : Name) (k : m α) : m α := do
+  for view in views do
+    for ctorView in view.ctors do
+      if ctorView.declName == ctorName then
+        return (← withRef ctorView.ref k)
+  k
+
 /-- Auxiliary function for `updateResultingUniverse` -/
-private partial def collectUniverses (r : Level) (rOffset : Nat) (numParams : Nat) (indTypes : List InductiveType) : TermElabM (Array Level) := do
+private partial def collectUniverses (views : Array InductiveView) (r : Level) (rOffset : Nat) (numParams : Nat) (indTypes : List InductiveType) : TermElabM (Array Level) := do
   let (_, us) ← go |>.run #[]
   return us
 where
   go : StateRefT (Array Level) TermElabM Unit :=
     indTypes.forM fun indType => indType.ctors.forM fun ctor =>
-      forallTelescopeReducing ctor.type fun ctorParams _ =>
-        for ctorParam in ctorParams[numParams:] do
-          let type ← inferType ctorParam
-          let u ← instantiateLevelMVars (← getLevel type)
-          accLevelAtCtor u r rOffset
+      withCtorRef views ctor.name do
+        forallTelescopeReducing ctor.type fun ctorParams _ =>
+          for ctorParam in ctorParams[numParams:] do
+            let type ← inferType ctorParam
+            let u ← instantiateLevelMVars (← getLevel type)
+            accLevelAtCtor u r rOffset
 
-private def updateResultingUniverse (numParams : Nat) (indTypes : List InductiveType) : TermElabM (List InductiveType) := do
+private def updateResultingUniverse (views : Array InductiveView) (numParams : Nat) (indTypes : List InductiveType) : TermElabM (List InductiveType) := do
   let r ← getResultingUniverse indTypes
   let rOffset : Nat   := r.getOffset
   let r       : Level := r.getLevelOffset
   unless r.isMVar do
     throwError "failed to compute resulting universe level of inductive datatype, provide universe explicitly: {r}"
-  let us ← collectUniverses r rOffset numParams indTypes
+  let us ← collectUniverses views r rOffset numParams indTypes
   trace[Elab.inductive] "updateResultingUniverse us: {us}, r: {r}, rOffset: {rOffset}"
   let rNew := mkResultUniverse us rOffset
   assignLevelMVar r.mvarId! rNew
@@ -483,16 +494,6 @@ def checkResultingUniverse (u : Level) : TermElabM Unit := do
     if !u.isZero && !u.isNeverZero then
       throwError "invalid universe polymorphic type, the resultant universe is not Prop (i.e., 0), but it may be Prop for some parameter values (solution: use 'u+1' or 'max 1 u'{indentD u}"
 
-/--
-  Execute `k` using the `Syntax` reference associated with constructor `ctorName`.
--/
-def withCtorRef (views : Array InductiveView) (ctorName : Name) (k : TermElabM α) : TermElabM α := do
-  for view in views do
-    for ctorView in view.ctors do
-      if ctorView.declName == ctorName then
-        return (← withRef ctorView.ref k)
-  k
-
 private def checkResultingUniverses (views : Array InductiveView) (numParams : Nat) (indTypes : List InductiveType) : TermElabM Unit := do
   let u := (← instantiateLevelMVars (← getResultingUniverse indTypes)).normalize
   checkResultingUniverse u
@@ -733,7 +734,7 @@ private def mkInductiveDecl (vars : Array Expr) (views : Array InductiveView) :
         let indTypes ← updateParams vars indTypes
         let indTypes ←
           if let some univToInfer := univToInfer? then
-            updateResultingUniverse numParams (← levelMVarToParam indTypes univToInfer)
+            updateResultingUniverse views numParams (← levelMVarToParam indTypes univToInfer)
           else
             checkResultingUniverses views numParams indTypes
             levelMVarToParam indTypes none
diff --git a/tests/lean/univInference.lean.expected.out b/tests/lean/univInference.lean.expected.out
index a05a382ca6..ad85bbc6a5 100644
--- a/tests/lean/univInference.lean.expected.out
+++ b/tests/lean/univInference.lean.expected.out
@@ -2,7 +2,7 @@ univInference.lean:25:0-27:9: error: failed to compute resulting universe level
 S6 : Sort (max w₁ w₂) → Type w₂ → Sort (max w₁ (w₂ + 1))
 univInference.lean:45:0-46:17: error: invalid universe polymorphic type, the resultant universe is not Prop (i.e., 0), but it may be Prop for some parameter values (solution: use 'u+1' or 'max 1 u'
   max u v
-univInference.lean:64:0-65:22: error: failed to compute resulting universe level of inductive datatype, provide universe explicitly
+univInference.lean:65:2-65:22: error: failed to compute resulting universe level of inductive datatype, provide universe explicitly
 univInference.lean:73:0-74:22: error: invalid universe polymorphic type, the resultant universe is not Prop (i.e., 0), but it may be Prop for some parameter values (solution: use 'u+1' or 'max 1 u'
   max u v
 univInference.lean:81:0-82:22: error: invalid universe polymorphic type, the resultant universe is not Prop (i.e., 0), but it may be Prop for some parameter values (solution: use 'u+1' or 'max 1 u'