Create own notebook documents service

packages/langium/src/lsp/default-lsp-module.ts

@@ -23,7 +23,7 @@ import { DefaultNodeKindProvider } from './node-kind-provider.js';
 import { DefaultReferencesProvider } from './references-provider.js';
 import { DefaultRenameProvider } from './rename-provider.js';
 import { DefaultWorkspaceSymbolProvider } from './workspace-symbol-provider.js';
-import { NormalizedTextDocuments } from './normalized-text-documents.js';
+import { NormalizedNotebookDocuments, NormalizedTextDocuments } from './normalized-text-documents.js';
 
 /**
  * Context required for creating the default language-specific dependency injection module.
@@ -96,7 +96,8 @@ export function createDefaultSharedLSPModule(context: DefaultSharedModuleContext
             FuzzyMatcher: () => new DefaultFuzzyMatcher(),
         },
         workspace: {
-            TextDocuments: () => new NormalizedTextDocuments(TextDocument)
+            TextDocuments: () => new NormalizedTextDocuments(TextDocument),
+            NotebookDocuments: () => new NormalizedNotebookDocuments(TextDocument)
         }
     };
 }

packages/langium/src/lsp/lsp-services.ts

@@ -34,7 +34,7 @@ import type { SignatureHelpProvider } from './signature-help-provider.js';
 import type { TypeHierarchyProvider } from './type-hierarchy-provider.js';
 import type { TypeDefinitionProvider } from './type-provider.js';
 import type { WorkspaceSymbolProvider } from './workspace-symbol-provider.js';
-import type { TextDocuments } from './normalized-text-documents.js';
+import type { NotebookDocuments, TextDocuments } from './normalized-text-documents.js';
 
 /**
  * Combined Core + LSP services of Langium (total services)
@@ -91,6 +91,7 @@ export type LangiumSharedLSPServices = {
     },
     readonly workspace: {
         readonly TextDocuments: TextDocuments<TextDocument>
+        readonly NotebookDocuments: NotebookDocuments<TextDocument>
     }
 };
 

packages/langium/src/lsp/normalized-text-documents.ts

@@ -6,11 +6,26 @@
 
 import type {
     Connection, DidOpenTextDocumentParams, DidChangeTextDocumentParams, DidCloseTextDocumentParams, TextDocumentsConfiguration, TextDocumentChangeEvent,
-    TextDocumentWillSaveEvent, RequestHandler, TextEdit, Event, WillSaveTextDocumentParams, CancellationToken, DidSaveTextDocumentParams
+    TextDocumentWillSaveEvent, RequestHandler, TextEdit, Event, WillSaveTextDocumentParams, CancellationToken, DidSaveTextDocumentParams,
+    NotebookCell,
+    NotebookDocument,
+    DocumentUri,
+    NotificationHandler
 } from 'vscode-languageserver';
 import { TextDocumentSyncKind, Disposable, Emitter } from 'vscode-languageserver';
 import type { URI } from '../utils/uri-utils.js';
 import { UriUtils } from '../utils/uri-utils.js';
+// For some reason, this isn't exported by vscode-languageserver
+import type { NotebookDocumentChangeEvent } from 'vscode-languageserver/lib/common/notebook.js';
+
+export type TextDocumentConnection = Pick<Connection,
+'onDidOpenTextDocument' |
+'onDidChangeTextDocument' |
+'onDidCloseTextDocument' |
+'onWillSaveTextDocument' |
+'onWillSaveTextDocumentWaitUntil' |
+'onDidSaveTextDocument'
+>;
 
 /**
  * A manager service that keeps track of all currently opened text documents.
@@ -91,7 +106,7 @@ export interface TextDocuments<T extends { uri: string }> {
      *
      * @param connection The connection to listen on.
      */
-    listen(connection: Connection): Disposable;
+    listen(connection: TextDocumentConnection): Disposable;
 }
 
 // Adapted from:
@@ -253,3 +268,279 @@ export class NormalizedTextDocuments<T extends { uri: string }> implements TextD
         return Disposable.create(() => { disposables.forEach(disposable => disposable.dispose()); });
     }
 }
+
+export interface NotebookDocuments<T extends { uri: string }> {
+    get cellTextDocuments(): TextDocuments<T>;
+    getCellTextDocument(cell: NotebookCell): T | undefined;
+    getNotebookDocument(uri: string | URI): NotebookDocument | undefined;
+    getNotebookCell(uri: string | URI): NotebookCell | undefined;
+    findNotebookDocumentForCell(cell: string | URI | NotebookCell): NotebookDocument | undefined;
+    get onDidOpen(): Event<NotebookDocument>;
+    get onDidSave(): Event<NotebookDocument>;
+    get onDidChange(): Event<NotebookDocumentChangeEvent>;
+    get onDidClose(): Event<NotebookDocument>;
+    /**
+     * Listens for `low level` notification on the given connection to
+     * update the notebook documents managed by this instance.
+     *
+     * Please note that the connection only provides handlers not an event model. Therefore
+     * listening on a connection will overwrite the following handlers on a connection:
+     * `onDidOpenNotebookDocument`, `onDidChangeNotebookDocument`, `onDidSaveNotebookDocument`,
+     *  and `onDidCloseNotebookDocument`.
+     *
+     * @param connection The connection to listen on.
+     */
+    listen(connection: Connection): Disposable;
+}
+
+export class NormalizedNotebookDocuments<T extends { uri: DocumentUri }> implements NotebookDocuments<T> {
+
+    private readonly notebookDocuments = new Map<DocumentUri, NotebookDocument>();
+    private readonly notebookCellMap = new Map<DocumentUri, [NotebookCell, NotebookDocument]>();
+
+    private readonly _onDidOpen = new Emitter<NotebookDocument>();
+    private readonly _onDidSave = new Emitter<NotebookDocument>();
+    private readonly _onDidChange = new Emitter<NotebookDocumentChangeEvent>();
+    private readonly _onDidClose = new Emitter<NotebookDocument>();
+
+    private readonly _cellTextDocuments: TextDocuments<T>;
+
+    constructor(configurationOrTextDocuments: TextDocumentsConfiguration<T> | TextDocuments<T>) {
+        if ('listen' in configurationOrTextDocuments) {
+            this._cellTextDocuments = configurationOrTextDocuments;
+        } else {
+            this._cellTextDocuments = new NormalizedTextDocuments<T>(configurationOrTextDocuments);
+        }
+    }
+
+    get cellTextDocuments(): TextDocuments<T> {
+        return this._cellTextDocuments;
+    }
+
+    getCellTextDocument(cell: NotebookCell): T | undefined {
+        return this._cellTextDocuments.get(cell.document);
+    }
+
+    getNotebookDocument(uri: string | URI): NotebookDocument | undefined {
+        return this.notebookDocuments.get(UriUtils.normalize(uri));
+    }
+
+    getNotebookCell(uri: DocumentUri): NotebookCell | undefined {
+        const value = this.notebookCellMap.get(uri);
+        return value && value[0];
+    }
+
+    findNotebookDocumentForCell(cell: string | URI | NotebookCell): NotebookDocument | undefined {
+        const key = typeof cell === 'string' || 'scheme' in cell ? cell : cell.document;
+        const value = this.notebookCellMap.get(UriUtils.normalize(key));
+        return value && value[1];
+    }
+
+    get onDidOpen(): Event<NotebookDocument> {
+        return this._onDidOpen.event;
+    }
+
+    get onDidSave(): Event<NotebookDocument> {
+        return this._onDidSave.event;
+    }
+
+    get onDidChange(): Event<NotebookDocumentChangeEvent> {
+        return this._onDidChange.event;
+    }
+
+    get onDidClose(): Event<NotebookDocument> {
+        return this._onDidClose.event;
+    }
+
+    /**
+     * Listens for `low level` notification on the given connection to
+     * update the notebook documents managed by this instance.
+     *
+     * Please note that the connection only provides handlers not an event model. Therefore
+     * listening on a connection will overwrite the following handlers on a connection:
+     * `onDidOpenNotebookDocument`, `onDidChangeNotebookDocument`, `onDidSaveNotebookDocument`,
+     *  and `onDidCloseNotebookDocument`.
+     *
+     * @param connection The connection to listen on.
+     */
+    listen(connection: Connection): Disposable {
+        const cellTextDocumentConnection = new CellTextDocumentConnection();
+        const disposables: Disposable[] = [];
+
+        disposables.push(this.cellTextDocuments.listen(cellTextDocumentConnection));
+        disposables.push(connection.notebooks.synchronization.onDidOpenNotebookDocument((params) => {
+            const uri = UriUtils.normalize(params.notebookDocument.uri);
+            this.notebookDocuments.set(uri, params.notebookDocument);
+            for (const cellTextDocument of params.cellTextDocuments) {
+                cellTextDocumentConnection.openTextDocument({ textDocument: cellTextDocument });
+            }
+            this.updateCellMap(params.notebookDocument);
+            this._onDidOpen.fire(params.notebookDocument);
+        }));
+        disposables.push(connection.notebooks.synchronization.onDidChangeNotebookDocument((params) => {
+            const uri = UriUtils.normalize(params.notebookDocument.uri);
+            const notebookDocument = this.notebookDocuments.get(uri);
+            if (notebookDocument === undefined) {
+                return;
+            }
+            notebookDocument.version = params.notebookDocument.version;
+            const oldMetadata = notebookDocument.metadata;
+            let metadataChanged: boolean = false;
+            const change = params.change;
+            if (change.metadata !== undefined) {
+                metadataChanged = true;
+                notebookDocument.metadata = change.metadata;
+            }
+
+            const opened: DocumentUri[] = [];
+            const closed: DocumentUri[] = [];
+            const data: Required<Required<Required<NotebookDocumentChangeEvent>['cells']>['changed']>['data'] = [];
+            const text: DocumentUri[] = [];
+            if (change.cells !== undefined) {
+                const changedCells = change.cells;
+                if (changedCells.structure !== undefined) {
+                    const array = changedCells.structure.array;
+                    notebookDocument.cells.splice(array.start, array.deleteCount, ...(array.cells !== undefined ? array.cells : []));
+                    // Additional open cell text documents.
+                    if (changedCells.structure.didOpen !== undefined) {
+                        for (const open of changedCells.structure.didOpen) {
+                            cellTextDocumentConnection.openTextDocument({ textDocument: open });
+                            opened.push(open.uri);
+                        }
+                    }
+                    // Additional closed cell test documents.
+                    if (changedCells.structure.didClose) {
+                        for (const close of changedCells.structure.didClose) {
+                            cellTextDocumentConnection.closeTextDocument({ textDocument: close });
+                            closed.push(close.uri);
+                        }
+                    }
+                }
+                if (changedCells.data !== undefined) {
+                    const cellUpdates: Map<string, NotebookCell> = new Map(changedCells.data.map(cell => [cell.document, cell]));
+                    for (let i = 0; i <= notebookDocument.cells.length; i++) {
+                        const change = cellUpdates.get(notebookDocument.cells[i].document);
+                        if (change !== undefined) {
+                            const old = notebookDocument.cells.splice(i, 1, change);
+                            data.push({ old: old[0], new: change });
+                            cellUpdates.delete(change.document);
+                            if (cellUpdates.size === 0) {
+                                break;
+                            }
+                        }
+                    }
+                }
+                if (changedCells.textContent !== undefined) {
+                    for (const cellTextDocument of changedCells.textContent) {
+                        cellTextDocumentConnection.changeTextDocument({ textDocument: cellTextDocument.document, contentChanges: cellTextDocument.changes });
+                        text.push(cellTextDocument.document.uri);
+                    }
+                }
+            }
+
+            // Update internal data structure.
+            this.updateCellMap(notebookDocument);
+
+            const changeEvent: NotebookDocumentChangeEvent = { notebookDocument };
+            if (metadataChanged) {
+                changeEvent.metadata = { old: oldMetadata, new: notebookDocument.metadata };
+            }
+
+            const added: NotebookCell[] = [];
+            for (const open of opened) {
+                added.push(this.getNotebookCell(open)!);
+            }
+            const removed: NotebookCell[] = [];
+            for (const close of closed) {
+                removed.push(this.getNotebookCell(close)!);
+            }
+            const textContent: NotebookCell[] = [];
+            for (const change of text) {
+                textContent.push(this.getNotebookCell(change)!);
+            }
+            if (added.length > 0 || removed.length > 0 || data.length > 0 || textContent.length > 0) {
+                changeEvent.cells = { added, removed, changed: { data, textContent } };
+            }
+            if (changeEvent.metadata !== undefined || changeEvent.cells !== undefined) {
+                this._onDidChange.fire(changeEvent);
+            }
+        }));
+        disposables.push(connection.notebooks.synchronization.onDidSaveNotebookDocument((params) => {
+            const notebookDocument = this.getNotebookDocument(params.notebookDocument.uri);
+            if (notebookDocument === undefined) {
+                return;
+            }
+            this._onDidSave.fire(notebookDocument);
+        }));
+        disposables.push(connection.notebooks.synchronization.onDidCloseNotebookDocument((params) => {
+            const uri = UriUtils.normalize(params.notebookDocument.uri);
+            const notebookDocument = this.notebookDocuments.get(uri);
+            if (notebookDocument === undefined) {
+                return;
+            }
+            this._onDidClose.fire(notebookDocument);
+            for (const cellTextDocument of params.cellTextDocuments) {
+                cellTextDocumentConnection.closeTextDocument({ textDocument: cellTextDocument });
+            }
+            this.notebookDocuments.delete(uri);
+            for (const cell of notebookDocument.cells) {
+                this.notebookCellMap.delete(cell.document);
+            }
+        }));
+        return Disposable.create(() => { disposables.forEach(disposable => disposable.dispose()); });
+    }
+
+    private updateCellMap(notebookDocument: NotebookDocument): void {
+        for (const cell of notebookDocument.cells) {
+            this.notebookCellMap.set(cell.document, [cell, notebookDocument]);
+        }
+    }
+}
+
+class CellTextDocumentConnection implements TextDocumentConnection {
+
+    private static readonly NULL_DISPOSE = Object.freeze({ dispose: () => { } });
+
+    private openHandler: NotificationHandler<DidOpenTextDocumentParams> | undefined;
+    private changeHandler: NotificationHandler<DidChangeTextDocumentParams> | undefined;
+    private closeHandler: NotificationHandler<DidCloseTextDocumentParams> | undefined;
+
+    onDidOpenTextDocument(handler: NotificationHandler<DidOpenTextDocumentParams>): Disposable {
+        this.openHandler = handler;
+        return Disposable.create(() => { this.openHandler = undefined; });
+    }
+
+    openTextDocument(params: DidOpenTextDocumentParams): void {
+        this.openHandler && this.openHandler(params);
+    }
+
+    onDidChangeTextDocument(handler: NotificationHandler<DidChangeTextDocumentParams>): Disposable {
+        this.changeHandler = handler;
+        return Disposable.create(() => { this.changeHandler = handler; });
+    }
+
+    changeTextDocument(params: DidChangeTextDocumentParams): void {
+        this.changeHandler && this.changeHandler(params);
+    }
+
+    onDidCloseTextDocument(handler: NotificationHandler<DidCloseTextDocumentParams>): Disposable {
+        this.closeHandler = handler;
+        return Disposable.create(() => { this.closeHandler = undefined; });
+    }
+
+    closeTextDocument(params: DidCloseTextDocumentParams): void {
+        this.closeHandler && this.closeHandler(params);
+    }
+
+    onWillSaveTextDocument(): Disposable {
+        return CellTextDocumentConnection.NULL_DISPOSE;
+    }
+
+    onWillSaveTextDocumentWaitUntil(): Disposable {
+        return CellTextDocumentConnection.NULL_DISPOSE;
+    }
+
+    onDidSaveTextDocument(): Disposable {
+        return CellTextDocumentConnection.NULL_DISPOSE;
+    }
+}