Added a map of open documents to the languageServerBase class. It is …

…responsible for applying any delta changes and updating the contents for each of the workspaces associated with that file. This allows source code contents to be shared between workspaces, which reduces memory usage. It also eliminates a correctness issue that can occur if new workspaces are added after files are open. It addresses #4994.

packages/pyright-internal/src/analyzer/backgroundAnalysisProgram.ts

@@ -8,7 +8,6 @@
  */
 
 import { CancellationToken } from 'vscode-languageserver';
-import { TextDocumentContentChangeEvent } from 'vscode-languageserver-textdocument';
 
 import { BackgroundAnalysisBase, IndexOptions, RefreshOptions } from '../backgroundAnalysisBase';
 import { ConfigOptions, ExecutionEnvironment } from '../common/configOptions';
@@ -99,8 +98,8 @@ export class BackgroundAnalysisProgram {
     }
 
     setFileOpened(filePath: string, version: number | null, contents: string, options: OpenFileOptions) {
-        this._backgroundAnalysis?.setFileOpened(filePath, version, [{ text: contents }], options);
-        this._program.setFileOpened(filePath, version, [{ text: contents }], options);
+        this._backgroundAnalysis?.setFileOpened(filePath, version, contents, options);
+        this._program.setFileOpened(filePath, version, contents, options);
     }
 
     getChainedFilePath(filePath: string): string | undefined {
@@ -112,12 +111,7 @@ export class BackgroundAnalysisProgram {
         this._program.updateChainedFilePath(filePath, chainedFilePath);
     }
 
-    updateOpenFileContents(
-        path: string,
-        version: number | null,
-        contents: TextDocumentContentChangeEvent[],
-        options: OpenFileOptions
-    ) {
+    updateOpenFileContents(path: string, version: number | null, contents: string, options: OpenFileOptions) {
         this._backgroundAnalysis?.setFileOpened(path, version, contents, options);
         this._program.setFileOpened(path, version, contents, options);
         this.markFilesDirty([path], /* evenIfContentsAreSame */ true);

packages/pyright-internal/src/analyzer/program.ts

@@ -9,7 +9,6 @@
  */
 
 import { CancellationToken, CompletionItem, DocumentSymbol } from 'vscode-languageserver';
-import { TextDocumentContentChangeEvent } from 'vscode-languageserver-textdocument';
 import { CompletionList } from 'vscode-languageserver-types';
 
 import { Commands } from '../commands/commands';
@@ -352,12 +351,7 @@ export class Program {
         return sourceFile;
     }
 
-    setFileOpened(
-        filePath: string,
-        version: number | null,
-        contents: TextDocumentContentChangeEvent[],
-        options?: OpenFileOptions
-    ) {
+    setFileOpened(filePath: string, version: number | null, contents: string, options?: OpenFileOptions) {
         let sourceFileInfo = this.getSourceFileInfo(filePath);
         if (!sourceFileInfo) {
             const importName = this._getImportNameForFile(filePath);
@@ -422,7 +416,7 @@ export class Program {
         if (sourceFileInfo) {
             sourceFileInfo.isOpenByClient = false;
             sourceFileInfo.isTracked = isTracked ?? sourceFileInfo.isTracked;
-            sourceFileInfo.sourceFile.setClientVersion(null, []);
+            sourceFileInfo.sourceFile.setClientVersion(null, '');
 
             // There is no guarantee that content is saved before the file is closed.
             // We need to mark the file dirty so we can re-analyze next time.
@@ -1632,7 +1626,7 @@ export class Program {
             const isTracked = info ? info.isTracked : true;
             const realFilePath = info ? info.sourceFile.getRealFilePath() : filePath;
 
-            cloned.setFileOpened(filePath, version, [{ text }], {
+            cloned.setFileOpened(filePath, version, text, {
                 chainedFilePath,
                 ipythonMode,
                 isTracked,
@@ -1718,7 +1712,7 @@ export class Program {
             program.setFileOpened(
                 fileInfo.sourceFile.getFilePath(),
                 version,
-                [{ text: fileInfo.sourceFile.getOpenFileContents()! }],
+                fileInfo.sourceFile.getOpenFileContents() ?? '',
                 {
                     chainedFilePath: fileInfo.chainedSourceFile?.sourceFile.getFilePath(),
                     ipythonMode: fileInfo.sourceFile.getIPythonMode(),

packages/pyright-internal/src/analyzer/service.ts

@@ -16,7 +16,6 @@ import {
     CompletionItem,
     DocumentSymbol,
 } from 'vscode-languageserver';
-import { TextDocumentContentChangeEvent } from 'vscode-languageserver-textdocument';
 
 import { BackgroundAnalysisBase, IndexOptions, RefreshOptions } from '../backgroundAnalysisBase';
 import { CancellationProvider, DefaultCancellationProvider } from '../common/cancellationUtils';
@@ -318,7 +317,7 @@ export class AnalyzerService {
     updateOpenFileContents(
         path: string,
         version: number | null,
-        contents: TextDocumentContentChangeEvent[],
+        contents: string,
         ipythonMode = IPythonMode.None,
         realFilePath?: string
     ) {

packages/pyright-internal/src/analyzer/sourceFile.ts

@@ -8,7 +8,6 @@
  */
 
 import { CancellationToken, CompletionItem, DocumentSymbol } from 'vscode-languageserver';
-import { TextDocument, TextDocumentContentChangeEvent } from 'vscode-languageserver-textdocument';
 import { isMainThread } from 'worker_threads';
 
 import * as SymbolNameUtils from '../analyzer/symbolNameUtils';
@@ -136,7 +135,8 @@ export class SourceFile {
 
     // Client's version of the file. Undefined implies that contents
     // need to be read from disk.
-    private _clientDocument: TextDocument | undefined;
+    private _clientDocumentContents: string | undefined;
+    private _clientDocumentVersion: number | undefined;
 
     // Version of file contents that have been analyzed.
     private _analyzedFileContentsVersion = -1;
@@ -561,7 +561,7 @@ export class SourceFile {
         // If this is an open file any content changes will be
         // provided through the editor. We can assume contents
         // didn't change without us knowing about them.
-        if (this._clientDocument) {
+        if (this._clientDocumentContents) {
             return false;
         }
 
@@ -635,11 +635,11 @@ export class SourceFile {
     }
 
     getClientVersion() {
-        return this._clientDocument?.version;
+        return this._clientDocumentVersion;
     }
 
     getOpenFileContents() {
-        return this._clientDocument?.getText();
+        return this._clientDocumentContents;
     }
 
     getFileContent(): string | undefined {
@@ -667,24 +667,22 @@ export class SourceFile {
         }
     }
 
-    setClientVersion(version: number | null, contents: TextDocumentContentChangeEvent[]): void {
+    setClientVersion(version: number | null, contents: string): void {
         if (version === null) {
-            this._clientDocument = undefined;
+            this._clientDocumentVersion = undefined;
+            this._clientDocumentContents = undefined;
         } else {
-            if (!this._clientDocument) {
-                this._clientDocument = TextDocument.create(this._filePath, 'python', version, '');
-            }
-            this._clientDocument = TextDocument.update(this._clientDocument, contents, version);
+            this._clientDocumentVersion = version;
+            this._clientDocumentContents = contents;
 
-            const fileContents = this._clientDocument.getText();
-            const contentsHash = StringUtils.hashString(fileContents);
+            const contentsHash = StringUtils.hashString(contents);
 
             // Have the contents of the file changed?
-            if (fileContents.length !== this._lastFileContentLength || contentsHash !== this._lastFileContentHash) {
+            if (contents.length !== this._lastFileContentLength || contentsHash !== this._lastFileContentHash) {
                 this.markDirty();
             }
 
-            this._lastFileContentLength = fileContents.length;
+            this._lastFileContentLength = contents.length;
             this._lastFileContentHash = contentsHash;
             this._isFileDeleted = false;
         }

packages/pyright-internal/src/backgroundAnalysisBase.ts

@@ -7,7 +7,6 @@
  */
 
 import { CancellationToken } from 'vscode-languageserver';
-import { TextDocumentContentChangeEvent } from 'vscode-languageserver-textdocument';
 import { MessageChannel, MessagePort, parentPort, threadId, Worker, workerData } from 'worker_threads';
 
 import { AnalysisCompleteCallback, AnalysisResults, analyzeProgram, nullCallback } from './analyzer/analysis';
@@ -71,12 +70,7 @@ export class BackgroundAnalysisBase {
         this.enqueueRequest({ requestType: 'ensurePartialStubPackages', data: { executionRoot } });
     }
 
-    setFileOpened(
-        filePath: string,
-        version: number | null,
-        contents: TextDocumentContentChangeEvent[],
-        options: OpenFileOptions
-    ) {
+    setFileOpened(filePath: string, version: number | null, contents: string, options: OpenFileOptions) {
         this.enqueueRequest({
             requestType: 'setFileOpened',
             data: { filePath, version, contents, options },

packages/pyright-internal/src/common/workspaceEditUtils.ts

@@ -16,6 +16,7 @@ import {
     WorkspaceEdit,
 } from 'vscode-languageserver';
 
+import { TextDocument } from 'vscode-languageserver-textdocument';
 import { SourceFileInfo } from '../analyzer/program';
 import { AnalyzerService } from '../analyzer/service';
 import { FileEditAction, FileEditActions, TextEditAction } from '../common/editAction';
@@ -148,11 +149,14 @@ export function applyDocumentChanges(clonedService: AnalyzerService, fileInfo: S
         );
     }
 
-    const version = (fileInfo.sourceFile.getClientVersion() ?? 0) + 1;
+    const version = fileInfo.sourceFile.getClientVersion() ?? 0;
+    const filePath = fileInfo.sourceFile.getFilePath();
+    const sourceDoc = TextDocument.create(filePath, 'python', version, fileInfo.sourceFile.getOpenFileContents() ?? '');
+
     clonedService.updateOpenFileContents(
-        fileInfo.sourceFile.getFilePath(),
-        version,
-        edits.map((t) => ({ range: t.range, text: t.newText })),
+        filePath,
+        version + 1,
+        TextDocument.applyEdits(sourceDoc, edits),
         fileInfo.sourceFile.getIPythonMode(),
         fileInfo.sourceFile.getRealFilePath()
     );
@@ -216,7 +220,7 @@ function _convertToWorkspaceEditWithDocumentChanges(
     };
 
     // Ordering of documentChanges are important.
-    // Make sure create operaiton happens before edits
+    // Make sure create operation happens before edits.
     for (const operation of editActions.fileOperations) {
         switch (operation.kind) {
             case 'create':

packages/pyright-internal/src/languageServerBase.ts

@@ -72,6 +72,7 @@ import {
 } from 'vscode-languageserver';
 import { ResultProgressReporter, attachWorkDone } from 'vscode-languageserver/lib/common/progress';
 
+import { TextDocument } from 'vscode-languageserver-textdocument';
 import { AnalysisResults } from './analyzer/analysis';
 import { BackgroundAnalysisProgram } from './analyzer/backgroundAnalysisProgram';
 import { CacheManager } from './analyzer/cacheManager';
@@ -336,6 +337,7 @@ export abstract class LanguageServerBase implements LanguageServerInterface {
 
     protected defaultClientConfig: any;
     protected workspaceFactory: WorkspaceFactory;
+    protected openFileMap = new Map<string, TextDocument>();
     protected cacheManager: CacheManager;
     protected fs: PyrightFileSystem;
     protected uriParser: UriParser;
@@ -1175,6 +1177,16 @@ export abstract class LanguageServerBase implements LanguageServerInterface {
             return;
         }
 
+        let doc = this.openFileMap.get(filePath);
+        if (doc) {
+            // We shouldn't get an open text document request for an already-opened doc.
+            this.console.error(`Received redundant open text document command for ${filePath}`);
+            doc = TextDocument.update(doc, [{ text: params.textDocument.text }], params.textDocument.version);
+        } else {
+            doc = TextDocument.create(filePath, 'python', params.textDocument.version, params.textDocument.text);
+        }
+        this.openFileMap.set(filePath, doc);
+
         // Send this open to all the workspaces that might contain this file.
         const workspaces = await this.getContainingWorkspacesForFile(filePath);
         workspaces.forEach((w) => {
@@ -1191,10 +1203,21 @@ export abstract class LanguageServerBase implements LanguageServerInterface {
             return;
         }
 
+        let doc = this.openFileMap.get(filePath);
+        if (!doc) {
+            // We shouldn't get a change text request for a closed doc.
+            this.console.error(`Received change text document command for closed file ${filePath}`);
+            return;
+        }
+
+        doc = TextDocument.update(doc, params.contentChanges, params.textDocument.version);
+        this.openFileMap.set(filePath, doc);
+        const newContents = doc.getText();
+
         // Send this change to all the workspaces that might contain this file.
         const workspaces = await this.getContainingWorkspacesForFile(filePath);
         workspaces.forEach((w) => {
-            w.service.updateOpenFileContents(filePath, params.textDocument.version, params.contentChanges, ipythonMode);
+            w.service.updateOpenFileContents(filePath, params.textDocument.version, newContents, ipythonMode);
         });
     }
 
@@ -1210,6 +1233,8 @@ export abstract class LanguageServerBase implements LanguageServerInterface {
         workspaces.forEach((w) => {
             w.service.setFileClosed(filePath);
         });
+
+        this.openFileMap.delete(filePath);
     }
 
     protected onDidChangeWatchedFiles(params: DidChangeWatchedFilesParams) {

packages/pyright-internal/src/tests/chainedSourceFiles.test.ts

@@ -152,7 +152,7 @@ test('modify chained files', async () => {
     assert.strictEqual(initialDiags.length, 0);
 
     // Change test1 content
-    service.updateOpenFileContents(data.markerPositions.get('changed')!.fileName, 2, [{ text: 'def foo5(): pass' }]);
+    service.updateOpenFileContents(data.markerPositions.get('changed')!.fileName, 2, 'def foo5(): pass');
     analyze(service.test_program);
 
     const finalDiags = await service.getDiagnosticsForRange(

packages/pyright-internal/src/tests/harness/fourslash/testState.ts

@@ -426,7 +426,7 @@ export class TestState {
         fileToOpen.fileName = normalizeSlashes(fileToOpen.fileName);
         this.activeFile = fileToOpen;
 
-        this.program.setFileOpened(this.activeFile.fileName, 1, [{ text: fileToOpen.content }]);
+        this.program.setFileOpened(this.activeFile.fileName, 1, fileToOpen.content);
     }
 
     openFiles(indexOrNames: (number | string)[]): void {
@@ -1358,7 +1358,7 @@ export class TestState {
             if (!this.program.getSourceFileInfo(fileName)) {
                 const file = this.testData.files.find((v) => v.fileName === fileName);
                 if (file) {
-                    this.program.setFileOpened(fileName, file.version, [{ text: file.content }]);
+                    this.program.setFileOpened(fileName, file.version, file.content);
                 }
             }
 

packages/pyright-internal/src/tests/sourceFile.test.ts

@@ -40,6 +40,6 @@ test('Empty Open file', () => {
         '# Content'
     );
 
-    state.workspace.service.updateOpenFileContents(marker.fileName, 1, [{ text: '' }]);
+    state.workspace.service.updateOpenFileContents(marker.fileName, 1, '');
     assert.strictEqual(state.workspace.service.test_program.getSourceFile(marker.fileName)?.getFileContent(), '');
 });

packages/pyright-internal/src/tests/testStateUtils.ts

@@ -77,7 +77,7 @@ export function applyFileOperations(state: TestState, fileEditActions: FileEditA
                 state.program.setFileClosed(renamed.oldFilePath);
             }
 
-            state.program.setFileOpened(openedFilePath, result.version + 1, [{ text: result.text }]);
+            state.program.setFileOpened(openedFilePath, result.version + 1, result.text);
         }
     }