From 691e73ca3ae878de8ca94068a151eb9860bcd555 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Thu, 27 Aug 2020 14:59:54 -0700
Subject: [PATCH] fix: collect metavars occurring in delayed assignments

---
 src/Lean/Meta/Basic.lean        |  4 ++++
 src/Lean/Meta/CollectMVars.lean | 34 ++++++++++++++++++++++++++++-----
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/src/Lean/Meta/Basic.lean b/src/Lean/Meta/Basic.lean
index 2651234224..4dac8900fd 100644
--- a/src/Lean/Meta/Basic.lean
+++ b/src/Lean/Meta/Basic.lean
@@ -327,6 +327,10 @@ def isDelayedAssigned (mvarId : MVarId) : m Bool := liftMetaM do
 mctx ← getMCtx;
 pure $ mctx.isDelayedAssigned mvarId
 
+def getDelayedAssignment? (mvarId : MVarId) : m (Option DelayedMetavarAssignment) := liftMetaM do
+mctx ← getMCtx;
+pure $ mctx.getDelayedAssignment? mvarId
+
 def hasAssignableMVar (e : Expr) : m Bool := liftMetaM do
 mctx ← getMCtx;
 pure $ mctx.hasAssignableMVar e
diff --git a/src/Lean/Meta/CollectMVars.lean b/src/Lean/Meta/CollectMVars.lean
index 2dfe156405..5398b376b5 100644
--- a/src/Lean/Meta/CollectMVars.lean
+++ b/src/Lean/Meta/CollectMVars.lean
@@ -9,11 +9,32 @@ import Lean.Meta.Basic
 namespace Lean
 namespace Meta
 
-/-- Collect unassigned metavariables -/
-def collectMVars (e : Expr) : StateRefT CollectMVars.State MetaM Unit := do
-e ← instantiateMVars e;
-s ← get;
-set $ e.collectMVars s
+/--
+  Collect unassigned metavariables occuring in the given expression.
+
+  Remark: if `e` contains `?m` and there is a `t` assigned to `?m`, we
+  collect unassigned metavariables occurring in `t`.
+
+  Remark: if `e` contains `?m` and `?m` is delayed assigned to some term `t`,
+  we collect `?m` and unassigned metavariables occurring in `t`.
+  We collect `?m` because it has not been assigned yet. -/
+partial def collectMVarsAux : Expr → StateRefT CollectMVars.State MetaM Unit
+| e => do
+  e ← instantiateMVars e;
+  s ← get;
+  let resultSavedSize := s.result.size;
+  let s := e.collectMVars s;
+  set s;
+  s.result.forFromM
+    (fun mvarId => do
+      d? ← getDelayedAssignment? mvarId;
+      match d? with
+      | none   => pure ()
+      | some d => collectMVarsAux d.val)
+    resultSavedSize
+
+def collectMVars (e : Expr) : StateRefT CollectMVars.State MetaM Unit :=
+collectMVarsAux e
 
 variables {m : Type → Type} [MonadLiftT MetaM m]
 
@@ -21,6 +42,7 @@ def getMVarsImp (e : Expr) : MetaM (Array MVarId) := do
 (_, s) ← (collectMVars e).run {};
 pure s.result
 
+/-- Return metavariables in occuring the given expression. See `collectMVars` -/
 def getMVars (e : Expr) : m (Array MVarId) :=
 liftM $ getMVarsImp e
 
@@ -34,5 +56,7 @@ pure s.result
 def getMVarsAtDecl (d : Declaration) : m (Array MVarId) :=
 liftM $ getMVarsAtDeclImp d
 
+
+
 end Meta
 end Lean