From 3d829c825ca3f9365310fd2368d1307077230d88 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Tue, 4 May 2021 22:07:21 -0700
Subject: [PATCH] fix: `Info` nodes for overloaded notation/declarations

See new test to understand the problem that has been fixed.
---
 src/Lean/Elab/App.lean                       |  2 +-
 src/Lean/Elab/Term.lean                      | 25 ++++++++++++--------
 tests/lean/interactive/amb.lean              | 21 ++++++++++++++++
 tests/lean/interactive/amb.lean.expected.out | 10 ++++++++
 4 files changed, 47 insertions(+), 11 deletions(-)
 create mode 100644 tests/lean/interactive/amb.lean
 create mode 100644 tests/lean/interactive/amb.lean.expected.out

diff --git a/src/Lean/Elab/App.lean b/src/Lean/Elab/App.lean
index 33d0ca0f3e..5ef6e0cfb8 100644
--- a/src/Lean/Elab/App.lean
+++ b/src/Lean/Elab/App.lean
@@ -783,9 +783,9 @@ private partial def elabAppFnId (fIdent : Syntax) (fExplicitUnivs : List Level)
   -- Set `errToSorry` to `false` if `funLVals` > 1. See comment above about the interaction between `errToSorry` and `observing`.
   withReader (fun ctx => { ctx with errToSorry := funLVals.length == 1 && ctx.errToSorry }) do
     funLVals.foldlM (init := acc) fun acc (f, fIdent, fields) => do
-      addTermInfo fIdent f
       let lvals' := toLVals fields (first := true)
       let s ← observing do
+        addTermInfo fIdent f
         let e ← elabAppLVals f (lvals' ++ lvals) namedArgs args expectedType? explicit ellipsis
         if overloaded then ensureHasType expectedType? e else pure e
       return acc.push s
diff --git a/src/Lean/Elab/Term.lean b/src/Lean/Elab/Term.lean
index f17fb995aa..1febd3053d 100644
--- a/src/Lean/Elab/Term.lean
+++ b/src/Lean/Elab/Term.lean
@@ -192,13 +192,14 @@ structure SavedState where
 protected def saveState : TermElabM SavedState := do
   pure { meta := (← Meta.saveState), «elab» := (← get) }
 
-def SavedState.restore (s : SavedState) : TermElabM Unit := do
+def SavedState.restore (s : SavedState) (restoreInfo : Bool := false) : TermElabM Unit := do
   let traceState ← getTraceState -- We never backtrack trace message
-  let infoState := (← get).infoState -- We also do not backtrack the info nodes
+  let infoState := (← get).infoState -- We also do not backtrack the info nodes when `restoreInfo == false`
   s.meta.restore
   set s.elab
   setTraceState traceState
-  modify fun s => { s with infoState := infoState }
+  unless restoreInfo do
+    modify fun s => { s with infoState := infoState }
 
 instance : MonadBacktrack SavedState TermElabM where
   saveState      := Term.saveState
@@ -220,22 +221,26 @@ def getMessageLog : TermElabM MessageLog :=
 
 /--
   Execute `x`, save resulting expression and new state.
-  If `x` fails, then it also stores exception and new state.
-  Remark: we do not capture `Exception.postpone`. -/
+  We remove any `Info` created by `x`.
+  The info nodes are committed when we execute `applyResult`.
+  We use `observing` to implement overloaded notation and decls.
+  We want to save `Info` nodes for the chosen alternative.
+-/
 @[inline] def observing (x : TermElabM α) : TermElabM (TermElabResult α) := do
   let s ← saveState
   try
     let e ← x
     let sNew ← saveState
-    s.restore
+    s.restore (restoreInfo := true)
     pure (EStateM.Result.ok e sNew)
   catch
     | ex@(Exception.error _ _) =>
       let sNew ← saveState
-      s.restore
+      s.restore (restoreInfo := true)
       pure (EStateM.Result.error ex sNew)
     | ex@(Exception.internal id _) =>
-      if id == postponeExceptionId then s.restore
+      if id == postponeExceptionId then
+        s.restore (restoreInfo := true)
       throw ex
 
 /--
@@ -243,8 +248,8 @@ def getMessageLog : TermElabM MessageLog :=
   We use this method to implement overloaded notation and symbols. -/
 @[inline] def applyResult (result : TermElabResult α) : TermElabM α :=
   match result with
-  | EStateM.Result.ok a r     => do r.restore; pure a
-  | EStateM.Result.error ex r => do r.restore; throw ex
+  | EStateM.Result.ok a r     => do r.restore (restoreInfo := true); pure a
+  | EStateM.Result.error ex r => do r.restore (restoreInfo := true); throw ex
 
 /--
   Execute `x`, but keep state modifications only if `x` did not postpone.
diff --git a/tests/lean/interactive/amb.lean b/tests/lean/interactive/amb.lean
new file mode 100644
index 0000000000..0a1bba1843
--- /dev/null
+++ b/tests/lean/interactive/amb.lean
@@ -0,0 +1,21 @@
+namespace Foo
+
+def f (n : Nat) : Bool :=
+  n == 0
+
+end Foo
+
+namespace Boo
+
+def f (n : String) : String :=
+  n ++ n
+
+end Boo
+
+open Foo
+open Boo
+
+def g := fun x => (f x : Bool)
+                 --^ textDocument/hover
+def h := fun x => (f x : String)
+                 --^ textDocument/hover
diff --git a/tests/lean/interactive/amb.lean.expected.out b/tests/lean/interactive/amb.lean.expected.out
new file mode 100644
index 0000000000..bd11b94bea
--- /dev/null
+++ b/tests/lean/interactive/amb.lean.expected.out
@@ -0,0 +1,10 @@
+{"textDocument": {"uri": "file://amb.lean"},
+ "position": {"line": 17, "character": 19}}
+{"range":
+ {"start": {"line": 17, "character": 19}, "end": {"line": 17, "character": 20}},
+ "contents": {"value": "```lean\nf : Nat → Bool\n```", "kind": "markdown"}}
+{"textDocument": {"uri": "file://amb.lean"},
+ "position": {"line": 19, "character": 19}}
+{"range":
+ {"start": {"line": 19, "character": 19}, "end": {"line": 19, "character": 20}},
+ "contents": {"value": "```lean\nf : String → String\n```", "kind": "markdown"}}