feat: improve support for match-expressions in grind (#6779)

This PR improves the support for `match`-expressions in the `grind`
tactic.

src/Init/Grind/Util.lean

@@ -12,11 +12,10 @@ namespace Lean.Grind
 def nestedProof (p : Prop) {h : p} : p := h
 
 /--
-Gadget for marking terms that should not be normalized by `grind`s simplifier.
-`grind` uses a simproc to implement this feature.
+Gadget for marking `match`-expressions that should not be reduced by the `grind` simplifier, but the discriminants should be normalized.
 We use it when adding instances of `match`-equations to prevent them from being simplified to true.
 -/
-def doNotSimp {α : Sort u} (a : α) : α := a
+def simpMatchDiscrsOnly {α : Sort u} (a : α) : α := a
 
 /-- Gadget for representing offsets `t+k` in patterns. -/
 def offset (a b : Nat) : Nat := a + b

src/Lean/Meta/Tactic/Grind.lean

@@ -27,7 +27,7 @@ import Lean.Meta.Tactic.Grind.CasesMatch
 import Lean.Meta.Tactic.Grind.Arith
 import Lean.Meta.Tactic.Grind.Ext
 import Lean.Meta.Tactic.Grind.MatchCond
-import Lean.Meta.Tactic.Grind.DoNotSimp
+import Lean.Meta.Tactic.Grind.MatchDiscrOnly
 
 namespace Lean
 

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

@@ -6,7 +6,7 @@ Authors: Leonardo de Moura
 prelude
 import Lean.Meta.Tactic.Grind.Types
 import Lean.Meta.Tactic.Grind.Intro
-import Lean.Meta.Tactic.Grind.DoNotSimp
+import Lean.Meta.Tactic.Grind.MatchDiscrOnly
 import Lean.Meta.Tactic.Grind.MatchCond
 import Lean.Meta.Tactic.Grind.Core
 
@@ -235,7 +235,8 @@ private def annotateMatchEqnType (prop : Expr) (initApp : Expr) : M Expr := do
   annotateEqnTypeConds prop fun prop => do
     let_expr f@Eq α lhs rhs := prop | return prop
     -- See comment at `Grind.EqMatch`
-    return mkApp4 (mkConst ``Grind.EqMatch f.constLevels!) α (← markAsDoNotSimp lhs) rhs initApp
+    let lhs ← markAsSimpMatchDiscrsOnly lhs
+    return mkApp4 (mkConst ``Grind.EqMatch f.constLevels!) α lhs rhs initApp
 
 /--
 Stores new theorem instance in the state.
@@ -282,6 +283,11 @@ private partial def instantiateTheorem (c : Choice) : M Unit := withDefault do w
         let some heq ← proveEq? vType mvarIdType
           | report
             return ()
+        /-
+        Some of the `cast`s will appear inside the `Grind.MatchCond` binders, and
+        we want their proofs to be properly wrapped.
+        -/
+        let heq := mkApp2 (mkConst ``Grind.nestedProof) (← mkEq vType mvarIdType) heq
         v ← mkAppM ``cast #[heq, v]
       unless (← mvarId.checkedAssign v) do
         report

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

@@ -0,0 +1,42 @@
+/-
+Copyright (c) 2025 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.Grind.Util
+import Init.Simproc
+import Lean.Meta.Tactic.Simp.Simproc
+import Lean.Meta.Tactic.Simp.Rewrite
+
+namespace Lean.Meta.Grind
+/--
+Returns `Grind.simpMatchDiscrsOnly e`. Recall that `Grind.simpMatchDiscrsOnly` is
+a gadget for instructing the `grind` simplifier to only normalize/simplify
+the discriminants of a `match`-expression. See `reduceSimpMatchDiscrsOnly`.
+-/
+def markAsSimpMatchDiscrsOnly (e : Expr) : MetaM Expr :=
+  mkAppM ``Grind.simpMatchDiscrsOnly #[e]
+
+builtin_simproc_decl reduceSimpMatchDiscrsOnly (Grind.simpMatchDiscrsOnly _) := fun e => do
+  let_expr Grind.simpMatchDiscrsOnly _ m ← e | return .continue
+  let .const declName _ := m.getAppFn
+    | return .done { expr := e }
+  let some info ← getMatcherInfo? declName
+    | return .done { expr := e }
+  if let some r ← Simp.simpMatchDiscrs? info m then
+    return .done { r with expr := (← markAsSimpMatchDiscrsOnly r.expr) }
+  return .done { expr := e }
+
+/-- Adds `reduceSimpMatchDiscrsOnly` to `s` -/
+def addSimpMatchDiscrsOnly (s : Simprocs) : CoreM Simprocs := do
+  s.add ``reduceSimpMatchDiscrsOnly (post := false)
+
+/-- Erases `Grind.simpMatchDiscrsOnly` annotations. -/
+def eraseSimpMatchDiscrsOnly (e : Expr) : CoreM Expr := do
+  let pre (e : Expr) := do
+    let_expr Grind.simpMatchDiscrsOnly _ a := e | return .continue e
+    return .continue a
+  Core.transform e (pre := pre)
+
+end Lean.Meta.Grind

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

@@ -9,7 +9,7 @@ import Lean.Meta.Tactic.Assert
 import Lean.Meta.Tactic.Simp.Main
 import Lean.Meta.Tactic.Grind.Util
 import Lean.Meta.Tactic.Grind.Types
-import Lean.Meta.Tactic.Grind.DoNotSimp
+import Lean.Meta.Tactic.Grind.MatchDiscrOnly
 import Lean.Meta.Tactic.Grind.MarkNestedProofs
 import Lean.Meta.Tactic.Grind.Canon
 
@@ -35,7 +35,7 @@ def simp (e : Expr) : GoalM Simp.Result := do
   let e' ← eraseIrrelevantMData e'
   let e' ← foldProjs e'
   let e' ← normalizeLevels e'
-  let e' ← eraseDoNotSimp e'
+  let e' ← eraseSimpMatchDiscrsOnly e'
   let e' ← canon e'
   let e' ← shareCommon e'
   trace[grind.simp] "{e}\n===>\n{e'}"

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

@@ -6,7 +6,7 @@ Authors: Leonardo de Moura
 prelude
 import Lean.Meta.Tactic.Simp.Simproc
 import Lean.Meta.Tactic.Grind.Simp
-import Lean.Meta.Tactic.Grind.DoNotSimp
+import Lean.Meta.Tactic.Grind.MatchDiscrOnly
 import Lean.Meta.Tactic.Grind.MatchCond
 import Lean.Meta.Tactic.Simp.BuiltinSimprocs.List
 
@@ -40,7 +40,7 @@ protected def getSimprocs : MetaM (Array Simprocs) := do
   We don't want it to be simplified to `[] = []`.
   -/
   let s := s.erase ``List.reduceReplicate
-  let s ← addDoNotSimp s
+  let s ← addSimpMatchDiscrsOnly s
   let s ← addMatchCond s
   return #[s]