Idea for copy on write

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

@@ -8,6 +8,7 @@
  */
 
 import { CancellationToken } from 'vscode-languageserver';
+import { TextDocumentContentChangeEvent } from 'vscode-languageserver-textdocument';
 
 import { BackgroundAnalysisBase, IndexOptions, RefreshOptions } from '../backgroundAnalysisBase';
 import { ConfigOptions, ExecutionEnvironment } from '../common/configOptions';
@@ -24,6 +25,7 @@ export class BackgroundAnalysisProgram {
     private _program: Program;
     private _disposed = false;
     private _onAnalysisCompletion: AnalysisCompleteCallback | undefined;
+    private _preEditAnalysis: BackgroundAnalysisBase | undefined;
 
     constructor(
         private _console: ConsoleInterface,
@@ -98,8 +100,8 @@ export class BackgroundAnalysisProgram {
     }
 
     setFileOpened(filePath: string, version: number | null, contents: string, options: OpenFileOptions) {
-        this._backgroundAnalysis?.setFileOpened(filePath, version, contents, options);
-        this._program.setFileOpened(filePath, version, contents, options);
+        this._backgroundAnalysis?.setFileOpened(filePath, version, [{ text: contents }], options);
+        this._program.setFileOpened(filePath, version, [{ text: contents }], options);
     }
 
     getChainedFilePath(filePath: string): string | undefined {
@@ -111,7 +113,12 @@ export class BackgroundAnalysisProgram {
         this._program.updateChainedFilePath(filePath, chainedFilePath);
     }
 
-    updateOpenFileContents(path: string, version: number | null, contents: string, options: OpenFileOptions) {
+    updateOpenFileContents(
+        path: string,
+        version: number | null,
+        contents: TextDocumentContentChangeEvent[],
+        options: OpenFileOptions
+    ) {
         this._backgroundAnalysis?.setFileOpened(path, version, contents, options);
         this._program.setFileOpened(path, version, contents, options);
         this.markFilesDirty([path], /* evenIfContentsAreSame */ true);
@@ -242,6 +249,20 @@ export class BackgroundAnalysisProgram {
         this._backgroundAnalysis?.shutdown();
     }
 
+    enterEditMode() {
+        // Turn off analysis while in edit mode.
+        this._preEditAnalysis = this._backgroundAnalysis;
+        this._backgroundAnalysis = undefined;
+
+        // Forward this request to the program.
+        this._program.enterEditMode();
+    }
+
+    leaveEditMode() {
+        this._backgroundAnalysis = this._preEditAnalysis;
+        this._preEditAnalysis = undefined;
+        return this._program.leaveEditMode();
+    }
     protected getIndices(): Indices | undefined {
         return undefined;
     }

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

@@ -9,6 +9,7 @@
  */
 
 import { CancellationToken, CompletionItem, DocumentSymbol } from 'vscode-languageserver';
+import { TextDocumentContentChangeEvent } from 'vscode-languageserver-textdocument';
 import { CompletionList } from 'vscode-languageserver-types';
 
 import { Commands } from '../commands/commands';
@@ -38,17 +39,17 @@ import {
 import { convertPositionToOffset, convertRangeToTextRange, convertTextRangeToRange } from '../common/positionUtils';
 import { computeCompletionSimilarity } from '../common/stringUtils';
 import { TextEditTracker } from '../common/textEditTracker';
-import { doRangesIntersect, getEmptyRange, Position, Range, TextRange } from '../common/textRange';
+import { Position, Range, TextRange, doRangesIntersect, getEmptyRange } from '../common/textRange';
 import { TextRangeCollection } from '../common/textRangeCollection';
 import { Duration, timingStats } from '../common/timing';
 import { applyTextEditsToString } from '../common/workspaceEditUtils';
 import {
-    AutoImporter,
     AutoImportOptions,
     AutoImportResult,
-    buildModuleSymbolsMap,
+    AutoImporter,
     ImportFormat,
     ModuleSymbolMap,
+    buildModuleSymbolsMap,
 } from '../languageService/autoImporter';
 import {
     AbbreviationMap,
@@ -84,15 +85,16 @@ import { Scope } from './scope';
 import { getScopeForNode } from './scopeUtils';
 import { IPythonMode, SourceFile } from './sourceFile';
 import { collectImportedByFiles, isUserCode } from './sourceFileInfoUtils';
-import { isStubFile, SourceMapper } from './sourceMapper';
+import { SourceMapper, isStubFile } from './sourceMapper';
 import { Symbol } from './symbol';
 import { isPrivateOrProtectedName } from './symbolNameUtils';
+import { createFileEditActions } from './textEditUtils';
 import { createTracePrinter } from './tracePrinter';
 import { PrintTypeOptions, TypeEvaluator } from './typeEvaluatorTypes';
 import { createTypeEvaluatorWithTracker } from './typeEvaluatorWithTracker';
 import { PrintTypeFlags } from './typePrinter';
-import { Type } from './types';
 import { TypeStubWriter } from './typeStubWriter';
+import { Type } from './types';
 
 const _maxImportDepth = 256;
 
@@ -187,6 +189,9 @@ export class Program {
     private _cacheManager: CacheManager;
     private _id: number;
     private static _nextId = 0;
+    private _editMode = false;
+    private _copiedOnWriteFiles = new Map<string, SourceFileInfo | undefined>();
+    private _copiedOnWriteChanges = new Map<string, FileEditAction[]>();
 
     constructor(
         initialImportResolver: ImportResolver,
@@ -240,6 +245,37 @@ export class Program {
         this._cacheManager.unregisterCacheOwner(this);
     }
 
+    enterEditMode() {
+        this._editMode = true;
+    }
+
+    leaveEditMode() {
+        this._editMode = false;
+        const edits: FileEditAction[] = [];
+        this._copiedOnWriteFiles.forEach((fileInfo, path) => {
+            const changes = this._copiedOnWriteChanges.get(path);
+            if (changes) {
+                edits.push(...changes);
+            }
+            const changedFile = this._sourceFileMap.get(path);
+            const index = changedFile ? this._sourceFileList.indexOf(changedFile) : -1;
+            if (fileInfo) {
+                this._sourceFileMap.set(path, fileInfo);
+                if (index >= 0) {
+                    this._sourceFileList[index] = fileInfo;
+                }
+            } else {
+                this._sourceFileMap.delete(path);
+                if (index >= 0) {
+                    this._sourceFileList.splice(index, 1);
+                }
+            }
+        });
+        this._copiedOnWriteChanges.clear();
+        this._copiedOnWriteFiles.clear();
+        return edits;
+    }
+
     setConfigOptions(configOptions: ConfigOptions) {
         this._configOptions = configOptions;
         this._importResolver.setConfigOptions(configOptions);
@@ -351,8 +387,50 @@ export class Program {
         return sourceFile;
     }
 
-    setFileOpened(filePath: string, version: number | null, contents: string, options?: OpenFileOptions) {
+    setFileOpened(
+        filePath: string,
+        version: number | null,
+        contents: TextDocumentContentChangeEvent[],
+        options?: OpenFileOptions
+    ) {
+        let chainedFilePath: string | undefined = options?.chainedFilePath;
+        const filePathKey = normalizePathCase(this.fileSystem, filePath);
         let sourceFileInfo = this.getSourceFileInfo(filePath);
+
+        // If in edit mode, we need to start over if this is the first edit.
+        if (this._editMode && !this._copiedOnWriteFiles.has(filePathKey)) {
+            // Save the old state of the file.
+            this._copiedOnWriteFiles.set(filePathKey, sourceFileInfo);
+
+            // Only save the changes if these are actually changes to a file. Otherwise
+            // initial creation of a file shouldn't count as a change.
+            if (sourceFileInfo) {
+                this._copiedOnWriteChanges.set(filePathKey, createFileEditActions(filePath, contents));
+            }
+            this._sourceFileMap.delete(filePathKey);
+            this._sourceFileList = this._sourceFileList.filter((f) => f !== sourceFileInfo);
+
+            // Chained file path needs to be maintained.
+            chainedFilePath = sourceFileInfo?.chainedSourceFile?.sourceFile.getFilePath();
+
+            // Add a new change that is the full text. This prefills the new source file.
+            contents = [
+                {
+                    text: sourceFileInfo?.sourceFile.getFileContent() || '',
+                },
+                ...contents,
+            ];
+
+            // Mark as not existing so we start over.
+            sourceFileInfo = undefined;
+        } else if (this._editMode && this._copiedOnWriteFiles.has(filePathKey)) {
+            const currentChanges = this._copiedOnWriteChanges.get(filePathKey) || [];
+            // File has already had changes applied. Apply more
+            this._copiedOnWriteChanges.set(filePathKey, [
+                ...currentChanges,
+                ...createFileEditActions(filePath, contents),
+            ]);
+        }
         if (!sourceFileInfo) {
             const importName = this._getImportNameForFile(filePath);
             const sourceFile = new SourceFile(
@@ -367,7 +445,6 @@ export class Program {
                 options?.ipythonMode ?? IPythonMode.None
             );
 
-            const chainedFilePath = options?.chainedFilePath;
             sourceFileInfo = {
                 sourceFile,
                 isTracked: options?.isTracked ?? false,
@@ -416,7 +493,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.
@@ -1626,7 +1703,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: text }], {
                 chainedFilePath,
                 ipythonMode,
                 isTracked,
@@ -1712,7 +1789,7 @@ export class Program {
             program.setFileOpened(
                 fileInfo.sourceFile.getFilePath(),
                 version,
-                fileInfo.sourceFile.getOpenFileContents() ?? '',
+                [{ text: fileInfo.sourceFile.getOpenFileContents() ?? '' }],
                 {
                     chainedFilePath: fileInfo.chainedSourceFile?.sourceFile.getFilePath(),
                     ipythonMode: fileInfo.sourceFile.getIPythonMode(),

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

@@ -16,14 +16,15 @@ import {
     CompletionItem,
     DocumentSymbol,
 } from 'vscode-languageserver';
+import { TextDocumentContentChangeEvent } from 'vscode-languageserver-textdocument';
 
 import { BackgroundAnalysisBase, IndexOptions, RefreshOptions } from '../backgroundAnalysisBase';
 import { CancellationProvider, DefaultCancellationProvider } from '../common/cancellationUtils';
 import { CommandLineOptions } from '../common/commandLineOptions';
 import { ConfigOptions, matchFileSpecs } from '../common/configOptions';
 import { ConsoleInterface, LogLevel, StandardConsole, log } from '../common/console';
 import { Diagnostic } from '../common/diagnostic';
-import { FileEditActions } from '../common/editAction';
+import { FileEditAction, FileEditActions } from '../common/editAction';
 import { Extensions, ProgramView } from '../common/extensibility';
 import { FileSystem, FileWatcher, FileWatcherEventType, ignoredWatchEventFunction } from '../common/fileSystem';
 import { Host, HostFactory, NoAccessHost } from '../common/host';
@@ -222,6 +223,14 @@ export class AnalyzerService {
         return service;
     }
 
+    enterEditMode() {
+        this._backgroundAnalysisProgram.enterEditMode();
+    }
+
+    leaveEditMode(): FileEditAction[] {
+        return this._backgroundAnalysisProgram.leaveEditMode();
+    }
+
     dispose() {
         if (!this._disposed) {
             // Make sure we dispose program, otherwise, entire program
@@ -317,7 +326,7 @@ export class AnalyzerService {
     updateOpenFileContents(
         path: string,
         version: number | null,
-        contents: string,
+        contents: TextDocumentContentChangeEvent[],
         ipythonMode = IPythonMode.None,
         realFilePath?: string
     ) {

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

@@ -8,6 +8,7 @@
  */
 
 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';
@@ -135,8 +136,7 @@ export class SourceFile {
 
     // Client's version of the file. Undefined implies that contents
     // need to be read from disk.
-    private _clientDocumentContents: string | undefined;
-    private _clientDocumentVersion: number | undefined;
+    private _clientDocument: TextDocument | 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._clientDocumentContents) {
+        if (this._clientDocument) {
             return false;
         }
 
@@ -635,11 +635,11 @@ export class SourceFile {
     }
 
     getClientVersion() {
-        return this._clientDocumentVersion;
+        return this._clientDocument?.version;
     }
 
     getOpenFileContents() {
-        return this._clientDocumentContents;
+        return this._clientDocument?.getText();
     }
 
     getFileContent(): string | undefined {
@@ -667,22 +667,24 @@ export class SourceFile {
         }
     }
 
-    setClientVersion(version: number | null, contents: string): void {
+    setClientVersion(version: number | null, contents: TextDocumentContentChangeEvent[]): void {
         if (version === null) {
-            this._clientDocumentVersion = undefined;
-            this._clientDocumentContents = undefined;
+            this._clientDocument = undefined;
         } else {
-            this._clientDocumentVersion = version;
-            this._clientDocumentContents = contents;
+            if (!this._clientDocument) {
+                this._clientDocument = TextDocument.create(this._filePath, 'python', version, '');
+            }
+            this._clientDocument = TextDocument.update(this._clientDocument, contents, version);
+            const fileContents = this._clientDocument.getText();
 
-            const contentsHash = StringUtils.hashString(contents);
+            const contentsHash = StringUtils.hashString(fileContents);
 
             // Have the contents of the file changed?
-            if (contents.length !== this._lastFileContentLength || contentsHash !== this._lastFileContentHash) {
+            if (fileContents.length !== this._lastFileContentLength || contentsHash !== this._lastFileContentHash) {
                 this.markDirty();
             }
 
-            this._lastFileContentLength = contents.length;
+            this._lastFileContentLength = fileContents.length;
             this._lastFileContentHash = contentsHash;
             this._isFileDeleted = false;
         }

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

@@ -0,0 +1,24 @@
+import { Range, TextDocumentContentChangeEvent } from 'vscode-languageserver';
+import { FileEditAction } from '../common/editAction';
+
+export function createFileEditActions(filePath: string, edits: TextDocumentContentChangeEvent[]): FileEditAction[] {
+    return edits.map((edit) => {
+        const range = (edit as any).range as Range;
+        if (range) {
+            return {
+                range,
+                replacementText: edit.text,
+                filePath,
+            };
+        } else {
+            return {
+                range: {
+                    start: { line: 0, character: 0 },
+                    end: { line: 0, character: 0 },
+                },
+                replacementText: edit.text,
+                filePath,
+            };
+        }
+    });
+}