From 9a7565d66c84bdef37f8cd021fe6ca6eebe7b425 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Sat, 21 Oct 2023 06:11:55 -0700
Subject: [PATCH] perf: closes #2552

---
 src/Lean/Meta/Tactic/UnifyEq.lean |  9 +++++--
 tests/lean/run/2552.lean          | 44 +++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 2 deletions(-)
 create mode 100644 tests/lean/run/2552.lean

diff --git a/src/Lean/Meta/Tactic/UnifyEq.lean b/src/Lean/Meta/Tactic/UnifyEq.lean
index ee63627e67..edf83975e0 100644
--- a/src/Lean/Meta/Tactic/UnifyEq.lean
+++ b/src/Lean/Meta/Tactic/UnifyEq.lean
@@ -53,8 +53,13 @@ def unifyEq? (mvarId : MVarId) (eqFVarId : FVarId) (subst : FVarSubst := {})
         /- Remark: we use `let rec` here because otherwise the compiler would generate an insane amount of code.
           We can remove the `rec` after we fix the eagerly inlining issue in the compiler. -/
         let rec substEq (symm : Bool) := do
-          /- Remark: `substCore` fails if the equation is of the form `x = x` -/
-          if let some (subst, mvarId) ← observing? (substCore mvarId eqFVarId symm subst) then
+          /-
+          Remark: `substCore` fails if the equation is of the form `x = x`
+          We use `(tryToSkip := true)` to ensure we avoid unnecessary `Eq.rec`s in user code.
+          Recall that `Eq.rec`s can negatively affect term reduction performance in the kernel and elaborator.
+          See issue https://github.com/leanprover/lean4/issues/2552 for an example.
+          -/
+          if let some (subst, mvarId) ← observing? (substCore mvarId eqFVarId symm subst (tryToSkip := true)) then
             return some { mvarId, subst }
           else if (← isDefEq a b) then
             /- Skip equality -/
diff --git a/tests/lean/run/2552.lean b/tests/lean/run/2552.lean
new file mode 100644
index 0000000000..3053252dbd
--- /dev/null
+++ b/tests/lean/run/2552.lean
@@ -0,0 +1,44 @@
+@[inline] def decidable_of_iff {b : Prop} (a : Prop) (h : a ↔ b) [Decidable a] : Decidable b :=
+  decidable_of_decidable_of_iff h
+
+theorem LE.le.lt_or_eq_dec {a b : Nat} (hab : a ≤ b) : a < b ∨ a = b :=
+by cases hab
+   · exact .inr rfl
+   · refine .inl ?_ ; refine Nat.succ_le_succ ?_ ; assumption
+
+section succeeds_using_match
+
+local instance decidableBallLT :
+  ∀ (n : Nat) (P : ∀ k, k < n → Prop) [∀ n h, Decidable (P n h)], Decidable (∀ n h, P n h)
+| 0, _, _ => isTrue fun _ => (by cases ·)
+| (n+1), P, H =>
+  match decidableBallLT n (P · <| Nat.le_succ_of_le ·) with
+  | isFalse h => isFalse (h fun _ _ => · _ _)
+  | isTrue h =>
+    match H n Nat.le.refl with
+    | isFalse p => isFalse (p <| · _ _)
+    | isTrue p => isTrue fun _ h' => (Nat.le_of_lt_succ h').lt_or_eq_dec.elim (h _) (· ▸ p)
+
+set_option maxHeartbeats 5000
+example : ∀ a, a < 9 → ∀ b, b < 9 → ∀ c, c < 9 → a ^ 2 + b ^ 2 + c ^ 2 ≠ 7 := by decide
+
+end succeeds_using_match
+
+section fails_with_timeout
+
+-- we change `match decidableBallLT` to `by cases decidableBallLT' ... exact`:
+local instance decidableBallLT' :
+  ∀ (n : Nat) (P : ∀ k, k < n → Prop) [∀ n h, Decidable (P n h)], Decidable (∀ n h, P n h)
+| 0, _, _ => isTrue fun _ => (by cases ·)
+| (n+1), P, H => by
+  cases decidableBallLT' n (P · <| Nat.le_succ_of_le ·) with
+  | isFalse h => exact isFalse (h fun _ _ => · _ _)
+  | isTrue h =>
+    exact match H n Nat.le.refl with
+    | isFalse p => isFalse (p <| · _ _)
+    | isTrue p => isTrue fun _ h' => (Nat.le_of_lt_succ h').lt_or_eq_dec.elim (h _) (· ▸ p)
+
+set_option maxHeartbeats 5000
+example : ∀ a, a < 9 → ∀ b, b < 9 → ∀ c, c < 9 → a ^ 2 + b ^ 2 + c ^ 2 ≠ 7 := by decide
+
+end fails_with_timeout