cleanup the PR

Author: JovanGerb

src/Lean/Meta/Basic.lean

@@ -376,6 +376,18 @@ We should also investigate the impact on memory consumption.
 -/
 abbrev DefEqCache := PersistentHashMap DefEqCacheKey Bool
 
+/--
+A `DefEqTransCache` is a `DefEqCache` that is only valid in the original `MetavarContext`.
+It stores of the `numAssignments` from that original `MetavarContext`.
+If the `numAssignments` in the `MetavarContext` has increased, we invalidate this cache.
+And when we revert the metavariable context in `checkpointDefEq`, if the `numAssignments`
+in the original `MetavarContext` is smaller than in the cache, we revert the cache to its original.
+ -/
+structure DefEqTransCache where
+  cache : DefEqCache := {}
+  numAssignments : Nat := 0
+  deriving Inhabited
+
 /--
 Cache datastructures for type inference, type class resolution, whnf, and definitional equality.
 -/
@@ -384,7 +396,7 @@ structure Cache where
   funInfo        : FunInfoCache := {}
   synthInstance  : SynthInstanceCache := {}
   whnf           : WhnfCache := {}
-  defEqTrans     : DefEqCache × Nat := ({}, 0) -- transient cache for terms containing mvars or using nonstandard configuration options, it is valid as long as the count matches `MetavarContext.numAssignments`.
+  defEqTrans     : DefEqTransCache := {} -- transient cache for terms containing mvars or using nonstandard configuration options, it is valid as long as the count matches `MetavarContext.numAssignments`.
   defEqPerm      : DefEqCache := {} -- permanent cache for terms not containing mvars and using standard configuration options
   deriving Inhabited
 
@@ -646,9 +658,9 @@ def resetCache : MetaM Unit :=
 
 @[inline] def modifyDefEqTransientCache (numAssignments : Nat) (f : DefEqCache → DefEqCache) : MetaM Unit :=
   modifyCache fun c =>
-    let (transCache, numAssignmentsOld) := c.defEqTrans
+    let ⟨transCache, numAssignmentsOld⟩ := c.defEqTrans
     let transCache := if numAssignments == numAssignmentsOld then transCache else {}
-    { c with defEqTrans := (f transCache, numAssignments) }
+    { c with defEqTrans := ⟨f transCache, numAssignments⟩ }
 
 @[inline] def modifyDefEqPermCache (f : DefEqCache → DefEqCache) : MetaM Unit :=
   modifyCache fun ⟨c1, c2, c3, c4, c5, defeqPerm⟩ => ⟨c1, c2, c3, c4, c5, f defeqPerm⟩
@@ -667,7 +679,7 @@ def mkInfoCacheKey (expr : Expr) (nargs? : Option Nat) : MetaM InfoCacheKey :=
   return { expr, nargs?, configKey := (← read).configKey }
 
 @[inline] def resetDefEqTransientCache : MetaM Unit :=
-  modify fun s => { s with cache.defEqTrans := ({}, s.mctx.numAssignments) }
+  modify fun s => { s with cache.defEqTrans := ⟨{}, s.mctx.numAssignments⟩ }
 
 @[inline] def resetDefEqPermCaches : MetaM Unit :=
   modifyDefEqPermCache fun _ => {}
@@ -2369,13 +2381,13 @@ partial def processPostponed (mayPostpone : Bool := true) (exceptionOnFailure :=
         return true
       else
         -- The transient cache needs to be reverted if it assumes an assignments that is being reverted.
-        let invalidCache := s.meta.mctx.numAssignments < (← get).cache.defEqTrans.2
-        s.restore (transCache := invalidCache)
+        let isInvalidCache := s.meta.mctx.numAssignments != (← get).cache.defEqTrans.numAssignments
+        s.restore (transCache := isInvalidCache)
         return false
     else
       -- The transient cache needs to be reverted if it assumes an assignments that is being reverted.
-      let invalidCache := s.meta.mctx.numAssignments < (← get).cache.defEqTrans.2
-      s.restore (transCache := invalidCache)
+      let isInvalidCache := s.meta.mctx.numAssignments != (← get).cache.defEqTrans.numAssignments
+      s.restore (transCache := isInvalidCache)
       return false
   catch ex =>
     s.restore
@@ -2416,8 +2428,8 @@ def isExprDefEq (t s : Expr) : MetaM Bool :=
       We have tried in the past to track when the result was independent of the `MetavarContext` state
       but it was not effective. It is more important to cache aggressively inside of a single `isDefEq`
       call because some of the heuristics create many similar subproblems.
-      See issue #1102 for an example that triggers an exponential blowup if we don't use this more
-      aggressive form of caching.
+      See issue #1102 and `tests/lean/run/defEqTransCache.lean` for examples that trigger an exponential blowup
+      if we don't use this more aggressive form of caching.
     -/
     resetDefEqTransientCache
     checkpointDefEq (mayPostpone := true) <| Meta.isExprDefEqAux t s

src/Lean/Meta/ExprDefEq.lean

@@ -316,7 +316,7 @@ private partial def isDefEqArgs (f : Expr) (args₁ args₂ : Array Expr) : Meta
     let info := finfo.paramInfo[i]!
     if info.isInstImplicit then
       unless (← withInferTypeConfig <| Meta.isExprDefEqAux a₁ a₂) do
-        return false
+       return false
     else
       unless (← Meta.isExprDefEqAux a₁ a₂) do
         return false
@@ -2085,7 +2085,7 @@ private def mkCacheKey (t s : Expr) : MetaM DefEqCacheKeyInfo := do
 private def getCachedResult (keyInfo : DefEqCacheKeyInfo) : MetaM LBool := do
   let cache ← match keyInfo.kind with
     | .transient numAssignments =>
-      let (cache, numAssignmentsCache) := (← get).cache.defEqTrans
+      let ⟨cache, numAssignmentsCache⟩ := (← get).cache.defEqTrans
       if numAssignments == numAssignmentsCache then
         pure cache
       else
@@ -2102,14 +2102,14 @@ private def cacheResult (keyInfo : DefEqCacheKeyInfo) (result : Bool) : MetaM Un
   | .transient numAssignmentsOld =>
     /-
     If the result is `false`, we cache it at `numAssignmentsOld`.
-    If the result is `true`, we check that the number of assignments hasn't increased.
+    If the result is `true`, we only cache it if the number of assignments hasn't increase.
     -/
     if !result then
       modifyDefEqTransientCache numAssignmentsOld fun c => c.insert key result
     else
       let numAssignmentsNew := (← getMCtx).numAssignments
       if numAssignmentsOld == numAssignmentsNew then
-        modifyDefEqTransientCache numAssignmentsNew fun c => c.insert key result
+        modifyDefEqTransientCache numAssignmentsOld fun c => c.insert key result
 
 private def whnfCoreAtDefEq (e : Expr) : MetaM Expr := do
   if backward.isDefEq.lazyWhnfCore.get (← getOptions) then

tests/lean/run/defEqTransCache.lean

@@ -2,11 +2,11 @@ import Lean
 /-!
 Previously, unification wouldn't be very careful with the `isDefEq` cache for terms containing metavariables.
 - This is mostly problematic because erasing the cache leads to exponential slowdowns (`test1` & `test2`)
-- but in some cases it leads to metavariable assignments leaking into places where they shouldn't be,
-  which either causes unification to fail where it should succeed (`test3`)
-  or to succeed where it is expected to fail.
-
+- but in some cases it lead to metavariable assignments leaking into places where they shouldn't be,
+  which either caused unification to fail where it should succeed (`test3`)
+  or to succeed where it is expected to fail (which happened in one mathlib proof).
 -/
+
 set_option maxHeartbeats 1000
 
 namespace test1
@@ -16,10 +16,10 @@ class A (n : Nat) where
 instance [A n] : A (n+1) where
   x := A.x n
 
-theorem test [A 0] : A.x 100 = sorry := sorry
+theorem test [A 0] : A.x 100 = 0 := sorry
 
--- Previously, this example was exponentially slow
-example [A 1] : A.x 100 = sorry := by
+-- This rewrite should fail. Previously, it failed exponentially slowly
+example [A 1] : A.x 100 = 0 := by
   fail_if_success rw [@test]
   sorry
 end test1
@@ -66,7 +66,7 @@ elab "unfold_head" e:term : term => do
   let e ← Elab.Term.elabTerm e none
   unfoldDefinition e
 
--- we use `unfold_head` in order to get the raw kernel projection `·.1` instead of the projection funtcion `A.x`.
+-- use `unfold_head` to get the raw kernel projection `·.1` instead of the projection funtcion `A.x`
 def test {α} (i : B α) : unfold_head i.toA.x := sorry
 
 -- Previously, in this example the unification failed,