fix: unknown identifier ranges (#8362)

This PR fixes a bug where the unknown identifier code actions wouldn't
work correctly for some unknown identifier error spans and adjusts
several unknown identifier spans to actually end on the identifier in
question.

The following additional adjustments are made:
- The fallback mechanism of the unknown identifier code actions is
removed, since it could produce severely incorrect suggestions for
unknown identifier errors on fields.
- A performance bug when using the code action to import all unknown
identifiers is fixed.
- A bug that occurs when the elaborator produces multiple overlapping
completion infos is fixed.
- A bug in the snapshot selection that could cause it to wait for
snapshots in snapshots with non-canonical syntax is fixed.
- Some invariants of the snapshot tree are documented.
- The snapshot tree formatting is adjusted to display the final info
tree again.

src/Lean/Elab/App.lean

@@ -1146,8 +1146,11 @@ inductive LValResolution where
   The `fullName` is the name of the recursive function, and `baseName` is the base name of the type to search for in the parameter list. -/
   | localRec (baseName : Name) (fullName : Name) (fvar : Expr)
 
-private def throwLValError (e : Expr) (eType : Expr) (msg : MessageData) : TermElabM α :=
-  throwError "{msg}{indentExpr e}\nhas type{indentExpr eType}"
+private def throwLValErrorAt (ref : Syntax) (e : Expr) (eType : Expr) (msg : MessageData) : TermElabM α :=
+  throwErrorAt ref "{msg}{indentExpr e}\nhas type{indentExpr eType}"
+
+private def throwLValError (e : Expr) (eType : Expr) (msg : MessageData) : TermElabM α := do
+  throwLValErrorAt (← getRef) e eType msg
 
 /--
 `findMethod? S fName` tries the for each namespace `S'` in the resolution order for `S` to resolve the name `S'.fname`.
@@ -1206,7 +1209,7 @@ private def resolveLValAux (e : Expr) (eType : Expr) (lval : LVal) : TermElabM L
           return LValResolution.projIdx structName (idx - 1)
       else
         throwLValError e eType m!"invalid projection, structure has only {numFields} field(s)"
-  | some structName, LVal.fieldName _ fieldName _ _ =>
+  | some structName, LVal.fieldName _ fieldName _ fullRef =>
     let env ← getEnv
     if isStructure env structName then
       if let some baseStructName := findField? env structName (Name.mkSimple fieldName) then
@@ -1223,10 +1226,10 @@ private def resolveLValAux (e : Expr) (eType : Expr) (lval : LVal) : TermElabM L
     if let some (baseStructName, fullName) ← findMethod? structName (.mkSimple fieldName) then
       return LValResolution.const baseStructName structName fullName
     let msg := mkUnknownIdentifierMessage m!"invalid field '{fieldName}', the environment does not contain '{Name.mkStr structName fieldName}'"
-    throwLValError e eType msg
-  | none, LVal.fieldName _ _ (some suffix) _ =>
+    throwLValErrorAt fullRef e eType msg
+  | none, LVal.fieldName _ _ (some suffix) fullRef =>
     if e.isConst then
-      throwUnknownConstant (e.constName! ++ suffix)
+      throwUnknownConstantAt fullRef (e.constName! ++ suffix)
     else
       throwInvalidFieldNotation e eType
   | _, _ => throwInvalidFieldNotation e eType
@@ -1511,7 +1514,7 @@ where
       else if let some (fvar, []) ← resolveLocalName idNew then
         return fvar
       else
-        throwUnknownIdentifier m!"invalid dotted identifier notation, unknown identifier `{idNew}` from expected type{indentExpr expectedType}"
+        throwUnknownIdentifierAt id m!"invalid dotted identifier notation, unknown identifier `{idNew}` from expected type{indentExpr expectedType}"
     catch
       | ex@(.error ..) =>
         match (← unfoldDefinition? resultType) with

src/Lean/Elab/Declaration.lean

@@ -311,7 +311,7 @@ def elabMutual : CommandElab := fun stx => do
         if (← Simp.isBuiltinSimproc name) then
           pure [name]
         else
-          throwUnknownConstant name
+          throwUnknownConstantAt ident name
     let declName ← ensureNonAmbiguous ident declNames
     Term.applyAttributes declName attrs
     for attrName in toErase do

src/Lean/Elab/Extra.lean

@@ -209,14 +209,14 @@ where
         | none => processLeaf s
 
   processBinOp (ref : Syntax) (kind : BinOpKind) (f lhs rhs : Syntax) := do
-    let some f ← resolveId? f | throwUnknownConstant f.getId
+    let some f ← resolveId? f | throwUnknownConstantAt f f.getId
     -- treat corresponding argument as leaf for `leftact/rightact`
     let lhs ← if kind == .leftact then processLeaf lhs else go lhs
     let rhs ← if kind == .rightact then processLeaf rhs else go rhs
     return .binop ref kind f lhs rhs
 
   processUnOp (ref : Syntax) (f arg : Syntax) := do
-    let some f ← resolveId? f | throwUnknownConstant f.getId
+    let some f ← resolveId? f | throwUnknownConstantAt f f.getId
     return .unop ref f (← go arg)
 
   processLeaf (s : Syntax) := do
@@ -547,7 +547,7 @@ def elabBinRelCore (noProp : Bool) (stx : Syntax) (expectedType? : Option Expr)
       let result ← toExprCore (← applyCoe tree maxType (isPred := true))
       trace[Elab.binrel] "result: {result}"
       return result
-  | none   => throwUnknownConstant stx[1].getId
+  | none   => throwUnknownConstantAt stx[1] stx[1].getId
 where
   /-- If `noProp == true` and `e` has type `Prop`, then coerce it to `Bool`. -/
   toBoolIfNecessary (e : Expr) : TermElabM Expr := do

src/Lean/Elab/Tactic/Simp.lean

@@ -202,7 +202,7 @@ def elabSimpArgs (stx : Syntax) (ctx : Simp.Context) (simprocs : Simp.SimprocsAr
                   if (← Simp.isBuiltinSimproc name) then
                     simprocs := simprocs.erase name
                   else
-                    withRef id <| throwUnknownConstant name
+                    throwUnknownConstantAt id name
             else if arg.getKind == ``Lean.Parser.Tactic.simpLemma then
               let post :=
                 if arg[0].isNone then

src/Lean/Elab/Term.lean

@@ -1979,7 +1979,7 @@ where
            isValidAutoBoundImplicitName n (relaxedAutoImplicit.get (← getOptions)) then
         throwAutoBoundImplicitLocal n
       else
-        throwUnknownIdentifier m!"unknown identifier '{Lean.mkConst n}'"
+        throwUnknownIdentifierAt stx m!"unknown identifier '{Lean.mkConst n}'"
     mkConsts candidates explicitLevels
 
 /--

src/Lean/Exception.lean

@@ -76,28 +76,43 @@ prompt the code action.
 -/
 def unknownIdentifierMessageTag : Name := `unknownIdentifier
 
+/-- Throw an error exception using the given message data and reference syntax. -/
+protected def throwErrorAt [Monad m] [MonadError m] (ref : Syntax) (msg : MessageData) : m α := do
+  withRef ref <| Lean.throwError msg
+
 /--
 Creates a `MessageData` that is tagged with `unknownIdentifierMessageTag`.
 This tag is used by the 'import unknown identifier' code action to detect messages that should
 prompt the code action.
+The end position of the range of an unknown identifier message should always point at the end of the
+unknown identifier.
 -/
 def mkUnknownIdentifierMessage (msg : MessageData) : MessageData :=
   MessageData.tagged unknownIdentifierMessageTag msg
 
 /--
 Throw an unknown identifier error message that is tagged with `unknownIdentifierMessageTag`.
+The end position of the range of `ref` should always point at the unknown identifier.
 See also `mkUnknownIdentifierMessage`.
 -/
-def throwUnknownIdentifier [Monad m] [MonadError m] (msg : MessageData) : m α :=
-  Lean.throwError <| mkUnknownIdentifierMessage msg
+def throwUnknownIdentifierAt [Monad m] [MonadError m] (ref : Syntax) (msg : MessageData) : m α :=
+  Lean.throwErrorAt ref <| mkUnknownIdentifierMessage msg
 
-/-- Throw an unknown constant error message. -/
-def throwUnknownConstant [Monad m] [MonadError m] (constName : Name) : m α :=
-  throwUnknownIdentifier m!"unknown constant '{.ofConstName constName}'"
+/--
+Throw an unknown constant error message.
+The end position of the range of `ref` should point at the unknown identifier.
+See also `mkUnknownIdentifierMessage`.
+-/
+def throwUnknownConstantAt [Monad m] [MonadError m] (ref : Syntax) (constName : Name) : m α := do
+  throwUnknownIdentifierAt ref m!"unknown constant '{.ofConstName constName}'"
 
-/-- Throw an error exception using the given message data and reference syntax. -/
-protected def throwErrorAt [Monad m] [MonadError m] (ref : Syntax) (msg : MessageData) : m α := do
-  withRef ref <| Lean.throwError msg
+/--
+Throw an unknown constant error message.
+The end position of the range of the current reference should point at the unknown identifier.
+See also `mkUnknownIdentifierMessage`.
+-/
+def throwUnknownConstant [Monad m] [MonadError m] (constName : Name) : m α := do
+  throwUnknownConstantAt (← getRef) constName
 
 /--
 Convert an `Except` into a `m` monadic action, where `m` is any monad that

src/Lean/Language/Basic.lean

@@ -82,6 +82,12 @@ structure SnapshotTask (α : Type) where
   `Syntax` processed by this `SnapshotTask`.
   The `Syntax` is used by the language server to determine whether to force this `SnapshotTask`
   when a request is made.
+  In general, the elaborator retains the following invariant:
+  If `stx?` is `none`, then this snapshot task (and all of its children) do not contain `InfoTree`
+  information that can be used in the language server, and so the language server will ignore it
+  when it is looking for an `InfoTree`.
+  Nonetheless, if `stx?` is `none`, then this snapshot task (and any of its children) may still
+  contain message log information.
   -/
   stx? : Option Syntax
   /--

src/Lean/Language/Lean.lean

@@ -687,7 +687,8 @@ where
           -- create a temporary snapshot tree containing all tasks but it
           let snaps := #[
             { stx? := stx', task := elabPromise.result!.map (sync := true) toSnapshotTree, cancelTk? := none },
-            { stx? := stx', task := resultPromise.result!.map (sync := true) toSnapshotTree, cancelTk? := none }] ++
+            { stx? := stx', task := resultPromise.result!.map (sync := true) toSnapshotTree, cancelTk? := none },
+            { stx? := stx', task := finishedPromise.result!.map (sync := true) toSnapshotTree, cancelTk? := none }] ++
             cmdState.snapshotTasks
           let tree := SnapshotTree.mk { diagnostics := .empty } snaps
           BaseIO.bindTask (← tree.waitAll) fun _ => do

src/Lean/Meta/Basic.lean

@@ -1221,7 +1221,7 @@ private def getConstTemp? (constName : Name) : MetaM (Option ConstantInfo) := do
   | some (info@(ConstantInfo.thmInfo _))  => getTheoremInfo info
   | some (info@(ConstantInfo.defnInfo _)) => getDefInfoTemp info
   | some info                             => pure (some info)
-  | none                                  => throwUnknownConstant constName
+  | none                                  => throwUnknownConstantAt (← getRef) constName
 
 private def isClassQuickConst? (constName : Name) : MetaM (LOption Name) := do
   if isClass (← getEnv) constName then

src/Lean/Meta/GetUnfoldableConst.lean

@@ -37,7 +37,7 @@ This is part of the implementation of `whnf`.
 External users wanting to look up names should be using `Lean.getConstInfo`.
 -/
 def getUnfoldableConst? (constName : Name) : MetaM (Option ConstantInfo) := do
-  let some ainfo := (← getEnv).findAsync? constName | throwUnknownConstant constName
+  let some ainfo := (← getEnv).findAsync? constName | throwUnknownConstantAt (← getRef) constName
   match ainfo.kind with
   | .thm =>
     if (← shouldReduceAll) then

src/Lean/ResolveName.lean

@@ -296,7 +296,7 @@ def resolveUniqueNamespace [Monad m] [MonadResolveName m] [MonadEnv m] [MonadErr
 /-- Helper function for `resolveGlobalConstCore`. -/
 def filterFieldList [Monad m] [MonadError m] (n : Name) (cs : List (Name × List String)) : m (List Name) := do
   let cs := cs.filter fun (_, fieldList) => fieldList.isEmpty
-  if cs.isEmpty then throwUnknownConstant n
+  if cs.isEmpty then throwUnknownConstantAt (← getRef) n
   return cs.map (·.1)
 
 /-- Given a name `n`, returns a list of possible interpretations for global constants.

src/Lean/Server/CodeActions/UnknownIdentifier.lean

@@ -23,13 +23,11 @@ structure UnknownIdentifierInfo where
 def waitUnknownIdentifierRanges (doc : EditableDocument) (requestedRange : String.Range)
     : BaseIO (Array String.Range) := do
   let text := doc.meta.text
-  let parsedSnaps := RequestM.findCmdParsedSnaps doc requestedRange |>.get
-  let msgLog := parsedSnaps.map Language.toSnapshotTree
-    |>.map (·.collectMessagesInRange requestedRange)
-    |>.map (·.get)
-    |>.foldl (· ++ ·) .empty
+  let some parsedSnap := RequestM.findCmdParsedSnap doc requestedRange.start |>.get
+    | return #[]
+  let msgLog := Language.toSnapshotTree parsedSnap |>.collectMessagesInRange requestedRange |>.get
   let mut ranges := #[]
-  for msg in msgLog.reportedPlusUnreported do
+  for msg in msgLog.unreported do
     if ! msg.data.hasTag (· == unknownIdentifierMessageTag) then
       continue
     let msgRange : String.Range := ⟨text.ofPosition msg.pos, text.ofPosition <| msg.endPos.getD msg.pos⟩
@@ -39,6 +37,20 @@ def waitUnknownIdentifierRanges (doc : EditableDocument) (requestedRange : Strin
     ranges := ranges.push msgRange
   return ranges
 
+def waitAllUnknownIdentifierRanges (doc : EditableDocument)
+    : BaseIO (Array String.Range) := do
+  let text := doc.meta.text
+  let msgLog : MessageLog := Language.toSnapshotTree doc.initSnap
+    |>.getAll.map (·.diagnostics.msgLog)
+    |>.foldl (· ++ ·) {}
+  let mut ranges := #[]
+  for msg in msgLog.unreported do
+    if ! msg.data.hasTag (· == unknownIdentifierMessageTag) then
+      continue
+    let msgRange : String.Range := ⟨text.ofPosition msg.pos, text.ofPosition <| msg.endPos.getD msg.pos⟩
+    ranges := ranges.push msgRange
+  return ranges
+
 structure Insertion where
   fullName : Name
   edit     : TextEdit
@@ -60,34 +72,6 @@ partial def collectOpenNamespaces (currentNamespace : Name) (openDecls : List Op
   openNamespaces := openNamespaces ++ openDeclNamespaces.toArray
   return openNamespaces
 
-def computeFallbackQuery?
-    (doc : EditableDocument)
-    (requestedRange : String.Range)
-    (unknownIdentifierRange : String.Range)
-    (infoTree : Elab.InfoTree)
-    : Option Query := do
-  let text := doc.meta.text
-  let info? := infoTree.smallestInfo? fun info => Id.run do
-    let some range := info.range?
-      | return false
-    return range.overlaps requestedRange (includeFirstStop := true) (includeSecondStop := true)
-  let some (ctx, _) := info?
-    | none
-  return {
-    identifier := text.source.extract unknownIdentifierRange.start unknownIdentifierRange.stop
-    openNamespaces := collectOpenNamespaces ctx.currNamespace ctx.openDecls
-    env := ctx.env
-    determineInsertion decl :=
-      let minimizedId := minimizeGlobalIdentifierInContext ctx.currNamespace ctx.openDecls decl
-      {
-        fullName := minimizedId
-        edit := {
-          range := text.utf8RangeToLspRange unknownIdentifierRange
-          newText := minimizedId.toString
-        }
-      }
-  }
-
 def computeIdQuery?
     (doc : EditableDocument)
     (ctx : Elab.ContextInfo)
@@ -130,6 +114,8 @@ def computeDotQuery?
       return none
   let some typeNames := typeNames?
     | return none
+  if typeNames.isEmpty then
+    return none
   return some {
     identifier := text.source.extract pos tailPos
     openNamespaces := typeNames.map (.allExcept · #[])
@@ -182,29 +168,32 @@ def computeDotIdQuery?
       }
   }
 
-def computeQuery?
-    (doc                    : EditableDocument)
-    (requestedRange         : String.Range)
-    (unknownIdentifierRange : String.Range)
-    : RequestM (Option Query) := do
+def computeQueries
+    (doc          : EditableDocument)
+    (requestedPos : String.Pos)
+    : RequestM (Array Query) := do
   let text := doc.meta.text
-  let some (stx, infoTree) := RequestM.findCmdDataAtPos doc unknownIdentifierRange.stop (includeStop := true) |>.get
-    | return none
-  let completionInfo? : Option ContextualizedCompletionInfo := do
-    let (completionPartitions, _) := findPrioritizedCompletionPartitionsAt text unknownIdentifierRange.stop stx infoTree
-    let highestPrioPartition ← completionPartitions[0]?
-    let (completionInfo, _) ← highestPrioPartition[0]?
-    return completionInfo
-  let some ⟨_, ctx, info⟩ := completionInfo?
-    | return computeFallbackQuery? doc requestedRange unknownIdentifierRange infoTree
-  match info with
-  | .id (stx := stx) (id := id) .. =>
-    return computeIdQuery? doc ctx stx id
-  | .dot (termInfo := ti) .. =>
-    return ← computeDotQuery? doc ctx ti
-  | .dotId stx id lctx expectedType? =>
-    return ← computeDotIdQuery? doc ctx stx id lctx expectedType?
-  | _ => return none
+  let some (stx, infoTree) := RequestM.findCmdDataAtPos doc requestedPos (includeStop := true) |>.get
+    | return #[]
+  let (completionPartitions, _) := findPrioritizedCompletionPartitionsAt text requestedPos stx infoTree
+  let mut queries := #[]
+  for partition in completionPartitions do
+    for (i, _) in partition do
+      let query? ←
+        match i.info with
+        | .id (stx := stx) (id := id) .. =>
+          pure <| computeIdQuery? doc i.ctx stx id
+        | .dot (termInfo := ti) .. =>
+          computeDotQuery? doc i.ctx ti
+        | .dotId stx id lctx expectedType? =>
+          computeDotIdQuery? doc i.ctx stx id lctx expectedType?
+        | _ =>
+          pure none
+      if let some query := query? then
+        queries := queries.push query
+    if ! queries.isEmpty then
+      break
+  return queries
 
 def importAllUnknownIdentifiersProvider : Name := `unknownIdentifiers
 
@@ -220,16 +209,14 @@ def importAllUnknownIdentifiersCodeAction (params : CodeActionParams) (kind : St
 }
 
 def handleUnknownIdentifierCodeAction
-    (id                      : JsonRpc.RequestID)
-    (params                  : CodeActionParams)
-    (requestedRange          : String.Range)
-    (unknownIdentifierRanges : Array String.Range)
+    (id             : JsonRpc.RequestID)
+    (params         : CodeActionParams)
+    (requestedRange : String.Range)
     : RequestM (Array CodeAction) := do
   let rc ← read
   let doc := rc.doc
   let text := doc.meta.text
-  let queries ← unknownIdentifierRanges.filterMapM fun unknownIdentifierRange =>
-    computeQuery? doc requestedRange unknownIdentifierRange
+  let queries ← computeQueries doc requestedRange.stop
   if queries.isEmpty then
     return #[]
   let responseTask ← RequestM.sendServerRequest LeanQueryModuleParams LeanQueryModuleResponse "$/lean/queryModule" {
@@ -255,13 +242,15 @@ def handleUnknownIdentifierCodeAction
   let importInsertionRange : Lsp.Range := ⟨importInsertionPos, importInsertionPos⟩
   let mut unknownIdentifierCodeActions := #[]
   let mut hasUnambigiousImportCodeAction := false
-  for q in queries, result in response.queryResults do
+  let some result := response.queryResults[0]?
+    | return #[]
+  for query in queries, result in response.queryResults do
     for ⟨mod, decl, isExactMatch⟩ in result do
-      let isDeclInEnv := q.env.contains decl
-      if ! isDeclInEnv && mod == q.env.mainModule then
+      let isDeclInEnv := query.env.contains decl
+      if ! isDeclInEnv && mod == query.env.mainModule then
         -- Don't offer any code actions for identifiers defined further down in the same file
         continue
-      let insertion := q.determineInsertion decl
+      let insertion := query.determineInsertion decl
       if ! isDeclInEnv then
         unknownIdentifierCodeActions := unknownIdentifierCodeActions.push {
           title := s!"Import {insertion.fullName} from {mod}"
@@ -301,8 +290,8 @@ def handleResolveImportAllUnknownIdentifiersCodeAction?
   let rc ← read
   let doc := rc.doc
   let text := doc.meta.text
-  let queries ← unknownIdentifierRanges.filterMapM fun unknownIdentifierRange =>
-    computeQuery? doc ⟨0, text.source.endPos⟩ unknownIdentifierRange
+  let queries ← unknownIdentifierRanges.flatMapM fun unknownIdentifierRange =>
+    computeQueries doc unknownIdentifierRange.stop
   if queries.isEmpty then
     return none
   let responseTask ← RequestM.sendServerRequest LeanQueryModuleParams LeanQueryModuleResponse "$/lean/queryModule" {

src/Lean/Server/FileWorker.lean

@@ -769,8 +769,7 @@ section MessageHandling
       if data.providerName != importAllUnknownIdentifiersProvider then
         return none
       return some <| ← RequestM.asTask do
-        let fileRange := ⟨0, st.doc.meta.text.source.endPos⟩
-        let unknownIdentifierRanges ← waitUnknownIdentifierRanges st.doc fileRange
+        let unknownIdentifierRanges ← waitAllUnknownIdentifierRanges st.doc
         if unknownIdentifierRanges.isEmpty then
           return { response := toJson params, isComplete := true }
         let action? ← handleResolveImportAllUnknownIdentifiersCodeAction? id params unknownIdentifierRanges
@@ -790,7 +789,7 @@ section MessageHandling
         let isSourceAction := params.context.only?.any fun only =>
             only.contains "source" || only.contains "source.organizeImports"
         if isSourceAction then
-          let unknownIdentifierRanges ← waitUnknownIdentifierRanges doc ⟨0, doc.meta.text.source.endPos⟩
+          let unknownIdentifierRanges ← waitAllUnknownIdentifierRanges doc
           if unknownIdentifierRanges.isEmpty then
             return r
           let .ok (codeActions : Array CodeAction) := fromJson? r.response
@@ -808,7 +807,7 @@ section MessageHandling
           -- we only do it when the user has stopped typing for a second.
           IO.sleep 1000
           RequestM.checkCancelled
-          let unknownIdentifierCodeActions ← handleUnknownIdentifierCodeAction id params requestedRange unknownIdentifierRanges
+          let unknownIdentifierCodeActions ← handleUnknownIdentifierCodeAction id params requestedRange
           return { r with response := toJson <| codeActions ++ unknownIdentifierCodeActions }
     | _ =>
       return task

src/Lean/Server/FileWorker/RequestHandling.lean

@@ -254,8 +254,10 @@ def findGoalsAt? (doc : EditableDocument) (hoverPos : String.Pos) : ServerTask (
   findCmdParsedSnap doc hoverPos |>.bindCostly fun
     | some cmdParsed =>
       let t := toSnapshotTree cmdParsed |>.foldSnaps [] fun snap oldGoals => Id.run do
-        let some (pos, tailPos, trailingPos) := getPositions snap
+        let some stx := snap.stx?
           | return .pure (oldGoals, .proceed (foldChildren := false))
+        let some (pos, tailPos, trailingPos) := getPositions stx
+          | return .pure (oldGoals, .proceed (foldChildren := true))
         let snapRange : String.Range := ⟨pos, trailingPos⟩
         -- When there is no trailing whitespace, we also consider snapshots directly before the
         -- cursor.
@@ -283,8 +285,7 @@ def findGoalsAt? (doc : EditableDocument) (hoverPos : String.Pos) : ServerTask (
     | none =>
       .pure none
 where
-  getPositions (snap : SnapshotTask SnapshotTree) : Option (String.Pos × String.Pos × String.Pos) := do
-    let stx ← snap.stx?
+  getPositions (stx : Syntax) : Option (String.Pos × String.Pos × String.Pos) := do
     let pos ← stx.getPos? (canonicalOnly := true)
     let tailPos ← stx.getTailPos? (canonicalOnly := true)
     let trailingPos? ← stx.getTrailingTailPos? (canonicalOnly := true)

src/Lean/Server/Requests.lean

@@ -44,6 +44,18 @@ def Lean.FileMap.rangeOverlapsRequestedRange
     (includeFirstStop := includeDocumentRangeStop || isDocumentRangeAtEOF)
     (includeSecondStop := includeRequestedRangeStop)
 
+def Lean.FileMap.rangeIncludesRequestedRange
+    (text : Lean.FileMap)
+    (documentRange : String.Range)
+    (requestedRange : String.Range)
+    (includeDocumentRangeStop := false)
+    (includeRequestedRangeStop := false)
+    : Bool :=
+  let isDocumentRangeAtEOF := documentRange.stop == text.source.endPos
+  documentRange.includes requestedRange
+    (includeSuperStop := includeDocumentRangeStop || isDocumentRangeAtEOF)
+    (includeSubStop := includeRequestedRangeStop)
+
 namespace Lean.Language
 
 open Lean.Server
@@ -88,13 +100,20 @@ that contains `hoverPos` in its whitespace, which is not necessarily the correct
 partial def SnapshotTree.findInfoTreeAtPos (text : FileMap) (tree : SnapshotTree)
     (hoverPos : String.Pos) (includeStop : Bool) : ServerTask (Option Elab.InfoTree) :=
   tree.foldSnaps (init := none) fun snap _ => Id.run do
-    let skipChild := .pure (none, .proceed (foldChildren := false))
     let some stx := snap.stx?
-      | return skipChild
+      -- One of the invariants of the snapshot tree is that `stx? = none` implies that
+      -- this entire subtree has no relevant `InfoTree` information, so we can safely discard it
+      -- here.
+      | return .pure (none, .proceed (foldChildren := false))
     let some range := stx.getRangeWithTrailing? (canonicalOnly := true)
-      | return skipChild
+      -- In the worst case, the `infoTreeSnap` of the `CommandParsedSnap` will have canonical
+      -- syntax that we can use here, so ignoring snapshots with non-canonical syntax can only
+      -- at worst break incrementality in request handlers.
+      | return .pure (none, .proceed (foldChildren := true))
     if ! text.rangeContainsHoverPos range hoverPos includeStop then
-      return skipChild
+      -- Subtrees of the snapshot tree always have syntax ranges that are contained in those of
+      -- their parents, so we can terminate early here.
+      return .pure (none, .proceed (foldChildren := false))
     return snap.task.asServerTask.mapCheap fun tree => Id.run do
       let some infoTree := tree.element.infoTree?
         | return (none, .proceed (foldChildren := true))
@@ -327,47 +346,40 @@ def withWaitFindSnapAtPos
     (notFoundX := throw ⟨.invalidParams, s!"no snapshot found at {lspPos}"⟩)
     (x := f)
 
-open Language.Lean in
-/-- Finds all `CommandParsedSnapshot`s overlapping `requestedRange`, asynchronously. -/
-partial def findCmdParsedSnaps (doc : EditableDocument) (requestedRange : String.Range)
-    : ServerTask (Array CommandParsedSnapshot) := Id.run do
-  let some headerParsed := doc.initSnap.result?
-    | .pure #[]
-  headerParsed.processedSnap.task.asServerTask.bindCheap fun headerProcessed => Id.run do
-    let some headerSuccess := headerProcessed.result?
-      | return .pure #[]
-    let firstCmdSnapTask : ServerTask CommandParsedSnapshot := headerSuccess.firstCmdSnap.task
-    firstCmdSnapTask.bindCheap (go · #[])
-where
-  go (cmdParsed : CommandParsedSnapshot) (acc : Array CommandParsedSnapshot)
-      : ServerTask (Array CommandParsedSnapshot) := Id.run do
-    if isAfterRequestedRange cmdParsed then
-      return .pure acc
-    let mut acc := acc
-    if overlapsRequestedRange cmdParsed then
-      acc := acc.push cmdParsed
-    match cmdParsed.nextCmdSnap? with
-    | some next =>
-      next.task.asServerTask.bindCheap (go · acc)
-    | none =>
-      return .pure acc
-
-  overlapsRequestedRange (cmdParsed : CommandParsedSnapshot) : Bool := Id.run do
-    let some range := cmdParsed.stx.getRangeWithTrailing? (canonicalOnly := true)
-      | return false
-    return doc.meta.text.rangeOverlapsRequestedRange range requestedRange
-      (includeDocumentRangeStop := false) (includeRequestedRangeStop := true)
-
-  isAfterRequestedRange (cmdParsed : CommandParsedSnapshot) : Bool := Id.run do
-    let some startPos := cmdParsed.stx.getPos? (canonicalOnly := true)
-      | return false
-    return requestedRange.stop < startPos
-
 open Language.Lean in
 /-- Finds the first `CommandParsedSnapshot` containing `hoverPos`, asynchronously. -/
 partial def findCmdParsedSnap (doc : EditableDocument) (hoverPos : String.Pos)
-    : ServerTask (Option CommandParsedSnapshot) :=
-  findCmdParsedSnaps doc ⟨hoverPos, hoverPos⟩ |>.mapCheap (·[0]?)
+    : ServerTask (Option CommandParsedSnapshot) := Id.run do
+  let some headerParsed := doc.initSnap.result?
+    | .pure none
+  headerParsed.processedSnap.task.asServerTask.bindCheap fun headerProcessed => Id.run do
+    let some headerSuccess := headerProcessed.result?
+      | return .pure none
+    let firstCmdSnapTask : ServerTask CommandParsedSnapshot := headerSuccess.firstCmdSnap.task
+    firstCmdSnapTask.bindCheap go
+where
+  go (cmdParsed : CommandParsedSnapshot) : ServerTask (Option CommandParsedSnapshot) := Id.run do
+    if containsHoverPos cmdParsed then
+      return .pure (some cmdParsed)
+    if isAfterHoverPos cmdParsed then
+      -- This should never happen in principle
+      -- (commands + trailing ws are consecutive and there is no unassigned space between them),
+      -- but it's always good to eliminate one additional assumption.
+      return .pure none
+    match cmdParsed.nextCmdSnap? with
+    | some next =>
+      next.task.asServerTask.bindCheap go
+    | none => .pure none
+
+  containsHoverPos (cmdParsed : CommandParsedSnapshot) : Bool := Id.run do
+    let some range := cmdParsed.stx.getRangeWithTrailing? (canonicalOnly := true)
+      | return false
+    return doc.meta.text.rangeContainsHoverPos range hoverPos (includeStop := false)
+
+  isAfterHoverPos (cmdParsed : CommandParsedSnapshot) : Bool := Id.run do
+    let some startPos := cmdParsed.stx.getPos? (canonicalOnly := true)
+      | return false
+    return hoverPos < startPos
 
 open Language in
 /--

src/Lean/Server/ServerTask.lean

@@ -62,6 +62,12 @@ def bindCheap (t : ServerTask α) (f : α → ServerTask β) : ServerTask β :=
 def bindCostly (t : ServerTask α) (f : α → ServerTask β) : ServerTask β :=
   t.task.bind (f · |>.task) (prio := .dedicated)
 
+def join (ts : Array (ServerTask α)) : ServerTask (Array α) := Id.run do
+  let mut r := ServerTask.pure #[]
+  for t in ts do
+    r := r.bindCheap fun acc => t.mapCheap (acc.push ·)
+  return r
+
 namespace BaseIO
 
 def asTask (act : BaseIO α) : BaseIO (ServerTask α) :=

src/Lean/Syntax.lean

@@ -22,8 +22,21 @@ protected structure String.Range where
 def String.Range.contains (r : String.Range) (pos : String.Pos) (includeStop := false) : Bool :=
   r.start <= pos && (if includeStop then pos <= r.stop else pos < r.stop)
 
-def String.Range.includes (super sub : String.Range) : Bool :=
-  super.start <= sub.start && super.stop >= sub.stop
+/--
+Checks whether `sub` is contained in `super`.
+`includeSuperStop` and `includeSubStop` control whether `super` and `sub` have
+an inclusive upper bound.
+-/
+def String.Range.includes (super sub : String.Range)
+    (includeSuperStop := false) (includeSubStop := false) : Bool :=
+  super.start <= sub.start && (
+    if includeSuperStop && !includeSubStop then
+      sub.stop.byteIdx <= super.stop.byteIdx + 1
+    else if !includeSuperStop && includeSubStop then
+      sub.stop < super.stop
+    else
+      sub.stop <= super.stop
+  )
 
 def String.Range.overlaps (first second : String.Range)
     (includeFirstStop := false) (includeSecondStop := false) : Bool :=

tests/lean/scientific.lean.expected.out

@@ -21,7 +21,7 @@ scientific.lean:22:9: error: missing exponent digits in scientific literal
 scientific.lean:23:9: error: missing exponent digits in scientific literal
 scientific.lean:24:9: error: missing exponent digits in scientific literal
 scientific.lean:25:9: error: missing exponent digits in scientific literal
-scientific.lean:26:6-26:8: error: invalid dotted identifier notation, unknown identifier `Nat.E` from expected type
+scientific.lean:26:7-26:8: error: invalid dotted identifier notation, unknown identifier `Nat.E` from expected type
   Nat
 scientific.lean:27:7-27:9: error: unknown identifier 'E3'
 scientific.lean:28:7: error: missing exponent digits in scientific literal