feat: client-side support for incremental diagnostics (#752)

This PR adds client-side support for
[lean4#13260](leanprover/lean4#13260).

Currently, none of the consumers of diagnostics in the VS Code extension
support incremental diagnostics, but adding client-side support for it
already means that the server-side serialization overhead is reduced and
we send less diagnostic data over the wire.

Author: mhuisi

vscode-lean4/src/diagnostics.ts

@@ -0,0 +1,120 @@
+import { CancellationToken, DiagnosticCollection, Disposable, EventEmitter } from 'vscode'
+import { DocumentUri } from 'vscode-languageserver-protocol'
+import { CoalescingSyncQueue } from './utils/coalescingSyncQueue'
+import { LeanDiagnostic, LeanPublishDiagnosticsParams, p2cConverter } from './utils/converters'
+
+export type DiagnosticChangeKind = 'replace' | 'append'
+
+export class DiagnosticChangeEvent {
+    readonly kind: DiagnosticChangeKind
+    readonly params: LeanPublishDiagnosticsParams
+    private readonly accumulated: LeanDiagnostic[]
+    private readonly amountAccumulated: number
+
+    constructor(kind: DiagnosticChangeKind, params: LeanPublishDiagnosticsParams, accumulated: LeanDiagnostic[]) {
+        this.kind = kind
+        this.params = params
+        this.accumulated = accumulated
+        this.amountAccumulated = accumulated.length
+    }
+
+    private accumulatedDiagnostics(): LeanDiagnostic[] {
+        return this.accumulated.slice(0, this.amountAccumulated)
+    }
+
+    accumulatedParams(): LeanPublishDiagnosticsParams {
+        return { ...this.params, diagnostics: this.accumulatedDiagnostics() }
+    }
+}
+
+type SyncQueueEntry = {
+    accumulatedParams: LeanPublishDiagnosticsParams
+    pendingKind: DiagnosticChangeKind
+    pendingBatch: LeanDiagnostic[]
+}
+
+function combineEntries(existing: SyncQueueEntry, incoming: SyncQueueEntry): SyncQueueEntry {
+    if (incoming.pendingKind === 'replace') {
+        return incoming
+    }
+    incoming.pendingKind satisfies 'append'
+    return {
+        accumulatedParams: incoming.accumulatedParams,
+        pendingKind: existing.pendingKind,
+        pendingBatch: [...existing.pendingBatch, ...incoming.pendingBatch],
+    }
+}
+
+export class LeanClientDiagnosticCollection implements Disposable {
+    readonly vsCodeCollection: DiagnosticCollection
+    private diags: Map<DocumentUri, LeanPublishDiagnosticsParams> = new Map()
+
+    private diagnosticsChangedEmitter = new EventEmitter<DiagnosticChangeEvent>()
+    onDidChangeDiagnostics = this.diagnosticsChangedEmitter.event
+
+    private syncQueue: CoalescingSyncQueue<SyncQueueEntry>
+
+    constructor(vsCodeCollection: DiagnosticCollection) {
+        this.vsCodeCollection = vsCodeCollection
+        this.syncQueue = new CoalescingSyncQueue(
+            (uri: string, entry: SyncQueueEntry, token: CancellationToken) => this.syncToCollection(uri, entry, token),
+            (existing, incoming) => combineEntries(existing, incoming)
+        )
+    }
+
+    private static determineChangeKind(
+        prev: LeanPublishDiagnosticsParams | undefined,
+        next: LeanPublishDiagnosticsParams,
+    ): DiagnosticChangeKind {
+        if (prev === undefined) {
+            return 'replace'
+        }
+        if (!next.isIncremental) {
+            return 'replace'
+        }
+        return 'append'
+    }
+
+    publishDiagnostics(params: LeanPublishDiagnosticsParams): void {
+        const prev = this.diags.get(params.uri)
+        const kind = LeanClientDiagnosticCollection.determineChangeKind(prev, params)
+
+        let accumulated: LeanDiagnostic[]
+        if (kind === 'append') {
+            accumulated = prev!.diagnostics
+            accumulated.push(...params.diagnostics)
+        } else {
+            accumulated = [...params.diagnostics]
+        }
+
+        const accumulatedParams = { ...params, diagnostics: accumulated }
+        this.diags.set(accumulatedParams.uri, accumulatedParams)
+
+        const entry: SyncQueueEntry = {
+            accumulatedParams,
+            pendingKind: kind,
+            pendingBatch: [...params.diagnostics],
+        }
+
+        this.syncQueue.enqueue(accumulatedParams.uri, entry)
+    }
+
+    private async syncToCollection(_uri: string, entry: SyncQueueEntry, token: CancellationToken): Promise<void> {
+        const nonSilentDiagnostics = entry.accumulatedParams.diagnostics.filter(d => !d.isSilent)
+        const vsUri = p2cConverter.asUri(entry.accumulatedParams.uri)
+        const vsDiags = await p2cConverter.asDiagnostics(nonSilentDiagnostics, token)
+        if (token.isCancellationRequested) {
+            return
+        }
+        this.vsCodeCollection.set(vsUri, vsDiags)
+        const collapsedParams = { ...entry.accumulatedParams, diagnostics: entry.pendingBatch }
+        this.diagnosticsChangedEmitter.fire(
+            new DiagnosticChangeEvent(entry.pendingKind, collapsedParams, entry.accumulatedParams.diagnostics),
+        )
+    }
+
+    dispose(): void {
+        this.syncQueue.dispose()
+        this.diagnosticsChangedEmitter.dispose()
+    }
+}

vscode-lean4/src/leanclient.ts

@@ -4,6 +4,7 @@ import {
     Disposable,
     EventEmitter,
     FileSystemWatcher,
+    languages,
     OutputChannel,
     Progress,
     ProgressLocation,
@@ -32,7 +33,6 @@ import {
 } from 'vscode-languageclient/node'
 
 import {
-    LeanDiagnostic,
     LeanFileProgressParams,
     LeanFileProgressProcessingInfo,
     LeanServerCapabilities,
@@ -52,6 +52,7 @@ import { logger } from './utils/logger'
 import * as fs from 'fs'
 import { glob } from 'glob'
 import * as semver from 'semver'
+import { DiagnosticChangeEvent, LeanClientDiagnosticCollection } from './diagnostics'
 import { formatCommandExecutionOutput } from './utils/batch'
 import {
     c2pConverter,
@@ -98,11 +99,13 @@ function logConfig(): LogConfig | undefined {
 }
 
 interface LeanClientCapabilties {
+    incrementalDiagnosticSupport?: boolean
     silentDiagnosticSupport?: boolean
     rpcWireFormat?: RpcWireFormat
 }
 
 const leanClientCapabilities: LeanClientCapabilties = {
+    incrementalDiagnosticSupport: true,
     silentDiagnosticSupport: true,
     rpcWireFormat: 'v1'
 }
@@ -140,6 +143,8 @@ export class LeanClient implements Disposable {
     private didChangeEmitter = new EventEmitter<DidChangeTextDocumentParams>()
     didChange = this.didChangeEmitter.event
 
+    private diagnosticCollection: LeanClientDiagnosticCollection | undefined
+
     private diagnosticsEmitter = new EventEmitter<LeanPublishDiagnosticsParams>()
     diagnostics = this.diagnosticsEmitter.event
 
@@ -540,7 +545,18 @@ export class LeanClient implements Disposable {
                     }
                 }
             })
+            this.diagnosticCollection?.dispose()
+            const vsCodeCollection = languages.createDiagnosticCollection('lean4')
+            this.diagnosticCollection = new LeanClientDiagnosticCollection(vsCodeCollection)
+            this.diagnosticCollection.onDidChangeDiagnostics((e: DiagnosticChangeEvent) => {
+                this.diagnosticsEmitter.fire(e.accumulatedParams())
+            })
+            this.client.onNotification('textDocument/publishDiagnostics', (params: LeanPublishDiagnosticsParams) => {
+                this.diagnosticCollection?.publishDiagnostics(params)
+            })
+
             await this.client.start()
+
             const rawVersion = this.client.initializeResult?.serverInfo?.version
             const version = rawVersion !== undefined ? semver.parse(rawVersion) : null
             if (version !== null && version.compare('0.2.0') < 0) {
@@ -709,6 +725,9 @@ export class LeanClient implements Disposable {
 
         this.noPrompt = false
         this.progress = new Map()
+        this.diagnosticCollection?.vsCodeCollection.dispose()
+        this.diagnosticCollection?.dispose()
+        this.diagnosticCollection = undefined
         this.client = undefined
         this.openServerDocuments = new Set()
         this.running = false
@@ -780,7 +799,7 @@ export class LeanClient implements Disposable {
     }
 
     getDiagnostics(): DiagnosticCollection | undefined {
-        return this.running ? this.client?.diagnostics : undefined
+        return this.running ? this.diagnosticCollection?.vsCodeCollection : undefined
     }
 
     get initializeResult(): InitializeResult | undefined {
@@ -874,18 +893,11 @@ export class LeanClient implements Disposable {
                     }
                     return next(type, param)
                 },
-                handleDiagnostics: (uri, diagnostics, next) => {
-                    const diagnosticsInVsCode = diagnostics.filter(d => !('isSilent' in d && d.isSilent))
-                    next(uri, diagnosticsInVsCode)
-                    const uri_ = c2pConverter.asUri(uri)
-                    const diagnostics_: LeanDiagnostic[] = []
-                    for (const d of diagnostics) {
-                        const d_: LeanDiagnostic = {
-                            ...c2pConverter.asDiagnostic(d),
-                        }
-                        diagnostics_.push(d_)
-                    }
-                    this.diagnosticsEmitter.fire({ uri: uri_, diagnostics: diagnostics_ })
+                handleDiagnostics: () => {
+                    // Suppress the default diagnostics handling by the LanguageClient.
+                    // We handle diagnostics ourselves via a custom notification handler
+                    // so that we can access the full LeanPublishDiagnosticsParams
+                    // (including version and isIncremental).
                 },
 
                 didOpen: async (doc, next) => {

vscode-lean4/src/utils/coalescingSyncQueue.ts

@@ -0,0 +1,79 @@
+import { CancellationToken, CancellationTokenSource, Disposable } from 'vscode'
+
+/**
+ * A queue that coalesces pending work by string key, processes one entry at a time,
+ * and supports cancellation of in-flight work when the same key is re-enqueued.
+ *
+ * When a key is enqueued while its worker is already running, the in-flight
+ * CancellationToken is cancelled. The worker should check the token and bail early.
+ * If the worker completes without cancellation, the entry is removed from the queue.
+ * If cancelled or errored, the entry stays so the next enqueue can update it.
+ */
+export class CoalescingSyncQueue<T> implements Disposable {
+    private queue: Map<string, T> = new Map()
+    private state: { kind: 'idle' } | { kind: 'busy'; key: string; tokenSource: CancellationTokenSource } = {
+        kind: 'idle',
+    }
+    private disposed = false
+
+    constructor(
+        private readonly worker: (key: string, value: T, token: CancellationToken) => Promise<void>,
+        private readonly merge?: (existing: T, incoming: T) => T,
+    ) {}
+
+    enqueue(key: string, value: T): void {
+        if (this.state.kind === 'busy' && this.state.key === key) {
+            this.state.tokenSource.cancel()
+        }
+
+        const existing = this.queue.get(key)
+        if (this.merge !== undefined && existing !== undefined) {
+            this.queue.set(key, this.merge(existing, value))
+        } else {
+            this.queue.set(key, value)
+        }
+
+        if (this.state.kind === 'idle') {
+            this.trigger()
+        }
+    }
+
+    private trigger(): void {
+        queueMicrotask(() => void this.work())
+    }
+
+    private async work(): Promise<void> {
+        if (this.state.kind === 'busy' || this.disposed) {
+            return
+        }
+
+        const next = this.queue.entries().next()
+        if (next.done) {
+            return
+        }
+        const [key, value] = next.value
+
+        const tokenSource = new CancellationTokenSource()
+        this.state = { kind: 'busy', key, tokenSource }
+        try {
+            await this.worker(key, value, tokenSource.token)
+            if (!tokenSource.token.isCancellationRequested && !this.disposed) {
+                this.queue.delete(key)
+            }
+        } finally {
+            tokenSource.dispose()
+            this.state = { kind: 'idle' }
+            if (this.queue.size > 0 && !this.disposed) {
+                this.trigger()
+            }
+        }
+    }
+
+    dispose(): void {
+        this.disposed = true
+        if (this.state.kind === 'busy') {
+            this.state.tokenSource.cancel()
+        }
+        this.queue.clear()
+    }
+}

vscode-lean4/src/utils/converters.ts

@@ -32,6 +32,7 @@ export interface LeanDiagnostic extends ls.Diagnostic {
 export interface LeanPublishDiagnosticsParams {
     uri: ls.DocumentUri
     version?: ls.integer
+    isIncremental?: boolean
     diagnostics: LeanDiagnostic[]
 }