From de8f8f0e28bffba8a4cfdbc5cf6f065d18ca4a9a Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Sat, 31 Oct 2020 19:13:46 -0700
Subject: [PATCH] feat: improve local context reduction approximation

---
 src/Lean/Meta/ExprDefEq.lean             | 30 +++++++++++++++++++++---
 tests/lean/run/checkAssignmentIssue.lean | 15 ++++++++++++
 2 files changed, 42 insertions(+), 3 deletions(-)
 create mode 100644 tests/lean/run/checkAssignmentIssue.lean

diff --git a/src/Lean/Meta/ExprDefEq.lean b/src/Lean/Meta/ExprDefEq.lean
index fdaaa37be1..558447f6cc 100644
--- a/src/Lean/Meta/ExprDefEq.lean
+++ b/src/Lean/Meta/ExprDefEq.lean
@@ -478,11 +478,35 @@ private def addAssignmentInfo (msg : MessageData) : CheckAssignmentM MessageData
         else
           let ctxMeta ← readThe Meta.Context
           if ctxMeta.config.ctxApprox && ctx.mvarDecl.lctx.isSubPrefixOf mvarDecl.lctx then
-            let mvarType ← check mvarDecl.type
             /- Create an auxiliary metavariable with a smaller context and "checked" type.
                Note that `mvarType` may be different from `mvarDecl.type`. Example: `mvarType` contains
-               a metavariable that we also need to reduce the context. -/
-            let newMVar ← mkAuxMVar ctx.mvarDecl.lctx ctx.mvarDecl.localInstances mvarType mvarDecl.numScopeArgs
+               a metavariable that we also need to reduce the context.
+
+               We remove from `ctx.mvarDecl.lctx` any variable that is not in `mvarDecl.lctx`
+               or in `ctx.fvars`. We don't need to remove the ones in `ctx.fvars` because
+               `elimMVarDeps` will take care of them.
+
+               First, we collect `toErase` the variables that need to be erased.
+               Notat that if a variable is `ctx.fvars`, but it depends on variable at `toErase`,
+               we must also erase it.
+            -/
+            let toErase := mvarDecl.lctx.foldl (init := #[]) fun toErase localDecl =>
+              if ctx.mvarDecl.lctx.contains localDecl.fvarId then
+                toErase
+              else if ctx.fvars.any fun fvar => fvar.fvarId! == localDecl.fvarId then
+                if mctx.findLocalDeclDependsOn localDecl fun fvarId => toErase.contains fvarId then
+                  -- localDecl depends on a variable that will be erased. So, we must add it to `toErase` too
+                  toErase.push localDecl.fvarId
+                else
+                  toErase
+              else
+                toErase.push localDecl.fvarId
+            let lctx := toErase.foldl (init := mvarDecl.lctx) fun lctx toEraseFVar =>
+              lctx.erase toEraseFVar
+            /- Compute new set of local instances. -/
+            let localInsts := mvarDecl.localInstances.filter fun localInst => toErase.contains localInst.fvar.fvarId!
+            let mvarType ← check mvarDecl.type
+            let newMVar ← mkAuxMVar lctx localInsts mvarType mvarDecl.numScopeArgs
             modifyThe Meta.State fun s => { s with mctx := s.mctx.assignExpr mvarId newMVar }
             pure newMVar
           else
diff --git a/tests/lean/run/checkAssignmentIssue.lean b/tests/lean/run/checkAssignmentIssue.lean
new file mode 100644
index 0000000000..b39f9f2e3a
--- /dev/null
+++ b/tests/lean/run/checkAssignmentIssue.lean
@@ -0,0 +1,15 @@
+instance {m} {σ} [Pure m] : Pure (StateT σ m) :=
+  let rec pure {α} (a s) := Pure.pure (a, s)
+  { pure := @pure }
+
+instance {m} {σ} [Pure m] : Pure (StateT σ m) :=
+  let rec pure {α} (a : α) (s : σ) := Pure.pure (a, s)
+  { pure := pure }
+
+instance {m} {σ} [Pure m] : Pure (StateT σ m) :=
+  let rec pure {α} (a : α) := fun s => Pure.pure (a, s)
+  { pure := pure }
+
+instance {m} {σ} [Pure m] : Pure (StateT σ m) :=
+  let rec pure {α} (a : α) : StateT σ m α := fun s => Pure.pure (a, s)
+  { pure := pure }