From ef7cccb92fc2134f8412df348488c30686b324de Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Mon, 12 May 2025 11:48:54 +0100
Subject: [PATCH 01/24] perf: better cache sharing in `isDefEq`

---
 src/Lean/Meta/Basic.lean     | 15 +++++++++++----
 src/Lean/Meta/ExprDefEq.lean |  2 +-
 src/Lean/MetavarContext.lean | 10 ++++++----
 3 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/src/Lean/Meta/Basic.lean b/src/Lean/Meta/Basic.lean
index dd01baa1e934..4df848fb391b 100644
--- a/src/Lean/Meta/Basic.lean
+++ b/src/Lean/Meta/Basic.lean
@@ -2178,7 +2178,7 @@ partial def processPostponed (mayPostpone : Bool := true) (exceptionOnFailure :=
   We currently try to postpone universe constraints as much as possible, even when by postponing them we
   are not sure whether `x` really succeeded or not.
 -/
-@[specialize] def checkpointDefEq (x : MetaM Bool) (mayPostpone : Bool := true) : MetaM Bool := do
+@[specialize] def checkpointDefEq (x : MetaM Bool) (mayPostpone : Bool := true) (shareCache : Bool := false) : MetaM Bool := do
   let s ← saveState
   /-
     It is not safe to use the `isDefEq` cache between different `isDefEq` calls.
@@ -2189,19 +2189,26 @@ partial def processPostponed (mayPostpone : Bool := true) (exceptionOnFailure :=
     See issue #1102 for an example that triggers an exponential blowup if we don't use this more
     aggressive form of caching.
   -/
-  modifyDefEqTransientCache fun _ => {}
+  unless shareCache do
+    modifyDefEqTransientCache fun _ => {}
   let postponed ← getResetPostponed
   try
+    let restore : MetaM Unit := do
+      if shareCache && (← getMCtx).numAssignments == s.meta.mctx.numAssignments then
+        /- There are no new metavariable assigments, so the cache is still valid. -/
+        ({ s with meta.cache := (← get).cache }).restore
+      else
+        s.restore
     if (← x) then
       if (← processPostponed mayPostpone) then
         let newPostponed ← getPostponed
         setPostponed (postponed ++ newPostponed)
         return true
       else
-        s.restore
+        restore
         return false
     else
-      s.restore
+      restore
       return false
   catch ex =>
     s.restore
diff --git a/src/Lean/Meta/ExprDefEq.lean b/src/Lean/Meta/ExprDefEq.lean
index 2a705e7827b4..215503fa8349 100644
--- a/src/Lean/Meta/ExprDefEq.lean
+++ b/src/Lean/Meta/ExprDefEq.lean
@@ -1341,7 +1341,7 @@ private def tryHeuristic (t s : Expr) : MetaM Bool := do
       TODO: instead of throwing an exception as soon as we get stuck, we should just set a flag.
       Then the entry-point for `isDefEq` checks the flag before returning `true`.
     -/
-    checkpointDefEq do
+    checkpointDefEq (shareCache := true) do
       isDefEqArgs tFn t.getAppArgs s.getAppArgs <&&>
         isListLevelDefEqAux tFn.constLevels! sFn.constLevels!
 
diff --git a/src/Lean/MetavarContext.lean b/src/Lean/MetavarContext.lean
index 3ba07b91e05b..3a2ee8ed46ca 100644
--- a/src/Lean/MetavarContext.lean
+++ b/src/Lean/MetavarContext.lean
@@ -331,6 +331,8 @@ structure MetavarContext where
   lAssignment    : PersistentHashMap LMVarId Level := {}
   /-- Assignment table for expression metavariables.-/
   eAssignment    : PersistentHashMap MVarId Expr := {}
+  /-- Number of assignments in `eAssignment` and `lAssignments`. -/
+  numAssignments : Nat := 0
   /-- Assignment table for delayed abstraction metavariables.
   For more information about delayed abstraction, see the docstring for `DelayedMetavarAssignment`. -/
   dAssignment    : PersistentHashMap MVarId DelayedMetavarAssignment := {}
@@ -489,11 +491,11 @@ def hasAssignableMVar [Monad m] [MonadMCtx m] : Expr → m Bool
   This is a low-level API, and it is safer to use `isLevelDefEq (mkLevelMVar mvarId) u`.
 -/
 def assignLevelMVar [MonadMCtx m] (mvarId : LMVarId) (val : Level) : m Unit :=
-  modifyMCtx fun m => { m with lAssignment := m.lAssignment.insert mvarId val }
+  modifyMCtx fun m => { m with lAssignment := m.lAssignment.insert mvarId val, numAssignments := m.numAssignments + 1 }
 
 @[export lean_assign_lmvar]
 def assignLevelMVarExp (m : MetavarContext) (mvarId : LMVarId) (val : Level) : MetavarContext :=
-  { m with lAssignment := m.lAssignment.insert mvarId val }
+  { m with lAssignment := m.lAssignment.insert mvarId val, numAssignments := m.numAssignments + 1 }
 
 /--
 Add `mvarId := x` to the metavariable assignment.
@@ -502,11 +504,11 @@ a cycle is being introduced, or whether the expression has the right type.
 This is a low-level API, and it is safer to use `isDefEq (mkMVar mvarId) x`.
 -/
 def _root_.Lean.MVarId.assign [MonadMCtx m] (mvarId : MVarId) (val : Expr) : m Unit :=
-  modifyMCtx fun m => { m with eAssignment := m.eAssignment.insert mvarId val }
+  modifyMCtx fun m => { m with eAssignment := m.eAssignment.insert mvarId val, numAssignments := m.numAssignments + 1 }
 
 @[export lean_assign_mvar]
 def assignExp (m : MetavarContext) (mvarId : MVarId) (val : Expr) : MetavarContext :=
-  { m with eAssignment := m.eAssignment.insert mvarId val }
+  { m with eAssignment := m.eAssignment.insert mvarId val, numAssignments := m.numAssignments + 1 }
 
 /--
 Add a delayed assignment for the given metavariable. You must make sure that

From f6881cfd9fb2c24c2104a3ad7eef69f4fe4c902e Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Mon, 12 May 2025 15:18:02 +0100
Subject: [PATCH 02/24] use `(shareCache := true)` everywhere in
 `ExprDefEq.lean`

---
 src/Lean/Meta/ExprDefEq.lean | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/Lean/Meta/ExprDefEq.lean b/src/Lean/Meta/ExprDefEq.lean
index 215503fa8349..b534894b7879 100644
--- a/src/Lean/Meta/ExprDefEq.lean
+++ b/src/Lean/Meta/ExprDefEq.lean
@@ -77,7 +77,7 @@ where
         trace[Meta.isDefEq.eta.struct] "failed, type is not a structure{indentExpr b}"
         return false
       else if (← isDefEq (← inferType a) (← inferType b)) then
-        checkpointDefEq do
+        checkpointDefEq (shareCache := true) do
           let args := b.getAppArgs
           let params := args[:ctorVal.numParams].toArray
           for h : i in [ctorVal.numParams : args.size] do
@@ -1069,7 +1069,7 @@ private partial def processAssignmentFOApprox (mvar : Expr) (args : Array Expr)
     else
       trace[Meta.isDefEq.foApprox] "{mvar} {args} := {v}"
       let v := v.headBeta
-      if (← checkpointDefEq <| processAssignmentFOApproxAux mvar args v) then
+      if (← checkpointDefEq (shareCache := true) <| processAssignmentFOApproxAux mvar args v) then
         pure true
       else
         match (← unfoldDefinition? v) with
@@ -1789,12 +1789,12 @@ private partial def isDefEqQuickMVarMVar (t s : Expr) : MetaM LBool := do
   if s.isMVar && !t.isMVar then
      /- Solve `?m t =?= ?n` by trying first `?n := ?m t`.
         Reason: this assignment is precise. -/
-     if (← checkpointDefEq (processAssignment s t)) then
+     if (← checkpointDefEq (shareCache := true) (processAssignment s t)) then
        return LBool.true
      else
        toLBoolM <| processAssignment t s
   else
-     if (← checkpointDefEq (processAssignment t s)) then
+     if (← checkpointDefEq (shareCache := true) (processAssignment t s)) then
        return LBool.true
      else
        toLBoolM <| processAssignment s t
@@ -1982,11 +1982,11 @@ private def isDefEqApp (t s : Expr) : MetaM Bool := do
   let sFn := s.getAppFn
   if tFn.isConst && sFn.isConst && tFn.constName! == sFn.constName! then
     /- See comment at `tryHeuristic` explaining why we process arguments before universe levels. -/
-    if (← checkpointDefEq (isDefEqArgs tFn t.getAppArgs s.getAppArgs <&&> isListLevelDefEqAux tFn.constLevels! sFn.constLevels!)) then
+    if (← checkpointDefEq (shareCache := true) (isDefEqArgs tFn t.getAppArgs s.getAppArgs <&&> isListLevelDefEqAux tFn.constLevels! sFn.constLevels!)) then
       return true
     else
       isDefEqOnFailure t s
-  else if (← checkpointDefEq (Meta.isExprDefEqAux tFn s.getAppFn <&&> isDefEqArgs tFn t.getAppArgs s.getAppArgs)) then
+  else if (← checkpointDefEq (shareCache := true) (Meta.isExprDefEqAux tFn s.getAppFn <&&> isDefEqArgs tFn t.getAppArgs s.getAppArgs)) then
     return true
   else
     isDefEqOnFailure t s

From f43191b539ed107b355736da3f7623b3900f4117 Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Mon, 12 May 2025 19:48:35 +0100
Subject: [PATCH 03/24] better management of `Lean.Meta.Cache.defEqTrans`

---
 src/Lean/Meta/Basic.lean     | 36 +++++++++++++++++++++---------------
 src/Lean/Meta/ExprDefEq.lean | 21 +++++++++++++--------
 2 files changed, 34 insertions(+), 23 deletions(-)

diff --git a/src/Lean/Meta/Basic.lean b/src/Lean/Meta/Basic.lean
index 4df848fb391b..406d4ec81941 100644
--- a/src/Lean/Meta/Basic.lean
+++ b/src/Lean/Meta/Basic.lean
@@ -367,7 +367,7 @@ structure Cache where
   funInfo        : FunInfoCache := {}
   synthInstance  : SynthInstanceCache := {}
   whnf           : WhnfCache := {}
-  defEqTrans     : DefEqCache := {} -- transient cache for terms containing mvars or using nonstandard configuration options, it is frequently reset.
+  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`.
   defEqPerm      : DefEqCache := {} -- permanent cache for terms not containing mvars and using standard configuration options
   deriving Inhabited
 
@@ -620,7 +620,11 @@ def resetCache : MetaM Unit :=
   modifyCache fun ⟨ic, c1, c2, c3, c4, c5⟩ => ⟨f ic, c1, c2, c3, c4, c5⟩
 
 @[inline] def modifyDefEqTransientCache (f : DefEqCache → DefEqCache) : MetaM Unit :=
-  modifyCache fun ⟨c1, c2, c3, c4, defeqTrans, c5⟩ => ⟨c1, c2, c3, c4, f defeqTrans, c5⟩
+  modify fun s =>
+    if s.mctx.numAssignments == s.cache.defEqTrans.2 then
+      { s with cache.defEqTrans.1 := f s.cache.defEqTrans.1 }
+    else
+      { s with cache.defEqTrans := (f {}, s.mctx.numAssignments) }
 
 @[inline] def modifyDefEqPermCache (f : DefEqCache → DefEqCache) : MetaM Unit :=
   modifyCache fun ⟨c1, c2, c3, c4, c5, defeqPerm⟩ => ⟨c1, c2, c3, c4, c5, f defeqPerm⟩
@@ -638,6 +642,9 @@ def mkDefEqCacheKey (lhs rhs : Expr) : MetaM DefEqCacheKey := do
 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) }
+
 @[inline] def resetDefEqPermCaches : MetaM Unit :=
   modifyDefEqPermCache fun _ => {}
 
@@ -2178,23 +2185,12 @@ partial def processPostponed (mayPostpone : Bool := true) (exceptionOnFailure :=
   We currently try to postpone universe constraints as much as possible, even when by postponing them we
   are not sure whether `x` really succeeded or not.
 -/
-@[specialize] def checkpointDefEq (x : MetaM Bool) (mayPostpone : Bool := true) (shareCache : Bool := false) : MetaM Bool := do
+@[specialize] def checkpointDefEq (x : MetaM Bool) (mayPostpone : Bool := true) : MetaM Bool := do
   let s ← saveState
-  /-
-    It is not safe to use the `isDefEq` cache between different `isDefEq` calls.
-    Reason: different configuration settings, and result depends on the state of the `MetavarContext`
-    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.
-  -/
-  unless shareCache do
-    modifyDefEqTransientCache fun _ => {}
   let postponed ← getResetPostponed
   try
     let restore : MetaM Unit := do
-      if shareCache && (← getMCtx).numAssignments == s.meta.mctx.numAssignments then
+      if (← getMCtx).numAssignments == s.meta.mctx.numAssignments then
         /- There are no new metavariable assigments, so the cache is still valid. -/
         ({ s with meta.cache := (← get).cache }).restore
       else
@@ -2243,6 +2239,16 @@ def isExprDefEq (t s : Expr) : MetaM Bool :=
     Remark: the kernel does *not* update the type of variables in the local context.
     -/
     resetDefEqPermCaches
+    /-
+      It is not safe to use the transient `isDefEq` cache between different `isDefEq` calls.
+      Reason: different configuration settings, and result depends on the state of the `MetavarContext`
+      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.
+    -/
+    resetDefEqTransientCache
     checkpointDefEq (mayPostpone := true) <| Meta.isExprDefEqAux t s
 
 /--
diff --git a/src/Lean/Meta/ExprDefEq.lean b/src/Lean/Meta/ExprDefEq.lean
index b534894b7879..95f7ddb72b40 100644
--- a/src/Lean/Meta/ExprDefEq.lean
+++ b/src/Lean/Meta/ExprDefEq.lean
@@ -77,7 +77,7 @@ where
         trace[Meta.isDefEq.eta.struct] "failed, type is not a structure{indentExpr b}"
         return false
       else if (← isDefEq (← inferType a) (← inferType b)) then
-        checkpointDefEq (shareCache := true) do
+        checkpointDefEq do
           let args := b.getAppArgs
           let params := args[:ctorVal.numParams].toArray
           for h : i in [ctorVal.numParams : args.size] do
@@ -1069,7 +1069,7 @@ private partial def processAssignmentFOApprox (mvar : Expr) (args : Array Expr)
     else
       trace[Meta.isDefEq.foApprox] "{mvar} {args} := {v}"
       let v := v.headBeta
-      if (← checkpointDefEq (shareCache := true) <| processAssignmentFOApproxAux mvar args v) then
+      if (← checkpointDefEq <| processAssignmentFOApproxAux mvar args v) then
         pure true
       else
         match (← unfoldDefinition? v) with
@@ -1341,7 +1341,7 @@ private def tryHeuristic (t s : Expr) : MetaM Bool := do
       TODO: instead of throwing an exception as soon as we get stuck, we should just set a flag.
       Then the entry-point for `isDefEq` checks the flag before returning `true`.
     -/
-    checkpointDefEq (shareCache := true) do
+    checkpointDefEq do
       isDefEqArgs tFn t.getAppArgs s.getAppArgs <&&>
         isListLevelDefEqAux tFn.constLevels! sFn.constLevels!
 
@@ -1789,12 +1789,12 @@ private partial def isDefEqQuickMVarMVar (t s : Expr) : MetaM LBool := do
   if s.isMVar && !t.isMVar then
      /- Solve `?m t =?= ?n` by trying first `?n := ?m t`.
         Reason: this assignment is precise. -/
-     if (← checkpointDefEq (shareCache := true) (processAssignment s t)) then
+     if (← checkpointDefEq (processAssignment s t)) then
        return LBool.true
      else
        toLBoolM <| processAssignment t s
   else
-     if (← checkpointDefEq (shareCache := true) (processAssignment t s)) then
+     if (← checkpointDefEq (processAssignment t s)) then
        return LBool.true
      else
        toLBoolM <| processAssignment s t
@@ -1982,11 +1982,11 @@ private def isDefEqApp (t s : Expr) : MetaM Bool := do
   let sFn := s.getAppFn
   if tFn.isConst && sFn.isConst && tFn.constName! == sFn.constName! then
     /- See comment at `tryHeuristic` explaining why we process arguments before universe levels. -/
-    if (← checkpointDefEq (shareCache := true) (isDefEqArgs tFn t.getAppArgs s.getAppArgs <&&> isListLevelDefEqAux tFn.constLevels! sFn.constLevels!)) then
+    if (← checkpointDefEq (isDefEqArgs tFn t.getAppArgs s.getAppArgs <&&> isListLevelDefEqAux tFn.constLevels! sFn.constLevels!)) then
       return true
     else
       isDefEqOnFailure t s
-  else if (← checkpointDefEq (shareCache := true) (Meta.isExprDefEqAux tFn s.getAppFn <&&> isDefEqArgs tFn t.getAppArgs s.getAppArgs)) then
+  else if (← checkpointDefEq (Meta.isExprDefEqAux tFn s.getAppFn <&&> isDefEqArgs tFn t.getAppArgs s.getAppArgs)) then
     return true
   else
     isDefEqOnFailure t s
@@ -2069,7 +2069,12 @@ private def mkCacheKey (t s : Expr) : MetaM DefEqCacheKeyInfo := do
 
 private def getCachedResult (keyInfo : DefEqCacheKeyInfo) : MetaM LBool := do
   let cache ← match keyInfo.kind with
-    | .transient => pure (← get).cache.defEqTrans
+    | .transient =>
+      let (cache, num) := (← get).cache.defEqTrans
+      if (← getMCtx).numAssignments == num then
+        pure cache
+      else
+        return .undef
     | .permanent => pure (← get).cache.defEqPerm
   match cache.find? keyInfo.key with
   | some val => return val.toLBool

From 19c4a796341ef0f5beb8e60bde9543f9e9323a53 Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Mon, 12 May 2025 20:19:40 +0100
Subject: [PATCH 04/24] retry CI


From 9964565a366fe430fd54fd92f237f3901d695590 Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Mon, 12 May 2025 21:52:04 +0100
Subject: [PATCH 05/24] retry CI again


From 676b0fc72944d32df05d44c752433e2f7dba07ed Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Mon, 12 May 2025 22:38:17 +0100
Subject: [PATCH 06/24] avoid `@[extern]` to resolve segfault

---
 src/Lean/Elab/Tactic/Try.lean | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/src/Lean/Elab/Tactic/Try.lean b/src/Lean/Elab/Tactic/Try.lean
index a05166ce3ac3..61bdc4568233 100644
--- a/src/Lean/Elab/Tactic/Try.lean
+++ b/src/Lean/Elab/Tactic/Try.lean
@@ -407,11 +407,10 @@ private def evalSuggestSimpTrace : TryTactic := fun tac => do (← getMainGoal).
       return tac
   | _ => throwUnsupportedSyntax
 
-@[extern "lean_eval_suggest_tactic"] -- forward definition to avoid mutual block
-opaque evalSuggest : TryTactic
+mutual
 
 /-- `evalSuggest` for `tac1 <;> tac2` -/
-private def evalSuggestChain (tac1 tac2 : TSyntax `tactic) : TryTacticM (TSyntax `tactic) := focus do
+private partial def evalSuggestChain (tac1 tac2 : TSyntax `tactic) : TryTacticM (TSyntax `tactic) := focus do
   unless (← read).terminal do
     throwError "invalid `<;>` occurrence in non-terminal position for `try?` script{indentD (← read).root}"
   let tac1 ← withNonTerminal do evalSuggest tac1
@@ -429,7 +428,7 @@ private def evalSuggestChain (tac1 tac2 : TSyntax `tactic) : TryTacticM (TSyntax
   mkChainResult tac1 tac2s
 
 /-- `evalSuggest` for a sequence of tactics. -/
-private def evalSuggestSeq (tacs : Array (TSyntax `tactic)) : TryTacticM (TSyntax `tactic) := do
+private partial def evalSuggestSeq (tacs : Array (TSyntax `tactic)) : TryTacticM (TSyntax `tactic) := do
   if (← read).terminal then
     let mut result := #[]
     for i in [:tacs.size - 1] do
@@ -440,10 +439,10 @@ private def evalSuggestSeq (tacs : Array (TSyntax `tactic)) : TryTacticM (TSynta
   else
     mkSeq (← tacs.mapM evalSuggest) (terminal := false)
 
-private def evalSuggestSeqCore (tacs : Array Syntax) : TryTacticM (TSyntax `tactic) := do
+private partial def evalSuggestSeqCore (tacs : Array Syntax) : TryTacticM (TSyntax `tactic) := do
   evalSuggestSeq (tacs.map fun tac => ⟨tac⟩)
 
-private def evalSuggestTacticSeq (s : TSyntax ``Parser.Tactic.tacticSeq) : TryTacticM (TSyntax `tactic) := do
+private partial def evalSuggestTacticSeq (s : TSyntax ``Parser.Tactic.tacticSeq) : TryTacticM (TSyntax `tactic) := do
   let tacs ← match s with
     | `(tacticSeq| { $t;* }) => pure t.getElems
     | `(tacticSeq| $t;*) => pure t.getElems
@@ -526,9 +525,7 @@ where
         else
           throw ex
 
--- `evalSuggest` implementation
-@[export lean_eval_suggest_tactic]
-private partial def evalSuggestImpl : TryTactic := fun tac => do
+partial def evalSuggest : TryTactic := fun tac => do
   trace[try.debug] "{tac}"
   -- TODO: Implement builtin cases using `[builtin_try_tactic]` after update-stage0
   match tac with
@@ -555,6 +552,8 @@ private partial def evalSuggestImpl : TryTactic := fun tac => do
           throwError "unsolved goals"
       return r
 
+end
+
 /-! `evalAndSuggest` frontend -/
 
 private def toSuggestion (t : TSyntax `tactic) : Tactic.TryThis.Suggestion :=

From 0422dcb3eda6bcb3bc31da0da6141453ddf0749f Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Mon, 12 May 2025 23:05:41 +0100
Subject: [PATCH 07/24] Revert "avoid `@[extern]` to resolve segfault"

This reverts commit 676b0fc72944d32df05d44c752433e2f7dba07ed.
---
 src/Lean/Elab/Tactic/Try.lean | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/src/Lean/Elab/Tactic/Try.lean b/src/Lean/Elab/Tactic/Try.lean
index 61bdc4568233..a05166ce3ac3 100644
--- a/src/Lean/Elab/Tactic/Try.lean
+++ b/src/Lean/Elab/Tactic/Try.lean
@@ -407,10 +407,11 @@ private def evalSuggestSimpTrace : TryTactic := fun tac => do (← getMainGoal).
       return tac
   | _ => throwUnsupportedSyntax
 
-mutual
+@[extern "lean_eval_suggest_tactic"] -- forward definition to avoid mutual block
+opaque evalSuggest : TryTactic
 
 /-- `evalSuggest` for `tac1 <;> tac2` -/
-private partial def evalSuggestChain (tac1 tac2 : TSyntax `tactic) : TryTacticM (TSyntax `tactic) := focus do
+private def evalSuggestChain (tac1 tac2 : TSyntax `tactic) : TryTacticM (TSyntax `tactic) := focus do
   unless (← read).terminal do
     throwError "invalid `<;>` occurrence in non-terminal position for `try?` script{indentD (← read).root}"
   let tac1 ← withNonTerminal do evalSuggest tac1
@@ -428,7 +429,7 @@ private partial def evalSuggestChain (tac1 tac2 : TSyntax `tactic) : TryTacticM
   mkChainResult tac1 tac2s
 
 /-- `evalSuggest` for a sequence of tactics. -/
-private partial def evalSuggestSeq (tacs : Array (TSyntax `tactic)) : TryTacticM (TSyntax `tactic) := do
+private def evalSuggestSeq (tacs : Array (TSyntax `tactic)) : TryTacticM (TSyntax `tactic) := do
   if (← read).terminal then
     let mut result := #[]
     for i in [:tacs.size - 1] do
@@ -439,10 +440,10 @@ private partial def evalSuggestSeq (tacs : Array (TSyntax `tactic)) : TryTacticM
   else
     mkSeq (← tacs.mapM evalSuggest) (terminal := false)
 
-private partial def evalSuggestSeqCore (tacs : Array Syntax) : TryTacticM (TSyntax `tactic) := do
+private def evalSuggestSeqCore (tacs : Array Syntax) : TryTacticM (TSyntax `tactic) := do
   evalSuggestSeq (tacs.map fun tac => ⟨tac⟩)
 
-private partial def evalSuggestTacticSeq (s : TSyntax ``Parser.Tactic.tacticSeq) : TryTacticM (TSyntax `tactic) := do
+private def evalSuggestTacticSeq (s : TSyntax ``Parser.Tactic.tacticSeq) : TryTacticM (TSyntax `tactic) := do
   let tacs ← match s with
     | `(tacticSeq| { $t;* }) => pure t.getElems
     | `(tacticSeq| $t;*) => pure t.getElems
@@ -525,7 +526,9 @@ where
         else
           throw ex
 
-partial def evalSuggest : TryTactic := fun tac => do
+-- `evalSuggest` implementation
+@[export lean_eval_suggest_tactic]
+private partial def evalSuggestImpl : TryTactic := fun tac => do
   trace[try.debug] "{tac}"
   -- TODO: Implement builtin cases using `[builtin_try_tactic]` after update-stage0
   match tac with
@@ -552,8 +555,6 @@ partial def evalSuggest : TryTactic := fun tac => do
           throwError "unsolved goals"
       return r
 
-end
-
 /-! `evalAndSuggest` frontend -/
 
 private def toSuggestion (t : TSyntax `tactic) : Tactic.TryThis.Suggestion :=

From be3f48e66df33b4d6d4e43569f6d609b7f5c4021 Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Tue, 13 May 2025 04:35:02 +0100
Subject: [PATCH 08/24] fix: bug that was already present in `isDefEqProj`:
 missing `checkPointDefEq`

---
 src/Lean/Meta/ExprDefEq.lean | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/Lean/Meta/ExprDefEq.lean b/src/Lean/Meta/ExprDefEq.lean
index 95f7ddb72b40..10961bbdd048 100644
--- a/src/Lean/Meta/ExprDefEq.lean
+++ b/src/Lean/Meta/ExprDefEq.lean
@@ -1925,11 +1925,11 @@ private def isDefEqProj : Expr → Expr → MetaM Bool
   | .proj m i t, .proj n j s => do
     if (← read).inTypeClassResolution then
       -- See comment at `inTypeClassResolution`
-      pure (i == j && m == n) <&&> Meta.isExprDefEqAux t s
+      pure (i == j && m == n) <&&> checkpointDefEq (Meta.isExprDefEqAux t s)
     else if !backward.isDefEq.lazyProjDelta.get (← getOptions) then
-      pure (i == j && m == n) <&&> Meta.isExprDefEqAux t s
+      pure (i == j && m == n) <&&> checkpointDefEq (Meta.isExprDefEqAux t s)
     else if i == j && m == n then
-      isDefEqProjDelta t s i
+      checkpointDefEq (isDefEqProjDelta t s i)
     else
       return false
   | .proj structName 0 s, v  => isDefEqSingleton structName s v

From 34a76d36ccef9d63530d5498f732f14e42437826 Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Tue, 13 May 2025 05:09:10 +0100
Subject: [PATCH 09/24] retry CI


From 2fd2cde32f329f143eed236769e31daef6ca7d0d Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Tue, 13 May 2025 09:49:10 +0100
Subject: [PATCH 10/24] retry CI


From aacf58db1878e9e2fe162ca8e5b02b72457f3fc8 Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Wed, 14 May 2025 19:46:46 +0100
Subject: [PATCH 11/24] move `Term.addTermInfo id` one line up in
 `withRWRulesSeq`

---
 src/Lean/Elab/Tactic/Rewrite.lean | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/Lean/Elab/Tactic/Rewrite.lean b/src/Lean/Elab/Tactic/Rewrite.lean
index 9f52fe3546b5..a8f7b2577b0f 100644
--- a/src/Lean/Elab/Tactic/Rewrite.lean
+++ b/src/Lean/Elab/Tactic/Rewrite.lean
@@ -62,8 +62,8 @@ def withRWRulesSeq (token : Syntax) (rwRulesSeqStx : Syntax) (x : (symm : Bool)
             let rec go : List Name →  TacticM Unit
               | [] => throwError "failed to rewrite using equation theorems for '{declName}'.{hint}"
               | eqThm::eqThms => (x symm (mkCIdentFrom id eqThm)) <|> go eqThms
-            go eqThms.toList
             discard <| Term.addTermInfo id (← mkConstWithFreshMVarLevels declName) (lctx? := ← getLCtx)
+            go eqThms.toList
         match term with
         | `($id:ident)  => processId id
         | `(@$id:ident) => processId id

From 4eee7f9a3976bbe13a423ee5d71060367f526eb6 Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Wed, 14 May 2025 19:49:05 +0100
Subject: [PATCH 12/24] Revert "move `Term.addTermInfo id` one line up in
 `withRWRulesSeq`"

This reverts commit aacf58db1878e9e2fe162ca8e5b02b72457f3fc8.
---
 src/Lean/Elab/Tactic/Rewrite.lean | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/Lean/Elab/Tactic/Rewrite.lean b/src/Lean/Elab/Tactic/Rewrite.lean
index a8f7b2577b0f..9f52fe3546b5 100644
--- a/src/Lean/Elab/Tactic/Rewrite.lean
+++ b/src/Lean/Elab/Tactic/Rewrite.lean
@@ -62,8 +62,8 @@ def withRWRulesSeq (token : Syntax) (rwRulesSeqStx : Syntax) (x : (symm : Bool)
             let rec go : List Name →  TacticM Unit
               | [] => throwError "failed to rewrite using equation theorems for '{declName}'.{hint}"
               | eqThm::eqThms => (x symm (mkCIdentFrom id eqThm)) <|> go eqThms
-            discard <| Term.addTermInfo id (← mkConstWithFreshMVarLevels declName) (lctx? := ← getLCtx)
             go eqThms.toList
+            discard <| Term.addTermInfo id (← mkConstWithFreshMVarLevels declName) (lctx? := ← getLCtx)
         match term with
         | `($id:ident)  => processId id
         | `(@$id:ident) => processId id

From 983efd6febcf36cce27c956a10f87a6b9c88c30d Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Wed, 14 May 2025 22:23:15 +0100
Subject: [PATCH 13/24] fix: oh, `Meta.SavedState.restore` doesn't restore the
 cache, so I have to do it.

---
 src/Lean/Meta/Basic.lean | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/src/Lean/Meta/Basic.lean b/src/Lean/Meta/Basic.lean
index 406d4ec81941..b60634cc5734 100644
--- a/src/Lean/Meta/Basic.lean
+++ b/src/Lean/Meta/Basic.lean
@@ -2189,22 +2189,22 @@ partial def processPostponed (mayPostpone : Bool := true) (exceptionOnFailure :=
   let s ← saveState
   let postponed ← getResetPostponed
   try
-    let restore : MetaM Unit := do
-      if (← getMCtx).numAssignments == s.meta.mctx.numAssignments then
-        /- There are no new metavariable assigments, so the cache is still valid. -/
-        ({ s with meta.cache := (← get).cache }).restore
-      else
-        s.restore
     if (← x) then
       if (← processPostponed mayPostpone) then
         let newPostponed ← getPostponed
         setPostponed (postponed ++ newPostponed)
         return true
       else
-        restore
+        if (← getMCtx).numAssignments != s.meta.mctx.numAssignments then
+          -- Some metavariable assigments are being reverted, so the transient cache won't be valid anymore.
+          resetDefEqTransientCache
+        s.restore
         return false
     else
-      restore
+      if (← getMCtx).numAssignments != s.meta.mctx.numAssignments then
+        -- Some metavariable assigments are being reverted, so the transient cache won't be valid anymore.
+        resetDefEqTransientCache
+      s.restore
       return false
   catch ex =>
     s.restore

From e1d19a568a3e6f20b89cad304a7025cfcc7d212c Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Thu, 15 May 2025 02:08:14 +0100
Subject: [PATCH 14/24] change how to `modifyDefEqTransientCache`, and add
 missing `checkpointDefEq`

---
 src/Lean/Meta/Basic.lean     | 10 ++++------
 src/Lean/Meta/ExprDefEq.lean | 36 +++++++++++++++++++-----------------
 2 files changed, 23 insertions(+), 23 deletions(-)

diff --git a/src/Lean/Meta/Basic.lean b/src/Lean/Meta/Basic.lean
index b60634cc5734..93903ee00201 100644
--- a/src/Lean/Meta/Basic.lean
+++ b/src/Lean/Meta/Basic.lean
@@ -619,12 +619,10 @@ def resetCache : MetaM Unit :=
 @[inline] def modifyInferTypeCache (f : InferTypeCache → InferTypeCache) : MetaM Unit :=
   modifyCache fun ⟨ic, c1, c2, c3, c4, c5⟩ => ⟨f ic, c1, c2, c3, c4, c5⟩
 
-@[inline] def modifyDefEqTransientCache (f : DefEqCache → DefEqCache) : MetaM Unit :=
-  modify fun s =>
-    if s.mctx.numAssignments == s.cache.defEqTrans.2 then
-      { s with cache.defEqTrans.1 := f s.cache.defEqTrans.1 }
-    else
-      { s with cache.defEqTrans := (f {}, s.mctx.numAssignments) }
+@[inline] def modifyDefEqTransientCache (numAssignments : Nat) (f : DefEqCache → DefEqCache) : MetaM Unit :=
+  modifyCache fun c =>
+    let (transCache, num) := c.defEqTrans
+    { c with defEqTrans := (f (if numAssignments == num then transCache else {}), numAssignments) }
 
 @[inline] def modifyDefEqPermCache (f : DefEqCache → DefEqCache) : MetaM Unit :=
   modifyCache fun ⟨c1, c2, c3, c4, c5, defeqPerm⟩ => ⟨c1, c2, c3, c4, c5, f defeqPerm⟩
diff --git a/src/Lean/Meta/ExprDefEq.lean b/src/Lean/Meta/ExprDefEq.lean
index 10961bbdd048..53c6a97e81b5 100644
--- a/src/Lean/Meta/ExprDefEq.lean
+++ b/src/Lean/Meta/ExprDefEq.lean
@@ -312,7 +312,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
@@ -1219,7 +1219,7 @@ private partial def processAssignment (mvarApp : Expr) (v : Expr) : MetaM Bool :
   ```
 -/
 private def processAssignment' (mvarApp : Expr) (v : Expr) : MetaM Bool := do
-  if (← processAssignment mvarApp v) then
+  if (← checkpointDefEq <| processAssignment mvarApp v) then
     return true
   else
     let vNew ← whnf v
@@ -1227,7 +1227,7 @@ private def processAssignment' (mvarApp : Expr) (v : Expr) : MetaM Bool := do
       if mvarApp == vNew then
         return true
       else
-        processAssignment mvarApp vNew
+        checkpointDefEq <| processAssignment mvarApp vNew
     else
       return false
 
@@ -1608,7 +1608,7 @@ private def isDefEqProofIrrel (t s : Expr) : MetaM LBool := do
 private def isDefEqMVarSelf (mvar : Expr) (args₁ args₂ : Array Expr) : MetaM Bool := do
   if args₁.size != args₂.size then
     pure false
-  else if (← isDefEqArgs mvar args₁ args₂) then
+  else if (← checkpointDefEq <| isDefEqArgs mvar args₁ args₂) then
     pure true
   else if !(← isAssignable mvar) then
     pure false
@@ -1846,8 +1846,8 @@ end
 
 private def isDefEqOnFailure (t s : Expr) : MetaM Bool := do
   withTraceNodeBefore `Meta.isDefEq.onFailure (return m!"{t} =?= {s}") do
-    unstuckMVar t (fun t => Meta.isExprDefEqAux t s) <|
-    unstuckMVar s (fun s => Meta.isExprDefEqAux t s) <|
+    unstuckMVar t (fun t => checkpointDefEq <| Meta.isExprDefEqAux t s) <|
+    unstuckMVar s (fun s => checkpointDefEq <| Meta.isExprDefEqAux t s) <|
     tryUnificationHints t s <||> tryUnificationHints s t
 
 /--
@@ -2046,12 +2046,12 @@ private def isExprDefEqExpensive (t : Expr) (s : Expr) : MetaM Bool := do
       isDefEqOnFailure t s
 
 inductive DefEqCacheKind where
-  | transient -- problem has mvars or is using nonstandard configuration, we should use transient cache
+  | transient (numAssignments : Nat) -- problem has mvars or is using nonstandard configuration, we should use transient cache
   | permanent -- problem does not have mvars and we are using standard config, we can use one persistent cache.
 
 private def getDefEqCacheKind (t s : Expr) : MetaM DefEqCacheKind := do
   if t.hasMVar || s.hasMVar || (← read).canUnfold?.isSome then
-    return .transient
+    return .transient (← getMCtx).numAssignments
   else
     return .permanent
 
@@ -2069,9 +2069,9 @@ private def mkCacheKey (t s : Expr) : MetaM DefEqCacheKeyInfo := do
 
 private def getCachedResult (keyInfo : DefEqCacheKeyInfo) : MetaM LBool := do
   let cache ← match keyInfo.kind with
-    | .transient =>
-      let (cache, num) := (← get).cache.defEqTrans
-      if (← getMCtx).numAssignments == num then
+    | .transient numAssignments =>
+      let (cache, numAssignments') := (← get).cache.defEqTrans
+      if numAssignments == numAssignments' then
         pure cache
       else
         return .undef
@@ -2084,14 +2084,16 @@ private def cacheResult (keyInfo : DefEqCacheKeyInfo) (result : Bool) : MetaM Un
   let key := keyInfo.key
   match keyInfo.kind with
   | .permanent => modifyDefEqPermCache fun c => c.insert key result
-  | .transient =>
+  | .transient numAssignments =>
     /-
-    We must ensure that all assigned metavariables in the key are replaced by their current assignments.
-    Otherwise, the key is invalid after the assignment is "backtracked".
-    See issue #1870 for an example.
+    we cache results with metavariables if the same result may show up again.
+    That is, if the unification attempt doesn't instantiate any metavariables,
+    either because it returns `false`, or because nothing was assigned.
     -/
-    let key ← mkDefEqCacheKey (← instantiateMVars key.lhs) (← instantiateMVars key.rhs)
-    modifyDefEqTransientCache fun c => c.insert key result
+    if !result then
+      modifyDefEqTransientCache numAssignments fun c => c.insert key result
+    else if numAssignments == (← getMCtx).numAssignments then
+      modifyDefEqTransientCache numAssignments fun c => c.insert key result
 
 private def whnfCoreAtDefEq (e : Expr) : MetaM Expr := do
   if backward.isDefEq.lazyWhnfCore.get (← getOptions) then

From 29307372978473081c0f38e5158ecbdbe4fdf4d8 Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Mon, 2 Jun 2025 18:06:37 +0100
Subject: [PATCH 15/24] perf: instead of resetting the transient cache to be
 empty, reset it to the previous value

---
 src/Lean/Meta/Basic.lean | 22 ++++++++++++----------
 1 file changed, 12 insertions(+), 10 deletions(-)

diff --git a/src/Lean/Meta/Basic.lean b/src/Lean/Meta/Basic.lean
index 93903ee00201..98c9e1c1388c 100644
--- a/src/Lean/Meta/Basic.lean
+++ b/src/Lean/Meta/Basic.lean
@@ -555,9 +555,13 @@ protected def saveState : MetaM SavedState :=
   return { core := (← Core.saveState), meta := (← get) }
 
 /-- Restore backtrackable parts of the state. -/
-def SavedState.restore (b : SavedState) : MetaM Unit := do
+def SavedState.restore (b : SavedState) (transCache : Bool := false) : MetaM Unit := do
   b.core.restore
-  modify fun s => { s with mctx := b.meta.mctx, zetaDeltaFVarIds := b.meta.zetaDeltaFVarIds, postponed := b.meta.postponed }
+  modify fun s => { s with
+    mctx := b.meta.mctx
+    zetaDeltaFVarIds := b.meta.zetaDeltaFVarIds
+    postponed := b.meta.postponed
+    cache.defEqTrans := if transCache then b.meta.cache.defEqTrans else s.cache.defEqTrans }
 
 @[specialize, inherit_doc Core.withRestoreOrSaveFull]
 def withRestoreOrSaveFull (reusableResult? : Option (α × SavedState)) (act : MetaM α) :
@@ -2193,16 +2197,14 @@ partial def processPostponed (mayPostpone : Bool := true) (exceptionOnFailure :=
         setPostponed (postponed ++ newPostponed)
         return true
       else
-        if (← getMCtx).numAssignments != s.meta.mctx.numAssignments then
-          -- Some metavariable assigments are being reverted, so the transient cache won't be valid anymore.
-          resetDefEqTransientCache
-        s.restore
+        -- If restoring the state causes a metavariable assignment to be reverted, the transient cache also needs to be reverted
+        let revertedMCtx := (← getMCtx).numAssignments != s.meta.mctx.numAssignments
+        s.restore (transCache := revertedMCtx)
         return false
     else
-      if (← getMCtx).numAssignments != s.meta.mctx.numAssignments then
-        -- Some metavariable assigments are being reverted, so the transient cache won't be valid anymore.
-        resetDefEqTransientCache
-      s.restore
+      -- If restoring the state causes a metavariable assignment to be reverted, the transient cache also needs to be reverted
+      let revertedMCtx := (← getMCtx).numAssignments != s.meta.mctx.numAssignments
+      s.restore (transCache := revertedMCtx)
       return false
   catch ex =>
     s.restore

From e05d9f94c363ebed2fc758e5bc3e653acfb7e51f Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Tue, 3 Jun 2025 04:01:10 +0100
Subject: [PATCH 16/24] perf: don't increment `numAssignments` in
 `instantiateMVars`

---
 src/Lean/MetavarContext.lean | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/Lean/MetavarContext.lean b/src/Lean/MetavarContext.lean
index 3a2ee8ed46ca..426566cb537e 100644
--- a/src/Lean/MetavarContext.lean
+++ b/src/Lean/MetavarContext.lean
@@ -493,9 +493,10 @@ def hasAssignableMVar [Monad m] [MonadMCtx m] : Expr → m Bool
 def assignLevelMVar [MonadMCtx m] (mvarId : LMVarId) (val : Level) : m Unit :=
   modifyMCtx fun m => { m with lAssignment := m.lAssignment.insert mvarId val, numAssignments := m.numAssignments + 1 }
 
+-- `assignLevelMVarExp` is only used in `instantiateMVars`, so it doesn't need to increment `numAssignments`
 @[export lean_assign_lmvar]
 def assignLevelMVarExp (m : MetavarContext) (mvarId : LMVarId) (val : Level) : MetavarContext :=
-  { m with lAssignment := m.lAssignment.insert mvarId val, numAssignments := m.numAssignments + 1 }
+  { m with lAssignment := m.lAssignment.insert mvarId val }
 
 /--
 Add `mvarId := x` to the metavariable assignment.
@@ -506,9 +507,10 @@ This is a low-level API, and it is safer to use `isDefEq (mkMVar mvarId) x`.
 def _root_.Lean.MVarId.assign [MonadMCtx m] (mvarId : MVarId) (val : Expr) : m Unit :=
   modifyMCtx fun m => { m with eAssignment := m.eAssignment.insert mvarId val, numAssignments := m.numAssignments + 1 }
 
+-- `assignExp` is only used in `instantiateMVars`, so it doesn't need to increment `numAssignments`
 @[export lean_assign_mvar]
 def assignExp (m : MetavarContext) (mvarId : MVarId) (val : Expr) : MetavarContext :=
-  { m with eAssignment := m.eAssignment.insert mvarId val, numAssignments := m.numAssignments + 1 }
+  { m with eAssignment := m.eAssignment.insert mvarId val }
 
 /--
 Add a delayed assignment for the given metavariable. You must make sure that

From 85f2f782e4e1ac54ed4f4389e514339c95232258 Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Tue, 3 Jun 2025 11:29:53 +0100
Subject: [PATCH 17/24] empty commit


From 244866dbd7647a16fcfedabeb10aa27da12850e8 Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Tue, 3 Jun 2025 17:15:44 +0100
Subject: [PATCH 18/24] add a test file

---
 tests/lean/run/defEqTransCache.lean | 75 +++++++++++++++++++++++++++++
 1 file changed, 75 insertions(+)
 create mode 100644 tests/lean/run/defEqTransCache.lean

diff --git a/tests/lean/run/defEqTransCache.lean b/tests/lean/run/defEqTransCache.lean
new file mode 100644
index 000000000000..a451d2aad35f
--- /dev/null
+++ b/tests/lean/run/defEqTransCache.lean
@@ -0,0 +1,75 @@
+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.
+
+-/
+namespace test1
+class A (n : Nat) where
+  x : Nat
+
+instance [A n] : A (n+1) where
+  x := A.x n
+
+theorem test [A 0] : A.x 100 = sorry := sorry
+
+-- Previously, this example was exponentially slow
+example [A 1] : A.x 100 = sorry := by
+  fail_if_success rw [@test]
+  sorry
+end test1
+
+
+namespace test2
+@[irreducible] def A : Type := Unit
+
+@[irreducible] def B : Type := Unit
+
+unseal B in
+@[coe] def AtoB (_a : A) : B := ()
+
+instance : Coe A B where coe := AtoB
+
+def h {α : Type} (a b : α) : Nat → α
+| 0 => a
+| n + 1 => h b a n
+
+def f {α : Type} (a b : α) : Nat → Prop
+| 0 => a = b
+| n + 1 => f (h a b n) (h b a n) n ∧ f (h a a n) (h b b n) n
+
+axiom foo {p} {α : Type} (a b : α) : f a b p
+
+variable (x : A) (y : B)
+-- Previously, this check was exponentially slow; now it is quadratically slow
+#check (foo (↑x) y : f (AtoB x) y 30)
+end test2
+
+
+namespace test3
+structure A (α : Type) where
+  x : Type
+  y : α
+
+structure B (α : Type) extends A α where
+  z : Nat
+
+def A.map {α β} (f : α → β) (a : A α) : A β := ⟨a.x, f a.y⟩
+
+open Lean Meta in
+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`.
+def test {α} (i : B α) : unfold_head i.toA.x := sorry
+
+-- Previously, in this example the unification failed,
+-- because some metavariable assignment wasn't reverted properly
+-- However, it is clearly the case that `(i.toA.map f).x` is the same as `i.toA.x`
+example (i : B α) (f : α → β) : (i.toA.map f).x := by
+  apply @test
+end test3

From 53a4d87a8c3b9eeda731f0b5a8d02b6cf0524468 Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Tue, 3 Jun 2025 20:53:40 +0100
Subject: [PATCH 19/24] `set_option maxHeartbeats 1000` in test

---
 tests/lean/run/defEqTransCache.lean | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/lean/run/defEqTransCache.lean b/tests/lean/run/defEqTransCache.lean
index a451d2aad35f..63eb49d18438 100644
--- a/tests/lean/run/defEqTransCache.lean
+++ b/tests/lean/run/defEqTransCache.lean
@@ -7,6 +7,8 @@ Previously, unification wouldn't be very careful with the `isDefEq` cache for te
   or to succeed where it is expected to fail.
 
 -/
+set_option maxHeartbeats 1000
+
 namespace test1
 class A (n : Nat) where
   x : Nat

From cc74f96f243177e469ef73afc3245fe1e854551e Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Tue, 3 Jun 2025 20:55:22 +0100
Subject: [PATCH 20/24] slightly improve caching

---
 src/Lean/Meta/Basic.lean     | 17 +++++++++--------
 src/Lean/Meta/ExprDefEq.lean | 17 ++++++++++-------
 2 files changed, 19 insertions(+), 15 deletions(-)

diff --git a/src/Lean/Meta/Basic.lean b/src/Lean/Meta/Basic.lean
index 98c9e1c1388c..b477def263f0 100644
--- a/src/Lean/Meta/Basic.lean
+++ b/src/Lean/Meta/Basic.lean
@@ -625,8 +625,9 @@ def resetCache : MetaM Unit :=
 
 @[inline] def modifyDefEqTransientCache (numAssignments : Nat) (f : DefEqCache → DefEqCache) : MetaM Unit :=
   modifyCache fun c =>
-    let (transCache, num) := c.defEqTrans
-    { c with defEqTrans := (f (if numAssignments == num then transCache else {}), numAssignments) }
+    let (transCache, numAssignmentsOld) := c.defEqTrans
+    let transCache := if numAssignments == numAssignmentsOld then transCache else {}
+    { 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⟩
@@ -2197,14 +2198,14 @@ partial def processPostponed (mayPostpone : Bool := true) (exceptionOnFailure :=
         setPostponed (postponed ++ newPostponed)
         return true
       else
-        -- If restoring the state causes a metavariable assignment to be reverted, the transient cache also needs to be reverted
-        let revertedMCtx := (← getMCtx).numAssignments != s.meta.mctx.numAssignments
-        s.restore (transCache := revertedMCtx)
+        -- The transient cache needs to be reverted if it assumes some assignments that are being reverted as well.
+        let invalidCache := s.meta.mctx.numAssignments < (← get).cache.defEqTrans.2
+        s.restore (transCache := invalidCache)
         return false
     else
-      -- If restoring the state causes a metavariable assignment to be reverted, the transient cache also needs to be reverted
-      let revertedMCtx := (← getMCtx).numAssignments != s.meta.mctx.numAssignments
-      s.restore (transCache := revertedMCtx)
+      -- The transient cache needs to be reverted if it assumes some assignments that are being reverted as well.
+      let invalidCache := s.meta.mctx.numAssignments < (← get).cache.defEqTrans.2
+      s.restore (transCache := invalidCache)
       return false
   catch ex =>
     s.restore
diff --git a/src/Lean/Meta/ExprDefEq.lean b/src/Lean/Meta/ExprDefEq.lean
index 53c6a97e81b5..c11d67db6170 100644
--- a/src/Lean/Meta/ExprDefEq.lean
+++ b/src/Lean/Meta/ExprDefEq.lean
@@ -2084,16 +2084,19 @@ private def cacheResult (keyInfo : DefEqCacheKeyInfo) (result : Bool) : MetaM Un
   let key := keyInfo.key
   match keyInfo.kind with
   | .permanent => modifyDefEqPermCache fun c => c.insert key result
-  | .transient numAssignments =>
+  | .transient numAssignmentsOld =>
     /-
-    we cache results with metavariables if the same result may show up again.
-    That is, if the unification attempt doesn't instantiate any metavariables,
-    either because it returns `false`, or because nothing was assigned.
+    If the result is `false`, we cache it at `numAssignmentsOld`.
+    If the result is `true`, and the number of assignments has increased,
+    then we need to instantiate metavariables and cache it at `numAssignmentsNew`
     -/
     if !result then
-      modifyDefEqTransientCache numAssignments fun c => c.insert key result
-    else if numAssignments == (← getMCtx).numAssignments then
-      modifyDefEqTransientCache numAssignments fun c => c.insert key result
+      modifyDefEqTransientCache numAssignmentsOld fun c => c.insert key result
+    else
+      let numAssignmentsNew := (← getMCtx).numAssignments
+      let key ← if numAssignmentsOld == numAssignmentsNew then pure key else
+        mkDefEqCacheKey (← instantiateMVars key.lhs) (← instantiateMVars key.rhs)
+      modifyDefEqTransientCache numAssignmentsNew fun c => c.insert key result
 
 private def whnfCoreAtDefEq (e : Expr) : MetaM Expr := do
   if backward.isDefEq.lazyWhnfCore.get (← getOptions) then

From 0e96d2c96f63fcc1016f6fb901faefa8f809fbca Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Wed, 4 Jun 2025 11:07:53 +0100
Subject: [PATCH 21/24] undo caching unifications that instantiate

---
 src/Lean/Meta/ExprDefEq.lean | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/src/Lean/Meta/ExprDefEq.lean b/src/Lean/Meta/ExprDefEq.lean
index c11d67db6170..ea64e1468282 100644
--- a/src/Lean/Meta/ExprDefEq.lean
+++ b/src/Lean/Meta/ExprDefEq.lean
@@ -2087,16 +2087,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`, and the number of assignments has increased,
-    then we need to instantiate metavariables and cache it at `numAssignmentsNew`
+    If the result is `true`, we check that the number of assignments hasn't increased.
     -/
     if !result then
       modifyDefEqTransientCache numAssignmentsOld fun c => c.insert key result
     else
       let numAssignmentsNew := (← getMCtx).numAssignments
-      let key ← if numAssignmentsOld == numAssignmentsNew then pure key else
-        mkDefEqCacheKey (← instantiateMVars key.lhs) (← instantiateMVars key.rhs)
-      modifyDefEqTransientCache numAssignmentsNew fun c => c.insert key result
+      if numAssignmentsOld == numAssignmentsNew then
+        modifyDefEqTransientCache numAssignmentsNew fun c => c.insert key result
 
 private def whnfCoreAtDefEq (e : Expr) : MetaM Expr := do
   if backward.isDefEq.lazyWhnfCore.get (← getOptions) then

From 8c514b453ec01eb6a1463ad6ec7a508bddf5275a Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Wed, 4 Jun 2025 18:03:02 +0100
Subject: [PATCH 22/24] avoid `isDefEq` in `ExprDefEq.lean`

---
 src/Lean/Meta/ExprDefEq.lean | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/Lean/Meta/ExprDefEq.lean b/src/Lean/Meta/ExprDefEq.lean
index ea64e1468282..24bc554c806a 100644
--- a/src/Lean/Meta/ExprDefEq.lean
+++ b/src/Lean/Meta/ExprDefEq.lean
@@ -76,7 +76,7 @@ where
       if !isStructureLike (← getEnv) ctorVal.induct then
         trace[Meta.isDefEq.eta.struct] "failed, type is not a structure{indentExpr b}"
         return false
-      else if (← isDefEq (← inferType a) (← inferType b)) then
+      else if (← checkpointDefEq <| Meta.isExprDefEqAux (← inferType a) (← inferType b)) then
         checkpointDefEq do
           let args := b.getAppArgs
           let params := args[:ctorVal.numParams].toArray
@@ -91,7 +91,7 @@ where
                 -- See comment at `isAbstractedUnassignedMVar`.
                 continue
             trace[Meta.isDefEq.eta.struct] "{a} =?= {b} @ [{j}], {proj} =?= {args[i]}"
-            unless (← isDefEq proj args[i]) do
+            unless (← Meta.isExprDefEqAux proj args[i]) do
               trace[Meta.isDefEq.eta.struct] "failed, unexpected arg #{i}, projection{indentExpr proj}\nis not defeq to{indentExpr args[i]}"
               return false
           return true

From 5f05587e05530650a279b4ba590cef61cdecb183 Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Wed, 4 Jun 2025 18:20:56 +0100
Subject: [PATCH 23/24] empty commit


From 9c0429d88fa6de5b2a99713acbf56d98533ee4ce Mon Sep 17 00:00:00 2001
From: Jovan Gerbscheid <jovan.gerbscheid@gmail.com>
Date: Sat, 7 Jun 2025 19:09:36 +0100
Subject: [PATCH 24/24] empty commit