fix: normalize universe levels in grind preprocessor

src/Lean/Meta/Basic.lean

@@ -605,6 +605,9 @@ variable [MonadControlT MetaM n] [Monad n]
 @[inline] def modifyCache (f : Cache → Cache) : MetaM Unit :=
   modify fun { mctx, cache, zetaDeltaFVarIds, postponed, diag } => { mctx, cache := f cache, zetaDeltaFVarIds, postponed, diag }
 
+def resetCache : MetaM Unit :=
+  modifyCache fun _ => {}
+
 @[inline] def modifyInferTypeCache (f : InferTypeCache → InferTypeCache) : MetaM Unit :=
   modifyCache fun ⟨ic, c1, c2, c3, c4, c5⟩ => ⟨f ic, c1, c2, c3, c4, c5⟩
 
@@ -1777,9 +1780,9 @@ private def withMVarContextImp (mvarId : MVarId) (x : MetaM α) : MetaM α := do
   withLocalContextImp mvarDecl.lctx mvarDecl.localInstances x
 
 /--
-  Execute `x` using the given metavariable `LocalContext` and `LocalInstances`.
-  The type class resolution cache is flushed when executing `x` if its `LocalInstances` are
-  different from the current ones. -/
+Executes `x` using the given metavariable `LocalContext` and `LocalInstances`.
+The type class resolution cache is flushed when executing `x` if its `LocalInstances` are
+different from the current ones. -/
 def _root_.Lean.MVarId.withContext (mvarId : MVarId) : n α → n α :=
   mapMetaM <| withMVarContextImp mvarId
 
@@ -1789,13 +1792,25 @@ private def withMCtxImp (mctx : MetavarContext) (x : MetaM α) : MetaM α := do
   try x finally setMCtx mctx'
 
 /--
-  `withMCtx mctx k` replaces the metavariable context and then executes `k`.
-  The metavariable context is restored after executing `k`.
+`withMCtx mctx k` replaces the metavariable context and then executes `k`.
+The metavariable context is restored after executing `k`.
 
-  This method is used to implement the type class resolution procedure. -/
+This method is used to implement the type class resolution procedure. -/
 def withMCtx (mctx : MetavarContext) : n α → n α :=
   mapMetaM <| withMCtxImp mctx
 
+/--
+`withoutModifyingMCtx k` executes `k` and then restores the metavariable context.
+-/
+def withoutModifyingMCtx : n α → n α :=
+  mapMetaM fun x => do
+    let mctx ← getMCtx
+    try
+      x
+    finally
+      resetCache
+      setMCtx mctx
+
 @[inline] private def approxDefEqImp (x : MetaM α) : MetaM α :=
   withConfig (fun config => { config with foApprox := true, ctxApprox := true, quasiPatternApprox := true}) x
 

src/Lean/Meta/Tactic/Grind/Preprocessor.lean

@@ -73,6 +73,7 @@ def introNext (goal : Goal) : PreM IntroResult := do
         let p' := r.expr
         let p' ← eraseIrrelevantMData p'
         let p' ← foldProjs p'
+        let p' ← normalizeLevels p'
         let p' ← canon p'
         let p' ← shareCommon p'
         let fvarId ← mkFreshFVarId
@@ -153,26 +154,38 @@ def ppGoals : PreM Format := do
   return r
 
 def preprocess (mvarId : MVarId) : PreM State := do
-  loop (← mkGoal mvarId)
-  if (← isTracingEnabledFor `grind.pre) then
-    trace[grind.pre] (← ppGoals)
-  get
-
-end Preprocessor
-
-open Preprocessor
-
-partial def main (mvarId : MVarId) (mainDeclName : Name) : MetaM (List MVarId) := do
   mvarId.ensureProp
   -- TODO: abstract metavars
   mvarId.ensureNoMVar
   let mvarId ← mvarId.clearAuxDecls
   let mvarId ← mvarId.revertAll
   mvarId.ensureNoMVar
-  let mvarId ← mvarId.abstractNestedProofs mainDeclName
+  let mvarId ← mvarId.abstractNestedProofs (← getMainDeclName)
   let mvarId ← mvarId.unfoldReducible
   let mvarId ← mvarId.betaReduce
-  let s ← preprocess mvarId |>.run |>.run mainDeclName
+  loop (← mkGoal mvarId)
+  if (← isTracingEnabledFor `grind.pre) then
+    trace[grind.pre] (← ppGoals)
+  get
+
+def preprocessAndProbe (mvarId : MVarId) (p : GoalM Unit) : PreM Unit := do
+  let s ← preprocess mvarId
+  s.goals.forM fun goal =>
+    discard <| GoalM.run' goal p
+
+end Preprocessor
+
+open Preprocessor
+
+def preprocessAndProbe (mvarId : MVarId) (mainDeclName : Name) (p : GoalM Unit) : MetaM Unit :=
+  withoutModifyingMCtx do
+    Preprocessor.preprocessAndProbe mvarId p |>.run |>.run mainDeclName
+
+def preprocess (mvarId : MVarId) (mainDeclName : Name) : MetaM Preprocessor.State :=
+  Preprocessor.preprocess mvarId |>.run |>.run mainDeclName
+
+def main (mvarId : MVarId) (mainDeclName : Name) : MetaM (List MVarId) := do
+  let s ← preprocess mvarId mainDeclName
   return s.goals.toList.map (·.mvarId)
 
 end Lean.Meta.Grind

src/Lean/Meta/Tactic/Grind/Types.lean

@@ -65,12 +65,15 @@ def getTrueExpr : GrindM Expr := do
 def getFalseExpr : GrindM Expr := do
   return (← get).falseExpr
 
+def getMainDeclName : GrindM Name :=
+  return (← read).mainDeclName
+
 /--
 Abtracts nested proofs in `e`. This is a preprocessing step performed before internalization.
 -/
 def abstractNestedProofs (e : Expr) : GrindM Expr := do
   let nextIdx := (← get).nextThmIdx
-  let (e, s') ← AbstractNestedProofs.visit e |>.run { baseName := (← read).mainDeclName } |>.run |>.run { nextIdx }
+  let (e, s') ← AbstractNestedProofs.visit e |>.run { baseName := (← getMainDeclName) } |>.run |>.run { nextIdx }
   modify fun s => { s with nextThmIdx := s'.nextIdx }
   return e
 
@@ -239,4 +242,18 @@ def mkGoal (mvarId : MVarId) : GrindM Goal := do
     mkENodeCore falseExpr (interpreted := true) (ctor := false) (generation := 0)
     mkENodeCore trueExpr (interpreted := true) (ctor := false) (generation := 0)
 
+def forEachENode (f : ENode → GoalM Unit) : GoalM Unit := do
+  -- We must sort because we are using pointer addresses to hash
+  let nodes := (← get).enodes.toArray.map (·.2)
+  let nodes := nodes.qsort fun a b => a.idx < b.idx
+  for n in nodes do
+    f n
+
+def filterENodes (p : ENode → GoalM Bool) : GoalM (Array ENode) := do
+  let ref ← IO.mkRef #[]
+  forEachENode fun n => do
+    if (← p n) then
+      ref.modify (·.push n)
+  ref.get
+
 end Lean.Meta.Grind

src/Lean/Meta/Tactic/Grind/Util.lean

@@ -124,4 +124,15 @@ def foldProjs (e : Expr) : MetaM Expr := do
       return .done e
   Meta.transform e (post := post)
 
+/--
+Normalizes universe levels in constants and sorts.
+-/
+def normalizeLevels (e : Expr) : CoreM Expr := do
+  let pre (e : Expr) := do
+    match e with
+    | .sort u => return .done <| e.updateSort! u.normalize
+    | .const _ us => return .done <| e.updateConst! (us.map Level.normalize)
+    | _ => return .continue
+  Core.transform e (pre := pre)
+
 end Lean.Meta.Grind

tests/lean/run/grindNormLevelIssue.lean

@@ -0,0 +1,21 @@
+import Lean
+
+def g {α : Sort u} (a : α) := a
+
+open Lean Meta Elab Tactic Grind in
+elab "grind_test" : tactic => withMainContext do
+  let declName := (← Term.getDeclName?).getD `_main
+  Meta.Grind.preprocessAndProbe (← getMainGoal) declName do
+    let nodes ← filterENodes fun e => return e.self.isAppOf ``g
+    logInfo (nodes.toList.map (·.self))
+
+-- `grind` final state must contain only two `g`-applications
+/--
+info: [g (a, b), g (g (a, b))]
+---
+warning: declaration uses 'sorry'
+-/
+#guard_msgs in
+example {β : Type v} {α : Type u} (a c : α) (b d : β) : g.{max u v + 1} (a, b) = (c, d) → g (g.{max (u+1) (v+1)} (a, b)) = (c, d) → False := by
+  grind_test
+  sorry