fix: Function.Injective initialization in grind (#11101)

This PR fixes an initialization issue for local `Function.Injective f`
hypotheses.

Closes #11088

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

@@ -11,36 +11,40 @@ import Init.Grind.Injective
 import Lean.Meta.Tactic.Grind.PropagatorAttr
 import Lean.Meta.Tactic.Grind.Simp
 namespace Lean.Meta.Grind
+
+def getInvFor? (f : Expr) (a : Expr) : GoalM (Option (Expr × Expr)) := do
+  let some info := (← get).inj.fns.find? { expr := f } | return none
+  if let some inv := info.inv? then
+    return some inv
+  let [u, _] := info.us | unreachable!
+  let nonEmpty := mkApp2 (mkConst ``Nonempty.intro [u]) info.α a
+  let inv := mkApp5 (mkConst ``Grind.leftInv info.us) info.α info.β f info.h nonEmpty
+  let inv ← preprocessLight inv
+  let args := inv.getAppArgs
+  let heq := mkAppN (mkConst ``Grind.leftInv_eq info.us) args
+  let inv? := some (inv, heq)
+  let info := { info with inv? }
+  modify fun s => { s with inj.fns := s.inj.fns.insert { expr := f } info }
+  return inv?
+
 /--
 If `e` is an application of the form `f a` where `f` is an injective function
 in `(← get).inj.fns`, asserts `f⁻¹ (f a) = a`
 -/
 public def mkInjEq (e : Expr) : GoalM Unit := do
   let .app f a := e | return ()
-  let some info := (← get).inj.fns.find? { expr := f } | return ()
-  let invApp := mkApp info.inv e
+  let some (inv, heq) ← getInvFor? f a | return ()
+  let invApp := mkApp inv e
   internalize invApp (← getGeneration e)
   trace[grind.inj.assert] "{invApp}, {a}"
-  pushEq invApp a <| mkApp info.heq a
+  pushEq invApp a <| mkApp heq a
 
 def initInjFn (us : List Level) (α β : Expr) (f : Expr) (h : Expr) : GoalM Unit := do
+  modify fun s => { s with inj.fns := s.inj.fns.insert { expr := f } { us, α, β, h } }
   let hidx := f.toHeadIndex
-  let mut first := true
   for e in (← get).appMap.findD hidx [] do
     if e.isApp && isSameExpr e.appFn! f then
-      if first then
-        initLeftInv e.appArg!
-        first := false
       mkInjEq e
-where
-  initLeftInv (a : Expr) : GoalM Unit := do
-    let [u, _] := us | unreachable!
-    let nonEmpty := mkApp2 (mkConst ``Nonempty.intro [u]) α a
-    let inv := mkApp5 (mkConst ``Grind.leftInv us) α β f h nonEmpty
-    let inv ← preprocessLight inv
-    let args := inv.getAppArgs
-    let heq := mkAppN (mkConst ``Grind.leftInv_eq us) args
-    modify fun s => { s with inj.fns := s.inj.fns.insert { expr := f } { inv, heq } }
 
 builtin_grind_propagator propagateInj ↓Function.Injective := fun e => do
   let_expr i@Function.Injective α β f := e | return ()

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

@@ -859,8 +859,15 @@ structure UnitLike.State where
   deriving Inhabited
 
 structure InjectiveInfo where
-  inv : Expr
-  heq : Expr
+  us   : List Level
+  α    : Expr
+  β    : Expr
+  h    : Expr
+  /--
+  Inverse function and a proof that `∀ a, inv (f a) = a`
+  **Note**: The following two fields are `none` if no `f`-application has been found yet.
+  -/
+  inv?  : Option (Expr × Expr) := none
   deriving Inhabited
 
 /-- State for injective theorem support. -/

tests/lean/run/grind_11088.lean

@@ -0,0 +1,3 @@
+example (f : α → β) (h : Function.Injective f)
+    : f a = f b → a = b := by
+  grind