fix: broken document symbols (#460)

This PR fixes an issue reported at
https://leanprover.zulipchat.com/#narrow/stream/113489-new-members/topic/VSCODE.20outline.20mode.20not.20working.20in.20imports/near/442000461
that caused the VS Code outline and other features that rely on document
symbols to not display anything when initially opening the document.

The synopsis is as follows (in case anyone stumbles upon related issues
in the future):

1. By default, VS Code triggers `onDidOpenTextDocument` when hovering
over an identifier while holding `Ctrl`, which in turn makes the
language client library emit a `didOpen` notification to the language
server (closely followed by `didClose` when moving the mouse away from
the identifier).
2. This puts unnecessary pressure on the language server since it starts
elaborating the file, especially given that VS Code does not emit any
further requests in this situation.
  microsoft/vscode#78453
  leanprover/lean4#367
3. Hence, [it was
suggested](microsoft/vscode#78453 (comment))
to filter these notifications in the language client middleware by only
sending these notifications for `window.visibleTextEditors`, which does
not contain "synthetic" documents like the one opened in the `Ctrl`
hover.
4. This doesn't work because `window.visibleTextEditors` is not
up-to-date yet when the corresponding `onDidOpenTextDocument` handler is
executed.
5. To work around this,
#17 filtered all
`didOpen` notifications emitted by the client and registered its own
handler to emit these notifications. This handler was executed *after*
`window.visibleTextEditors` had been updated and so could use
`window.visibleTextEditors` to decide for which files `didOpen` should
be emitted. See
microsoft/vscode-languageserver-node#848 for
the currently still open issue to properly fix this in the language
client library.
6. This worked fine until we added language server support for requests
that are commonly emitted right after the file is opened (e.g. document
symbols).
7. Since the language client registers its own handlers for emitting
requests and vscode-lean4 manages `didOpen` notifications, there's now
the potential for a race condition: If the handlers from different
sources are called in the wrong order and a file is opened after a
request is emitted, the language server may respond to the request with
an error since it doesn't know about the document yet.
8. VS Code does not frequently issue these requests if the file doesn't
change, so the empty state sticks around until the document is changed
(or even longer for some requests).
9. This causes an empty document symbols view when initially opening the
document, which is especially critical when using 'Go to Definition'.

In VS Code version 1.67 (long after the initial vscode-lean4 work-around
for this issue), a new `window.tabGroups` API was added that allows
querying all documents that are currently open in VS Code tabs.
`window.tabGroups` is fortunately already up-to-date when
`onDidOpenTextDocument` is called, so we can now use this information to
directly filter the `didOpen` notifications in the language client
middleware without registering our own handler that introduces a race
condition. This fixes the bug.

One strange thing I noticed while debugging this issue is that VS Code
actually emits multiple requests of the same kind when the document is
first opened. One of these typically fails because the document is not
open yet whereas the second one typically succeeds. But for some reason
(I suspect a bug), VS Code does not actually end up using the result of
the successful request in this case. This lead to a very confusing
situation during debugging where both the language server and the
language client library seemed to return a correct result, but VS Code
still produced an empty outline.

I also suspect that this issue is one cause of other non-deterministic
issues that we have previously encountered in the language server. For
example, while we already fixed this issue using other means, the
non-deterministic behavior of semantic highlighting sometimes being
broken (leanprover/lean4#2977) may also have
been partially caused by this as VS Code queries the semantic tokens for
the full document right after opening the document as well and only
rarely re-queries them afterwards.

This PR also replaces our own tracking of documents opened in the
language server with the open documents tracked by the language client
library. I suspect that the language client library has a bug here (the
set of `openDocuments` in the language client library will probably also
contain documents that were filtered in the `didOpen` middleware if the
document was opened while starting the client), but it shouldn't affect
us as we only filter `didOpen` notifications for documents that can only
be opened after the language client was already launched.

vscode-lean4/src/leanclient.ts

@@ -1,5 +1,4 @@
 import {
-    Diagnostic,
     DiagnosticCollection,
     Disposable,
     DocumentHighlight,
@@ -32,7 +31,6 @@ import * as ls from 'vscode-languageserver-protocol'
 
 import { LeanFileProgressParams, LeanFileProgressProcessingInfo, ServerStoppedReason } from '@leanprover/infoview-api'
 import {
-    automaticallyBuildDependencies,
     getElaborationDelay,
     getFallBackToStringOccurrenceHighlighting,
     serverArgs,
@@ -45,7 +43,8 @@ import { logger } from './utils/logger'
 // @ts-ignore
 import { SemVer } from 'semver'
 import { c2pConverter, p2cConverter, patchConverters } from './utils/converters'
-import { ExtUri, parseExtUri, toExtUri } from './utils/exturi'
+import { collectAllOpenLeanDocumentUris } from './utils/docInfo'
+import { ExtUri, extUriEquals, parseExtUri, toExtUri } from './utils/exturi'
 import {
     displayError,
     displayErrorWithOptionalInput,
@@ -59,10 +58,6 @@ const escapeRegExp = (s: string) => s.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')
 
 export type ServerProgress = Map<ExtUri, LeanFileProgressProcessingInfo[]>
 
-export function getFullRange(diag: Diagnostic): Range {
-    return (diag as any)?.fullRange || diag.range
-}
-
 export class LeanClient implements Disposable {
     running: boolean
     private client: LanguageClient | undefined
@@ -112,9 +107,6 @@ export class LeanClient implements Disposable {
     private serverFailedEmitter = new EventEmitter<string>()
     serverFailed = this.serverFailedEmitter.event
 
-    /** Files which are open. */
-    private isOpen: Map<string, TextDocument> = new Map()
-
     constructor(folderUri: ExtUri, outputChannel: OutputChannel, elanDefaultToolchain: string) {
         this.outputChannel = outputChannel // can be null when opening adhoc files.
         this.folderUri = folderUri
@@ -233,11 +225,6 @@ export class LeanClient implements Disposable {
                 this.staleDepNotifier = this.diagnostics(params => this.checkForImportsOutdatedError(params))
             }
 
-            // tell the new client about the documents that are already open!
-            for (const key of this.isOpen.keys()) {
-                const doc = this.isOpen.get(key)
-                if (doc) this.notifyDidOpen(doc)
-            }
             // if we got this far then the client is happy so we are running!
             this.running = true
         } catch (error) {
@@ -340,30 +327,14 @@ export class LeanClient implements Disposable {
         return 'Success'
     }
 
-    async openLean4Document(doc: TextDocument) {
-        if (this.isOpen.has(doc.uri.toString())) return
-
-        this.isOpen.set(doc.uri.toString(), doc)
-
-        if (!this.running) return // there was a problem starting lean server.
-
-        // didOpenEditor may have also changed the language, so we fire the
-        // event here because the InfoView should be wired up to receive it now.
-        this.didSetLanguageEmitter.fire(doc.languageId)
+    notifyDidOpen(doc: TextDocument) {
+        if (this.client === undefined) {
+            return
+        }
 
-        this.notifyDidOpen(doc)
-    }
+        const params = this.client.code2ProtocolConverter.asOpenTextDocumentParams(doc)
 
-    notifyDidOpen(doc: TextDocument) {
-        void this.client?.sendNotification('textDocument/didOpen', {
-            textDocument: {
-                uri: doc.uri.toString(),
-                languageId: doc.languageId,
-                version: 1,
-                text: doc.getText(),
-            },
-            dependencyBuildMode: automaticallyBuildDependencies() ? 'always' : 'never',
-        })
+        void this.client.sendNotification('textDocument/didOpen', params)
     }
 
     isInFolderManagedByThisClient(uri: ExtUri): boolean {
@@ -441,15 +412,7 @@ export class LeanClient implements Disposable {
                 uri,
             },
         })
-        void this.client?.sendNotification('textDocument/didOpen', {
-            textDocument: {
-                uri,
-                languageId: 'lean4',
-                version: 1,
-                text: doc.getText(),
-            },
-            dependencyBuildMode: automaticallyBuildDependencies() ? 'always' : 'once',
-        })
+        this.notifyDidOpen(doc)
         this.restartedWorkerEmitter.fire(uri)
     }
 
@@ -544,7 +507,6 @@ export class LeanClient implements Disposable {
             middleware: {
                 handleDiagnostics: (uri, diagnostics, next) => {
                     next(uri, diagnostics)
-                    if (!this.client) return
                     const uri_ = c2pConverter.asUri(uri)
                     const diagnostics_ = []
                     for (const d of diagnostics) {
@@ -556,22 +518,31 @@ export class LeanClient implements Disposable {
                     this.diagnosticsEmitter.fire({ uri: uri_, diagnostics: diagnostics_ })
                 },
 
-                didOpen: async () => {
-                    // Note: as per the LSP spec: An open notification must not be sent more than once
-                    // without a corresponding close notification send before. This means open and close
-                    // notification must be balanced and the max open count for a particular textDocument
-                    // is one.  So this even does nothing the notification is handled by the
-                    // openLean4Document method below after the 'lean4' languageId is established and
-                    // it has weeded out documents opened to invisible editors (like 'git:' schemes and
-                    // invisible editors created for Ctrl+Hover events.  A side effect of unbalanced
-                    // open/close notification is leaking 'lean --worker' processes.
-                    // See https://github.com/microsoft/vscode/issues/78453).
-                    return
+                didOpen: async (doc, next) => {
+                    const docUri = toExtUri(doc.uri)
+                    if (!docUri) {
+                        return // This should never happen since the glob we launch the client for ensures that all uris are ext uris
+                    }
+
+                    const openDocUris: ExtUri[] = collectAllOpenLeanDocumentUris()
+                    const docIsOpen = openDocUris.some(openDocUri => extUriEquals(openDocUri, docUri))
+
+                    if (!docIsOpen) {
+                        // The language client library emits a `didOpen` notification when hovering over an identifier while holding `Ctrl` in order to provide a preview for the line that the definition is on.
+                        // In Lean, this is very expensive and hence does not make much sense, so we filter these notification here.
+                        // Should VS Code decide to send requests to a file that was filtered here, the language server will respond with an error, which VS Code will silently discard and interpret as having received an empty response.
+                        // See https://github.com/microsoft/vscode/issues/78453 (the solution suggested in the thread is wrong, but `collectAllOpenLeanDocumentUris` works).
+                        return
+                    }
+
+                    await next(doc)
+
+                    // Opening the document may have set the language ID.
+                    this.didSetLanguageEmitter.fire(doc.languageId)
                 },
 
                 didChange: async (data, next) => {
                     await next(data)
-                    if (!this.running || !this.client) return // there was a problem starting lean server.
                     const params = c2pConverter.asChangeTextDocumentParams(
                         data,
                         data.document.uri,
@@ -581,11 +552,30 @@ export class LeanClient implements Disposable {
                 },
 
                 didClose: async (doc, next) => {
-                    if (!this.isOpen.delete(doc.uri.toString())) {
+                    const docUri = toExtUri(doc.uri)
+                    if (!docUri) {
+                        return // This should never happen since the glob we launch the client for ensures that all uris are ext uris
+                    }
+
+                    // There is a bug I noticed in the language client library where `openDocuments` will also contain documents
+                    // that were filtered in the `didOpen` middleware if the document was opened while starting the client (as of 8.1.0).
+                    // Fortunately for us, our middleware should only filter synthetic documents that can only be opened
+                    // after launching the language server first, so this should never be an issue.
+                    const openDocuments = Array.from(
+                        this.client?.getFeature('textDocument/didOpen').openDocuments ?? [],
+                    )
+                    const docIsOpen = openDocuments.some(openDoc => {
+                        const openDocUri = toExtUri(openDoc.uri)
+                        return openDocUri !== undefined && extUriEquals(openDocUri, docUri)
+                    })
+                    if (!docIsOpen) {
+                        // Do not send `didClose` if we filtered the corresponding `didOpen` (see comment in the `didOpen` middleware).
+                        // The language server is only resilient against requests for closed files, not the `didClose` notification itself.
                         return
                     }
+
                     await next(doc)
-                    if (!this.running || !this.client) return // there was a problem starting lean server.
+
                     const params = c2pConverter.asCloseTextDocumentParams(doc)
                     this.didCloseEmitter.fire(params)
                 },

vscode-lean4/src/utils/clientProvider.ts

@@ -1,5 +1,5 @@
 import { LeanFileProgressProcessingInfo, ServerStoppedReason } from '@leanprover/infoview-api'
-import { Disposable, EventEmitter, OutputChannel, TextDocument, TextEditor, commands, window, workspace } from 'vscode'
+import { Disposable, EventEmitter, OutputChannel, TextDocument, commands, window, workspace } from 'vscode'
 import {
     checkAll,
     checkIsLakeInstalledCorrectly,
@@ -135,15 +135,6 @@ export class LeanClientProvider implements Disposable {
         this.processingInstallChanged = false
     }
 
-    private getVisibleEditor(uri: ExtUri): TextEditor | undefined {
-        for (const editor of window.visibleTextEditors) {
-            if (uri.equalsUri(editor.document.uri)) {
-                return editor
-            }
-        }
-        return undefined
-    }
-
     private restartFile() {
         if (!this.activeClient || !this.activeClient.isRunning()) {
             displayError('No active client.')
@@ -185,20 +176,7 @@ export class LeanClientProvider implements Disposable {
             return
         }
 
-        if (!this.getVisibleEditor(uri)) {
-            // Sometimes VS code opens a document that has no editor yet.
-            // For example, this happens when the vs code opens files to get git
-            // information using a "git:" Uri scheme:
-            //  git:/d%3A/Temp/lean_examples/Foo/Foo/Hello.lean.git?%7B%22path%22%3A%22d%3A%5C%5CTemp%5C%5Clean_examples%5C%5CFoo%5C%5CFoo%5C%5CHello.lean%22%2C%22ref%22%3A%22%22%7D
-            return
-        }
-
-        const [_, client] = await this.ensureClient(uri)
-        if (!client) {
-            return
-        }
-
-        await client.openLean4Document(document)
+        await this.ensureClient(uri)
     }
 
     // Find the client for a given document.

vscode-lean4/src/utils/converters.ts

@@ -16,6 +16,7 @@ import { createConverter as createC2PConverter } from 'vscode-languageclient/lib
 import { createConverter as createP2CConverter } from 'vscode-languageclient/lib/common/protocolConverter'
 import * as async from 'vscode-languageclient/lib/common/utils/async'
 import * as ls from 'vscode-languageserver-protocol'
+import { automaticallyBuildDependencies } from '../config'
 
 interface Lean4Diagnostic extends ls.Diagnostic {
     fullRange: ls.Range
@@ -66,6 +67,14 @@ export function patchConverters(p2cConverter: Protocol2CodeConverter, c2pConvert
     }
     c2pConverter.asDiagnostics = async (diags, token) => async.map(diags, d => c2pConverter.asDiagnostic(d), token)
 
+    // eslint-disable-next-line @typescript-eslint/unbound-method
+    const oldC2pAsOpenTextDocumentParams = c2pConverter.asOpenTextDocumentParams
+    c2pConverter.asOpenTextDocumentParams = doc => {
+        const params = oldC2pAsOpenTextDocumentParams.apply(this, [doc])
+        params.dependencyBuildMode = automaticallyBuildDependencies() ? 'always' : 'once'
+        return params
+    }
+
     // eslint-disable-next-line @typescript-eslint/unbound-method
     const oldP2CAsWorkspaceEdit = p2cConverter.asWorkspaceEdit
     p2cConverter.asWorkspaceEdit = async function (

vscode-lean4/src/utils/docInfo.ts

@@ -0,0 +1,39 @@
+import { TabInputText, TextDocument, window, workspace } from 'vscode'
+import { ExtUri, toExtUri } from './exturi'
+
+export function collectAllOpenLeanDocuments(): TextDocument[] {
+    const documentsByUri: Map<string, TextDocument> = new Map<string, TextDocument>()
+    for (const doc of workspace.textDocuments) {
+        documentsByUri.set(doc.uri.toString(), doc)
+    }
+
+    const visibleDocs: TextDocument[] = []
+    for (const tab of window.tabGroups.all.flatMap(group => group.tabs)) {
+        if (!(tab.input instanceof TabInputText)) {
+            continue
+        }
+        const uri = toExtUri(tab.input.uri)
+        if (uri === undefined) {
+            continue
+        }
+
+        const doc = documentsByUri.get(uri.toString())
+        if (doc === undefined || doc.languageId !== 'lean4') {
+            continue
+        }
+
+        visibleDocs.push(doc)
+    }
+
+    return visibleDocs
+}
+
+export function collectAllOpenLeanDocumentUris(): ExtUri[] {
+    return collectAllOpenLeanDocuments().map(doc => {
+        const uri = toExtUri(doc.uri)
+        if (uri === undefined) {
+            throw new Error()
+        }
+        return uri
+    })
+}