chore: improve heuristic

src/Init/Grind/Tactics.lean

@@ -51,6 +51,8 @@ structure Config where
   failures : Nat := 1
   /-- Maximum number of heartbeats (in thousands) the canonicalizer can spend per definitional equality test. -/
   canonHeartbeats : Nat := 1000
+  /-- If `ext` is `true`, `grind` uses extensionality theorems available in the environment. -/
+  ext : Bool := true
   deriving Inhabited, BEq
 
 end Lean.Grind

src/Lean/Meta/Tactic/Grind.lean

@@ -25,6 +25,7 @@ import Lean.Meta.Tactic.Grind.EMatch
 import Lean.Meta.Tactic.Grind.Main
 import Lean.Meta.Tactic.Grind.CasesMatch
 import Lean.Meta.Tactic.Grind.Arith
+import Lean.Meta.Tactic.Grind.Ext
 
 namespace Lean
 

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

@@ -266,7 +266,7 @@ private partial def instantiateTheorem (c : Choice) : M Unit := withDefault do w
   for mvar in mvars, bi in bis do
     if bi.isInstImplicit && !(← mvar.mvarId!.isAssigned) then
       let type ← inferType mvar
-      unless (← synthesizeInstance mvar type) do
+      unless (← synthesizeInstanceAndAssign mvar type) do
         reportIssue m!"failed to synthesize instance when instantiating {← thm.origin.pp}{indentExpr type}"
         return ()
   let proof := mkAppN proof mvars
@@ -278,10 +278,6 @@ private partial def instantiateTheorem (c : Choice) : M Unit := withDefault do w
       reportIssue m!"failed to instantiate {← thm.origin.pp}, failed to instantiate non propositional argument with type{indentExpr (← inferType mvarBad)}"
     let proof ← mkLambdaFVars (binderInfoForMVars := .default) mvars (← instantiateMVars proof)
     addNewInstance thm.origin proof c.gen
-where
-  synthesizeInstance (x type : Expr) : MetaM Bool := do
-    let .some val ← trySynthInstance type | return false
-    isDefEq x val
 
 /-- Process choice stack until we don't have more choices to be processed. -/
 private def processChoices : M Unit := do

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

@@ -7,9 +7,7 @@ prelude
 import Lean.Meta.AppBuilder
 
 namespace Lean.Meta.Grind
-/-!
-A basic "equality resolution" procedure to make Kim happy.
--/
+/-! A basic "equality resolution" procedure. -/
 
 private def eqResCore (prop proof : Expr) : MetaM (Option (Expr × Expr)) := withNewMCtxDepth do
   let (ms, _, type) ← forallMetaTelescopeReducing prop

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

@@ -0,0 +1,40 @@
+/-
+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 Lean.Meta.Tactic.Grind.Types
+
+namespace Lean.Meta.Grind
+/-! Extensionality theorems support. -/
+
+def instantiateExtTheorem (thm : Ext.ExtTheorem) (e : Expr) : GoalM Unit := withNewMCtxDepth do
+  unless (← getGeneration e) < (← getMaxGeneration) do return ()
+  let c ← mkConstWithFreshMVarLevels thm.declName
+  let (mvars, bis, type) ← withDefault <| forallMetaTelescopeReducing (← inferType c)
+  unless (← isDefEq e type) do
+    reportIssue m!"failed to apply extensionality theorem `{thm.declName}` for {indentExpr e}\nis not definitionally equal to{indentExpr type}"
+    return ()
+  -- Instantiate type class instances
+  for mvar in mvars, bi in bis do
+    if bi.isInstImplicit && !(← mvar.mvarId!.isAssigned) then
+      let type ← inferType mvar
+      unless (← synthesizeInstanceAndAssign mvar type) do
+        reportIssue m!"failed to synthesize instance when instantiating extensionality theorem `{thm.declName}` for {indentExpr e}"
+        return ()
+  -- Remark: `proof c mvars` has type `e`
+  let proof ← instantiateMVars (mkAppN c mvars)
+  -- `e` is equal to `False`
+  let eEqFalse ← mkEqFalseProof e
+  -- So, we use `Eq.mp` to build a `proof` of `False`
+  let proof ← mkEqMP eEqFalse proof
+  let mvars ← mvars.filterM fun mvar => return !(← mvar.mvarId!.isAssigned)
+  let proof' ← instantiateMVars (← mkLambdaFVars mvars proof)
+  let prop' ← inferType proof'
+  if proof'.hasMVar || prop'.hasMVar then
+    reportIssue m!"failed to apply extensionality theorem `{thm.declName}` for {indentExpr e}\nresulting terms contain metavariables"
+    return ()
+  addNewFact proof' prop' ((← getGeneration e) + 1)
+
+end Lean.Meta.Grind

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

@@ -8,6 +8,7 @@ import Init.Grind
 import Lean.Meta.Tactic.Grind.Proof
 import Lean.Meta.Tactic.Grind.PropagatorAttr
 import Lean.Meta.Tactic.Grind.Simp
+import Lean.Meta.Tactic.Grind.Ext
 import Lean.Meta.Tactic.Grind.Internalize
 
 namespace Lean.Meta.Grind
@@ -133,6 +134,19 @@ builtin_grind_propagator propagateEqDown ↓Eq := fun e => do
   if (← isEqTrue e) then
     let_expr Eq _ a b := e | return ()
     pushEq a b <| mkOfEqTrueCore e (← mkEqTrueProof e)
+  else if (← isEqFalse e) then
+    let_expr Eq α lhs rhs := e | return ()
+    let thms ← getExtTheorems α
+    if !thms.isEmpty then
+      /-
+      Heuristic for lists: If `lhs` or `rhs` are contructors we do not apply extensionality theorems.
+      For example, we don't want to apply the extensionality theorem to things like `xs ≠ []`.
+      TODO: polish this hackish heuristic later.
+      -/
+      if α.isAppOf ``List && ((← getRootENode lhs).ctor || (← getRootENode rhs).ctor) then
+        return ()
+      for thm in (← getExtTheorems α) do
+        instantiateExtTheorem thm e
 
 /-- Propagates `EqMatch` downwards -/
 builtin_grind_propagator propagateEqMatchDown ↓Grind.EqMatch := fun e => do

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

@@ -13,6 +13,7 @@ import Lean.Meta.CongrTheorems
 import Lean.Meta.AbstractNestedProofs
 import Lean.Meta.Tactic.Simp.Types
 import Lean.Meta.Tactic.Util
+import Lean.Meta.Tactic.Ext
 import Lean.Meta.Tactic.Grind.ENodeKey
 import Lean.Meta.Tactic.Grind.Attr
 import Lean.Meta.Tactic.Grind.Arith.Types
@@ -395,6 +396,8 @@ structure Goal where
   users when `grind` fails.
   -/
   issues     : List MessageData := []
+  /-- Cached extensionality theorems for types. -/
+  extThms    : PHashMap ENodeKey (Array Ext.ExtTheorem) := {}
   deriving Inhabited
 
 def Goal.admit (goal : Goal) : MetaM Unit :=
@@ -886,4 +889,25 @@ def markCaseSplitAsResolved (e : Expr) : GoalM Unit := do
     trace_goal[grind.split.resolved] "{e}"
     modify fun s => { s with resolvedSplits := s.resolvedSplits.insert { expr := e } }
 
+/--
+Returns extensionality theorems for the given type if available.
+If `Config.ext` is `false`, the result is `#[]`.
+-/
+def getExtTheorems (type : Expr) : GoalM (Array Ext.ExtTheorem) := do
+  unless (← getConfig).ext do return #[]
+  if let some thms := (← get).extThms.find? { expr := type } then
+    return thms
+  else
+    let thms ← Ext.getExtTheorems type
+    modify fun s => { s with extThms := s.extThms.insert { expr := type } thms }
+    return thms
+
+/--
+Helper function for instantiating a type class `type`, and
+then using the result to perform `isDefEq x val`.
+-/
+def synthesizeInstanceAndAssign (x type : Expr) : MetaM Bool := do
+  let .some val ← trySynthInstance type | return false
+  isDefEq x val
+
 end Lean.Meta.Grind

tests/lean/run/grind_t1.lean

@@ -293,3 +293,27 @@ example {α} (f : α → Type) (a : α) (h : ∀ x, Nonempty (f x)) : Nonempty (
 
 example {α β} (f : α → β) (a : α) : ∃ a', f a' = f a := by
   grind
+
+open List in
+example : (replicate n a).map f = replicate n (f a) := by
+  grind only [Option.map_some', Option.map_none', getElem?_map, getElem?_replicate]
+
+open List in
+example : (replicate n a).map f = replicate n (f a) := by
+  -- Should fail since extensionality is disabled
+  fail_if_success grind -ext only [Option.map_some', Option.map_none', getElem?_map, getElem?_replicate]
+  sorry
+
+@[ext] structure S where
+  a : Nat
+  b : Bool
+
+example (x y : S) : x.a = y.a → y.b = x.b → x = y := by
+  grind
+
+example (x y : S) : x.a = y.a → y.b = x.b → x = y := by
+  fail_if_success grind -ext
+  sorry
+
+example (x : S) : x.a = 10 → false ≠ x.b → x = { a := 10, b := true } := by
+  grind