From 86072873f48ed7bd17f89ea05343243d3981d75b Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Fri, 19 Mar 2021 15:14:48 +0100
Subject: [PATCH] feat: server: allow more than one ContextInfo

---
 src/Lean/Server/FileWorker.lean |  4 +--
 src/Lean/Server/InfoUtils.lean  | 53 +++++++++++++--------------------
 2 files changed, 23 insertions(+), 34 deletions(-)

diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index d35fcf0b34..5206c315fe 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -578,10 +578,10 @@ section RequestHandling
     highlightId (stx : Syntax) : ReaderT SemanticTokensContext (StateT SemanticTokensState IO) _ := do
       if let (some pos, some tailPos) := (stx.getPos?, stx.getTailPos?) then
         for t in (← read).infoState.trees do
-          for t in t.smallestNodes (fun i => match i.pos? with
+          for (ci, i) in t.deepestNodes (fun i => match i.pos? with
               | some ipos => pos <= ipos && ipos < tailPos
               | _         => false) do
-            if let Elab.InfoTree.context ci (Elab.InfoTree.node (Elab.Info.ofTermInfo ti) _) := t then
+            if let Elab.Info.ofTermInfo ti := i then
               match ti.expr with
               | Expr.fvar .. => addToken ti.stx SemanticTokenType.variable
               | _            => if ti.stx.getPos?.get! > pos then addToken ti.stx SemanticTokenType.property
diff --git a/src/Lean/Server/InfoUtils.lean b/src/Lean/Server/InfoUtils.lean
index 101a835b7d..dcc33bdf8a 100644
--- a/src/Lean/Server/InfoUtils.lean
+++ b/src/Lean/Server/InfoUtils.lean
@@ -10,29 +10,21 @@ import Lean.Util.Sorry
 
 namespace Lean.Elab
 
-/-- Find the deepest node matching `p` in the first subtree which contains a matching node.
-The result is wrapped in all outer `ContextInfo`s. -/
-partial def InfoTree.smallestNode? (p : Info → Bool) : InfoTree → Option InfoTree
-  | context i t => context i <$> t.smallestNode? p
-  | n@(node i cs) =>
-    let cs := cs.map (·.smallestNode? p)
-    let cs := cs.filter (·.isSome)
-    if !cs.isEmpty then cs.get! 0
-    else if p i then some n
-    else none
-  | _ => none
-
-/-- For every branch, find the deepest node in that branch matching `p`
-and return all of them. Each result is wrapped in all outer `ContextInfo`s. -/
-partial def InfoTree.smallestNodes (p : Info → Bool) : InfoTree → List InfoTree
-  | context i t => t.smallestNodes p |>.map (context i)
+/--
+  For every branch, find the deepest node in that branch matching `p`
+  with a surrounding context (the innermost one) and return all of them. -/
+partial def InfoTree.deepestNodes (p : Info → Bool) : InfoTree → List (ContextInfo × Info) :=
+  go none
+where go ctx?
+  | context ctx t => go ctx t
   | n@(node i cs) =>
     let cs := cs.toList
-    let ccs := cs.map (smallestNodes p)
+    let ccs := cs.map (go ctx?)
     let cs := ccs.join
     if !cs.isEmpty then cs
-    else if p i then [n]
-    else []
+    else match ctx?, p i with
+      | some ctx, true => [(ctx, i)]
+      | _,        _    => []
   | _ => []
 
 def Info.stx : Info → Syntax
@@ -48,13 +40,11 @@ def Info.tailPos? (i : Info) : Option String.Pos :=
   i.stx.getTailPos? (originalOnly := true)
 
 def InfoTree.smallestInfo? (p : Info → Bool) (t : InfoTree) : Option (ContextInfo × Info) :=
-  let ts := t.smallestNodes p
+  let ts := t.deepestNodes p
 
-  let infos : List (Nat × ContextInfo × Info) := ts.filterMap fun
-    | context ci (node i _) =>
-      let diff := i.tailPos?.get! - i.pos?.get!
-      some (diff, ci, i)
-    | _ => none
+  let infos := ts.map fun (ci, i) =>
+    let diff := i.tailPos?.get! - i.pos?.get!
+    (diff, ci, i)
 
   infos.toArray.getMax? (fun a b => a.1 > b.1) |>.map fun (_, ci, i) => (ci, i)
 
@@ -113,17 +103,16 @@ def Info.fmtHover? (ci : ContextInfo) (i : Info) : IO (Option Format) := do
     | _ => return none
 
 /-- Return a flattened list of smallest-in-span tactic info nodes, sorted by position. -/
-partial def InfoTree.smallestTacticStates (t : InfoTree) : Array (Nat × ContextInfo × TacticInfo) :=
+partial def InfoTree.smallestTacticStates (t : InfoTree) : Array (String.Pos × ContextInfo × TacticInfo) :=
   let ts := tacticLeaves t
   let ts := ts.filterMap fun
-    | context ci (node i@(Info.ofTacticInfo ti) _) => some (i.pos?.get!, ci, ti)
+    | (ci, i@(Info.ofTacticInfo ti)) => some (i.pos?.get!, ci, ti)
     | _ => none
   ts.toArray.qsort fun a b => a.1 < b.1
-
-  where tacticLeaves (t : InfoTree) : List InfoTree :=
-          t.smallestNodes fun
-            | i@(Info.ofTacticInfo _) => i.pos?.isSome ∧ i.tailPos?.isSome
-            | _ => false
+where tacticLeaves t :=
+  t.deepestNodes fun
+    | i@(Info.ofTacticInfo _) => i.pos?.isSome ∧ i.tailPos?.isSome
+    | _ => false
 
 partial def InfoTree.goalsAt? (t : InfoTree) (hoverPos : String.Pos) : Option (ContextInfo × TacticInfo) :=
   let ts := t.smallestTacticStates