From db47664d4a8586b46923eb095398bc48975e6aaa Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Thu, 7 Jul 2022 15:39:58 -0700
Subject: [PATCH] fix: discrepancy between `isDefEq` and `whnf` for
 transparency mode `instances`

---
 src/Lean/Meta/ExprDefEq.lean     | 19 +++++++++++++++++--
 src/Lean/Meta/WHNF.lean          |  2 +-
 tests/lean/run/defEqVsWhnfI.lean | 25 +++++++++++++++++++++++++
 3 files changed, 43 insertions(+), 3 deletions(-)
 create mode 100644 tests/lean/run/defEqVsWhnfI.lean

diff --git a/src/Lean/Meta/ExprDefEq.lean b/src/Lean/Meta/ExprDefEq.lean
index bdd1ea281b..dbaf8a53dc 100644
--- a/src/Lean/Meta/ExprDefEq.lean
+++ b/src/Lean/Meta/ExprDefEq.lean
@@ -1249,8 +1249,8 @@ private def unfoldNonProjFnDefEq (tInfo sInfo : ConstantInfo) (t s : Expr) : Met
   11- Otherwise, unfold `t` and `s` and continue.
   Remark: 9&10&11 are implemented by `unfoldComparingHeadsDefEq` -/
 private def isDefEqDelta (t s : Expr) : MetaM LBool := do
-  let tInfo? ← isDeltaCandidate? t.getAppFn
-  let sInfo? ← isDeltaCandidate? s.getAppFn
+  let tInfo? ← isDeltaCandidate? t
+  let sInfo? ← isDeltaCandidate? s
   match tInfo?, sInfo? with
   | none,       none       => pure LBool.undef
   | some tInfo, none       => unfold t (pure LBool.undef) fun t => isDefEqLeft tInfo.name t s
@@ -1603,6 +1603,20 @@ private def isDefEqUnitLike (t : Expr) (s : Expr) : MetaM Bool := do
     else
       return false
 
+/-
+  The `whnf` procedure has support for unfolding class projections when the
+  transparency mode is set to `.instances`. This method ensures the behavior
+  of `whnf` and `isDefEq` is consistent in this transparency mode.
+-/
+private def isDefEqProjInst (t : Expr) (s : Expr) : MetaM LBool := do
+  if (← getTransparency) != .instances then return .undef
+  let t? ← unfoldProjInstWhenIntances? t
+  let s? ← unfoldProjInstWhenIntances? s
+  if t?.isSome || s?.isSome then
+    toLBoolM <| Meta.isExprDefEqAux (t?.getD t) (s?.getD s)
+  else
+    return .undef
+
 private def isExprDefEqExpensive (t : Expr) (s : Expr) : MetaM Bool := do
   if (← (isDefEqEta t s <||> isDefEqEta s t)) then return true
   -- TODO: investigate whether this is the place for putting this check
@@ -1622,6 +1636,7 @@ private def isExprDefEqExpensive (t : Expr) (s : Expr) : MetaM Bool := do
     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
diff --git a/src/Lean/Meta/WHNF.lean b/src/Lean/Meta/WHNF.lean
index 7bc422c619..66c526c19e 100644
--- a/src/Lean/Meta/WHNF.lean
+++ b/src/Lean/Meta/WHNF.lean
@@ -604,7 +604,7 @@ mutual
     Recall that class instance projections are not marked with `[reducible]` because we want them to be
     in "reducible canonical form".
   -/
-  private partial def unfoldProjInstWhenIntances? (e : Expr) : MetaM (Option Expr) := do
+  partial def unfoldProjInstWhenIntances? (e : Expr) : MetaM (Option Expr) := do
     if (← getTransparency) != TransparencyMode.instances then
       return none
     else
diff --git a/tests/lean/run/defEqVsWhnfI.lean b/tests/lean/run/defEqVsWhnfI.lean
new file mode 100644
index 0000000000..fd1422f57f
--- /dev/null
+++ b/tests/lean/run/defEqVsWhnfI.lean
@@ -0,0 +1,25 @@
+import Lean
+
+class CountParts_ (S : Type u) where
+  μ : Type v
+  α : Type w
+  φ : S → μ → α
+
+instance : CountParts_ String where
+  μ := Char
+  α := Nat
+  φ haystack needle := haystack.foldl (fun acc x => acc + if x == needle then 1 else 0) 0
+
+class CountParts (S : Type u) (μ : Type v) (α : Type w) where
+  φ : S → μ → α
+
+open Lean Elab Meta Term in
+def test : TermElabM Unit := do
+  let e ← elabTerm (← `(CountParts_.α String)) none
+  let nat := Lean.mkConst ``Nat
+  assert! (← whnf e) == nat
+  assert! (← whnfI e) == nat
+  assert! (← whnfR e) != nat
+  assert! (← isDefEq e nat)
+  assert! (← withReducibleAndInstances <| isDefEq e nat)
+  assert! !(← withReducible <| isDefEq e nat)