From b3cc19b819d2d4bf79c6b95d043acffcb04255e5 Mon Sep 17 00:00:00 2001
From: Johannes Rieken <johannes.rieken@gmail.com>
Date: Tue, 10 Nov 2020 09:45:44 +0100
Subject: [PATCH] :lipstick: use async-await for code lens provider logic

---
 .../api/browser/mainThreadLanguageFeatures.ts | 19 ++---
 .../api/common/extHostLanguageFeatures.ts     | 84 +++++++++----------
 2 files changed, 47 insertions(+), 56 deletions(-)

diff --git a/src/vs/workbench/api/browser/mainThreadLanguageFeatures.ts b/src/vs/workbench/api/browser/mainThreadLanguageFeatures.ts
index 64a3f19a3be..33824992142 100644
--- a/src/vs/workbench/api/browser/mainThreadLanguageFeatures.ts
+++ b/src/vs/workbench/api/browser/mainThreadLanguageFeatures.ts
@@ -165,16 +165,15 @@ export class MainThreadLanguageFeatures implements MainThreadLanguageFeaturesSha
 	$registerCodeLensSupport(handle: number, selector: IDocumentFilterDto[], eventHandle: number | undefined): void {
 
 		const provider = <modes.CodeLensProvider>{
-			provideCodeLenses: (model: ITextModel, token: CancellationToken): Promise<modes.CodeLensList | undefined> => {
-				return this._proxy.$provideCodeLenses(handle, model.uri, token).then(listDto => {
-					if (!listDto) {
-						return undefined;
-					}
-					return {
-						lenses: listDto.lenses,
-						dispose: () => listDto.cacheId && this._proxy.$releaseCodeLenses(handle, listDto.cacheId)
-					};
-				});
+			provideCodeLenses: async (model: ITextModel, token: CancellationToken): Promise<modes.CodeLensList | undefined> => {
+				const listDto = await this._proxy.$provideCodeLenses(handle, model.uri, token);
+				if (!listDto) {
+					return undefined;
+				}
+				return {
+					lenses: listDto.lenses,
+					dispose: () => listDto.cacheId && this._proxy.$releaseCodeLenses(handle, listDto.cacheId)
+				};
 			},
 			resolveCodeLens: (_model: ITextModel, codeLens: modes.CodeLens, token: CancellationToken): Promise<modes.CodeLens | undefined> => {
 				return this._proxy.$resolveCodeLens(handle, codeLens, token);
diff --git a/src/vs/workbench/api/common/extHostLanguageFeatures.ts b/src/vs/workbench/api/common/extHostLanguageFeatures.ts
index e4fab08b3fb..77c0b391594 100644
--- a/src/vs/workbench/api/common/extHostLanguageFeatures.ts
+++ b/src/vs/workbench/api/common/extHostLanguageFeatures.ts
@@ -116,65 +116,57 @@ class CodeLensAdapter {
 		private readonly _provider: vscode.CodeLensProvider
 	) { }
 
-	provideCodeLenses(resource: URI, token: CancellationToken): Promise<extHostProtocol.ICodeLensListDto | undefined> {
+	async provideCodeLenses(resource: URI, token: CancellationToken): Promise<extHostProtocol.ICodeLensListDto | undefined> {
 		const doc = this._documents.getDocument(resource);
 
-		return asPromise(() => this._provider.provideCodeLenses(doc, token)).then(lenses => {
-
-			if (!lenses || token.isCancellationRequested) {
-				return undefined;
-			}
-
-			const cacheId = this._cache.add(lenses);
-			const disposables = new DisposableStore();
-			this._disposables.set(cacheId, disposables);
-
-			const result: extHostProtocol.ICodeLensListDto = {
-				cacheId,
-				lenses: [],
-			};
-
-			for (let i = 0; i < lenses.length; i++) {
-				result.lenses.push({
-					cacheId: [cacheId, i],
-					range: typeConvert.Range.from(lenses[i].range),
-					command: this._commands.toInternal(lenses[i].command, disposables)
-				});
-			}
-
-			return result;
-		});
+		const lenses = await this._provider.provideCodeLenses(doc, token);
+		if (!lenses || token.isCancellationRequested) {
+			return undefined;
+		}
+		const cacheId = this._cache.add(lenses);
+		const disposables = new DisposableStore();
+		this._disposables.set(cacheId, disposables);
+		const result: extHostProtocol.ICodeLensListDto = {
+			cacheId,
+			lenses: [],
+		};
+		for (let i = 0; i < lenses.length; i++) {
+			result.lenses.push({
+				cacheId: [cacheId, i],
+				range: typeConvert.Range.from(lenses[i].range),
+				command: this._commands.toInternal(lenses[i].command, disposables)
+			});
+		}
+		return result;
 	}
 
-	resolveCodeLens(symbol: extHostProtocol.ICodeLensDto, token: CancellationToken): Promise<extHostProtocol.ICodeLensDto | undefined> {
+	async resolveCodeLens(symbol: extHostProtocol.ICodeLensDto, token: CancellationToken): Promise<extHostProtocol.ICodeLensDto | undefined> {
 
 		const lens = symbol.cacheId && this._cache.get(...symbol.cacheId);
 		if (!lens) {
-			return Promise.resolve(undefined);
+			return undefined;
 		}
 
-		let resolve: Promise<vscode.CodeLens | undefined | null>;
+		let resolvedLens: vscode.CodeLens | undefined | null;
 		if (typeof this._provider.resolveCodeLens !== 'function' || lens.isResolved) {
-			resolve = Promise.resolve(lens);
+			resolvedLens = lens;
 		} else {
-			resolve = asPromise(() => this._provider.resolveCodeLens!(lens, token));
+			resolvedLens = await this._provider.resolveCodeLens(lens, token);
+		}
+		if (!resolvedLens) {
+			resolvedLens = lens;
 		}
 
-		return resolve.then(newLens => {
-			if (token.isCancellationRequested) {
-				return undefined;
-			}
-
-			const disposables = symbol.cacheId && this._disposables.get(symbol.cacheId[0]);
-			if (!disposables) {
-				// We've already been disposed of
-				return undefined;
-			}
-
-			newLens = newLens || lens;
-			symbol.command = this._commands.toInternal(newLens.command || CodeLensAdapter._badCmd, disposables);
-			return symbol;
-		});
+		if (token.isCancellationRequested) {
+			return undefined;
+		}
+		const disposables = symbol.cacheId && this._disposables.get(symbol.cacheId[0]);
+		if (!disposables) {
+			// disposed in the meantime
+			return undefined;
+		}
+		symbol.command = this._commands.toInternal(resolvedLens.command ?? CodeLensAdapter._badCmd, disposables);
+		return symbol;
 	}
 
 	releaseCodeLenses(cachedId: number): void {