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();
+    }
+
+    exitEditMode() {
+        this._backgroundAnalysis = this._preEditAnalysis;
+        this._preEditAnalysis = undefined;
+        return this._program.exitEditMode();
+    }
     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,15 @@ 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 { 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 +188,7 @@ export class Program {
     private _cacheManager: CacheManager;
     private _id: number;
     private static _nextId = 0;
+    private _editMode = false;
 
     constructor(
         initialImportResolver: ImportResolver,
@@ -240,6 +242,36 @@ export class Program {
         this._cacheManager.unregisterCacheOwner(this);
     }
 
+    enterEditMode() {
+        // Keep track of edit mode so we can apply it to new source files.
+        this._editMode = true;
+
+        // Tell all source files we're in edit mode.
+        this._sourceFileList.forEach((sourceFile) => {
+            sourceFile.sourceFile.enterEditMode();
+        });
+    }
+
+    exitEditMode() {
+        // Tell all source files we're no longer in edit mode. Gather
+        // up all of their edits.
+        const edits: FileEditAction[] = [];
+        this._sourceFileList.forEach((sourceFile) => {
+            const sourceFileEdits = sourceFile.sourceFile.exitEditMode();
+            if (sourceFileEdits.length > 0) {
+                // This means this source file was modified. We need to recompute its imports after
+                // we put it back to how it was.
+                this._updateSourceFileImports(sourceFile, this.configOptions);
+            }
+            edits.push(...sourceFileEdits);
+        });
+
+        // Stop applying edit mode to new source files.
+        this._editMode = false;
+
+        return edits;
+    }
+
     setConfigOptions(configOptions: ConfigOptions) {
         this._configOptions = configOptions;
         this._importResolver.setConfigOptions(configOptions);
@@ -331,6 +363,7 @@ export class Program {
             importName,
             isThirdPartyImport,
             isInPyTypedPackage,
+            this._editMode,
             this._console,
             this._logTracker
         );
@@ -351,7 +384,12 @@ 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 sourceFileInfo = this.getSourceFileInfo(filePath);
         if (!sourceFileInfo) {
             const importName = this._getImportNameForFile(filePath);
@@ -361,12 +399,12 @@ export class Program {
                 importName,
                 /* isThirdPartyImport */ false,
                 /* isInPyTypedPackage */ false,
+                this._editMode,
                 this._console,
                 this._logTracker,
                 options?.realFilePath,
                 options?.ipythonMode ?? IPythonMode.None
             );
-
             const chainedFilePath = options?.chainedFilePath;
             sourceFileInfo = {
                 sourceFile,
@@ -416,7 +454,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 +1664,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 +1750,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(),
@@ -2501,6 +2539,7 @@ export class Program {
                         importName,
                         importInfo.isThirdPartyImport,
                         importInfo.isPyTypedPresent,
+                        this._editMode,
                         this._console,
                         this._logTracker
                     );
@@ -2644,6 +2683,7 @@ export class Program {
             importName,
             /* isThirdPartyImport */ false,
             /* isInPyTypedPackage */ false,
+            this._editMode,
             this._console,
             this._logTracker
         );

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,22 @@ export class AnalyzerService {
         return service;
     }
 
+    async useEditMode(callback: () => Promise<void>, token: CancellationToken) {
+        let edits: FileEditAction[] = [];
+        const disposable = token.onCancellationRequested(() => {
+            edits = [];
+            this._backgroundAnalysisProgram.exitEditMode();
+        });
+        this._backgroundAnalysisProgram.enterEditMode();
+        try {
+            await callback();
+        } finally {
+            disposable.dispose();
+            edits = this._backgroundAnalysisProgram.exitEditMode();
+        }
+        return edits;
+    }
+
     dispose() {
         if (!this._disposed) {
             // Make sure we dispose program, otherwise, entire program
@@ -317,7 +334,7 @@ export class AnalyzerService {
     updateOpenFileContents(
         path: string,
         version: number | null,
-        contents: string,
+        contents: TextDocumentContentChangeEvent[],
         ipythonMode = IPythonMode.None,
         realFilePath?: string
     ) {