Have text models capture an undo-redo snapshot when the first edit occurs and...

Have text models capture an undo-redo snapshot when the first edit occurs and on dispose restore that snapshot

src/vs/editor/common/model/textModel.ts

@@ -35,7 +35,7 @@ import { VSBufferReadableStream, VSBuffer } from 'vs/base/common/buffer';
 import { TokensStore, MultilineTokens, countEOL, MultilineTokens2, TokensStore2 } from 'vs/editor/common/model/tokensStore';
 import { Color } from 'vs/base/common/color';
 import { EditorTheme } from 'vs/editor/common/view/viewContext';
-import { IUndoRedoService } from 'vs/platform/undoRedo/common/undoRedo';
+import { IUndoRedoService, ResourceEditStackSnapshot } from 'vs/platform/undoRedo/common/undoRedo';
 import { TextChange } from 'vs/editor/common/model/textChange';
 import { Constants } from 'vs/base/common/uint';
 
@@ -278,6 +278,7 @@ export class TextModel extends Disposable implements model.ITextModel {
 	 * Unlike, versionId, this can go down (via undo) or go to previous values (via redo)
 	 */
 	private _alternativeVersionId: number;
+	private _initialUndoRedoSnapshot: ResourceEditStackSnapshot | null;
 	private readonly _isTooLargeForSyncing: boolean;
 	private readonly _isTooLargeForTokenization: boolean;
 
@@ -351,6 +352,7 @@ export class TextModel extends Disposable implements model.ITextModel {
 
 		this._versionId = 1;
 		this._alternativeVersionId = 1;
+		this._initialUndoRedoSnapshot = null;
 
 		this._isDisposed = false;
 		this._isDisposing = false;
@@ -719,6 +721,11 @@ export class TextModel extends Disposable implements model.ITextModel {
 		return this._alternativeVersionId;
 	}
 
+	public getInitialUndoRedoSnapshot(): ResourceEditStackSnapshot | null {
+		this._assertNotDisposed();
+		return this._initialUndoRedoSnapshot;
+	}
+
 	public getOffsetAt(rawPosition: IPosition): number {
 		this._assertNotDisposed();
 		let position = this._validatePosition(rawPosition.lineNumber, rawPosition.column, StringOffsetValidationType.Relaxed);
@@ -744,6 +751,10 @@ export class TextModel extends Disposable implements model.ITextModel {
 		this._alternativeVersionId = newAlternativeVersionId;
 	}
 
+	public _overwriteInitialUndoRedoSnapshot(newInitialUndoRedoSnapshot: ResourceEditStackSnapshot | null): void {
+		this._initialUndoRedoSnapshot = newInitialUndoRedoSnapshot;
+	}
+
 	public getValue(eol?: model.EndOfLinePreference, preserveBOM: boolean = false): string {
 		this._assertNotDisposed();
 		const fullModelRange = this.getFullModelRange();
@@ -1187,6 +1198,9 @@ export class TextModel extends Disposable implements model.ITextModel {
 		try {
 			this._onDidChangeDecorations.beginDeferredEmit();
 			this._eventEmitter.beginDeferredEmit();
+			if (this._initialUndoRedoSnapshot === null) {
+				this._initialUndoRedoSnapshot = this._undoRedoService.createSnapshot(this.uri);
+			}
 			this._commandManager.pushEOL(eol);
 		} finally {
 			this._eventEmitter.endDeferredEmit();
@@ -1311,6 +1325,9 @@ export class TextModel extends Disposable implements model.ITextModel {
 
 			this._trimAutoWhitespaceLines = null;
 		}
+		if (this._initialUndoRedoSnapshot === null) {
+			this._initialUndoRedoSnapshot = this._undoRedoService.createSnapshot(this.uri);
+		}
 		return this._commandManager.pushEditOperation(beforeCursorState, editOperations, cursorStateComputer);
 	}
 

src/vs/editor/common/services/modelServiceImpl.ts

@@ -24,7 +24,7 @@ import { RunOnceScheduler } from 'vs/base/common/async';
 import { CancellationTokenSource } from 'vs/base/common/cancellation';
 import { IThemeService } from 'vs/platform/theme/common/themeService';
 import { ILogService } from 'vs/platform/log/common/log';
-import { IUndoRedoService, IUndoRedoElement, IPastFutureElements } from 'vs/platform/undoRedo/common/undoRedo';
+import { IUndoRedoService, IUndoRedoElement, IPastFutureElements, ResourceEditStackSnapshot } from 'vs/platform/undoRedo/common/undoRedo';
 import { StringSHA1 } from 'vs/base/common/hash';
 import { SingleModelEditStackElement, MultiModelEditStackElement, EditStackElement, isEditStackElement } from 'vs/editor/common/model/editStack';
 import { Schemas } from 'vs/base/common/network';
@@ -51,7 +51,7 @@ function computeModelSha1(model: ITextModel): string {
 
 
 class ModelData implements IDisposable {
-	public readonly model: ITextModel;
+	public readonly model: TextModel;
 
 	private _languageSelection: ILanguageSelection | null;
 	private _languageSelectionListener: IDisposable | null;
@@ -59,7 +59,7 @@ class ModelData implements IDisposable {
 	private readonly _modelEventListeners = new DisposableStore();
 
 	constructor(
-		model: ITextModel,
+		model: TextModel,
 		onWillDispose: (model: ITextModel) => void,
 		onDidChangeLanguage: (model: ITextModel, e: IModelLanguageChangedEvent) => void
 	) {
@@ -138,6 +138,7 @@ function isEditStackElements(elements: IUndoRedoElement[]): elements is EditStac
 class DisposedModelInfo {
 	constructor(
 		public readonly uri: URI,
+		public readonly initialUndoRedoSnapshot: ResourceEditStackSnapshot | null,
 		public readonly time: number,
 		public readonly sharesUndoRedoStack: boolean,
 		public readonly heapSize: number,
@@ -362,7 +363,9 @@ export class ModelServiceImpl extends Disposable implements IModelService {
 			while (disposedModels.length > 0 && this._disposedModelsHeapSize > maxModelsHeapSize) {
 				const disposedModel = disposedModels.shift()!;
 				this._removeDisposedModel(disposedModel.uri);
-				this._undoRedoService.removeElements(disposedModel.uri);
+				if (disposedModel.initialUndoRedoSnapshot !== null) {
+					this._undoRedoService.restoreSnapshot(disposedModel.initialUndoRedoSnapshot);
+				}
 			}
 		}
 	}
@@ -390,9 +393,12 @@ export class ModelServiceImpl extends Disposable implements IModelService {
 				if (sha1IsEqual) {
 					model._overwriteVersionId(disposedModelData.versionId);
 					model._overwriteAlternativeVersionId(disposedModelData.alternativeVersionId);
+					model._overwriteInitialUndoRedoSnapshot(disposedModelData.initialUndoRedoSnapshot);
 				}
 			} else {
-				this._undoRedoService.removeElements(resource);
+				if (disposedModelData.initialUndoRedoSnapshot !== null) {
+					this._undoRedoService.restoreSnapshot(disposedModelData.initialUndoRedoSnapshot);
+				}
 			}
 		}
 		const modelId = MODEL_ID(model.uri);
@@ -541,7 +547,10 @@ export class ModelServiceImpl extends Disposable implements IModelService {
 
 		if (!maintainUndoRedoStack) {
 			if (!sharesUndoRedoStack) {
-				this._undoRedoService.removeElements(resource);
+				const initialUndoRedoSnapshot = modelData.model.getInitialUndoRedoSnapshot();
+				if (initialUndoRedoSnapshot !== null) {
+					this._undoRedoService.restoreSnapshot(initialUndoRedoSnapshot);
+				}
 			}
 			modelData.model.dispose();
 			return;
@@ -550,7 +559,10 @@ export class ModelServiceImpl extends Disposable implements IModelService {
 		const maxMemory = ModelServiceImpl.MAX_MEMORY_FOR_CLOSED_FILES_UNDO_STACK;
 		if (!sharesUndoRedoStack && heapSize > maxMemory) {
 			// the undo stack for this file would never fit in the configured memory, so don't bother with it.
-			this._undoRedoService.removeElements(resource);
+			const initialUndoRedoSnapshot = modelData.model.getInitialUndoRedoSnapshot();
+			if (initialUndoRedoSnapshot !== null) {
+				this._undoRedoService.restoreSnapshot(initialUndoRedoSnapshot);
+			}
 			modelData.model.dispose();
 			return;
 		}
@@ -559,7 +571,7 @@ export class ModelServiceImpl extends Disposable implements IModelService {
 
 		// We only invalidate the elements, but they remain in the undo-redo service.
 		this._undoRedoService.setElementsValidFlag(resource, false, (element) => (isEditStackElement(element) && element.matchesResource(resource)));
-		this._insertDisposedModel(new DisposedModelInfo(resource, Date.now(), sharesUndoRedoStack, heapSize, computeModelSha1(model), model.getVersionId(), model.getAlternativeVersionId()));
+		this._insertDisposedModel(new DisposedModelInfo(resource, modelData.model.getInitialUndoRedoSnapshot(), Date.now(), sharesUndoRedoStack, heapSize, computeModelSha1(model), model.getVersionId(), model.getAlternativeVersionId()));
 
 		modelData.model.dispose();
 	}

src/vs/platform/undoRedo/common/undoRedo.ts

@@ -52,6 +52,13 @@ export interface UriComparisonKeyComputer {
 	getComparisonKey(uri: URI): string | null;
 }
 
+export class ResourceEditStackSnapshot {
+	constructor(
+		public readonly resource: URI,
+		public readonly elements: number[]
+	) { }
+}
+
 export interface IUndoRedoService {
 	readonly _serviceBrand: undefined;
 
@@ -81,6 +88,9 @@ export interface IUndoRedoService {
 	 */
 	removeElements(resource: URI): void;
 
+	createSnapshot(resource: URI): ResourceEditStackSnapshot;
+	restoreSnapshot(snapshot: ResourceEditStackSnapshot): void;
+
 	canUndo(resource: URI): boolean;
 	undo(resource: URI): Promise<void> | void;
 

src/vs/platform/undoRedo/common/undoRedoService.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import * as nls from 'vs/nls';
-import { IUndoRedoService, IWorkspaceUndoRedoElement, UndoRedoElementType, IUndoRedoElement, IPastFutureElements, UriComparisonKeyComputer } from 'vs/platform/undoRedo/common/undoRedo';
+import { IUndoRedoService, IWorkspaceUndoRedoElement, UndoRedoElementType, IUndoRedoElement, IPastFutureElements, UriComparisonKeyComputer, ResourceEditStackSnapshot } from 'vs/platform/undoRedo/common/undoRedo';
 import { URI } from 'vs/base/common/uri';
 import { onUnexpectedError } from 'vs/base/common/errors';
 import { registerSingleton } from 'vs/platform/instantiation/common/extensions';
@@ -20,7 +20,10 @@ function getResourceLabel(resource: URI): string {
 	return resource.scheme === Schemas.file ? resource.fsPath : resource.path;
 }
 
+let stackElementCounter = 0;
+
 class ResourceStackElement {
+	public readonly id = (++stackElementCounter);
 	public readonly type = UndoRedoElementType.Resource;
 	public readonly actual: IUndoRedoElement;
 	public readonly label: string;
@@ -46,7 +49,7 @@ class ResourceStackElement {
 	}
 
 	public toString(): string {
-		return `[VALID] ${this.actual}`;
+		return `[${this.id}] [${this.isValid ? 'VALID' : 'INVALID'}] ${this.actual}`;
 	}
 }
 
@@ -105,6 +108,7 @@ class RemovedResources {
 }
 
 class WorkspaceStackElement {
+	public readonly id = (++stackElementCounter);
 	public readonly type = UndoRedoElementType.Workspace;
 	public readonly actual: IWorkspaceUndoRedoElement;
 	public readonly label: string;
@@ -151,7 +155,7 @@ class WorkspaceStackElement {
 	}
 
 	public toString(): string {
-		return `[VALID] ${this.actual}`;
+		return `[${this.id}] [${this.invalidatedResources ? 'INVALID' : 'VALID'}] ${this.actual}`;
 	}
 }
 
@@ -263,6 +267,54 @@ class ResourceEditStack {
 		this.versionId++;
 	}
 
+	public createSnapshot(resource: URI): ResourceEditStackSnapshot {
+		const elements: number[] = [];
+
+		for (let i = 0, len = this._past.length; i < len; i++) {
+			elements.push(this._past[i].id);
+		}
+		for (let i = this._future.length - 1; i >= 0; i--) {
+			elements.push(this._future[i].id);
+		}
+
+		return new ResourceEditStackSnapshot(resource, elements);
+	}
+
+	public restoreSnapshot(snapshot: ResourceEditStackSnapshot): void {
+		const snapshotLength = snapshot.elements.length;
+		let isOK = true;
+		let snapshotIndex = 0;
+		let removePastAfter = -1;
+		for (let i = 0, len = this._past.length; i < len; i++, snapshotIndex++) {
+			const element = this._past[i];
+			if (isOK && (snapshotIndex >= snapshotLength || element.id !== snapshot.elements[snapshotIndex])) {
+				isOK = false;
+				removePastAfter = 0;
+			}
+			if (!isOK && element.type === UndoRedoElementType.Workspace) {
+				element.removeResource(this.resourceLabel, this.strResource, RemovedResourceReason.ExternalRemoval);
+			}
+		}
+		let removeFutureBefore = -1;
+		for (let i = this._future.length - 1; i >= 0; i--, snapshotIndex++) {
+			const element = this._future[i];
+			if (isOK && (snapshotIndex >= snapshotLength || element.id !== snapshot.elements[snapshotIndex])) {
+				isOK = false;
+				removeFutureBefore = i;
+			}
+			if (!isOK && element.type === UndoRedoElementType.Workspace) {
+				element.removeResource(this.resourceLabel, this.strResource, RemovedResourceReason.ExternalRemoval);
+			}
+		}
+		if (removePastAfter !== -1) {
+			this._past = this._past.slice(0, removePastAfter);
+		}
+		if (removeFutureBefore !== -1) {
+			this._future = this._future.slice(removeFutureBefore + 1);
+		}
+		this.versionId++;
+	}
+
 	public getElements(): IPastFutureElements {
 		const past: IUndoRedoElement[] = [];
 		const future: IUndoRedoElement[] = [];
@@ -550,6 +602,32 @@ export class UndoRedoService implements IUndoRedoService {
 		return false;
 	}
 
+	public createSnapshot(resource: URI): ResourceEditStackSnapshot {
+		const strResource = this.getUriComparisonKey(resource);
+		if (this._editStacks.has(strResource)) {
+			const editStack = this._editStacks.get(strResource)!;
+			return editStack.createSnapshot(resource);
+		}
+		return new ResourceEditStackSnapshot(resource, []);
+	}
+
+	public restoreSnapshot(snapshot: ResourceEditStackSnapshot): void {
+		const strResource = this.getUriComparisonKey(snapshot.resource);
+		if (this._editStacks.has(strResource)) {
+			const editStack = this._editStacks.get(strResource)!;
+			editStack.restoreSnapshot(snapshot);
+
+			if (!editStack.hasPastElements() && !editStack.hasFutureElements()) {
+				// the edit stack is now empty, just remove it entirely
+				editStack.dispose();
+				this._editStacks.delete(strResource);
+			}
+		}
+		if (DEBUG) {
+			this._print('restoreSnapshot');
+		}
+	}
+
 	public getElements(resource: URI): IPastFutureElements {
 		const strResource = this.getUriComparisonKey(resource);
 		if (this._editStacks.has(strResource)) {