From 6b2d2ffac60e71625433b98d618a8285108fa569 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Mon, 4 Jul 2022 05:49:54 -0700
Subject: [PATCH] fix: preserve `usedAssignment` flag when replacing
 `MetavarContext`

---
 src/Lean/Meta/Basic.lean     | 16 ++++------------
 src/Lean/MetavarContext.lean |  4 ++--
 2 files changed, 6 insertions(+), 14 deletions(-)

diff --git a/src/Lean/Meta/Basic.lean b/src/Lean/Meta/Basic.lean
index 2dd67ccc4c..d59f4055f9 100644
--- a/src/Lean/Meta/Basic.lean
+++ b/src/Lean/Meta/Basic.lean
@@ -307,8 +307,8 @@ protected def saveState : MetaM SavedState :=
 /-- Restore backtrackable parts of the state. -/
 def SavedState.restore (b : SavedState) : MetaM Unit := do
   Core.restore b.core
-  let usedAssignment := (← getMCtx).usedAssignment
-  modify fun s => { s with mctx := { b.meta.mctx with usedAssignment }, zetaFVarIds := b.meta.zetaFVarIds, postponed := b.meta.postponed }
+  setMCtx b.meta.mctx -- Recall that `setMCtx` propagate `usedAssignment`
+  modify fun s => { s with zetaFVarIds := b.meta.zetaFVarIds, postponed := b.meta.postponed }
 
 instance : MonadBacktrack SavedState MetaM where
   saveState      := Meta.saveState
@@ -363,9 +363,6 @@ def getLocalInstances : MetaM LocalInstances :=
 def getConfig : MetaM Config :=
   return (← read).config
 
-def setMCtx (mctx : MetavarContext) : MetaM Unit :=
-  modify fun s => { s with mctx := mctx }
-
 def resetZetaFVarIds : MetaM Unit :=
   modify fun s => { s with zetaFVarIds := {} }
 
@@ -1144,13 +1141,8 @@ private def withNewMCtxDepthImp (x : MetaM α) : MetaM α := do
   try
     x
   finally
-    -- TODO: document why we need to restore defEqCache
-    modify fun s => { s with
-      mctx               := saved.mctx
-      cache.defEqDefault := saved.cache.defEqDefault
-      cache.defEqAll     := saved.cache.defEqAll
-      postponed          := saved.postponed
-    }
+    setMCtx saved.mctx -- Recall that `setMCtx` propagate `usedAssignment`
+    modify fun s => { s with postponed := saved.postponed }
 
 /--
   Save cache and `MetavarContext`, bump the `MetavarContext` depth, execute `x`,
diff --git a/src/Lean/MetavarContext.lean b/src/Lean/MetavarContext.lean
index 29648aaceb..e84ec985c4 100644
--- a/src/Lean/MetavarContext.lean
+++ b/src/Lean/MetavarContext.lean
@@ -294,8 +294,8 @@ instance (m n) [MonadLift m n] [MonadMCtx m] : MonadMCtx n where
 def markUsedAssignment [MonadMCtx m] : m Unit :=
   modifyMCtx fun mctx => { mctx with usedAssignment := true }
 
-abbrev setMCtx [MonadMCtx m] (mctx : MetavarContext) : m Unit :=
-  modifyMCtx fun _ => mctx
+abbrev setMCtx [MonadMCtx m] (mctx' : MetavarContext) : m Unit :=
+  modifyMCtx fun mctx => { mctx' with usedAssignment := mctx'.usedAssignment || mctx.usedAssignment }
 
 abbrev getLevelMVarAssignment? [Monad m] [MonadMCtx m] (mvarId : MVarId) : m (Option Level) := do
   let result? := (← getMCtx).lAssignment.find? mvarId