perf: add introSubstEq shortcut (#12190)

This PR adds the `introSubstEq` MetaM tactic, as an optimization over
`intro h; subst h` that avoids introducing `h : a = b` if it can be
avoided,
which is the case when `b` can be reverted without reverting anything
else. Speeds up the generation of `injEq` theorem.

src/Init/Core.lean

@@ -932,6 +932,14 @@ noncomputable def HEq.ndrec.{u1, u2} {α : Sort u2} {a : α} {motive : {β : Sor
 noncomputable def HEq.ndrecOn.{u1, u2} {α : Sort u2} {a : α} {motive : {β : Sort u2} → β → Sort u1} {β : Sort u2} {b : β} (h : a ≍ b) (m : motive a) : motive b :=
   h.rec m
 
+/-- `HEq.ndrec` specialized to homogeneous heterogeneous equality -/
+noncomputable def HEq.homo_ndrec.{u1, u2} {α : Sort u2} {a : α} {motive : α → Sort u1} (m : motive a) {b : α} (h : a ≍ b) : motive b :=
+  (eq_of_heq h).ndrec m
+
+/-- `HEq.ndrec` specialized to homogeneous heterogeneous equality, symmetric variant -/
+noncomputable def HEq.homo_ndrec_symm.{u1, u2} {α : Sort u2} {a : α} {motive : α → Sort u1} (m : motive a) {b : α} (h : b ≍ a) : motive b :=
+  (eq_of_heq h).ndrec_symm m
+
 /-- `HEq.ndrec` variant -/
 noncomputable def HEq.elim {α : Sort u} {a : α} {p : α → Sort v} {b : α} (h₁ : a ≍ b) (h₂ : p a) : p b :=
   eq_of_heq h₁ ▸ h₂

src/Init/Prelude.lean

@@ -322,6 +322,10 @@ For more information: [Equality](https://lean-lang.org/theorem_proving_in_lean4/
 @[symm] theorem Eq.symm {α : Sort u} {a b : α} (h : Eq a b) : Eq b a :=
   h ▸ rfl
 
+/-- Non-dependent recursor for the equality type (symmetric variant) -/
+@[simp] abbrev Eq.ndrec_symm.{u1, u2} {α : Sort u2} {a : α} {motive : α → Sort u1} (m : motive a) {b : α} (h : Eq b a) : motive b :=
+  h.symm.ndrec m
+
 /--
 Equality is transitive: if `a = b` and `b = c` then `a = c`.
 

src/Lean/AuxRecursor.lean

@@ -33,6 +33,7 @@ def isAuxRecursor (env : Environment) (declName : Name) : Bool :=
   -- TODO: use `markAuxRecursor` when they are defined
   -- An attribute is not a good solution since we don't want users to control what is tagged as an auxiliary recursor.
   || declName == ``Eq.ndrec
+  || declName == ``Eq.ndrec_symm
   || declName == ``Eq.ndrecOn
 
 def isAuxRecursorWithSuffix (env : Environment) (declName : Name) (suffix : String) : Bool :=

src/Lean/Meta/Injective.lean

@@ -157,8 +157,7 @@ private def mkInjectiveEqTheoremValue (ctorVal : ConstructorVal) (targetType : E
             | throwError "unexpected number of goals after applying `Lean.and_imp`"
         mvarId₂ := mvarId₂'
       | _ => pure ()
-      let (h, mvarId₂') ← mvarId₂.intro1
-      (_, mvarId₂) ← substEq mvarId₂' h
+      (_, mvarId₂) ← introSubstEq mvarId₂
     try mvarId₂.refl catch _ => throwError (injTheoremFailureHeader ctorVal.name)
     mkLambdaFVars xs mvar
 

src/Lean/Meta/Match/CaseValues.lean

@@ -62,9 +62,8 @@ public def caseValues (mvarId : MVarId) (fvarId : FVarId) (values : Array Expr)
       let (thenMVarId, elseMVarId) ← caseValue mvarId fvarId v (hNamePrefix.appendIndexAfter i)
       appendTagSuffix thenMVarId ((`case).appendIndexAfter i)
       let thenMVarId ← thenMVarId.tryClearMany hs
-      let (thenH, thenMVarId) ← thenMVarId.intro1P
-      let (subst, thenMVarId) ← substCore thenMVarId thenH (symm := false) {} (clearH := true)
-      let subgoals := subgoals.push { mvarId := thenMVarId, newHs := #[], subst := subst }
+      let (subst, thenMVarId) ← introSubstEq thenMVarId (substLHS := true)
+      let subgoals := subgoals.push { mvarId := thenMVarId, newHs := #[], subst }
       let (hs', elseMVarId) ←
         if needHyps then
           let (elseH, elseMVarId) ← elseMVarId.intro1P

src/Lean/Meta/Match/MatchEqs.lean

@@ -72,8 +72,7 @@ partial def proveCondEqThm (matchDeclName : Name) (type : Expr)
   if heqNum > 0 then
     mvarId := (← mvarId.introN heqPos).2
     for _ in *...heqNum do
-      let (h, mvarId') ← mvarId.intro1
-      mvarId ← subst mvarId' h
+      (_, mvarId) ← introSubstEq mvarId
     trace[Meta.Match.matchEqs] "proveCondEqThm after subst{mvarId}"
   mvarId := (← mvarId.intros).2
   try mvarId.refl

src/Lean/Meta/Tactic/Subst.lean

@@ -212,6 +212,54 @@ partial def subst (mvarId : MVarId) (h : FVarId) : MetaM MVarId :=
           subst mvarId' h'
       | none => substVar mvarId h
 
+/--
+Given a goal `(a = b) → goal[b]`, creates a new goal `goal[a]`, clearing `b`.
+
+This is essentially `intro h; subst h`, but in the case that `b` is a free variable and has no
+forward dependencies implements this without introducing the equality, which can make a difference
+in terms of performance.
+
+If `substLHS = true`, assume `(a = b) → goal[a]` and create goal `goal[b]`, clearing `a`.
+
+Also handles heterogeneous equalities in cases where `eq_of_heq` would apply.
+-/
+def introSubstEq (mvarId : MVarId) (substLHS := false) : MetaM (FVarSubst × MVarId) := do
+  mvarId.checkNotAssigned  `introSubstEq
+  try commitIfNoEx do mvarId.withContext do
+    let goalType ← mvarId.getType'
+    let some (heq, body) := goalType.arrow? | throwError "not an arrow type"
+    let (α, a, b, ndrec) ←
+      match_expr heq with
+      | Eq α a b =>
+        if substLHS then
+          pure (α, b, a, ``Eq.ndrec_symm)
+        else
+          pure (α, a, b, ``Eq.ndrec)
+      | HEq α a β b =>
+        unless (← isDefEq α β) do throwError "hetereogenenous equality isn't homogeneous"
+        if substLHS then
+          pure (α, b, a, ``HEq.homo_ndrec_symm)
+        else
+          pure (α, a, b, ``HEq.homo_ndrec)
+      | _ => throwError "not an equality"
+    unless b.isFVar do throwError "equality rhs not a free variable"
+    let (reverted, mvarId) ← mvarId.revert #[b.fvarId!]
+    unless reverted.size = 1 do throwError "variable {b} has forward dependencies"
+    let motive ← mkLambdaFVars #[b] body
+    let goal := motive.beta #[a]
+    let e ← mkFreshExprSyntheticOpaqueMVar goal (tag := (← mvarId.getTag))
+    let u1 ← getLevel goal
+    let u2 ← getLevel α
+    mvarId.assign <| mkApp4 (mkConst ndrec [u1, u2]) α a motive e
+    let subst : FVarSubst := FVarSubst.empty.insert b.fvarId! a
+    return (subst, e.mvarId!)
+  catch e =>
+    trace[Meta.Tactic.subst] "introSubstEq falling back to intro\n{e.toMessageData}\n{mvarId}"
+    if (← mvarId.isAssigned) then throwError "introSubstEq: now assigned?"
+    let (h, mvarId) ← mvarId.intro1
+    let (subst, mvarId) ← substEq mvarId h
+    return (subst, mvarId)
+
 /--
 Given `x`, try to find an equation of the form `heq : x = rhs` or `heq : lhs = x`,
 and runs `substCore` on it. Returns `none` if no such equation is found, or if `substCore` fails.

tests/bench/big_struct.lean

@@ -202,3 +202,4 @@ structure M where
   b197 : Unit
   b198 : Unit
   b199 : Unit
+