feat: improve grind failure message (leanprover#6633)

This PR improves the failure message produced by the `grind` tactic. We
now include information about asserted facts, propositions that are
known to be true and false, and equivalence classes.

src/Init/Grind.lean

@@ -11,3 +11,4 @@ import Init.Grind.Cases
 import Init.Grind.Propagator
 import Init.Grind.Util
 import Init.Grind.Offset
+import Init.Grind.PP

src/Init/Grind/PP.lean

@@ -0,0 +1,30 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+prelude
+import Init.NotationExtra
+
+namespace Lean.Grind
+/-!
+This is a hackish module for hovering node information in the `grind` tactic state.
+-/
+
+inductive NodeDef where
+  | unit
+
+set_option linter.unusedVariables false in
+def node_def (_ : Nat) {α : Sort u} {a : α} : NodeDef := .unit
+
+@[app_unexpander node_def]
+def nodeDefUnexpander : PrettyPrinter.Unexpander := fun stx => do
+  match stx with
+  | `($_ $id:num) => return mkIdent <| Name.mkSimple $ "#" ++ toString id.getNat
+  | _ => throw ()
+
+@[app_unexpander NodeDef]
+def NodeDefUnexpander : PrettyPrinter.Unexpander := fun _ => do
+  return mkIdent <| Name.mkSimple "NodeDef"
+
+end Lean.Grind

src/Init/Grind/Util.lean

@@ -9,7 +9,7 @@ import Init.Core
 namespace Lean.Grind
 
 /-- A helper gadget for annotating nested proofs in goals. -/
-def nestedProof (p : Prop) (h : p) : p := h
+def nestedProof (p : Prop) {h : p} : p := h
 
 /--
 Gadget for marking terms that should not be normalized by `grind`s simplifier.
@@ -28,7 +28,7 @@ When `EqMatch a b origin` is `True`, we mark `origin` as a resolved case-split.
 -/
 def EqMatch (a b : α) {_origin : α} : Prop := a = b
 
-theorem nestedProof_congr (p q : Prop) (h : p = q) (hp : p) (hq : q) : HEq (nestedProof p hp) (nestedProof q hq) := by
+theorem nestedProof_congr (p q : Prop) (h : p = q) (hp : p) (hq : q) : HEq (@nestedProof p hp) (@nestedProof q hq) := by
   subst h; apply HEq.refl
 
 end Lean.Grind

src/Lean/Elab/Tactic/Grind.lean

@@ -35,9 +35,9 @@ def elabGrindPattern : CommandElab := fun stx => do
   | _ => throwUnsupportedSyntax
 
 def grind (mvarId : MVarId) (config : Grind.Config) (mainDeclName : Name) (fallback : Grind.Fallback) : MetaM Unit := do
-  let mvarIds ← Grind.main mvarId config mainDeclName fallback
-  unless mvarIds.isEmpty do
-    throwError "`grind` failed\n{goalsToMessageData mvarIds}"
+  let goals ← Grind.main mvarId config mainDeclName fallback
+  unless goals.isEmpty do
+    throwError "`grind` failed\n{← Grind.goalsToMessageData goals}"
 
 private def elabFallback (fallback? : Option Term) : TermElabM (Grind.GoalM Unit) := do
   let some fallback := fallback? | return (pure ())

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

@@ -118,7 +118,7 @@ private partial def addEqStep (lhs rhs proof : Expr) (isHEq : Bool) : GoalM Unit
   unless (← isInconsistent) do
     if valueInconsistency then
       closeGoalWithValuesEq lhsRoot.self rhsRoot.self
-  trace_goal[grind.debug] "after addEqStep, {← ppState}"
+  trace_goal[grind.debug] "after addEqStep, {← (← get).ppState}"
   checkInvariants
 where
   go (lhs rhs : Expr) (lhsNode rhsNode lhsRoot rhsRoot : ENode) (flipped : Bool) : GoalM Unit := do
@@ -192,22 +192,27 @@ where
     processTodo
 
 /-- Adds a new equality `lhs = rhs`. It assumes `lhs` and `rhs` have already been internalized. -/
-def addEq (lhs rhs proof : Expr) : GoalM Unit := do
+private def addEq (lhs rhs proof : Expr) : GoalM Unit := do
   addEqCore lhs rhs proof false
 
-
 /-- Adds a new heterogeneous equality `HEq lhs rhs`. It assumes `lhs` and `rhs` have already been internalized. -/
-def addHEq (lhs rhs proof : Expr) : GoalM Unit := do
+private def addHEq (lhs rhs proof : Expr) : GoalM Unit := do
   addEqCore lhs rhs proof true
 
+/-- Save asserted facts for pretty printing goal. -/
+private def storeFact (fact : Expr) : GoalM Unit := do
+  modify fun s => { s with facts := s.facts.push fact }
+
 /-- Internalizes `lhs` and `rhs`, and then adds equality `lhs = rhs`. -/
 def addNewEq (lhs rhs proof : Expr) (generation : Nat) : GoalM Unit := do
+  storeFact (← mkEq lhs rhs)
   internalize lhs generation
   internalize rhs generation
   addEq lhs rhs proof
 
 /-- Adds a new `fact` justified by the given proof and using the given generation. -/
 def add (fact : Expr) (proof : Expr) (generation := 0) : GoalM Unit := do
+  storeFact fact
   trace_goal[grind.assert] "{fact}"
   if (← isInconsistent) then return ()
   resetNewEqs

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

@@ -72,8 +72,8 @@ def all (goals : List Goal) (f : Goal → GrindM (List Goal)) : GrindM (List Goa
 private def simple (goals : List Goal) : GrindM (List Goal) := do
   applyToAll (assertAll >> ematchStar >> (splitNext >> assertAll >> ematchStar).iterate) goals
 
-def main (mvarId : MVarId) (config : Grind.Config) (mainDeclName : Name) (fallback : Fallback) : MetaM (List MVarId) := do
-  let go : GrindM (List MVarId) := do
+def main (mvarId : MVarId) (config : Grind.Config) (mainDeclName : Name) (fallback : Fallback) : MetaM (List Goal) := do
+  let go : GrindM (List Goal) := do
     let goals ← initCore mvarId
     let goals ← simple goals
     let goals ← goals.filterMapM fun goal => do
@@ -83,7 +83,7 @@ def main (mvarId : MVarId) (config : Grind.Config) (mainDeclName : Name) (fallba
       if (← goal.mvarId.isAssigned) then return none
       return some goal
     trace[grind.debug.final] "{← ppGoals goals}"
-    return goals.map (·.mvarId)
+    return goals
   go.run mainDeclName config fallback
 
 end Lean.Meta.Grind

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

@@ -5,62 +5,107 @@ Authors: Leonardo de Moura
 -/
 prelude
 import Init.Grind.Util
+import Init.Grind.PP
 import Lean.Meta.Tactic.Grind.Types
 
 namespace Lean.Meta.Grind
 
 /-- Helper function for pretty printing the state for debugging purposes. -/
-def ppENodeRef (e : Expr) : GoalM Format := do
-  let some n ← getENode? e | return "_"
-  return f!"#{n.idx}"
+def Goal.ppENodeRef (goal : Goal) (e : Expr) : MetaM MessageData := do
+  let some n := goal.getENode? e | return "_"
+  let type ← inferType e
+  let u ← getLevel type
+  let d := mkApp3 (mkConst ``Grind.node_def [u]) (toExpr n.idx) type e
+  return m!"{d}"
+
+@[inherit_doc Goal.ppENodeRef]
+def ppENodeRef (e : Expr) : GoalM MessageData := do
+  (← get).ppENodeRef e
 
 /-- Helper function for pretty printing the state for debugging purposes. -/
-def ppENodeDeclValue (e : Expr) : GoalM Format := do
+private def Goal.ppENodeDeclValue (goal : Goal) (e : Expr) : MetaM MessageData := do
   if e.isApp && !(← isLitValue e) then
     e.withApp fun f args => do
       let r ← if f.isConst then
-        ppExpr f
+        pure m!"{f}"
       else
-        ppENodeRef f
+        goal.ppENodeRef f
       let mut r := r
       for arg in args do
-        r := r ++ " " ++ (← ppENodeRef arg)
+        r := r ++ " " ++ (← goal.ppENodeRef arg)
       return r
   else
     ppExpr e
 
 /-- Helper function for pretty printing the state for debugging purposes. -/
-def ppENodeDecl (e : Expr) : GoalM Format := do
-  let mut r := f!"{← ppENodeRef e} := {← ppENodeDeclValue e}"
-  let n ← getENode e
+private def Goal.ppENodeDecl (goal : Goal) (e : Expr) : MetaM MessageData := do
+  let mut r := m!"{← goal.ppENodeRef e} := {← goal.ppENodeDeclValue e}"
+  let n ← goal.getENode e
   unless isSameExpr e n.root do
-    r := r ++ f!" ↦ {← ppENodeRef n.root}"
+    r := r ++ m!" ↦ {← goal.ppENodeRef n.root}"
   if n.interpreted then
     r := r ++ ", [val]"
   if n.ctor then
     r := r ++ ", [ctor]"
   if grind.debug.get (← getOptions) then
-    if let some target ← getTarget? e then
-      r := r ++ f!" ↝ {← ppENodeRef target}"
+    if let some target := goal.getTarget? e then
+      r := r ++ m!" ↝ {← goal.ppENodeRef target}"
   return r
 
 /-- Pretty print goal state for debugging purposes. -/
-def ppState : GoalM Format := do
-  let mut r := f!"Goal:"
-  let nodes ← getENodes
+def Goal.ppState (goal : Goal) : MetaM MessageData := do
+  let mut r := m!"Goal:"
+  let nodes := goal.getENodes
   for node in nodes do
-    r := r ++ "\n" ++ (← ppENodeDecl node.self)
-  let eqcs ← getEqcs
+    r := r ++ "\n" ++ (← goal.ppENodeDecl node.self)
+  let eqcs := goal.getEqcs
   for eqc in eqcs do
     if eqc.length > 1 then
-      r := r ++ "\n" ++ "{" ++ (Format.joinSep (← eqc.mapM ppENodeRef) ", ") ++  "}"
+      r := r ++ "\n" ++ "{" ++ (MessageData.joinSep (← eqc.mapM goal.ppENodeRef) ", ") ++  "}"
   return r
 
-def ppGoals (goals : List Goal) : GrindM Format := do
-  let mut r := f!""
+def ppGoals (goals : List Goal) : MetaM MessageData := do
+  let mut r := m!""
   for goal in goals do
-    let (f, _) ← GoalM.run goal ppState
-    r := r ++ Format.line ++ f
+    let m ← goal.ppState
+    r := r ++ Format.line ++ m
   return r
 
+private def ppExprArray (cls : Name) (header : String) (es : Array Expr) (clsElem : Name := Name.mkSimple "_") : MessageData :=
+  let es := es.map fun e => .trace { cls := clsElem} m!"{e}" #[]
+  .trace { cls } header es
+
+private def ppEqcs (goal : Goal) : MetaM (Array MessageData) := do
+   let mut trueEqc?  : Option MessageData := none
+   let mut falseEqc? : Option MessageData := none
+   let mut otherEqcs : Array MessageData := #[]
+   for eqc in goal.getEqcs do
+     if Option.isSome <| eqc.find? (·.isTrue) then
+       let eqc := eqc.filter fun e => !e.isTrue
+       unless eqc.isEmpty do
+         trueEqc? := ppExprArray `eqc "True propositions" eqc.toArray `prop
+     else if Option.isSome <| eqc.find? (·.isFalse) then
+       let eqc := eqc.filter fun e => !e.isFalse
+       unless eqc.isEmpty do
+         falseEqc? := ppExprArray `eqc "False propositions" eqc.toArray `prop
+     else if let e :: _ :: _ := eqc then
+       -- We may want to add a flag to pretty print equivalence classes of nested proofs
+       unless (← isProof e) do
+         otherEqcs := otherEqcs.push <| .trace { cls := `eqc } (.group ("{" ++ (MessageData.joinSep (eqc.map toMessageData) ("," ++ Format.line)) ++  "}")) #[]
+   let mut result := #[]
+   if let some trueEqc := trueEqc? then result := result.push trueEqc
+   if let some falseEqc := falseEqc? then result := result.push falseEqc
+   unless otherEqcs.isEmpty do
+     result := result.push <| .trace { cls := `eqc } "Equivalence classes" otherEqcs
+   return result
+
+def goalToMessageData (goal : Goal) : MetaM MessageData := goal.mvarId.withContext do
+  let mut m : Array MessageData := #[.ofGoal goal.mvarId]
+  m := m.push <| ppExprArray `facts "Asserted facts" goal.facts.toArray `prop
+  m := m ++ (← ppEqcs goal)
+  addMessageContextFull <| MessageData.joinSep m.toList ""
+
+def goalsToMessageData (goals : List Goal) : MetaM MessageData :=
+  return MessageData.joinSep (← goals.mapM goalToMessageData) m!"\n"
+
 end Lean.Meta.Grind