From 6263141f7bb7f64c6ec0c7cd98cbd28d1bc435f5 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Wed, 17 Feb 2021 15:20:33 -0800
Subject: [PATCH] fix: preserve `InfoTree` when tactics fail

cc @Kha @Vtec234
---
 src/Lean/Elab/InfoTree.lean     | 20 ++++++++++++--------
 src/Lean/Elab/Tactic/Basic.lean |  2 ++
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/src/Lean/Elab/InfoTree.lean b/src/Lean/Elab/InfoTree.lean
index d6e0357fa7..68b2935eea 100644
--- a/src/Lean/Elab/InfoTree.lean
+++ b/src/Lean/Elab/InfoTree.lean
@@ -195,8 +195,13 @@ def mkInfoNode (info : Info) : m Unit := do
   else
     x
 
-@[inline] def withInfoContext [MonadFinally m] (x : m α) (mkInfo : m Info) : m α :=
-  withInfoContext' x fun _ => return Sum.inl (← mkInfo)
+@[inline] def withInfoContext [MonadFinally m] (x : m α) (mkInfo : m Info) : m α := do
+  if (← getInfoState).enabled then
+    let treesSaved ← getResetInfoTrees
+    Prod.fst <$> MonadFinally.tryFinally' x fun _ => do
+      modifyInfoTrees fun trees => treesSaved.push <|  InfoTree.node (← mkInfo) trees
+  else
+    x
 
 def getInfoHoleIdAssignment? (mvarId : MVarId) : m (Option InfoTree) :=
   return (← getInfoState).assignment[mvarId]
@@ -218,12 +223,11 @@ def withMacroExpansionInfo [MonadFinally m] [Monad m] [MonadInfoTree m] [MonadLC
 @[inline] def withInfoHole [MonadFinally m] [Monad m] [MonadInfoTree m] (mvarId : MVarId) (x : m α) : m α := do
   if (← getInfoState).enabled then
     let treesSaved ← getResetInfoTrees
-    Prod.fst <$> MonadFinally.tryFinally' x fun a? => do
-      match a? with
-      | none   => modifyInfoTrees fun _ => treesSaved
-      | some a => modifyInfoState fun s =>
-        assert! s.trees.size == 1 -- if size is not one, then API is being misused.
-        { s with trees := treesSaved, assignment := s.assignment.insert mvarId s.trees[0] }
+    Prod.fst <$> MonadFinally.tryFinally' x fun a? => modifyInfoState fun s =>
+      if s.trees.size > 0 then
+        { s with trees := treesSaved, assignment := s.assignment.insert mvarId s.trees[s.trees.size - 1] }
+      else
+        { s with trees := treesSaved }
   else
     x
 
diff --git a/src/Lean/Elab/Tactic/Basic.lean b/src/Lean/Elab/Tactic/Basic.lean
index f0c8cf97ef..05539355bd 100644
--- a/src/Lean/Elab/Tactic/Basic.lean
+++ b/src/Lean/Elab/Tactic/Basic.lean
@@ -45,9 +45,11 @@ def saveBacktrackableState : TacticM BacktrackableState := do
 def BacktrackableState.restore (b : BacktrackableState) : TacticM Unit := do
   setEnv b.env
   setMCtx b.mctx
+  let infoState ← getInfoState -- we do not backtrack info state
   let msgLog ← Term.getMessageLog -- we do not backtrack the message log
   set b.term
   Term.setMessageLog msgLog
+  modifyInfoState fun _ => infoState
   modify fun s => { s with goals := b.goals }
 
 @[inline] protected def tryCatch {α} (x : TacticM α) (h : Exception → TacticM α) : TacticM α := do