feat: structure auto-completion & partial InfoTrees

src/Init/Meta.lean

@@ -374,6 +374,9 @@ partial def structEq : Syntax → Syntax → Bool
 instance : BEq Lean.Syntax := ⟨structEq⟩
 instance : BEq (Lean.TSyntax k) := ⟨(·.raw == ·.raw)⟩
 
+/--
+Finds the first `SourceInfo` from the back of `stx` or `none` if no `SourceInfo` can be found.
+-/
 partial def getTailInfo? : Syntax → Option SourceInfo
   | atom info _   => info
   | ident info .. => info
@@ -382,14 +385,39 @@ partial def getTailInfo? : Syntax → Option SourceInfo
   | node info _ _    => info
   | _             => none
 
+/--
+Finds the first `SourceInfo` from the back of `stx` or `SourceInfo.none`
+if no `SourceInfo` can be found.
+-/
 def getTailInfo (stx : Syntax) : SourceInfo :=
   stx.getTailInfo?.getD SourceInfo.none
 
+/--
+Finds the trailing size of the first `SourceInfo` from the back of `stx`.
+If no `SourceInfo` can be found or the first `SourceInfo` from the back of `stx` contains no
+trailing whitespace, the result is `0`.
+-/
 def getTrailingSize (stx : Syntax) : Nat :=
   match stx.getTailInfo? with
   | some (SourceInfo.original (trailing := trailing) ..) => trailing.bsize
   | _ => 0
 
+/--
+Finds the trailing whitespace substring of the first `SourceInfo` from the back of `stx`.
+If no `SourceInfo` can be found or the first `SourceInfo` from the back of `stx` contains
+no trailing whitespace, the result is `none`.
+-/
+def getTrailing? (stx : Syntax) : Option Substring :=
+  stx.getTailInfo.getTrailing?
+
+/--
+Finds the tail position of the trailing whitespace of the first `SourceInfo` from the back of `stx`.
+If no `SourceInfo` can be found or the first `SourceInfo` from the back of `stx` contains
+no trailing whitespace and lacks a tail position, the result is `none`.
+-/
+def getTrailingTailPos? (stx : Syntax) (canonicalOnly := false) : Option String.Pos :=
+  stx.getTailInfo.getTrailingTailPos? canonicalOnly
+
 /--
   Return substring of original input covering `stx`.
   Result is meaningful only if all involved `SourceInfo.original`s refer to the same string (as is the case after parsing). -/

src/Init/Prelude.lean

@@ -3654,7 +3654,8 @@ namespace SourceInfo
 
 /--
 Gets the position information from a `SourceInfo`, if available.
-If `originalOnly` is true, then `.synthetic` syntax will also return `none`.
+If `canonicalOnly` is true, then `.synthetic` syntax with `canonical := false`
+will also return `none`.
 -/
 def getPos? (info : SourceInfo) (canonicalOnly := false) : Option String.Pos :=
   match info, canonicalOnly with
@@ -3665,7 +3666,8 @@ def getPos? (info : SourceInfo) (canonicalOnly := false) : Option String.Pos :=
 
 /--
 Gets the end position information from a `SourceInfo`, if available.
-If `originalOnly` is true, then `.synthetic` syntax will also return `none`.
+If `canonicalOnly` is true, then `.synthetic` syntax with `canonical := false`
+will also return `none`.
 -/
 def getTailPos? (info : SourceInfo) (canonicalOnly := false) : Option String.Pos :=
   match info, canonicalOnly with
@@ -3674,6 +3676,24 @@ def getTailPos? (info : SourceInfo) (canonicalOnly := false) : Option String.Pos
   | synthetic (endPos := endPos) .., false => some endPos
   | _,                               _     => none
 
+/--
+Gets the substring representing the trailing whitespace of a `SourceInfo`, if available.
+-/
+def getTrailing? (info : SourceInfo) : Option Substring :=
+  match info with
+  | original (trailing := trailing) .. => some trailing
+  | _                                  => none
+
+/--
+Gets the end position information of the trailing whitespace of a `SourceInfo`, if available.
+If `canonicalOnly` is true, then `.synthetic` syntax with `canonical := false`
+will also return `none`.
+-/
+def getTrailingTailPos? (info : SourceInfo) (canonicalOnly := false) : Option String.Pos :=
+  match info.getTrailing? with
+  | some trailing => some trailing.stopPos
+  | none          => info.getTailPos? canonicalOnly
+
 end SourceInfo
 
 /--
@@ -3972,7 +3992,6 @@ position information.
 def getPos? (stx : Syntax) (canonicalOnly := false) : Option String.Pos :=
   stx.getHeadInfo.getPos? canonicalOnly
 
-
 /--
 Get the ending position of the syntax, if possible.
 If `canonicalOnly` is true, non-canonical `synthetic` nodes are treated as not carrying

src/Lean/Data/Lsp/Basic.lean

@@ -365,6 +365,7 @@ structure TextDocumentRegistrationOptions where
 
 inductive MarkupKind where
   | plaintext | markdown
+  deriving DecidableEq, Hashable
 
 instance : FromJson MarkupKind := ⟨fun
   | str "plaintext" => Except.ok MarkupKind.plaintext
@@ -378,7 +379,7 @@ instance : ToJson MarkupKind := ⟨fun
 structure MarkupContent where
   kind  : MarkupKind
   value : String
-  deriving ToJson, FromJson
+  deriving ToJson, FromJson, DecidableEq, Hashable
 
 /-- Reference to the progress of some in-flight piece of work.
 

src/Lean/Data/Lsp/LanguageFeatures.lean

@@ -25,7 +25,7 @@ inductive CompletionItemKind where
   | unit | value | enum | keyword | snippet
   | color | file | reference | folder | enumMember
   | constant | struct | event | operator | typeParameter
-  deriving Inhabited, DecidableEq, Repr
+  deriving Inhabited, DecidableEq, Repr, Hashable
 
 instance : ToJson CompletionItemKind where
   toJson a := toJson (a.toCtorIdx + 1)
@@ -39,11 +39,11 @@ structure InsertReplaceEdit where
   newText : String
   insert  : Range
   replace : Range
-  deriving FromJson, ToJson
+  deriving FromJson, ToJson, BEq, Hashable
 
 inductive CompletionItemTag where
   | deprecated
-  deriving Inhabited, DecidableEq, Repr
+  deriving Inhabited, DecidableEq, Repr, Hashable
 
 instance : ToJson CompletionItemTag where
   toJson t := toJson (t.toCtorIdx + 1)
@@ -73,7 +73,7 @@ structure CompletionItem where
   commitCharacters? : string[]
   command? : Command
   -/
-  deriving FromJson, ToJson, Inhabited
+  deriving FromJson, ToJson, Inhabited, BEq, Hashable
 
 structure CompletionList where
   isIncomplete : Bool

src/Lean/Elab/App.lean

@@ -1399,8 +1399,8 @@ private def elabAppLValsAux (namedArgs : Array NamedArg) (args : Array Arg) (exp
   let rec loop : Expr → List LVal → TermElabM Expr
   | f, []          => elabAppArgs f namedArgs args expectedType? explicit ellipsis
   | f, lval::lvals => do
-    if let LVal.fieldName (fullRef := fullRef) .. := lval then
-      addDotCompletionInfo fullRef f expectedType?
+    if let LVal.fieldName (ref := ref) .. := lval then
+      addDotCompletionInfo ref f expectedType?
     let hasArgs := !namedArgs.isEmpty || !args.isEmpty
     let (f, lvalRes) ← resolveLVal f lval hasArgs
     match lvalRes with
@@ -1650,6 +1650,14 @@ private def getSuccesses (candidates : Array (TermElabResult Expr)) : TermElabM
 -/
 private def mergeFailures (failures : Array (TermElabResult Expr)) : TermElabM α := do
   let exs := failures.map fun | .error ex _ => ex | _ => unreachable!
+  let trees := failures.map (fun | .error _ s => s.meta.core.infoState.trees | _ => unreachable!)
+    |>.filterMap (·[0]?)
+  -- Retain partial `InfoTree` subtrees in an `.ofChoiceInfo` node in case of multiple failures.
+  -- This ensures that the language server still has `Info` to work with when multiple overloaded
+  -- elaborators fail.
+  withInfoContext (mkInfo := pure <| .ofChoiceInfo { elaborator := .anonymous, stx := ← getRef }) do
+    for tree in trees do
+      pushInfoTree tree
   throwErrorWithNestedErrors "overloaded" exs
 
 private def elabAppAux (f : Syntax) (namedArgs : Array NamedArg) (args : Array Arg) (ellipsis : Bool) (expectedType? : Option Expr) : TermElabM Expr := do

src/Lean/Elab/BuiltinTerm.lean

@@ -42,16 +42,15 @@ private def elabOptLevel (stx : Syntax) : TermElabM Level :=
 @[builtin_term_elab «completion»] def elabCompletion : TermElab := fun stx expectedType? => do
   /- `ident.` is ambiguous in Lean, we may try to be completing a declaration name or access a "field". -/
   if stx[0].isIdent then
-    /- If we can elaborate the identifier successfully, we assume it is a dot-completion. Otherwise, we treat it as
-       identifier completion with a dangling `.`.
-       Recall that the server falls back to identifier completion when dot-completion fails. -/
+    -- Add both an `id` and a `dot` `CompletionInfo` and have the language server figure out which
+    -- one to use.
+    addCompletionInfo <| CompletionInfo.id stx stx[0].getId (danglingDot := true) (← getLCtx) expectedType?
     let s ← saveState
     try
       let e ← elabTerm stx[0] none
       addDotCompletionInfo stx e expectedType?
     catch _ =>
       s.restore
-      addCompletionInfo <| CompletionInfo.id stx stx[0].getId (danglingDot := true) (← getLCtx) expectedType?
     throwErrorAt stx[1] "invalid field notation, identifier or numeral expected"
   else
     elabPipeCompletion stx expectedType?
@@ -328,7 +327,7 @@ private def mkSilentAnnotationIfHole (e : Expr) : TermElabM Expr := do
 @[builtin_term_elab withAnnotateTerm] def elabWithAnnotateTerm : TermElab := fun stx expectedType? => do
   match stx with
   | `(with_annotate_term $stx $e) =>
-    withInfoContext' stx (elabTerm e expectedType?) (mkTermInfo .anonymous (expectedType? := expectedType?) stx)
+    withTermInfoContext' .anonymous stx (expectedType? := expectedType?) (elabTerm e expectedType?)
   | _ => throwUnsupportedSyntax
 
 private unsafe def evalFilePathUnsafe (stx : Syntax) : TermElabM System.FilePath :=

src/Lean/Elab/Extra.lean

@@ -327,15 +327,18 @@ private def toExprCore (t : Tree) : TermElabM Expr := do
   | .term _ trees e =>
     modifyInfoState (fun s => { s with trees := s.trees ++ trees }); return e
   | .binop ref kind f lhs rhs =>
-    withRef ref <| withInfoContext' ref (mkInfo := mkTermInfo .anonymous ref) do
-      mkBinOp (kind == .lazy) f (← toExprCore lhs) (← toExprCore rhs)
+    withRef ref <|
+      withTermInfoContext' .anonymous ref do
+        mkBinOp (kind == .lazy) f (← toExprCore lhs) (← toExprCore rhs)
   | .unop ref f arg =>
-    withRef ref <| withInfoContext' ref (mkInfo := mkTermInfo .anonymous ref) do
-      mkUnOp f (← toExprCore arg)
+    withRef ref <|
+      withTermInfoContext' .anonymous ref do
+        mkUnOp f (← toExprCore arg)
   | .macroExpansion macroName stx stx' nested =>
-    withRef stx <| withInfoContext' stx (mkInfo := mkTermInfo macroName stx) do
-      withMacroExpansion stx stx' do
-        toExprCore nested
+    withRef stx <|
+      withTermInfoContext' macroName stx <|
+        withMacroExpansion stx stx' <|
+          toExprCore nested
 
 /--
   Auxiliary function to decide whether we should coerce `f`'s argument to `maxType` or not.

src/Lean/Elab/InfoTree/Main.lean

@@ -139,12 +139,16 @@ def TermInfo.runMetaM (info : TermInfo) (ctx : ContextInfo) (x : MetaM α) : IO
 
 def TermInfo.format (ctx : ContextInfo) (info : TermInfo) : IO Format := do
   info.runMetaM ctx do
-    let ty : Format ← try
-      Meta.ppExpr (← Meta.inferType info.expr)
-    catch _ =>
-      pure "<failed-to-infer-type>"
+    let ty : Format ←
+      try
+        Meta.ppExpr (← Meta.inferType info.expr)
+      catch _ =>
+        pure "<failed-to-infer-type>"
     return f!"{← Meta.ppExpr info.expr} {if info.isBinder then "(isBinder := true) " else ""}: {ty} @ {formatElabInfo ctx info.toElabInfo}"
 
+def PartialTermInfo.format (ctx : ContextInfo) (info : PartialTermInfo) : Format :=
+  f!"Partial term @ {formatElabInfo ctx info.toElabInfo}"
+
 def CompletionInfo.format (ctx : ContextInfo) (info : CompletionInfo) : IO Format :=
   match info with
   | .dot i (expectedType? := expectedType?) .. => return f!"[.] {← i.format ctx} : {expectedType?}"
@@ -191,9 +195,13 @@ def FieldRedeclInfo.format (ctx : ContextInfo) (info : FieldRedeclInfo) : Format
 def OmissionInfo.format (ctx : ContextInfo) (info : OmissionInfo) : IO Format := do
   return f!"Omission @ {← TermInfo.format ctx info.toTermInfo}\nReason: {info.reason}"
 
+def ChoiceInfo.format (ctx : ContextInfo) (info : ChoiceInfo) : Format :=
+  f!"Choice @ {formatElabInfo ctx info.toElabInfo}"
+
 def Info.format (ctx : ContextInfo) : Info → IO Format
   | ofTacticInfo i         => i.format ctx
   | ofTermInfo i           => i.format ctx
+  | ofPartialTermInfo i    => pure <| i.format ctx
   | ofCommandInfo i        => i.format ctx
   | ofMacroExpansionInfo i => i.format ctx
   | ofOptionInfo i         => i.format ctx
@@ -204,10 +212,12 @@ def Info.format (ctx : ContextInfo) : Info → IO Format
   | ofFVarAliasInfo i      => pure <| i.format
   | ofFieldRedeclInfo i    => pure <| i.format ctx
   | ofOmissionInfo i       => i.format ctx
+  | ofChoiceInfo i         => pure <| i.format ctx
 
 def Info.toElabInfo? : Info → Option ElabInfo
   | ofTacticInfo i         => some i.toElabInfo
   | ofTermInfo i           => some i.toElabInfo
+  | ofPartialTermInfo i    => some i.toElabInfo
   | ofCommandInfo i        => some i.toElabInfo
   | ofMacroExpansionInfo _ => none
   | ofOptionInfo _         => none
@@ -218,6 +228,7 @@ def Info.toElabInfo? : Info → Option ElabInfo
   | ofFVarAliasInfo _      => none
   | ofFieldRedeclInfo _    => none
   | ofOmissionInfo i       => some i.toElabInfo
+  | ofChoiceInfo i         => some i.toElabInfo
 
 /--
   Helper function for propagating the tactic metavariable context to its children nodes.
@@ -311,24 +322,36 @@ def realizeGlobalNameWithInfos (ref : Syntax) (id : Name) : CoreM (List (Name ×
       addConstInfo ref n
   return ns
 
-/-- Use this to descend a node on the infotree that is being built.
+/--
+Adds a node containing the `InfoTree`s generated by `x` to the `InfoTree`s in `m`.
+
+If `x` succeeds and `mkInfo` yields an `Info`, the `InfoTree`s of `x` become subtrees of a node
+containing the `Info` produced by `mkInfo`, which is then added to the `InfoTree`s in `m`.
+If `x` succeeds and `mkInfo` yields an `MVarId`, the `InfoTree`s of `x` are discarded and a `hole`
+node is added to the `InfoTree`s in `m`.
+If `x` fails, the `InfoTree`s of `x` become subtrees of a node containing the `Info` produced by
+`mkInfoOnError`, which is then added to the `InfoTree`s in `m`.
 
-It saves the current list of trees `t₀` and resets it and then runs `x >>= mkInfo`, producing either an `i : Info` or a hole id.
-Running `x >>= mkInfo` will modify the trees state and produce a new list of trees `t₁`.
-In the `i : Info` case, `t₁` become the children of a node `node i t₁` that is appended to `t₀`.
- -/
-def withInfoContext' [MonadFinally m] (x : m α) (mkInfo : α → m (Sum Info MVarId)) : m α := do
+The `InfoTree`s in `m` are reset before `x` is executed and restored with the addition of a new tree
+after `x` is executed.
+-/
+def withInfoContext'
+    [MonadFinally m]
+    (x : m α)
+    (mkInfo : α → m (Sum Info MVarId))
+    (mkInfoOnError : m Info) :
+    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 =>
-        let info ← mkInfo a
-        modifyInfoTrees fun trees =>
-          match info with
-          | Sum.inl info   => treesSaved.push <| InfoTree.node info trees
-          | Sum.inr mvarId => treesSaved.push <| InfoTree.hole mvarId
+      let info ← do
+        match a? with
+        | none => pure <| .inl <| ← mkInfoOnError
+        | some a => mkInfo a
+      modifyInfoTrees fun trees =>
+        match info with
+        | Sum.inl info   => treesSaved.push <| InfoTree.node info trees
+        | Sum.inr mvarId => treesSaved.push <| InfoTree.hole mvarId
   else
     x
 

src/Lean/Elab/InfoTree/Types.lean

@@ -70,6 +70,18 @@ structure TermInfo extends ElabInfo where
   isBinder : Bool := false
   deriving Inhabited
 
+/--
+Used instead of `TermInfo` when a term couldn't successfully be elaborated,
+and so there is no complete expression available.
+
+The main purpose of `PartialTermInfo` is to ensure that the sub-`InfoTree`s of a failed elaborator
+are retained so that they can still be used in the language server.
+-/
+structure PartialTermInfo extends ElabInfo where
+  lctx : LocalContext -- The local context when the term was elaborated.
+  expectedType? : Option Expr
+  deriving Inhabited
+
 structure CommandInfo extends ElabInfo where
   deriving Inhabited
 
@@ -79,7 +91,7 @@ inductive CompletionInfo where
   | dot (termInfo : TermInfo) (expectedType? : Option Expr)
   | id (stx : Syntax) (id : Name) (danglingDot : Bool) (lctx : LocalContext) (expectedType? : Option Expr)
   | dotId (stx : Syntax) (id : Name) (lctx : LocalContext) (expectedType? : Option Expr)
-  | fieldId (stx : Syntax) (id : Name) (lctx : LocalContext) (structName : Name)
+  | fieldId (stx : Syntax) (id : Option Name) (lctx : LocalContext) (structName : Name)
   | namespaceId (stx : Syntax)
   | option (stx : Syntax)
   | endSection (stx : Syntax) (scopeNames : List String)
@@ -165,10 +177,18 @@ regular delaboration settings.
 structure OmissionInfo extends TermInfo where
   reason : String
 
+/--
+Indicates that all overloaded elaborators failed. The subtrees of a `ChoiceInfo` node are the
+partial `InfoTree`s of those failed elaborators. Retaining these partial `InfoTree`s helps
+the language server provide interactivity even when all overloaded elaborators failed.
+-/
+structure ChoiceInfo extends ElabInfo where
+
 /-- Header information for a node in `InfoTree`. -/
 inductive Info where
   | ofTacticInfo (i : TacticInfo)
   | ofTermInfo (i : TermInfo)
+  | ofPartialTermInfo (i : PartialTermInfo)
   | ofCommandInfo (i : CommandInfo)
   | ofMacroExpansionInfo (i : MacroExpansionInfo)
   | ofOptionInfo (i : OptionInfo)
@@ -179,6 +199,7 @@ inductive Info where
   | ofFVarAliasInfo (i : FVarAliasInfo)
   | ofFieldRedeclInfo (i : FieldRedeclInfo)
   | ofOmissionInfo (i : OmissionInfo)
+  | ofChoiceInfo (i : ChoiceInfo)
   deriving Inhabited
 
 /-- The InfoTree is a structure that is generated during elaboration and used