From 702c9875eedadf0d3c008159d669d359845a9c28 Mon Sep 17 00:00:00 2001
From: Gabriel Ebner <gebner@gebner.org>
Date: Tue, 3 Jan 2023 14:30:55 -0800
Subject: [PATCH] perf: whnf projections during defeq

---
 src/Lean/Meta/ExprDefEq.lean |  42 ++++----
 src/Lean/Meta/WHNF.lean      |   1 -
 tests/lean/run/1986.lean     | 204 +++++++++++++++++++++++++++++++++++
 3 files changed, 226 insertions(+), 21 deletions(-)
 create mode 100644 tests/lean/run/1986.lean

diff --git a/src/Lean/Meta/ExprDefEq.lean b/src/Lean/Meta/ExprDefEq.lean
index 9e1f2f8fd1a..73b328ab224 100644
--- a/src/Lean/Meta/ExprDefEq.lean
+++ b/src/Lean/Meta/ExprDefEq.lean
@@ -1732,24 +1732,19 @@ private def isExprDefEqExpensive (t : Expr) (s : Expr) : MetaM Bool := do
   -- TODO: investigate whether this is the place for putting this check
   if (← (isDefEqEtaStruct t s <||> isDefEqEtaStruct s t)) then return true
   if (← isDefEqProj t s) then return true
-  let t' ← whnfCore t
-  let s' ← whnfCore s
-  if t != t' || s != s' then
-    Meta.isExprDefEqAux t' s'
+  whenUndefDo (isDefEqNative t s) do
+  whenUndefDo (isDefEqNat t s) do
+  whenUndefDo (isDefEqOffset t s) do
+  whenUndefDo (isDefEqDelta t s) do
+  if t.isConst && s.isConst then
+    if t.constName! == s.constName! then isListLevelDefEqAux t.constLevels! s.constLevels! else return false
+  else if (← pure t.isApp <&&> pure s.isApp <&&> isDefEqApp t s) then
+    return true
   else
-    whenUndefDo (isDefEqNative t s) do
-    whenUndefDo (isDefEqNat t s) do
-    whenUndefDo (isDefEqOffset t s) do
-    whenUndefDo (isDefEqDelta t s) do
-    if t.isConst && s.isConst then
-      if t.constName! == s.constName! then isListLevelDefEqAux t.constLevels! s.constLevels! else return false
-    else if (← pure t.isApp <&&> pure s.isApp <&&> isDefEqApp t s) then
-      return true
-    else
-      whenUndefDo (isDefEqProjInst t s) do
-      whenUndefDo (isDefEqStringLit t s) do
-      if (← isDefEqUnitLike t s) then return true else
-      isDefEqOnFailure t s
+    whenUndefDo (isDefEqProjInst t s) do
+    whenUndefDo (isDefEqStringLit t s) do
+    if (← isDefEqUnitLike t s) then return true else
+    isDefEqOnFailure t s
 
 private def mkCacheKey (t : Expr) (s : Expr) : Expr × Expr :=
   if Expr.quickLt t s then (t, s) else (s, t)
@@ -1774,9 +1769,16 @@ partial def isExprDefEqAuxImpl (t : Expr) (s : Expr) : MetaM Bool := withIncRecD
   checkMaxHeartbeats "isDefEq"
   whenUndefDo (isDefEqQuick t s) do
   whenUndefDo (isDefEqProofIrrel t s) do
-  -- We perform `whnfCore` again with `deltaAtProj := true` at `isExprDefEqExpensive` after `isDefEqProj`
-  let t' ← whnfCore t (deltaAtProj := false)
-  let s' ← whnfCore s (deltaAtProj := false)
+  /-
+    We also reduce projections here to prevent expensive defeq checks when unifying TC operations.
+    When unifying e.g. `@Neg.neg α (@Field.toNeg α inst1) =?= @Neg.neg α (@Field.toNeg α inst2)`,
+    we only want to unify negation (and not all other field operations as well).
+    Unifying the field instances slowed down unification: https://github.com/leanprover/lean4/issues/1986
+    We used to *not* reduce projections here, to support unifying `(?a).1 =?= (x, y).1`.
+    NOTE: this still seems to work because we don't eagerly unfold projection definitions to primitive projections.
+  -/
+  let t' ← whnfCore t
+  let s' ← whnfCore s
   if t != t' || s != s' then
     isExprDefEqAuxImpl t' s'
   else
diff --git a/src/Lean/Meta/WHNF.lean b/src/Lean/Meta/WHNF.lean
index faadca91fea..e5db1e2e04e 100644
--- a/src/Lean/Meta/WHNF.lean
+++ b/src/Lean/Meta/WHNF.lean
@@ -490,7 +490,6 @@ expand let-expressions, expand assigned meta-variables.
 
 The parameter `deltaAtProj` controls how to reduce projections `s.i`. If `deltaAtProj == true`,
 then delta reduction is used to reduce `s` (i.e., `whnf` is used), otherwise `whnfCore`.
-We only set this flag to `false` when implementing `isDefEq`.
 
 If `simpleReduceOnly`, then `iota` and projection reduction are not performed.
 Note that the value of `deltaAtProj` is irrelevant if `simpleReduceOnly = true`.
diff --git a/tests/lean/run/1986.lean b/tests/lean/run/1986.lean
new file mode 100644
index 00000000000..68cf03faaba
--- /dev/null
+++ b/tests/lean/run/1986.lean
@@ -0,0 +1,204 @@
+instance {ι : Type u} {α : ι → Type v} [∀ i, LE (α i)] : LE (∀ i, α i) where
+  le x y := ∀ i, x i ≤ y i
+
+class Top (α : Type u) where
+  top : α
+
+class Bot (α : Type u) where
+  bot : α
+
+notation "⊤" => Top.top
+
+notation "⊥" => Bot.bot
+
+class Preorder (α : Type u) extends LE α, LT α where
+  le_refl : ∀ a : α, a ≤ a
+  le_trans : ∀ a b c : α, a ≤ b → b ≤ c → a ≤ c
+  lt := λ a b => a ≤ b ∧ ¬ b ≤ a
+  lt_iff_le_not_le : ∀ a b : α, a < b ↔ (a ≤ b ∧ ¬ b ≤ a)
+
+class PartialOrder (α : Type u) extends Preorder α :=
+(le_antisymm : ∀ a b : α, a ≤ b → b ≤ a → a = b)
+
+def Set (α : Type u) := α → Prop
+
+def setOf {α : Type u} (p : α → Prop) : Set α :=
+p
+
+namespace Set
+
+protected def Mem (a : α) (s : Set α) : Prop :=
+s a
+
+instance : Membership α (Set α) :=
+⟨Set.Mem⟩
+
+def range (f : ι → α) : Set α :=
+  setOf (λ x => ∃ y, f y = x)
+
+end Set
+
+class InfSet (α : Type _) where
+  infₛ : Set α → α
+
+class SupSet (α : Type _) where
+  supₛ : Set α → α
+
+export SupSet (supₛ)
+
+export InfSet (infₛ)
+
+open Set
+
+def supᵢ {α : Type _} [SupSet α] {ι} (s : ι → α) : α :=
+  supₛ (range s)
+
+def infᵢ {α : Type _} [InfSet α] {ι} (s : ι → α) : α :=
+  infₛ (range s)
+
+class HasSup (α : Type u) where
+  sup : α → α → α
+
+class HasInf (α : Type u) where
+  inf : α → α → α
+
+@[inherit_doc]
+infixl:68 " ⊔ " => HasSup.sup
+
+@[inherit_doc]
+infixl:69 " ⊓ " => HasInf.inf
+
+class SemilatticeSup (α : Type u) extends HasSup α, PartialOrder α where
+  protected le_sup_left : ∀ a b : α, a ≤ a ⊔ b
+  protected le_sup_right : ∀ a b : α, b ≤ a ⊔ b
+  protected sup_le : ∀ a b c : α, a ≤ c → b ≤ c → a ⊔ b ≤ c
+
+class SemilatticeInf (α : Type u) extends HasInf α, PartialOrder α where
+  protected inf_le_left : ∀ a b : α, a ⊓ b ≤ a
+  protected inf_le_right : ∀ a b : α, a ⊓ b ≤ b
+  protected le_inf : ∀ a b c : α, a ≤ b → a ≤ c → a ≤ b ⊓ c
+
+class Lattice (α : Type u) extends SemilatticeSup α, SemilatticeInf α
+
+class CompleteSemilatticeInf (α : Type _) extends PartialOrder α, InfSet α where
+  infₛ_le : ∀ s, ∀ a, a ∈ s → infₛ s ≤ a
+  le_infₛ : ∀ s a, (∀ b, b ∈ s → a ≤ b) → a ≤ infₛ s
+
+class CompleteSemilatticeSup (α : Type _) extends PartialOrder α, SupSet α where
+  le_supₛ : ∀ s, ∀ a, a ∈ s → a ≤ supₛ s
+  supₛ_le : ∀ s a, (∀ b, b ∈ s → b ≤ a) → supₛ s ≤ a
+
+class CompleteLattice (α : Type _) extends Lattice α, CompleteSemilatticeSup α,
+  CompleteSemilatticeInf α, Top α, Bot α where
+  protected le_top : ∀ x : α, x ≤ ⊤
+  protected bot_le : ∀ x : α, ⊥ ≤ x
+
+class Frame (α : Type _) extends CompleteLattice α where
+
+class Coframe (α : Type _) extends CompleteLattice α where
+  infᵢ_sup_le_sup_infₛ (a : α) (s : Set α) : (infᵢ (λ b => a ⊔ b)) ≤ a ⊔ infₛ s
+  -- should be ⨅ b ∈ s but I had problems with notation
+
+class CompleteDistribLattice (α : Type _) extends Frame α where
+  infᵢ_sup_le_sup_infₛ : ∀ a s, (infᵢ (λ b => a ⊔ b)) ≤ a ⊔ infₛ s
+  -- similarly this is not quite right mathematically but this doesn't matter
+
+-- See note [lower instance priority]
+instance (priority := 100) CompleteDistribLattice.toCoframe {α : Type _} [CompleteDistribLattice α] :
+    Coframe α :=
+  { ‹CompleteDistribLattice α› with }
+
+class OrderTop (α : Type u) [LE α] extends Top α where
+  le_top : ∀ a : α, a ≤ ⊤
+
+class OrderBot (α : Type u) [LE α] extends Bot α where
+  bot_le : ∀ a : α, ⊥ ≤ a
+
+class BoundedOrder (α : Type u) [LE α] extends OrderTop α, OrderBot α
+
+instance(priority := 100) CompleteLattice.toBoundedOrder  {α : Type _} [h : CompleteLattice α] :
+    BoundedOrder α :=
+  { h with }
+
+namespace Pi
+
+variable {ι : Type _} {α' : ι → Type _}
+
+instance [∀ i, Bot (α' i)] : Bot (∀ i, α' i) :=
+  ⟨fun _ => ⊥⟩
+
+instance [∀ i, Top (α' i)] : Top (∀ i, α' i) :=
+  ⟨fun _ => ⊤⟩
+
+protected instance LE {ι : Type u} {α : ι → Type v} [∀ i, LE (α i)] : LE (∀ i, α i) where
+  le x y := ∀ i, x i ≤ y i
+
+instance Preorder {ι : Type u} {α : ι → Type v} [∀ i, Preorder (α i)] : Preorder (∀ i, α i) :=
+  { Pi.LE with
+  le_refl := sorry
+  le_trans := sorry
+  lt_iff_le_not_le := sorry }
+
+instance PartialOrder {ι : Type u} {α : ι → Type v} [∀ i, PartialOrder (α i)] :
+    PartialOrder (∀ i, α i) :=
+  { Pi.Preorder with
+  le_antisymm := sorry }
+
+instance semilatticeSup [∀ i, SemilatticeSup (α' i)] : SemilatticeSup (∀ i, α' i) where
+  le_sup_left _ _ _ := SemilatticeSup.le_sup_left _ _
+  le_sup_right _ _ _ := SemilatticeSup.le_sup_right _ _
+  sup_le _ _ _ ac bc i := SemilatticeSup.sup_le _ _ _ (ac i) (bc i)
+
+instance semilatticeInf [∀ i, SemilatticeInf (α' i)] : SemilatticeInf (∀ i, α' i) where
+  inf_le_left _ _ _ := SemilatticeInf.inf_le_left _ _
+  inf_le_right _ _ _ := SemilatticeInf.inf_le_right _ _
+  le_inf _ _ _ ac bc i := SemilatticeInf.le_inf _ _ _ (ac i) (bc i)
+
+instance lattice [∀ i, Lattice (α' i)] : Lattice (∀ i, α' i) :=
+{ Pi.semilatticeSup, Pi.semilatticeInf with }
+
+instance orderTop [∀ i, LE (α' i)] [∀ i, OrderTop (α' i)] : OrderTop (∀ i, α' i) :=
+  { inferInstanceAs (Top (∀ i, α' i)) with le_top := sorry }
+
+instance orderBot [∀ i, LE (α' i)] [∀ i, OrderBot (α' i)] : OrderBot (∀ i, α' i) :=
+  { inferInstanceAs (Bot (∀ i, α' i)) with bot_le := sorry }
+
+instance boundedOrder [∀ i, LE (α' i)] [∀ i, BoundedOrder (α' i)] : BoundedOrder (∀ i, α' i) :=
+{ Pi.orderTop, Pi.orderBot with }
+
+instance SupSet {α : Type _} {β : α → Type _} [∀ i, SupSet (β i)] : SupSet (∀ i, β i) :=
+  ⟨fun s i => supᵢ (λ (f : {f : ∀ i, β i // f ∈ s}) => f.1 i)⟩
+
+instance InfSet {α : Type _} {β : α → Type _} [∀ i, InfSet (β i)] : InfSet (∀ i, β i) :=
+  ⟨fun s i => infᵢ (λ (f : {f : ∀ i, β i // f ∈ s}) => f.1 i)⟩
+
+instance completeLattice {α : Type _} {β : α → Type _} [∀ i, CompleteLattice (β i)] :
+    CompleteLattice (∀ i, β i) :=
+  { Pi.boundedOrder, Pi.lattice with
+    le_supₛ := sorry
+    infₛ_le := sorry
+    supₛ_le := sorry
+    le_infₛ := sorry
+  }
+
+instance frame {ι : Type _} {π : ι → Type _} [∀ i, Frame (π i)] : Frame (∀ i, π i) :=
+  { Pi.completeLattice with }
+
+instance coframe {ι : Type _} {π : ι → Type _} [∀ i, Coframe (π i)] : Coframe (∀ i, π i) :=
+  { Pi.completeLattice with infᵢ_sup_le_sup_infₛ := sorry }
+
+end Pi
+
+-- very quick (instantaneous) in Lean 4
+instance Pi.completeDistribLattice' {ι : Type _} {π : ι → Type _}
+    [∀ i, CompleteDistribLattice (π i)] : CompleteDistribLattice (∀ i, π i) :=
+CompleteDistribLattice.mk (Pi.coframe.infᵢ_sup_le_sup_infₛ)
+
+-- takes around 2 seconds wall clock time on my PC (but very quick in Lean 3)
+set_option maxHeartbeats 400 -- make sure it stays fast
+set_option synthInstance.maxHeartbeats 400
+instance Pi.completeDistribLattice'' {ι : Type _} {π : ι → Type _}
+    [∀ i, CompleteDistribLattice (π i)] : CompleteDistribLattice (∀ i, π i) :=
+  { Pi.frame, Pi.coframe with }
+-- quick Lean 3 version:
+-- https://github.com/leanprover-community/mathlib/blob/b26e15a46f1a713ce7410e016d50575bb0bc3aa4/src/order/complete_boolean_algebra.lean#L210