Validate workspace undo again after async prompt

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

@@ -376,22 +376,29 @@ export class UndoRedoService implements IUndoRedoService {
 		}
 	}
 
-	private _verifyWorkspaceUndo(element: WorkspaceStackElement): WorkspaceVerificationResult {
+	private _getAffectedEditStacks(element: WorkspaceStackElement): ResourceEditStack[] {
+		const affectedEditStacks: ResourceEditStack[] = [];
+		for (const strResource of element.strResources) {
+			affectedEditStacks.push(this._editStacks.get(strResource)!);
+		}
+		return affectedEditStacks;
+	}
+
+	private _checkWorkspaceUndo(resource: URI, element: WorkspaceStackElement, affectedEditStacks: ResourceEditStack[]): WorkspaceVerificationError | null {
 		if (element.removedResources) {
+			this._splitPastWorkspaceElement(element, element.removedResources);
 			const message = nls.localize('cannotWorkspaceUndo', "Could not undo '{0}' across all files. {1}", element.label, element.removedResources.createMessage());
-			return new InvalidWorkspaceVerificationResult(message, element.removedResources);
+			this._notificationService.info(message);
+			return new WorkspaceVerificationError(this.undo(resource));
 		}
 		if (element.invalidatedResources) {
+			this._splitPastWorkspaceElement(element, element.invalidatedResources);
 			const message = nls.localize('cannotWorkspaceUndo', "Could not undo '{0}' across all files. {1}", element.label, element.invalidatedResources.createMessage());
-			return new InvalidWorkspaceVerificationResult(message, element.invalidatedResources);
+			this._notificationService.info(message);
+			return new WorkspaceVerificationError(this.undo(resource));
 		}
 
 		// this must be the last past element in all the impacted resources!
-		const affectedEditStacks: ResourceEditStack[] = [];
-		for (const strResource of element.strResources) {
-			affectedEditStacks.push(this._editStacks.get(strResource)!);
-		}
-
 		const cannotUndoDueToResources: URI[] = [];
 		for (const editStack of affectedEditStacks) {
 			if (editStack.past.length === 0 || editStack.past[editStack.past.length - 1] !== element) {
@@ -400,29 +407,27 @@ export class UndoRedoService implements IUndoRedoService {
 		}
 
 		if (cannotUndoDueToResources.length > 0) {
+			this._splitPastWorkspaceElement(element, null);
 			const paths = cannotUndoDueToResources.map(r => r.scheme === Schemas.file ? r.fsPath : r.path);
 			const message = nls.localize('cannotWorkspaceUndoDueToChanges', "Could not undo '{0}' across all files because changes were made to {1}", element.label, paths.join(', '));
-			return new InvalidWorkspaceVerificationResult(message, null);
+			this._notificationService.info(message);
+			return new WorkspaceVerificationError(this.undo(resource));
 		}
 
-		return new ValidWorkspaceVerificationResult(affectedEditStacks);
+		return null;
 	}
 
 	private _workspaceUndo(resource: URI, element: WorkspaceStackElement): Promise<void> | void {
-		const verificationResult = this._verifyWorkspaceUndo(element);
-		if (verificationResult.type === WorkspaceVerificationResultType.Invalid) {
-			// cannot apply the workspace undo
-			this._splitPastWorkspaceElement(element, verificationResult.ignoreResources);
-			this._notificationService.info(verificationResult.message);
-			return this.undo(resource);
+		const affectedEditStacks = this._getAffectedEditStacks(element);
+		const verificationError = this._checkWorkspaceUndo(resource, element, affectedEditStacks);
+		if (verificationError) {
+			return verificationError.returnValue;
 		}
-
-		const affectedEditStacks = verificationResult.affectedEditStacks;
 		return this._confirmAndExecuteWorkspaceUndo(resource, element, affectedEditStacks);
 	}
 
 	private async _confirmAndExecuteWorkspaceUndo(resource: URI, element: WorkspaceStackElement, affectedEditStacks: ResourceEditStack[]): Promise<void> {
-		return this._dialogService.show(
+		const result = await this._dialogService.show(
 			Severity.Info,
 			nls.localize('confirmWorkspace', "Would you like to undo '{0}' across all files?", element.label),
 			[
@@ -433,21 +438,29 @@ export class UndoRedoService implements IUndoRedoService {
 			{
 				cancelId: 2
 			}
-		).then((result) => {
-			if (result.choice === 2) {
-				// cancel
-				return;
-			} else if (result.choice === 0) {
-				for (const editStack of affectedEditStacks) {
-					editStack.past.pop();
-					editStack.future.push(element);
-				}
-				return this._safeInvoke(element, () => element.actual.undo());
-			} else {
-				this._splitPastWorkspaceElement(element, null);
-				return this.undo(resource);
+		);
+
+		// At this point, it is possible that the element has been made invalid in the meantime (due to the confirmation await)
+		const verificationError = this._checkWorkspaceUndo(resource, element, affectedEditStacks);
+		if (verificationError) {
+			return verificationError.returnValue;
+		}
+
+		if (result.choice === 2) {
+			// cancel
+			return;
+		}
+
+		if (result.choice === 0) {
+			for (const editStack of affectedEditStacks) {
+				editStack.past.pop();
+				editStack.future.push(element);
 			}
-		});
+			return this._safeInvoke(element, () => element.actual.undo());
+		} else {
+			this._splitPastWorkspaceElement(element, null);
+			return this.undo(resource);
+		}
 	}
 
 	private _resourceUndo(editStack: ResourceEditStack, element: ResourceStackElement): Promise<void> | void {
@@ -490,23 +503,22 @@ export class UndoRedoService implements IUndoRedoService {
 		return false;
 	}
 
-	private _verifyWorkspaceRedo(element: WorkspaceStackElement): WorkspaceVerificationResult {
+	private _checkWorkspaceRedo(resource: URI, element: WorkspaceStackElement, affectedEditStacks: ResourceEditStack[]): WorkspaceVerificationError | null {
 		if (element.removedResources) {
+			this._splitFutureWorkspaceElement(element, element.removedResources);
 			const message = nls.localize('cannotWorkspaceRedo', "Could not redo '{0}' across all files. {1}", element.label, element.removedResources.createMessage());
-			return new InvalidWorkspaceVerificationResult(message, element.removedResources);
+			this._notificationService.info(message);
+			return new WorkspaceVerificationError(this.redo(resource));
 		}
 		if (element.invalidatedResources) {
+			this._splitFutureWorkspaceElement(element, element.invalidatedResources);
 			const message = nls.localize('cannotWorkspaceRedo', "Could not redo '{0}' across all files. {1}", element.label, element.invalidatedResources.createMessage());
-			return new InvalidWorkspaceVerificationResult(message, element.invalidatedResources);
+			this._notificationService.info(message);
+			return new WorkspaceVerificationError(this.redo(resource));
 		}
 
 		// this must be the last future element in all the impacted resources!
-		let affectedEditStacks: ResourceEditStack[] = [];
-		for (const strResource of element.strResources) {
-			affectedEditStacks.push(this._editStacks.get(strResource)!);
-		}
-
-		let cannotRedoDueToResources: URI[] = [];
+		const cannotRedoDueToResources: URI[] = [];
 		for (const editStack of affectedEditStacks) {
 			if (editStack.future.length === 0 || editStack.future[editStack.future.length - 1] !== element) {
 				cannotRedoDueToResources.push(editStack.resource);
@@ -514,24 +526,23 @@ export class UndoRedoService implements IUndoRedoService {
 		}
 
 		if (cannotRedoDueToResources.length > 0) {
+			this._splitFutureWorkspaceElement(element, null);
 			const paths = cannotRedoDueToResources.map(r => r.scheme === Schemas.file ? r.fsPath : r.path);
 			const message = nls.localize('cannotWorkspaceRedoDueToChanges', "Could not redo '{0}' across all files because changes were made to {1}", element.label, paths.join(', '));
-			return new InvalidWorkspaceVerificationResult(message, null);
+			this._notificationService.info(message);
+			return new WorkspaceVerificationError(this.redo(resource));
 		}
 
-		return new ValidWorkspaceVerificationResult(affectedEditStacks);
+		return null;
 	}
 
 	private _workspaceRedo(resource: URI, element: WorkspaceStackElement): Promise<void> | void {
-		const verificationResult = this._verifyWorkspaceRedo(element);
-		if (verificationResult.type === WorkspaceVerificationResultType.Invalid) {
-			// cannot apply the workspace redo
-			this._splitFutureWorkspaceElement(element, verificationResult.ignoreResources);
-			this._notificationService.info(verificationResult.message);
-			return this.redo(resource);
+		const affectedEditStacks = this._getAffectedEditStacks(element);
+		const verificationError = this._checkWorkspaceRedo(resource, element, affectedEditStacks);
+		if (verificationError) {
+			return verificationError.returnValue;
 		}
 
-		const affectedEditStacks = verificationResult.affectedEditStacks;
 		for (const editStack of affectedEditStacks) {
 			editStack.future.pop();
 			editStack.past.push(element);
@@ -571,21 +582,8 @@ export class UndoRedoService implements IUndoRedoService {
 	}
 }
 
-const enum WorkspaceVerificationResultType {
-	Invalid = 0,
-	Valid = 1
-}
-class InvalidWorkspaceVerificationResult {
-	public readonly type = WorkspaceVerificationResultType.Invalid;
-	constructor(
-		public readonly message: string,
-		public readonly ignoreResources: RemovedResources | null
-	) { }
-}
-class ValidWorkspaceVerificationResult {
-	public readonly type = WorkspaceVerificationResultType.Valid;
-	constructor(public readonly affectedEditStacks: ResourceEditStack[]) { }
+class WorkspaceVerificationError {
+	constructor(public readonly returnValue: Promise<void> | void) { }
 }
-type WorkspaceVerificationResult = InvalidWorkspaceVerificationResult | ValidWorkspaceVerificationResult;
 
 registerSingleton(IUndoRedoService, UndoRedoService);