Eagerly delete diagnostic on quick fix (#166417)

Fixes #166415

extensions/typescript-language-features/src/languageFeatures/diagnostics.ts

@@ -34,6 +34,7 @@ export const enum DiagnosticKind {
 }
 
 class FileDiagnostics {
+
 	private readonly _diagnostics = new Map<DiagnosticKind, ReadonlyArray<vscode.Diagnostic>>();
 
 	constructor(
@@ -61,7 +62,7 @@ class FileDiagnostics {
 		return true;
 	}
 
-	public getDiagnostics(settings: DiagnosticSettings): vscode.Diagnostic[] {
+	public getAllDiagnostics(settings: DiagnosticSettings): vscode.Diagnostic[] {
 		if (!settings.getValidate(this.language)) {
 			return [];
 		}
@@ -73,6 +74,12 @@ class FileDiagnostics {
 		];
 	}
 
+	public delete(toDelete: vscode.Diagnostic): void {
+		for (const [type, diags] of this._diagnostics) {
+			this._diagnostics.set(type, diags.filter(diag => !diagnosticsEquals(diag, toDelete)));
+		}
+	}
+
 	private getSuggestionDiagnostics(settings: DiagnosticSettings) {
 		const enableSuggestions = settings.getEnableSuggestions(this.language);
 		return this.get(DiagnosticKind.Suggestion).filter(x => {
@@ -177,14 +184,14 @@ export class DiagnosticsManager extends Disposable {
 	public setValidate(language: DiagnosticLanguage, value: boolean) {
 		const didUpdate = this._settings.setValidate(language, value);
 		if (didUpdate) {
-			this.rebuild();
+			this.rebuildAll();
 		}
 	}
 
 	public setEnableSuggestions(language: DiagnosticLanguage, value: boolean) {
 		const didUpdate = this._settings.setEnableSuggestions(language, value);
 		if (didUpdate) {
-			this.rebuild();
+			this.rebuildAll();
 		}
 	}
 
@@ -217,11 +224,19 @@ export class DiagnosticsManager extends Disposable {
 		this._currentDiagnostics.set(file, diagnostics);
 	}
 
-	public delete(resource: vscode.Uri): void {
+	public deleteAllDiagnosticsInFile(resource: vscode.Uri): void {
 		this._currentDiagnostics.delete(resource);
 		this._diagnostics.delete(resource);
 	}
 
+	public deleteDiagnostic(resource: vscode.Uri, diagnostic: vscode.Diagnostic): void {
+		const fileDiagnostics = this._diagnostics.get(resource);
+		if (fileDiagnostics) {
+			fileDiagnostics.delete(diagnostic);
+			this.rebuildFile(fileDiagnostics);
+		}
+	}
+
 	public getDiagnostics(file: vscode.Uri): ReadonlyArray<vscode.Diagnostic> {
 		return this._currentDiagnostics.get(file) || [];
 	}
@@ -239,13 +254,17 @@ export class DiagnosticsManager extends Disposable {
 		}
 
 		const fileDiagnostics = this._diagnostics.get(file);
-		this._currentDiagnostics.set(file, fileDiagnostics ? fileDiagnostics.getDiagnostics(this._settings) : []);
+		this._currentDiagnostics.set(file, fileDiagnostics ? fileDiagnostics.getAllDiagnostics(this._settings) : []);
 	}
 
-	private rebuild(): void {
+	private rebuildAll(): void {
 		this._currentDiagnostics.clear();
 		for (const fileDiagnostic of this._diagnostics.values) {
-			this._currentDiagnostics.set(fileDiagnostic.file, fileDiagnostic.getDiagnostics(this._settings));
+			this.rebuildFile(fileDiagnostic);
 		}
 	}
+
+	private rebuildFile(fileDiagnostic: FileDiagnostics) {
+		this._currentDiagnostics.set(fileDiagnostic.file, fileDiagnostic.getAllDiagnostics(this._settings));
+	}
 }

extensions/typescript-language-features/src/languageFeatures/quickFix.ts

@@ -20,6 +20,11 @@ import * as typeConverters from '../utils/typeConverters';
 import { DiagnosticsManager } from './diagnostics';
 import FileConfigurationManager from './fileConfigurationManager';
 
+type ApplyCodeActionCommand_args = {
+	readonly resource: vscode.Uri;
+	readonly diagnostic: vscode.Diagnostic;
+	readonly action: Proto.CodeFixAction;
+};
 
 class ApplyCodeActionCommand implements Command {
 	public static readonly ID = '_typescript.applyCodeActionCommand';
@@ -27,12 +32,11 @@ class ApplyCodeActionCommand implements Command {
 
 	constructor(
 		private readonly client: ITypeScriptServiceClient,
+		private readonly diagnosticManager: DiagnosticsManager,
 		private readonly telemetryReporter: TelemetryReporter,
 	) { }
 
-	public async execute(
-		action: Proto.CodeFixAction
-	): Promise<boolean> {
+	public async execute({ resource, action, diagnostic }: ApplyCodeActionCommand_args): Promise<boolean> {
 		/* __GDPR__
 			"quickFix.execute" : {
 				"owner": "mjbvz",
@@ -46,6 +50,7 @@ class ApplyCodeActionCommand implements Command {
 			fixName: action.fixName
 		});
 
+		this.diagnosticManager.deleteDiagnostic(resource, diagnostic);
 		return applyCodeActionCommands(this.client, action.commands, nulToken);
 	}
 }
@@ -212,7 +217,7 @@ class TypeScriptQuickFixProvider implements vscode.CodeActionProvider<VsCodeCode
 		private readonly diagnosticsManager: DiagnosticsManager,
 		telemetryReporter: TelemetryReporter
 	) {
-		commandManager.register(new ApplyCodeActionCommand(client, telemetryReporter));
+		commandManager.register(new ApplyCodeActionCommand(client, diagnosticsManager, telemetryReporter));
 		commandManager.register(new ApplyFixAllCodeAction(client, telemetryReporter));
 
 		this.supportedCodeActionProvider = new SupportedCodeActionProvider(client);
@@ -223,26 +228,32 @@ class TypeScriptQuickFixProvider implements vscode.CodeActionProvider<VsCodeCode
 		_range: vscode.Range,
 		context: vscode.CodeActionContext,
 		token: vscode.CancellationToken
-	): Promise<VsCodeCodeAction[]> {
+	): Promise<VsCodeCodeAction[] | undefined> {
 		const file = this.client.toOpenedFilePath(document);
 		if (!file) {
-			return [];
+			return;
 		}
 
 		const fixableDiagnostics = await this.supportedCodeActionProvider.getFixableDiagnosticsForContext(context);
-		if (!fixableDiagnostics.size) {
-			return [];
+		if (!fixableDiagnostics.size || token.isCancellationRequested) {
+			return;
 		}
 
 		if (this.client.bufferSyncSupport.hasPendingDiagnostics(document.uri)) {
-			return [];
+			return;
 		}
 
 		await this.formattingConfigurationManager.ensureConfigurationForDocument(document, token);
+		if (token.isCancellationRequested) {
+			return;
+		}
 
 		const results = new CodeActionSet();
 		for (const diagnostic of fixableDiagnostics.values) {
 			await this.getFixesForDiagnostic(document, file, diagnostic, results, token);
+			if (token.isCancellationRequested) {
+				return;
+			}
 		}
 
 		const allActions = Array.from(results.values);
@@ -303,12 +314,13 @@ class TypeScriptQuickFixProvider implements vscode.CodeActionProvider<VsCodeCode
 		diagnostic: vscode.Diagnostic,
 		tsAction: Proto.CodeFixAction
 	): CodeActionSet {
-		results.addAction(this.getSingleFixForTsCodeAction(diagnostic, tsAction));
-		this.addFixAllForTsCodeAction(results, document, file, diagnostic, tsAction as Proto.CodeFixAction);
+		results.addAction(this.getSingleFixForTsCodeAction(document.uri, diagnostic, tsAction));
+		this.addFixAllForTsCodeAction(results, document.uri, file, diagnostic, tsAction as Proto.CodeFixAction);
 		return results;
 	}
 
 	private getSingleFixForTsCodeAction(
+		resource: vscode.Uri,
 		diagnostic: vscode.Diagnostic,
 		tsAction: Proto.CodeFixAction
 	): VsCodeCodeAction {
@@ -317,7 +329,7 @@ class TypeScriptQuickFixProvider implements vscode.CodeActionProvider<VsCodeCode
 		codeAction.diagnostics = [diagnostic];
 		codeAction.command = {
 			command: ApplyCodeActionCommand.ID,
-			arguments: [tsAction],
+			arguments: [<ApplyCodeActionCommand_args>{ action: tsAction, diagnostic, resource }],
 			title: ''
 		};
 		return codeAction;
@@ -325,7 +337,7 @@ class TypeScriptQuickFixProvider implements vscode.CodeActionProvider<VsCodeCode
 
 	private addFixAllForTsCodeAction(
 		results: CodeActionSet,
-		document: vscode.TextDocument,
+		resource: vscode.Uri,
 		file: string,
 		diagnostic: vscode.Diagnostic,
 		tsAction: Proto.CodeFixAction,
@@ -335,7 +347,7 @@ class TypeScriptQuickFixProvider implements vscode.CodeActionProvider<VsCodeCode
 		}
 
 		// Make sure there are multiple diagnostics of the same type in the file
-		if (!this.diagnosticsManager.getDiagnostics(document.uri).some(x => {
+		if (!this.diagnosticsManager.getDiagnostics(resource).some(x => {
 			if (x === diagnostic) {
 				return false;
 			}

extensions/typescript-language-features/src/typescriptServiceClient.ts

@@ -178,7 +178,7 @@ export default class TypeScriptServiceClient extends Disposable implements IType
 		this.diagnosticsManager = new DiagnosticsManager('typescript', onCaseInsenitiveFileSystem);
 		this.bufferSyncSupport.onDelete(resource => {
 			this.cancelInflightRequestsForResource(resource);
-			this.diagnosticsManager.delete(resource);
+			this.diagnosticsManager.deleteAllDiagnosticsInFile(resource);
 		}, null, this._disposables);
 
 		this.bufferSyncSupport.onWillChange(resource => {