From 1cdad2be467af9bb5d7c0719ea7bb4416ad4eac0 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Sat, 2 Oct 2021 20:19:37 -0700
Subject: [PATCH] fix: missing info trees at `let_mvar%` elaborator

---
 src/Lean/Elab/BuiltinTerm.lean |  3 ++-
 src/Lean/Elab/Term.lean        | 35 +++++++++++++++++++++++++++++++++-
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/src/Lean/Elab/BuiltinTerm.lean b/src/Lean/Elab/BuiltinTerm.lean
index eac7f88a2a..823568acde 100644
--- a/src/Lean/Elab/BuiltinTerm.lean
+++ b/src/Lean/Elab/BuiltinTerm.lean
@@ -104,7 +104,8 @@ private def elabOptLevel (stx : Syntax) : TermElabM Level :=
        let e ← elabTerm e none
        let mvar ← mkFreshExprMVar (← inferType e) MetavarKind.syntheticOpaque n.getId
        assignExprMVar mvar.mvarId! e
-       elabTerm b expectedType?
+       -- We use `mkSaveInfoAnnotation` to make sure the info trees for `e` are saved even if `b` is a metavariable.
+       return mkSaveInfoAnnotation (← elabTerm b expectedType?)
   | _ => throwUnsupportedSyntax
 
 private def getMVarFromUserName (ident : Syntax) : MetaM Expr := do
diff --git a/src/Lean/Elab/Term.lean b/src/Lean/Elab/Term.lean
index d45a8377f9..e97ad934a0 100644
--- a/src/Lean/Elab/Term.lean
+++ b/src/Lean/Elab/Term.lean
@@ -1007,6 +1007,37 @@ private def postponeElabTerm (stx : Syntax) (expectedType? : Option Expr) : Term
 def getSyntheticMVarDecl? (mvarId : MVarId) : TermElabM (Option SyntheticMVarDecl) :=
   return (← get).syntheticMVars.find? fun d => d.mvarId == mvarId
 
+/--
+  Create an auxiliary annotation to make sure we create a `Info` even if `e` is a metavariable.
+  See `mkTermInfo`.
+
+  We use this functions because some elaboration functions elaborate subterms that may not be immediately
+  part of the resulting term. Example:
+  ```
+  let_mvar% ?m := b; wait_if_type_mvar% ?m; body
+  ```
+  If the type of `b` is not known, then `wait_if_type_mvar% ?m; body` is postponed and just return a fresh
+  metavariable `?n`. The elaborator for
+  ```
+  let_mvar% ?m := b; wait_if_type_mvar% ?m; body
+  ```
+  returns `mkSaveInfoAnnotation ?n` to make sure the info nodes created when elaborating `b` are "saved".
+  This is a bit hackish, but elaborators like `let_mvar%` are rare.
+-/
+def mkSaveInfoAnnotation (e : Expr) : Expr :=
+  if e.isMVar then
+    mkAnnotation `save_info e
+  else
+    e
+
+def isSaveInfoAnnotation? (e : Expr) : Option Expr :=
+  annotation? `save_info e
+
+partial def removeSaveInfoAnnotation (e : Expr) : Expr :=
+  match isSaveInfoAnnotation? e with
+  | some e => removeSaveInfoAnnotation e
+  | _ => e
+
 def mkTermInfo (elaborator : Name) (stx : Syntax) (e : Expr) (expectedType? : Option Expr := none) (lctx? : Option LocalContext := none) (isBinder := false) : TermElabM (Sum Info MVarId) := do
   let isHole? : TermElabM (Option MVarId) := do
     match e with
@@ -1017,8 +1048,10 @@ def mkTermInfo (elaborator : Name) (stx : Syntax) (e : Expr) (expectedType? : Op
       | _                                                   => return none
     | _ => pure none
   match (← isHole?) with
-  | none        => return Sum.inl <| Info.ofTermInfo { elaborator, lctx := lctx?.getD (← getLCtx), expr := e, stx, expectedType?, isBinder }
   | some mvarId => return Sum.inr mvarId
+  | none =>
+    let e := removeSaveInfoAnnotation e
+    return Sum.inl <| Info.ofTermInfo { elaborator, lctx := lctx?.getD (← getLCtx), expr := e, stx, expectedType?, isBinder }
 
 def addTermInfo (stx : Syntax) (e : Expr) (expectedType? : Option Expr := none) (lctx? : Option LocalContext := none) (elaborator := Name.anonymous) (isBinder := false) : TermElabM Unit := do
   withInfoContext' (pure ()) (fun _ => mkTermInfo elaborator stx e expectedType? lctx? isBinder) |> discard