feat: improve error message when stuck solving universe constraints

closes #343

src/Lean/Elab/Term.lean

@@ -5,6 +5,7 @@ Authors: Leonardo de Moura
 -/
 import Lean.ResolveName
 import Lean.Util.Sorry
+import Lean.Util.ReplaceExpr
 import Lean.Structure
 import Lean.Meta.ExprDefEq
 import Lean.Meta.AppBuilder
@@ -1469,6 +1470,55 @@ unsafe def evalExpr (α) (typeName : Name) (value : Expr) : TermElabM α :=
     addAndCompile decl
     evalConst α name
 
+private def throwStuckAtUniverseCnstr : TermElabM Unit := do
+  let entries ← getPostponed
+  let mut found : Std.HashSet (Level × Level) := {}
+  let mut uniqueEntries := #[]
+  for entry in entries do
+    let mut lhs := entry.lhs
+    let mut rhs := entry.rhs
+    if Level.normLt rhs lhs then
+      (lhs, rhs) := (rhs, lhs)
+    unless found.contains (lhs, rhs) do
+      found := found.insert (lhs, rhs)
+      uniqueEntries := uniqueEntries.push entry
+  for i in [1:uniqueEntries.size] do
+    logErrorAt uniqueEntries[i].ref (← mkMessage uniqueEntries[i])
+  throwErrorAt uniqueEntries[0].ref (← mkMessage uniqueEntries[0])
+where
+  /- Annotate any constant and sort in `e` that satisfies `p` with `pp.universes true` -/
+  exposeRelevantUniverses (e : Expr) (p : Level → Bool) : Expr :=
+    e.replace fun e =>
+      match e with
+      | Expr.const _ us _ => if us.any p then some (e.setPPUniverses true) else none
+      | Expr.sort u _     => if p u then some (e.setPPUniverses true) else none
+      | _                 => none
+
+  mkMessage (entry : PostponedEntry) : TermElabM MessageData := do
+    match entry.ctx? with
+    | none =>
+      return m!"stuck at solving universe constraints{indentD m!"{entry.lhs} =?= {entry.rhs}"}"
+    | some ctx =>
+      withLCtx ctx.lctx ctx.localInstances do
+        let s   := entry.lhs.collectMVars entry.rhs.collectMVars
+        /- `p u` is true if it contains a universe metavariable in `s` -/
+        let p (u : Level) := u.any fun | Level.mvar m _ => s.contains m | _ => false
+        let lhs := exposeRelevantUniverses (← instantiateMVars ctx.lhs) p
+        let rhs := exposeRelevantUniverses (← instantiateMVars ctx.rhs) p
+        addMessageContext m!"stuck at solving universe constraints{indentD m!"{entry.lhs} =?= {entry.rhs}"}\nwhile trying to unify{indentExpr lhs}\nwith{indentExpr rhs}"
+
+@[specialize] def withoutPostponingUniverseConstraints (x : TermElabM α) : TermElabM α := do
+  let postponed ← getResetPostponed
+  try
+    let a ← x
+    unless (← processPostponed (mayPostpone := false)) do
+      throwStuckAtUniverseCnstr
+    setPostponed postponed
+    return a
+  catch ex =>
+    setPostponed postponed
+    throw ex
+
 end Term
 
 builtin_initialize

src/Lean/Expr.lean

@@ -976,6 +976,9 @@ def setOption (e : Expr) (optionName : Name) [KVMap.Value α] (val : α) : Expr
 def setPPExplicit (e : Expr) (flag : Bool) :=
   e.setOption `pp.explicit flag
 
+def setPPUniverses (e : Expr) (flag : Bool) :=
+  e.setOption `pp.universes flag
+
 /- If `e` is an application `f a_1 ... a_n` annotate `f`, `a_1` ... `a_n` with `pp.explicit := false`,
    and annotate `e` with `pp.explicit := true`. -/
 def setAppPPExplicit (e : Expr) : Expr :=

src/Lean/Level.lean

@@ -531,6 +531,26 @@ abbrev LevelSet := HashSet Level
 abbrev PersistentLevelSet := PHashSet Level
 abbrev PLevelSet := PersistentLevelSet
 
+def Level.collectMVars (u : Level) (s : NameSet := {}) : NameSet :=
+  match u with
+  | succ v _   => collectMVars v s
+  | max u v _  => collectMVars u (collectMVars v s)
+  | imax u v _ => collectMVars u (collectMVars v s)
+  | mvar n _   => s.insert n
+  | _          => s
+
+def Level.find? (u : Level) (p : Level → Bool) : Option Level :=
+  if p u then
+    some u
+  else match u with
+    | succ v _   => v.find? p
+    | max u v _  => u.find? p <|> v.find? p
+    | imax u v _ => u.find? p <|> v.find? p
+    | _          => none
+
+def Level.any (u : Level) (p : Level → Bool) : Bool :=
+  u.find? p |>.isSome
+
 end Lean
 
 abbrev Nat.toLevel (n : Nat) : Lean.Level :=

src/Lean/Meta/Basic.lean

@@ -96,9 +96,27 @@ structure Cache where
   whnfAll       : WhnfCache := {} -- cache for closed terms and `TransparencyMode.all`
   deriving Inhabited
 
+/--
+ "Context" for a postponed universe constraint.
+ `lhs` and `rhs` are the surrounding `isDefEq` call when the postponed constraint was created.
+-/
+structure DefEqContext where
+  lhs            : Expr
+  rhs            : Expr
+  lctx           : LocalContext
+  localInstances : LocalInstances
+
+/--
+  Auxiliary structure for representing postponed universe constraints.
+  Remark: the fields `ref` and `rootDefEq?` are used for error message generation only.
+  Remark: we may consider improving the error message generation in the future.
+-/
 structure PostponedEntry where
-  lhs       : Level
-  rhs       : Level
+  ref  : Syntax -- We save the `ref` at entry creation time
+  lhs  : Level
+  rhs  : Level
+  ctx? : Option DefEqContext -- Context for the surrounding `isDefEq` call when entry was created
+  deriving Inhabited
 
 structure State where
   mctx        : MetavarContext := {}
@@ -109,9 +127,11 @@ structure State where
   deriving Inhabited
 
 structure Context where
-  config         : Config         := {}
-  lctx           : LocalContext   := {}
-  localInstances : LocalInstances := #[]
+  config         : Config               := {}
+  lctx           : LocalContext         := {}
+  localInstances : LocalInstances       := #[]
+  /-- Not `none` when inside of an `isDefEq` test. See `PostponedEntry`. -/
+  defEqCtx?      : Option DefEqContext  := none
 
 abbrev MetaM  := ReaderT Context $ StateRefT State CoreM
 

src/Lean/Meta/LevelDefEq.lean

@@ -81,8 +81,8 @@ private def solveSelfMax (mvarId : MVarId) (v : Level) : MetaM Unit := do
   let n ← mkFreshLevelMVar
   assignLevelMVar mvarId <| mkMaxArgsDiff mvarId v n
 
-private def postponeIsLevelDefEq (lhs : Level) (rhs : Level) : MetaM Unit :=
-  modifyPostponed fun postponed => postponed.push { lhs := lhs, rhs := rhs }
+private def postponeIsLevelDefEq (lhs : Level) (rhs : Level) : MetaM Unit := do
+  modifyPostponed fun postponed => postponed.push { lhs := lhs, rhs := rhs, ref := (← getRef), ctx? := (← read).defEqCtx? }
 
 mutual
 
@@ -158,7 +158,7 @@ private def getNumPostponed : MetaM Nat := do
 
 open Std (PersistentArray)
 
-private def getResetPostponed : MetaM (PersistentArray PostponedEntry) := do
+def getResetPostponed : MetaM (PersistentArray PostponedEntry) := do
   let ps ← getPostponed
   setPostponed {}
   return ps
@@ -167,33 +167,33 @@ private def processPostponedStep : MetaM Bool :=
   traceCtx `Meta.isLevelDefEq.postponed.step do
     let ps ← getResetPostponed
     for p in ps do
-      unless (← isLevelDefEqAux p.lhs p.rhs) do
+      unless (← withReader (fun ctx => { ctx with defEqCtx? := p.ctx? }) <| isLevelDefEqAux p.lhs p.rhs) do
         return false
     return true
 
-private partial def processPostponed (mayPostpone : Bool := true) : MetaM Bool := do
-if (← getNumPostponed) == 0 then
-  return true
-else
-  traceCtx `Meta.isLevelDefEq.postponed do
-    let rec loop : MetaM Bool := do
-      let numPostponed ← getNumPostponed
-      if numPostponed == 0 then
-        return true
-      else
-        trace[Meta.isLevelDefEq.postponed] "processing #{numPostponed} postponed is-def-eq level constraints"
-        if !(← processPostponedStep) then
-          return false
+partial def processPostponed (mayPostpone : Bool := true) : MetaM Bool := do
+  if (← getNumPostponed) == 0 then
+    return true
+  else
+    traceCtx `Meta.isLevelDefEq.postponed do
+      let rec loop : MetaM Bool := do
+        let numPostponed ← getNumPostponed
+        if numPostponed == 0 then
+          return true
         else
-          let numPostponed' ← getNumPostponed
-          if numPostponed' == 0 then
-            return true
-          else if numPostponed' < numPostponed then
-            loop
+          trace[Meta.isLevelDefEq.postponed] "processing #{numPostponed} postponed is-def-eq level constraints"
+          if !(← processPostponedStep) then
+            return false
           else
-            trace[Meta.isLevelDefEq.postponed] "no progress solving pending is-def-eq level constraints"
-            return mayPostpone
-    loop
+            let numPostponed' ← getNumPostponed
+            if numPostponed' == 0 then
+              return true
+            else if numPostponed' < numPostponed then
+              loop
+            else
+              trace[Meta.isLevelDefEq.postponed] "no progress solving pending is-def-eq level constraints"
+              return mayPostpone
+      loop
 
 private def restore (env : Environment) (mctx : MetavarContext) (postponed : PersistentArray PostponedEntry) : MetaM Unit := do
   setEnv env
@@ -226,34 +226,6 @@ private def restore (env : Environment) (mctx : MetavarContext) (postponed : Per
     restore env mctx postponed
     throw ex
 
-private def postponedToMessageData (ps : PersistentArray PostponedEntry) : MessageData := do
-  let mut found : Std.HashSet (Level × Level) := {}
-  let mut r := MessageData.nil
-  for p in ps do
-    let mut lhs := p.lhs
-    let mut rhs := p.rhs
-    if Level.normLt rhs lhs then
-      (lhs, rhs) := (rhs, lhs)
-    unless found.contains (lhs, rhs) do
-      found := found.insert (lhs, rhs)
-      r := m!"{r}\n{lhs} =?= {rhs}"
-  return r
-
-@[specialize] def withoutPostponingUniverseConstraintsImp {α} (x : MetaM α) : MetaM α := do
-  let postponed ← getResetPostponed
-  try
-    let a ← x
-    unless (← processPostponed (mayPostpone := false)) do
-      throwError "stuck at solving universe constraints{MessageData.nestD (postponedToMessageData (← getPostponed))}"
-    setPostponed postponed
-    return a
-  catch ex =>
-    setPostponed postponed
-    throw ex
-
-@[inline] def withoutPostponingUniverseConstraints {α m} [MonadControlT MetaM m] [Monad m] : m α → m α :=
-  mapMetaM <| withoutPostponingUniverseConstraintsImp
-
 def isLevelDefEq (u v : Level) : MetaM Bool :=
   traceCtx `Meta.isLevelDefEq do
     let b ← commitWhen (mayPostpone := true) <| Meta.isLevelDefEqAux u v
@@ -261,7 +233,7 @@ def isLevelDefEq (u v : Level) : MetaM Bool :=
     return b
 
 def isExprDefEq (t s : Expr) : MetaM Bool :=
-  traceCtx `Meta.isDefEq do
+  traceCtx `Meta.isDefEq <| withReader (fun ctx => { ctx with defEqCtx? := some { lhs := t, rhs := s, lctx := ctx.lctx, localInstances := ctx.localInstances } }) do
     let b ← commitWhen (mayPostpone := true) <| Meta.isExprDefEqAux t s
     trace[Meta.isDefEq] "{t} =?= {s} ... {if b then "success" else "failure"}"
     return b

tests/lean/343.lean

@@ -0,0 +1,30 @@
+structure CatIsh where
+  Obj : Type o
+  Hom : Obj → Obj → Type m
+
+infixr:75 " ~> " => (CatIsh.Hom _)
+
+structure FunctorIsh (C D : CatIsh) where
+  onObj : C.Obj → D.Obj
+  onHom : ∀ {s d : C.Obj}, (s ~> d) → (onObj s ~> onObj d)
+
+def Catish : CatIsh :=
+  {
+    Obj := CatIsh
+    Hom := FunctorIsh
+  }
+
+universes m o
+unif_hint (mvar : CatIsh) where
+  Catish.{m,o} =?= mvar |- mvar.Obj =?= CatIsh.{m,o}
+
+structure CtxSyntaxLayerParamsObj where
+  Ct : CatIsh
+
+def CtxSyntaxLayerParams : CatIsh :=
+  {
+    Obj := CtxSyntaxLayerParamsObj
+    Hom := sorry
+  }
+
+def CtxSyntaxLayerTy := CtxSyntaxLayerParams ~> Catish

tests/lean/343.lean.expected.out

@@ -0,0 +1,9 @@
+343.lean:27:11-27:16: warning: declaration uses 'sorry'
+343.lean:30:24-30:54: error: stuck at solving universe constraints
+  max (?u+1) (?u+1) =?= max (?u+1) (?u+1)
+while trying to unify
+  CatIsh.Obj.{max (max ?u ?u) (?u + 1) (?u + 1),
+      max ((max ?u ?u) + 1) ((max (?u + 1) (?u + 1)) + 1)}
+    Catish.{max ?u ?u, max (?u + 1) (?u + 1)}
+with
+  CatIsh.{max ?u ?u, max (?u + 1) (?u + 1)}

tests/lean/docStr.lean.expected.out

@@ -208,11 +208,11 @@ namedPattern :=
       endPos := { line := 161, column := 29 },
       endCharUtf16 := 29 } }
 Lean.Meta.forallTelescopeReducing :=
-  { range := { pos := { line := 679, column := 0 },
+  { range := { pos := { line := 699, column := 0 },
       charUtf16 := 0,
-      endPos := { line := 680, column := 58 },
+      endPos := { line := 700, column := 58 },
       endCharUtf16 := 58 },
-    selectionRange := { pos := { line := 679, column := 4 },
+    selectionRange := { pos := { line := 699, column := 4 },
       charUtf16 := 4,
-      endPos := { line := 679, column := 27 },
+      endPos := { line := 699, column := 27 },
       endCharUtf16 := 27 } }